import sys, json, math
from pathlib import Path
import numpy as np
import torch
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
import bench
import geometric_track as gt

SEEDS = tuple(range(8))
LR_GRID = [0.001, 0.003, 0.01]
EPOCHS = 18
BATCH = 128
BUDGET = 8


def ds_for(seed, mask):
    d = gt.get_dataset(seed, 400, 120)
    for k in ('xtr','xte'):
        a = d[k].reshape(-1, gt.N_CANDIDATES, 3).copy()
        a[:, ~mask, :] = 0.0
        d[k] = torch.tensor(a.reshape(len(a), -1), dtype=torch.float32)
    d['ytr'] = torch.tensor(d['ytr'], dtype=torch.float32)
    d['yte'] = torch.tensor(d['yte'], dtype=torch.float32)
    d['input_shape'] = (gt.N_CANDIDATES * 3,)
    d['out_dim'] = 1
    return d


def random_mask(seed):
    r = np.random.default_rng(10000 + int(seed))
    m = np.zeros(gt.N_CANDIDATES, dtype=bool)
    m[r.choice(gt.N_CANDIDATES, BUDGET, replace=False)] = True
    return m


def area_grad(x):
    a,b,c = [float(v) for v in x]
    s=(a+b+c)/2
    ar=math.sqrt(max(s*(s-a)*(s-b)*(s-c),1e-8))
    return np.array([a*(b*b+c*c-a*a), b*(a*a+c*c-b*b), c*(a*a+b*b-c*c)])/(8*ar)


def greedy_mask(seed):
    d=gt.get_dataset(seed, 400, 120)
    x=d['xtr'].reshape(-1,gt.N_CANDIDATES,3)
    rows=[]
    for h in range(gt.N_CANDIDATES):
        g=area_grad(x[:,h,:].mean(0))
        row=np.zeros(len(gt.PAIRS)); row[gt.TRI_EDGES[h]]=g
        rows.append(row)
    rows=np.asarray(rows); selected=[]; G=1e-3*np.eye(len(gt.PAIRS))
    for _ in range(BUDGET):
        inv=np.linalg.inv(G)
        rem=[i for i in range(len(rows)) if i not in selected]
        gains=[math.log1p(float(rows[i]@inv@rows[i])) for i in rem]
        j=rem[int(np.argmax(gains))]; selected.append(j); G += np.outer(rows[j],rows[j])
    m=np.zeros(gt.N_CANDIDATES,dtype=bool); m[selected]=True
    return m, rows


def train(mask_fn, lr, seed, keep=False):
    mask = mask_fn(seed)
    d=ds_for(seed, mask)
    model=bench.make_model('mlp_tiny', d['input_shape'], 1)
    net, metric, hist=bench.train_model(model,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_: None)
    return (float(metric), net, d, mask) if keep else float(metric)


def eval_cfg(mask_fn, lr, seeds=SEEDS):
    vals=[train(mask_fn,lr,s) for s in seeds]
    return {'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}


def signature():
    # Signature is measured on independently trained seed-0 systems.
    bm=lambda s: random_mask(s)
    im=lambda s: greedy_mask(s)[0]
    b, bn, bd, bmask=train(bm,0.003,0,True)
    i, inn, idd, imask=train(im,0.003,0,True)
    def sens(net,d):
        device=next(net.parameters()).device
        z=d['xte'][:64].to(device).clone().requires_grad_(True)
        out=net(z).sum(); grad=torch.autograd.grad(out,z)[0].detach().reshape(-1,gt.N_CANDIDATES,3)
        return grad.norm(dim=2).mean(0).cpu().numpy()
    bs=sens(bn,bd); ins=sens(inn,idd)
    return {'trained_seed':0,'baseline_selected_gradient':float(bs[bmask].mean()),'baseline_omitted_gradient':float(bs[~bmask].mean()),'idea_selected_gradient':float(ins[imask].mean()),'idea_omitted_gradient':float(ins[~imask].mean()),'predicted_gain':'selected simplex coordinates should carry larger trained output sensitivity than omitted coordinates','confirmed':bool(ins[imask].mean()>ins[~imask].mean())}


def main():
    # Baseline sweep and idea sweep share exactly the same learning-rate grid.
    grid=[{'lr':lr,'budget':BUDGET,'router':'random'} for lr in LR_GRID]
    base_block=bench.sweep_baseline(
        lambda cfg: (lambda seed: train(lambda s: random_mask(s), cfg['lr'], seed)),
        grid, seeds=tuple(range(4)))
    best=base_block['best_cfg']
    base_full=eval_cfg(lambda s: random_mask(s),best['lr'],SEEDS)
    base_block['full']=base_full
    idea_runs=[]
    for lr in LR_GRID:
        r=eval_cfg(lambda s: greedy_mask(s)[0],lr,SEEDS)
        idea_runs.append((lr,r))
    idea_lr, idea_best=min(idea_runs,key=lambda z:z[1]['mean'])
    sig=signature()
    report=bench.make_report('geometric_triangle_area','mlp_tiny',base_block,idea_best,extra={'custom_track':{'name':'geometric_triangle_area','file':'geometric_track.py','domain':'geometric_graph'},'idea_sweep':[{'cfg':{'lr':lr,'budget':BUDGET,'router':'greedy_jacobian'},'mean':r['mean']} for lr,r in idea_runs],'mechanism_signature':sig})
    report['protocol_notes']={'epochs':EPOCHS,'batch':BATCH,'budget':BUDGET,'candidates':gt.N_CANDIDATES,'same_architecture':True,'full_8_paired_seeds':True}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()