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, make_model, train_model, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
LR_GRID = [1e-3, 3e-3, 1e-2]
WD_GRID = [0.0, 1e-4]
MU = 0.30
LF = 0.20
EPOCHS = 30


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


class DissipativeRNN(nn.Module):
    """Same rnn_small core; only output parameterization is changed.

    The last action u is the action-coordinate analogue of a diffusion state.
    The residual uses a state/history-dependent center and direct tanh(u), so
    |f(u)-f(u')| <= LF |u-u'| exactly (for fixed observed history).
    """
    def __init__(self, input_shape, out_dim, mu=MU, lf=LF):
        super().__init__()
        self.core = make_model('rnn_small', input_shape, out_dim)
        self.mu, self.lf = float(mu), float(lf)

    def forward(self, x):
        # core prediction supplies a learned history/state-dependent center;
        # remove the current action from the core input before forming it.
        seq = x.view(x.shape[0], -1, 3)
        u = seq[:, -1:, 2:3]
        masked = seq.clone(); masked[:, -1, 2] = 0.0
        center = self.core(masked.reshape(x.shape[0], -1))
        return -self.mu * u.reshape(x.shape[0], 1) + self.lf * torch.tanh(center + u.reshape(x.shape[0], 1))


def train_one(kind, seed, lr, wd=0.0, return_model=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=4000, n_test=1000)
    if kind == 'baseline':
        model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    else:
        model = DissipativeRNN(ds['input_shape'], ds['out_dim'])
    trained, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128,
                                      weight_decay=wd, log=lambda *_: None)
    if return_model:
        return float(metric), trained, ds
    return float(metric)


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


def idea_eval(cfg):
    return evaluate(lambda seed: train_one('idea', seed, cfg['lr'], cfg['weight_decay']), SEEDS)


def math_check():
    # Independent numerical sanity check of the exact scalar inequality.
    rng = np.random.RandomState(123)
    u = rng.randn(200000); up = rng.randn(200000); c = rng.randn(200000)
    f = LF*np.tanh(c+u); fp = LF*np.tanh(c+up)
    b = -MU*u+f; bp = -MU*up+fp
    da = u-up; ratio = (b-bp)*da/(da*da+1e-12)
    lip = np.abs(f-fp)/(np.abs(da)+1e-12)
    return {'mu': MU, 'Lf': LF, 'theoretical_bound': -(MU-LF),
            'max_residual_fd_lipschitz': float(lip.max()),
            'max_one_sided_ratio': float(ratio.max()),
            'violation_fraction': float(np.mean(ratio > -(MU-LF)+1e-10)),
            'confirmed': bool(lip.max() <= LF + 1e-8 and ratio.max() <= -(MU-LF)+1e-8)}


def mechanism_signature():
    # Measure trained-model behavior, not an analytical identity.
    mb, _, ds = train_one('baseline', 0, 3e-3, 0.0, True)
    mi, model, ds = train_one('idea', 0, 3e-3, 0.0, True)
    model.eval(); dev = next(model.parameters()).device
    x = ds['xte'][:1000].clone().to(dev)
    xp = x.clone(); xp[:, -1] += 0.5
    with torch.no_grad():
        y = model(x); yp = model(xp)
        # Recover learned residual f=b+mu*u and compare finite differences.
        u = x[:, -1:]; up = xp[:, -1:]
        r = y + MU*u; rp = yp + MU*up
        db = (y-yp).squeeze(1); da = (u-up).squeeze(1)
        rr = ((db*da)/(da*da+1e-12)).cpu().numpy()
        ll = (torch.abs(r-rp)/(torch.abs(u-up)+1e-12)).cpu().numpy().ravel()
        # baseline model is separately trained so measure it explicitly
    _, bm, _ = train_one('baseline', 0, 3e-3, 0.0, True)
    bm.eval()
    with torch.no_grad():
        yb=bm(x); ybp=bm(xp)
        rb=((yb-ybp).squeeze(1)*da/(da*da+1e-12)).cpu().numpy()
    return {'mu': MU, 'Lf': LF, 'idea_observed_residual_fd_lipschitz_max': float(ll.max()),
            'idea_observed_one_sided_ratio_max': float(rr.max()),
            'baseline_observed_one_sided_ratio_max': float(rb.max()),
            'predicted_bound': -(MU-LF),
            'confirmed': bool(ll.max() <= LF+1e-5 and rr.max() <= -(MU-LF)+1e-5)}


def main():
    math = math_check()
    grid = [{'lr': lr, 'weight_decay': wd} for lr in LR_GRID for wd in WD_GRID]
    base = sweep_baseline(baseline_factory, grid, seeds=SWEEP_SEEDS)
    # Union parity: all idea lr settings were included in baseline sweep;
    # idea evaluates the best baseline lr plus two nearby grid settings.
    best_lr = base['best_cfg']['lr']; ordered = sorted(LR_GRID, key=lambda z: abs(np.log(z/best_lr)))
    idea_grid = [{'lr': z, 'weight_decay': base['best_cfg']['weight_decay']} for z in ordered[:3]]
    idea_runs = [(cfg, idea_eval(cfg)) for cfg in idea_grid]
    best_cfg, idea = min(idea_runs, key=lambda p: p[1]['mean'])
    sig = mechanism_signature()
    report = make_report('dynamics', 'rnn_small', base, idea,
                         {'mechanism_signature': sig,
                          'idea_grid': [{'cfg': c, 'mean': r['mean']} for c,r in idea_runs],
                          'math_sanity_check': math,
                          'custom_track': None})
    report['idea_best_cfg'] = best_cfg
    report['protocol_note'] = 'Dynamics is structurally matched: controlled pendulum windows contain (theta, omega, action), and the intervention acts on the action coordinate while retaining the shared GRU core.'
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()