import sys, json, time
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 make_model, evaluate, sweep_baseline, make_report
from categorical_track import get_dataset

SEEDS = (0,1,2,3,4,5,6,7)
SWEEP_SEEDS = (0,1,2,3)
EPOCHS, NTR, NTE, BATCH = 8, 400, 200, 128

def device():
    return torch.device('cpu')

def tensor_ds(seed):
    d = get_dataset(seed, NTR, NTE)
    return d

def standard_train(seed, lr):
    d = tensor_ds(seed); dev = device()
    try:
        net = make_model('mlp_tiny', d['input_shape'], d['out_dim']).to(dev)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        x = torch.tensor(d['xtr'], dtype=torch.float32, device=dev); y=torch.tensor(d['ytr'],dtype=torch.long,device=dev)
        torch.manual_seed(1000+seed); g=None
        net.train()
        for _ in range(EPOCHS):
            for ix in torch.randperm(len(x)).split(BATCH):
                ix=ix.to(dev); loss=nn.functional.cross_entropy(net(x[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval(); xt=torch.tensor(d['xte'],dtype=torch.float32,device=dev); yt=torch.tensor(d['yte'],dtype=torch.long,device=dev)
        with torch.no_grad(): err=float((net(xt).argmax(1)!=yt).float().mean().cpu())
        return err
    except Exception:
        return float('nan')

def tau_refine(net, x, h, gen):
    b=x.shape[0]; d,k=8,3; z=x.reshape(b,d,k); old=z.argmax(2)
    with torch.no_grad():
        u0=-torch.log_softmax(net(x),1)[:,1]
        states=[]
        for i in range(d):
            for aa in range(k):
                zz=z.clone(); zz[:,i,:]=0; zz[:,i,aa]=1
                states.append(zz)
        neigh=torch.stack(states,1)
        energies=-torch.log_softmax(net(neigh.reshape(b*d*k,-1)),1)[:,1].reshape(b,d,k)
        same=torch.arange(k,device=x.device)[None,None,:]==old[:,:,None]
        energies=torch.where(same,u0[:,None,None],energies)
        q=0.8*torch.exp(torch.clamp(-0.5*(energies-u0[:,None,None]),-8,8))
        q=torch.where(same,torch.zeros_like(q),q)
        prob=-torch.expm1(-h*q)
        active=torch.rand(prob.shape,device=x.device)<prob
        weights=(q*active).reshape(b*d,k); total=weights.sum(1)
        safe=weights/(total[:,None]+1e-12)
        safe=torch.where((total>0)[:,None],safe,torch.full_like(safe,1.0/k))
        picks=torch.multinomial(safe,1).reshape(b,d); has=(total>0).reshape(b,d)
        out=z.clone()
        for i in range(d):
            rows=torch.where(has[:,i])[0]
            if len(rows):
                out[rows,i,:]=0; out[rows,i,picks[rows,i]]=1
        return out.reshape(b,-1),float(prob.sum().cpu()),float(has.sum().cpu())

def tau_train(seed, lr, h, collect=False):
    d=tensor_ds(seed); dev=device(); stats=[]
    try:
        net=make_model('mlp_tiny',d['input_shape'],d['out_dim']).to(dev); opt=torch.optim.Adam(net.parameters(),lr=lr)
        x=torch.tensor(d['xtr'],dtype=torch.float32,device=dev); y=torch.tensor(d['ytr'],dtype=torch.long,device=dev)
        torch.manual_seed(2000+seed); g=None
        net.train()
        for _ in range(EPOCHS):
            for ix in torch.randperm(len(x),generator=g).split(BATCH):
                ix=ix.to(dev); xb=x[ix]; xb,pr,ob=tau_refine(net,xb,h,g); stats.append((pr,ob))
                loss=nn.functional.cross_entropy(net(xb),y[ix]); opt.zero_grad(); loss.backward(); opt.step()
        net.eval(); xt=torch.tensor(d['xte'],dtype=torch.float32,device=dev); yt=torch.tensor(d['yte'],dtype=torch.long,device=dev)
        with torch.no_grad(): err=float((net(xt).argmax(1)!=yt).float().mean().cpu())
        if collect: return err, stats
        return err
    except Exception as e:
        if collect: raise
        print("IDEA_ERROR", repr(e))
        return float('nan')

def main():
    # Search-space parity: every idea lr appears in baseline grid.
    lrs=[0.001,0.003,0.01]; hs=[0.02,0.05,0.10]
    base=sweep_baseline(lambda c: lambda s: standard_train(s,c['lr']), [{'lr':v} for v in lrs], seeds=SWEEP_SEEDS)
    idea_rows=[]
    for h in hs:
        r=evaluate(lambda s,h=h: tau_train(s,base['best_cfg']['lr'],h), seeds=SEEDS)
        idea_rows.append({'h':h,'result':r})
    best=min(idea_rows,key=lambda q:q['result']['mean'])
    idea=best['result']
    # NN-scale mechanism signature on all paired trained-model runs, not an analytic toy identity.
    sig=[]
    for s in SEEDS:
        er, st=tau_train(s,base['best_cfg']['lr'],best['h'],True)
        pp=np.mean([a for a,b in st]); oo=np.mean([b for a,b in st])
        sig.append({'seed':s,'predicted_mean_counts':float(pp),'observed_updates':float(oo)})
    pred=float(np.mean([v['predicted_mean_counts'] for v in sig])); obs=float(np.mean([v['observed_updates'] for v in sig]))
    signature={'quantity':'per-leap active proposal count versus resolved coordinate updates','predicted':pred,'observed':obs,'relative_error':abs(pred-obs)/max(pred,1e-9),'confirmed':bool(abs(pred-obs)/max(pred,1e-9)<0.10),'per_seed':sig}
    extra={'custom_track':{'name':'categorical_energy','file':'categorical_track.py','domain':'discrete_energy_sampling'},'idea_sweep':idea_rows,'mechanism_signature':signature}
    rep=make_report('categorical_energy','mlp_tiny',base,idea,extra)
    rep['baseline']['sweep_space']=[{'lr':v} for v in lrs]
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()