import sys, json, random
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, reload_custom_tracks
import importlib.util
_spec=importlib.util.spec_from_file_location('matched_track','/home/maxwelhelp/all/math2nn/bench/custom_tracks/correlated_multitask_regression.py')
_track=importlib.util.module_from_spec(_spec); _spec.loader.exec_module(_track)

def get_ds(seed):
    d0=_track.get_dataset(seed=seed,n_train=400,n_test=400)
    return {'track':TRACK,'task':d0.get('task','regression'),'metric':d0.get('metric','mse'),
            'xtr':torch.as_tensor(np.asarray(d0['xtr']),dtype=torch.float32),
            'ytr':torch.as_tensor(np.asarray(d0['ytr']),dtype=torch.float32),
            'xte':torch.as_tensor(np.asarray(d0['xte']),dtype=torch.float32),
            'yte':torch.as_tensor(np.asarray(d0['yte']),dtype=torch.float32),
            'input_shape':tuple(np.asarray(d0['xtr']).shape[1:]),'out_dim':int(np.asarray(d0['ytr']).shape[1])}


TRACK = 'correlated_multitask_regression'
MODEL = 'mlp_tiny'
SEEDS = [11,22,33,44,55,66,77,88]
SWEEP_SEEDS = [11,22,33]
LR_GRID = [1e-3, 3e-3, 1e-2]
EPOCHS = 22
PROBE_EPOCHS = 4
BATCH = 128
WIDTH = 64

class SharedMTL(nn.Module):
    def __init__(self, p, width=WIDTH):
        super().__init__()
        self.trunk = nn.Sequential(nn.Linear(p,width), nn.ReLU(), nn.Linear(width,width), nn.ReLU())
        self.heads = nn.ModuleList([nn.Linear(width,1) for _ in range(6)])
    def features(self, x): return self.trunk(x)
    def forward(self, x):
        z = self.features(x)
        return torch.cat([h(z) for h in self.heads], 1)

class PartitionMTL(nn.Module):
    def __init__(self, p, groups, width=WIDTH):
        super().__init__()
        self.groups = list(map(int, groups)); self.ng = max(self.groups)+1
        self.trunks = nn.ModuleList([nn.Sequential(nn.Linear(p,width), nn.ReLU(), nn.Linear(width,width), nn.ReLU()) for _ in range(self.ng)])
        self.heads = nn.ModuleList([nn.Linear(width,1) for _ in range(6)])
    def forward(self, x):
        return torch.cat([self.heads[i](self.trunks[self.groups[i]](x)) for i in range(6)], 1)

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)

def loss_for(out, y): return ((out-y)**2).mean()

def probe_groups(ds, seed):
    seed_all(seed)
    p = ds['xtr'].shape[1]
    trunk = nn.Sequential(nn.Linear(p, WIDTH), nn.ReLU(), nn.Linear(WIDTH, WIDTH), nn.ReLU())
    heads = nn.ModuleList([nn.Linear(WIDTH,1) for _ in range(6)])
    opt = torch.optim.Adam(list(trunk.parameters())+list(heads.parameters()), lr=3e-3)
    x,y=ds['xtr'],ds['ytr']
    for _ in range(PROBE_EPOCHS):
        for j in range(0,len(x),BATCH):
            z=trunk(x[j:j+BATCH]); out=torch.cat([h(z) for h in heads],1)
            loss=loss_for(out,y[j:j+BATCH]); opt.zero_grad(); loss.backward(); opt.step()
    emb=[]
    idx=torch.arange(min(256,len(x)))
    for i,h in enumerate(heads):
        trunk.zero_grad(set_to_none=True); h.zero_grad(set_to_none=True)
        z=trunk(x[idx]); li=((h(z)-y[idx,i:i+1])**2).mean(); li.backward()
        act=z.detach().mean(0).numpy()
        grad=np.concatenate([p.grad.detach().numpy().ravel() for p in trunk.parameters() if p.grad is not None])
        grad=grad/(np.linalg.norm(grad)+1e-12)
        e=np.r_[act,grad]; e=e/(np.linalg.norm(e)+1e-12); emb.append(e)
    E=np.asarray(emb); D=1-E@E.T; np.fill_diagonal(D,0)
    m=2; core=np.zeros(6)
    for i in range(6): core[i]=np.sort(np.delete(D[i],i))[m-1]
    MR=np.maximum(D,np.maximum(core[:,None],core[None,:]))
    edges=sorted((MR[i,j],i,j) for i in range(6) for j in range(i))
    par=list(range(6))
    def find(a):
        while par[a]!=a:
            par[a]=par[par[a]]; a=par[a]
        return a
    mst=[]
    for w,i,j in edges:
        a,b=find(i),find(j)
        if a!=b: par[a]=b; mst.append((w,i,j))
    ws=np.array([z[0] for z in mst]); cut=float('inf')
    if len(ws)>=2:
        s=np.sort(ws); k=int(np.argmax(np.diff(s)))
        if s[k+1] > 1.15*np.median(s): cut=float(s[k+1])
    par=list(range(6))
    for w,i,j in mst:
        if w<cut:
            a,b=find(i),find(j)
            if a!=b: par[a]=b
    roots={}; groups=[]
    for i in range(6):
        r=find(i); roots.setdefault(r,len(roots)); groups.append(roots[r])
    return groups, float(np.mean(D)), float(np.mean(core))

def train_baseline(cfg, seed, return_aux=False):
    seed_all(seed); ds=get_ds(seed)
    net=SharedMTL(ds['xtr'].shape[1]); net,metric,hist=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
    if return_aux: return metric, {'groups':[0]*6,'probe':None,'model':net,'ds':ds}
    return metric

def train_idea(cfg, seed, return_aux=False):
    seed_all(seed); ds=get_ds(seed)
    groups,md,mc=probe_groups(ds,seed)
    seed_all(seed)
    net=PartitionMTL(ds['xtr'].shape[1],groups)
    net,metric,hist=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
    if return_aux: return metric, {'groups':groups,'probe_mean_distance':md,'probe_mean_core':mc,'model':net,'ds':ds}
    return metric

def signature(seed, base_metric, idea_metric, groups):
    # Re-test the mechanism on trained systems: task-gradient cosine conflict.
    def conflicts(kind):
        ds=get_ds(seed); seed_all(seed)
        _,aux=(train_baseline({'lr':3e-3},seed,True) if kind=='base' else train_idea({'lr':3e-3},seed,True))
        net,ds=aux['model'],aux['ds']; device=next(net.parameters()).device; x,y=ds['xtr'][:128].to(device),ds['ytr'][:128].to(device)
        gs=[]
        for i in range(6):
            net.zero_grad(set_to_none=True); out=net(x); ((out[:,i:i+1]-y[:128,i:i+1])**2).mean().backward()
            ps=[]
            for p in net.parameters():
                if p.grad is not None: ps.append(p.grad.detach().flatten())
            gs.append(torch.cat(ps))
        C=torch.stack(gs); cs=[]
        for i in range(6):
            for j in range(i): cs.append(float(torch.dot(C[i],C[j])/(C[i].norm()*C[j].norm()+1e-12)))
        return float(np.mean(np.asarray(cs)<0))
    cb=conflicts('base'); ci=conflicts('idea')
    return {'predicted':'partitioning reduces negative cross-task gradient cosine','baseline_conflict_rate':cb,'idea_conflict_rate':ci,'observed_reduction':cb-ci,'confirmed':bool(ci < cb)}

def main():
    reload_custom_tracks()
    grid=[{'lr':x} for x in LR_GRID]
    base=sweep_baseline(lambda cfg: lambda s: train_baseline(cfg,s),grid,seeds=SWEEP_SEEDS)
    idea_grid={x: evaluate(lambda s,c={'lr':x}: train_idea(c,s),seeds=SWEEP_SEEDS) for x in LR_GRID}
    best_lr=min(idea_grid,key=lambda x: idea_grid[x]['mean'])
    idea=evaluate(lambda s: train_idea({'lr':best_lr},s),seeds=SEEDS)
    # Also run two nearby settings on the same sweep union, already evaluated above.
    sig=signature(SEEDS[0],base['full']['mean'],idea['mean'],train_idea({'lr':best_lr},SEEDS[0],True)[1]['groups'])
    rep=make_report(TRACK,MODEL,base,idea,{'track_rationale':'multi-task regression with six task heads and latent task groups; custom track is structurally matched','idea_grid':idea_grid,'selected_idea_lr':best_lr,'probe_epochs':PROBE_EPOCHS,'groups_seed0':train_idea({'lr':best_lr},SEEDS[0],True)[1]['groups'],'mechanism_signature':sig})
    rep['parameter_note']='Partitioned trunks replicate the trunk per discovered cluster; heads remain separate. Equal optimizer/data/epoch budget, but inference parameter count can differ.'
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()