"""Lecture 3: the numbers Chapter 3 works by hand, and the invariants the code must keep.

Every assertion here is either a quantity the chapter computes in the text or a property the
derivation guarantees. Nothing trains, so the whole file runs in a few seconds.
"""

import math

import pytest
import torch

from cis6270.metrics import mode_report
from lecture_3 import (
    ddpm, forward_diffusion, guidance, latent_diffusion, samplers, score_matching, sde,
)


# %% The forward process, Sections 3.2 through 3.10


def test_euler_maruyama_step_matches_the_worked_numbers():
    landed, drift, noise = sde.euler_maruyama_step(2.0, -1.0, 0.5, 0.04, 1.5)
    assert landed == pytest.approx(2.11)
    assert (drift, noise) == pytest.approx((-0.04, 0.15))


def test_three_draws_average_above_the_drift_prediction():
    endpoints = [sde.euler_maruyama_step(2.0, -1.0, 0.5, 0.04, e)[0] for e in (1.5, -0.5, 0.2)]
    assert endpoints == pytest.approx([2.11, 1.91, 1.98])
    assert sum(endpoints) / 3.0 == pytest.approx(2.00, abs=1e-6)  # The drift alone reaches 1.96.


def test_increment_scales_as_the_square_root_of_the_step():
    assert sde.increment_scale(0.04)["std"] == pytest.approx(0.2)
    assert sde.increment_scale(0.02)["std"] == pytest.approx(0.1414, abs=1e-4)
    assert sde.increment_scale(0.04, -1.0, 0.5)["noise_over_drift"] == pytest.approx(2.5)


def test_ornstein_uhlenbeck_law_at_log_two():
    law = sde.ornstein_uhlenbeck_law(2.0, 1.0, math.sqrt(2.0), math.log(2.0))
    assert law["mean"] == pytest.approx(1.0)
    assert law["variance"] == pytest.approx(0.75)
    assert law["std"] == pytest.approx(0.8660, abs=1e-4)
    assert law["stationary_variance"] == pytest.approx(1.0)


def test_ito_recovers_the_variance_and_the_chain_rule_alone_does_not():
    generator = torch.Generator().manual_seed(0)
    check = sde.check_ito_on_square(8192, 200, 0.25, generator)
    assert check["ito_mean"] == pytest.approx(0.25, abs=0.02)
    assert check["realized_mean_square"] == pytest.approx(0.25, abs=0.02)
    assert abs(check["no_curvature_mean"]) < 0.02  # The curvature term carries all the growth.


def test_fokker_planck_holds_on_the_spreading_gaussian():
    grid = torch.linspace(-3.0, 3.0, 41)
    for t in (0.0, 0.5, 2.0):
        residual = sde.fokker_planck_residual(grid, torch.full_like(grid, t))
        assert float(residual.abs().max()) < 1e-6


def test_peak_of_the_gaussian_falls_at_the_rate_the_equation_gives():
    peak = sde.peak_extrapolation(0.1)
    assert peak["p_0"] == pytest.approx(0.3989, abs=1e-4)
    assert peak["rate"] == pytest.approx(-0.1995, abs=1e-4)
    assert peak["linear"] == pytest.approx(0.3790, abs=1e-4)
    assert peak["exact"] == pytest.approx(0.3804, abs=1e-4)
    assert peak["gap"] == pytest.approx(0.0014, abs=1e-4)


# %% Score matching, Sections 3.14 through 3.17


def test_posterior_average_at_the_running_observation():
    report = score_matching.posterior_report(7.4, (10.0, 9.0), (0.3, 0.7), epsilon_theta=-0.8)
    assert report["weights"] == pytest.approx([0.2156, 0.7844], abs=1e-4)
    assert report["clean_mean"] == pytest.approx(9.2156, abs=1e-4)
    assert report["conditional_scores"] == pytest.approx([5 / 3, -5 / 9], abs=1e-4)
    assert report["marginal"] == pytest.approx([-0.0765, -0.0765], abs=1e-4)
    assert report["residual"] == pytest.approx(0.835, abs=1e-3)
    assert report["score_hat"] == pytest.approx(4 / 3, abs=1e-4)
    assert report["x0"] == pytest.approx([9.85, 9.85, 9.2156], abs=1e-4)


def test_one_parameter_fit_is_minimized_at_the_inverse_variance():
    assert score_matching.one_parameter_loss(0.5, 1.0) == pytest.approx(-0.375)
    assert score_matching.one_parameter_loss(1.0, 1.0) == pytest.approx(-0.5)
    assert score_matching.one_parameter_loss(0.25, 4.0) == pytest.approx(-0.125)
    grid = [0.05 * i for i in range(1, 41)]
    best = min(grid, key=lambda a: score_matching.one_parameter_loss(a, 4.0))
    assert best == pytest.approx(0.25)  # a = 1 / v with v = 4.


def test_the_three_objectives_share_one_minimizer():
    generator = torch.Generator().manual_seed(1)
    clean = score_matching.sample_clean(32768, generator)
    epsilon = torch.randn(clean.shape, generator=generator)
    found = score_matching.minimizers(clean, epsilon)
    assert found["explicit"] == pytest.approx(1.0)
    assert found["implicit"] == pytest.approx(1.0, abs=0.03)
    assert found["denoising"] == pytest.approx(1.0, abs=0.03)


def test_integration_by_parts_holds_on_the_exact_field():
    generator = torch.Generator().manual_seed(2)
    clean = score_matching.sample_clean(32768, generator)
    points = score_matching.A_T * clean + score_matching.SIGMA_T * torch.randn(
        clean.shape, generator=generator,
    )
    identity = score_matching.integration_by_parts_gap(score_matching.exact_score, points)
    assert identity["gap"] < 0.05 * abs(identity["left"])


def test_hutchinson_estimates_the_trace_the_exact_loop_computes():
    generator = torch.Generator().manual_seed(3)
    means = torch.randn(2, score_matching.TRACE_DIM, generator=generator) * 1.5
    points = torch.randn(8, score_matching.TRACE_DIM, generator=generator)

    def field(z):
        return score_matching.high_dimensional_score(z, means, 1.0)

    exact, passes = score_matching.exact_trace(field, points)
    assert passes == score_matching.TRACE_DIM  # One backward pass per coordinate.
    estimate, draws = score_matching.hutchinson_trace(field, points, 256, generator)
    error = float(draws.std(dim=0).mean()) / math.sqrt(256)
    assert float((estimate - exact).mean().abs()) < 4.0 * error


# %% The schedule, the corruption and DDPM, Sections 3.11, 3.12, 3.19 and 3.24 through 3.28


def test_schedule_matches_the_continuous_curve():
    assert float(forward_diffusion.alphas_cumprod[207]) == pytest.approx(0.6402, abs=1e-4)
    assert float(forward_diffusion.a[207]) == pytest.approx(0.8001, abs=1e-4)
    assert float(forward_diffusion.sigma[207]) == pytest.approx(0.5999, abs=1e-4)
    assert forward_diffusion.continuous_signal(0.207) == pytest.approx(0.7997, abs=1e-4)
    assert float(forward_diffusion.alphas_cumprod[1000]) == pytest.approx(4.036e-5, rel=1e-3)
    assert float(forward_diffusion.a[1000]) == pytest.approx(0.0064, abs=1e-4)
    assert forward_diffusion.continuous_signal(1.0) == pytest.approx(0.0066, abs=1e-4)
    assert float(forward_diffusion.a[500]) == pytest.approx(0.2803, abs=1e-4)


def test_every_copy_of_the_schedule_agrees():
    """No lecture file imports another, so each rebuilds the schedule. The copies must match."""
    for name in ("betas", "alphas", "alphas_cumprod"):
        for other in (ddpm, latent_diffusion):
            assert torch.equal(getattr(forward_diffusion, name), getattr(other, name)), name
    assert torch.equal(ddpm.posterior_variance[1:], latent_diffusion.posterior_variance[1:])
    # The samplers read the same path off the continuous curve, and the two agree at a fine grid.
    for grid in (samplers.make_schedule(1000), guidance.schedule_tensors(1000)):
        assert float(grid["alphas_cumprod"][207].sqrt()) == pytest.approx(0.8, abs=5e-3)
        assert torch.allclose(grid["alphas_cumprod"], ddpm.alphas_cumprod, atol=1e-3)


def test_noise_levels_of_the_continuous_schedule():
    for signal, time, rate in ((0.98, 0.0589, 1.2721), (0.8, 0.2068, 4.2157),
                               (0.35, 0.4544, 9.1420), (0.04, 0.7994, 16.0073)):
        located = forward_diffusion.level_time(signal)
        assert located["t"] == pytest.approx(time, abs=1e-4)
        assert located["beta_t"] == pytest.approx(rate, abs=1e-3)
        assert located["sigma_t"] == pytest.approx(math.sqrt(1.0 - signal ** 2), abs=1e-9)
    assert forward_diffusion.level_time(0.8)["sigma_t"] == pytest.approx(0.6, abs=1e-4)


def test_dropped_elbo_weight_at_step_five_hundred():
    weight = ddpm.elbo_weight(
        ddpm.betas[500], ddpm.posterior_variance[500],
        ddpm.alphas[500], ddpm.alphas_cumprod[500],
    )
    assert float(ddpm.posterior_variance[500]) == pytest.approx(0.01003, abs=1e-5)
    assert float(weight) == pytest.approx(0.005508, abs=1e-6)
    assert ddpm.elbo_weight(0.2, 1 / 9, 0.8, 0.64) == pytest.approx(0.625, abs=1e-6)


def test_one_pixel_through_the_forward_code():
    a_k = float(forward_diffusion.a[207])
    sigma_k = float(forward_diffusion.sigma[207])
    assert a_k * 0.5 + sigma_k * (-1.0) == pytest.approx(-0.1998, abs=1e-4)


def test_q_sample_uses_the_supplied_noise_exactly():
    x_0 = torch.randn(4, 1, 28, 28)
    noise = torch.randn_like(x_0)
    k = torch.tensor([1, 207, 500, 1000])
    for module in (forward_diffusion, ddpm):
        x_k, returned = module.q_sample(x_0, k, noise)
        abar = module.alphas_cumprod[k].view(-1, 1, 1, 1)
        assert torch.allclose(x_k, abar.sqrt() * x_0 + (1 - abar).sqrt() * noise, atol=1e-6)
        assert returned is noise


def test_corruption_law_matches_the_gaussian_it_is_drawn_from():
    generator = torch.Generator().manual_seed(0)
    x_0 = torch.randn(1, 1, 28, 28, generator=generator)
    levels = [1, 207, 500, 1000]
    law = forward_diffusion.measure_law(x_0, levels, 1024, generator)
    assert law["mean_gap"] < 0.05  # Centered on a_k x_0, to one over the square root of 1024.
    assert law["variance_gap"] < 0.05  # Spread sigma_k^2.
    # Four times the draws halves the gap, and a systematic error in a coefficient would not.
    finer = forward_diffusion.measure_law(x_0, levels, 4096, generator)
    assert finer["mean_gap"] < 0.6 * law["mean_gap"]


def test_worked_training_batch_at_abar_point_eight():
    batch = forward_diffusion.worked_batch(
        0.8, (2.0, -1.0, 0.5), (0.5, -0.2, 1.0), (0.4, -0.1, 0.8),
    )
    assert batch["a_k"] == pytest.approx(0.8944, abs=1e-4)
    assert batch["sigma_k"] == pytest.approx(0.4472, abs=1e-4)
    assert batch["x_k"] == pytest.approx([2.0125, -0.9839, 0.8944], abs=1e-3)
    assert batch["loss"] == pytest.approx(0.02, abs=1e-6)
    assert batch["score_targets"] == pytest.approx([-1.118, 0.447, -2.236], abs=1e-3)
    assert batch["x_0_hat"] == pytest.approx([2.05, -1.05, 0.60], abs=1e-2)


def test_two_step_reverse_transition_agrees_in_both_forms():
    step = ddpm.two_step_transition(0.2, 7.4, -1.0, 0.3)
    assert step["x_0_hat"] == pytest.approx(10.0, abs=1e-6)
    assert step["mean_from_noise"] == pytest.approx(8.646, abs=1e-3)
    assert step["mean_from_clean"] == pytest.approx(step["mean_from_noise"], abs=1e-6)
    assert step["weights"] == pytest.approx([0.4969, 0.4969], abs=1e-4)
    assert step["variance"] == pytest.approx(1 / 9)
    assert step["x_{k-1}"] == pytest.approx(8.746, abs=1e-3)
    assert step["elbo_weight"] == pytest.approx(0.625, abs=1e-6)


def test_unet_keeps_the_image_shape_and_the_last_reverse_step_adds_no_noise():
    torch.manual_seed(0)
    model = ddpm.UNet(in_channels=1, base_channels=8, time_embed_dim=32)
    x = torch.randn(2, 1, 28, 28)
    k = torch.tensor([1, 1])
    assert model(x, k).shape == x.shape
    first, second = ddpm.p_sample(model, x, k), ddpm.p_sample(model, x, k)
    assert torch.allclose(first, second)  # keep = 0 at k = 1, so the step is the mean.


# %% The three samplers, Sections 3.22, 3.26, 3.32, 3.34 and 3.35


def test_one_backward_step_under_both_dynamics():
    step = samplers.one_step_comparison(1.2, 2.0, 0.5, 0.8, 0.1)
    assert step["score"] == pytest.approx(-1.6)
    assert step["velocity"] == pytest.approx(0.4)
    assert step["ode_next"] == pytest.approx(1.16)
    assert step["sde_drift"] == pytest.approx(2.0)
    assert step["sde_mean"] == pytest.approx(1.0)
    assert step["sde_noise_std"] == pytest.approx(0.4472, abs=1e-4)


def test_langevin_step_from_the_worked_state():
    assert samplers.langevin_step(2.0, -2.0, 0.01, 0.5) == pytest.approx(2.0507, abs=1e-4)


def test_schedule_stays_a_valid_discretization_at_any_grid_size():
    for num_steps in (20, 200, 1000):
        schedule = samplers.make_schedule(num_steps)
        assert float(schedule["betas"][1:].max()) < 1.0
        assert float(schedule["alphas_cumprod"][0]) == pytest.approx(1.0)
        assert float(schedule["alphas_cumprod"][num_steps]) < 1e-4  # Terminal law is Gaussian.
    thousand = samplers.make_schedule(1000)
    assert float(thousand["alphas_cumprod"][207].sqrt()) == pytest.approx(0.8, abs=5e-3)


def test_every_sampler_finds_all_four_modes():
    budget = 60
    for name, points in (
        ("ancestral", samplers.ancestral_sample(
            samplers.ExactScore(), samplers.make_schedule(budget), 512,
            torch.Generator().manual_seed(0))),
        ("flow", samplers.probability_flow_sample(
            samplers.ExactScore(), 512, budget, torch.Generator().manual_seed(0))),
        ("langevin", samplers.annealed_langevin_sample(
            samplers.ExactScore(), 512, [0.7 - 0.69 * i / 19 for i in range(20)], 3, 0.15,
            torch.Generator().manual_seed(0))),
    ):
        measured = mode_report(points)
        assert measured["modes_found"] == 4, name
        assert measured["mode_entropy"] > 1.30, name  # log 4 = 1.386 when the modes are even.


def test_each_sampler_reports_the_evaluations_it_spent():
    score = samplers.ExactScore()
    samplers.probability_flow_sample(score, 64, 37, torch.Generator().manual_seed(0))
    assert score.calls == 37
    score = samplers.ExactScore()
    samplers.annealed_langevin_sample(score, 64, [0.5, 0.2, 0.05], 4, 0.1,
                                      torch.Generator().manual_seed(0))
    assert score.calls == 12  # Three levels, four iterations each.


# %% Guidance, Sections 3.37 through 3.45


def test_worked_guided_score():
    score, gradient = torch.tensor([[-1.0, -0.1]]), torch.tensor([[1.3, 1.1]])
    at_one = guidance.guided_score(score, gradient, 1.0)[0]
    at_two = guidance.guided_score(score, gradient, 2.0)[0]
    assert at_one.tolist() == pytest.approx([0.3, 1.0], abs=1e-6)
    assert float(at_one.norm()) == pytest.approx(1.0440, abs=1e-4)
    assert at_two.tolist() == pytest.approx([1.6, 2.1], abs=1e-6)
    assert float(at_two.norm()) == pytest.approx(2.6401, abs=1e-4)


def test_classifier_free_step_and_the_score_it_implies():
    eps = guidance.classifier_free_field(torch.tensor(0.50), torch.tensor(0.10), 2.0)
    assert float(eps) == pytest.approx(-0.30, abs=1e-6)
    assert float(-eps / 0.5) == pytest.approx(0.6, abs=1e-6)


def test_guided_backward_step_agrees_in_both_parameterizations():
    guided_eps = float(guidance.guided_noise(torch.tensor(0.8), torch.tensor(0.6), 2.0, 0.5))
    assert guided_eps == pytest.approx(0.2, abs=1e-6)
    velocity = -0.5 * 2.0 * (1.2 + (-guided_eps / 0.5))  # The collected form of the velocity.
    assert velocity == pytest.approx(-0.8, abs=1e-6)
    assert 1.2 - 0.1 * velocity == pytest.approx(1.28, abs=1e-6)


def test_classifier_gradient_equals_the_difference_of_the_two_fields():
    points = torch.randn(64, 2, generator=torch.Generator().manual_seed(0))
    gradient = guidance.classifier_gradient(points, 0.8, 0)
    difference = (guidance.conditional_score(points, 0.8, 0)
                  - guidance.unconditional_score(points, 0.8))
    assert torch.allclose(gradient, difference, atol=1e-6)


def test_weights_are_projected_onto_the_simplex():
    weights = guidance.normalize_weights((-1.0, 3.0, 1.0), torch.device("cpu"))
    assert float(weights.sum()) == pytest.approx(1.0)
    assert float(weights[0]) == 0.0  # A negative weight is clamped away.


def test_scalarized_rewards_reverse_the_ranking_with_the_weights():
    first, second = (1.2, 0.2), (0.5, 1.1)
    assert guidance.scalarize(first, (0.7, 0.3)) == pytest.approx(0.90, abs=1e-4)
    assert guidance.scalarize(second, (0.7, 0.3)) == pytest.approx(0.68, abs=1e-4)
    assert guidance.scalarize(first, (0.3, 0.7)) == pytest.approx(0.50, abs=1e-4)
    assert guidance.scalarize(second, (0.3, 0.7)) == pytest.approx(0.92, abs=1e-4)


def test_competing_gradients_cancel_at_the_stated_weights():
    combined = 0.6 * torch.tensor([1.0, 0.8]) + 0.4 * torch.tensor([0.5, -1.2])
    assert combined.tolist() == pytest.approx([0.80, 0.0], abs=1e-6)


def test_selection_weights_and_the_guidance_schedule():
    weights = (2.0 * torch.tensor([1.2, 0.8, 0.3])).softmax(0)
    assert weights.tolist() == pytest.approx([0.6193, 0.2783, 0.1024], abs=1e-4)
    assert guidance.schedule_value(0.6) == pytest.approx(0.96)
    assert guidance.schedule_value(0.0) == 0.0 and guidance.schedule_value(1.0) == 0.0


def test_guidance_moves_mass_onto_the_target_mode():
    schedule = guidance.schedule_tensors(60)
    fractions = []
    for w in (0.0, 0.25, 1.0):
        points = guidance.sample_diffusion(
            schedule, 512, guidance.TARGET_LABEL, w=w,
            generator=torch.Generator().manual_seed(0),
        )
        fractions.append(mode_report(points)["mode_fractions"][guidance.TARGET_LABEL])
    assert fractions[0] < 0.35  # Unconditional, so about a quarter.
    assert fractions[0] < fractions[1] < fractions[2]
    assert fractions[2] > 0.9


def test_reward_guidance_prefers_the_right_hand_modes():
    schedule = guidance.schedule_tensors(60)
    points = guidance.sample_diffusion(
        schedule, 512, guidance.TARGET_LABEL, eta=2.0, lambdas=(1.0, 0.0),
        generator=torch.Generator().manual_seed(0),
    )
    fractions = mode_report(points)["mode_fractions"]
    assert fractions[0] + fractions[1] > 0.9  # The two modes at x = +2.


# %% Latent diffusion, Sections 3.21, 3.29 and 3.46


def test_latent_coordinate_corrupted_and_carried_back():
    abar = float(latent_diffusion.alphas_cumprod[500])
    a_500, sigma_500 = math.sqrt(abar), math.sqrt(1.0 - abar)
    assert (a_500, sigma_500) == pytest.approx((0.2803, 0.9599), abs=1e-4)
    z_500 = a_500 * 1.2 + sigma_500 * 0.4
    assert z_500 == pytest.approx(0.7204, abs=1e-4)
    eps_hat = (z_500 - a_500 * 1.05) / sigma_500
    assert eps_hat == pytest.approx(0.4438, abs=1e-4)
    beta_k, alpha_k = float(latent_diffusion.betas[500]), float(latent_diffusion.alphas[500])
    mean = (z_500 - beta_k / sigma_500 * eps_hat) / math.sqrt(alpha_k)
    assert mean == pytest.approx(0.7193, abs=1e-4)
    variance = float(latent_diffusion.posterior_variance[500])
    assert math.sqrt(variance) == pytest.approx(0.1002, abs=1e-4)


def test_masked_corruption_leaves_padding_at_zero():
    z0 = torch.randn(3, 8, 16)
    mask = torch.ones(3, 8)
    mask[:, 5:] = 0.0  # The last three positions are padding.
    z_k, noise = latent_diffusion.q_sample_latent(z0, torch.tensor([1, 500, 1000]), mask)
    assert torch.count_nonzero(z_k[:, 5:, :]) == 0
    assert noise.shape == z0.shape


def test_latent_loss_counts_valid_residues_only():
    z0 = torch.zeros(2, 4, 3)
    z0_hat = torch.ones(2, 4, 3)  # Unit error at every coordinate.
    mask = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]])
    assert float(latent_diffusion.latent_loss(z0_hat, z0, mask)) == pytest.approx(1.0)
    z0_hat[:, 2:, :] = 100.0  # A large error, but only on padding.
    assert float(latent_diffusion.latent_loss(z0_hat, z0, mask)) == pytest.approx(1.0)


def test_denoiser_returns_a_latent_of_the_same_shape():
    torch.manual_seed(0)
    model = latent_diffusion.LatentDenoiser(embed_dim=16, num_layers=1, num_heads=2, max_len=8)
    z_k = torch.randn(2, 8, 16)
    mask = torch.ones(2, 8)
    assert model(z_k, torch.tensor([7, 900]), mask).shape == z_k.shape


def test_composition_proxies_on_a_known_peptide():
    charge, polar = latent_diffusion.composition_proxies("KRDE")
    assert charge == pytest.approx(0.0)  # Two positive and two negative over four residues.
    assert polar == pytest.approx(1.0)  # All four are in the polar set.
    assert latent_diffusion.composition_proxies("AAAA") == pytest.approx((0.0, 0.0))


def test_enforce_polar_reaches_the_minimum_at_the_cheapest_positions():
    torch.manual_seed(0)
    logits = torch.randn(2, 12, 20)
    logits[:, :, latent_diffusion.POLAR_COLUMNS] -= 5.0  # Make the argmax non-polar everywhere.
    choices = logits.argmax(dim=-1)
    sequences = latent_diffusion.enforce_polar(logits, choices, 5)
    for sequence in sequences:
        assert len(sequence) == 12
        assert sum(a in latent_diffusion.POLAR_RESIDUES for a in sequence) >= 5
    unchanged = latent_diffusion.enforce_polar(logits, choices, 0)  # Nothing is missing.
    assert unchanged == ["".join(latent_diffusion.AMINO_ACIDS[int(i)] for i in row)
                         for row in choices]


def test_the_cheapest_two_substitutions_cost_what_the_notes_says():
    costs = torch.tensor([0.9, 1.4, 2.6, 3.1, 3.3, 4.0, 4.2, 5.5, 6.1])
    cheapest, _ = costs.topk(2, largest=False)
    assert cheapest.tolist() == pytest.approx([0.9, 1.4])
    assert float(cheapest.sum()) == pytest.approx(2.3, abs=1e-5)