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

SEEDS = tuple(range(8))
SWEEP = tuple(range(4))
EPOCHS = 15
NTR, NTE = 1000, 400

# Nominal cheap structured solver for the actuated pendulum. It holds each
# observed control over the same four microsteps used by the bench generator.
def background(x, g=9.81, damp=0.25):
    z = x.reshape(-1, 8, 3)
    th, om = z[:, 0, 0], z[:, 0, 1]
    states = []
    for k in range(8):
        u = z[:, k, 2]
        states.append(torch.stack((th, om, u), 1))
        for _ in range(4):
            om = om + (-g/10 * torch.sin(th) - damp*om + 2*u) * (.05/4)
            th = th + om * (.05/4)
    return torch.stack(states, 1), th.unsqueeze(1)

class ResidualRNN(nn.Module):
    """Same rnn_small GRU+head, with analytic background and defect gate."""
    def __init__(self, gate_mid=0.08, gate_slope=3.0):
        super().__init__()
        self.core = make_model('rnn_small', (24,), 1)
        self.gate_mid, self.gate_slope = gate_mid, gate_slope
    def forward(self, x):
        bgseq, bgfinal = background(x)
        raw = x.reshape(-1, 8, 3)
        residual = raw.clone()
        residual[:, :, :2] = raw[:, :, :2] - bgseq[:, :, :2]
        # Discrete background defect: observed acceleration minus nominal one.
        dt = .05
        obs_acc = (raw[:, 1:, 1] - raw[:, :-1, 1]) / dt
        nom_acc = -9.81/10*torch.sin(raw[:, :-1, 0]) - .25*raw[:, :-1, 1] + 2*raw[:, :-1, 2]
        defect = torch.sqrt(torch.mean((obs_acc-nom_acc)**2, dim=1) + 1e-8)
        gate = torch.sigmoid(self.gate_slope*(torch.log(defect+1e-6)-math.log(self.gate_mid))).unsqueeze(1)
        closure = self.core(residual.reshape(x.shape[0], -1))
        return bgfinal + gate * closure

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

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

def baseline_fn(cfg):
    def run(seed):
        seed_all(seed); d=ds_for(seed); net=make_model('rnn_small', d['input_shape'], d['ytr'].shape[-1])
        _, metric, _ = train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,weight_decay=cfg.get('weight_decay',0.0),log=lambda *_:None)
        return metric
    return run

def idea_fn(cfg):
    def run(seed):
        seed_all(seed); d=ds_for(seed); net=ResidualRNN(cfg.get('gate_mid', 0.08), cfg.get('gate_slope', 3.0))
        _, metric, _ = train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,weight_decay=cfg.get('weight_decay',0.0),log=lambda *_:None)
        return metric
    return run

def signature(cfg):
    seed_all(0); d=ds_for(0); net=ResidualRNN(cfg['gate_mid'],cfg['gate_slope'])
    net,_,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,log=lambda *_:None)
    dev=next(net.parameters()).device
    x=d['xte'].to(dev)
    with torch.no_grad():
        bgseq,bgf=background(x); raw=x.reshape(-1,8,3)
        acc=(raw[:,1:,1]-raw[:,:-1,1])/.05
        nom=-.981*torch.sin(raw[:,:-1,0])-.25*raw[:,:-1,1]+2*raw[:,:-1,2]
        defect=torch.sqrt(torch.mean((acc-nom)**2,1)+1e-8)
        res=raw.clone()
        res[:,:,:2]=raw[:,:,:2]-bgseq[:,:,:2]
        closure=net.core(res.reshape(x.shape[0],-1)).abs().mean(1)
        bgerr=((bgf-d['yte'].to(dev))**2).sqrt().squeeze(1)
        order=torch.argsort(defect); lo=order[:len(order)//3]; hi=order[-len(order)//3:]
        corr=float(np.corrcoef(defect.cpu(),closure.cpu())[0,1])
        return {'trained_model':True,'defect_low_closure':float(closure[lo].mean()),'defect_high_closure':float(closure[hi].mean()),'defect_closure_corr':corr,'background_rmse':float(bgerr.mean()),'prediction':'closure should rise with defect and fall as background becomes accurate','confirmed':bool(closure[hi].mean()>closure[lo].mean() and corr>0)}

def main():
    grid=[{'lr':1e-3,'gate_mid':0.08,'gate_slope':3.0},{'lr':3e-3,'gate_mid':0.08,'gate_slope':3.0},{'lr':1e-2,'gate_mid':0.08,'gate_slope':3.0}]
    base=sweep_baseline(baseline_fn,grid,seeds=SWEEP)
    # Search-space parity: identical three learning rates on both sides.
    idea_trials=[]
    for cfg in grid:
        r=evaluate(idea_fn(cfg),seeds=SEEDS); idea_trials.append({'cfg':cfg,'result':r})
    best=min(idea_trials,key=lambda q:q['result']['mean'])
    rep=make_report('dynamics','rnn_small',base,best['result'],extra=signature(best['cfg']))
    rep['idea_sweep']=idea_trials
    rep['protocol']={'paired_seeds':list(SEEDS),'sweep_seeds':list(SWEEP),'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'same_architecture':True,'baseline_grid':grid,'idea_grid':grid,'intervention':'analytic pendulum background, residual coordinates, defect sigmoid gate'}
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()