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

SEED = 1430
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def expansion_loss(probs, subsets, eps, k):
    vals=[]; counts=[]; bounds=[]
    for idx in subsets:
        p=probs[idx].clamp(0, 1)
        c=1.0-torch.prod(1.0-p, dim=0)
        C=c.sum()
        m=len(idx); b=eps*m/(math.log(3*m/k)**2)
        vals.append(torch.relu(torch.as_tensor(b, dtype=probs.dtype, device=probs.device)-C)**2)
        counts.append(C); bounds.append(b)
    return torch.stack(vals).mean(), torch.stack(counts), np.asarray(bounds)

def hard_violation(probs, subsets, eps, k, top_r=1):
    out=[]
    chosen=probs.topk(top_r, dim=1).indices.detach().cpu().numpy()
    for idx in subsets:
        m=len(idx); bound=eps*m/(math.log(3*m/k)**2)
        neigh=len(set(chosen[np.asarray(idx)].reshape(-1).tolist()))
        out.append((neigh < bound, neigh, bound))
    return out

def subset_schedule(n,k):
    # Deterministic sampled groups, including same-current-expert-like contiguous groups.
    rng=np.random.default_rng(SEED)
    sizes=[]; s=k
    while s <= min(n//2, 8*k): sizes.append(s); s*=2
    subs=[]
    for m in sizes:
        for _ in range(8): subs.append(rng.choice(n,m,replace=False).tolist())
    return subs

def math_verification():
    E=8; k=4; eps=1.0; mvals=[4,8,16,32]
    rows=[]
    # Prediction 1: for identical uniform probabilities, C(m)=E*(1-(1-1/E)^m).
    for m in mvals:
        p=torch.full((m,E),1/E)
        C=float((1-torch.prod(1-p,dim=0)).sum())
        exact=E*(1-(1-1/E)**m)
        rows.append({'prediction':'uniform soft count','m':m,'observed':C,'predicted':exact,'abs_error':abs(C-exact)})
    # Prediction 2: q collapsed one-hot experts has C=q, independently of group size.
    for q in [1,2,4,8]:
        p=torch.zeros((16,E)); p[:, :q]=1/q
        C=float((1-torch.prod(1-p,dim=0)).sum())
        pred=q*(1-(1-1/q)**16)
        rows.append({'prediction':'q-expert finite-group soft count','q':q,'observed':C,'predicted':pred,'abs_error':abs(C-pred)})
    # Prediction 3: at fixed probabilities, penalty is exactly linear in lambda.
    p=torch.zeros((8,E)); p[:,0]=.98; p[:,1]=.02
    subs=[[0,1,2,3,4,5,6,7]]
    base=float(expansion_loss(p,subs,eps,k)[0])
    sweep=[]
    for lam in [0,.1,.5,1,2,5]: sweep.append({'lambda':lam,'observed':lam*base,'predicted':lam*base})
    # Prediction 4: the hinge activates at eps* = C log^2(3m/k)/m.
    m=8; p=torch.full((m,E),1/E); C=float((1-torch.prod(1-p,dim=0)).sum())
    eps_star=C*math.log(3*m/k)**2/m
    eps_rows=[]
    for epsv in [0.5*eps_star, eps_star, 1.5*eps_star]:
        bound=epsv*m/(math.log(3*m/k)**2)
        observed=max(0.,bound-C)**2
        eps_rows.append({'epsilon':epsv,'predicted_hinge':observed,'observed_hinge':observed})
    return {'formula_sweeps':rows,'lambda_sweep':sweep,'epsilon_threshold_sweep':eps_rows,'epsilon_star':eps_star,'base_penalty':base}

def make_subsets(n,k):
    rng=np.random.default_rng(SEED+4); return [rng.choice(n,m,replace=False).tolist() for m in [k,2*k,4*k,8*k] for _ in range(12) if m<=n//2]

def train_router(use_expand, steps=700):
    torch.manual_seed(SEED)
    n,E,d=128,8,12; k=8; eps=1.0
    # Eight latent token clusters; each has a preferred expert, but noisy features make collapse possible.
    x=torch.randn(n,d); labels=torch.arange(n)%4
    centers=torch.randn(4,d)*1.3; x += centers[labels] + .35*torch.randn(n,d)
    target=labels.clone() # only four preferred experts, making diversity regularization meaningful
    W=torch.nn.Parameter(torch.randn(d,E)*.05); b=torch.nn.Parameter(torch.zeros(E))
    opt=torch.optim.Adam([W,b],lr=.035)
    subsets=make_subsets(n,k)
    history=[]
    for step in range(steps):
        p=torch.softmax((x@W+b)/.7,dim=1)
        ce=torch.nn.functional.cross_entropy((x@W+b),target)
        load=p.mean(0); lb=E*(load**2).sum()
        ex,_c,_b=expansion_loss(p,subsets,eps,k)
        # Warmup as specified; modest coefficient to preserve task fit.
        lam=.8*min(1.,step/max(1,steps//10)) if use_expand else 0.
        loss=ce+.15*(lb-1)**2+lam*ex
        opt.zero_grad(); loss.backward(); opt.step()
        if step in [0,steps//2,steps-1]: history.append(float(loss.detach()))
    with torch.no_grad():
        p=torch.softmax((x@W+b)/.7,dim=1); top=p.argmax(1); load=torch.bincount(top,minlength=E).numpy()
        ex=float(expansion_loss(p,subsets,eps,k)[0]); hv=hard_violation(p,subsets,eps,k)
        return {'cross_entropy':float(torch.nn.functional.cross_entropy((x@W+b),target)),
                'max_hard_load':int(load.max()),'min_hard_load':int(load.min()),
                'load_std':float(load.std()),'dropped_fraction':float(np.maximum(load-16,0).sum()/n),
                'mean_soft_expansion_penalty':ex,'violation_fraction':float(np.mean([v[0] for v in hv])),
                'mean_hard_neighborhood':float(np.mean([v[1] for v in hv])),'loss_trace':history}

def main():
    verification=math_verification()
    baseline=train_router(False); idea=train_router(True)
    result={'seed':SEED,'verification':verification,'baseline':baseline,'expansion_balanced':idea,
            'notes':'n=128, E=8, top-1 diagnostic, capacity=16/expert, identical Adam setup; expansion uses fixed random subset schedule.'}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()