import sys, json, math, random
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
from bench.protocol import DEFAULT_SEEDS


def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)


def controller_rate(h, alpha, eps, hc, delta):
    z = max(-60., min(60., (h-hc)/delta))
    b = 1./(1.+math.exp(-z))
    return alpha*(eps+(1.-eps)*b)


def flat_state(model):
    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])


def assign_state(model, v):
    k=0
    with torch.no_grad():
        for p in model.parameters():
            n=p.numel(); p.copy_(v[k:k+n].view_as(p)); k += n


def predict(model, theta, x):
    # functional_call avoids mutating the center while scoring each candidate.
    from torch.func import functional_call
    params={n:v for (n,_),v in zip(model.named_parameters(), split_state(model, theta))}
    return functional_call(model, params, (x,))


def split_state(model, theta):
    out=[]; k=0
    for p in model.parameters():
        n=p.numel(); out.append(theta[k:k+n].view_as(p)); k += n
    return out


def run(seed, mode, cfg, signature=False):
    seed_all(seed)
    ds=get_dataset('tabular', seed, n_train=400, n_test=400)
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        model=make_model('mlp_tiny', tuple(ds['xtr'].shape[1:]), 1).to(device)
        x=ds['xtr'].to(device); y=ds['ytr'].to(device).view(-1,1)
        xt=ds['xte'].to(device); yt=ds['yte'].to(device).view(-1,1)
        theta=flat_state(model)
        # Same number of candidate evaluations and perturbation scale for both methods.
        N=16; sigma=cfg['sigma']; tau=cfg['tau0']; T=cfg['steps']
        alpha=cfg['alpha']; eps=cfg['epsilon']; hc=cfg['hc']; delta=cfg['delta']
        hs=[]; rates=[]; ess=[]; losses=[]
        for _ in range(T):
            noise=torch.randn((N,theta.numel()), device=device)
            cand=theta[None,:] + sigma*noise
            vals=[]
            for i in range(N):
                pred=predict(model,cand[i],x)
                vals.append(((pred-y)**2).mean())
            lv=torch.stack(vals)
            use_tau=tau
            z=-(lv-lv.min())/max(use_tau,1e-8)
            w=torch.softmax(torch.clamp(z,-80,0),dim=0)
            H=-(w*torch.log(w.clamp_min(1e-12))).sum()
            h=float(H.item()/math.log(N)); e=float(torch.exp(H).item())
            theta=(w[:,None]*cand).sum(0)
            if mode=='feedback': r=controller_rate(h,alpha,eps,hc,delta)
            else: r=alpha
            tau *= math.exp(-r)
            hs.append(h); rates.append(r); ess.append(e); losses.append(float(((predict(model,theta,x)-y)**2).item()))
        assign_state(model,theta)
        with torch.no_grad(): metric=float(((model(xt)-yt)**2).mean().item())
        if signature:
            high=np.mean([r for h,r in zip(hs,rates) if h>hc+2*delta]) if any(h>hc+2*delta for h in hs) else float('nan')
            low=np.mean([r for h,r in zip(hs,rates) if h<hc-2*delta]) if any(h<hc-2*delta for h in hs) else float('nan')
            pred_high=alpha; pred_low=alpha*eps
            return {'metric':metric,'signature':{'predicted_transition_ess':N**hc,'observed_ess_at_closest_hc':ess[int(np.argmin(np.abs(np.asarray(hs)-hc)))], 'predicted_high_entropy_rate':pred_high,'observed_high_entropy_rate':high,'predicted_low_entropy_rate':pred_low,'observed_low_entropy_rate':low,'mean_entropy':float(np.mean(hs)),'confirmed':bool(np.isfinite(high) and np.isfinite(low) and abs(high-pred_high)/pred_high<.25 and abs(low-pred_low)/pred_low<.5)}}
        return {'metric':metric}
    except RuntimeError:
        # Shared-space fallback required for constrained CUDA environments.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        return run_cpu(seed,mode,cfg,signature)


def run_cpu(seed,mode,cfg,signature=False):
    old=torch.cuda.is_available
    # Re-run deterministically on CPU by temporarily selecting tensors explicitly.
    seed_all(seed); ds=get_dataset('tabular',seed,n_train=400,n_test=400)
    model=make_model('mlp_tiny',tuple(ds['xtr'].shape[1:]),1)
    x,y=ds['xtr'],ds['ytr'].view(-1,1); xt,yt=ds['xte'],ds['yte'].view(-1,1)
    theta=flat_state(model); N=16; tau=cfg['tau0']; hs=[]; rates=[]; ess=[]
    for _ in range(cfg['steps']):
        cand=theta[None,:]+cfg['sigma']*torch.randn((N,theta.numel())); lv=torch.stack([((predict(model,cand[i],x)-y)**2).mean() for i in range(N)])
        w=torch.softmax(torch.clamp(-(lv-lv.min())/max(tau,1e-8),-80,0),0); H=-(w*torch.log(w.clamp_min(1e-12))).sum(); h=float(H/math.log(N)); r=controller_rate(h,cfg['alpha'],cfg['epsilon'],cfg['hc'],cfg['delta']) if mode=='feedback' else cfg['alpha']; theta=(w[:,None]*cand).sum(0); tau*=math.exp(-r); hs.append(h); rates.append(r); ess.append(float(torch.exp(H)))
    assign_state(model,theta); metric=float(((model(xt)-yt)**2).mean())
    if not signature:return {'metric':metric}
    hi=np.mean([r for h,r in zip(hs,rates) if h>cfg['hc']+2*cfg['delta']]) if any(h>cfg['hc']+2*cfg['delta'] for h in hs) else float('nan'); lo=np.mean([r for h,r in zip(hs,rates) if h<cfg['hc']-2*cfg['delta']]) if any(h<cfg['hc']-2*cfg['delta'] for h in hs) else float('nan')
    return {'metric':metric,'signature':{'predicted_transition_ess':N**cfg['hc'],'observed_ess_at_closest_hc':ess[int(np.argmin(np.abs(np.asarray(hs)-cfg['hc'])))],'predicted_high_entropy_rate':cfg['alpha'],'observed_high_entropy_rate':hi,'predicted_low_entropy_rate':cfg['alpha']*cfg['epsilon'],'observed_low_entropy_rate':lo,'mean_entropy':float(np.mean(hs)),'confirmed':bool(np.isfinite(hi) and np.isfinite(lo))}}


def cfg(lr):
    return {'tau0':1.0,'alpha':0.05,'epsilon':0.02,'hc':0.5,'delta':0.05,'sigma':0.03,'steps':12,'lr':lr}


def main():
    grid=[{"lr":v} for v in (0.001,0.003,0.01)]
    def base_factory(c):
        return lambda seed: run(seed, "fixed", cfg(c["lr"]))["metric"]
    base=sweep_baseline(base_factory, grid, seeds=(0,1,2,3))
    best=base["best_cfg"]
    idea_vals=[run(s, "feedback", cfg(best["lr"]))["metric"] for s in DEFAULT_SEEDS]
    idea_res={"mean":float(np.mean(idea_vals)), "std":float(np.std(idea_vals)), "per_seed":idea_vals, "n":len(idea_vals), "config":best}
    nearby={str(v):[run(s,"feedback",cfg(v))["metric"] for s in DEFAULT_SEEDS] for v in (0.001,0.003,0.01)}
    sig=run(0, "feedback", cfg(best["lr"]), signature=True)["signature"]
    report=make_report("tabular", "mlp_tiny", base, idea_res, extra=sig)
    report["idea_nearby_settings"]=nearby
    report["protocol_notes"]=("Optimizer ideas are matched to tabular Friedman regression; both systems use identical mlp_tiny, population, perturbation scale, steps, data, and paired seeds. Baseline fixed exponential cooling was swept over the shared lr union.")
    report["custom_track"]=None
    with open("bench_report.json","w") as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()