import sys, json, math, 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, make_report, sweep_baseline, evaluate

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
LRS = [1e-3, 3e-3, 6e-3]
EPOCHS = 18
BATCH = 128
H = 64
STEP = 0.12
LAMBDA = 0.10


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    try:
        if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
    except Exception:
        pass


def tangent_project(x, q):
    p = torch.where(x <= -1.0, torch.clamp(q, min=0.0), q)
    return torch.where(x >= 1.0, torch.clamp(p, max=0.0), p)


class ResidualRNN(nn.Module):
    def __init__(self, out_dim=1):
        super().__init__()
        self.inp = nn.Linear(3, H)
        self.W = nn.Linear(H, H, bias=False)
        self.bias = nn.Parameter(torch.zeros(H))
        self.head = nn.Linear(H, out_dim)
        nn.init.normal_(self.W.weight, std=0.08 / math.sqrt(H))

    def hidden(self, x, z0=None):
        s = x.view(x.shape[0], -1, 3)
        z = torch.zeros(x.shape[0], H, device=x.device, dtype=x.dtype) if z0 is None else z0
        for t in range(s.shape[1]):
            z = z + STEP * torch.tanh(self.W(z) + self.inp(s[:, t]) + self.bias)
        return z

    def forward(self, x):
        return self.head(self.hidden(x))


class ProjectedDissipativeRNN(nn.Module):
    def __init__(self, out_dim=1):
        super().__init__()
        self.inp = nn.Linear(3, H)
        self.L = nn.Parameter(torch.randn(H, H) * (0.08 / math.sqrt(H)))
        self.bias = nn.Parameter(torch.zeros(H))
        self.head = nn.Linear(H, out_dim)

    def matrix(self):
        return self.L.T @ self.L + LAMBDA * torch.eye(H, device=self.L.device, dtype=self.L.dtype)

    def hidden(self, x, z0=None):
        s = x.view(x.shape[0], -1, 3)
        z = torch.zeros(x.shape[0], H, device=x.device, dtype=x.dtype) if z0 is None else z0
        M = self.matrix()
        for t in range(s.shape[1]):
            q = -(z @ M.T) + self.inp(s[:, t]) + self.bias
            z = torch.clamp(z + STEP * tangent_project(z, q), -1.0, 1.0)
        return z

    def forward(self, x):
        return self.head(self.hidden(x))


def train_one(kind, seed, lr, keep=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    model = ResidualRNN(ds['out_dim']) if kind == 'baseline' else ProjectedDissipativeRNN(ds['out_dim'])
    net, metric, history = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    if net is None:
        raise RuntimeError('bench training failed')
    return (float(metric), net, ds) if keep else float(metric)


def baseline_factory(cfg):
    return lambda seed: train_one('baseline', seed, cfg['lr'])


def idea_factory(cfg):
    return lambda seed: train_one('idea', seed, cfg['lr'])


def idea_sweep():
    tried = []
    for cfg in [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LRS]:
        r = evaluate(idea_factory(cfg), SWEEP_SEEDS)
        tried.append({'cfg': cfg, 'mean': r['mean']})
    best = min(tried, key=lambda a: a['mean'])
    return {'best_cfg': best['cfg'], 'sweep': tried}


def signature(base, idea):
    ratios_b, ratios_i, boxes = [], [], []
    for (mb, db), (mi, di) in zip(base, idea):
        x = db['xte'][:32]
        device = next(mb.parameters()).device
        x = x.to(device)
        zb = torch.zeros(x.shape[0], H, device=device)
        zi = torch.zeros(x.shape[0], H, device=device)
        with torch.no_grad():
            delta = 1e-3 * torch.randn_like(zb)
            rb = (mb.hidden(x, zb + delta) - mb.hidden(x, zb)).norm(dim=1) / (delta.norm(dim=1) + 1e-12)
            ri = (mi.hidden(x, zi + delta) - mi.hidden(x, zi)).norm(dim=1) / (delta.norm(dim=1) + 1e-12)
            ratios_b.extend(rb.cpu().numpy().tolist()); ratios_i.extend(ri.cpu().numpy().tolist())
            boxes.append(float(mi.hidden(x).abs().max().cpu()))
    mean_b, mean_i = float(np.mean(ratios_b)), float(np.mean(ratios_i))
    return {'prediction': 'projected dissipative dynamics reduces initial-state perturbation amplification and keeps hidden states in [-1,1]', 'baseline_perturbation_ratio': mean_b, 'idea_perturbation_ratio': mean_i, 'baseline_to_idea_ratio': mean_i / (mean_b + 1e-12), 'idea_max_abs_hidden': float(max(boxes)), 'confirmed': bool(mean_i < mean_b and max(boxes) <= 1.00001)}


def main():
    grid = [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LRS]
    base = sweep_baseline(baseline_factory, grid, seeds=SWEEP_SEEDS)
    isweep = idea_sweep()
    bcfg, icfg = base['best_cfg'], isweep['best_cfg']
    bfull = evaluate(baseline_factory(bcfg), SEEDS)
    if icfg['lr'] == bcfg['lr']:
        ifull = evaluate(idea_factory(icfg), SEEDS)
    else:
        ifull = evaluate(idea_factory(icfg), SEEDS)
    bmodels = [train_one('baseline', s, bcfg['lr'], keep=True)[1:] for s in SEEDS]
    imodels = [train_one('idea', s, icfg['lr'], keep=True)[1:] for s in SEEDS]
    bblock = {'sweep': base['sweep'], 'best_config': bcfg, 'full': bfull}
    rep = make_report('dynamics', 'rnn_small_matched_residual', bblock, {'config': icfg, **ifull}, {'mechanism_signature': signature(bmodels, imodels), 'track_rationale': 'Dynamics is the built-in structural match for stability/control and Lyapunov-style contraction.'})
    rep['idea_sweep'] = isweep
    rep['parameterization'] = {'hidden': H, 'step': STEP, 'lambda': LAMBDA, 'same_input_and_head': True, 'lr_union_tested_on_both': LRS}
    with open('bench_report.json', 'w') as f: json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))

if __name__ == '__main__':
    main()