import sys, json, random
from pathlib import Path
import numpy as np
import torch
from scipy.stats import norm
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, evaluate, sweep_baseline, make_report

TRACK='tabular'
MODEL='mlp_tiny_factorized_first_layer'
LR_GRID=[0.001,0.003,0.006]
SEEDS=tuple(range(8))
RANK=8
KEEP=4
EPOCHS=20
BATCH=128
PRUNE_EPOCH=11


def hac_stats(samples, delta, lag=8):
    x=np.asarray(samples, dtype=np.float64)
    mean=x.mean(0); c=x-mean; v=np.mean(c*c,0)
    lag=min(int(lag),len(x)-1)
    for k in range(1,lag+1):
        gamma=np.mean(c[k:]*c[:-k],axis=0)
        v += 2.0*(1.0-k/(lag+1.0))*gamma
    v=np.maximum(v,1e-12)
    se=np.sqrt(v/len(x)); z=(mean-delta)/se; p=norm.cdf(z)
    return mean,se,p


class FactorizedMLP(torch.nn.Module):
    # W = W0 + A B is a rank-one component decomposition.
    def __init__(self, input_dim, hidden=64, rank=RANK, seed=0):
        super().__init__()
        g=torch.Generator().manual_seed(seed+100003)
        self.register_buffer('w0', torch.randn(input_dim,hidden,generator=g)*0.08)
        self.register_buffer('b0', torch.zeros(hidden))
        self.A=torch.nn.Parameter(torch.randn(input_dim,rank,generator=g)*0.03)
        self.B=torch.nn.Parameter(torch.randn(rank,hidden,generator=g)*0.03)
        self.out=torch.nn.Linear(hidden,1)
    def forward(self,x):
        return torch.relu(x @ (self.w0+self.A@self.B)+self.b0) @ self.out.weight.t()+self.out.bias
    def prune(self, idx):
        with torch.no_grad():
            oldA=self.A.detach().clone(); oldB=self.B.detach().clone()
            self.A=torch.nn.Parameter(oldA[:,idx].clone())
            self.B=torch.nn.Parameter(oldB[idx,:].clone())


def train_one(seed, lr, mode, delta_fraction=0.25, collect_signature=False):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    d=get_dataset(TRACK, seed, n_train=400, n_test=1000)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net=FactorizedMLP(int(np.prod(d['input_shape'])),seed=seed).to(device)
        x=d['xtr'].to(device); y=d['ytr'].to(device)
        xt=d['xte'].to(device); yt=d['yte'].to(device)
        opt=torch.optim.Adam(net.parameters(),lr=lr)
        history=[]; selected=None; pre_proxy=None
        for ep in range(EPOCHS):
            perm=torch.randperm(x.shape[0],device=device)
            for start in range(0,x.shape[0],BATCH):
                ix=perm[start:start+BATCH]; xb=x[ix]; yb=y[ix]
                w=net.w0+net.A@net.B
                z=xb@w+net.b0; z.retain_grad()
                pred=torch.relu(z) @ net.out.weight.t()+net.out.bias
                loss=((pred-yb)**2).mean()
                opt.zero_grad(); loss.backward()
                with torch.no_grad():
                    # x_t,j = <grad_W, a_j b_j^T>, using grad_W = X^T grad_z.
                    gw=xb.t()@z.grad
                    vals=[]
                    for j in range(net.A.shape[1]):
                        vals.append(float(torch.abs((gw* (net.A[:,j:j+1]@net.B[j:j+1,:])).sum()).cpu()))
                    history.append(vals)
                opt.step()
            if ep==PRUNE_EPOCH:
                h=np.asarray(history)
                if mode=='confidence':
                    positive=h[h>0]
                    delta=float(np.median(positive)*delta_fraction) if positive.size else 0.0
                    mean,se,p=hac_stats(h,delta)
                    idx=np.argsort(p)[-KEEP:]
                    score_info={'mean':mean.tolist(),'se':se.tolist(),'p':p.tolist(),'delta':delta}
                else:
                    idx=np.argsort(h[-1])[-KEEP:]
                    score_info={'latest':h[-1].tolist()}
                idx=np.sort(idx)
                pre_proxy=float(np.asarray(h[-1])[idx].sum())
                net.prune(idx)
                opt=torch.optim.Adam(net.parameters(),lr=lr)
                selected=idx.tolist()
        with torch.no_grad():
            pred=net(xt); mse=float(((pred-yt)**2).mean().cpu())
            terms=[]
            hidden=torch.relu(xt@(net.w0+net.A@net.B)+net.b0)
            for j in range(net.A.shape[1]):
                terms.append(float(torch.abs((xt@net.A[:,j:j+1])@net.B[j:j+1,:]).mean().cpu()))
            observed=float(np.sum(terms))
        return {'metric':mse,'rank':int(net.A.shape[1]),'selected':selected,
                'predicted_proxy':pre_proxy,'observed_response':observed,
                'score_info':score_info if 'score_info' in locals() else {}}
    except Exception:
        if device=='cuda':
            torch.cuda.empty_cache()
            old=torch.cuda.is_available
        raise


def make_train(mode, cfg):
    return lambda seed: train_one(seed,float(cfg['lr']),mode,float(cfg.get('delta_fraction',0.25)))['metric']


def main():
    # Baseline is tuned by the official sweep on exactly the union of idea lrs.
    grid=[{'lr':lr,'delta_fraction':df} for lr in LR_GRID for df in [0.0]]
    base=sweep_baseline(lambda cfg: make_train('latest',cfg),grid)
    idea_runs={}
    idea_summaries=[]
    # Same lr grid, plus three a-priori confidence thresholds.
    for cfg in [{'lr':lr,'delta_fraction':df} for lr in LR_GRID for df in [0.15,0.25,0.50]]:
        res=evaluate(make_train('confidence',cfg),SEEDS)
        idea_runs[str(cfg)]=res
        idea_summaries.append({'cfg':cfg,'mean':res['mean']})
    best_cfg=min(idea_summaries,key=lambda q:q['mean'])['cfg']
    idea=idea_runs[str(best_cfg)]
    # Re-run best baseline is already full-seed result returned by sweep_baseline.
    example=train_one(0,float(best_cfg['lr']),'confidence',float(best_cfg['delta_fraction']))
    signature={'predicted_quantity':'sum of retained rank-one gradient contributions at pruning','observed_quantity':'sum of retained rank-one absolute responses on held-out inputs','predicted_example':example['predicted_proxy'],'observed_example':example['observed_response'],'confirmed':False}
    report=make_report(TRACK,MODEL,base,idea,signature)
    report['idea_sweep']=idea_summaries
    report['selected_idea_cfg']=best_cfg
    report['protocol_note']='Official registered tabular track; baseline sweep and idea use shared lr union, 8 final paired seeds.'
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()