Blame for tests/test_lecture_5.py

"""The worked numbers of Chapter 5, asserted against the files of ``lecture_5/``.

Every number checked here is printed in the notes. Where the chapter runs a simulation, the test
asserts the analytic value the simulation estimates, with a tolerance wide enough for the sample
size the file uses by default.
"""

from __future__ import annotations

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

import pytest
import torch

from cis6270.nets import SequenceDenoiser

ROOT = Path(__file__).resolve().parents[1]


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


dfm = load("dfm")
velocity = load("velocity")
corrector = load("corrector")
kinetic = load("kinetic_optimal")
edits = load("edit_flows")
dirichlet = load("dirichlet")
fisher = load("fisher")
gumbel = load("gumbel")
guidance = load("guidance")
objectives = load("multi_objective")
rectify = load("rectify")


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


# %% Section 5.2 and 5.3: a velocity, its two rules, and the master equation


def test_worked_master_equation():
    u, p, p_dot, arriving, leaving = velocity.worked_master_equation()
    assert velocity.obeys_the_two_rules(u)
    assert torch.allclose(u, velocity.set_diagonal(u))
    close(p_dot.tolist(), [-0.1, -0.2, 0.7, -0.4])
    assert abs(float(p_dot.sum())) < 1e-9
    assert abs(arriving - 0.7) < 1e-9 and abs(leaving) < 1e-9
    close((p + 0.1 * p_dot).tolist(), [0.09, 0.18, 0.37, 0.36])


def test_hazard_and_survival():
    assert abs(velocity.hazard(0.5) - 2.0) < 1e-9
    assert abs(velocity.hazard(0.9) - 10.0) < 1e-9
    assert abs(velocity.survival(4000) - 0.1) < 2e-3  # S(t) = 1 - kappa_t.


def test_continuous_and_discrete_side_by_side():
    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])
    moved, drawn, lam = velocity.side_by_side(2, p1, t=0.5, h=0.1)
    close(moved.tolist(), [0.02, 0.12, 0.84, 0.02])
    close(drawn.tolist(), [0.02, 0.12, 0.84, 0.02])
    assert abs(lam - 0.2) < 1e-9


# %% Sections 5.5 to 5.11: the mixture path, the velocity, the loss and the sampler


def test_dfm_worked_numbers():
    report = dfm.worked_numbers()
    assert report["mixture_draw"] == "ATGA"
    close(report["velocity_at_G"], [0.2, 1.2, -1.6, 0.2])
    close(report["euler_step_at_G"], [0.02, 0.12, 0.84, 0.02])
    assert abs(report["redraw_probability"] - 0.2) < 1e-9
    assert abs(report["example_loss_nats"] - 4.780) < 1e-3
    close(report["mask_posterior_C_G"], [0.625, 0.375])
    assert report["mask_posterior_unchanged"]  # Proposition 6 of Discrete Flow Matching.
    assert abs(report["posterior_entropy_nats"] - 0.6616) < 5e-4
    assert abs(report["kl_gap_to_uniform"] - 0.0316) < 5e-4


def test_conditional_velocity_generates_the_conditional_path():
    """Section 5.6.2: the two terms of the master equation on the mixture path."""
    kappa, x0, x1 = 0.4, 0, 1
    path = torch.zeros(4, dtype=torch.double)
    path[x0], path[x1] = 1.0 - kappa, kappa
    total = sum(path[z] * dfm.conditional_velocity(z, x1, kappa) for z in (x0, x1))
    derivative = torch.zeros(4, dtype=torch.double)
    derivative[x1], derivative[x0] = 1.0, -1.0  # kappa_dot_t times the two indicators.
    assert torch.allclose(total, derivative, atol=1e-9)


def test_corrupt_reaches_both_endpoints():
    torch.manual_seed(0)
    process = dfm.DiscreteFlowMatching(4)
    x1 = torch.randint(0, 4, (64, 16))
    assert torch.equal(process.corrupt(x1, torch.ones(64)), x1)  # kappa_t = 1 gives the data.
    source = process.corrupt(x1, torch.zeros(64))  # kappa_t = 0 gives the source.
    assert float((source == x1).double().mean()) < 0.4  # A uniform draw agrees about a quarter of the time.
    halfway = process.corrupt(x1, torch.full((64,), 0.5))
    assert 0.55 < float((halfway == x1).double().mean()) < 0.75  # Heads with probability one half.


def test_two_sources():
    assert torch.equal(dfm.DiscreteFlowMatching(4, mask=4).source((3, 5)),
                       torch.full((3, 5), 4, dtype=torch.long))  # The all-mask source.
    uniform = dfm.DiscreteFlowMatching(4).source((200, 20))
    assert int(uniform.min()) == 0 and int(uniform.max()) == 3  # Every base appears.


# %% Sections 5.13 and 5.14: zero-flux velocities and the corrector


def test_zero_flux_velocity_on_three_states():
    p, u, v, with_u, with_both = corrector.three_state_example()
    assert torch.allclose(corrector.net_flux(v, p), torch.zeros(3), atol=1e-12)
    close(with_u.tolist(), [-0.2, -0.3, 0.5])
    close(with_both.tolist(), [-0.2, -0.3, 0.5])
    assert abs(float(-u[0, 0]) - 1.0) < 1e-12 and abs(float(-(u + v)[0, 0]) - 5.0) < 1e-12


def test_corrector_step_table():
    """Section 5.14.3 at t = 0.4 with h = 0.01 and the denoiser (0.1, 0.6, 0.2, 0.1)."""
    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])
    expected = {
        0.0: ([0.00167, 0.01000, 0.00333, 0.00167], 0.98333),
        2.0: ([0.00300, 0.01800, 0.00600, 0.00300], 0.97000),
        10.0: ([0.00833, 0.05000, 0.01667, 0.00833], 0.91667),
    }
    for eta, (probabilities, stay) in expected.items():
        bar_u = corrector.corrector_velocity(p1, corrector.MASK, 0.4, eta)
        close((0.01 * bar_u[:4]).tolist(), probabilities, tolerance=1e-5)
        assert abs(float(1.0 + 0.01 * bar_u[corrector.MASK]) - stay) < 1e-5
        assert abs(corrector.masked_fraction_derivative(0.4, eta) + 1.0) < 1e-12


def test_corrector_rates_and_admissible_steps():
    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])
    bar_u = corrector.corrector_velocity(p1, corrector.MASK, 0.4, 2.0)
    assert abs(float(-bar_u[corrector.MASK]) - 3.0) < 1e-12  # (1 + 2 x 0.4)/(1 - 0.4).
    visible = torch.zeros(4)
    visible[1] = 1.0
    assert abs(float(corrector.corrector_velocity(visible, 1, 0.4, 2.0)[corrector.MASK]) - 2.0) < 1e-12
    close([corrector.admissible_step(0.9, eta) for eta in (0.0, 2.0, 10.0)],
          [0.1, 0.0357, 0.01], tolerance=1e-4)


def test_jump_counts_are_one_plus_eta():
    """Section 5.14.3: the masked fraction stays at 1 - t and the jumps come to 1 + eta."""
    p1 = torch.tensor([0.25, 0.25, 0.25, 0.25])
    for eta in (0.0, 2.0, 10.0):
        half, mean = corrector.simulate_jumps(p1, eta, positions=10_000, steps=300)
        assert abs(half - 0.5) < 0.03
        assert abs(mean - (1.0 + eta)) < 0.25


def test_parallel_step_error():
    assert abs(corrector.joint_reveal_rate(64) - 1.0 / 64) < 1e-12
    assert abs(corrector.pair_survival(2.0) - 0.1353) < 5e-4
    assert abs(corrector.pair_survival(10.0) - 4.54e-5) < 1e-7
    measured = corrector.simulate_pair(0.0, samples=40_000, steps=64)
    assert abs(measured - 0.5 / 64) < 2e-3  # Half of the joint reveals land outside the data.


# %% Section 5.16: kinetic-optimal paths


def test_kinetic_optimal_recovers_the_mixture_velocity():
    q = torch.tensor([0.1, 0.6, 0.2, 0.1])
    p, p_dot = kinetic.mask_path(q, kappa_t=0.5)
    u = kinetic.kinetic_optimal(p, p_dot)
    assert abs(float(p_dot[kinetic.MASK] / p[kinetic.MASK]) + 2.0) < 1e-12
    assert abs(float(p_dot[1] / p[1]) - 2.0) < 1e-12
    assert abs(float(u[1, kinetic.MASK]) - 1.2) < 1e-12  # 0.30 times 4.
    assert abs(float(u[1, kinetic.MASK]) - kinetic.mixture_rate(0.5) * 0.6) < 1e-12
    assert abs(float(u[1, 2])) < 1e-12  # No probability moves between two visible tokens.
    assert torch.allclose(u @ p, p_dot, atol=1e-12)


def test_worked_metric_path_for_dna():
    beta_t, beta_dot = kinetic.beta_schedule(0.5)
    assert abs(beta_t - 1.0) < 1e-12 and abs(beta_dot - 4.0) < 1e-12
    p = kinetic.metric_path(2, beta_t)
    close(p.tolist(), [0.2245, 0.0826, 0.6103, 0.0826])
    u = kinetic.metric_velocity(2, beta_t, beta_dot)
    assert abs(float(u[0, 1]) - 0.898) < 5e-4  # C to A.
    assert abs(float(u[2, 1]) - 4.882) < 5e-4  # C to G.
    assert abs(float(u[2, 0]) - 2.441) < 5e-4  # A to G.
    assert abs(float(-u[1, 1]) - 5.780) < 1e-3  # The total rate out of C.
    close((torch.eye(4)[1] + 0.05 * u[:, 1]).tolist(), [0.0449, 0.7110, 0.2441, 0.0])
    close((u @ p).tolist(), [-0.3997, -0.4774, 1.3546, -0.4774])
    close(kinetic.path_derivative(2, 0.5).tolist(), [-0.3997, -0.4774, 1.3546, -0.4774], 1e-3)


def test_energy_of_two_schedules():
    assert abs(kinetic.schedule_energy(*kinetic.LINEAR, 600, 1e-6) - 27.63) < 5e-3
    assert abs(kinetic.schedule_energy(*kinetic.LINEAR, 600, 1e-9) - 41.45) < 5e-3
    assert abs(kinetic.schedule_energy(*kinetic.GEODESIC, 600, 1e-6) - math.pi ** 2) < 1e-3
    assert abs(kinetic.geodesic_schedule(0.25) - 0.1464) < 5e-4
    assert abs(kinetic.geodesic_schedule(0.75) - 0.8536) < 5e-4
    assert abs(kinetic.geodesic_schedule(0.5, 0.25) - 0.6667) < 5e-4


# %% Section 5.17: Edit Flows


def test_alignment_and_column_wise_path():
    z0 = [0, 1, 2, 3, edits.BLANK]  # (A, C, G, T, eps).
    z1 = [0, edits.BLANK, 2, 0, 1]  # (A, eps, G, A, C).
    assert edits.column_kinds(z0, z1) == ["none", "del", "none", "sub", "ins"]
    drawn = edits.draw_state(z0, z1, [True, False, True, True, True])
    assert drawn == [0, 1, 2, 0, 1]  # ACGAC, the draw of probability 0.5 ** 3 = 0.125.
    assert abs(edits.expected_length([0, 1, 2, 3], [0, 2, 0, 1], 0.5) - 4.0) < 1e-12
    recovered = edits.align([0, 1, 2, 3], [0, 2, 0, 1])
    assert len(recovered[0]) == len(recovered[1])  # Any alignment pads to a common width.


def test_alignment_posterior_and_averaged_rates():
    first = ([0, 1, edits.BLANK], [0, 0, 1])  # Substitute C with A, then insert C.
    second = ([0, edits.BLANK, 1], [0, 0, 1])  # Insert A, and leave the last column alone.
    for kappa_t, expected in ((0.25, 0.4286), (0.5, 1 / 3), (0.75, 0.2)):
        weights = edits.alignment_posterior([first, second], [0, 1], kappa_t)
        assert abs(weights[0] - expected) < 5e-4
    posterior, rates = edits.averaged_rates([first, second], [0, 1], 0.5)
    assert abs(rates["sub(2,A)"] - 0.667) < 1e-3
    assert abs(rates["ins(1,A)"] - 1.333) < 1e-3
    assert abs(rates["ins(2,C)"] - 0.667) < 1e-3
    assert abs(sum(rates.values()) - 2.667) < 1e-3


def test_edit_loss_and_sampler_step():
    assert abs(edits.edit_loss(3.0, [0.8], kappa_t=0.5) - 3.446) < 1e-3
    grid = [0.01 * i for i in range(1, 801)]
    assert abs(edits.loss_minimizer(2.0, grid) - 2.0) < 1e-6
    intensities = {"ins": [0.2, 0.3, 0.4], "del": [0.1, 0.9, 0.2], "sub": [0.4, 0.3, 0.5]}
    insert, change = edits.step_probabilities(intensities, 0.05)
    close(insert, [0.010, 0.015, 0.020])
    close(change, [0.025, 0.060, 0.035])
    exact, first_order = edits.untouched_probability(intensities, 0.05)
    assert abs(exact - 0.8452) < 5e-4 and abs(first_order - 0.835) < 5e-4
    assert abs(exact - first_order - 0.0102) < 5e-4
    assert abs(edits.expected_length_change(intensities, 0.05) + 0.015) < 1e-9
    assert edits.edit_rates(intensities)[0] == 11  # Three intensities and two token distributions.


# %% Section 5.20: Dirichlet Flow Matching


def test_dirichlet_path_and_speeds():
    assert abs(dirichlet.dirichlet_density(torch.full((4,), 0.25), 0, 0.0) - 6.0) < 1e-9
    assert dirichlet.straight_line_corners(torch.tensor([0.34, 0.33, 0.33, 0.0]), 0.3) == 3
    assert dirichlet.straight_line_corners(torch.tensor([0.70, 0.30, 0.0, 0.0]), 0.6) == 1
    closed = dirichlet.speed_two_letters(0.3, 1.0)
    assert abs(closed - 0.2580) < 5e-5
    assert abs(float(dirichlet.speed(0.3, 1.0, K=2)) - closed) < 1e-6
    x = torch.tensor([0.40, 0.30, 0.20, 0.10])
    close(dirichlet.speed_vector(x, 1.0).tolist(), [0.232, 0.199, 0.159, 0.104], 1e-3)


def test_dirichlet_velocity_and_posterior_average():
    x = torch.tensor([0.40, 0.30, 0.20, 0.10])
    toward_a = dirichlet.conditional_velocity(x, 0, 1.0)
    close(toward_a.tolist(), [0.139, -0.070, -0.046, -0.023], 1e-3)
    assert abs(float(toward_a.sum())) < 1e-12
    p_hat = torch.tensor([0.60, 0.25, 0.10, 0.05])
    v = dirichlet.marginal_velocity(x, p_hat, 1.0)
    close(v.tolist(), [0.0553, -0.0133, -0.0262, -0.0158], 2e-4)
    close((x + 0.1 * v).tolist(), [0.4055, 0.2987, 0.1974, 0.0984])
    assert abs(float((x + 0.1 * v).sum()) - 1.0) < 1e-12


def test_dirichlet_continuity():
    time_term, space_term = dirichlet.continuity_terms(torch.tensor([0.5, 0.3, 0.2]), 0, 0.8)
    assert abs(time_term - 0.6355) < 5e-4
    assert abs(space_term + 0.6355) < 5e-4
    assert abs(time_term + space_term) < 1e-5


# %% Section 5.21: Fisher-Flow


def test_two_moves_under_two_metrics():
    near = torch.tensor([0.5, 0.5]), torch.tensor([0.4, 0.6])
    rare = torch.tensor([0.01, 0.99]), torch.tensor([0.001, 0.999])
    assert abs(fisher.euclidean_distance(*near) - 0.1414) < 5e-4
    assert abs(fisher.euclidean_distance(*rare) - 0.0127) < 5e-4
    assert abs(fisher.fisher_rao_distance(*near) - 0.201) < 1e-3
    assert abs(fisher.fisher_rao_distance(*rare) - 0.137) < 1e-3
    orthogonal = fisher.fisher_rao_distance(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0]))
    assert abs(orthogonal - math.pi) < 1e-6  # Any two letters are the farthest apart.


def test_sphere_and_geodesic():
    p0 = torch.full((4,), 0.25)
    p1 = fisher.smooth(torch.eye(4)[0])
    close(p1.tolist(), [0.91, 0.03, 0.03, 0.03])
    x0, x1 = fisher.to_sphere(p0), fisher.to_sphere(p1)
    close(x1.tolist(), [0.954, 0.173, 0.173, 0.173], 1e-3)
    assert abs(float((x0 * x1).sum()) - 0.737) < 1e-3
    assert abs(math.acos(float((x0 * x1).sum())) - 0.743) < 1e-3
    assert abs(fisher.fisher_rao_distance(p0, p1) - 1.485) < 1e-3
    mid = fisher.geodesic(x0, x1, torch.tensor(0.5))
    close(mid.tolist(), [0.780, 0.361, 0.361, 0.361], 1e-3)
    close(fisher.to_simplex(mid).tolist(), [0.609, 0.130, 0.130, 0.130], 1e-3)
    remaining = fisher.log_map(mid, x1)
    close(remaining.tolist(), [0.2323, -0.1672, -0.1672, -0.1672], 1e-3)
    assert abs(float(remaining.norm()) - 0.3713) < 5e-4
    assert abs(float(fisher.conditional_velocity(mid, x1, 0.5).norm()) - 0.743) < 1e-3
    stepped = fisher.exp_map(mid, 0.1 * fisher.conditional_velocity(mid, x1, 0.5))
    close(stepped.tolist(), [0.824, 0.327, 0.327, 0.327], 1e-3)
    close(fisher.to_simplex(stepped).tolist(), [0.680, 0.107, 0.107, 0.107], 1e-3)
    assert abs(float(stepped.norm()) - 1.0) < 1e-10


def test_tangent_projection():
    x0, p1 = fisher.to_sphere(torch.full((4,), 0.25)), fisher.smooth(torch.eye(4)[0])
    mid = fisher.geodesic(x0, fisher.to_sphere(p1), torch.tensor(0.5))
    raw = torch.tensor([0.3, -0.2, 0.5, 0.1])
    assert abs(float((mid * raw).sum()) - 0.378) < 1e-3
    tangent = fisher.project(mid, raw)
    close(tangent.tolist(), [0.005, -0.337, 0.363, -0.037], 1e-3)
    assert abs(float((mid * tangent).sum())) < 1e-12


# %% Section 5.22: Gumbel-Softmax Flow Matching


def test_gumbel_max_and_temperature():
    pi = torch.tensor([0.5, 0.2, 0.2, 0.1])
    torch.manual_seed(0)
    frequencies = gumbel.gumbel_max(pi, 200_000)
    close(frequencies.tolist(), pi.tolist(), tolerance=5e-3)
    close([gumbel.temperature(t) for t in (0.0, 0.5, 0.9, 1.0)], [10.0, 2.2313, 0.6721, 0.4979], 5e-4)


def test_annealed_path_and_replicator():
    logits = torch.tensor([0.0, -0.45, 1.40, -0.05])
    expected = {
        0.0: [0.2438, 0.2331, 0.2805, 0.2426],
        0.5: [0.2142, 0.1751, 0.4012, 0.2095],
        0.9: [0.0955, 0.0489, 0.7669, 0.0887],
        1.0: [0.0528, 0.0214, 0.8781, 0.0477],
    }
    for t, values in expected.items():
        close(gumbel.annealed_path(logits, t).tolist(), values)
    x = gumbel.annealed_path(logits, 0.5)
    assert abs(float((x * logits).sum()) - 0.4724) < 5e-4
    noisy = gumbel.replicator_velocity(x, logits, 0.5)
    close(noisy.tolist(), [-0.136, -0.217, 0.500, -0.147], 1e-3)
    assert abs(float(noisy.sum())) < 1e-12
    close(gumbel.inference_velocity(x, 2, 0.5).tolist(), [-0.116, -0.094, 0.323, -0.113], 1e-3)


def test_gumbel_marginal_velocity():
    late = torch.tensor([0.138, 0.084, 0.647, 0.131])
    p_hat = torch.tensor([0.05, 0.10, 0.70, 0.15])
    coefficients = (gumbel.LAMBDA / gumbel.temperature(0.8)) * late * p_hat
    close(coefficients.tolist(), [0.0228, 0.0278, 1.4976, 0.0650], 5e-4)
    v = gumbel.marginal_velocity(late, p_hat, 0.8)
    close(v.tolist(), [-0.200, -0.108, 0.454, -0.146], 1e-3)
    close((late + 0.01 * v).tolist(), [0.1360, 0.0829, 0.6515, 0.1295])


# %% Sections 5.24 to 5.29 and 5.33: guidance


def test_guided_base():
    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4])
    rates = p1 / (1.0 - 0.6)
    scores = torch.tensor([0.10, 0.35, 0.55, 0.25])
    close(rates.tolist(), [0.25, 0.50, 0.75, 1.00])
    assert abs(float(rates.sum()) - 2.50) < 1e-12
    close((scores / 0.20).tolist(), [0.50, 1.75, 2.75, 1.25])
    guided = guidance.tilt(rates, scores, 0.20)
    close(guided.tolist(), [0.125, 0.875, 2.0625, 1.250])
    assert abs(float(guided.sum()) - 4.3125) < 1e-12
    close((guided / guided.sum()).tolist(), [0.029, 0.203, 0.478, 0.290], 1e-3)
    close((p1 * scores).tolist(), [0.010, 0.070, 0.165, 0.100])
    h = guidance.h_value(p1, scores)
    assert abs(h - 0.345) < 1e-12
    assert abs(float(guidance.tilt(rates, scores, h).sum()) - 2.50) < 1e-12
    assert guidance.reachable(100, 21) == 2001


def test_guidance_strength_and_classifier_free():
    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4])
    rates, scores = p1 / 0.4, torch.tensor([0.10, 0.35, 0.55, 0.25])
    harder = guidance.tilt(rates, scores, 0.20, gamma=2.0)
    close(harder.tolist(), [0.0625, 1.5312, 5.6719, 1.5625], 1e-3)
    assert abs(float(harder.sum()) - 8.828) < 1e-3
    close((harder / harder.sum()).tolist(), [0.007, 0.173, 0.642, 0.177], 1e-3)
    mixed = guidance.classifier_free(guidance.tilt(rates, scores, 0.20), rates, gamma=2.0)
    assert torch.allclose(mixed, harder, atol=1e-12)


def test_taylor_expansion():
    gradient = torch.tensor([-2.10, -1.15, -0.75, -1.30, -1.61])
    approximate = guidance.taylor_log_ratio(gradient, guidance.MASK)[:4]
    close(approximate.tolist(), [-0.49, 0.46, 0.86, 0.31])
    exact = (torch.tensor([0.10, 0.35, 0.55, 0.25]) / 0.20).log()
    close(exact.tolist(), [-0.693, 0.560, 1.012, 0.223], 1e-3)
    rates = torch.tensor([0.25, 0.50, 0.75, 1.00]) * approximate.exp()
    close((rates / rates.sum()).tolist(), [0.038, 0.194, 0.434, 0.334], 1e-3)


def test_guided_step_never_remasks_a_visible_position():
    """The denoiser in the sampler is a distribution over the data tokens, the mask excluded."""
    torch.manual_seed(0)
    model = SequenceDenoiser(guidance.V, 8, width=32, depth=1)
    prop = guidance.PropertyModel(8, width=32)
    xt = torch.randint(0, 4, (16, 8))  # Every position already holds a base.
    stepped = guidance.guided_step(model, prop, xt, torch.full((16,), 0.9), 0.5, gamma=1.0)
    assert int((stepped == guidance.MASK).sum()) == 0  # A revealed token never returns to the mask.


# %% Sections 5.30 to 5.32: several objectives at once


def test_pareto_front():
    scores = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]])
    assert objectives.pareto_front(scores) == [0, 1]  # The third and fourth are dominated.
    close((scores @ torch.tensor([0.6, 0.4]))[:2].tolist(), [0.60, 0.64])
    close((scores @ torch.tensor([0.8, 0.2]))[:2].tolist(), [0.70, 0.62])  # The order reverses.


def test_mogdfm_step():
    changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]])
    delta_s, ranks, direction = objectives.combined_score(
        changes, torch.tensor([1.2, 0.8]), torch.tensor([0.6, 0.4]))
    close(ranks.tolist(), [0.500, 0.700, 0.850, 0.450])
    close(direction.tolist(), [0.0, 0.100, 0.090, -0.028])
    close(delta_s.tolist(), [-1.511, 1.541, 2.297, -2.327], 1e-3)
    close(delta_s.exp().tolist(), [0.221, 4.667, 9.948, 0.098], 1e-2)
    base = torch.tensor([0.3, 0.3, 0.1]) / 0.3  # The pretrained rates out of A toward C, G and T.
    tilted = base * delta_s[1:].exp()
    close(tilted.tolist(), [4.667, 9.948, 0.033], 1e-2)
    angles, cosines = objectives.cone_angles(changes[1:], torch.tensor([0.6, 0.4]))
    close(cosines.tolist(), [0.6727, 0.7894, -0.3032], 5e-4)
    close(angles.tolist(), [47.7, 37.9, 107.7], 0.05)
    survivors = tilted[angles <= 60.0]
    close((survivors / survivors.sum()).tolist(), [0.319, 0.681], 1e-3)
    assert abs(objectives.update_cone(60.0, 1 / 3) - 64.1) < 0.05


def test_pcomole_rollout():
    omega, reference = torch.tensor([0.5, 0.5]), torch.zeros(2)
    first = [torch.tensor(s) for s in ((0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5))]
    value, rewards = objectives.rollout_value(first, [1, 1, 0, 1], omega, reference, 2.0)
    close(rewards, [1.682, 2.054, 0.0, 1.822], 1e-3)
    assert abs(value - 1.390) < 1e-3
    second = [torch.tensor(s) for s in ((1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1))]
    other, _ = objectives.rollout_value(second, [1, 1, 1, 1], omega, reference, 2.0)
    assert abs(other - 1.181) < 1e-3
    products = torch.tensor([0.6 * value, 1.2 * other, 0.0])
    close(products.tolist(), [0.834, 1.417, 0.0], 1e-3)
    close((products / products.sum()).tolist(), [0.371, 0.629, 0.0], 1e-3)


# %% Sections 5.36 to 5.41: rectification and AReUReDi


def test_total_correlation_on_two_bases():
    source = torch.full((rectify.STATES,), 1.0 / rectify.STATES)
    q, support = rectify.complementary_pairs()
    coupling = q[None, :].repeat(rectify.STATES, 1)
    assert abs(rectify.total_correlation(coupling, source) - math.log(4)) < 1e-9
    assert abs(rectify.hamming(coupling, source) - 1.5) < 1e-9
    assert abs(rectify.one_step_success(coupling, source, support) - 0.25) < 1e-9


def test_redi_rounds():
    source = torch.full((rectify.STATES,), 1.0 / rectify.STATES)
    q, support = rectify.complementary_pairs()
    coupling = q[None, :].repeat(rectify.STATES, 1)
    coupling = rectify.recouple(coupling, source, steps=300)
    assert abs(rectify.total_correlation(coupling, source) - 1.101) < 2e-3
    assert abs(rectify.hamming(coupling, source) - 1.022) < 2e-3
    assert abs(rectify.one_step_success(coupling, source, support) - 0.403) < 2e-3
    columns = [a * 4 + b for a, b in support]
    close([float(coupling[0, c]) for c in columns], [0.400, 0.400, 0.100, 0.100], 2e-3)
    close([float(coupling[3, c]) for c in columns], [0.755, 0.082, 0.082, 0.082], 2e-3)
    coupling = rectify.recouple(coupling, source, steps=300)
    assert abs(rectify.total_correlation(coupling, source) - 0.843) < 2e-3
    assert abs(rectify.one_step_success(coupling, source, support) - 0.525) < 2e-3


def test_areuredi_proposal_and_acceptance():
    scores = torch.tensor([[0.60, 0.80], [0.70, 0.76], [0.50, 0.90], [0.64, 0.62]])
    weights = torch.tensor([0.5, 0.5])
    assert abs(rectify.anneal(11, 20) - 12.0) < 1e-12
    tchebycheff = rectify.tchebycheff(scores, weights)
    close(tchebycheff.tolist(), [0.300, 0.350, 0.250, 0.310])
    tilts = (12.0 * tchebycheff).exp()
    close(tilts.tolist(), [36.60, 66.69, 20.09, 41.26], 5e-2)
    model = torch.tensor([0.40, 0.25, 0.20, 0.15])
    q_current, z_current, ratios = rectify.proposal(model, tilts, 0)
    close(ratios.tolist(), [1.000, 1.822, 0.549, 1.127], 1e-3)
    close(rectify.barker(ratios).tolist(), [0.500, 0.646, 0.354, 0.530], 1e-3)
    assert abs(z_current - 0.5118) < 5e-4
    close(q_current.tolist(), [0.391, 0.315, 0.138, 0.155], 1e-3)
    _, z_proposed, _ = rectify.proposal(model, tilts, 1)
    assert abs(z_proposed - 0.3704) < 5e-4
    assert abs(rectify.acceptance(z_current, z_proposed) - 1.0) < 1e-12
    assert abs(rectify.acceptance(z_proposed, z_current) - 0.724) < 1e-3


def test_areuredi_concentration():
    base = torch.tensor([0.50, 0.25, 0.15, 0.10])
    objective = torch.tensor([0.20, 0.30, 0.45, 0.45])
    tilted = rectify.tilted_distribution(base, objective, 20.0)
    close(tilted.tolist(), [0.0127, 0.0468, 0.5643, 0.3762], 1e-3)
    assert abs(float(tilted[2:].sum()) - 0.9405) < 1e-3
    bound, gap = rectify.concentration_bound(base, objective, 20.0)
    assert abs(gap - 0.15) < 1e-9
    assert abs(bound * float(base[1]) - 0.0498) < 5e-4
    assert bound * float(base[1]) >= float(tilted[1])
    sharper = rectify.tilted_distribution(base, objective, 60.0)
    assert abs(float(sharper[2:].sum()) - 0.9999) < 1e-3
    assert abs(float(sharper[1]) - 1.234e-4) < 1e-6


# %% The command line of every file


@pytest.mark.parametrize("script", [
    "velocity", "kinetic_optimal", "rectify", "edit_flows", "multi_objective",
])
def test_script_runs(script):
    """The five files that need no training run end to end at the smallest grid."""
    result = subprocess.run(
        [sys.executable, f"lecture_5/{script}.py", "--steps", "20", "--quiet"],
        cwd=ROOT, capture_output=True, text=True, env={**os.environ, "PYTHONPATH": str(ROOT)},
    )
    assert result.returncode == 0, result.stderr[-2000:]
    assert "nan" not in result.stdout.lower()