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

SEEDS = tuple(range(8))
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 5e-3}]
EPOCHS = 20
BATCH = 128
B = 1.0

# Finite-candidate KKT shield for the scalar feasible interval [-B,B].
# Candidates are the empty active set z and the two one-constraint boundaries.
def interval_shield(z, bound=B):
    z = torch.as_tensor(z)
    return torch.minimum(torch.maximum(z, torch.as_tensor(-bound, dtype=z.dtype, device=z.device)),
                         torch.as_tensor(bound, dtype=z.dtype, device=z.device))

def kkt_math_check(seed=2906):
    # Also check the 2-D disk+box finite candidate construction against SLSQP.
    from scipy.optimize import minimize
    rng = np.random.default_rng(seed); box=np.array([1.35,1.10]); errs=[]; viol=[]; interior=[]
    def g(r,R): return np.array([r[0]-box[0],-r[0]-box[0],r[1]-box[1],-r[1]-box[1],r@r-R*R])
    for _ in range(100):
        z=rng.uniform(-2,2,2); R=rng.uniform(.65,1.3)
        cand=[z.copy()]
        for i in range(2):
            for s in (-1,1):
                q=z.copy(); q[i]=s*box[i]; cand.append(q)
        nz=np.linalg.norm(z)
        if nz>R: cand.append(z*R/nz)
        for s0 in (-1,1):
            for s1 in (-1,1): cand.append(np.array([s0*box[0],s1*box[1]]))
        for i in range(2):
            j=1-i; rem=R*R-box[i]**2
            if rem>=0:
                for s in (-1,1):
                    q=np.zeros(2); q[i]=box[i]; q[j]=s*np.sqrt(rem); cand.append(q)
                    q=np.zeros(2); q[i]=-box[i]; q[j]=s*np.sqrt(rem); cand.append(q)
        feas=[(0.5*np.sum((q-z)**2),q) for q in cand if np.max(g(q,R))<=1e-9]
        if not feas: continue
        got=min(feas,key=lambda a:a[0])[1]
        ref=minimize(lambda q:.5*np.sum((q-z)**2), np.clip(z,-box,box),
          constraints=[{'type':'ineq','fun':lambda q,R=R:R*R-q@q},
                       {'type':'ineq','fun':lambda q:box-q},{'type':'ineq','fun':lambda q:box+q}],
          method='SLSQP', options={'ftol':1e-12,'maxiter':300})
        if ref.success: errs.append(np.linalg.norm(got-ref.x))
        viol.append(max(0.,np.max(g(got,R))))
        if np.max(g(z,R))<=0: interior.append(np.linalg.norm(got-z))
    return {'max_error_vs_slsqp':float(max(errs) if errs else np.nan),
            'max_constraint_violation':float(max(viol) if viol else np.nan),
            'interior_identity_max_error':float(max(interior) if interior else np.nan),
            'n_reference':len(errs)}

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

def train_metric(track, s, lr, shielded):
    seed_all(s)
    d=get_dataset(track, seed=s, n_train=400, n_test=400)
    base=make_model('rnn_small', d['input_shape'], d['out_dim'])
    if shielded:
        class Shielded(torch.nn.Module):
            def __init__(self, net): super().__init__(); self.net=net
            def forward(self,x): return interval_shield(self.net(x))
        net=Shielded(base)
    else: net=base
    _, metric, _ = train_model(net,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    return float(metric)

def cfg_fn(shielded):
    def make(cfg):
        return lambda s: train_metric('dynamics', int(s), cfg['lr'], shielded)
    return make

def observed_signature(s, lr):
    seed_all(s); d=get_dataset('dynamics',seed=s,n_train=400,n_test=400)
    raw=make_model('rnn_small',d['input_shape'],d['out_dim'])
    trained,_,_=train_model(raw,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    raw = trained if trained is not None else raw
    raw.eval()
    dev=next(raw.parameters()).device
    with torch.no_grad(): z=raw(d['xte'].to(dev)); r=interval_shield(z); z=z.detach().cpu().numpy().ravel(); r=r.detach().cpu().numpy().ravel()
    interior=np.abs(z)<=B
    return {'bound':B,'n_test':len(z),'interior_fraction':float(interior.mean()),
            'interior_identity_max_error':float(np.max(np.abs(r[interior]-z[interior])) if interior.any() else 0.),
            'max_constraint_violation':float(np.max(np.maximum(np.abs(r)-B,0))),
            'mean_shaping_distance':float(np.mean(np.abs(r-z))),
            'predicted_interior_error':0.0,'predicted_max_violation':0.0,
            'confirmed': bool((np.max(np.abs(r[interior]-z[interior])) if interior.any() else 0.)<1e-7 and np.max(np.maximum(np.abs(r)-B,0))<1e-7)}

def main():
    t=time.time(); math=kkt_math_check()
    # Official sweep (first four seeds), plus parity full evaluations for every lr.
    sb=sweep_baseline(cfg_fn(False),GRID,seeds=(0,1,2,3))
    baseline_all=[]
    for c in GRID: baseline_all.append({'cfg':c,'full':evaluate(cfg_fn(False)(c),SEEDS)})
    best=min(baseline_all,key=lambda q:q['full']['mean'])
    base_block={'best_cfg':best['cfg'],'sweep':sb['sweep'],'full':best['full'],
                'parity_full':baseline_all}
    idea_all=[]
    for c in GRID: idea_all.append({'cfg':c,'full':evaluate(cfg_fn(True)(c),SEEDS)})
    ib=min(idea_all,key=lambda q:q['full']['mean'])
    sig=observed_signature(0,best['cfg']['lr'])
    sig['math_check']=math
    sig['trained_baseline_mean']=best['full']['mean']
    sig['trained_idea_best_mean']=ib['full']['mean']
    rep=make_report('dynamics','rnn_small',base_block,ib['full'],sig)
    rep['idea_sweep']=idea_all; rep['runtime_sec']=time.time()-t
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()