import json, math, random
import numpy as np
import torch

SEED = 620
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DTYPE = torch.float64


def gm_utility(load, C, eps=1e-12):
    # Stable log-space implementation of prod_j (c_j dot load)^(1/d).
    factors = load @ C.T
    return torch.exp(torch.log(factors.clamp_min(eps)).mean())


def concavity_check():
    E, d = 5, 5
    C = torch.full((d, E), 0.05, dtype=DTYPE)
    C += torch.eye(E, dtype=DTYPE)
    worst = 1.0
    jensen_gaps = []
    for _ in range(1000):
        x = torch.rand(E, dtype=DTYPE) + .02
        y = torch.rand(E, dtype=DTYPE) + .02
        t = float(torch.rand(1))
        lhs = gm_utility(t*x + (1-t)*y, C).item()
        rhs = t*gm_utility(x, C).item() + (1-t)*gm_utility(y, C).item()
        gap = rhs-lhs
        jensen_gaps.append(gap)
    # Hessian eigenvalues at an interior point, with unconstrained positive loads.
    x = (torch.rand(E, dtype=DTYPE)+.2).requires_grad_(True)
    H = torch.autograd.functional.hessian(lambda z: gm_utility(z, C), x).detach().numpy()
    eig = np.linalg.eigvalsh((H+H.T)/2)
    return {"max_jensen_violation": float(max(jensen_gaps)),
            "mean_jensen_gap": float(np.mean(jensen_gaps)),
            "max_hessian_eigenvalue": float(eig.max()),
            "min_hessian_eigenvalue": float(eig.min())}


def optimize(method, seed=SEED, steps=500):
    torch.manual_seed(seed)
    T, E, K = 256, 8, 2
    # Two task groups have distinct preferred experts, creating a real
    # specialization-vs-balance tradeoff.
    group = torch.arange(T) % K
    target = torch.zeros(T, E, dtype=DTYPE)
    target[:, :] = 0.025/(E-1)
    target[group == 0, 0] = .975
    target[group == 1, 1] = .975
    # Router logits are per-token parameters: this isolates the auxiliary
    # balancing behavior from architecture and optimizer confounders.
    logits = (0.1*torch.randn(T, E, dtype=DTYPE)).requires_grad_(True)
    opt = torch.optim.Adam([logits], lr=.12)
    C = torch.full((E, E), .03, dtype=DTYPE) + torch.eye(E, dtype=DTYPE)
    history=[]
    for step in range(steps):
        opt.zero_grad()
        r = torch.softmax(logits, dim=1)
        load = r.mean(0)
        task = -(target * torch.log(r.clamp_min(1e-12))).sum(1).mean()
        entropy = -(r * torch.log(r.clamp_min(1e-12))).sum(1).mean()
        switch = E * (load*load).sum() # differentiable proxy for Switch load loss
        if method == "entropy":
            aux = -0.08*entropy
        elif method == "switch":
            aux = 0.08*switch
        elif method == "gm":
            aux = -0.55*torch.log(gm_utility(load, C).clamp_min(1e-12))
        elif method == "gm_symmetric":
            Cs = torch.ones((2,E), dtype=DTYPE)
            aux = -0.55*torch.log(gm_utility(load, Cs))
        loss = task + aux
        loss.backward(); opt.step()
        if step % 50 == 0 or step == steps-1:
            with torch.no_grad():
                hard = r.argmax(1)
                counts = torch.bincount(hard, minlength=E).double()/T
                cv = counts.std(unbiased=False)/(counts.mean()+1e-12)
                history.append({"step":step,"loss":float(loss),"task":float(task),
                                "entropy":float(entropy),"cv_soft":float(load.std(unbiased=False)/load.mean()),
                                "cv_hard":float(cv),"drop_rate_capacity_1":float(torch.clamp(counts*E-1,min=0).sum()/E)})
    return history[-1], history


def main():
    conc = concavity_check()
    results={}
    for method in ["entropy","switch","gm_symmetric","gm"]:
        final, hist=optimize(method)
        results[method]={"final":final,"history":hist}
    # Directly demonstrate the literal c=1 degeneracy.
    load=torch.tensor([.01,.05,.1,.12,.15,.17,.19,.21],dtype=DTYPE,requires_grad=True)
    C=torch.ones((2,8),dtype=DTYPE)
    u=gm_utility(load,C); u.backward()
    results["symmetric_gradient"]={"utility":float(u),"gradient_norm":float(load.grad.norm()),"gradient":load.grad.detach().tolist()}
    out={"seed":SEED,"concavity":conc,"results":results}
    with open("results.json","w") as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__ == "__main__": main()