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

TRACK='directional_polytope'; MODEL='mlp_tiny'
SEEDS=tuple(range(8)); BOUNDS=(-1.0,1.0)

def legendre_box(z, eps=0.1, lo=-1., hi=1.):
    # argmin z*w + eps[-log(w-lo)-log(hi-w)] independently.
    # Solves z + eps[-1/(w-lo)+1/(hi-w)] = 0 in closed form.
    a = z
    L = hi-lo
    disc = torch.sqrt((a*L)**2 + 4.0*eps*eps)
    # numerically stable positive root for s=w-lo: a*s^2 + ...
    # equivalent centered solution: q=w-(lo+hi)/2, q = (eps*a? )
    # stationarity with center c and radius r: z + eps*(-1/(r+q)+1/(r-q))=0
    # => z*(r^2-q^2)+2 eps*q=0; q=(sqrt(eps^2+z^2 r^2)-eps)/z
    r=L/2; c=(lo+hi)/2
    root=torch.sqrt(eps*eps + (a*r)**2)
    q=torch.where(a.abs()>1e-8, -a*r*r/(root+eps), torch.zeros_like(a))
    return (c+q).clamp(min=lo+1e-7,max=hi-1e-7)

def train(seed, lr, kind, penalty=1.0, eps=0.1, epochs=20, return_net=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d=get_dataset(TRACK, seed, n_train=400, n_test=200)
    net=make_model(MODEL,d['input_shape'],d['out_dim'])
    opt=torch.optim.Adam(net.parameters(),lr=lr)
    xtr,ytr,xte,yte=d['xtr'],d['ytr'],d['xte'],d['yte']
    if d.get('out_dim', 1) > 1 and ytr.shape[0] != xtr.shape[0]:
        ytr=ytr.reshape(xtr.shape[0], d['out_dim']); yte=yte.reshape(xte.shape[0], d['out_dim'])
    for ep in range(epochs):
        p=torch.randperm(len(xtr)); net.train()
        for i in range(0,len(xtr),128):
            ix=p[i:i+128]; raw=net(xtr[ix])
            if kind=='idea': out=legendre_box(raw,eps)
            else: out=raw
            loss=((out-ytr[ix])**2).mean()
            if kind=='baseline':
                loss=loss + penalty*(torch.relu(out-1).square()+torch.relu(-1-out).square()).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    net.eval()
    with torch.no_grad():
        raw=net(xte); out=legendre_box(raw,eps) if kind=='idea' else raw
        metric=((out-yte)**2).mean().item()
    if return_net: return metric,net,d
    return metric

def base_factory(cfg):
    return lambda seed: train(seed,cfg['lr'],'baseline',cfg['penalty'])
def idea_factory(cfg):
    return lambda seed: train(seed,cfg['lr'],'idea',eps=cfg['eps'])

def signature():
    metric,net,d=train(0,0.01,'idea',eps=.1,return_net=True)
    z=net(d['xte'][:1]).detach().requires_grad_(True)
    w=legendre_box(z,.1)
    # diagonal observed J via autograd and analytic barrier inverse Hessian
    obs=[]
    for j in range(w.shape[1]):
        obs.append(torch.autograd.grad(w[0,j],z,retain_graph=True)[0][0,j].item())
    zz=z.detach()[0]; ww=legendre_box(zz,.1); h=.1*((1/(ww+1)**2)+(1/(1-ww)**2))
    pred=(-1/h).numpy(); obs=np.array(obs)
    rel=float(np.max(np.abs(obs-pred)/(np.abs(pred)+1e-12)))
    return {'quantity':'diagonal Jacobian of trained feasibility head', 'predicted_mean_abs':float(np.mean(np.abs(pred))), 'observed_mean_abs':float(np.mean(np.abs(obs))), 'max_relative_error':rel, 'confirmed':bool(rel<1e-5)}

def main():
    # Union parity: all lr and baseline method knobs used by either side are swept here.
    grid=[{'lr':lr,'penalty':p} for lr in (.003,.01,.03) for p in (0.0,1.0,10.0)]
    t=time.time(); base=sweep_baseline(base_factory,grid,seeds=(0,1,2,3))
    # explicit idea 3-config sweep at the same lr union; report best by 4-seed mean
    igrid=[{'lr':lr,'eps':e} for lr,e in ((base['best_cfg']['lr'],.05),(base['best_cfg']['lr'],.1),(base['best_cfg']['lr'],.2))]
    itried=[]
    for cfg in igrid:
        vals=[idea_factory(cfg)(s) for s in (0,1,2,3)]
        itried.append({'cfg':cfg,'mean':float(np.mean(vals))})
    ibest=min(itried,key=lambda x:x['mean'])['cfg']
    idea={'best_cfg':ibest,'sweep':itried,'full':__import__('bench').evaluate(idea_factory(ibest),SEEDS)}
    rep=make_report(TRACK,MODEL,base,idea['full'],signature())
    rep['idea_sweep']=idea['sweep']; rep['idea_best_cfg']=idea['best_cfg']
    rep['runtime_seconds']=time.time()-t; rep['protocol_note']='Registered directional_polytope; same MLP, end-to-end trained systems; box log barrier feasibility.'
    print(json.dumps(rep,indent=2))
    open('bench_report.json','w').write(json.dumps(rep,indent=2))
if __name__=='__main__': main()