import sys, json, random
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 24
BATCH = 128
H = 4.0
ALPHA = 0.5


def math_check():
    hs = np.array([0.5, 1., 2., 4., 8., 32.])
    q = hs / (1.0 + ALPHA * hs)
    # Contractive scalar test x'=-x: amplification is |1-step|.
    euler = np.abs(1.0 - hs)
    damped = np.abs(1.0 - q)
    return {'h': hs.tolist(), 'q': q.tolist(),
            'q_bound': 1.0 / ALPHA,
            'euler_amplification': euler.tolist(),
            'damped_amplification': damped.tolist(),
            'predicted_bounded_effective_step': bool(np.all(q < 1.0 / ALPHA)),
            'damped_nonexpansive_test': bool(np.all(damped <= 1.0 + 1e-12)),
            'euler_nonexpansive_test': bool(np.all(euler <= 1.0 + 1e-12))}


class ResidualDynamics(nn.Module):
    """Eight-step residual state model; only the integration rule differs."""
    def __init__(self, mode, hidden=64, nominal_h=H, alpha=ALPHA):
        super().__init__()
        self.mode, self.h, self.alpha = mode, nominal_h, alpha
        self.inp = nn.Linear(3, hidden)
        self.field = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(),
                                   nn.Linear(hidden, hidden))
        self.head = nn.Linear(hidden, 1)
        self.last = {}

    def vector_field(self, z, u):
        # Shared feature map; input injection is identical in both systems.
        return self.field(z + self.inp(u))

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        z = torch.zeros(x.shape[0], self.inp.out_features, device=x.device, dtype=x.dtype)
        q = self.h / (1.0 + self.alpha * self.h)
        max_norm, max_update, max_resid = 0., 0., 0.
        for k in range(seq.shape[1]):
            u = seq[:, k]
            if self.mode == 'baseline':
                f = self.vector_field(z, u)
                zn = z + self.h * f
                resid = torch.zeros((), device=z.device)
            else:
                f1 = self.vector_field(z, u)
                mid = z + 0.5 * q * f1
                f2 = self.vector_field(mid, u)
                zn = z + q * f2
                resid = (zn - (z + q * f1)).norm(dim=1).mean()
            upd = (zn - z).norm(dim=1).mean()
            max_norm = max(max_norm, float(zn.detach().norm(dim=1).max()))
            max_update = max(max_update, float(upd.detach()))
            max_resid = max(max_resid, float(resid.detach()))
            z = zn
        self.last = {'max_activation': max_norm, 'max_update': max_update,
                     'stage_residual': max_resid,
                     'finite': bool(torch.isfinite(z).all().item())}
        return self.head(z)


def factory(mode, lr, seed):
    # Explicit deterministic pairing: same initialization for both systems.
    torch.manual_seed(1000 + int(seed))
    np.random.seed(1000 + int(seed)); random.seed(1000 + int(seed))
    ds = get_dataset('dynamics', seed=int(seed), n_train=400, n_test=200)
    model = ResidualDynamics(mode)
    net, metric, history = train_model(model, ds, epochs=EPOCHS, lr=lr,
                                       batch=BATCH, weight_decay=0.0,
                                       log=lambda *_: None)
    if net is None or metric is None:
        return float('inf'), {'failed': True}
    stats = dict(net.last)
    stats['test_metric'] = float(metric)
    stats['nan'] = not bool(stats.pop('finite', False))
    return float(metric), stats


def make_fn(mode, cfg):
    def run(seed):
        val, stats = factory(mode, cfg['lr'], seed)
        RUN_STATS.setdefault(mode, {}).setdefault(str(cfg['lr']), {})[str(seed)] = stats
        return val
    return run


def signature(base_res, idea_res):
    b = [v for v in base_res.values() if v]
    i = [v for v in idea_res.values() if v]
    br = float(np.mean([x['max_update'] for x in b])) if b else float('nan')
    ir = float(np.mean([x['max_update'] for x in i])) if i else float('nan')
    # Prediction is damping of the actual per-step update at the trained NN scale.
    ratio = ir / br if br > 0 else float('nan')
    return {'prediction': 'denominator block has smaller trained-state update proxy at h=4',
            'nominal_h': H, 'alpha': ALPHA, 'observed_baseline_update': br,
            'observed_idea_update': ir, 'observed_update_ratio': ratio,
            'predicted_upper_ratio_from_q_over_h': (H/(1+ALPHA*H))/H,
            'confirmed': bool(np.isfinite(ratio) and ratio < 0.85)}


if __name__ == '__main__':
    RUN_STATS = {}
    check = math_check()
    grid = [{'lr': x} for x in LRS]
    base = sweep_baseline(lambda cfg: make_fn('baseline', cfg), grid, seeds=SWEEP_SEEDS)
    # Evaluate all shared learning rates on the idea side; the best is selected only
    # after the same-sized, parity-complete search.
    idea_trials = []
    for cfg in grid:
        r = evaluate(make_fn('idea', cfg), SEEDS)
        idea_trials.append({'cfg': cfg, 'result': r})
    best = min(idea_trials, key=lambda a: a['result']['mean'])
    rep = make_report('dynamics', 'residual_rnn_small', base, best['result'], {
        'math_check': check,
        'idea_config': best['cfg'],
        'idea_sweep': idea_trials,
        'mechanism_signature': signature(
            RUN_STATS.get('baseline', {}).get(str(base['best_cfg']['lr']), {}),
            RUN_STATS.get('idea', {}).get(str(best['cfg']['lr']), {}))})
    with open('bench_report.json', 'w') as f:
        json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))