import sys, json, random
from pathlib import Path
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, sweep_baseline, make_report

TRACK = 'dynamics'
MODEL = 'rnn_small'
SEEDS = tuple(range(8))
LRS = [1e-3, 3e-3, 6e-3]
EPOCHS = 12
NTRAIN, NTEST = 1000, 300

class Field(nn.Module):
    def __init__(self, d, width=64):
        super().__init__()
        self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, width), nn.GELU(), nn.Linear(width, d))
    def forward(self, x):
        return self.net(x)

class CyclicRNN(nn.Module):
    def __init__(self, out_dim=1, step=0.15, random_order=True):
        super().__init__()
        self.rnn = nn.GRU(3, 64, batch_first=True)
        self.f1, self.f2 = Field(64), Field(64)
        self.head = nn.Linear(64, out_dim)
        self.step, self.random_order = step, random_order
        self._no_cudnn = False
    def encode(self, x):
        seq = x.view(x.shape[0], -1, 3)
        try:
            _, h = self.rnn(seq)
        except RuntimeError:
            self._no_cudnn = True
        if self._no_cudnn:
            old = torch.backends.cudnn.enabled
            torch.backends.cudnn.enabled = False
            try:
                _, h = self.rnn(seq)
            finally:
                torch.backends.cudnn.enabled = old
        return h[-1]
    def cyclic(self, z, reverse=False):
        u, v = self.f1(z), self.f2(z)
        mean = (u + v) * 0.5
        u, v = u - mean, v - mean
        if not reverse:
            z = z + self.step * u
            z = z + 0.5 * self.step * (self.f2(z) - self.f1(z))
        else:
            z = z + self.step * v
            z = z + 0.5 * self.step * (self.f1(z) - self.f2(z))
        return z
    def forward(self, x):
        z = self.encode(x)
        if self.training and self.random_order:
            reverse = bool(torch.rand((), device=z.device) < 0.5)
        else:
            reverse = False
        return self.head(self.cyclic(z, reverse))

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)

def ds(seed):
    return get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST)

def base_model(d):
    # Exact canonical rnn_small architecture.
    from bench import make_model
    return make_model(MODEL, d['input_shape'], d['out_dim'])

def idea_model(d, step):
    return CyclicRNN(d['out_dim'], step=step, random_order=True)

def run_one(kind, cfg, seed):
    seed_all(seed)
    d = ds(seed)
    model = base_model(d) if kind == 'baseline' else idea_model(d, cfg['step'])
    _, metric, _ = train_model(model, d, epochs=EPOCHS, lr=cfg['lr'], batch=128, log=lambda *_: None)
    if metric is None: return float('nan')
    return float(metric)

def eval_cfg(kind, cfg, seeds=SEEDS):
    vals = [run_one(kind, cfg, int(s)) for s in seeds]
    vals = [v for v in vals if np.isfinite(v)]
    return {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals)}

def main():
    # Baseline sweep explicitly covers every lr used by the idea and its nearby settings.
    baseline_grid = [{'lr': lr, 'step': 0.0} for lr in LRS]
    base = sweep_baseline(lambda cfg: (lambda seed: run_one('baseline', cfg, seed)), baseline_grid, seeds=(0,1,2,3))
    # Re-run baseline at every union lr on all paired seeds for fair reporting.
    base_full_by_cfg = []
    for cfg in baseline_grid:
        r = eval_cfg('baseline', cfg)
        base_full_by_cfg.append({'cfg': cfg, **r})
    best_cfg = min(base_full_by_cfg, key=lambda x: x['mean'])['cfg']
    base = {'best_cfg': best_cfg, 'sweep': [{'cfg': x['cfg'], 'mean': x['mean']} for x in base_full_by_cfg], 'full': next(x for x in base_full_by_cfg if x['cfg'] == best_cfg)}
    idea_grid = [{'lr': best_cfg['lr'], 'step': 0.10}, {'lr': best_cfg['lr'], 'step': 0.15}, {'lr': best_cfg['lr'], 'step': 0.22}]
    idea_runs = [{'cfg': cfg, **eval_cfg('idea', cfg)} for cfg in idea_grid]
    idea_best = min(idea_runs, key=lambda x: x['mean'])
    idea_result = idea_best
    # Signature is measured on trained models, not an analytic toy field.
    seed = 0; seed_all(seed); d = ds(seed); model = idea_model(d, idea_best['cfg']['step'])
    model, _, _ = train_model(model, d, epochs=EPOCHS, lr=idea_best['cfg']['lr'], batch=128, log=lambda *_: None)
    model.eval(); x = d['xte'][:128]
    with torch.no_grad():
        z = model.encode(x.to(next(model.parameters()).device)); ab = model.cyclic(z, False); ba = model.cyclic(z, True)
        gap = (ab-ba).norm(dim=1).mean().item()
        e = torch.tensor([0.01, 0.02, 0.04], device=z.device)
        gaps = []
        for h in e:
            old = model.step; model.step = float(h); gaps.append((model.cyclic(z,False)-model.cyclic(z,True)).norm(dim=1).mean().item()); model.step = old
    slope = float(np.polyfit(np.log(e.cpu().numpy()), np.log(np.maximum(gaps,1e-12)), 1)[0])
    signature = {'order_gap_mean': float(gap), 'steps': e.cpu().numpy().tolist(), 'gaps': gaps, 'predicted_exponent': 2.0, 'observed_exponent': slope, 'confirmed': bool(gap > 1e-8 and 1.5 < slope < 2.5)}
    idea_report = dict(idea_result)
    idea_report['best_cfg'] = idea_best['cfg']
    idea_report['sweep'] = [{'cfg': x['cfg'], 'mean': x['mean']} for x in idea_runs]
    report = make_report(TRACK, MODEL, base, idea_report, extra=signature)
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()