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

SEEDS=tuple(range(8)); GRID=[{'lr':1e-3,'epochs':12},{'lr':3e-3,'epochs':12},{'lr':6e-3,'epochs':12}]; BATCH=128; CHUNKS=8
class Node:
    def __init__(self,a=None,b=None,leaf=None): self.a,self.b,self.leaf=a,b,leaf
def balanced(ids):
    ids=list(ids)
    if len(ids)==1:return Node(leaf=ids[0])
    m=len(ids)//2; return Node(balanced(ids[:m]),balanced(ids[m:]))
def huffman(w):
    h=[(float(x),i,Node(leaf=i)) for i,x in enumerate(w)]; heapq.heapify(h); serial=len(h)
    while len(h)>1:
        a,_,x=heapq.heappop(h); b,_,y=heapq.heappop(h); heapq.heappush(h,(a+b,serial,Node(x,y))); serial+=1
    return h[0][2]
def depths(t,d=0,o=None):
    o={} if o is None else o
    if t.leaf is not None:o[t.leaf]=d
    else:depths(t.a,d+1,o);depths(t.b,d+1,o)
    return o
def seed_all(s):
    random.seed(s);np.random.seed(s);torch.manual_seed(s)
def round_stages(x,n):
    # Shape-preserving proxy for n internal low-precision additions. Adding zero
    # isolates rounding while avoiding the invalid operation of summing disjoint chunks.
    z=torch.zeros_like(x)
    for _ in range(n): x=(x.to(torch.float16)+z.to(torch.float16)).to(torch.float32)
    return x
def run(seed,cfg,adaptive,signature=False):
    seed_all(seed); ds=get_dataset('tabular',seed=seed,n_train=400,n_test=400)
    net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']); opt=torch.optim.AdamW(net.parameters(),lr=cfg['lr'])
    xtr,ytr=ds['xtr'],ds['ytr']; rng=np.random.default_rng(seed+991); em=np.ones(CHUNKS); eq=np.ones(CHUNKS); last=None; errors=[]
    net.train()
    for _ in range(cfg['epochs']):
        perm=rng.permutation(len(xtr))
        for st in range(0,len(xtr),BATCH):
            idx=torch.as_tensor(perm[st:st+BATCH]); loss=F.mse_loss(net(xtr[idx]),ytr[idx]); opt.zero_grad();loss.backward()
            ps=[p for p in net.parameters() if p.grad is not None]; flat=torch.cat([p.grad.detach().reshape(-1) for p in ps]); raw=list(torch.tensor_split(flat,CHUNKS))
            means=np.array([float(x.mean()) for x in raw]); secs=np.array([float((x*x).mean()) for x in raw]); em=.9*em+.1*means;eq=.9*eq+.1*secs;var=np.maximum(eq-em*em,1e-12)
            tree=huffman(var) if adaptive else balanced(range(CHUNKS)); last=depths(tree)
            approx=torch.cat([round_stages(x,last[i]) for i,x in enumerate(raw)]); errors.append(float(((approx-flat)**2).mean()))
            pos=0
            for p in ps:n=p.numel();p.grad.copy_(approx[pos:pos+n].reshape_as(p));pos+=n
            opt.step()
    with torch.no_grad(): metric=float(((net(ds['xte'])-ds['yte'])**2).mean())
    if signature:return metric,{'last_depths':{str(k):int(v) for k,v in last.items()},'observed_reduction_mse':float(np.mean(errors)),'n_steps':len(errors)}
    return metric
def baseline_fn(c):return lambda s:run(s,c,False)
def idea_fn(c):return lambda s:run(s,c,True)
def main():
    base=sweep_baseline(baseline_fn,GRID); tried=[]
    for c in GRID:tried.append({'cfg':c,'mean':float(np.mean([idea_fn(c)(s) for s in range(4)]))})
    best=min(tried,key=lambda x:x['mean'])['cfg']; idea=evaluate(idea_fn(best),seeds=SEEDS)
    ss=[run(s,best,True,True)[1] for s in SEEDS]; da=[x['last_depths'] for x in ss]; md=float(np.mean([min(map(int,d.values())) for d in da]))
    sig={'predicted':'high-variance chunks have shallower Huffman depth','observed_mean_min_depth':md,'balanced_depth':3,'trained_model_measurements':da,'observed_reduction_mse_mean':float(np.mean([x['observed_reduction_mse'] for x in ss])),'confirmed':bool(md<3)}
    rep=make_report('tabular','mlp_tiny',base,idea,{'idea_sweep':tried,'mechanism_signature':sig});rep['idea']['best_cfg']=best;rep['idea']['sweep']=tried;Path('bench_report.json').write_text(json.dumps(rep,indent=2));print(json.dumps(rep,indent=2))
if __name__=='__main__':main()