import os, sys, json, math
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, sweep_baseline, make_report
from bench.protocol import evaluate

SEED0 = 2803
EPOCHS = 12
NTR, NTE = 1200, 400
LRS = [1e-3, 3e-3, 1e-2]
TARGET_RHO = 0.95
LAMBDA = 0.05

class RandomAttractingRegressor(nn.Module):
    def __init__(self, hidden=64, candidates=2, target_rho=0.95):
        super().__init__()
        self.hidden, self.k, self.target_rho = hidden, candidates, target_rho
        self.W = nn.Parameter(torch.empty(candidates, hidden, hidden))
        self.U = nn.Parameter(torch.empty(candidates, hidden, 3))
        self.b = nn.Parameter(torch.zeros(candidates, hidden))
        self.gate = nn.Linear(3, candidates)
        self.head = nn.Linear(hidden, 1)
        for i in range(candidates):
            nn.init.orthogonal_(self.W[i])
        with torch.no_grad():
            self.W[0].mul_(0.82); self.W[1].mul_(1.14)
        nn.init.xavier_uniform_(self.U)
        nn.init.zeros_(self.gate.weight)
        nn.init.constant_(self.gate.bias, 0.0)

    def gains(self):
        return torch.linalg.matrix_norm(self.W, ord=2, dim=(-2, -1))

    def forward(self, x, return_aux=False, h0=None):
        # x is [batch, 24], eight (theta, omega, control) observations.
        seq = x.view(x.shape[0], -1, 3)
        h = x.new_zeros(x.shape[0], self.hidden) if h0 is None else h0
        states, probs = [], []
        for t in range(seq.shape[1]):
            xt = seq[:, t]
            p = torch.softmax(self.gate(xt), dim=-1)
            cand = torch.tanh(torch.einsum('kij,bj->bki', self.W, h) +
                              torch.einsum('kij,bj->bki', self.U, xt) + self.b)
            h = (p.unsqueeze(-1) * cand).sum(1)
            states.append(h); probs.append(p)
        out = self.head(h)
        if return_aux:
            return out, torch.stack(states, 1), torch.stack(probs, 1)
        return out

    def contraction_penalty(self, probs):
        eg = (probs * self.gains()).sum(-1)
        excess = torch.log(eg + 1e-8) - math.log(self.target_rho)
        return torch.relu(excess).square().mean()

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

def data(seed):
    return get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)

def train_random(cfg, seed, lam, keep=False):
    seed_all(SEED0 + seed * 101 + int(cfg['lr'] * 1e6))
    ds = data(seed)
    model = RandomAttractingRegressor(target_rho=TARGET_RHO)
    devices = ['cuda', 'cpu'] if torch.cuda.is_available() else ['cpu']
    for dev in devices:
        try:
            net = model.to(dev)
            opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
            xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev)
            for ep in range(EPOCHS):
                net.train(); perm = torch.randperm(len(xtr), device=dev)
                for j in range(0, len(xtr), 128):
                    ix = perm[j:j+128]
                    pred, _, p = net(xtr[ix], return_aux=True)
                    loss = ((pred-ytr[ix])**2).mean() + lam * net.contraction_penalty(p)
                    opt.zero_grad(); loss.backward(); opt.step()
            net.eval()
            with torch.no_grad():
                pred = net(ds['xte'].to(dev))
                metric = float(((pred-ds['yte'].to(dev)) ** 2).mean())
            if keep:
                torch.save(net.state_dict(), 'idea_model.pt' if lam else 'baseline_model.pt')
            return metric, net, ds, dev
        except RuntimeError:
            if dev == 'cuda':
                continue
    return float('nan'), None, ds, 'cpu'

def baseline_metric(cfg, seed):
    return train_random(cfg, seed, 0.0)[0]

def idea_train(cfg, seed, keep=False):
    return train_random(cfg, seed, LAMBDA, keep)

def idea_metric(cfg, seed):
    return idea_train(cfg, seed)[0]

def signature(cfg, seed=0):
    metric, net, ds, dev = idea_train(cfg, seed, keep=True)
    if net is None: return {'confirmed': False, 'error': 'training failed'}
    x = ds['xte'][:64].to(dev)
    with torch.no_grad():
        _, states, probs = net(x, return_aux=True)
        # Same inputs, two nearby initial states; re-run explicitly for observed contraction.
        h0 = torch.zeros(x.shape[0], net.hidden, device=dev); h1 = h0.clone(); h1[:,0] = 1.0
        _, s0, _ = net(x, return_aux=True, h0=h0)
        _, s1, _ = net(x, return_aux=True, h0=h1)
        d = (s0-s1).norm(dim=-1).mean(0).cpu().numpy() + 1e-12
        slope = float(np.polyfit(np.arange(len(d))[-4:], np.log(d)[-4:], 1)[0])
        gains = net.gains().cpu().numpy()
        pp = probs.mean((0,1)).cpu().numpy()
        predicted = float(np.log(np.sum(pp*gains)))
    return {'predicted_log_expected_gain': predicted, 'observed_log_distance_slope': slope,
            'relative_slope_error': abs(slope-predicted)/(abs(predicted)+1e-8),
            'mean_route_probabilities': pp.tolist(), 'candidate_spectral_gains': gains.tolist(),
            'test_mse': metric, 'confirmed': bool(abs(slope-predicted)/(abs(predicted)+1e-8) < 0.35)}

def main():
    grid = [{'lr': v} for v in LRS]
    base = sweep_baseline(lambda c: lambda s: baseline_metric(c, s), grid)
    idea_runs = []
    for cfg in grid:
        r = evaluate(lambda s, c=cfg: idea_metric(c, s))
        idea_runs.append({'cfg': cfg, 'result': r})
    best = min(idea_runs, key=lambda z: z['result']['mean'])
    idea = best['result']
    sig = signature(best['cfg'], 0)
    report = make_report('dynamics', 'rnn_small', base, idea,
                         {'predicted_vs_observed': sig, 'idea_sweep': idea_runs,
                          'track_justification': 'Dynamics is the built-in structural match for recurrent stability/control.'})
    report['protocol'] = {'epochs': EPOCHS, 'n_train': NTR, 'n_test': NTE,
                          'lr_union': LRS, 'idea_lambda': LAMBDA, 'target_rho': TARGET_RHO}
    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()