import json
import math
from pathlib import Path
import numpy as np

SEED = 1362
rng = np.random.default_rng(SEED)


def normalize(x):
    return x / np.linalg.norm(x, axis=-1, keepdims=True)


def rope(x, positions, omegas):
    # x: [n,d], d even; rotations act independently on each pair.
    y = x.copy()
    for l, w in enumerate(omegas):
        a, b = 2*l, 2*l+1
        th = w * positions
        c, s = np.cos(th), np.sin(th)
        y[:, a] = c*x[:, a] - s*x[:, b]
        y[:, b] = s*x[:, a] + c*x[:, b]
    return y


def attention(x, positions, omegas, beta):
    q = rope(x, positions, omegas)
    scores = q @ q.T
    z = np.exp(beta * (scores - scores.max(axis=1, keepdims=True)))
    a = z / z.sum(axis=1, keepdims=True)
    return a, scores


def spherical_step(x, a, eta):
    m = a @ x
    v = m - (m*x).sum(axis=1, keepdims=True)*x
    return normalize(x + eta*v)


def angular_diameter(x):
    c = np.clip(x @ x.T, -1.0, 1.0)
    return float(np.arccos(c).max())


def make_consensus(n, d, positions, omegas, beta):
    u = normalize(rng.normal(size=(1,d)))[0]
    x = np.repeat(u[None,:], n, axis=0)
    a, scores = attention(x, positions, omegas, beta)
    return x, a, scores, u


def stability_trial(a, z, u, eta, steps=8, eps=1e-6):
    # At consensus, exp(beta*scores) is symmetric and A is reversible.
    # If S=D^(1/2) A D^(-1/2), S w=lambda w, then
    # A(D^(-1/2)w)=lambda(D^(-1/2)w). Use the worst transverse mode.
    dz = np.sqrt(z)
    s = dz[:, None] * a / dz[None, :]
    vals, vecs = np.linalg.eigh((s + s.T) / 2)
    lam_min = float(vals[0])
    mode = vecs[:, 0] / dz
    mode /= np.linalg.norm(mode)
    tangent = np.zeros_like(u)
    tangent[0] = 1.0
    if abs(np.dot(tangent, u)) > .9:
        tangent[:] = 0; tangent[1] = 1.0
    tangent -= np.dot(tangent, u)*u
    tangent /= np.linalg.norm(tangent)
    x = normalize(u[None, :] + eps*mode[:, None]*tangent[None, :])
    initial = np.linalg.norm(x - u[None, :])
    for _ in range(steps):
        x = spherical_step(x, a, eta)
    final = np.linalg.norm(x - u[None, :])
    return final / initial, lam_min


def main():
    n, d = 12, 8
    positions = np.arange(n, dtype=float)
    omegas = np.array([1.0, 0.37, 0.13, 0.047])
    betas = [0.5, 1.0, 2.0, 3.0]
    results = {"seed": SEED, "n": n, "d": d, "predictions": {}, "rows": []}

    # Prediction 1: scores in [-1,1] imply exact softmax floor.
    x_random = normalize(rng.normal(size=(n,d)))
    a_rand, scores_rand = attention(x_random, positions, omegas, beta=2.0)
    alpha = math.exp(-4.0)/n
    floor_observed = float(a_rand.min())
    floor_pass = floor_observed + 1e-14 >= alpha
    delta = max(0.5*np.abs(a_rand[i]-a_rand[j]).sum() for i in range(n) for j in range(n))
    delta_bound = 1-math.exp(-4.0)
    results["predictions"]["kernel_floor"] = {
        "predicted_alpha": alpha, "observed_min_weight": floor_observed,
        "predicted_dobrushin_bound": delta_bound, "observed_dobrushin": float(delta),
        "pass": bool(floor_pass and delta <= delta_bound + 1e-12)
    }

    # Prediction 2: norm-preserving update remains on the sphere.
    x0 = normalize(rng.normal(size=(n,d))); aa, _ = attention(x0, positions, omegas, beta=1.5)
    x1 = spherical_step(x0, aa, eta=2.0)
    results["predictions"]["sphere_invariance"] = {
        "max_norm_error": float(np.max(np.abs(np.linalg.norm(x1,axis=1)-1))),
        "pass": bool(np.max(np.abs(np.linalg.norm(x1,axis=1)-1)) < 1e-12)
    }

    # Prediction 3: transverse Euler boundary eta_c=2/(1-lambda_min).
    # At consensus the score matrix is symmetric, so the similarity transform is A itself.
    for beta in betas:
        xcons, a, scores, u = make_consensus(n,d,positions,omegas,beta)
        z = np.exp(beta * scores).sum(axis=1)
        dz = np.sqrt(z)
        sim = dz[:, None] * a / dz[None, :]
        ev = np.linalg.eigvalsh((sim + sim.T) / 2)
        lam_min = float(ev[0])
        eta_c = 2.0/(1.0-lam_min)
        # sweep relative to predicted boundary; classify by growth of worst mode
        ratios = []
        for rel in [0.70, 0.90, 0.99, 1.01, 1.10, 1.30]:
            eta = rel*eta_c
            ratio, lam_mode = stability_trial(a, z, u, eta)
            ratios.append({"relative_eta":rel, "growth_ratio":float(ratio), "grows":bool(ratio>1.01)})
        # nearest grid transition (first point with clear growth)
        growing = [r["relative_eta"] for r in ratios if r["grows"]]
        observed_rel = min(growing) if growing else None
        # Direct linear prediction for this eigenmode is |1+eta*(lambda_min-1)|.
        predicted_at_110 = abs(1+1.10*eta_c*(lam_min-1))**8
        results["rows"].append({
            "beta":beta, "lambda_min":lam_min, "predicted_eta_c":eta_c,
            "observed_transition_relative_grid":observed_rel,
            "predicted_growth_ratio_at_1.10xc":float(predicted_at_110),
            "sweep":ratios,
            "dobrushin_bound":1-math.exp(-2*beta),
            "dobrushin_observed":float(max(0.5*np.abs(a[i]-a[j]).sum() for i in range(n) for j in range(n)))
        })
    # Explicitly check the claimed boundary: all below 0.99 should contract, all 1.10 should grow.
    below_ok = all(r["sweep"][2]["growth_ratio"] < 1 for r in results["rows"])
    above_ok = all(r["sweep"][4]["growth_ratio"] > 1 for r in results["rows"])
    results["predictions"]["euler_stability"] = {
        "prediction": "eta<2/(1-lambda_min) contracts; eta>boundary grows",
        "all_0.99x_contract": bool(below_ok), "all_1.10x_grow": bool(above_ok),
        "pass": bool(below_ok and above_ok)
    }
    # A compact baseline comparison: unconstrained residual has norm growth, spherical does not.
    xb = x0.copy()
    for _ in range(20): xb = xb + 2.0*(aa@xb)
    xs = x0.copy()
    for _ in range(20): xs = spherical_step(xs, aa, 2.0)
    results["baseline_vs_idea"] = {
        "unconstrained_max_norm_after_20": float(np.linalg.norm(xb,axis=1).max()),
        "spherical_max_norm_after_20": float(np.linalg.norm(xs,axis=1).max()),
        "initial_angular_diameter": angular_diameter(x0),
        "spherical_angular_diameter_after_20": angular_diameter(xs)
    }
    out = Path("results.json")
    out.write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()