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

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

# Degree-calibrated flow: f=-a*r^m*z + (I-zz'/||z||^2)h(z).
def tangent_project(z, h, eps=1e-12):
    q = np.sum(z*z, axis=-1, keepdims=True)
    return h - z * (np.sum(z*h, axis=-1, keepdims=True) / (q + eps))

def calibrated_f(z, W, a, m):
    r = 0.5*np.sum(z*z, axis=-1, keepdims=True)
    h = np.tanh(z @ W.T)
    return -0.5*a*(r**m)*z + tangent_project(z, h)

def simulate(z0, W, a, m, dt, steps, calibrated=True):
    z = z0.copy()
    norms, radial_errors = [], []
    for _ in range(steps):
        r = 0.5*np.sum(z*z, axis=1, keepdims=True)
        if calibrated:
            f = calibrated_f(z, W, a, m)
            allowed = -a*(r**(m+1))
        else:
            f = np.tanh(z @ W.T)
            allowed = np.zeros_like(r)
        radial = np.sum(z*f, axis=1, keepdims=True)
        radial_errors.append(float(np.max(radial-allowed)))
        norms.append(float(np.max(np.linalg.norm(z, axis=1))))
        z = z + dt*f
    norms.append(float(np.max(np.linalg.norm(z, axis=1))))
    return np.asarray(norms), np.asarray(radial_errors)

def slope(x, y):
    return float(np.polyfit(np.log(x), np.log(np.maximum(y, 1e-30)), 1)[0])

def main():
    random.seed(SEED); np.random.seed(SEED)
    d, n, a, steps = 8, 256, 1.0, 30000
    # Small initial states avoid numerical stiffness for m>0 while retaining a clear tail.
    z0 = rng.normal(size=(n,d)); z0 *= rng.uniform(.6,1.4,size=(n,1)) / np.linalg.norm(z0,axis=1,keepdims=True)
    W = rng.normal(scale=.7/math.sqrt(d), size=(d,d))

    decay = {}
    for m in (0,1,2):
        z = z0.copy(); rs=[]
        dt=.01
        for _ in range(steps):
            rs.append(float(np.mean(.5*np.sum(z*z,axis=1))))
            z += dt*calibrated_f(z,W,a,m)
        rs=np.asarray(rs)
        # Fit sufficiently late, but before float underflow / Euler floor.
        lo, hi = (100, 1500) if m == 0 else (3000, 25000)
        if m == 0:
            fit = float(np.polyfit(np.arange(lo,hi)*dt, np.log(rs[lo:hi]), 1)[0])
            expected = -a
            observed = fit
        else:
            fit = slope(np.arange(lo,hi)*dt, rs[lo:hi])
            expected = -1.0/m
            observed = fit
        decay[str(m)] = {"observed_log_r_slope": observed, "expected": expected,
                         "initial_r": float(rs[0]), "final_r": float(rs[-1]),
                         "max_abs_radial_identity_error": None}
        # Exact derivative check, independent of Euler discretization.
        zcheck = rng.normal(size=(1000,d)); f=calibrated_f(zcheck,W,a,m)
        rr=.5*np.sum(zcheck*zcheck,axis=1,keepdims=True)
        err=np.max(np.abs(np.sum(zcheck*f,axis=1,keepdims=True)+a*rr**(m+1)))
        decay[str(m)]["max_abs_radial_identity_error"] = float(err)

    # Stability control: same learned field and Euler step, with/without radial calibration.
    stability={}
    for dt in (0.01, 0.1, 0.5, 1.0):
        base_n, base_e = simulate(z0,W,a,1,dt,300,False)
        idea_n, idea_e = simulate(z0,W,a,1,dt,300,True)
        stability[str(dt)]={
            "baseline_final_max_norm":float(base_n[-1]),
            "idea_final_max_norm":float(idea_n[-1]),
            "baseline_peak_norm":float(np.max(base_n)),
            "idea_peak_norm":float(np.max(idea_n)),
            "idea_max_continuous_radial_violation":float(np.max(idea_e)),
            "baseline_max_radial_derivative":float(np.max(base_e))
        }
    # After input removal, calibrated state shrinks while unconstrained residual state persists/grows.
    result={"seed":SEED,"dimension":d,"batch":n,"a":a,"decay":decay,"stability":stability,
            "note":"Continuous radial identity is tested exactly; trajectory tests use explicit Euler."}
    Path("results.json").write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__ == '__main__': main()