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

EPOCHS = 14
SEEDS = tuple(range(8))
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]

class AdapterNet(nn.Module):
    def __init__(self, groups, rank=3, seed=0):
        super().__init__(); torch.manual_seed(seed)
        self.groups = [list(g) for g in groups]
        self.trunk = nn.Sequential(nn.Linear(12,48), nn.ReLU())
        self.us = nn.ParameterList([nn.Parameter(torch.randn(48,rank)*.04) for _ in self.groups])
        self.vs = nn.ParameterList([nn.Parameter(torch.randn(rank,48)*.04) for _ in self.groups])
        self.head = nn.Linear(48,6)
    def forward(self,x):
        h=self.trunk(x); y=h.new_zeros((x.shape[0],6))
        for g,u,v in zip(self.groups,self.us,self.vs):
            y[:,g]=self.head(h+(h@u)@v)[:,g]
        return y

def as_tensors(seed):
    d=get_dataset(seed,800,400)
    return {k:(torch.from_numpy(v) if isinstance(v,np.ndarray) else v) for k,v in d.items()}

def tc(cov,eps=1e-6):
    cov=np.asarray(cov,float); diag=np.maximum(np.diag(cov),eps)
    sign,ld=np.linalg.slogdet(cov+eps*np.eye(len(diag)))
    return float(.5*(np.log(diag).sum()-ld)) if sign>0 else 0.0

def cluster_residuals(residual,k=3,alpha=.10):
    s=np.cov(residual,rowvar=False,ddof=1); s=(1-alpha)*s+alpha*np.diag(np.diag(s))
    def merge(a,b):
        ab=list(a)+list(b)
        return tc(s[np.ix_(ab,ab)])-tc(s[np.ix_(a,a)])-tc(s[np.ix_(b,b)])
    groups=[[i] for i in range(s.shape[0])]
    while len(groups)>k:
        _,i,j=max((merge(a,b),i,j) for i,a in enumerate(groups) for j,b in enumerate(groups) if j>i)
        groups[i]=sorted(groups[i]+groups[j]); del groups[j]
    return sorted(groups),s

def fit_metric(seed,lr,mode,return_sig=False):
    d=as_tensors(seed); n=len(d['xtr']); cut=n//5
    # Held-out warm-up residual buffer: no gradient updates use this buffer.
    warm=AdapterNet([list(range(6))],rank=9,seed=seed)
    wd={'xtr':d['xtr'][:cut],'ytr':d['ytr'][:cut],'xte':d['xtr'][cut:],'yte':d['ytr'][cut:],'task':'regression'}
    warm,_,_=train_model(warm,wd,epochs=5,lr=lr,batch=128)
    with torch.no_grad():
        dev=next(warm.parameters()).device
        pred=warm(d['xtr'][cut:].to(dev)).cpu()
        residual=(d['ytr'][cut:]-pred).numpy()
    discovered,cov=cluster_residuals(residual,3)
    groups=[list(range(6))] if mode=='baseline' else discovered
    rank=9 if mode=='baseline' else 3
    ds={'xtr':d['xtr'][cut:],'ytr':d['ytr'][cut:],'xte':d['xte'],'yte':d['yte'],'task':'regression'}
    net=AdapterNet(groups,rank=rank,seed=seed)
    _,metric,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=128)
    if return_sig: return metric,{'groups':discovered,'all_tc':tc(cov),'within_tc':float(np.mean([tc(cov[np.ix_(g,g)]) for g in discovered]))}
    return metric

def baseline_fn(cfg): return lambda seed: fit_metric(seed,cfg['lr'],'baseline')
def idea_fn(cfg): return lambda seed: fit_metric(seed,cfg['lr'],'idea')

def main():
    base=sweep_baseline(baseline_fn,GRID,seeds=(0,1,2,3))
    candidates=[(evaluate(idea_fn(c),seeds=SEEDS),c) for c in GRID]
    idea,cfg=min(candidates,key=lambda z:z[0]['mean'])
    sig=[fit_metric(s,cfg['lr'],'idea',True)[1] for s in SEEDS]
    rec=float(np.mean([x['groups']==[[0,1],[2,3],[4,5]] for x in sig]))
    extra={'prediction':'residual-TC clustering discovers correlated task pairs and lowers within-cluster residual dependence','predicted_mean_all_tc':float(np.mean([x['all_tc'] for x in sig])),'observed_mean_within_cluster_tc':float(np.mean([x['within_tc'] for x in sig])),'pair_recovery_rate':rec,'confirmed':bool(rec>=0.5)}
    rep=make_report('custom_correlated_multitask_regression','mlp_tiny',base,idea,extra)
    rep['custom_track']={'name':'correlated_multitask_regression','file':'custom_multitask_track.py','domain':'multi_task_learning'}
    rep['idea_sweep']=[{'cfg':c,'mean':r['mean']} for r,c in candidates]
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()