"""Independent analytic limits for a constant-inertia saturated wind toy model."""
from pathlib import Path
import json
import numpy as np
from scipy.integrate import solve_ivp

# Omega is rad/day; time is Gyr; k is day^2/Gyr. Radians are dimensionless.
k = 1.4
saturation = 2 * np.pi / 3
reports = []
for initial_period in [.5, 3., 8.]:
    omega0 = 2 * np.pi / initial_period
    transition = max(0., np.log(omega0 / saturation) / (k * saturation ** 2))
    times = np.linspace(0, 4, 801)
    def exact(time):
        if omega0 <= saturation:
            return 1 / np.sqrt(omega0 ** -2 + 2 * k * time)
        if time <= transition:
            return omega0 * np.exp(-k * saturation ** 2 * time)
        return 1 / np.sqrt(saturation ** -2 + 2 * k * (time - transition))
    def derivative(time, value):
        omega = value[0]
        return [-k * omega * min(omega ** 2, saturation ** 2)]
    solution = solve_ivp(derivative, [0, 4], [omega0], t_eval=times, method='DOP853', rtol=1e-11, atol=1e-12)
    expected = np.array([exact(time) for time in times])
    error = float(max(abs(solution.y[0] / expected - 1)))
    assert solution.success and error < 1e-8
    assert np.all(np.diff(solution.y[0]) <= 0)
    reports.append({'initial_period_days': initial_period, 'saturated_transition_gyr': transition, 'maximum_relative_analytic_error': error,
                    'final_period_days': float(2 * np.pi / solution.y[0, -1]), 'solver_forward_evaluations': solution.nfev})
# A prescribed contracting inertia with zero external torque must conserve I*Omega.
times = np.linspace(0, 1, 200)
solution = solve_ivp(lambda time, omega: [omega[0] / .3], [0, 1], [1.], t_eval=times, rtol=1e-11, atol=1e-12, method='DOP853')
inertia = np.exp(-times / .3)
conservation = float(max(abs(inertia * solution.y[0] - 1)))
assert conservation < 1e-9
report = {'status': 'PASS analytic saturated/unsaturated matching and zero-torque variable-inertia conservation',
          'units': {'time': 'Gyr', 'period': 'days', 'omega': 'radians/day', 'k': 'day^2/Gyr', 'inertia': 'arbitrary reference units'},
          'equation': 'd(I*Omega)/dt = torque; at constant I, dOmega/dt=-k*Omega*min(Omega^2,Omega_sat^2)',
          'k': k, 'saturation_period_days': 3., 'cases': reports,
          'zero_torque_contracting_inertia_maximum_relative_J_error': conservation,
          'scope': 'Mechanism/units prototype, not a stellar structure track or calibrated population model. Constant inertia excludes pre-main-sequence contraction and subgiants; optional contraction is a separate exact conservation check.'}
Path('results/angular-momentum-results.json').write_text(json.dumps(report, indent=2) + '\n')
print(json.dumps(report, indent=2))