import json
from pathlib import Path
import numpy as np
from scipy.optimize import root

SEED = 1255
np.random.seed(SEED)
I = np.array([1.0, 2.0, 4.0])
Iinv = 1.0 / I
# Levi-Civita symbols: eps[0,1,2] = +1.
E = np.zeros((3, 3, 3))
E[0, 1, 2] = E[1, 2, 0] = E[2, 0, 1] = 1.0
E[0, 2, 1] = E[2, 1, 0] = E[1, 0, 2] = -1.0


def H(z):
    return 0.5 * np.sum(Iinv * z[3:] ** 2)


def grad_H(z):
    g = np.zeros(6)
    g[3:] = Iinv * z[3:]
    return g


def poisson_J(z):
    p = z[3:]
    J = np.zeros((6, 6))
    J[:3, 3:] = np.eye(3)
    J[3:, :3] = -np.eye(3)
    J[3:, 3:] = np.einsum('a,aij->ij', p, E)
    return J


def vector_field(z, use_coadjoint=True):
    p = z[3:]
    v = Iinv * p
    # p_a c^a_ij dH/dp_j; this is the explicit Lie-Poisson term.
    coad = np.einsum('a,aij,j->i', p, E, v) if use_coadjoint else np.zeros(3)
    return np.r_[v, coad]


def euler_step(z, dt, use_coadjoint=True):
    return z + dt * vector_field(z, use_coadjoint)


def midpoint_step(z, dt, use_coadjoint=True):
    def residual(zn):
        zm = 0.5 * (z + zn)
        return zn - z - dt * vector_field(zm, use_coadjoint)
    sol = root(residual, z + dt * vector_field(z, use_coadjoint), method='hybr')
    if not sol.success or np.linalg.norm(residual(sol.x), ord=np.inf) > 1e-9:
        raise RuntimeError('implicit midpoint failed: ' + sol.message)
    return sol.x


def rollout(z0, dt, steps, method='midpoint', use_coadjoint=True):
    z = z0.copy()
    out = [z.copy()]
    for _ in range(steps):
        if method == 'euler':
            z = euler_step(z, dt, use_coadjoint)
        else:
            z = midpoint_step(z, dt, use_coadjoint)
        out.append(z.copy())
    return np.asarray(out)


def loglog_slope(xs, ys):
    return float(np.polyfit(np.log(xs), np.log(np.maximum(ys, 1e-30)), 1)[0])


def main():
    z0 = np.r_[np.array([0.2, -0.1, 0.3]), np.array([0.7, 1.1, 1.5])]
    # Core algebraic check: J is antisymmetric and hence grad H^T J grad H = 0.
    zcheck = np.r_[np.array([0.4, -0.2, 0.1]), np.array([0.8, -1.3, 1.7])]
    J = poisson_J(zcheck)
    g = grad_H(zcheck)
    antisym = float(np.max(np.abs(J + J.T)))
    energy_derivative = float(abs(g @ J @ g))

    # Prediction 1: coadjoint effect vanishes exactly at p=0.
    zzero = np.r_[z0[:3], np.zeros(3)]
    a0 = rollout(zzero, 0.05, 100, 'midpoint', True)
    b0 = rollout(zzero, 0.05, 100, 'midpoint', False)
    zero_momentum_difference = float(np.max(np.abs(a0 - b0)))

    # Prediction 2: at fixed time, the coadjoint-vs-ablated difference is O(alpha^2)
    # for small momentum scale alpha, because the extra vector field is quadratic in p.
    alphas = np.array([0.125, 0.25, 0.5, 1.0])
    alpha_errors = []
    for alpha in alphas:
        za = np.r_[z0[:3], alpha * z0[3:]]
        full = rollout(za, 0.01, 100, 'midpoint', True)
        ablated = rollout(za, 0.01, 100, 'midpoint', False)
        alpha_errors.append(np.linalg.norm(full[-1, 3:] - ablated[-1, 3:]))
    alpha_errors = np.asarray(alpha_errors)
    alpha_slope = loglog_slope(alphas, alpha_errors)

    # Prediction 3: midpoint energy error scales quadratically with dt over fixed T,
    # while Euler has first-order drift.
    T = 10.0
    dts = np.array([0.1, 0.05, 0.025, 0.0125])
    midpoint_energy_errors, euler_energy_errors = [], []
    for dt in dts:
        n = int(round(T / dt))
        zm = rollout(z0, dt, n, 'midpoint', True)[-1]
        ze = rollout(z0, dt, n, 'euler', True)[-1]
        midpoint_energy_errors.append(abs(H(zm) - H(z0)))
        euler_energy_errors.append(abs(H(ze) - H(z0)))
    midpoint_energy_errors = np.asarray(midpoint_energy_errors)
    euler_energy_errors = np.asarray(euler_energy_errors)
    midpoint_slope = loglog_slope(dts, midpoint_energy_errors)
    euler_slope = loglog_slope(dts, euler_energy_errors)

    # Secondary comparison: long-rollout energy stability and ablation phase error.
    long_steps, long_dt = 1000, 0.02
    mid = rollout(z0, long_dt, long_steps, 'midpoint', True)
    eu = rollout(z0, long_dt, long_steps, 'euler', True)
    abl = rollout(z0, long_dt, long_steps, 'midpoint', False)
    long_mid_energy = float(np.max(np.abs(np.array([H(x) for x in mid]) - H(z0))))
    long_euler_energy = float(np.max(np.abs(np.array([H(x) for x in eu]) - H(z0))))
    ablation_terminal_error = float(np.linalg.norm(mid[-1, 3:] - abl[-1, 3:]))

    result = {
        'seed': SEED,
        'inertia': I.tolist(),
        'math_check': {
            'max_J_plus_JT': antisym,
            'abs_gradH_J_gradH': energy_derivative,
            'tolerance': 1e-12,
        },
        'predictions': {
            'zero_momentum_effect_predicted': 0.0,
            'zero_momentum_observed_max_difference': zero_momentum_difference,
            'momentum_scaling_predicted_exponent': 2.0,
            'momentum_scaling_observed_exponent': alpha_slope,
            'momentum_scales': alphas.tolist(),
            'momentum_effects': alpha_errors.tolist(),
            'midpoint_energy_dt_predicted_exponent': 2.0,
            'midpoint_energy_dt_observed_exponent': midpoint_slope,
            'euler_energy_dt_predicted_exponent': 1.0,
            'euler_energy_dt_observed_exponent': euler_slope,
            'dts': dts.tolist(),
            'midpoint_energy_errors': midpoint_energy_errors.tolist(),
            'euler_energy_errors': euler_energy_errors.tolist(),
        },
        'secondary_comparison': {
            'midpoint_max_energy_error_T10': long_mid_energy,
            'euler_max_energy_error_T10': long_euler_energy,
            'midpoint_full_vs_no_coadjoint_terminal_momentum_error': ablation_terminal_error,
        },
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()