"""The worked numbers of Chapter 6, asserted against the Lecture 6 files.

Every number checked here appears in the chapter text, and the file under test computes each one
and stores none of them. We never run the training loops, and the tests call the closed-form and
numerical functions the lectures build their reports from.
"""

from __future__ import annotations

import importlib.util
import math
import sys
from pathlib import Path

import pytest
import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))


def load(name: str):
    """Import one lecture file by path, so that ``lecture_6`` needs no package marker."""
    spec = importlib.util.spec_from_file_location(f"lecture_6_{name}",
                                                  ROOT / "lecture_6" / f"{name}.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def close(values, expected, tolerance=5e-4):
    """Compare a list of numbers against the chapter's, coordinate by coordinate."""
    assert len(values) == len(expected)
    for got, want in zip(values, expected):
        assert got == pytest.approx(want, abs=tolerance)


# %% 01. flow_map.py, Sections 6.3 to 6.12
def test_flow_map_listing():
    flow_map = load("flow_map")
    x = torch.tensor(0.7, dtype=torch.float64)
    assert float(flow_map.flow(x, 0.0, 1.0)) == pytest.approx(4.1, abs=1e-9)
    assert float(flow_map.flow(flow_map.flow(x, 0.0, 0.5), 0.5, 1.0)) == pytest.approx(4.1, abs=1e-9)
    assert float(flow_map.flow(flow_map.flow(x, 0.0, 1.0), 1.0, 0.0)) == pytest.approx(0.7, abs=1e-9)
    from torch.func import jacrev

    assert float(jacrev(lambda z: flow_map.flow(z, 0.0, 1.0))(x)) == pytest.approx(3.0, abs=1e-9)


def test_flow_map_axioms_and_liouville():
    flow_map = load("flow_map")
    x = torch.tensor(0.7, dtype=torch.float64)
    identity, cocycle, tangent = flow_map.check_axioms(x, 200)
    assert identity < 1e-12 and cocycle < 1e-10 and tangent < 1e-4
    assert flow_map.liouville_log_det(0.0, 1.0, 4000) == pytest.approx(math.log(3.0), abs=1e-6)
    assert float(flow_map.jacobian_by_ode(x, 0.0, 1.0, 4000)) == pytest.approx(3.0, abs=1e-3)


def test_flow_map_transport_and_likelihood():
    flow_map = load("flow_map")
    source = flow_map.normal_density(0.7, 0.0, 1.0)
    data = flow_map.normal_density(4.1, 2.0, 3.0)
    assert source == pytest.approx(0.3123, abs=1e-4)
    assert data == pytest.approx(0.1041, abs=1e-4)
    assert source / data == pytest.approx(3.0, abs=1e-6)
    assert math.log(source) - math.log(3.0) == pytest.approx(-2.2625, abs=1e-4)
    assert math.log(data) == pytest.approx(-2.2625, abs=1e-4)


# %% 02. residuals.py, Sections 6.16 and 6.17
def test_stay_still_residuals():
    residuals = load("residuals")
    report = residuals.stay_still_residuals()
    assert report["stay-still Lagrangian"] == pytest.approx(1.0, abs=1e-9)
    assert report["stay-still Eulerian"] == pytest.approx(1.0, abs=1e-9)
    assert report["stay-still composition"] == pytest.approx(0.0, abs=1e-12)
    assert report["stay-still diagonal gap"] == pytest.approx(0.0, abs=1e-12)


def test_identity_axiom_is_free():
    residuals = load("residuals")
    torch.manual_seed(0)
    model = residuals.Map(dim=2, width=16)
    x, t = torch.randn(8, 2), torch.rand(8)
    assert torch.allclose(model.X(x, t, t), x)


# %% 03. meanflow.py, Sections 6.13 and 6.18
def test_meanflow_scalar_identity():
    meanflow = load("meanflow")
    report = meanflow.scalar_check()
    assert report["v_0,1(1)"] == pytest.approx(math.e - 1.0, abs=1e-5)
    assert report["b(1, 0)"] == pytest.approx(1.0, abs=1e-6)
    assert report["partial_s v"] == pytest.approx(-1.0, abs=1e-5)
    assert report["grad_x v_0,1"] == pytest.approx(math.e - 1.0, abs=1e-5)
    assert report["d/ds v_0,1"] == pytest.approx(math.e - 2.0, abs=1e-5)
    assert report["identity gap"] == pytest.approx(0.0, abs=1e-5)
    assert report["one map evaluation"] == pytest.approx(math.e, abs=1e-5)
    assert report["one Euler step"] == pytest.approx(2.0, abs=1e-6)


def test_meanflow_correction_vanishes_on_the_gaussian_flow():
    meanflow = load("meanflow")
    report = meanflow.gaussian_check()
    assert report["v_0,t(0.7) spread"] == pytest.approx(0.0, abs=1e-6)
    assert report["v_0,1(0.7)"] == pytest.approx(3.4, abs=1e-5)
    assert report["b(0.7, 0)"] == pytest.approx(3.4, abs=1e-5)


# %% 04. shortcut.py, Section 6.20
def test_shortcut_doubling():
    shortcut = load("shortcut")
    report = shortcut.scalar_doubling()
    assert report["s(1, 0, d)"] == pytest.approx(1.136102, abs=1e-6)
    assert report["state after the first step"] == pytest.approx(1.284025, abs=1e-6)
    assert report["s(x', d, d)"] == pytest.approx(1.458783, abs=1e-6)
    assert report["average of the two"] == pytest.approx(1.297443, abs=1e-6)
    assert report["s(1, 0, 2d) direct"] == pytest.approx(1.297443, abs=1e-6)
    assert report["doubling gap"] == pytest.approx(0.0, abs=1e-12)
    assert report["state after the doubled step"] == pytest.approx(math.exp(0.5), abs=1e-9)
    assert shortcut.ladder()[0] == 1.0 / 128.0 and shortcut.ladder()[-1] == 1.0


# %% 05. consistency.py, Section 6.20
def test_consistency_boundary_parameterization():
    consistency = load("consistency")
    report = consistency.scalar_consistency()
    assert report["f(0.5)"] == pytest.approx(math.exp(0.5), abs=1e-6)
    assert report["f(0.9)"] == pytest.approx(math.exp(0.1), abs=1e-6)
    assert report["f(1.0)"] == pytest.approx(1.0, abs=1e-9)
    assert report["c_out(1.0)"] == pytest.approx(0.0, abs=1e-12)
    assert report["v_0.5,1(1)"] == pytest.approx(1.297443, abs=1e-6)


def test_consistency_teacher_floor():
    consistency = load("consistency")
    report = consistency.teacher_floor()
    assert report["f(1, 0.5)"] == pytest.approx(1.648721, abs=1e-6)
    assert report["f(1.01, 0.51)"] == pytest.approx(1.648639, abs=1e-6)
    assert report["distillation residual"] == pytest.approx(8.189e-5, rel=1e-3)
    assert report["local Euler error"] == pytest.approx(5.0167e-5, rel=1e-3)
    assert report["carried to the data end"] == pytest.approx(
        report["distillation residual"], rel=1e-6)


def test_consistency_trajectory_head():
    consistency = load("consistency")
    report = consistency.scalar_ctm()
    assert report["a_0.0,0.5"] == pytest.approx(0.5, abs=1e-12)
    assert report["g_0.0,0.5"] == pytest.approx(2.297443, abs=1e-6)
    assert report["a_0.0,1.0"] == pytest.approx(0.0, abs=1e-12)
    assert report["g_0.0,1.0"] == pytest.approx(math.e, abs=1e-6)
    assert report["two routes apart"] == pytest.approx(0.0, abs=1e-12)


# %% 06. fmlm.py, Sections 6.26 and 6.27
def test_fmlm_denoiser_cocycle():
    fmlm = load("fmlm")
    report = fmlm.scalar_cocycle()
    assert report["gamma"] == pytest.approx(0.2, abs=1e-12)
    assert report["delta_s,t"] == pytest.approx(2.096158, abs=1e-6)
    assert report["delta_s,u"] == pytest.approx(1.932957, abs=1e-6)
    assert report["delta_u,t"] == pytest.approx(2.136959, abs=1e-6)
    assert report["two legs combined"] == pytest.approx(2.096158, abs=1e-6)
    assert report["convex weights"] == (0.25, 0.75)
    assert report["convex form"] == pytest.approx(math.exp(0.6), abs=1e-6)
    assert report["X_s,t(x)"] == pytest.approx(math.exp(0.6), abs=1e-9)


def test_fmlm_worked_kl():
    fmlm = load("fmlm")
    report = fmlm.worked_kl()
    assert report["worked gamma"] == pytest.approx(0.2, abs=1e-12)
    close(report["composed target"], [0.06, 0.12, 0.76, 0.06], 1e-6)
    assert report["target mass"] == pytest.approx(1.0, abs=1e-6)
    close(report["the four terms"], [-0.030650, -0.026777, 0.118823, -0.030650], 1e-6)
    assert report["worked KL"] == pytest.approx(0.030747, abs=1e-6)


def test_fmlm_gradient_comparison():
    """Section 6.27.4, the two gradients in the logits at the chapter's target and student."""
    fmlm = load("fmlm")
    report = fmlm.gradient_comparison()
    close(report["KL gradient"], [0.04, 0.03, -0.11, 0.04], 1e-6)
    close(report["squared-error gradient"], [0.0099, 0.01335, -0.03315, 0.0099], 1e-6)
    close(report["squared over KL"], [0.2475, 0.445, 0.3014, 0.2475], 1e-4)  # A quarter to a half.
    close(report["saturated student"], [1e-6, 0.1666665, 0.7222215, 0.111111], 1e-6)
    assert report["saturated KL gradient"] == pytest.approx(-0.0600, abs=1e-4)
    # Section 6.27.4 quotes -4.62e-08, its own expression
    #   d/dz_i (1/2) sum_j (q_j - p_j)^2 = q_i (q_i - p_i) - q_i sum_j q_j (q_j - p_j)
    # at q_1 = 1e-06 with the other three renormalized. The file prints that expression and
    # differentiates the loss automatically as well, and the two routes agree to four figures.
    assert report["saturated squared-error gradient, in 1e-08"] == pytest.approx(-4.6171, abs=1e-3)
    assert report["orders of magnitude between them"] == pytest.approx(6.11, abs=0.01)


def test_fmlm_decoding_clock():
    fmlm = load("fmlm")
    tau, start = fmlm.clock(fmlm.BOOK_VOCAB)
    assert start == pytest.approx(0.999967, abs=1e-5)
    assert tau(0.60) == pytest.approx(0.006, abs=5e-4)
    assert tau(0.75) == pytest.approx(0.141, abs=5e-4)
    assert tau(0.80) == pytest.approx(0.456, abs=1e-3)
    assert tau(0.85) == pytest.approx(0.931, abs=1e-3)
    grid = fmlm.even_in_tau(fmlm.BOOK_VOCAB)
    close(grid, [0.0, 0.773490, 0.804481, 0.828095, 1.0], 5e-5)
    close([tau(t) for t in (0.25, 0.5, 0.75)], [0.00009, 0.00124, 0.14109], 1e-5)


def test_fmlm_clock_weights():
    fmlm = load("fmlm")
    rounded = [0.0, 0.773, 0.804, 0.828, 1.0]  # The chapter quotes the grid to three decimals.
    pairs = [((1 - b) / (1 - a), (b - a) / (1 - a)) for a, b in zip(rounded[:-1], rounded[1:])]
    close([p[0] for p in pairs], [0.2270, 0.8634, 0.8776, 0.0], 1e-4)
    close([p[1] for p in pairs], [0.7730, 0.1366, 0.1224, 1.0], 1e-4)
    times = [torch.tensor([v]) for v in rounded[1:4]]
    assert float(fmlm.composition_weight(*times)) == pytest.approx(0.494620, abs=1e-5)
    early = [torch.tensor([v]) for v in (0.0, 0.773, 0.804)]
    assert float(fmlm.composition_weight(*early)) == pytest.approx(0.830144, abs=1e-5)


def test_fmlm_endpoint_slice_is_trained():
    """Section 6.26 asks for t = 1 on purpose, since a continuous draw never lands there."""
    fmlm = load("fmlm")
    assert 0.0 < fmlm.ENDPOINT_FRACTION < 1.0
    tokens, _ = fmlm.synthetic_dna(num_sequences=64, length=16, seed=0)
    model = fmlm.TwoTimeMLP(dim=16 * 4, width=16)
    torch.manual_seed(0)
    loss_fn = fmlm.make_loss(model, tokens, 16, 4, 256, torch.device("cpu"))
    assert torch.isfinite(loss_fn(0))  # The rows at t = 1 go through the loss like any other.
    # The map at t = 1 is the prediction itself, so the one-step sampler reads that slice alone.
    x = torch.randn(2, 16, 4)
    s, t = torch.zeros(2), torch.ones(2)
    prediction = fmlm.predict(model, x, s, t, 16, 4)
    assert torch.allclose(fmlm.step_map(x, prediction, s, t), prediction)


# %% 07. categorical_map.py, Section 6.28
def test_categorical_worked_cancellation():
    categorical = load("categorical_map")
    report = categorical.worked_cancellation()
    close(report["state x"], [-0.2591, -0.6675, -0.0467], 1e-4)
    close(report["X_s,t(x)"], [0.1611, -0.2420, 0.0943], 1e-4)
    close(report["pi_s,t(x)"], [0.5813, 0.1835, 0.2352], 1e-4)
    close(report["pi_t,t(X_s,t(x))"], [0.6927, 0.0917, 0.2157], 1e-4)
    assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12)
    close(report["d pi / dt"], [0.5570, -0.4592, -0.0979], 1e-4)
    close(report["central difference"], [0.5570, -0.4592, -0.0979], 1e-4)
    close(report["gap a"], [-0.111, 0.092, 0.020], 1e-3)
    close(report["drift b"], [0.111, -0.092, -0.020], 1e-3)


def test_categorical_bound_is_positive_where_the_residual_is_zero():
    categorical = load("categorical_map")
    report = categorical.worked_cancellation()
    assert report["residual |a + b|^2"] < 1e-20
    assert report["4 KL"] == pytest.approx(0.156, abs=1e-3)
    assert report["2 |b|^2"] == pytest.approx(0.042, abs=1e-3)
    assert report["the bound"] == pytest.approx(0.199, abs=2e-3)
    assert report["teacher entropy"] == pytest.approx(0.8043, abs=1e-3)
    assert report["cross-entropy"] == pytest.approx(0.8434, abs=1e-3)
    assert report["cross-entropy"] == pytest.approx(
        0.25 * report["4 KL"] + report["teacher entropy"], abs=1e-6)


def test_categorical_averaged_bound():
    categorical = load("categorical_map")
    report = categorical.averaged_bound(1000, seed=0)
    assert report["averaged residual"] < 1e-18
    assert report["averaged 4 KL"] == pytest.approx(0.39, abs=0.06)
    assert report["averaged 2 |b|^2"] == pytest.approx(0.095, abs=0.03)


# %% 08. discrete_map.py, Sections 6.26 and 6.29
def test_discrete_derivative_target_leaves_the_simplex():
    discrete = load("discrete_map")
    report = discrete.worked_derivative()
    assert report["derivative prefactor"] == pytest.approx(6.6667, abs=1e-3)
    close(report["derivative target"], [-0.3333, -0.6667, 1.3333, -0.3333], 1e-4)
    assert report["target sum"] == pytest.approx(0.0, abs=1e-6)
    assert report["negative entries"] == 3


def test_discrete_logit_repair():
    discrete = load("discrete_map")
    report = discrete.worked_repair()
    assert report["weighted average of dz/dt"] == pytest.approx(0.0, abs=1e-6)
    close(report["d psi/dt"], [1.92, -1.14, -0.78], 1e-4)
    close(report["additive target"], [-0.284, 0.828, 0.456], 1e-4)
    assert report["additive target sum"] == pytest.approx(1.0, abs=1e-6)
    assert report["additive negative entries"] == 1
    close(report["multiplicative factor"], [2.920, 0.544, 0.480], 1e-4)
    close(report["quotient"], [0.0342, 1.1029, 0.6250], 1e-4)
    assert report["quotient mass"] == pytest.approx(1.7622, abs=1e-4)
    close(report["repaired target"], [0.0194, 0.6259, 0.3547], 1e-4)
    close(report["log of the factor"], [1.0716, -0.6088, -0.7340], 1e-4)
    assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12)
    assert report["gamma^E_s,t"] == pytest.approx(0.8, abs=1e-12)


# %% 09. stochastic_map.py, Section 6.31
def test_stochastic_weak_floor():
    stochastic = load("stochastic_map")
    report = stochastic.weak_floor(draws=200_000, seed=0)
    assert report["K_1"] == pytest.approx(3.194528, abs=1e-5)
    assert report["sqrt K_1"] == pytest.approx(1.7873, abs=1e-4)
    assert report["sqrt K_0.5"] == pytest.approx(0.9269, abs=1e-4)
    assert report["one full step"] == pytest.approx(2.2453, abs=1e-3)
    assert report["two half steps"] == pytest.approx(0.4673, abs=1e-3)
    assert report["gap"] == pytest.approx(1.778, abs=1e-3)
    assert report["squared gap"] == pytest.approx(3.161, abs=1e-3)
    assert report["floor 2 K_1"] == pytest.approx(math.e ** 2 - 1.0, abs=1e-9)
    assert report["carried variance equals K_1"] == pytest.approx(report["K_1"], abs=1e-9)
    assert report["the draw that would close it"] == pytest.approx(-0.195, abs=1e-3)
    assert report["Monte Carlo floor"] == pytest.approx(6.389, abs=0.05)


def test_stochastic_chen_relations():
    stochastic = load("stochastic_map")
    report = stochastic.worked_chen()
    assert report["I0 over [0,1]"] == pytest.approx(0.50, abs=1e-9)
    assert report["I1 over [0,1]"] == pytest.approx(-0.20, abs=1e-9)
    assert report["fused variance at h=1"] == pytest.approx(1.0 / 3.0, abs=1e-9)


def test_stochastic_truncation():
    stochastic = load("stochastic_map")
    report = stochastic.truncation_error(paths=512, grid=2048, seed=0)
    for n, expected in enumerate((1.0, 1.0 / 3.0, 1.0 / 5.0)):
        assert report[f"coefficient variance I{n}"] == pytest.approx(expected, rel=0.15)
    errors = [report[f"pathwise error with {k}"] for k in (1, 2, 3)]
    assert errors[0] > 5.0 * errors[1] > 25.0 * errors[2]


# %% 10. expanding.py, Section 6.32
def test_expanding_local_clocks():
    expanding = load("expanding")
    report = expanding.worked_clocks()
    close(report["local clocks at t=0.8"], [0.800, 0.800, 0.692, 0.333], 1e-3)
    close(report["clock speeds"], [1.0, 1.0, 1.538, 3.333], 1e-3)
    close(report["source share"], [0.200, 0.200, 0.308, 0.667], 1e-3)
    assert report["last clock from its speed"] == pytest.approx(0.667, abs=1e-3)
    assert report["live coordinates at 0.8"] == 4 and report["live coordinates at 0.6"] == 3


def test_expanding_composition_weight():
    expanding = load("expanding")
    report = expanding.worked_omega()
    assert report["omega at (0.3, 0.6, 0.9)"] == pytest.approx(0.125, abs=1e-6)
    assert report["omega for a birth at 0.0"] == pytest.approx(0.125, abs=1e-6)
    assert report["omega for a birth at 0.7"] == pytest.approx(0.0, abs=1e-9)
    assert 0.0 < report["omega for a birth at 0.35"] < report["omega for a birth at 0.0"]


# %% 11. posterior_map.py, Sections 6.33 and 6.34
def test_posterior_tilt_and_jensen_gap():
    posterior = load("posterior_map")
    report = posterior.two_point()
    close(report["tilted target"], [0.1192, 0.8808], 1e-4)
    close(report["posterior at x=0.7"], [0.9820, 0.0180], 1e-4)
    close(report["tilted posterior"], [0.8808, 0.1192], 1e-4)
    assert report["V_t(0.7)"] == pytest.approx(0.1088, abs=1e-4)
    assert report["V over the data"] == pytest.approx(1.4338, abs=1e-4)
    assert report["single-endpoint estimate"] == pytest.approx(1.0, abs=1e-9)
    assert report["Jensen gap"] == pytest.approx(0.434, abs=1e-3)


def test_posterior_value_bias_and_gradient():
    posterior = load("posterior_map")
    report = posterior.value_bias()
    close([report[f"E[V_{k}]"] for k in (1, 2, 4, 8)], [1.000, 1.217, 1.340, 1.393], 1e-3)
    close([report[f"bias at K={k}"] for k in (1, 2, 4, 8)], [0.434, 0.217, 0.093, 0.041], 1e-3)
    weights, weighted, plain = posterior.value_gradient(
        torch.tensor([0.2, 1.5, 0.9]), torch.tensor([0.4, -1.2, 0.5]))
    close(weights.tolist(), [0.1496, 0.5491, 0.3013], 1e-4)
    assert weighted == pytest.approx(-0.4483, abs=1e-4)
    assert plain == pytest.approx(-0.1000, abs=1e-6)


def test_posterior_renoising_and_meta_fixed_point():
    posterior = load("posterior_map")
    report = posterior.renoise_example()
    assert report["renoise scale"] == pytest.approx(0.75, abs=1e-9)
    assert report["renoise variance"] == pytest.approx(0.1375, abs=1e-9)
    assert report["renoise standard deviation"] == pytest.approx(0.3708, abs=1e-4)
    assert report["proposal weight on the far mode"] == pytest.approx(0.420, abs=1e-3)
    assert report["penalty at z=2"] == pytest.approx(10.125, abs=1e-6)
    assert report["penalty at z=0"] == pytest.approx(6.125, abs=1e-6)
    assert report["corrected weight on the far mode"] == pytest.approx(0.0131, abs=1e-4)
    meta = posterior.stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]),
                                               torch.tensor([0.4, 1.0, 3.6]))
    assert meta["E[w]"] == pytest.approx(8.0 / 3.0, abs=1e-6)
    assert meta["E[G]"] == pytest.approx(5.0 / 3.0, abs=1e-6)
    assert meta["stop-gradient fixed point"] == pytest.approx(0.625, abs=1e-6)
    assert meta["per-sample ratio average"] == pytest.approx(0.540, abs=1e-6)
    assert meta["gap"] == pytest.approx(-0.085, abs=1e-3)
    assert meta["covariance of G/w and w"] == pytest.approx(0.2267, abs=1e-4)
    assert meta["gap from the covariance"] == pytest.approx(meta["gap"], abs=1e-6)