"""The worked numbers of Chapter 2, asserted against the Lecture 2 scripts.

Every figure here is quoted in the chapter: the Euler landings at one, two and four steps, the flux
balance on [-1, 1], both sides of the continuity equation at x = 2, the single endpoint pair, the
average at the crossing, the nine-unit gap between the two objectives, the table of three
interpolants at t = 0.25, and the parameter count of the MNIST U-Net. Nothing in this file trains.
"""

import math
import sys
from pathlib import Path

import pytest
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # Run from a fresh clone.

from lecture_2 import flow_matching_2d as fm2d  # noqa: E402
from lecture_2 import interpolants as interp  # noqa: E402
from lecture_2 import velocity_fields as vf  # noqa: E402

LOG2 = math.log(2.0)


# %% 01. Velocity fields, trajectories and the flow map (Sections 2.2 to 2.4)


def test_trajectories_match_the_chapter():
    report = vf.trajectories()
    assert report["constant_field_trajectory"] == [1.0, 1.5, 2.0, 2.5, 3.0]
    assert report["expansion_x_half_from_two"] == pytest.approx(2.828, abs=1e-3)
    assert report["expansion_endpoints"] == [-2.0, 0.0, 2.0]
    assert report["expansion_outer_spacing"] == pytest.approx(4.0)
    assert report["contraction_endpoints"] == [-0.5, 0.0, 0.5]
    assert report["contraction_outer_spacing"] == pytest.approx(1.0)


def test_rotation_preserves_the_radius():
    report = vf.trajectories()
    assert report["rotation_quarter_turn"] == [0.0, 1.0]
    assert report["rotation_half_turn"] == [-1.0, 0.0]
    assert report["rotation_radius"] == pytest.approx(1.0, abs=1e-6)
    assert report["rotation_squared_norm_rate"] == pytest.approx(0.0, abs=1e-6)


def test_flow_map_is_the_identity_at_zero_and_doubles_at_one():
    report = vf.flow_map_report()
    assert report["flow_map_identity_error"] == pytest.approx(0.0, abs=1e-6)
    assert report["flow_map_jacobian_det"] == pytest.approx(2.0, abs=1e-6)
    assert report["trajectory_gap_minimum"] == pytest.approx(2.0, abs=1e-6)


# %% 02. Euler's method (Section 2.5)


def test_four_euler_steps_land_at_the_chapter_value():
    landing = vf.euler(lambda x, t: x, torch.tensor(1.0), 4)  # dx/dt = x from x_0 = 1.
    assert float(landing) == pytest.approx(2.44140625, abs=1e-6)


def test_euler_ladder_errors_and_ratios():
    report = vf.euler_ladder(4)
    assert report["euler_step_counts"] == [1, 2, 4]
    assert report["euler_landings"] == [2.0, 2.25, pytest.approx(2.4414, abs=1e-4)]
    assert report["euler_errors"] == [pytest.approx(0.7183, abs=1e-4),
                                      pytest.approx(0.4683, abs=1e-4),
                                      pytest.approx(0.2769, abs=1e-4)]
    assert report["euler_error_halving_ratios"] == [pytest.approx(1.53, abs=5e-3),
                                                    pytest.approx(1.69, abs=5e-3)]
    assert report["euler_relative_error_four_steps"] == pytest.approx(0.1019, abs=1e-4)


def test_the_euler_error_closes_like_one_over_the_step_count():
    report = vf.euler_ladder(256)  # The error times the step count settles at e/2 = 1.359.
    scaled = report["euler_error_times_steps"]
    assert scaled[-1] == pytest.approx(math.e / 2.0, abs=0.02)
    assert report["euler_error_halving_ratios"][-1] == pytest.approx(2.0, abs=0.02)


# %% 03. The pushforward, the flux and the continuity equation (Sections 2.6 and 2.7)


def test_the_expansion_field_doubles_the_standard_deviation():
    report = vf.pushforward_moments(4096, seed=0)
    assert report["pushforward_mean"] == pytest.approx(0.0, abs=0.05)
    assert report["pushforward_variance"] == pytest.approx(4.0, abs=0.2)


def test_flux_balance_on_the_unit_interval():
    report = vf.flux_balance()
    assert report["density_at_one"] == pytest.approx(0.2420, abs=1e-4)
    assert report["flux_at_minus_one"] == pytest.approx(-0.1677, abs=1e-4)
    assert report["flux_at_one"] == pytest.approx(0.1677, abs=1e-4)
    assert report["flux_balance"] == pytest.approx(-0.3354, abs=1e-4)
    assert report["enclosed_mass"] == pytest.approx(0.6827, abs=1e-4)
    # The balance and the derivative of the enclosed mass agree to four figures.
    assert report["enclosed_mass_rate"] == pytest.approx(report["flux_balance"], abs=1e-4)


def test_both_sides_of_the_continuity_equation_at_x_equals_two():
    report = vf.continuity_check()
    assert report["density_at_two"] == pytest.approx(0.05399, abs=1e-5)
    assert report["density_gradient_at_two"] == pytest.approx(-0.10798, abs=1e-5)
    assert report["velocity_divergence"] == pytest.approx(LOG2, abs=1e-6)
    assert report["continuity_time_derivative"] == pytest.approx(0.11227, abs=1e-5)
    assert report["continuity_minus_flux_divergence"] == pytest.approx(0.11227, abs=1e-5)


def test_log_density_along_a_trajectory_falls_by_log_two():
    report = vf.log_density_along_trajectory()
    assert report["log_density_rate"] == pytest.approx(-LOG2)
    assert report["log_density_change"] == pytest.approx(-LOG2, abs=1e-6)
    assert report["log_density_accumulated"] == pytest.approx(-LOG2)
    assert report["prior_density_at_origin"] == pytest.approx(0.3989, abs=1e-4)
    assert report["final_density_at_origin"] == pytest.approx(0.1995, abs=1e-4)
    assert report["density_ratio_at_origin"] == pytest.approx(2.0, abs=1e-6)


# %% 04. The conditional path and the crossing (Sections 2.12 and 2.13)


def test_the_single_endpoint_pair():
    report = fm2d.worked_pair()
    assert report["worked_pair_path"] == [-2.0, -0.75, 0.5, 1.75, 3.0]
    assert report["worked_pair_velocity"] == pytest.approx(5.0)


def test_the_average_at_the_crossing():
    report = fm2d.crossing_average()
    assert report["crossing_state"] == pytest.approx(2.0)
    assert report["crossing_velocities"] == [4.0, -2.0]
    assert report["crossing_average_velocity"] == pytest.approx(1.0)


def test_the_two_losses_differ_by_nine_at_the_crossing():
    report = fm2d.crossing_gap()
    assert report["crossing_conditional_loss"] == [10.0, 9.0, 10.0]
    assert report["crossing_marginal_loss"] == [1.0, 0.0, 1.0]
    assert report["crossing_constant_gap"] == [9.0, 9.0, 9.0]
    assert report["crossing_target_variance"] == pytest.approx(9.0)


# %% 05. The marginal velocity and the equivalence (Sections 2.13 and 2.15)


def test_the_marginal_velocity_at_the_two_ends():
    x = torch.randn(64, 2)
    # v_0(x) = E[X_1] - x = -x, and v_1(x) = x - E[X_0] = x, since the coupling is independent.
    assert torch.allclose(fm2d.marginal_velocity(x, torch.zeros(64)), -x, atol=1e-5)
    assert torch.allclose(fm2d.marginal_velocity(x, torch.ones(64)), x, atol=1e-5)


def test_the_marginal_velocity_transports_the_prior_onto_the_data():
    from cis6270.metrics import mean_log_density, mode_report

    torch.manual_seed(0)
    x_t = torch.randn(512, 2)  # X_0 ~ p_0.
    for step in range(50):  # Euler on the exact field, with no network anywhere.
        t = torch.full((512,), step / 50.0)
        x_t = x_t + fm2d.marginal_velocity(x_t, t) / 50.0
    assert mode_report(x_t)["modes_found"] == 4
    assert mean_log_density(x_t) > -3.0  # The data themselves measure -2.14.


def test_the_losses_decompose_exactly_on_any_batch():
    from cis6270.nets import VelocityMLP

    torch.manual_seed(0)
    model = VelocityMLP(dim=2, width=32)
    x_0, x_1 = torch.randn(256, 2), torch.randn(256, 2) + 2.0
    t = torch.rand(256)
    loss_cfm, loss_fm, gap, cross = fm2d.objectives(model, x_0, x_1, t)
    # We expand the square and get L_CFM = L_FM + cross + gap term by term, on every batch.
    assert float(loss_cfm) == pytest.approx(float(loss_fm + cross + gap), rel=1e-5)


def test_the_two_objectives_share_a_gradient():
    from cis6270.data import gaussian_mixture_2d
    from cis6270.nets import VelocityMLP

    torch.manual_seed(0)
    model = VelocityMLP(dim=2, width=64)
    data = gaussian_mixture_2d(4096, seed=0)
    report = fm2d.equivalence_check(model, data, 8192)
    assert report["loss_cfm"] == pytest.approx(report["loss_fm_plus_gap"], abs=0.3)
    assert abs(report["loss_cross_term"]) < 0.3  # Zero in expectation, small on a large batch.
    assert report["gradient_cosine"] > 0.99  # The residual is the batch's own sampling error.


# %% 06. The three interpolants (Section 2.18)


def test_every_schedule_reaches_the_right_endpoints():
    for name in interp.SCHEDULES:
        assert interp.endpoint_and_table(name)[f"{name}_endpoint_error"] == pytest.approx(0.0,
                                                                                          abs=1e-6)


def test_the_table_of_three_interpolants_at_one_quarter():
    expected = {  # Table 2.1, at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
        "linear": ([1.50, 0.50], [-2.0, 2.0]),
        "quadratic": ([1.875, 0.125], [-1.0, 1.0]),
        "trigonometric": ([1.848, 0.765], [-1.202, 2.902]),
    }
    for name, (position, velocity) in expected.items():
        report = interp.endpoint_and_table(name)
        assert report[f"{name}_position_at_quarter"] == pytest.approx(position, abs=1e-3)
        assert report[f"{name}_velocity_at_quarter"] == pytest.approx(velocity, abs=1e-3)


def test_only_the_trigonometric_schedule_keeps_identity_covariance():
    trig = interp.covariance_check("trigonometric", 8192, seed=0)
    assert trig["trigonometric_coefficient_squares"] == [1.0, 1.0, 1.0, 1.0, 1.0]
    assert trig["trigonometric_covariance_deviation"] < 0.05
    linear = interp.covariance_check("linear", 8192, seed=0)
    assert linear["linear_coefficient_squares"] == [1.0, 0.625, 0.5, 0.625, 1.0]
    assert linear["linear_covariance_deviation"] > 0.4  # Cov(X_{1/2}) is half the identity.


def test_the_quadratic_schedule_has_covered_a_sixteenth_at_one_quarter():
    # beta_{1/4} = (1/4)^2 = 0.0625, so the state has covered 6.25 percent of the segment.
    _, beta, _, _ = interp.quadratic(torch.tensor([0.25]))
    assert float(beta) == pytest.approx(0.0625)
    _, beta_half, _, _ = interp.quadratic(torch.tensor([0.5]))
    assert float(beta_half) == pytest.approx(0.25)  # A quarter of the distance at the midpoint.


# %% 07. The MNIST U-Net (Section 2.20)


def test_the_unet_has_the_parameter_count_the_chapter_quotes():
    from lecture_2 import flow_matching_mnist as mnist

    model = mnist.FlowUNet()  # The default base width of 32, so the widths are 32, 64 and 128.
    assert sum(p.numel() for p in model.parameters()) == 471_265