"""
Unit tests for the simulation's physics core -- no Gradio/UI involved,
so these run in any CI environment with just numpy + scipy installed.

Run with: python -m unittest test_pendulum_physics.py -v
"""

import math
import unittest

import numpy as np

from pendulum_physics import (
    period_exact,
    period_series_correction,
    period_small_angle,
    simulate,
    virtual_experiment,
)


class TestSmallAngleLimit(unittest.TestCase):
    def test_exact_period_converges_to_small_angle_limit(self):
        """At a vanishingly small amplitude, the exact nonlinear period
        must converge to the SHM analytic period."""
        g, L = 9.80665, 0.30
        T0 = period_small_angle(g, L)
        T_tiny = period_exact(math.radians(0.01), g, L)
        self.assertAlmostEqual(T_tiny, T0, places=5)

    def test_series_correction_matches_exact_elliptic_at_15_degrees(self):
        g, L = 9.81, 0.30
        theta0 = math.radians(15)
        T0 = period_small_angle(g, L)
        T_exact = period_exact(theta0, g, L)
        exact_fraction = T_exact / T0 - 1.0
        series_fraction = period_series_correction(theta0)
        self.assertAlmostEqual(exact_fraction, series_fraction, places=4)

    def test_correction_grows_with_amplitude(self):
        c10 = period_series_correction(math.radians(10))
        c30 = period_series_correction(math.radians(30))
        c90 = period_series_correction(math.radians(90))
        self.assertLess(c10, c30)
        self.assertLess(c30, c90)


class TestIntegratorPhysics(unittest.TestCase):
    def test_energy_conserved_without_damping(self):
        res = simulate(math.radians(30), 9.81, 0.3, damping=0.0, t_max=8.0, dt=0.001)
        energy = res["energy"]
        relative_drift = (energy.max() - energy.min()) / energy.mean()
        self.assertLess(relative_drift, 1e-6)

    def test_energy_monotonically_decreases_with_damping(self):
        res = simulate(math.radians(30), 9.81, 0.3, damping=0.8, t_max=8.0, dt=0.001)
        energy = res["energy"]
        # sampled every 200 steps to ignore within-cycle KE/PE exchange noise
        sampled = energy[::200]
        diffs = np.diff(sampled)
        self.assertTrue((diffs <= 1e-9).all())

    def test_zero_amplitude_stays_at_rest(self):
        res = simulate(0.0, 9.81, 0.3, damping=0.0, t_max=2.0, dt=0.01)
        self.assertTrue(np.allclose(res["theta"], 0.0, atol=1e-9))
        self.assertTrue(np.allclose(res["omega"], 0.0, atol=1e-9))

    def test_period_from_integration_matches_elliptic_prediction(self):
        """Cross-check: find the first zero-crossing time of theta(t) in
        the numerical integration and confirm a quarter period matches the
        analytic elliptic-integral period prediction."""
        g, L = 9.81, 0.3
        theta0 = math.radians(20)
        res = simulate(theta0, g, L, damping=0.0, t_max=5.0, dt=0.0002)
        theta = res["theta"]
        t = res["t"]
        sign_changes = np.where(np.diff(np.sign(theta)) != 0)[0]
        self.assertGreater(len(sign_changes), 0)
        t_quarter_period_numeric = t[sign_changes[0]]
        T_exact = period_exact(theta0, g, L)
        self.assertAlmostEqual(t_quarter_period_numeric, T_exact / 4.0, delta=0.01)


class TestVirtualExperiment(unittest.TestCase):
    def test_zero_noise_still_shows_small_angle_systematic_bias(self):
        """With timing noise switched off, virtual_experiment is NOT expected
        to recover g_true exactly: the measurement method itself assumes the
        small-angle (SHM) relation g=4*pi^2*L/T^2, but the ball is timed on
        its *exact* nonlinear trajectory -- so a small, deterministic,
        amplitude-dependent bias survives even with perfect timing. This is
        a feature, not a bug: it is exactly the ~0.19% (at 10 deg) systematic
        error that motivates keeping release angles small in vision_lab/.
        """
        g_true = 9.81
        theta0 = math.radians(10)
        rng = np.random.default_rng(0)
        trial = virtual_experiment(0.3, g_true, theta0, n_oscillations=20,
                                    timing_noise_s=0.0, rng=rng)
        expected_bias_factor = (1.0 + period_series_correction(theta0)) ** 2
        expected_g = g_true / expected_bias_factor
        self.assertAlmostEqual(trial.g_measured, expected_g, places=6)
        # And that bias should be small but non-zero at this amplitude.
        self.assertLess(abs(trial.g_measured - g_true), 0.1)
        self.assertGreater(abs(trial.g_measured - g_true), 1e-6)

    def test_more_oscillations_reduce_scatter_across_many_trials(self):
        g_true = 9.81
        rng10 = np.random.default_rng(1)
        rng40 = np.random.default_rng(1)
        trials_10 = [virtual_experiment(0.3, g_true, math.radians(10), 10, 0.02, rng10) for _ in range(300)]
        trials_40 = [virtual_experiment(0.3, g_true, math.radians(10), 40, 0.02, rng40) for _ in range(300)]
        std_10 = np.std([t.g_measured for t in trials_10])
        std_40 = np.std([t.g_measured for t in trials_40])
        self.assertLess(std_40, std_10)

    def test_measured_g_centers_on_true_g_over_many_trials(self):
        rng = np.random.default_rng(7)
        trials = [virtual_experiment(0.3, 9.81, math.radians(10), 20, 0.02, rng) for _ in range(500)]
        mean_g = np.mean([t.g_measured for t in trials])
        self.assertAlmostEqual(mean_g, 9.81, delta=0.05)


if __name__ == "__main__":
    unittest.main()