import sys, json, math, importlib.util
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, evaluate, sweep_baseline, make_report
HERE = Path(__file__).resolve().parent
sp = importlib.util.spec_from_file_location('poisson_field_track', HERE/'poisson_field_track.py')
track = importlib.util.module_from_spec(sp); sp.loader.exec_module(track)

def wquant(v, w, mass):
    o = np.argsort(v, kind='mergesort'); c = np.cumsum(w[o]/w.sum())
    return float(v[o][np.searchsorted(c, mass, side='left')])

def cq(scores, alpha=.1):
    a = np.sort(np.asarray(scores)); k = min(len(a), int(math.ceil((len(a)+1)*(1-alpha))))
    return float(a[k-1])

def run_once(seed, lr, mode, gamma=.1, collect=False):
    raw = track.get_dataset(seed, 400, 200)
    xfit,yfit = raw['xtr'][:280],raw['ytr'][:280]
    xcal,ycal = raw['xtr'][280:],raw['ytr'][280:]
    ds={'xtr':torch.tensor(xfit),'ytr':torch.tensor(yfit),'xte':torch.tensor(raw['xte']),'yte':torch.tensor(raw['yte']), 'task':'regression','metric':'mse','input_shape':(6,),'out_dim':64}
    torch.manual_seed(10000+seed); np.random.seed(10000+seed)
    net=make_model('mlp_tiny',ds['input_shape'],64)
    net,_,_=train_model(net,ds,epochs=18,lr=lr,batch=128)
    device=next(net.parameters()).device
    with torch.no_grad():
        pc=net(torch.tensor(xcal, device=device)).cpu().numpy(); pt=net(torch.tensor(raw['xte'], device=device)).cpu().numpy()
    rc=np.abs(ycal-pc); rt=np.abs(raw['yte']-pt)
    scale=np.maximum(np.quantile(np.abs(yfit-net(torch.tensor(xfit, device=device)).detach().cpu().numpy()),.75,axis=0),.01)
    normc=rc/scale; normt=rt/scale
    weights=np.ones(normc.shape[1])/normc.shape[1]
    if mode=='max': scores=normc.max(axis=1); q=cq(scores)
    else:
        scores=np.array([wquant(row,weights,1-gamma) for row in normc]); q=cq(scores)
    frac=(normt<=q).mean(axis=1)
    mse=float(np.mean((pt-raw['yte'])**2))
    out={'mse':mse,'q':q,'mean_width':float(2*q*scale.mean()),'event_rate':float(np.mean(frac>=1-gamma)),'mean_fraction':float(frac.mean())}
    return out if collect else mse

if __name__ == '__main__':
    lrs = [0.0015, 0.003, 0.006]
    seeds = tuple(range(8))
    base = sweep_baseline(
        lambda cfg: (lambda s: run_once(s, cfg['lr'], 'max')),
        [{'lr': lr} for lr in lrs], seeds=(0,1,2,3))
    # Idea uses the same three learning rates: search-space parity is exact.
    idea_runs = []
    for lr in lrs:
        r = evaluate(lambda s: run_once(s, lr, 'quantile', gamma=.10), seeds=seeds)
        idea_runs.append({'cfg': {'lr': lr, 'gamma': .10}, 'result': r})
    best = min(idea_runs, key=lambda z: z['result']['mean'])
    idea = best['result']
    # Measure the trained-model mechanism on all paired test predictions.
    sig = []
    for s in seeds:
        b = run_once(s, base['best_cfg']['lr'], 'max', collect=True)
        q = run_once(s, best['cfg']['lr'], 'quantile', gamma=.10, collect=True)
        sig.append({'seed': s, 'max_width': b['mean_width'], 'quantile_width': q['mean_width'],
                    'max_event_rate': b['event_rate'], 'quantile_event_rate': q['event_rate']})
    ms = {
        'prediction': 'spatial quantile should reduce mean band width for gamma=.10 while retaining approximately 90% domain-event rate',
        'observed_mean_max_width': float(np.mean([x['max_width'] for x in sig])),
        'observed_mean_quantile_width': float(np.mean([x['quantile_width'] for x in sig])),
        'observed_mean_max_event_rate': float(np.mean([x['max_event_rate'] for x in sig])),
        'observed_mean_quantile_event_rate': float(np.mean([x['quantile_event_rate'] for x in sig])),
        'confirmed': bool(np.mean([x['quantile_width'] for x in sig]) < np.mean([x['max_width'] for x in sig]) and np.mean([x['quantile_event_rate'] for x in sig]) >= .85),
        'per_seed': sig,
    }
    report = make_report('poisson_field_operator', 'mlp_tiny', base, idea,
        {'mechanism_signature': ms, 'idea_sweep': idea_runs,
         'custom_track': {'name': 'poisson_field_operator', 'file': 'poisson_field_track.py', 'domain': 'pde'}})
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))