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

SEEDS=tuple(range(8))
# Same union on both sides; baseline sweep uses all values tried by idea.
GRID=[{'lr':1e-3,'iters':8},{'lr':3e-3,'iters':16},{'lr':6e-3,'iters':24}]

class EquilibriumRNN(nn.Module):
    """Small implicit tanh recurrent layer followed by a readout.
    Both systems have exactly this architecture and parameters; only solve policy differs.
    """
    def __init__(self, in_dim=3, hidden=32, out_dim=1, iters=16, certified=False):
        super().__init__(); self.hidden=hidden; self.iters=iters; self.certified=certified
        self.in_proj=nn.Linear(in_dim,hidden); self.h_proj=nn.Linear(hidden,hidden)
        self.head=nn.Linear(hidden,out_dim)
        self.last_stats={'q':[], 'certified':0, 'fallback':0}
    def _map(self,h, x): return torch.tanh(self.in_proj(x)+self.h_proj(h))
    def _solve(self,x,n):
        h=torch.zeros(x.shape[0],self.hidden,device=x.device)
        for _ in range(n): h=self._map(h,x)
        return h
    def _certificate(self,x,h):
        # Conservative local interval Jacobian bound on a box h +/- radius.
        # For tanh(Ax+Bh), |d tanh| <= sech^2 of interval preactivation;
        # use a conservative finite-radius analytic bound.
        rad=.15
        with torch.no_grad():
            pre=self.in_proj(x)+self.h_proj(h)
            br=rad*torch.sum(torch.abs(self.h_proj.weight),dim=1)
            lo=pre-br; hi=pre+br
            # max sech^2 over interval, exact for intervals crossing zero
            near=torch.minimum(torch.abs(lo),torch.abs(hi))
            near=torch.where((lo<=0)&(hi>=0),torch.zeros_like(near),near)
            sech2=1/torch.cosh(near).clamp_min(1e-6)**2
            # induced infinity norm of interval Jacobian
            q=float(torch.max(sech2[:,None]*torch.abs(self.h_proj.weight).sum(dim=1)).item())
            # residual and a simple Krawczyk enclosure radius; conservative scalar bound
            res=torch.max(torch.abs(self._map(h,x)-h),dim=1).values
            margin=(res/(1-max(q,0.0)) if q<1 else torch.full_like(res,float('inf')))
            ok=(q<0.8) & (margin < rad*.5)
        return q, bool(torch.all(ok).item())
    def forward(self,x):
        seq=x.view(x.shape[0],-1,3)
        # equilibrium conditioning uses the complete observed sequence, not a new readout
        u=seq.mean(dim=1)
        self.last_stats={'q':[], 'certified':0, 'fallback':0}
        if not self.certified:
            h=self._solve(u,self.iters)
        else:
            # inexpensive approximate center, then interval certificate; fallback is standard solve
            h0=self._solve(u,min(8,self.iters))
            q,ok=self._certificate(u,h0)
            self.last_stats={'q':[q], 'certified':int(ok), 'fallback':int(not ok)}
            h=h0 if ok else self._solve(u,self.iters)
        return self.head(h)

def run_one(kind,cfg,seed,collect=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d=get_dataset('dynamics',seed,n_train=400,n_test=200)
    net=EquilibriumRNN(3,32,1,iters=int(cfg['iters']),certified=(kind=='idea'))
    net,metric,_=train_model(net,d,epochs=10,lr=float(cfg['lr']),batch=128,log=lambda *_:None)
    if net is None: return float('inf')
    if collect:
        with torch.no_grad():
            dev=next(net.parameters()).device
            _=net(d['xte'][:128].to(dev))
        return float(metric), dict(net.last_stats)
    return float(metric)

def factory(kind):
    return lambda cfg: (lambda seed: run_one(kind,cfg,seed))

def main():
    t=time.time()
    base=sweep_baseline(factory('base'),GRID,seeds=(0,1,2,3))
    idea_cfgs=[]
    for cfg in GRID:
        vals=evaluate(factory('idea')(cfg),SEEDS)
        idea_cfgs.append({'cfg':cfg,'result':vals})
    best=min(idea_cfgs,key=lambda z:z['result']['mean'])
    sig=[]
    for s in SEEDS:
        val,st=run_one('idea',best['cfg'],s,collect=True)
        sig.append({'seed':s,'q_observed':st['q'][0] if st['q'] else None,
                    'certified_batches':st['certified'],'fallback_batches':st['fallback']})
    qs=[z['q_observed'] for z in sig if z['q_observed'] is not None]
    certified=[z['certified_batches'] for z in sig]
    # Prediction: q<0.8 should certify; this is measured on trained models.
    low=[q for q in qs if q<.8]; high=[q for q in qs if q>=.8]
    signature={'prediction':'trained-model local q below 0.8 predicts certification; q near/above 1 predicts fallback',
      'n_models':len(qs),'predicted_low_q_count':len(low),'observed_certified_low_q_count':sum(c>0 for q,c in zip(qs,certified) if q<.8),
      'mean_q':float(np.mean(qs)) if qs else None,'min_q':float(np.min(qs)) if qs else None,
      'max_q':float(np.max(qs)) if qs else None,'q_values':qs,'certified_flags':certified,
      'confirmed':bool(low and all(c>0 for q,c in zip(qs,certified) if q<.8) and all(c==0 for q,c in zip(qs,certified) if q>=.8))}
    idea=dict(best['result']); idea['best_cfg']=best['cfg']; idea['sweep']=[{'cfg':z['cfg'],'mean':z['result']['mean']} for z in idea_cfgs]
    rep=make_report('dynamics','custom_equilibrium_rnn',base,idea,{'mechanism_signature':signature,
      'architecture_note':'Both arms train the same implicit tanh recurrent layer; baseline always iterates, idea certifies then falls back.',
      'runtime_sec':time.time()-t})
    json.dump(rep,open('bench_report.json','w'),indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()