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

SEEDS = tuple(range(8))
TRACE = []

class PositiveGainRNN(nn.Module):
    def __init__(self, hidden=64, mode='fixed', gain=1.0, a=1.0, b=0.2, dt=0.2):
        super().__init__()
        self.hidden, self.mode, self.gain = hidden, mode, gain
        self.a, self.b, self.dt = a, b, dt
        self.inp = nn.Linear(3, hidden)
        self.rec = nn.Linear(hidden, hidden)
        self.bias = nn.Parameter(torch.zeros(hidden))
        self.head = nn.Linear(hidden, 1)
        self.last_trace = None

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        h = torch.zeros(x.shape[0], self.hidden, device=x.device, dtype=x.dtype)
        w = torch.full((x.shape[0], 1), float(self.gain), device=x.device, dtype=x.dtype)
        ys, ws = [], []
        for k in range(seq.shape[1]):
            # Positive state and positive output projection, matching the intended plant.
            h = torch.relu((1.0 - self.dt) * h + self.dt * (self.inp(seq[:, k]) + w * self.rec(h) + self.bias))
            y = h.mean(dim=1, keepdim=True)
            if self.mode == 'pito':
                w = torch.relu(w + self.dt * (self.b - self.a * y * w))
            ys.append(y.detach())
            ws.append(w.detach())
        self.last_trace = (torch.stack(ys, 1), torch.stack(ws, 1))
        return self.head(h)

def set_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

def run_one(seed, cfg, mode, keep_trace=False):
    set_seed(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=400)
    if mode == 'fixed':
        net = PositiveGainRNN(mode='fixed', gain=cfg['gain'])
    else:
        net = PositiveGainRNN(mode='pito', gain=cfg.get('gain', 1.0), a=cfg['a'], b=cfg['b'])
    net, metric, hist = train_model(net, ds, epochs=20, lr=cfg['lr'], batch=128, log=lambda *_: None)
    if net is None or metric is None: return float('nan')
    if keep_trace:
        dev = next(net.parameters()).device
        with torch.no_grad(): net(ds['xte'][:128].to(dev))
        if net.last_trace is not None: TRACE.append((net.a, net.b, net.last_trace))
    return float(metric)

def baseline_factory(cfg):
    return lambda seed: run_one(seed, cfg, 'fixed')

def idea_factory(cfg):
    return lambda seed: run_one(seed, cfg, 'pito', keep_trace=True)

def signature():
    # Re-test the theorem's prediction on traces of trained idea models. Select
    # a high-output tail threshold, then fit log(w - b/(aV)) where positive.
    rates, preds, thresholds = [], [], []
    for a, b, (ys, ws) in TRACE:
        y = ys.cpu().numpy().reshape(-1); w = ws.cpu().numpy().reshape(-1)
        V = float(np.quantile(y, .75))
        if V <= 1e-5: continue
        floor = b/(a*V); pred = a*V
        mask = (y >= V) & (w > floor + 1e-5)
        if mask.sum() >= 8:
            t = np.arange(len(w), dtype=float)[mask] * .2
            slope = -float(np.polyfit(t, np.log(w[mask]-floor), 1)[0])
            if np.isfinite(slope) and slope > 0:
                rates.append(slope); preds.append(pred); thresholds.append(V)
    if not rates:
        return {'prediction': 'decay slope = a*V when y>=V', 'predicted_rate': None,
                'observed_rate': None, 'relative_error': None, 'n_traces': 0,
                'confirmed': False}
    rel = abs(float(np.mean(rates))-float(np.mean(preds)))/float(np.mean(preds))
    return {'prediction': 'decay slope = a*V when y>=V',
            'predicted_rate': float(np.mean(preds)), 'observed_rate': float(np.mean(rates)),
            'relative_error': float(rel), 'threshold_V': float(np.mean(thresholds)),
            'n_traces': len(rates), 'confirmed': bool(rel <= .20)}

def main():
    # Union of all learning rates is shared by both methods. Baseline's central
    # method knob (fixed recurrent gain) is swept as well.
    lrs = [1e-3, 3e-3, 1e-2]
    base_grid = [{'lr': lr, 'gain': g} for lr in lrs for g in [0.5, 1.0, 1.5]]
    base = sweep_baseline(baseline_factory, base_grid, seeds=(0,1,2,3))
    idea_grid = [{'lr': base['best_cfg']['lr'], 'a': 1.0, 'b': .2, 'gain': 1.0},
                 {'lr': 1e-3, 'a': 1.0, 'b': .2, 'gain': 1.0},
                 {'lr': 1e-2, 'a': 1.0, 'b': .2, 'gain': 1.0}]
    idea_runs = []
    for cfg in idea_grid:
        TRACE.clear()
        r = evaluate(idea_factory(cfg), seeds=SEEDS)
        idea_runs.append((cfg, r, signature()))
    best_cfg, idea, sig = min(idea_runs, key=lambda z: z[1]['mean'])
    # Ensure signature corresponds to selected configuration by rerunning it.
    TRACE.clear(); idea = evaluate(idea_factory(best_cfg), seeds=SEEDS); sig = signature()
    report = make_report('dynamics', 'positive_rnn_fixed_gain_vs_pito', base, idea,
                         {'prediction': sig, 'selected_idea_cfg': best_cfg,
                          'matched_structure': 'dynamics/control'})
    report['idea_sweep'] = [{'cfg': c, 'result': r, 'mechanism_signature': s}
                            for c, r, s in idea_runs]
    report['custom_track'] = None
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()