"""Lecture 7: every number the chapter works by hand, asserted against the file that prints it.

Chapter 7 computes almost all of its results on paper, so these tests are the chapter itself read
back. The tolerances are the number of digits the chapter quotes.
"""

from __future__ import annotations

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

import pytest
import torch

LECTURE = Path(__file__).resolve().parent.parent / "lecture_7"


def load(name: str):
    """Import one lecture file by path, since ``lecture_7`` is a folder of scripts, not a package."""
    spec = importlib.util.spec_from_file_location(f"lecture_7_{name}", LECTURE / f"{name}.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def tensor(*values):
    return torch.tensor(values, dtype=torch.float64)


# %% Sections 7.1 to 7.6, optimal_transport.py


def test_four_points_and_two_pairings():
    ot = load("optimal_transport")
    pairings = ot.two_pairings()
    assert pairings["arrows_ordered"] == pytest.approx(2.0)
    assert pairings["arrows_crossed"] == pytest.approx(10.0)
    assert pairings["mass_ordered"] == pytest.approx(1.0)
    assert pairings["mass_crossed"] == pytest.approx(5.0)
    assert pairings["mass_independent"] == pytest.approx(3.0)
    assert pairings["C"].tolist() == [[1.0, 9.0], [1.0, 1.0]]


def test_two_by_two_plan_and_prices():
    ot = load("optimal_transport")
    assert ot.feasible_interval() == (0.0, 0.3)
    C = ot.cost_matrix(ot.SOURCE_X, ot.TARGET_Y)
    assert ot.cost_along_the_segment(0.0, C) == pytest.approx(5.8)
    assert ot.cost_along_the_segment(0.3, C) == pytest.approx(3.4)
    plan = ot.feasible_plan(0.3)
    assert plan.tolist() == [[0.3, 0.3], [0.0, 0.4]]
    assert ot.is_coupling(plan, ot.SOURCE_A, ot.TARGET_B)
    f, g = ot.prices_from_loaded_routes(C, plan)
    assert f.tolist() == [0.0, -8.0] and g.tolist() == [1.0, 9.0]
    assert ot.slack(C, f, g).tolist() == [[0.0, 0.0], [8.0, 0.0]]
    assert ot.dual_value(ot.SOURCE_A, f, ot.TARGET_B, g) == pytest.approx(3.4)
    assert ot.plan_cost(C, plan) ** 0.5 == pytest.approx(1.844, abs=5e-4)


def test_no_crossing_and_three_points():
    ot = load("optimal_transport")
    expanded, collected = ot.crossing_difference(0.0, 2.0, 1.0, 3.0)
    assert expanded == pytest.approx(collected) and collected == pytest.approx(-8.0)
    x = tensor(4.0, 1.0, 6.0)
    y = tensor(9.0, 2.0, 3.0)
    costs = [c for c, _ in ot.monge_permutations(x.sort().values, y.sort().values)]
    assert costs == pytest.approx([11 / 3, 17 / 3, 35 / 3, 45 / 3, 77 / 3, 81 / 3])
    assert ot.wasserstein_on_the_line(x, y, 2.0, 3072) == pytest.approx(11 / 3)


def test_gaussian_closed_form():
    ot = load("optimal_transport")
    assert ot.gaussian_w2_squared(0.0, 1.0, 3.0, 2.0) == pytest.approx(10.0)
    assert math.sqrt(10.0) == pytest.approx(3.162, abs=5e-4)
    points = tensor(-1.0, 0.0, 1.0)
    assert ot.brenier_map(points, 0.0, 1.0, 3.0, 2.0).tolist() == [1.0, 3.0, 5.0]
    source = torch.special.ndtri(ot.quantile_levels(3072))
    mapped = ot.brenier_map(source, 0.0, 1.0, 3.0, 2.0)
    assert float(((mapped - source) ** 2).mean()) == pytest.approx(10.0, abs=1e-3)


# %% Sections 7.9 and 7.10, sinkhorn.py


def test_entropies_of_the_three_couplings():
    sk = load("sinkhorn")
    independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B)
    assert independent.flatten().tolist() == pytest.approx([0.18, 0.42, 0.12, 0.28])
    assert sk.entropy(sk.EXACT_PLAN) == pytest.approx(1.089, abs=5e-4)
    assert sk.entropy(independent) == pytest.approx(1.284, abs=5e-4)
    assert sk.entropy(sk.SOURCE_A) == pytest.approx(0.673, abs=5e-4)
    assert sk.entropy(sk.TARGET_B) == pytest.approx(0.611, abs=5e-4)
    assert sk.relative_entropy(sk.EXACT_PLAN, independent) == pytest.approx(0.195, abs=5e-4)


def test_sinkhorn_first_sweep_by_hand():
    sk = load("sinkhorn")
    K = sk.gibbs_kernel(sk.COST, 2.0)
    assert K.flatten().tolist() == pytest.approx([0.606531, 0.011109, 0.606531, 0.606531],
                                                 abs=5e-7)
    v = torch.ones(2, dtype=torch.float64)
    assert (K @ v).tolist() == pytest.approx([0.617640, 1.213061], abs=5e-7)
    u = sk.SOURCE_A / (K @ v)
    assert u.tolist() == pytest.approx([0.971440, 0.329744], abs=5e-7)
    assert (K.T @ u).tolist() == pytest.approx([0.789208, 0.210792], abs=5e-7)
    v = sk.TARGET_B / (K.T @ u)
    assert v.tolist() == pytest.approx([0.38013, 3.32081], abs=5e-6)
    plan = sk.scaling_form(u, K, v)
    assert plan.flatten().tolist() == pytest.approx([0.2240, 0.0358, 0.0760, 0.6642], abs=5e-5)
    assert plan.sum(dim=0).tolist() == pytest.approx([0.3, 0.7])
    assert plan.sum(dim=1).tolist() == pytest.approx([0.2598, 0.7402], abs=5e-5)
    after = sk.project_rows(plan, sk.SOURCE_A)
    assert after.sum(dim=0).tolist() == pytest.approx([0.558, 0.442], abs=5e-4)


def test_sinkhorn_convergence_and_epsilon():
    sk = load("sinkhorn")
    K = sk.gibbs_kernel(sk.COST, 2.0)
    u, v, errors = sk.sinkhorn_sweeps(K, sk.SOURCE_A, sk.TARGET_B, 40)
    assert errors[:5] == pytest.approx([0.680, 0.382, 0.155, 0.047, 0.012], abs=5e-4)
    assert sk.sweeps_to_tolerance(K, sk.SOURCE_A, sk.TARGET_B) == 17
    plan = sk.scaling_form(u, K, v)
    assert plan.flatten().tolist() == pytest.approx(
        [0.293122, 0.306878, 0.006878, 0.393122], abs=5e-7)
    assert float((sk.COST * plan).sum()) == pytest.approx(3.45502, abs=5e-6)
    assert sk.entropy(plan) == pytest.approx(1.124, abs=5e-4)
    independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B)
    assert sk.entropic_objective(plan, sk.COST, 2.0) == pytest.approx(1.208, abs=5e-4)
    assert sk.entropic_objective(sk.EXACT_PLAN, sk.COST, 2.0) == pytest.approx(1.222, abs=5e-4)
    assert sk.entropic_objective(independent, sk.COST, 2.0) == pytest.approx(1.792, abs=5e-4)
    assert sk.relative_entropy(plan, independent) == pytest.approx(0.160, abs=5e-4)
    f = sk.multipliers_from_scalings(u, 2.0)
    g = sk.multipliers_from_scalings(v, 2.0)
    assert f.tolist() == pytest.approx([3.766, -3.739], abs=5e-4)
    assert g.tolist() == pytest.approx([-3.220, 4.872], abs=5e-4)
    expected = {10.0: (4.0540165, 7), 2.0: (3.4550204, 17), 0.5: (3.4000004, 28),
                0.1: (3.4000000, 80)}
    for epsilon, (cost, sweeps) in expected.items():
        kernel = sk.gibbs_kernel(sk.COST, epsilon)
        su, sv, _ = sk.sinkhorn_sweeps(kernel, sk.SOURCE_A, sk.TARGET_B, 400)
        assert float((sk.COST * sk.scaling_form(su, kernel, sv)).sum()) == pytest.approx(
            cost, abs=5e-8)
        assert sk.sweeps_to_tolerance(kernel, sk.SOURCE_A, sk.TARGET_B) == sweeps


# %% Sections 7.7 to 7.7.5, displacement.py


def test_one_particle_and_the_averaged_bound():
    dp = load("displacement")
    times = torch.linspace(0.0, 1.0, 1025, dtype=torch.float64)
    assert dp.action_of_a_path(dp.steady_particle(times), times) == pytest.approx(4.0)
    assert dp.action_of_a_path(dp.dawdling_particle(times), times) == pytest.approx(8.0)
    x, y = tensor(0.0, 2.0), tensor(1.0, 3.0)
    ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=torch.float64)
    crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=torch.float64)
    independent = torch.full((2, 2), 0.25, dtype=torch.float64)
    assert dp.average_squared_displacement(x, y, ordered) == pytest.approx(1.0)
    assert dp.average_squared_displacement(x, y, independent) == pytest.approx(3.0)
    assert dp.average_squared_displacement(x, y, crossed) == pytest.approx(5.0)


def test_displacement_interpolation_and_its_velocity():
    dp = load("displacement")
    mean, std = dp.displacement_law(0.5)
    assert (mean, std) == pytest.approx((1.5, 1.5))
    assert dp.gaussian_w2(0.0, 1.0, mean, std) == pytest.approx(math.sqrt(2.5))
    assert math.sqrt(2.5) == pytest.approx(1.581, abs=5e-4)
    assert dp.gaussian_w2(0.0, 1.0, 3.0, 2.0) == pytest.approx(3.162, abs=5e-4)
    grid = torch.linspace(-6.0, 9.0, 15001, dtype=torch.float64)
    peak = float(grid[int(dp.mixture_density(grid).argmax())])
    assert peak == pytest.approx(0.13, abs=5e-3)
    half = torch.tensor(0.5, dtype=torch.float64)
    assert float(dp.velocity_field(torch.tensor(1.5, dtype=torch.float64), half)) == pytest.approx(3.0)
    assert float(dp.velocity_field(torch.tensor(3.0, dtype=torch.float64), half)) == pytest.approx(4.0)
    for t in (0.0, 0.25, 0.5, 0.75, 1.0):
        assert dp.action_density(t, 4096) == pytest.approx(10.0, abs=1e-3)
    assert abs(dp.continuity_residual(1.5, 0.5)) < 1e-12


# %% Sections 7.8 and 7.8.1, ot_flow_matching.py


def test_batch_coupling_costs():
    fm = load("ot_flow_matching")
    x = torch.tensor([[0.0], [2.0]])
    y = torch.tensor([[1.0], [3.0]])
    C = fm.batch_cost_matrix(x, y)
    assert C.flatten().tolist() == pytest.approx([1.0, 9.0, 1.0, 1.0])
    assert float(C.mean()) == pytest.approx(3.0)
    pairing = fm.optimal_pairing(x, y)
    assert pairing.tolist() == [0, 1]
    assert float(C[torch.arange(2), pairing].mean()) == pytest.approx(1.0)
    assert fm.vertex_count(256, 256) == {"unknowns": 65536, "constraints": 512,
                                         "loaded_routes": 511}
    assert fm.conditional_velocity_energy(x, fm.couple(x, y, "optimal")) == pytest.approx(1.0)


def test_flow_matching_pipeline_runs():
    fm = load("ot_flow_matching")
    from cis6270.nets import VelocityMLP

    torch.manual_seed(0)
    model = VelocityMLP(dim=2, width=16)
    x0 = torch.randn(8, 2)
    x1 = torch.randn(8, 2)
    assert float(fm.flow_matching_loss(model, x0, fm.couple(x0, x1, "optimal"))) > 0.0
    samples, straightness = fm.integrate(model, x0, 5)
    assert samples.shape == (8, 2) and straightness >= 0.0


# %% Sections 7.12 to 7.15.3, static_bridge.py


def test_eight_paths_and_the_chain_rule():
    sb = load("static_bridge")
    law = sb.path_law(sb.PI_0, [sb.A, sb.A])
    assert law[(0, 0, 1)] == pytest.approx(0.105)
    assert law[(0, 0, 0)] == pytest.approx(0.245)
    assert law[(1, 1, 1)] == pytest.approx(0.180)
    assert sum(law.values()) == pytest.approx(1.0)
    table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=torch.float64)
    uniform = torch.full((2, 2), 0.25, dtype=torch.float64)
    total, marginal, conditional = sb.chain_rule_terms(table, uniform)
    assert total == pytest.approx(0.2180, abs=5e-5)
    assert marginal == pytest.approx(0.0201, abs=5e-5)
    assert conditional == pytest.approx(0.1979, abs=5e-5)
    assert marginal + conditional == pytest.approx(total)


def test_proportional_fitting_on_the_chain():
    sb = load("static_bridge")
    R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2)
    assert R01.flatten().tolist() == pytest.approx([0.305, 0.195, 0.26, 0.24])
    assert float((R01.sum(dim=0) - sb.PI_1).abs().sum()) == pytest.approx(0.73)
    f, g, errors = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 4)
    assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3)
    assert f.tolist() == pytest.approx([1.072002, 0.937019], abs=5e-7)
    assert g.tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7)
    plan = f[:, None] * R01 * g[None, :]
    assert plan.flatten().tolist() == pytest.approx(
        [0.114605, 0.385395, 0.085395, 0.414605], abs=5e-7)
    assert sb.kl_vectors(plan.flatten(), R01.flatten()) == pytest.approx(0.2820, abs=5e-5)
    independent = sb.PI_0[:, None] * sb.PI_1[None, :]
    assert sb.kl_vectors(independent.flatten(), R01.flatten()) == pytest.approx(0.2847, abs=5e-5)


def test_h_transform_and_the_bridge_path_law():
    sb = load("static_bridge")
    R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2)
    f, g, _ = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 40)
    potentials = sb.backward_potential(sb.A, g, 2)
    assert potentials[2].tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7)
    assert potentials[1].tolist() == pytest.approx([0.798453, 1.246389], abs=5e-7)
    assert potentials[0].tolist() == pytest.approx([0.932834, 1.067215], abs=5e-7)
    tilted = [sb.tilted_matrix(sb.A, potentials[k], potentials[k + 1]) for k in range(2)]
    assert tilted[0].flatten().tolist() == pytest.approx(
        [0.599160, 0.400840, 0.299266, 0.700734], abs=5e-7)
    assert tilted[1].flatten().tolist() == pytest.approx(
        [0.307297, 0.692703, 0.112490, 0.887510], abs=5e-7)
    assert tilted[0].sum(dim=1).tolist() == pytest.approx([1.0, 1.0])
    middle = sb.PI_0 @ tilted[0]
    assert middle.tolist() == pytest.approx([0.449213, 0.550787], abs=5e-7)
    assert (middle @ tilted[1]).tolist() == pytest.approx([0.2, 0.8])
    assert (1.0 / potentials[0]).tolist() == pytest.approx(f.tolist(), abs=1e-9)
    plan = f[:, None] * R01 * g[None, :]
    bridge = sb.fill_with_reference_bridges(plan, sb.A)
    reference = sb.path_law(sb.PI_0, [sb.A, sb.A])
    assert bridge[(0, 0, 0)] == pytest.approx(0.092060, abs=5e-7)
    assert bridge[(0, 0, 1)] == pytest.approx(0.207520, abs=5e-7)
    assert bridge[(1, 1, 1)] == pytest.approx(0.310954, abs=5e-7)
    ratios = {round(bridge[p] / reference[p], 6) for p in reference}
    assert ratios == {0.375755, 1.976383, 0.328441, 1.727522}
    assert sb.relative_entropy(bridge, reference) == pytest.approx(0.281961, abs=5e-7)
    assert sum(v for k, v in bridge.items() if k[1] == 0) == pytest.approx(0.449213, abs=5e-7)


def test_brownian_reference_recovers_the_entropic_plan():
    sb = load("static_bridge")
    plan = sb.brownian_reference_plan(1.0)
    assert plan.flatten().tolist() == pytest.approx([0.2931, 0.3069, 0.0069, 0.3931], abs=5e-5)


# %% Sections 7.16 to 7.18, dynamic_bridge.py


def test_constant_control_two_ways():
    db = load("dynamic_bridge")
    assert db.control_energy_entropy(2.0, 1.0) == pytest.approx(2.0)
    assert db.gaussian_entropy_gap(2.0, 1.0) == pytest.approx(2.0)
    generator = torch.Generator().manual_seed(0)
    _, log_ratio = db.simulate_controlled(2.0, 200, 8192, 1.0, generator)
    assert float(log_ratio.mean()) == pytest.approx(2.0, abs=0.1)


def test_the_bridge_potential_solves_both_equations():
    db = load("dynamic_bridge")
    value = db.hopf_cole(db.bridge_potential)
    for x, t in ((-0.5, 0.2), (0.5, 0.5), (2.5, 0.8)):
        assert abs(db.backward_residual(db.bridge_potential, x, t)) < 1e-12
        assert abs(db.hjb_residual(value, x, t)) < 1e-12
        drift = float(db.bridge_drift(torch.tensor(x, dtype=torch.float64),
                                      torch.tensor(t, dtype=torch.float64)))
        assert db.optimal_control(value, x, t) == pytest.approx(drift, abs=1e-10)


def test_brownian_bridge_numbers():
    db = load("dynamic_bridge")
    mean, variance = db.bridge_law(0.25)
    assert mean == pytest.approx(0.5)
    assert variance == pytest.approx(0.1875)
    assert math.sqrt(variance) == pytest.approx(0.4330, abs=5e-5)
    drift = float(db.bridge_drift(torch.tensor(0.6732, dtype=torch.float64),
                                  torch.tensor(0.25, dtype=torch.float64)))
    assert drift == pytest.approx(1.7691, abs=5e-5)
    assert db.bridge_law(0.5)[1] == pytest.approx(0.25)
    for epsilon, spread in ((0.25, 0.25), (1.0, 0.5), (4.0, 1.0)):
        assert math.sqrt(db.bridge_law(0.5, epsilon=epsilon)[1]) == pytest.approx(spread)


# %% Section 7.19.1, dsb.py


def test_dsb_half_steps_on_the_chain():
    dsb = load("dsb")
    reference = dsb.reference_path_law(dsb.PI_0, dsb.A, 2)
    assert dsb.marginal(reference, 1).tolist() == pytest.approx([0.55, 0.45])
    assert dsb.marginal(reference, 2).tolist() == pytest.approx([0.565, 0.435])
    B2 = dsb.reverse_step(reference, 1)
    assert B2.flatten().tolist() == pytest.approx(
        [0.681416, 0.318584, 0.379310, 0.620690], abs=5e-7)
    assert float(dsb.reverse_step(reference, 0)[0, 0]) == pytest.approx(0.636364, abs=5e-7)
    half = dsb.fix_terminal(reference, dsb.PI_1)
    assert dsb.marginal(half, 2).tolist() == pytest.approx([0.2, 0.8])
    assert dsb.marginal(half, 1).tolist() == pytest.approx([0.439731, 0.560269], abs=5e-7)
    assert dsb.marginal(half, 0).tolist() == pytest.approx([0.466585, 0.533415], abs=5e-7)
    assert float((dsb.marginal(half, 0) - dsb.PI_0).abs().sum()) == pytest.approx(0.0668, abs=5e-5)
    full = dsb.fix_initial(half, dsb.PI_0)
    assert dsb.one_step_conditional(full, 0).flatten().tolist() == pytest.approx(
        [0.599738, 0.400262, 0.299771, 0.700229], abs=5e-7)
    assert dsb.one_step_conditional(full, 1).flatten().tolist() == pytest.approx(
        [0.309924, 0.690076, 0.113725, 0.886275], abs=5e-7)
    assert dsb.marginal(full, 2).tolist() == pytest.approx([0.201967, 0.798033], abs=5e-7)
    gap = max(float((dsb.one_step_conditional(full, 0) - dsb.TILTED_0).abs().max()),
              float((dsb.one_step_conditional(full, 1) - dsb.TILTED_1).abs().max()))
    assert gap < 0.003
    _, errors, _ = dsb.proportional_fitting(reference, dsb.PI_0, dsb.PI_1, 4)
    assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3)
    assert errors[0] / errors[1] == pytest.approx(187.5, abs=0.1)


# %% Section 7.19.2, dsbm.py


def test_markov_drift_at_one_point():
    dm = load("dsbm")
    assert dm.bridge_drift(3.0, 1.0, 0.5) == pytest.approx(4.0)
    assert dm.bridge_drift(0.0, 1.0, 0.5) == pytest.approx(-2.0)
    endpoints = tensor(3.0, 0.0)
    weights = tensor(0.5, 0.5)
    assert dm.markov_drift(endpoints, weights, 1.0, 0.5) == pytest.approx(1.0)
    assert dm.least_squares_derivative(1.0, endpoints, weights, 1.0, 0.5) == pytest.approx(0.0)
    assert dm.least_squares_derivative(0.0, endpoints, weights, 1.0, 0.5) == pytest.approx(-2.0)


def test_one_imf_round_on_the_chain():
    dm = load("dsbm")
    bridges = dm.reference_bridge_table(dm.A)
    assert float(bridges[0, 0, 0]) == pytest.approx(0.803279, abs=5e-7)
    assert float(bridges[0, 0, 1]) == pytest.approx(0.538462, abs=5e-7)
    target = dm.bridge_coupling(dm.A, dm.PI_0, dm.PI_1)
    independent = dm.PI_0[:, None] * dm.PI_1[None, :]
    assert independent.flatten().tolist() == pytest.approx([0.1, 0.4, 0.1, 0.4])
    assert float((independent - target).abs().sum()) == pytest.approx(0.0584, abs=5e-5)
    coupling, gaps, first = dm.imf_rounds(independent, dm.A, 9, target)
    assert first[0].flatten().tolist() == pytest.approx(
        [0.591425, 0.408575, 0.307692, 0.692308], abs=5e-7)
    assert first[1].flatten().tolist() == pytest.approx(
        [0.298457, 0.701543, 0.119588, 0.880412], abs=5e-7)
    projected = torch.diag(dm.PI_0) @ first[0] @ first[1]
    assert projected.flatten().tolist() == pytest.approx(
        [0.112688, 0.387312, 0.087312, 0.412688], abs=5e-7)
    assert projected.sum(dim=0).tolist() == pytest.approx([0.2, 0.8])
    assert projected.sum(dim=1).tolist() == pytest.approx([0.5, 0.5])
    assert gaps[0] == pytest.approx(0.0077, abs=5e-5)
    assert gaps[3] / gaps[4] == pytest.approx(7.37, abs=5e-3)
    assert gaps[8] < 1e-9


# %% Section 7.19.3, sf2m.py


def test_the_two_regression_targets():
    sm = load("sf2m")
    x0, x1 = tensor(0.0)[0], tensor(2.0)[0]
    z = tensor(0.4)[0]
    for t_value, mean, sample, drift, flow, coefficient in (
        (0.25, 0.5, 0.6732, 1.7691, 2.2309, 4 / 3),
        (0.75, 1.5, 1.6732, 1.3072, 1.7691, -4 / 3),
    ):
        t = torch.tensor(t_value, dtype=torch.float64)
        assert float(sm.bridge_mean(x0, x1, t)) == pytest.approx(mean)
        assert float(sm.bridge_std(t)) == pytest.approx(0.4330, abs=5e-5)
        x = sm.bridge_sample(x0, x1, t, z)
        assert float(x) == pytest.approx(sample, abs=5e-5)
        assert float(sm.conditional_drift(x1, x, t)) == pytest.approx(drift, abs=5e-5)
        assert float(sm.conditional_score(x0, x1, x, t)) == pytest.approx(-0.9238, abs=5e-5)
        assert float(sm.flow_velocity(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5)
        assert float(sm.flow_velocity_collected(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5)
        assert float(sm.flow_coefficient(t)) == pytest.approx(coefficient)
        assert float(sm.scaled_score(x0, x1, x, t)) == pytest.approx(-0.4, abs=5e-7)


def test_the_batch_plan_of_sf2m():
    sm = load("sf2m")
    x, y = tensor(0.0, 2.0), tensor(1.0, 3.0)
    uniform = tensor(0.5, 0.5)
    plan = sm.entropic_plan(x, y, uniform, uniform, 1.0)
    assert plan.sum(dim=0).tolist() == pytest.approx([0.5, 0.5])
    assert float(plan[0, 0] + plan[1, 1]) > float(plan[0, 1] + plan[1, 0])


# %% Sections 7.21 to 7.23.3, discrete_bridge.py


def test_rate_error_is_not_symmetric():
    db = load("discrete_bridge")
    assert db.scaling_cost(0.5) == pytest.approx(0.1534, abs=5e-5)
    assert db.scaling_cost(2.0) == pytest.approx(0.3863, abs=5e-5)
    assert float(db.rate_error(tensor(1.0)[0], tensor(1.0)[0])) == pytest.approx(0.0)


def test_two_state_continuous_time_bridge():
    db = load("discrete_bridge")
    assert torch.matrix_exp(db.Q).flatten().tolist() == pytest.approx(
        [0.6328, 0.3672, 0.5507, 0.4493], abs=5e-5)
    R01 = db.endpoint_law(db.Q, db.PI_0)
    assert R01.flatten().tolist() == pytest.approx([0.3164, 0.1836, 0.2754, 0.2246], abs=5e-5)
    assert R01.sum(dim=0).tolist() == pytest.approx([0.5918, 0.4082], abs=5e-5)
    f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1)
    assert f.tolist() == pytest.approx([1.0716, 0.9373], abs=5e-5)
    assert g.tolist() == pytest.approx([0.3349, 1.9643], abs=5e-5)
    assert float(g[1] / g[0]) == pytest.approx(5.865, abs=5e-4)
    assert plan.flatten().tolist() == pytest.approx([0.1136, 0.3864, 0.0864, 0.4136], abs=5e-5)
    assert torch.matrix_exp(db.Q / 2).flatten().tolist() == pytest.approx(
        [0.7146, 0.2854, 0.4281, 0.5719], abs=5e-5)
    phi = db.backward_potential(db.Q, g, 0.5)
    assert phi.tolist() == pytest.approx([0.7999, 1.2667], abs=5e-5)
    tilted = db.tilted_generator(db.Q, phi)
    assert tilted.flatten().tolist() == pytest.approx(
        [-1.5836, 1.5836, 0.9472, -0.9472], abs=5e-5)
    assert db.backward_potential(db.Q, g, 0.0).tolist() == pytest.approx([0.9331, 1.0669], abs=5e-5)
    at_zero = db.tilted_generator(db.Q, db.backward_potential(db.Q, g, 0.0))
    assert float(at_zero[0, 1]) == pytest.approx(1.1433, abs=5e-5)
    assert float(at_zero[1, 0]) == pytest.approx(1.3120, abs=5e-5)
    at_one = db.tilted_generator(db.Q, g)
    assert float(at_one[0, 1]) == pytest.approx(5.8652, abs=5e-5)
    assert float(at_one[1, 0]) == pytest.approx(0.2557, abs=5e-5)
    marginal = db.bridge_marginal(db.Q, db.PI_0, f, g, 0.5)
    assert marginal.tolist() == pytest.approx([0.4668, 0.5332], abs=5e-5)
    assert db.bridge_marginal(db.Q, db.PI_0, f, g, 1.0).tolist() == pytest.approx([0.2, 0.8])
    reverse = db.reverse_rates(tilted, marginal)
    assert float(reverse[1, 0]) == pytest.approx(1.3862, abs=5e-5)
    assert float(reverse[0, 1]) == pytest.approx(1.0821, abs=5e-5)
    assert float(marginal[0] * tilted[0, 1]) == pytest.approx(
        float(marginal[1] * reverse[1, 0]), abs=5e-5)


def test_euler_sampler_and_the_cost_computed_twice():
    db = load("discrete_bridge")
    R01 = db.endpoint_law(db.Q, db.PI_0)
    f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1)
    errors = {}
    for steps, expected in ((10, 0.0223), (20, 0.0106), (40, 0.0051), (100, 0.0021)):
        terminal, leaving = db.euler_marginal(db.Q, db.PI_0, f, g, steps)
        errors[steps] = float((terminal - db.PI_1).abs().sum() / 2)
        assert errors[steps] == pytest.approx(expected, abs=1e-4)
        if steps == 10:
            assert leaving[:3] == pytest.approx([0.1143, 0.1187, 0.1245], abs=5e-5)
            assert leaving[-1] == pytest.approx(0.3649, abs=5e-5)
    assert db.relative_entropy(plan, R01) == pytest.approx(0.3235, abs=5e-5)
    integral, at_half = db.path_cost(db.Q, db.PI_0, f, g)
    assert at_half == pytest.approx(0.1300, abs=5e-5)
    assert integral == pytest.approx(0.3235, abs=5e-5)
    expected = {0.5: 0.3148, 1.0: 0.3235, 2.0: 0.3337, 4.0: 0.3348}
    for lam, value in expected.items():
        law = db.endpoint_law(lam * db.Q, db.PI_0)
        _, _, scaled = db.sinkhorn(law, db.PI_0, db.PI_1)
        assert db.relative_entropy(scaled, law) == pytest.approx(value, abs=5e-5)
    assert db.relative_entropy(db.PI_1, db.STATIONARY) == pytest.approx(0.3348, abs=5e-5)


# %% Section 7.24, ddsbm.py


def test_ddsbm_single_bond():
    dd = load("ddsbm")
    kernel = dd.replacement_kernel(0.4, dd.PRIOR)
    assert kernel.flatten().tolist() == pytest.approx(
        [0.6, 0.2, 0.2, 0.2, 0.6, 0.2, 0.2, 0.2, 0.6])
    reference = torch.diag(dd.BOND_0) @ kernel
    assert reference.sum(dim=0).tolist() == pytest.approx([0.48, 0.28, 0.24])
    assert (dd.BOND_1 / reference.sum(dim=0)).tolist() == pytest.approx(
        [0.2083, 0.7143, 2.9167], abs=5e-5)
    f, g, plan, errors = dd.sinkhorn(reference, dd.BOND_0, dd.BOND_1, 200)
    assert f.tolist() == pytest.approx([1.1779, 0.9506, 0.5062], abs=5e-5)
    assert g.tolist() == pytest.approx([0.1842, 0.6918, 3.0004], abs=5e-5)
    assert errors[-1] < 1e-12
    assert float(plan[0, 2]) == pytest.approx(0.495, abs=5e-4)
    assert dd.held_mass(plan) == pytest.approx(0.261, abs=5e-4)
    assert float((dd.BOND_0 * dd.BOND_1).sum()) == pytest.approx(0.18)
    assert float(torch.minimum(dd.BOND_0, dd.BOND_1).sum()) == pytest.approx(0.4)
    expected = {0.2: 0.2250, 0.4: 0.2612, 0.6: 0.2926, 0.8: 0.3248, 0.99: 0.3812}
    for r, value in expected.items():
        held = dd.sinkhorn(torch.diag(dd.BOND_0) @ dd.replacement_kernel(r, dd.PRIOR),
                           dd.BOND_0, dd.BOND_1, 200)[2]
        assert dd.held_mass(held) == pytest.approx(value, abs=5e-5)


def test_ddsbm_pinned_rates_and_the_posterior_average():
    dd = load("ddsbm")
    beta = dd.resampling_rate(0.4)
    assert beta == pytest.approx(0.9163, abs=5e-5)
    rates = dd.replacement_generator(beta, dd.PRIOR)
    assert float(rates[0, 1]) == pytest.approx(0.3054, abs=5e-5)
    assert rates.sum(dim=1).abs().max() < 1e-12
    assert dd.half_survival_time(0.4) == pytest.approx(math.log(0.8) / math.log(0.4))
    assert dd.half_survival_time(0.4) == pytest.approx(0.2435, abs=5e-5)
    pinned = dd.replacement_kernel(0.5, dd.PRIOR)
    assert float(pinned[0, 0]) == pytest.approx(2 / 3, abs=5e-5)
    assert float(pinned[0, 1]) == pytest.approx(1 / 6, abs=5e-5)
    toward = dd.pinned_rate(rates, pinned, 0, 2, 2)
    away = dd.pinned_rate(rates, pinned, 2, 0, 2)
    between = dd.pinned_rate(rates, pinned, 0, 1, 2)
    assert toward == pytest.approx(1.2217, abs=5e-5)
    assert away == pytest.approx(0.0764, abs=5e-5)
    assert between == pytest.approx(0.3054, abs=5e-5)
    assert toward / away == pytest.approx(16.0, abs=5e-4)
    assert dd.rate_error(toward, 1.0) == pytest.approx(0.0229, abs=5e-5)
    posterior = 0.6 * toward + 0.4 * between
    assert posterior == pytest.approx(0.8552, abs=5e-5)
    rates_two = [toward, between]
    weights = [0.6, 0.4]
    assert dd.posterior_loss(rates_two, weights, posterior) == pytest.approx(0.1357, abs=5e-5)
    assert dd.posterior_loss(rates_two, weights, 1.0) == pytest.approx(0.1467, abs=5e-5)
    assert dd.posterior_loss(rates_two, weights, 0.5) == pytest.approx(0.2395, abs=5e-5)


def test_ddsbm_factorization_counts():
    dd = load("ddsbm")
    assert dd.feature_count(9) == (9, 36, 45)
    assert dd.state_space_size(9, 5, 4) == pytest.approx(9.2e27, rel=5e-3)
    assert dd.total_jump_rate(36, dd.resampling_rate(0.4), 4) == pytest.approx(24.74, abs=5e-3)


# %% Section 7.25, csbm.py


def test_csbm_references():
    cs = load("csbm")
    A = cs.uniform_reference(0.35, 2)
    assert A.flatten().tolist() == pytest.approx([0.65, 0.35, 0.35, 0.65])
    assert torch.linalg.matrix_power(A, 3).flatten().tolist() == pytest.approx(
        [0.5135, 0.4865, 0.4865, 0.5135], abs=5e-5)
    assert cs.uniform_reference(0.35, 5).sum(dim=1).tolist() == pytest.approx([1.0] * 5)
    ordered = cs.ordered_reference(1.0, 5)
    assert ordered[2].tolist() == pytest.approx(
        [0.0545, 0.2442, 0.4026, 0.2442, 0.0545], abs=5e-5)
    assert cs.transition_loss(tensor(0.3160, 0.6840), tensor(0.25, 0.75)) == pytest.approx(
        0.0110, abs=5e-5)


def test_csbm_dimf_on_sixteen_paths():
    cs = load("csbm")
    A = cs.uniform_reference(cs.ALPHA, 2)
    plan = cs.sinkhorn(torch.diag(cs.PI_0) @ torch.linalg.matrix_power(A, 3), cs.PI_0, cs.PI_1)
    assert plan.flatten().tolist() == pytest.approx(
        [0.10432, 0.39568, 0.09568, 0.40432], abs=5e-6)
    assert float(cs.bridge_table(A, 3)[0, 0, 0, 0]) == pytest.approx(0.5348, abs=5e-5)
    gap, start, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1)
    assert float(target[0, 0, 0, 0]) == pytest.approx(0.0558, abs=5e-5)
    assert float(target[1, 1, 1, 1]) == pytest.approx(0.2162, abs=5e-5)
    assert float(start[0, 0, 0, 0]) == pytest.approx(0.0535, abs=5e-5)
    assert float(start[1, 1, 1, 1]) == pytest.approx(0.2139, abs=5e-5)
    assert gap == pytest.approx(0.0173, abs=5e-5)
    for axis, value in ((1, [0.4754, 0.5246]), (2, [0.4107, 0.5893])):
        marginal = start.sum(dim=[d for d in range(4) if d != axis])
        assert marginal.tolist() == pytest.approx(value, abs=5e-5)
    gaps, conditionals = cs.d_imf(cs.PI_0[:, None] * cs.PI_1[None, :], A, 3, 5, target)
    assert conditionals[0].flatten().tolist() == pytest.approx(
        [0.6243, 0.3757, 0.3265, 0.6735], abs=5e-5)
    assert conditionals[1].flatten().tolist() == pytest.approx(
        [0.5626, 0.4374, 0.2729, 0.7271], abs=5e-5)
    assert conditionals[2].flatten().tolist() == pytest.approx(
        [0.3160, 0.6840, 0.1192, 0.8808], abs=5e-5)
    assert gaps == pytest.approx([2.9e-4, 5.0e-6, 8.4e-8, 1.4e-9, 2.4e-11], rel=0.05)
    assert gaps[0] / gaps[1] == pytest.approx(59.0, abs=1.0)


def test_csbm_mixing_and_length():
    cs = load("csbm")
    independent = cs.PI_0[:, None] * cs.PI_1[None, :]
    expected = {0.2: (0.13589, 5.0), 0.4: (0.00512, 284.3)}
    for alpha, (start, ratio) in expected.items():
        A = cs.uniform_reference(alpha, 2)
        gap, _, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1)
        assert gap == pytest.approx(start, abs=5e-5)
        later, _ = cs.d_imf(independent, A, 3, 2, target)
        assert later[0] / later[1] == pytest.approx(ratio, rel=0.01)
    A = cs.uniform_reference(cs.ALPHA, 2)
    for steps, (start, ratio) in {2: (0.0574, 7.3), 5: (0.00156, 4674.0)}.items():
        gap, _, target = cs.starting_gap(A, steps, cs.PI_0, cs.PI_1)
        assert gap == pytest.approx(start, abs=5e-5)
        later, _ = cs.d_imf(independent, A, steps, 2, target)
        assert later[0] / later[1] == pytest.approx(ratio, rel=0.01)


# %% Sections 7.26 and 7.27, branched.py


def test_two_branches_worked_out():
    br = load("branched")
    linear = lambda t: 0.3 * t
    late = lambda t: 0.3 * t ** 2
    assert float(br.branch_weights(linear, torch.tensor(0.5))) == pytest.approx(0.15)
    assert float(br.branch_weights(linear, torch.tensor(1.0))) == pytest.approx(0.3)
    near_zero = br.weighted_energy(2.0, lambda t: 1.0 - linear(t))
    near_one = br.weighted_energy(-2.0, linear)
    assert near_zero == pytest.approx(1.7, abs=5e-5)
    assert near_one == pytest.approx(0.3, abs=5e-5)
    assert near_zero + near_one == pytest.approx(2.0, abs=5e-5)
    far_linear = near_zero + br.weighted_energy(-4.0, linear)
    far_late = br.weighted_energy(2.0, lambda t: 1.0 - late(t)) + br.weighted_energy(-4.0, late)
    assert far_linear == pytest.approx(2.9, abs=5e-5)
    assert far_late == pytest.approx(2.6, abs=5e-5)


def test_the_interpolant():
    br = load("branched")
    grid = torch.tensor([0.0, 0.5, 1.0])
    assert br.interpolant(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([0.0, 1.25, 2.0])
    assert br.interpolant_velocity(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([3.0, 2.0, 1.0])
    assert br.interpolant_energy(0.0, 2.0, 1.0) == pytest.approx(2.1667, abs=5e-5)
    assert br.interpolant_energy(0.0, 2.0, 0.0) == pytest.approx(2.0, abs=5e-5)


def test_entangled_tilt_and_parameterization():
    br = load("branched")
    normalizer = 0.1 * 1.0 + 0.9 * 0.01
    assert normalizer == pytest.approx(0.109)
    tilted = 0.1 / normalizer
    assert tilted == pytest.approx(0.917, abs=5e-4)
    assert tilted / 0.1 == pytest.approx(9.17, abs=5e-3)
    reward = (1.0 - tilted) * math.log(0.01)
    assert reward == pytest.approx(-0.380, abs=5e-4)
    assert reward - math.log(normalizer) == pytest.approx(1.836, abs=5e-4)
    assert -math.log(normalizer) == pytest.approx(2.216, abs=5e-4)
    assert br.NOISE == pytest.approx(1.4142, abs=5e-5)
    step = (-0.3 + 0.5) * 0.01 - 1.0 * 0.8 * 0.01 + br.NOISE * math.sqrt(0.01) * 0.4
    assert step == pytest.approx(0.0506, abs=5e-5)
    assert 0.5 / br.NOISE == pytest.approx(0.3536, abs=5e-5)
    assert 0.5 * (0.5 / br.NOISE) ** 2 * 0.01 == pytest.approx(0.000625, abs=5e-7)
    force = br.compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]),
                            torch.tensor([[[1.0, 0.0]]]))[0, 0]
    assert force.tolist() == pytest.approx([0.5, -0.4])
    assert float(force.norm()) == pytest.approx(0.640, abs=5e-4)
    weights = torch.softmax(torch.tensor([-1.0, 0.0, 1.0]), dim=0)
    assert weights.tolist() == pytest.approx([0.090, 0.245, 0.665], abs=5e-4)


def test_entangled_training_lowers_the_loss():
    br = load("branched")
    torch.manual_seed(0)
    generator = torch.Generator().manual_seed(0)
    force = br.BiasForce(width=16)
    rollout = br.roll_out(force, 16, 8, torch.device("cpu"), generator)
    weights = br.importance_weights(rollout)
    assert float(weights.sum()) == pytest.approx(1.0)
    assert br.force_check(rollout["positions"][-1]) < 1e-5
    loss = br.cross_entropy_loss(force, rollout, weights)
    assert torch.isfinite(loss)


# %% Every file runs end to end


@pytest.mark.parametrize("name,flags", [
    ("optimal_transport", ["--steps", "300"]),
    ("sinkhorn", ["--steps", "20"]),
    ("displacement", ["--steps", "64"]),
    ("ot_flow_matching", ["--steps", "3", "--samples", "64", "--sample-steps", "5"]),
    ("static_bridge", ["--steps", "3"]),
    ("dynamic_bridge", ["--steps", "20"]),
    ("dsb", ["--steps", "2"]),
    ("dsbm", ["--steps", "3"]),
    ("sf2m", ["--steps", "21"]),
    ("discrete_bridge", ["--steps", "4"]),
    ("ddsbm", ["--steps", "20"]),
    ("csbm", ["--steps", "3"]),
    ("branched", ["--steps", "2", "--samples", "16"]),
])
def test_the_script_runs(name, flags, monkeypatch, capsys):
    module = load(name)
    monkeypatch.setattr(sys, "argv", [name] + flags + ["--quiet"])
    module.main()
    assert capsys.readouterr().out.strip()