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

SEEDS = tuple(range(8)); LRS = [1e-3, 3e-3, 1e-2]; EPOCHS = 18; CODES = 16

class STEQuant(nn.Module):
    def __init__(self, edges, values):
        super().__init__(); self.register_buffer('edges',torch.tensor(edges,dtype=torch.float32)); self.register_buffer('values',torch.tensor(values,dtype=torch.float32))
    def forward(self,x):
        z=torch.relu(x); idx=torch.bucketize(z.detach(),self.edges[1:-1]); q=self.values[idx]; return z+(q-z).detach()

def lloyd_edges_values(a,n):
    x=np.asarray(a,dtype=np.float64).ravel(); xmax=max(float(np.quantile(x,.999)),1e-3); x=np.clip(x,0,xmax); c=np.quantile(x,(np.arange(n)+.5)/n)
    for _ in range(30):
        lab=np.searchsorted((c[:-1]+c[1:])/2,x); nc=c.copy()
        for k in range(n):
            v=x[lab==k]
            if len(v): nc[k]=v.mean()
        if np.max(abs(nc-c))<1e-5: break
        c=nc
    return np.r_[0,(c[:-1]+c[1:])/2,xmax],c

def utility_edges_values(a,n):
    x=np.asarray(a,dtype=np.float64).ravel(); xmax=max(float(np.quantile(x,.999)),1e-3); grid=np.linspace(0,xmax,1025)
    hist,_=np.histogram(np.clip(x,0,xmax),bins=grid); mids=(grid[:-1]+grid[1:])/2
    p=hist.astype(float)/max(hist.sum(),1); s=max(float(np.median(x)),1e-3); qp=1/((s+mids)*np.log1p(xmax/s)); g=np.sqrt(p*qp+1e-12)
    # g is a per-bin mass proxy; assign it to the corresponding right edge.
    c=np.r_[0,np.cumsum(g)]; c/=c[-1]
    internal=np.interp(np.arange(1,n)/n,c,grid)
    return np.r_[0,internal,xmax],np.r_[0,internal]

def make_net(ds,kind,seed):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed); probe=nn.Sequential(nn.Linear(10,64),nn.ReLU())
    with torch.no_grad(): act=probe(ds['xtr']).numpy()
    e,v=utility_edges_values(act,CODES) if kind=='idea' else lloyd_edges_values(act,CODES)
    class Net(nn.Module):
        def __init__(self):
            super().__init__(); self.l1=nn.Linear(10,64); self.q=STEQuant(e,v); self.l2=nn.Linear(64,64); self.l3=nn.Linear(64,1)
        def forward(self,x): return self.l3(torch.relu(self.l2(self.q(self.l1(x)))))
    return Net(),e,v

def run_one(seed,kind,lr):
    ds=get_dataset('tabular',seed,n_train=400,n_test=400); net,e,v=make_net(ds,kind,seed); net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=128)
    if net is None:return float('nan'),{}
    dev=next(net.parameters()).device
    with torch.no_grad(): z=torch.relu(net.l1(ds['xte'].to(dev))); q=net.q(z); up=float((q>z+1e-6).float().mean().cpu()); gap=float((z-q).mean().cpu())
    return metric,{'upward_fraction':up,'mean_left_gap':gap,'codes':CODES}

def baseline_factory(cfg): return lambda seed:run_one(seed,'baseline',cfg['lr'])[0]

def main():
    grid=[{'lr':x} for x in LRS]; base=sweep_baseline(baseline_factory,grid,seeds=SEEDS[:4]); best=base.get('best_config',base.get('best',grid[1])); best_lr=float(best.get('lr',best.get('learning_rate',3e-3)))
    idea_all=[]
    for lr in [best_lr]+[x for x in LRS if x!=best_lr][:2]:
        vals=[]; aux=[]
        for s in SEEDS:
            m,a=run_one(s,'idea',lr); vals.append(m); aux.append(a)
        idea_all.append({'lr':lr,'per_seed':vals,'aux':aux,'mean':float(np.nanmean(vals))})
    chosen=min(idea_all,key=lambda z:z['mean']); bvals=[]; baux=[]
    for s in SEEDS:
        m,a=run_one(s,'baseline',chosen['lr']); bvals.append(m); baux.append(a)
    base_block={'sweep':base,'best_config':{'lr':chosen['lr']},'full':{'per_seed':bvals,'aux':baux,'mean':float(np.nanmean(bvals))}}
    idea_block={'config':{'lr':chosen['lr']},'per_seed':chosen['per_seed'],'aux':chosen['aux'],'mean':chosen['mean'],'sweep':idea_all}
    up=float(np.mean([a['upward_fraction'] for a in chosen['aux']])); gap=float(np.mean([a['mean_left_gap'] for a in chosen['aux']]))
    sig={'predicted':{'upward_fraction':0.0,'left_gap_nonnegative':True},'observed':{'mean_upward_fraction':up,'mean_left_gap':gap},'confirmed':bool(up==0.0 and gap>=0)}
    rep=make_report('tabular','mlp_tiny',base_block,idea_block,{'mechanism_signature':sig}); Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__':main()