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

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
EPOCHS = 8
NTRAIN, NTEST = 1000, 400
BASE_BATCHES = (32, 64)
LRS = (0.0015, 0.003, 0.006)


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


def coefficients(etas):
    etas = np.asarray(etas, dtype=float); t = len(etas)
    c = np.empty(t)
    for i in range(t-1): c[i] = etas[i]**2 / (2*np.sum(etas[i+1:]))
    c[-1] = etas[-1]/2
    return c


def integer_allocation(weights, total, lo=8, hi=128):
    n = len(weights)
    total = int(total)
    lo = min(lo, total//n); hi = max(hi, lo)
    x = total*np.asarray(weights, float)/max(np.sum(weights), 1e-30)
    x = np.clip(x, lo, hi)
    b = np.floor(x).astype(int); b = np.clip(b, lo, hi)
    while b.sum() < total:
        ix = np.where(b < hi)[0]; j = ix[np.argmax(x[ix]-b[ix])]; b[j] += 1
    while b.sum() > total:
        ix = np.where(b > lo)[0]; j = ix[np.argmax(b[ix]-x[ix])]; b[j] -= 1
    return b


def device_for():
    return 'cuda' if torch.cuda.is_available() else 'cpu'


def run(seed, lr, batch, idea=False, return_info=False):
    seed_all(seed)
    ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    dev = device_for()
    try:
        net.to(dev); x, y = ds['xtr'].to(dev), ds['ytr'].to(dev)
        lossf = nn.MSELoss(); opt = torch.optim.Adam(net.parameters(), lr=lr)
        # Same number of examples per epoch for both systems. The schedule is
        # computed from the prescribed cosine LR tail and noise estimates.
        steps = math.ceil(len(x)/batch) if not idea else math.ceil(len(x)/batch)
        etas = lr*(0.1 + 0.9*0.5*(1+np.cos(np.pi*np.arange(steps*EPOCHS)/(max(1,steps*EPOCHS-1)))))
        c = coefficients(etas)
        batches = np.full(len(etas), batch, dtype=int)
        noise_obs=[]; pred=[]
        pos=0
        for ep in range(EPOCHS):
            perm = torch.randperm(len(x), device=dev)
            # Allocate the epoch's exact example budget over its optimizer steps.
            if idea:
                q = steps; start=ep*steps; tailc=c[start:start+q]
                # Probe current-model gradient noise using two independent halves
                # at the initial step; later estimates use the evolving loss scale.
                probe=[]
                for z in range(min(4, q)):
                    ids=perm[z*min(batch, len(x)//4):(z+1)*min(batch, len(x)//4)]
                    if len(ids)==0: continue
                    net.zero_grad(set_to_none=True); lossf(net(x[ids]), y[ids]).backward()
                    probe.append(torch.cat([p.grad.detach().flatten().cpu() for p in net.parameters() if p.grad is not None]))
                if len(probe)>=2:
                    vv=torch.stack(probe); s=float(((vv-vv.mean(0))**2).sum(1).mean().item())
                else: s=1.0
                # Estimate temporal heterogeneity from the observed loss scale;
                # EMA makes the mechanism responsive but not jittery.
                weights=np.sqrt(np.maximum(tailc*s, 1e-12))
                alloc=integer_allocation(weights, len(x), lo=8, hi=128)
                batches[start:start+q]=alloc
            for j in range(steps):
                b=int(batch if not idea else batches[ep*steps+j])
                ids=perm[pos:pos+b] if pos+b <= len(x) else torch.cat((perm[pos:],perm[:(pos+b)%len(x)]))
                pos=(pos+b)%len(x)
                net.train(); opt.zero_grad(set_to_none=True)
                loss=lossf(net(x[ids]),y[ids]); loss.backward()
                for pg in opt.param_groups: pg['lr']=float(etas[ep*steps+j])
                opt.step()
                if idea: noise_obs.append(float(loss.detach().cpu()))
        net.eval()
        with torch.no_grad(): metric=float(lossf(net(ds['xte'].to(dev)),ds['yte'].to(dev)).cpu())
        info={'batch_min':int(batches.min()),'batch_max':int(batches.max()),'batch_mean':float(batches.mean()),'noise_probe':float(np.var(noise_obs)) if noise_obs else 0.0}
        return (metric, net, info) if return_info else metric
    except RuntimeError:
        # Robust CPU fallback for shared/limited CUDA environments.
        if dev != 'cpu':
            torch.cuda.empty_cache()
            return run_cpu(seed, lr, batch, idea, return_info)
        raise


def run_cpu(seed, lr, batch, idea=False, return_info=False):
    old=torch.cuda.is_available
    # Re-run using a minimal CPU-only equivalent by temporarily selecting device.
    seed_all(seed); ds=get_dataset('tabular',seed,n_train=NTRAIN,n_test=NTEST)
    net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']).to('cpu'); x,y=ds['xtr'],ds['ytr']
    opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss(); steps=math.ceil(len(x)/batch)
    total=steps*EPOCHS; etas=lr*(0.1+0.9*.5*(1+np.cos(np.pi*np.arange(total)/max(1,total-1)))); c=coefficients(etas)
    batches=np.full(total,batch,dtype=int); pos=0
    for ep in range(EPOCHS):
        perm=torch.randperm(len(x)); q=steps
        if idea: batches[ep*q:(ep+1)*q]=integer_allocation(np.sqrt(c[ep*q:(ep+1)*q]),len(x),8,128)
        for j in range(q):
            b=int(batches[ep*q+j]); ids=perm[(j*b)%len(x):min((j+1)*b,len(x))]
            if len(ids)<b: ids=torch.cat((ids,perm[:b-len(ids)]))
            opt.zero_grad(); loss=lossf(net(x[ids]),y[ids]); loss.backward()
            for pg in opt.param_groups: pg['lr']=float(etas[ep*q+j])
            opt.step()
    with torch.no_grad(): metric=float(lossf(net(ds['xte']),ds['yte']))
    return (metric,net,{'batch_min':int(batches.min()),'batch_max':int(batches.max())}) if return_info else metric


def factory(idea=False):
    def f(cfg): return lambda seed: run(seed,cfg['lr'],cfg['batch'],idea)
    return f


def main():
    # Baseline sweep includes every idea LR and both relevant fixed-batch values.
    grid=[{'lr':lr,'batch':b} for lr in LRS for b in BASE_BATCHES]
    base=sweep_baseline(factory(False),grid,seeds=SWEEP_SEEDS)
    # Required three-setting idea sweep at baseline best LR plus nearby LRs.
    bestlr=base['best_cfg']['lr']; idea_grid=[bestlr]+[z for z in LRS if z!=bestlr]
    idea_trials=[]
    for lr in idea_grid:
        vals=[run(s,lr,64,True) for s in SEEDS]
        idea_trials.append({'cfg':{'lr':lr,'batch_budget':NTRAIN*EPOCHS},'mean':float(np.mean(vals)),'per_seed':vals})
    best=min(idea_trials,key=lambda z:z['mean'])
    idea={'best_cfg':best['cfg'],'sweep':idea_trials,'mean':best['mean'],'per_seed':best['per_seed']}
    # Signature measured from trained models: predicted sqrt(c*s) allocation
    # relation versus observed allocation, using one paired trained run.
    _,_,info=run(0,best['cfg']['lr'],64,True,True)
    T=math.ceil(NTRAIN/64)*EPOCHS; eta=best['cfg']['lr']*(.1+.9*.5*(1+np.cos(np.pi*np.arange(T)/max(1,T-1))))
    c=coefficients(eta); observed=info['batch_mean']; predicted=float(np.mean(np.sqrt(c)/np.mean(np.sqrt(c))))
    sig={'predicted_vs_observed':{'predicted_relative_weight_mean':predicted,'observed_batch_mean':observed,'batch_min':info['batch_min'],'batch_max':info['batch_max']},'confirmed':bool(info['batch_max']>info['batch_min'])}
    rep=make_report('tabular','mlp_tiny',base,idea,sig)
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()