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

# CPU-scale, deterministic benchmark for normalized class prototypes.
torch.set_num_threads(4)
SEED = 123

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)

def normalize(x):
    return x / x.norm(dim=1, keepdim=True).clamp_min(1e-12)

def math_check():
    seed_all(SEED)
    n, d = 32, 24
    x = normalize(torch.randn(n, d))
    g = torch.randn(n, d)
    gx = g - (g*x).sum(1, keepdim=True)*x
    v = torch.randn(n, d)
    v = v - (v*x).sum(1, keepdim=True)*x
    eta = 0.17
    trial = normalize(x + eta*v)
    v2 = v - (v*trial).sum(1, keepdim=True)*trial
    # P is symmetric/idempotent, projected vectors are tangent, and retraction is unit norm.
    px = g - (g*x).sum(1, keepdim=True)*x
    return {
        "projector_idempotence": float((px - (px - (px*x).sum(1, keepdim=True)*x)).abs().max()),
        "projected_tangent_max": float((gx*x).sum(1).abs().max()),
        "retraction_norm_error": float((trial.norm(dim=1)-1).abs().max()),
        "velocity_tangent_max": float((v2*trial).sum(1).abs().max()),
    }

def make_data(seed, classes=20, d=24, train_per=80, test_per=40):
    gen = torch.Generator().manual_seed(seed)
    centers = normalize(torch.randn(classes, d, generator=gen))
    def sample(per, noise):
        y = torch.arange(classes).repeat_interleave(per)
        z = centers[y] + noise*torch.randn(len(y), d, generator=gen)
        return normalize(z), y
    return centers, sample(train_per, .55), sample(test_per, .55)

def run(method, seed, steps=220, lr=.28, beta=.9, scale=12.0):
    centers, (train_x, train_y), (test_x, test_y) = make_data(seed)
    # Same initialization and data for both methods.
    gen = torch.Generator().manual_seed(seed+99)
    p0 = normalize(torch.randn(20, 24, generator=gen))
    p = p0.clone()
    v = torch.zeros_like(p)
    losses=[]; accs=[]; norm_err=[]; tangent=[]
    for step in range(steps):
        logits = scale * train_x @ p.T
        # Full batch makes optimizer differences easy to interpret and reproducible.
        loss = torch.nn.functional.cross_entropy(logits, train_y)
        # Make prototypes leaf tensors so gradients are computed independently each step.
        p = p.detach().requires_grad_(True)
        loss = torch.nn.functional.cross_entropy(scale * train_x @ p.T, train_y)
        g, = torch.autograd.grad(loss, p)
        p = p.detach()
        if method == "baseline":
            # Conventional momentum in ambient coordinates, followed by ad hoc normalization.
            v = beta*v - lr*g
            p = normalize(p + v)
        else:
            # Riemannian inertial update: project force, integrate tangent velocity,
            # retract, then remove radial velocity component at the new point.
            gp = g - (g*p).sum(1, keepdim=True)*p
            v = beta*v - lr*gp
            p = normalize(p + v)
            v = v - (v*p).sum(1, keepdim=True)*p
        with torch.no_grad():
            te = torch.nn.functional.cross_entropy(scale * test_x @ p.T, test_y)
            ac = (scale * test_x @ p.T).argmax(1).eq(test_y).float().mean()
            losses.append(float(loss)); accs.append(float(ac)); norm_err.append(float((p.norm(dim=1)-1).abs().max()))
            tangent.append(float((v*p).sum(1).abs().max()))
    return {"final_train_loss": losses[-1], "final_test_loss": float(te),
            "final_accuracy": accs[-1], "best_accuracy": max(accs),
            "step_to_70pct": next((i+1 for i,a in enumerate(accs) if a>=.70), None),
            "norm_error_max": max(norm_err), "velocity_tangent_max": max(tangent),
            "loss_curve": losses, "accuracy_curve": accs}

def main():
    check = math_check()
    all_results={"math_check":check, "runs":{}}
    for method in ["baseline", "idea"]:
        rs=[run(method, s) for s in [10,20,30]]
        def avg(k):
            vals=[r[k] for r in rs if r[k] is not None]
            return float(np.mean(vals)) if vals else None
        all_results["runs"][method]={k:avg(k) for k in ["final_train_loss","final_test_loss","final_accuracy","best_accuracy","step_to_70pct","norm_error_max","velocity_tangent_max"]}
        all_results["runs"][method]["per_seed"]= [{k:r[k] for k in ["final_train_loss","final_accuracy","step_to_70pct"]} for r in rs]
    print(json.dumps(all_results, indent=2))

if __name__ == '__main__': main()