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

TRACK='poisson_boundary'; MODEL='mlp_tiny'; SEEDS=tuple(range(8))
A=0.75; S=0.10; DELTA=1e-5
# Union of all learning rates tried by either system; baseline also sweeps its decisive boundary knob.
LRS=[1e-3, 3e-3, 1e-2]; BASE_GRID=[{'lr':lr,'boundary_weight':bw} for lr in LRS for bw in (1.0,10.0,30.0)]
IDEA_GRID=[{'lr':lr,'lambda_g':lam} for lr,lam in zip(LRS,(0.0,0.001,0.01))]
EPOCHS=35; BATCH=128

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 device():
    try:
        d=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        if d.type=='cuda': torch.zeros(1,device=d)
        return d
    except Exception: return torch.device('cpu')

def dist_disk(x):
    return (1.0-torch.sqrt((x*x).sum(1,keepdim=True)).clamp(max=1.0)).clamp_min(0.0)

def train_one(kind, cfg, seed, collect=False):
    seed_all(seed); ds=get_dataset(seed,n_train=400,n_test=400); dev=device()
    net=make_model(MODEL,ds['input_shape'],ds['out_dim']).to(dev)
    xtr=torch.as_tensor(ds['xtr'],dtype=torch.float32,device=dev); ytr=torch.as_tensor(ds['ytr'],dtype=torch.float32,device=dev).reshape(-1,1); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
    # fixed boundary points make the baseline a conventional Dirichlet-penalty fit.
    ang=torch.linspace(0,2*np.pi,128,device=dev)[:-1]; xb=torch.stack((torch.cos(ang),torch.sin(ang)),1)
    hist=[]
    for ep in range(EPOCHS):
        net.train(); perm=torch.randperm(len(xtr),device=dev); total=0.
        for i in range(0,len(xtr),BATCH):
            idx=perm[i:i+BATCH]; x=xtr[idx]
            if kind=='idea':
                x=x.detach().requires_grad_(True); v=net(x); d=dist_disk(x)
                u=(torch.sqrt(d*d+DELTA*DELTA)-DELTA).pow(A)*v
                grad=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
                wd=(torch.sqrt(d*d+DELTA*DELTA)-DELTA).clamp_min(DELTA).pow(1-A+S)*grad
                loss=((u-ytr[idx])**2).mean()+cfg['lambda_g']*(wd.square().sum(1)).mean()
            else:
                u=net(x); loss=((u-ytr[idx])**2).mean()+cfg['boundary_weight']*(net(xb)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step(); total += float(loss.detach())*len(idx)
        hist.append(total/len(xtr))
    net.eval();
    with torch.no_grad():
        xte=torch.as_tensor(ds['xte'],dtype=torch.float32,device=dev); yte=torch.as_tensor(ds['yte'],dtype=torch.float32,device=dev).reshape(-1,1); pred=net(xte)
        if kind=='idea': pred=dist_disk(xte).clamp_min(DELTA).pow(A)*pred
        metric=float(((pred-yte)**2).mean())
    out={'metric':metric,'final_loss':hist[-1],'history':hist}
    if collect:
        # Behavioural signature: fitted models' near-boundary scaling, not an analytic identity.
        with torch.no_grad():
            th=torch.linspace(0.13,2*np.pi-0.13,256,device=dev); vals=[]
            for dd in (0.02,0.10):
                xx=torch.stack(((1-dd)*torch.cos(th),(1-dd)*torch.sin(th)),1)
                z=net(xx)
                if kind=='idea': z=dd**A*z
                vals.append(z.abs().median().item()+1e-8)
            observed_ratio=vals[0]/vals[1]
        out['near_boundary_ratio']=observed_ratio
    return out

def baseline_fn(cfg): return lambda seed: train_one('baseline',cfg,seed)['metric']
def idea_fn(cfg): return lambda seed: train_one('idea',cfg,seed)['metric']

def main():
    t=time.time()
    base=sweep_baseline(baseline_fn,BASE_GRID,seeds=(0,1,2,3))
    # Explicitly run all three idea settings; report the best on the same sweep seeds,
    # then evaluate that selected setting on all eight paired seeds.
    idea_trials=[]
    for cfg in IDEA_GRID:
        r=evaluate(idea_fn(cfg),seeds=(0,1,2,3)); idea_trials.append({'cfg':cfg,'mean':r['mean']})
    best_cfg=min(idea_trials,key=lambda z:z['mean'])['cfg']; idea_full=evaluate(idea_fn(best_cfg),SEEDS)
    rep=make_report(TRACK,MODEL,base,idea_full,extra={
      'custom_track':{'name':TRACK,'file':'fractional_dirichlet_track.py','domain':'pde'},
      'idea_sweep':idea_trials,
      'protocol':{'epochs':EPOCHS,'batch':BATCH,'a':A,'s':S,'delta':DELTA,'paired_seeds':list(SEEDS)},
      'mechanism_signature':{
        'quantity':'median |u| at d=0.02 divided by median |u| at d=0.10 on trained models',
        'predicted_idea_ratio':float((0.02/0.10)**A),
        'predicted_baseline_ratio':'not constrained (free output with boundary penalty)',
        'observed_idea':train_one('idea',best_cfg,0,True)['near_boundary_ratio'],
        'observed_baseline':train_one('baseline',base['best_cfg'],0,True)['near_boundary_ratio'],
        'confirmed':False
      }, 'runtime_seconds':time.time()-t})
    # Signature confirmation requires quantitative agreement within 25%.
    sig=rep['mechanism_signature']; sig['confirmed']=abs(sig['observed_idea']-sig['predicted_idea_ratio'])/sig['predicted_idea_ratio']<0.25
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()