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

SEED = 2742
D = 8
BETA = 2.0
KAPPA = 1.5
MU = np.zeros(D); MU[0] = 1.0

def proj(x, v):
    return v - np.sum(x*v, axis=-1, keepdims=True)*x

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

def sphere_run(method, h, T=0.25, ntraj=2000, amp=0.0, seed=0):
    r = np.random.default_rng(seed)
    n = int(round(T/h)); h = T/n
    x = np.zeros((ntraj,D)); x[:,0] = 1.0
    maxviol = np.zeros(ntraj); energy = np.zeros(ntraj)
    noise_scale = math.sqrt(2.0*h/BETA)
    for _ in range(n):
        drift = KAPPA*proj(x, MU[None,:])
        u = amp*proj(x, MU[None,:])
        z = r.normal(size=x.shape)
        if method in ('intrinsic','projected_each'):
            z = proj(x,z)
        xnew = x + h*(drift+u) + noise_scale*z
        energy += np.sum(u*u,axis=1)*h
        if method in ('intrinsic','projected_each'):
            x = normalize(xnew)
            # Constraint is evaluated on the state actually used by the method.
            state_violation = np.abs(np.linalg.norm(x,axis=1)-1.0)
        else:
            x = xnew
            state_violation = np.abs(np.linalg.norm(x,axis=1)-1.0)
        maxviol = np.maximum(maxviol, state_violation)
    if method == 'projected_final':
        x = normalize(x)
    return dict(max_violation=float(np.max(maxviol)), median_violation=float(np.median(maxviol)),
      mean_final_norm_error=float(np.mean(np.abs(np.linalg.norm(x,axis=1)-1))),
      mean_alignment=float(np.mean(x[:,0])), mean_control_energy=float(np.mean(energy)),
      q95_violation=float(np.quantile(maxviol,.95)))

def one_step_violation(h, ntraj=100000, seed=0):
    r=np.random.default_rng(seed)
    x=np.zeros((ntraj,D)); x[:,0]=1
    z=r.normal(size=x.shape)
    y=x+math.sqrt(2*h/BETA)*z
    e=np.abs(np.linalg.norm(y,axis=1)-1)
    return float(np.median(e)), float(np.quantile(e,.95))

def main():
    r=np.random.default_rng(SEED+1)
    x=normalize(r.normal(size=(100,D))); v=r.normal(size=(100,D)); p=proj(x,v)
    projector_resid=max(float(np.max(np.abs(np.sum(x*p,axis=1)))),float(np.max(np.abs(proj(x,p)-p))))
    hs=[1/25,1/50,1/100,1/200]
    intrinsic=[sphere_run('intrinsic',h,seed=100+i) for i,h in enumerate(hs)]
    med=[]; q95=[]
    for i,h in enumerate(hs):
        a,b=one_step_violation(h,seed=300+i); med.append(a); q95.append(b)
    slope=float(np.polyfit(np.log(hs),np.log(med),1)[0])
    x0=normalize(r.normal(size=(10000,D))); tangent=proj(x0,MU[None,:]); base=float(np.mean(np.sum(tangent*tangent,axis=1)))
    amps=[.25,.5,1.,2.,4.]
    costs=[BETA/4*a*a*base for a in amps]
    cost_slope=float(np.polyfit(np.log(amps),np.log(costs),1)[0])
    endpoint={
      'intrinsic':sphere_run('intrinsic',1/200,T=.25,amp=1,seed=2024),
      'projected_each':sphere_run('projected_each',1/200,T=.25,amp=1,seed=2024),
      'projected_final':sphere_run('projected_final',1/200,T=.25,amp=1,seed=2024),
      'ambient':sphere_run('ambient',1/200,T=.25,amp=1,seed=2024)}
    response=[sphere_run('intrinsic',1/200,T=.25,amp=a,seed=900+i)['mean_alignment'] for i,a in enumerate([0,.5,1,2,4])]
    report={'seed':SEED,'dimension':D,'beta':BETA,'kappa':KAPPA,
      'algebra_check':{'max_projector_tangent_or_idempotence_residual':projector_resid},
      'prediction_checks':{
       'P1_intrinsic_constraint':{'prediction':'post-retraction intermediate |norm-1| stays at floating-point precision, independent of h','h_values':hs,'observed_max_each':[q['max_violation'] for q in intrinsic],'observed_max_over_sweep':max(q['max_violation'] for q in intrinsic)},
       'P2_ambient_local_scaling':{'prediction':'one-step unconstrained radial norm violation scales as h^0.5','h_values':hs,'median_violation':med,'q95_violation':q95,'observed_loglog_exponent':slope},
       'P3_fixed_state_control_cost':{'prediction':'beta/4 E||u||^2 scales exactly as amplitude^2','amplitudes':amps,'costs':costs,'observed_loglog_exponent':cost_slope,'intrinsic_endpoint_alignment_by_amplitude':response}},
      'endpoint_comparison_equal_steps':endpoint}
    Path('results.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()