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

SEED = 2548

def seed_all(seed=SEED):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

def dpd_loss(logits, y, alpha, weights=None):
    # Exact categorical A_alpha=sum_c p_c^(1+alpha), averaged over examples/components.
    p = torch.softmax(logits.float(), dim=-1)
    a = float(alpha)
    A = (p.pow(1.0+a)).sum(dim=-1)
    observed = p.gather(-1, y.long().unsqueeze(-1)).squeeze(-1).clamp_min(1e-12)
    loss = A - (1.0 + 1.0/a) * observed.pow(a)
    if weights is not None:
        loss = loss * torch.as_tensor(weights, device=loss.device, dtype=loss.dtype)
    return loss.mean()

def observed_term_loss(logits, y, alpha):
    p = torch.softmax(logits.float(), dim=-1)
    q = p.gather(-1, y.long().unsqueeze(-1)).squeeze(-1).clamp_min(1e-12)
    return -(1.0 + 1.0/alpha) * q.pow(alpha)

def math_sweep():
    # Prediction 1: observed-score DPD/CE gradient norm ratio=(1+alpha)q^alpha.
    qs = np.array([1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1])
    alphas = [0.1, 0.3, 0.5, 1.0]
    rows=[]
    for a in alphas:
        ratios=[]
        for q in qs:
            # logits chosen so class 0 has probability q, remaining mass uniform
            logits = torch.tensor([[math.log(q), math.log((1-q)/2), math.log((1-q)/2)]], requires_grad=True)
            y=torch.tensor([0])
            g= torch.autograd.grad(observed_term_loss(logits,y,a).sum(), logits)[0].norm().item()
            ce= torch.autograd.grad((-torch.log(torch.softmax(logits, -1)[:,0])).sum(), logits, retain_graph=True)[0].norm().item()
            ratios.append(g/ce)
        slope=np.polyfit(np.log(qs), np.log(ratios), 1)[0]
        pred=np.array([(1+a)*q**a for q in qs])
        rel=float(np.max(np.abs(np.array(ratios)-pred)/(pred+1e-12)))
        rows.append({'alpha':a,'max_relative_error':rel,'loglog_slope':float(slope),'predicted_slope':a,
                     'ratio_at_q_1e-3':float(ratios[2]),'predicted_ratio_at_q_1e-3':float((1+a)*1e-3**a)})
    # Prediction 2: alpha=0 limit of full DPD gradient approaches CE gradient.
    logits=torch.tensor([[1.2,-.7,.1]],dtype=torch.float64,requires_grad=True); y=torch.tensor([2])
    ce=torch.autograd.grad((-torch.log_softmax(logits,-1)[:,2]).sum(),logits,retain_graph=True)[0]
    limit=[]
    for a in [0.5,0.2,0.1,0.05,0.02,0.01]:
        l=dpd_loss(logits.float(),y,a)
        g=torch.autograd.grad(l,logits,retain_graph=True)[0].double()
        limit.append({'alpha':a,'gradient_relative_error':float((g-ce).norm()/ce.norm())})
    return {'observed_score_scaling':rows,'alpha_to_zero_full_gradient':limit}

def make_data(n=1800, noise=0.20, seed=SEED):
    rng=np.random.default_rng(seed)
    centers=np.array([[-1.4,-1.0],[1.4,-1.0],[0,1.5]],dtype=np.float32)
    y=rng.integers(0,3,n); x=centers[y]+rng.normal(0,.85,(n,2)).astype(np.float32)
    clean=y.copy(); noisy=y.copy(); mask=rng.random(n)<noise
    for i in np.where(mask)[0]: noisy[i]=rng.integers(0,2) if y[i]==2 else 2 # deterministic different class
    idx=rng.permutation(n); tr=idx[:1200]; va=idx[1200:]
    return torch.tensor(x[tr]),torch.tensor(noisy[tr]),torch.tensor(clean[tr]),torch.tensor(x[va]),torch.tensor(clean[va])

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,3))
    def forward(self,x): return self.net(x)

def train(kind, alpha=.5, epochs=80):
    x,y,yclean,xv,yv=make_data()
    model=TinyNet(); opt=torch.optim.Adam(model.parameters(),lr=.025)
    for _ in range(epochs):
        perm=torch.randperm(len(x))
        for j in range(0,len(x),64):
            ix=perm[j:j+64]; logits=model(x[ix])
            loss=nn.functional.cross_entropy(logits,y[ix]) if kind=='ce' else dpd_loss(logits,y[ix],alpha)
            opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        pred=model(xv).argmax(1); acc=(pred==yv).float().mean().item()
        train_pred=model(x).argmax(1); train_acc=(train_pred==yclean).float().mean().item()
    # Per-example gradient norm for clean-vs-corrupted labels at the trained model.
    corrupted=(y!=yclean); vals=[]
    for i in range(len(x)):
        model.zero_grad(set_to_none=True); z=model(x[i:i+1])
        l=nn.functional.cross_entropy(z,y[i:i+1]) if kind=='ce' else dpd_loss(z,y[i:i+1],alpha)
        l.backward(); vals.append(sum((p.grad.detach()**2).sum().item() for p in model.parameters() if p.grad is not None)**.5)
    vals=np.array(vals)
    return {'clean_validation_accuracy':acc,'clean_train_accuracy':train_acc,
            'corrupted_gradient_norm':float(vals[corrupted.numpy()].mean()),
            'clean_gradient_norm':float(vals[~corrupted.numpy()].mean()),
            'gradient_corrupted_to_clean':float(vals[corrupted.numpy()].mean()/vals[~corrupted.numpy()].mean())}

def main():
    seed_all(); math_results=math_sweep()
    # same data/model initialization stream is reset for an apples-to-apples comparison.
    seed_all(); ce=train('ce'); seed_all(); dpd=train('dpd',.5)
    out={'seed':SEED,'math_verification':math_results,'toy_training':{'baseline_CE':ce,'idea_DPD_alpha_0.5':dpd}}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()