import json, random, sys
import numpy as np
import torch
import torch.nn as 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)
EPOCHS = 12
NTR, NTE = 400, 200
LR_GRID = [1e-3, 3e-3, 1e-2]


def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class VariableBaseline(nn.Module):
    """Standard variable-only predictor: one learned state readout."""
    def __init__(self, width=48):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(24, width), nn.Tanh(),
                                     nn.Linear(width, width), nn.Tanh())
        self.readout = nn.Linear(width, 1)

    def forward(self, x):
        return self.readout(self.encoder(x))


class EquationAddressable(nn.Module):
    """Implicit bipartite layer with three mechanism residuals and two variables.

    f0 and f1 both address variable z0 through different learned mechanisms;
    f2 addresses z1 and couples to z0. The forward pass minimizes all residuals
    simultaneously by differentiable damped residual-gradient iterations.
    """
    def __init__(self, width=48, steps=8, step_size=0.20):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(24, width), nn.Tanh(),
                                     nn.Linear(width, width), nn.Tanh())
        self.mechanisms = nn.ModuleList([
            nn.Sequential(nn.Linear(width + 2, width), nn.Tanh(), nn.Linear(width, 1))
            for _ in range(3)
        ])
        self.steps = int(steps)
        self.step_size = float(step_size)

    def residuals(self, z, h):
        q = torch.cat([h, z], dim=1)
        a = [m(q)[:, 0] for m in self.mechanisms]
        # Bipartite incidence: f0--z0, f1--(z0,z1), f2--z1.
        return torch.stack([z[:, 0] - a[0],
                            z[:, 0] + 0.5 * z[:, 1] - a[1],
                            z[:, 1] - a[2]], dim=1)

    def solve(self, h, intervention=None, xi=None, steps=None):
        # Residual-gradient equilibrium updates need a graph even during eval.
        with torch.enable_grad():
            return self._solve_grad(h, intervention, xi, steps)

    def _solve_grad(self, h, intervention=None, xi=None, steps=None):
        z = torch.zeros(h.shape[0], 2, device=h.device, dtype=h.dtype)
        nsteps = self.steps if steps is None else int(steps)
        for _ in range(nsteps):
            z.requires_grad_(True)
            r = self.residuals(z, h)
            if intervention is not None:
                j, v = intervention
                rr = r.clone()
                rr[:, j] = z[:, v] - xi
            else:
                rr = r
            loss = 0.5 * (rr * rr).sum()
            grad = torch.autograd.grad(loss, z, create_graph=True)[0]
            z = z - self.step_size * grad
        return z

    def forward(self, x):
        h = self.encoder(x)
        return self.solve(h)[:, :1]


def train_one(kind, seed, lr, steps=8, step_size=0.20, return_model=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
    if kind == 'baseline':
        net = VariableBaseline()
    else:
        net = EquationAddressable(steps=steps, step_size=step_size)
    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128,
                                    weight_decay=0.0, log=lambda *_: None)
    if net is None:
        raise RuntimeError('benchmark training failed')
    if return_model:
        return float(metric), net, ds
    return float(metric)


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


def idea_factory(cfg):
    return lambda seed: train_one('idea', seed, float(cfg['lr']),
                                  steps=int(cfg['steps']), step_size=float(cfg['step_size']))


def signature(idea_model, baseline_model, ds):
    """Re-test the stage-1 prediction on trained models, not analytic equations."""
    device = next(idea_model.parameters()).device
    x = ds['xte'].to(device)[:64]
    with torch.enable_grad():
        h = idea_model.encoder(x)
        with torch.no_grad():
            obs = idea_model.solve(h)
        xi = obs[:, 0].median().detach()
        # Two distinct equation replacements fix the same trained variable value.
        a = idea_model.solve(h, intervention=(0, 0), xi=xi, steps=12).detach()
        b = idea_model.solve(h, intervention=(1, 0), xi=xi, steps=12).detach()
    with torch.no_grad():
        base_pred = baseline_model(x)[:, 0]
    target_a = float((a[:, 0] - xi).abs().mean())
    target_b = float((b[:, 0] - xi).abs().mean())
    downstream = float((a[:, 1] - b[:, 1]).abs().mean())
    ordinary = float((obs[:, 0] - base_pred).abs().mean())
    return {
        'prediction': 'same xi imposed by different equation replacements yields different downstream state',
        'n_samples': 64,
        'xi': float(xi),
        'trained_model_target_error_f0': target_a,
        'trained_model_target_error_f1': target_b,
        'trained_model_downstream_separation': downstream,
        'baseline_vs_observational_prediction_abs_gap': ordinary,
        'confirmed': bool(max(target_a, target_b) < 0.08 and downstream > 0.01)
    }


def main():
    # Baseline and idea use the same lr union; baseline is swept on four seeds.
    grid = [{'lr': lr} for lr in LR_GRID]
    base = sweep_baseline(baseline_factory, grid, seeds=SWEEP_SEEDS)
    best_lr = float(base['best_cfg']['lr'])
    idea_grid = [
        {'lr': best_lr, 'steps': 8, 'step_size': 0.20},
        {'lr': LR_GRID[max(0, LR_GRID.index(best_lr)-1)], 'steps': 8, 'step_size': 0.20},
        {'lr': LR_GRID[min(len(LR_GRID)-1, LR_GRID.index(best_lr)+1)], 'steps': 8, 'step_size': 0.20},
    ]
    # Deduplicate if the best is at an edge while retaining three nearby trials when possible.
    unique = []
    for c in idea_grid:
        if c not in unique: unique.append(c)
    idea_grid = unique
    idea_runs = []
    for cfg in idea_grid:
        r = evaluate(idea_factory(cfg), seeds=SEEDS)
        idea_runs.append({'cfg': cfg, 'result': r})
    chosen = min(idea_runs, key=lambda q: q['result']['mean'])
    # Refit one paired seed for the behavior signature using the selected setting.
    _, im, ds = train_one('idea', 0, chosen['cfg']['lr'], chosen['cfg']['steps'], chosen['cfg']['step_size'], True)
    _, bm, _ = train_one('baseline', 0, best_lr, return_model=True)
    sig = signature(im, bm, ds)
    report = make_report('dynamics', 'rnn_small', base, chosen['result'], {
        'mechanism_signature': sig,
        'track_justification': 'Dynamics is structurally matched: the built-in task is an actuated pendulum rollout whose target depends on coupled state-transition mechanisms.',
        'baseline_sweep_union_lr': LR_GRID,
        'idea_sweep': idea_runs,
        'idea_selected_cfg': chosen['cfg'],
        'parameter_counts': {'baseline': sum(p.numel() for p in bm.parameters()), 'idea': sum(p.numel() for p in im.parameters())}
    })
    report['comparison']['idea_sweep_results'] = idea_runs
    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()