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

SEED = 1726
torch.set_default_dtype(torch.float64)
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def normalize(z):
    return z / (torch.linalg.vector_norm(z, dim=-1, keepdim=True) + 1e-12)

def legendre_all(x, L):
    vals = [torch.ones_like(x)]
    if L >= 1: vals.append(x)
    for l in range(2, L + 1):
        vals.append(((2*l-1)*x*vals[-1] - (l-1)*vals[-2]) / l)
    return vals

def spectrum(u, L=4):
    # Addition theorem: sum_m Y_lm(u_i)Y*_lm(u_j)=(2l+1)P_l(dot)/(4pi).
    # Hence P_l=(1/(4pi))*sum_ij P_l(u_i dot u_j).
    dots = torch.clamp(u @ u.T, -1.0, 1.0)
    return torch.stack([v.sum()/(4*math.pi) for v in legendre_all(dots, L)])

def spec_loss_z(z, L=4):
    u = normalize(z); p = spectrum(u, L)
    q = p[1:] / p[0].detach()
    return (q*q).sum(), p, q

def explicit_l1_power(u):
    a = math.sqrt(3/(4*math.pi)) * u.sum(dim=0)
    return (a*a).sum()/3.0

def optimize(lam, steps=700, seed=SEED, N=64, L=4):
    torch.manual_seed(seed)
    z = torch.randn(N, 3)*0.35; z[:,2] += 1.0; z.requires_grad_()
    opt = torch.optim.Adam([z], lr=0.035); trace = {}
    for step in range(steps+1):
        u = normalize(z); task = -u[:,2].mean(); sl,p,q = spec_loss_z(z,L)
        if step in (0,10,50,200,700): trace[str(step)] = [float(task),float(sl),float(q[0])]
        if step < steps:
            opt.zero_grad(); (task + lam*sl).backward(); opt.step()
    with torch.no_grad():
        u=normalize(z); task=-u[:,2].mean(); sl,p,q=spec_loss_z(z,L)
        dirs=normalize(torch.randn(100,3)); counts=torch.stack([(u@d >= math.cos(math.pi/4)).double().sum() for d in dirs])
    return {'lambda':lam,'task':float(task),'spec_loss':float(sl),'q':q.tolist(),'cap_var':float(counts.var(unbiased=False)),'trace':trace}

def numerical_checks():
    torch.manual_seed(8); z=(torch.randn(11,3)*.4).requires_grad_()
    u=normalize(z); p=spectrum(u,4); p1=explicit_l1_power(u)
    sl,_,_=spec_loss_z(z,4); sl.backward(); g=z.grad.detach().clone()
    d=torch.randn_like(z); eps=1e-5
    with torch.no_grad():
        plus=spec_loss_z(z.detach()+eps*d,4)[0]; minus=spec_loss_z(z.detach()-eps*d,4)[0]
    fd=float((plus-minus)/(2*eps)); ad=float((g*d).sum())
    return {'addition_theorem_l1_abs_error':float(abs(p1-p[1])),'gradient_fd':fd,'gradient_autodiff':ad,'gradient_relative_error':abs(fd-ad)/(abs(fd)+1e-12)}

def gradient_scaling():
    torch.manual_seed(91); z=(torch.randn(64,3)*.35); z[:,2]+=1; z.requires_grad_()
    sl,_,_=spec_loss_z(z,4); sl.backward(); gs=z.grad.norm().item()
    z2=z.detach().clone().requires_grad_(); (-normalize(z2)[:,2].mean()).backward(); gt=z2.grad.norm().item()
    return {'spectrum_grad_norm':gs,'task_grad_norm':gt,'lambda_times_grad_ratio':{str(x):x*gs/gt for x in [0,.01,.1,1]}}

def main():
    sweep=[optimize(x) for x in [0,.001,.003,.01,.03,.1,.3,1]]
    conv=[optimize(.1,steps=s)['spec_loss'] for s in [0,10,50,200,700]]
    print(json.dumps({'checks':numerical_checks(),'gradient_scaling':gradient_scaling(),'sweep':sweep,'convergence_lambda_0.1':conv,'notes':'q_l=P_l/P_0, target q_1..q_4=0; task favors north-pole collapse.'},indent=2))

if __name__=='__main__': main()