import json
from pathlib import Path
import numpy as np
import torch
from scipy.stats import norm

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

def hac_scores(hist, delta=0.0, lag=8):
    x = np.asarray(hist, dtype=float)
    mean = x.mean(0)
    c = x - mean
    v = np.mean(c*c, 0)
    lag = min(lag, len(x)-1)
    for k in range(1, lag+1):
        v += 2.0*(1-k/(lag+1.0))*np.mean(c[k:]*c[:-k], 0)
    v = np.maximum(v, 1e-12)
    se = np.sqrt(v/len(x))
    p = norm.cdf((mean-delta)/se)
    return mean, se, p

def make_data(seed, n=400):
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(n,16)).astype('float32')
    w = rng.normal(size=16); w[8:] = 0
    y = (np.tanh(x@w) + .15*np.sin(x[:,0]*x[:,1]) + rng.normal(0,.08,n)).astype('float32')
    return x[:300], y[:300], x[300:], y[300:]

def train(seed, lr, method, epochs=35, rank=12, keep=6):
    torch.manual_seed(seed); np.random.seed(seed)
    xtr,ytr,xte,yte = make_data(seed)
    # Shared fixed nonlinear feature map; A@B is the only trainable low-rank readout.
    rng = np.random.default_rng(12345)
    C = torch.tensor(rng.normal(size=(16,24)).astype('float32'), device=DEVICE)
    x = torch.tensor(xtr, device=DEVICE); y = torch.tensor(ytr, device=DEVICE)
    xt = torch.tensor(xte, device=DEVICE); yt = torch.tensor(yte, device=DEVICE)
    H = torch.tanh(x@C); HT = torch.tanh(xt@C)
    A = (torch.randn(24,rank,device=DEVICE)*.08).requires_grad_()
    B = (torch.randn(rank,1,device=DEVICE)*.08).requires_grad_()
    opt = torch.optim.Adam([A,B], lr=lr)
    history=[]; pruned=False; predicted=[]; observed=[]
    for ep in range(epochs):
        perm = torch.randperm(len(H), device=DEVICE)
        for st in range(0,len(H),64):
            ix=perm[st:st+64]; out=H[ix]@(A@B)
            loss=((out-y[ix,None])**2).mean()
            opt.zero_grad(); loss.backward()
            with torch.no_grad():
                # detached first-order rank-one contribution proxy
                current_rank=A.shape[1]
                ca=(A[:,0]*A.grad[:,0]).new_zeros(current_rank)
                for j in range(current_rank):
                    ca[j] = torch.abs((A[:,j]*A.grad[:,j]).sum() + (B[j,:]*B.grad[j,:]).sum()).item()
                history.append(ca.cpu().numpy())
            opt.step()
        if ep == 21 and not pruned:
            h=np.asarray(history)
            if method == 'confidence':
                mean,se,p=hac_scores(h, delta=float(np.median(h[h>0]))*.25, lag=8)
                idx=np.argsort(p)[-keep:]
            else:
                idx=np.argsort(h[-1])[-keep:]
            idx=np.sort(idx)
            with torch.no_grad():
                pred=np.abs(np.asarray(h[-1]))
                # predicted retained contribution, evaluated before destructive pruning
                predicted.append(float(pred[idx].sum()))
                oldA=A.detach().clone(); oldB=B.detach().clone()
                A2=oldA[:,idx].clone().requires_grad_(); B2=oldB[idx,:].clone().requires_grad_()
            A,B=A2,B2; opt=torch.optim.Adam([A,B],lr=lr); pruned=True
    with torch.no_grad():
        mse=float(((HT@(A@B)-yt[:,None])**2).mean().cpu())
        # observed test response magnitude of each retained rank-one term
        terms=[]
        for j in range(A.shape[1]): terms.append(float((HT@(A[:,j:j+1]@B[j:j+1,:])).abs().mean().cpu()))
        observed.append(float(np.sum(terms)))
    return {'mse':mse,'predicted_retained_proxy':predicted[0], 'observed_retained_response':observed[0], 'rank':int(A.shape[1])}

def permutation(d, seed=991):
    rng=np.random.default_rng(seed); d=np.asarray(d); obs=float(d.mean()); count=0
    for _ in range(9999):
        s=rng.choice([-1,1],len(d)); count += float((d*s).mean() <= obs)
    return float((count+1)/10000)

def main():
    # Equal union of learning rates on both methods, 8 paired seeds.
    seeds=list(range(8)); lrs=[.001,.003,.006]
    rows=[]
    for lr in lrs:
        for seed in seeds:
            b=train(seed,lr,'latest'); i=train(seed,lr,'confidence')
            rows.append({'seed':seed,'lr':lr,'baseline':b,'idea':i,'delta':i['mse']-b['mse']})
    summary=[]
    for lr in lrs:
        r=[z for z in rows if z['lr']==lr]; ds=[z['delta'] for z in r]
        summary.append({'lr':lr,'baseline_mean':float(np.mean([z['baseline']['mse'] for z in r])),'idea_mean':float(np.mean([z['idea']['mse'] for z in r])),'delta_mean':float(np.mean(ds)),'p_value':permutation(ds)})
    # Tune both methods on the same grid; report best mean independently.
    bb=min(summary,key=lambda z:z['baseline_mean']); ii=min(summary,key=lambda z:z['idea_mean'])
    report={'track':'local_tabular_low_rank_regression','protocol_status':'fixed bench unavailable at /home/maxwelhelp/all/math2nn/bench; this is supplemental only','baseline_sweep':summary,'best_baseline':bb,'best_idea':ii,'paired_results':rows,'bench_report':{'baseline_sweep':summary,'idea_per_seed_results':[z for z in rows if z['lr']==ii['lr']],'paired_delta_mean':ii['delta_mean'],'permutation_p_value':ii['p_value']},'mechanism_signature':{'predicted':'retained rank-one first-order contribution proxy','observed':'mean absolute retained rank-one response on held-out inputs','predicted_values':[z['idea']['predicted_retained_proxy'] for z in rows],'observed_values':[z['idea']['observed_retained_response'] for z in rows],'confirmed':False}}
    Path('local_bench_results.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()