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

SEED = 1448

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 RobustMoE(nn.Module):
    """Matched recurrent MoE. beta=0 is ordinary learned softmax routing."""
    def __init__(self, beta=0.0, hidden=64, experts=4):
        super().__init__()
        self.beta = float(beta); self.experts = experts
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self.router = nn.Linear(hidden, experts)
        self.transitions = nn.Parameter(torch.zeros(experts, experts))
        self.heads = nn.ModuleList([nn.Linear(hidden, 1) for _ in range(experts)])
        self.register_buffer('expert_bias', torch.tensor([-0.15, -0.05, 0.05, 0.15]))

    def _robustness(self, last, pred):
        # Smooth STL-style G(position >= -1.5 AND |omega| <= 3).
        # A short constant-velocity rollout makes the temporal minimum differentiable.
        th, om = last[:, 0], last[:, 1]
        hs = torch.arange(1, 7, device=last.device, dtype=last.dtype)[None, :]
        future_th = pred[:, :, None] + 0.05 * hs * om[:, None, None]
        r_pos = future_th + 1.5
        r_vel = 3.0 - torch.abs(om[:, None, None].expand_as(future_th))
        atoms = torch.minimum(r_pos, r_vel)
        tau = 0.12
        return -tau * torch.logsumexp(-atoms / tau, dim=2)

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        _, h = self.rnn(seq)
        h = h[-1]
        pred = torch.cat([head(h) for head in self.heads], dim=1)
        rho = self._robustness(seq[:, -1, :], pred)
        pi = torch.softmax(self.router(h), dim=1)
        # T[j,i] = softmax_i(A[j,i] + beta*rho_i), then pi' = pi @ T.
        T = torch.softmax(self.transitions[None, :, :] + self.beta * rho[:, None, :], dim=2)
        post = torch.bmm(pi[:, None, :], T).squeeze(1)
        return (post * pred).sum(1, keepdim=True)


def make_train(cfg):
    def fn(seed):
        seed_all(seed)
        d = get_dataset('dynamics', seed, n_train=800, n_test=400)
        net = RobustMoE(beta=cfg.get('beta', 0.0))
        _, metric, _ = train_model(net, d, epochs=12, lr=cfg['lr'], batch=128,
                                   weight_decay=0.0, log=lambda *a, **k: None)
        return metric
    return fn


def signature(seed, beta, lr):
    seed_all(seed); d = get_dataset('dynamics', seed, n_train=800, n_test=400)
    net = RobustMoE(beta=beta)
    net, _, _ = train_model(net, d, epochs=12, lr=lr, batch=128, log=lambda *a, **k: None)
    net = net.cpu(); net.eval(); device = torch.device('cpu'); x = d['xte'][:256]
    with torch.no_grad():
        seq=x.view(x.shape[0],-1,3); _,h=net.rnn(seq); h=h[-1]
        pred=torch.cat([q(h) for q in net.heads],1)
        rho=net._robustness(seq[:,-1,:],pred)
        pi=torch.softmax(net.router(h),1)
        T=torch.softmax(net.transitions[None,:,:]+beta*rho[:,None,:],2)
        post=torch.bmm(pi[:,None,:],T).squeeze(1)
        # Across samples, test the claimed positive robustness dependence after
        # controlling for the learned transition-logit difference.
        i,j=0,1
        observed=(torch.log(post[:,i]+1e-8)-torch.log(post[:,j]+1e-8)).cpu().numpy()
        delta=(rho[:,i]-rho[:,j]).cpu().numpy()
        slope=float(np.polyfit(delta, observed, 1)[0])
        predicted=float(beta)
        corr=float(np.corrcoef(delta, observed)[0,1])
    return {'beta': beta, 'predicted_slope': predicted,
            'observed_slope': slope, 'robustness_logodds_correlation': corr,
            'n_samples': len(delta), 'confirmed': bool(beta > 0 and slope > 0 and corr > 0.05)}


def main():
    # Cheap algebra sanity check is independent of trained scoring.
    A=np.array([0.4,-0.3,0.1,-0.2]); r=np.array([.2,-.75,.55,-.1]); b=np.linspace(0,4,9)
    logits=A[None,:]+b[:,None]*r[None,:]
    odds=logits[:,2]-logits[:,1]
    math_check={'predicted_slope':float(r[2]-r[1]),
                'observed_slope':float(np.polyfit(b,odds,1)[0]),
                'max_identity_error':float(np.max(np.abs(odds-(A[2]-A[1]+b*(r[2]-r[1])))))}
    # Baseline sweep includes every lr used in the shared search space.
    grid=[{'lr':v} for v in (1e-3,3e-3,6e-3)]
    base=sweep_baseline(make_train, grid)
    best_lr=base['best_cfg']['lr']
    # Same-sized idea sweep, with beta=0 retained as matched control.
    idea_grid=[{'lr':best_lr,'beta':v} for v in (0.0,0.75,1.5)]
    tried=[]
    for cfg in idea_grid:
        r=evaluate(make_train(cfg), seeds=(0,1,2,3))
        tried.append({'cfg':cfg,'mean':r['mean']})
    best_beta=min(tried,key=lambda z:z['mean'])['cfg']['beta']
    idea=evaluate(make_train({'lr':best_lr,'beta':best_beta}))
    # Include sweep evidence without changing the canonical report schema.
    extra={'track_choice':'dynamics: pendulum control and long-term stability structure',
           'math_check':math_check,
           'idea_sweep':tried,
           'mechanism_signature':signature(0,best_beta,best_lr)}
    report=make_report('dynamics','rnn_small',base,idea,extra)
    report['idea']['selected_cfg']={'lr':best_lr,'beta':best_beta}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()