import sys, json, random
from pathlib import Path
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset
from custom_inverse_gain_track import META

SEEDS=tuple(range(8)); SWEEP=tuple(range(4)); EPOCHS=15
# Union is used on both sides: baseline and idea see every lr tried.
GRID=[{'lr':1e-3},{'lr':2e-3},{'lr':3e-3}]

class DirectGRU(nn.Module):
    def __init__(self):
        super().__init__(); self.rnn=nn.GRU(3,64,batch_first=True); self.head=nn.Linear(64,1)
    def forward(self,x):
        _,h=self.rnn(x.view(x.shape[0],-1,3)); return self.head(h[-1])

class StructuredGRU(nn.Module):
    def __init__(self, qlo=.5, qhi=4.0):
        super().__init__(); self.rnn=nn.GRU(3,64,batch_first=True); self.qhead=nn.Linear(64,1)
        self.qlo=qlo; self.qhi=qhi
    def forward(self,x, return_q=False):
        seq=x.view(x.shape[0],-1,3); _,h=self.rnn(seq); q=self.qlo+(self.qhi-self.qlo)*torch.sigmoid(self.qhead(h[-1]))
        last=seq[:,-1,:]; prev=seq[:,-2,0]
        delta=(last[:,0]-prev)/.05; z=1.5*last[:,2]-1.5*last[:,0]-delta
        q=q[:,0]; out=(last[:,1]+q*z).unsqueeze(1)
        return (out,q,z) if return_q else out

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)

def ds_for(seed):
    d=get_dataset('inverse_gain_control', seed, 400, 400)
    for k in ('xtr','ytr','xte','yte'): d[k]=torch.as_tensor(d[k],dtype=torch.float32)
    return d

def run(kind, seed, lr):
    seed_all(seed+917)
    d=ds_for(seed)
    model=DirectGRU() if kind=='baseline' else StructuredGRU()
    _,metric,_=train_model(model,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
    return float(metric)

def baseline_fn(cfg):
    return lambda seed: run('baseline',seed,cfg['lr'])
def idea_fn(cfg):
    return lambda seed: run('idea',seed,cfg['lr'])

# Canonical baseline sweep (four seeds), then full paired evaluation.
base=sweep_baseline(baseline_fn,GRID,seeds=SWEEP)
# Idea sweep has exactly the same three configurations and seed budget.
idea_trials=[]
for cfg in GRID:
    r=evaluate(idea_fn(cfg),SWEEP)
    idea_trials.append({'cfg':cfg,'mean':r['mean']})
best_idea_cfg=min(idea_trials,key=lambda z:z['mean'])['cfg']
idea_full=evaluate(idea_fn(best_idea_cfg),SEEDS)
rep=make_report('inverse_gain_control','rnn_small',base,idea_full,extra={})
rep['custom_track']={'name':META['name'],'file':'custom_inverse_gain_track.py','domain':META['domain']}
rep['idea_sweep']=idea_trials
rep['idea_best_cfg']=best_idea_cfg
rep['protocol_note']='Both systems use identical GRU(3,64), Adam, epochs, batch, data, and the same lr union; only the output parameterization differs.'

# NN-scale mechanism signature from a freshly trained held-out model, not a toy graph.
# For each test sample, infer observed q from the held-out expert target and measured z.
seed=0; lr=best_idea_cfg['lr']; seed_all(seed+917); d=ds_for(seed); net=StructuredGRU()
net,_,_=train_model(net,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
net=net.cpu(); net.eval()
with torch.no_grad():
    pred,q,z=net(d['xte'],return_q=True)
    pred=pred[:,0]; target=d['yte'][:,0]; qobs=torch.where(z.abs()>1e-5,(target-d['xte'].view(-1,8,3)[:,-1,1])/z,torch.zeros_like(z))
    lhs=pred-target; rhs=(q-qobs)*z
    mask=z.abs()>1e-4
    x=rhs[mask].numpy(); y=lhs[mask].numpy()
    slope=float(np.dot(x,y)/max(np.dot(x,x),1e-12)); corr=float(np.corrcoef(x,y)[0,1])
    rel=float(torch.mean(torch.abs(lhs[mask]-rhs[mask]))/(torch.mean(torch.abs(lhs[mask]))+1e-8))
    q_mae=float(torch.mean(torch.abs(q-qobs)))
sig={'n_test':int(mask.sum()),'observed_action_error_mean_abs':float(np.mean(np.abs(y))),
     'predicted_factor_error_mean_abs':float(np.mean(np.abs(x))),'slope_observed_on_predicted':slope,
     'correlation':corr,'relative_residual':rel,'inverse_gain_mae':q_mae,
     'prediction':'action_error=(qhat-q_observed)*z','confirmed':bool(abs(slope-1)<0.05 and corr>0.995 and rel<0.05)}
rep['mechanism_signature']=sig
Path('bench_report.json').write_text(json.dumps(rep,indent=2))
print(json.dumps(rep,indent=2))