import json, math, sys
from pathlib import Path
import numpy as np
import torch

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import make_model, train_model, sweep_baseline, make_report, evaluate
from custom_strain_track import get_dataset

SEEDS = tuple(range(8))
GRID = [dict(lr=1e-3, weight_decay=0.0), dict(lr=3e-3, weight_decay=0.0), dict(lr=1e-2, weight_decay=1e-4)]
EPOCHS = 18
BATCH = 128

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

def ico_axes():
    p=(1+math.sqrt(5))/2
    a=np.array([(0,1,p),(1,p,0),(p,0,1),(0,1,-p),(1,-p,0),(p,0,-1)], np.float32)
    return torch.tensor(a/math.sqrt(1+p*p), dtype=torch.float32)
AXES = ico_axes()

def idea_features(x):
    # x contains row-major F (9) and six observed directional measurements.
    F=x[:, :9].reshape(-1,3,3); y=x[:,9:]
    v=AXES.to(x.device)
    w=torch.einsum('bij,nj->bni', F, v)
    A=torch.stack([w[:,:,0]*w[:,:,0], w[:,:,1]*w[:,:,1], w[:,:,2]*w[:,:,2],
                   math.sqrt(2)*w[:,:,0]*w[:,:,1], math.sqrt(2)*w[:,:,0]*w[:,:,2],
                   math.sqrt(2)*w[:,:,1]*w[:,:,2]], dim=-1)
    eye=torch.eye(6, device=x.device, dtype=x.dtype).expand(x.shape[0],6,6)
    c=torch.linalg.solve(A.transpose(1,2)@A + 1e-5*eye, A.transpose(1,2)@y.unsqueeze(-1)).squeeze(-1)
    # trace-free postprocessing in Frobenius-preserving coordinates.
    tr=(c[:,0]+c[:,1]+c[:,2])/3
    c=c.clone(); c[:,:3]=c[:,:3]-tr[:,None]
    return torch.cat([x[:,:9], c], dim=1)

def prepare(seed, transformed):
    d=get_dataset(seed,400,400)
    out={'track':'strain_congruence','task':'regression','metric':'mse','input_shape':(15,), 'out_dim':1}
    for k in ('xtr','xte','ytr','yte'):
        out[k]=torch.as_tensor(d[k], dtype=torch.float32)
    if transformed:
        out['xtr']=idea_features(out['xtr']); out['xte']=idea_features(out['xte'])
    return out

cache={}
def run(side,cfg,seed,keep=False):
    key=(side,tuple(sorted(cfg.items())),int(seed))
    if key in cache: return cache[key][0]
    torch.manual_seed(10000+int(seed)); np.random.seed(10000+int(seed))
    ds=prepare(seed, side=='idea')
    net=make_model('mlp_tiny', ds['input_shape'], 1)
    try:
        net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'])
    except Exception:
        # train_model normally performs this fallback; preserve an explicit CPU retry.
        ds={k:(v.cpu() if torch.is_tensor(v) else v) for k,v in ds.items()}
        torch.manual_seed(10000+int(seed)); net=make_model('mlp_tiny',(15,),1)
        net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'])
    cache[key]=(float(metric),net.cpu(),ds)
    return float(metric)

def make_fn(side):
    return lambda cfg: (lambda seed: run(side,cfg,int(seed)))

def main():
    # Baseline sweep uses the same union of settings that is tested for the idea.
    base=sweep_baseline(make_fn('baseline'),GRID,seeds=(0,1,2,3))
    idea_trials=[]
    for cfg in GRID:
        r=evaluate(make_fn('idea')(cfg),SEEDS)
        idea_trials.append({'cfg':cfg,'result':r})
    best=min(idea_trials,key=lambda z:z['result']['mean'])
    rep=make_report('strain_congruence','mlp_tiny',base,best['result'],extra={})
    rep['idea_sweep']=idea_trials
    rep['custom_track']={'name':'strain_congruence','file':'custom_strain_track.py','domain':'dynamics'}
    # Mechanism signature: trained-system errors stratified by observed deformation condition.
    cond_rows=[]
    bc=base['best_cfg']; ic=best['cfg']
    for s in SEEDS:
        bm=cache[('baseline',tuple(sorted(bc.items())),s)][1]
        im=cache[('idea',tuple(sorted(ic.items())),s)][1]
        ds=cache[('idea',tuple(sorted(ic.items())),s)][2]
        raw=prepare(s,False); F=raw['xte'][:,:9].reshape(-1,3,3)
        cond=torch.linalg.cond(F).numpy(); bins=[cond<1.5,(cond>=1.5)&(cond<3),(cond>=3)]
        with torch.no_grad():
            pb=bm(raw['xte']).squeeze().numpy(); pi=im(ds['xte']).squeeze().numpy(); yy=raw['yte'].squeeze().numpy()
        for name,mask in zip(('low','mid','high'),bins):
            if mask.any(): cond_rows.append((name,float(np.mean((pb[mask]-yy[mask])**2)),float(np.mean((pi[mask]-yy[mask])**2))))
    sig={}
    for name in ('low','mid','high'):
        z=[r for r in cond_rows if r[0]==name]
        if z: sig[name]={'baseline_mse':float(np.mean([r[1] for r in z])),'idea_mse':float(np.mean([r[2] for r in z]))}
    high=sig.get('high',{'baseline_mse':float('nan'),'idea_mse':float('nan')})
    rep['mechanism_signature']={'prediction':'icosahedral reconstruction should reduce learned prediction error under ill-conditioned SL(3) deformations','condition_bins':sig,'confirmed':bool(np.isfinite(high['idea_mse']) and high['idea_mse']<high['baseline_mse'])}
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()