import json, random, importlib.util, sys
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, validate_and_promote, reload_custom_tracks

TRACK='planar_ising_autoregressive'; SEEDS=tuple(range(8))
GRID=[{'lr':lr,'epochs':5,'batch':128} for lr in (.001,.003,.01)]
HERE=Path(__file__).resolve().parent
spec=importlib.util.spec_from_file_location('local_ising_track', HERE/'ising_teacher_track.py')
MOD=importlib.util.module_from_spec(spec); spec.loader.exec_module(MOD)

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

def train_one(seed,cfg,soft,capture=False):
    seed_all(seed); raw=MOD.get_dataset(seed,400,200)
    xtr=torch.tensor(raw['xtr']); ytr=torch.tensor(raw['ytr']); xte=torch.tensor(raw['xte'])
    target=torch.tensor(raw['qtr'] if soft else raw['ytr'])
    net=make_model('mlp_tiny',tuple(xtr.shape[1:]),1)
    try:
        dev='cuda' if torch.cuda.is_available() else 'cpu'; net.to(dev)
        xtr,target=xtr.to(dev),target.to(dev); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
        for _ in range(cfg['epochs']):
            p=torch.randperm(len(xtr),device=dev)
            for st in range(0,len(p),cfg['batch']):
                ix=p[st:st+cfg['batch']]; loss=nn.functional.mse_loss(net(xtr[ix]),target[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): pred=net(xte.to(dev)).cpu().numpy().ravel()
    except RuntimeError:
        seed_all(seed); net=make_model('mlp_tiny',tuple(xtr.shape[1:]),1); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
        for _ in range(cfg['epochs']):
            p=torch.randperm(len(xtr))
            for st in range(0,len(p),cfg['batch']):
                ix=p[st:st+cfg['batch']]; loss=nn.functional.mse_loss(net(xtr[ix]),target[ix]); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): pred=net(xte).numpy().ravel()
    metric=float(np.mean((pred-raw['yte'].ravel())**2)); qerr=float(np.mean((pred-raw['qte'].ravel())**2))
    return {'metric':metric,'q_mse':qerr,'pred_mean':float(pred.mean())} if capture else metric

def main():
    # Promotion makes the track registered in the shared bench; never edit bench directly.
    if TRACK not in __import__('bench').all_track_names():
        if not validate_and_promote(HERE/'ising_teacher_track.py',TRACK,'exp2721 Kac-Ward Exact Teacher Stage-2'):
            raise RuntimeError('custom track promotion failed')
    reload_custom_tracks()
    base=sweep_baseline(lambda c:(lambda s:train_one(s,c,False)),GRID)
    idea_sweep=[]
    for c in GRID:
        idea_sweep.append({'cfg':c,'result':evaluate(lambda s,c=c:train_one(s,c,True),seeds=SEEDS)})
    chosen=min(idea_sweep,key=lambda z:z['result']['mean']); idea=chosen['result']; icfg=chosen['cfg']
    b0=train_one(0,base['best_cfg'],False,True); i0=train_one(0,icfg,True,True)
    sig={'prediction':'exact soft conditional labels reduce trained-NN conditional probability MSE versus sampled labels','baseline_seed0':b0,'idea_seed0':i0,'q_mse_reduction':b0['q_mse']-i0['q_mse'],'confirmed':bool(i0['q_mse']<b0['q_mse'])}
    report=make_report(TRACK,'mlp_tiny',base,idea,extra=sig)
    report['idea_sweep']=idea_sweep
    report['custom_track']={'name':TRACK,'file':'ising_teacher_track.py','domain':'sequence-level sampling'}
    (HERE/'bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()