import sys, json, time
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

SEEDS=tuple(range(8))
GRID=[{'lr':1e-3,'epochs':10},{'lr':3e-3,'epochs':10},{'lr':1e-2,'epochs':10}]

def fit_ls(X,Y,w=None):
    if w is None: w=np.ones(len(X))
    A=X.T@(w[:,None]*X); B=Y.T@(w[:,None]*X)
    return B@np.linalg.pinv(A)

def coreset(X,Y,tol=1e-10):
    # Preserve normal equations at the full-data least-squares optimum.
    W=fit_ls(X,Y); n,d=X.shape; m=Y.shape[1]
    R=((Y-X@W.T)[:,:,None]*X[:,None,:]).reshape(n,m*d)
    w=np.ones(n); active=list(range(n)); rank=np.linalg.matrix_rank(X,tol=1e-6*np.linalg.norm(X,2))
    target=max(1,(m+1)*rank)
    while len(active)>target:
        block=np.asarray(active[:min(len(active),m*d+1)])
        _,s,vh=np.linalg.svd(R[block].T,full_matrices=True)
        c=vh[-1]
        if np.linalg.norm(c)<tol: break
        if np.all(c<=tol): c=-c
        pos=c>tol
        if not np.any(pos): break
        t=np.min(w[block[pos]]/c[pos]); w[block]-=t*c
        w[np.abs(w)<1e-10]=0
        active=[i for i in active if w[i]>1e-10]
    idx=np.asarray(active); return idx,w[idx],W,rank

def seed_run(seed,cfg,idea):
    torch.manual_seed(seed); np.random.seed(seed)
    ds=get_dataset('tabular',seed,n_train=400,n_test=200)
    # Canonical training path; the coreset intervention is applied to the trained readout.
    net,_,hist=train_model(make_model('mlp_tiny',ds['input_shape'],ds['out_dim']),ds,
                            epochs=cfg['epochs'],lr=cfg['lr'],batch=128)
    net.eval()
    with torch.no_grad():
        # mlp_tiny is Sequential: Linear-ReLU-Linear-ReLU-Linear; use penultimate activations.
        dev=next(net.parameters()).device
        xtr=ds['xtr'].to(dev); xte=ds['xte'].to(dev)
        emb=net[:-1](xtr).detach().cpu().numpy()
        et=net[:-1](xte).detach().cpu().numpy()
    ytr=ds['ytr'].detach().cpu().numpy().reshape(-1,1)
    yte=ds['yte'].detach().cpu().numpy().reshape(-1,1)
    if idea:
        idx,w,W,r=coreset(emb,ytr)
        W2=fit_ls(emb[idx],ytr[idx],w)
        support=len(idx)
    else:
        W2=fit_ls(emb,ytr); idx=np.arange(len(emb)); w=np.ones(len(emb)); r=np.linalg.matrix_rank(emb); support=len(idx)
    pred=et@W2.T
    mse=float(np.mean((pred-yte)**2))
    fullW=fit_ls(emb,ytr)
    normal=np.linalg.norm(((ytr[idx]-emb[idx]@fullW.T).T@(w[:,None]*emb[idx])) if idea else (ytr-emb@fullW.T).T@emb)
    denom=max(np.linalg.norm(ytr.T@emb),1e-12)
    return mse, {'support':support,'rank':int(r),'normal_residual_rel':float(normal/denom),'embedding_dim':int(emb.shape[1]),'train_final_loss':float(hist[-1]) if hist else None}

def main():
    # Common cache makes baseline and idea paired while preserving each system's own readout.
    cache={}
    def fn(cfg,idea):
        def run(seed):
            key=(seed,cfg['lr'],cfg['epochs'])
            if key not in cache: cache[key]={}
            # train separately per system to satisfy system parity; deterministic initialization/data
            v,meta=seed_run(seed,cfg,idea)
            cache[key][('idea' if idea else 'base')]=meta
            return v
        return run
    base=sweep_baseline(lambda c: fn(c,False),GRID)
    # Same grid on idea side: parity, and choose best using its 4-seed mean.
    tried=[]
    best=None; bm=float('inf')
    for c in GRID:
        r=evaluate(fn(c,True),seeds=(0,1,2,3)); tried.append({'cfg':c,'mean':r['mean']})
        if r['mean']<bm: bm=r['mean']; best=c
    idea=evaluate(fn(best,True),seeds=SEEDS)
    sig=[]
    for s in SEEDS:
        # metadata was overwritten only by same-side repeated call, but deterministic and same values.
        sig.append(cache[(s,best['lr'],best['epochs'])].get('idea',{}))
    mean_support=float(np.mean([x['support'] for x in sig])); mean_rank=float(np.mean([x['rank'] for x in sig]))
    predicted=(1+1)*mean_rank; observed=mean_support
    report=make_report('tabular','mlp_tiny',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,
      {'prediction':'support <= (m+1)r and near exact normal-equation preservation at trained embedding scale',
       'predicted_support_bound':predicted,'observed_mean_support':observed,
       'observed_mean_normal_residual_rel':float(np.mean([x['normal_residual_rel'] for x in sig])),
       'confirmed':bool(observed <= predicted+1e-6 and np.mean([x['normal_residual_rel'] for x in sig])<1e-6),
       'idea_sweep':tried,'per_seed':sig})
    report['idea']['best_cfg']=best
    report['runtime_sec']=None
    json.dump(report,open('bench_report.json','w'),indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()