import sys, json, random, math
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset as bench_get_dataset, train_model, evaluate, sweep_baseline, make_report
from custom_incidence_actions import MAX_ATOMS, MAX_ACTIONS, D

SEEDS = tuple(range(8)); SWEEP = tuple(range(4))

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)

class IncidenceHead(nn.Module):
    def __init__(self, width=32, alpha=0.0):
        super().__init__(); self.alpha=alpha
        self.atom = nn.Sequential(nn.Linear(D,width), nn.ReLU(), nn.Linear(width,1))
    def forward(self, x):
        atoms=x[:, :MAX_ATOMS*D].view(-1,MAX_ATOMS,D)
        A=x[:, MAX_ATOMS*D:].view(-1,MAX_ACTIONS,MAX_ATOMS)
        s=self.atom(atoms).squeeze(-1)
        counts=A.sum(-1)
        z=torch.bmm(A,s.unsqueeze(-1)).squeeze(-1)
        z=z/(counts.clamp_min(1e-6)**self.alpha)
        return z.masked_fill(counts.eq(0), -1e9)

class PaddedActionMLP(nn.Module):
    def __init__(self, width=32):
        super().__init__()
        # Standard padded action head: each legal row receives its padded aggregate
        # atom representation and a count, then a shared MLP emits its action logit.
        self.net=nn.Sequential(nn.Linear(D+1,width),nn.ReLU(),nn.Linear(width,1))
    def forward(self,x):
        atoms=x[:, :MAX_ATOMS*D].view(-1,MAX_ATOMS,D)
        A=x[:, MAX_ATOMS*D:].view(-1,MAX_ACTIONS,MAX_ATOMS)
        counts=A.sum(-1)
        feat=torch.bmm(A,atoms)
        feat=feat/(counts.unsqueeze(-1).clamp_min(1.0))
        z=self.net(torch.cat([feat, counts.unsqueeze(-1)],-1)).squeeze(-1)
        return z.masked_fill(counts.eq(0), -1e9)

def run(kind,cfg,seed,return_model=False):
    seed_all(seed); ds=bench_get_dataset('incidence_actions', seed, 400, 160)
    model=IncidenceHead(alpha=cfg.get('alpha',0.0)) if kind=='idea' else PaddedActionMLP()
    model, metric, hist=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=0.0)
    if return_model: return metric,model,ds
    return metric

def trial(kind): return lambda cfg: (lambda seed: run(kind,cfg,seed))

def main():
    # Union parity: every idea lr is included in this baseline sweep.
    grid=[{'lr':lr,'epochs':12} for lr in (0.001,0.003,0.006)]
    base=sweep_baseline(trial('base'),grid,seeds=SWEEP)
    best=base['best_cfg']
    idea_grid=[{'lr':best['lr'],'epochs':best['epochs'],'alpha':a} for a in (0.0,0.5,1.0)]
    tried=[]
    for c in idea_grid:
        r=evaluate(trial('idea')(c),seeds=SWEEP)
        tried.append({'cfg':c,'mean':r['mean']})
    bestidea=min(idea_grid,key=lambda c: next(q['mean'] for q in tried if q['cfg']==c))
    idea=evaluate(trial('idea')(bestidea),seeds=SEEDS)
    idea.update({'best_cfg':bestidea,'sweep':tried})
    # Re-test the quantitative identity on outputs of one actually trained model.
    metric,m,ds=run('idea',bestidea,0,True); m.eval()
    with torch.no_grad():
        dev=next(m.parameters()).device
        x=ds['xte'][:64].to(dev); atoms=x[:,:MAX_ATOMS*D].view(-1,MAX_ATOMS,D)
        A=x[:,MAX_ATOMS*D:].view(-1,MAX_ACTIONS,MAX_ATOMS)
        s=m.atom(atoms).squeeze(-1); counts=A.sum(-1)
        observed=m(x); predicted=torch.bmm(A,s.unsqueeze(-1)).squeeze(-1)/(counts.clamp_min(1e-6)**bestidea['alpha'])
        valid=counts.gt(0)
        err=(observed[valid]-predicted[valid]).abs()
    sig={'prediction':'trained incidence logits equal A times trained atomic scores (alpha-normalized)','n_observed_logits':int(valid.sum()),'predicted_max_abs_error':float(err.max()),'observed_mean_abs_error':float(err.mean()),'confirmed':bool(float(err.max())<1e-5)}
    report=make_report('incidence_actions','custom_incidence_head',base,idea,{'custom_track':{'name':'incidence_actions','file':'incidence_actions.py','domain':'structured_action_classification'},'baseline_sweep_union_lrs':[c['lr'] for c in grid],'idea_sweep':tried,'mechanism_signature':sig})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()