import json, math, random
from pathlib import Path
import numpy as np

SEED=1234
rng=np.random.default_rng(SEED)


def cosine_dist(a,b):
    a=np.asarray(a,float); b=np.asarray(b,float)
    return 1-float(a@b)/(np.linalg.norm(a)*np.linalg.norm(b)+1e-12)


def mutual_reachability(D,m=2):
    n=len(D); cores=np.zeros(n)
    for i in range(n):
        vals=np.sort(np.delete(D[i],i)); cores[i]=vals[min(m-1,len(vals)-1)]
    MR=np.maximum(D,np.maximum(cores[:,None],cores[None,:])); np.fill_diagonal(MR,0)
    return cores,MR


def mst_cut_clusters(MR):
    # Kruskal MST, cut the largest unusually separated edge (density hierarchy toy).
    n=len(MR); edges=sorted((MR[i,j],i,j) for i in range(n) for j in range(i))
    par=list(range(n))
    def find(x):
        while par[x]!=x: par[x]=par[par[x]]; x=par[x]
        return x
    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([x[0] for x in mst])
    if len(ws)<2: return np.arange(n)
    gaps=np.diff(np.sort(ws)); k=int(np.argmax(gaps))
    cut=float(np.sort(ws)[k+1])
    # only split if gap is meaningful; this is the singleton/noise-safe hierarchy cut
    if cut <= (np.median(ws)+1e-12)*1.15: cut=float('inf')
    par=list(range(n))
    for w,i,j in mst:
        if w<cut:
            a,b=find(i),find(j)
            if a!=b: par[a]=b
    roots={}; out=[]
    for i in range(n):
        r=find(i); roots.setdefault(r,len(roots)); out.append(roots[r])
    return np.array(out)


def math_checks():
    u=np.array([1.,0.,0.]); v=np.array([0.,1.,0.]);
    exact={"same":cosine_dist(u,u),"orthogonal":cosine_dist(u,v),"antipodal":cosine_dist(u,-u)}
    # For e=normalize(mu + sigma*z), small-noise prediction E[d] ~= (d-1)sigma^2.
    d=8; mu=np.zeros(d); mu[0]=1
    rows=[]
    for s in [0.01,0.03,0.06,0.10,0.18,0.30]:
        vals=[]
        for _ in range(30000):
            a=mu+s*rng.normal(size=d); b=mu+s*rng.normal(size=d)
            vals.append(cosine_dist(a,b))
        obs=float(np.mean(vals)); pred=(d-1)*s*s
        rows.append({"sigma":s,"observed":obs,"predicted":pred,"ratio":obs/pred})
    return exact,rows


def clustering_check():
    d=8; centers=np.zeros((2,d)); centers[0,0]=1; centers[1,1]=1
    rows=[]
    for s in [0.01,0.05,0.10,0.20,0.35,0.60]:
        E=[]; labels=[]
        for c in range(2):
            for _ in range(8):
                E.append(centers[c]+s*rng.normal(size=d)); labels.append(c)
        E=np.asarray(E); E/=np.linalg.norm(E,axis=1,keepdims=True)
        D=1-E@E.T; np.fill_diagonal(D,0)
        core,MR=mutual_reachability(D,2); pred=float(np.mean([D[i,j] for i in range(8) for j in range(8,16)]))
        within=float(np.mean([D[i,j] for i in range(16) for j in range(i) if labels[i]==labels[j]]))
        # density hierarchy output; match up to cluster label permutation
        got=mst_cut_clusters(MR)
        same=(got[:8,None]==got[None,:8]).mean() # diagnostic only
        # pairwise purity, invariant to numeric cluster labels
        purity=np.mean([max(np.mean(got[np.array(labels)==c]==q) for q in set(got)) for c in [0,1]])
        rows.append({"sigma":s,"within_distance":within,"between_distance":pred,"between_minus_within":pred-within,"cluster_purity":float(purity),"mean_core":float(core.mean())})
    return rows


def gradient_embeddings(angles):
    # At the zero-output probe point, average trunk gradient is proportional to target vector.
    E=np.array([[math.cos(a),math.sin(a)] for a in angles],float)
    E/=np.linalg.norm(E,axis=1,keepdims=True)
    return E


def rank1_loss(angles, groups):
    # Exact population optimum for a scalar shared trunk and scalar task heads.
    # For each group, residual is the smaller eigenvalue of sum_i w_i w_i^T.
    total=0.
    for g in sorted(set(groups)):
        W=np.array([[math.cos(angles[i]),math.sin(angles[i])] for i in range(len(angles)) if groups[i]==g])
        ev=np.linalg.eigvalsh(W.T@W)
        total += float(ev[0])
    return total/len(angles)


def training_check():
    # Two task families: discovery should recover them from probe gradient embeddings.
    angles=np.array([0.00,0.10,0.18,0.25,1.30,1.40,1.48,1.58])
    E=gradient_embeddings(angles); D=1-E@E.T; np.fill_diagonal(D,0)
    core,MR=mutual_reachability(D,m=2); discovered=mst_cut_clusters(MR)
    shared=np.zeros(len(angles),dtype=int)
    oracle=np.array([0]*4+[1]*4)
    # Compare equal analytic population objective; add finite-sample SGD confirmation.
    return {"angles":angles.tolist(),"discovered_groups":discovered.tolist(),"probe_core":core.tolist(),
            "shared_population_loss":rank1_loss(angles,shared),
            "discovered_population_loss":rank1_loss(angles,discovered),
            "oracle_population_loss":rank1_loss(angles,oracle),
            "conflict_sweep":conflict_sweep()}


def conflict_sweep():
    rows=[]
    for theta in [0,.2,.4,.6,.8,1.0,1.2,math.pi/2]:
        a=np.array([0.,theta]); shared=rank1_loss(a,[0,0]); part=rank1_loss(a,[0,1]);
        pred=(1-abs(math.cos(theta)))/2 # per-task average, two unit tasks
        rows.append({"angle":theta,"shared_loss":shared,"partition_loss":part,"observed_gain":shared-part,"predicted_gain":pred})
    return rows


def main():
    exact,noise=math_checks(); clusters=clustering_check(); train=training_check()
    out={"seed":SEED,"math_exact":exact,"noise_scaling":noise,"density_sweep":clusters,"training":train,
         "predictions":{"P1":"cosine distance is 0/1/2 for same/orthogonal/antipodal normalized embeddings",
         "P2":"for d=8 and small sigma, E[distance] ~= 7 sigma^2",
         "P3":"for two unit linear tasks at angle theta, partition gain per task = (1-|cos(theta)|)/2"}}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()