import sys, json, time, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_report, evaluate, sweep_baseline

HERE = Path(__file__).resolve().parent
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

class FactorRNN(nn.Module):
    def __init__(self, hidden=32):
        super().__init__()
        self.inp = nn.Linear(3, hidden)
        self.f1 = nn.Parameter(0.42 * torch.eye(hidden) + 0.035 * torch.randn(hidden, hidden))
        self.f2 = nn.Parameter(0.92 * torch.eye(hidden) + 0.035 * torch.randn(hidden, hidden))
        self.head = nn.Linear(hidden, 1)

    def transition(self):
        return self.f2 @ self.f1

    def envelope(self):
        return self.f2.abs() @ self.f1.abs()

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        h = torch.zeros(x.shape[0], self.f1.shape[0], device=x.device, dtype=x.dtype)
        A = self.transition()
        for t in range(seq.shape[1]):
            h = torch.tanh(h @ A.T + self.inp(seq[:, t]))
        return self.head(h)


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 fit(seed, lr, epochs, idea=False, beta=0.0, target=0.92):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    model = FactorRNN().to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    xtr, ytr = ds['xtr'].to(DEVICE), ds['ytr'].to(DEVICE)
    xte, yte = ds['xte'].to(DEVICE), ds['yte'].to(DEVICE)
    # Equal full-batch budget for both systems; CPU fallback handles CUDA errors.
    def run(dev):
        nonlocal model, opt, xtr, ytr, xte, yte
        model = model.to(dev); xtr, ytr, xte, yte = xtr.to(dev), ytr.to(dev), xte.to(dev), yte.to(dev)
        opt = torch.optim.Adam(model.parameters(), lr=lr)
        model.train()
        for _ in range(epochs):
            opt.zero_grad(set_to_none=True)
            pred = model(xtr)
            loss = (pred-ytr).pow(2).mean()
            if idea:
                P = model.envelope()
                # Cheap conservative infinity-norm envelope constraint plus cancellation reward.
                rows = P.sum(dim=1)
                env_pen = torch.relu(rows.max() - target).pow(2)
                A = model.transition()
                cancel = (P.sum() - A.abs().sum()) / (P.sum() + 1e-6)
                loss = loss + 0.08 * env_pen + beta * cancel
            if not torch.isfinite(loss):
                raise RuntimeError('nonfinite loss')
            loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0); opt.step()
        model.eval()
        with torch.no_grad():
            metric = float((model(xte)-yte).pow(2).mean().cpu())
        return metric
    try:
        metric = run(DEVICE)
    except Exception:
        DEVICE_CPU = 'cpu'
        metric = run(DEVICE_CPU)
    return metric, model, ds


def train_fn(cfg, idea=False):
    def f(seed):
        return fit(seed, cfg['lr'], cfg['epochs'], idea, cfg.get('beta', 0.0), cfg.get('target', 0.92))[0]
    return f


def signature(seed, cfg):
    metric, model, ds = fit(seed, cfg['lr'], cfg['epochs'], True, cfg.get('beta', 0.0), cfg.get('target', 0.92))
    with torch.no_grad():
        A = model.transition(); P = model.envelope()
        # Re-test the NN-scale prediction |Ax| <= P|x| on observed hidden probes.
        rng = torch.Generator(device=next(model.parameters()).device).manual_seed(seed+991)
        x = torch.randn(256, A.shape[0], generator=rng, device=A.device)
        lhs = (x @ A.T).abs(); rhs = x.abs() @ P.T
        violation = float((lhs-rhs).max().cpu())
        ratio = float((lhs/(rhs+1e-8)).max().cpu())
        row_env = float(P.sum(1).max().cpu())
        observed_A = float(A.abs().sum().cpu()); observed_P = float(P.sum().cpu())
    return {'metric': metric, 'max_component_violation': violation,
            'max_observed_ratio': ratio, 'predicted_bound': 'nonpositive violation',
            'observed_envelope_row_sum': row_env, 'observed_sum_abs_A': observed_A,
            'observed_sum_P': observed_P,
            'confirmed': bool(violation <= 2e-5 and ratio <= 1.00002)}


def main():
    # Shared architecture and equal epochs. The union of learning rates is evaluated
    # on both systems; baseline tuning uses the harness four-seed sweep, followed by
    # full eight-seed evaluations for every shared lr.
    base_grid = [{'lr': lr, 'epochs': 18, 'beta': 0.0}
                 for lr in (0.0015, 0.003, 0.006)]
    idea_grid = [{'lr': lr, 'epochs': 18, 'beta': beta, 'target': 0.92}
                 for lr, beta in ((0.0015, 0.005), (0.003, 0.01), (0.006, 0.02))]
    seeds = tuple(range(8))
    t0 = time.perf_counter()
    base = sweep_baseline(train_fn, base_grid)
    # Full paired evaluations at every shared lr establish search-space parity.
    base_full_by_lr = []
    for cfg in base_grid:
        r = evaluate(train_fn(cfg, False), seeds=seeds)
        base_full_by_lr.append({'cfg': cfg, **r})
    best_base_entry = min(base_full_by_lr, key=lambda z: z['mean'])
    base['full_all_union'] = base_full_by_lr
    base['best_cfg'] = best_base_entry['cfg']
    base['full'] = {k: best_base_entry[k] for k in ('mean', 'std', 'per_seed', 'n')}

    idea_all = []
    for cfg in idea_grid:
        r = evaluate(train_fn(cfg, True), seeds=seeds)
        idea_all.append({'cfg': cfg, **r})
    best_idea = min(idea_all, key=lambda z: z['mean'])
    idea = {k: best_idea[k] for k in ('mean', 'std', 'per_seed', 'n')}
    idea_cfg = best_idea['cfg']
    sig = signature(0, idea_cfg)
    report = make_report('dynamics', 'factor_rnn', base, idea,
                         {'track_structure': 'controlled pendulum multi-step dynamics', **sig,
                          'config': idea_cfg, 'wall_seconds': time.perf_counter()-t0})
    report['idea_grid'] = idea_all
    report['device'] = DEVICE
    report['note'] = ('Baseline and idea are end-to-end trained systems with identical '
                      'FactorRNN architecture; only envelope loss differs. All shared '
                      'learning rates were evaluated on the same eight paired seeds.')
    (HERE/'bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()