"""Discrete Diffusion Schrodinger Bridge Matching: one categorical feature of a molecular graph.

A molecule is a graph with a categorical feature on every atom and on every pair of atoms, and the
reference resamples each feature independently, so one feature with S categories specifies the
whole construction. We scale the reference kernel to the two laws to get the bridge coupling for
that feature, and we pin a path at a known terminal graph to get the rate a network regresses onto.
Both stay usable on a state space of 10^27 graphs because the reference is a product over the
features, so every sum over graphs becomes a sum over the categories of one of them.

Chapter 7, Section 7.24 of the notes. Runs on one bond with three types, no bond, single and
double, at the survival probability r = 0.4. Prints the kernel and its generator, the bridge
coupling and the mass it leaves unchanged, the pinned rates and the posterior average a network
learns, and the counting over the forty-five features of a nine-atom graph.

    python lecture_7/ddsbm.py              # 200 Sinkhorn sweeps on the bond coupling
    python lecture_7/ddsbm.py --steps 1    # the first sweep alone, computed by hand
"""

# %% 01. Imports and configuration
from __future__ import annotations

import math

import torch

from cis6270.runner import common_parser, print_report, save_run, set_seed

DTYPE = torch.float64  # The chapter quotes four figures of the coupling and of the pinned rates.

PRIOR = torch.full((3,), 1.0 / 3.0, dtype=DTYPE)  # m, uniform over the three bond types.
BOND_0 = torch.tensor([0.7, 0.2, 0.1], dtype=DTYPE)  # pi_0, mostly no bond.
BOND_1 = torch.tensor([0.1, 0.2, 0.7], dtype=DTYPE)  # pi_1, mostly double.
SURVIVAL = 0.4  # r = alpha_bar(1) / alpha_bar(0) over the whole unit interval.
ATOMS = 9  # Heavy atoms in the molecule the chapter counts.
ATOM_TYPES, BOND_TYPES = 5, 4  # Categories of a node feature and of a bond feature.


# %% 02. The reference process on one feature
#
# The kernel keeps the category with probability r and redraws it from the prior m otherwise. Over
# a short step r = 1 - h beta(t), so the rate out of x toward a different category y is the
# resampling rate beta(t) times the chance m(y) that the redraw lands on y. Section 7.24.1.


def replacement_kernel(r: float, prior: torch.Tensor) -> torch.Tensor:
    """K_st(x, y) = r 1_{x=y} + (1 - r) m(y), the uniform-replacement kernel with a general prior."""
    return r * torch.eye(len(prior), dtype=DTYPE) + (1.0 - r) * prior[None, :]  # [S, S].


def replacement_generator(beta: float, prior: torch.Tensor) -> torch.Tensor:
    """Q_t(x, y) = beta(t) m(y) off the diagonal, with the diagonal -beta(t) (1 - m(x))."""
    rates = beta * prior[None, :].repeat(len(prior), 1)  # [S, S].
    rates = rates - torch.diag(rates.diagonal())  # Clear the diagonal before closing the row.
    return rates - torch.diag(rates.sum(dim=1))


def resampling_rate(r: float) -> float:
    """beta = -log alpha_bar(1) for the schedule alpha_bar(t) = r^t, constant in time."""
    return -math.log(r)


# %% 03. A bridge for a single bond
#
# We scale diag(pi_0) K_01 to the two marginals and get the coupling. The mass it leaves on the
# diagonal sits between the mass an independent coupling happens to leave and the largest any
# coupling can leave, and r fixes where between the two it lands. Section 7.24.2.


def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int
             ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[float]]:
    """Scale a positive matrix to the two marginals, returning f, g, the plan and the column errors."""
    f = torch.ones_like(start)
    g = torch.ones_like(end)
    errors: list[float] = []
    for _ in range(sweeps):
        g = end / (matrix.T @ f)  # Column sums become exactly pi_1.
        f = start / (matrix @ g)  # Row sums become exactly pi_0, and the columns move again.
        errors.append(float(((f[:, None] * matrix * g[None, :]).sum(dim=0) - end).abs().sum()))
    return f, g, f[:, None] * matrix * g[None, :], errors


def held_mass(plan: torch.Tensor) -> float:
    """The trace of the plan, the fraction of bonds whose type the bridge does not change."""
    return float(plan.diagonal().sum())


# %% 04. Training with the terminal graph fixed
#
# We pin the endpoint with an h-transform, the reference probability of arriving at x_1 in the
# place of the potential, so a move onto the terminal type has its rate multiplied and a move off
# it divided. Several terminal types can occur at one state, and the minimizer of the rate error
# over a single number is their posterior average. Section 7.24.3.


def pinned_rate(generator: torch.Tensor, kernel: torch.Tensor, x: int, y: int, x1: int) -> float:
    """Q_t(x, y; x_1) = Q_t(x, y) K_t1(y, x_1) / K_t1(x, x_1), the rate of the bridge ending at x_1."""
    return float(generator[x, y] * kernel[y, x1] / kernel[x, x1])


def half_survival_time(r: float) -> float:
    """The time where alpha_bar(1) / alpha_bar(t) = r^{1-t} equals 1/2, on alpha_bar(t) = r^t."""
    return 1.0 - math.log(0.5) / math.log(r)  # At r = 0.4 this is log 0.8 / log 0.4 = 0.2435.


def rate_error(controlled: float, network: float) -> float:
    """Q^pin log(Q^pin / Q^theta) - Q^pin + Q^theta, the Bregman divergence of u log u on one edge."""
    return controlled * math.log(controlled / network) - controlled + network


def posterior_loss(rates: list[float], weights: list[float], network: float) -> float:
    """The expected rate error at a state whose terminal type is uncertain, over the posterior."""
    return sum(w * rate_error(q, network) for q, w in zip(rates, weights))


# %% 05. The factorization over a graph
#
# The reference applies the same kernel independently at every node and every pair of nodes, so the
# generator changes one feature at a time. The total jump rate out of a graph is then the sum of
# the per-feature rates, and in the pinned rate every factor cancels but the one of the feature
# that moves. Section 7.24.4.


def feature_count(atoms: int) -> tuple[int, int, int]:
    """The node features, the bond features and their total, for a molecule with this many atoms."""
    bonds = atoms * (atoms - 1) // 2  # One feature per unordered pair.
    return atoms, bonds, atoms + bonds


def state_space_size(atoms: int, atom_types: int, bond_types: int) -> float:
    """atom_types^nodes times bond_types^bonds, the number of graphs the sums would run over."""
    nodes, bonds, _ = feature_count(atoms)
    return float(atom_types ** nodes) * float(bond_types ** bonds)


def total_jump_rate(features: int, beta: float, categories: int) -> float:
    """The sum of the per-feature rates out of a graph, and it fixes the sampler's step size."""
    return features * beta * (1.0 - 1.0 / categories)  # A redraw lands elsewhere with 1 - m(x).


# %% 06. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=200)  # --steps is the number of Sinkhorn sweeps on the bond coupling.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet
    sweeps = max(args.steps, 1)

    kernel = replacement_kernel(SURVIVAL, PRIOR)
    beta = resampling_rate(SURVIVAL)
    rates = replacement_generator(beta, PRIOR)
    reference = torch.diag(BOND_0) @ kernel
    if loud:
        print(f"One bond with three types at r = {SURVIVAL}, beta = {beta:.4f}, "
              f"off-diagonal rate {float(rates[0, 1]):.4f}")
        print(f"  K_01 = {[[round(float(z), 1) for z in row] for row in kernel]}")
        print(f"  diag(pi_0) K_01 has column sums "
              f"{[round(float(z), 2) for z in reference.sum(dim=0)]}, so the first g is "
              f"{[round(float(z), 4) for z in BOND_1 / reference.sum(dim=0)]}")

    f, g, plan, errors = sinkhorn(reference, BOND_0, BOND_1, sweeps)
    independent = float((BOND_0 * BOND_1).sum())
    largest = float(torch.minimum(BOND_0, BOND_1).sum())
    held = {r: held_mass(sinkhorn(torch.diag(BOND_0) @ replacement_kernel(r, PRIOR),
                                  BOND_0, BOND_1, 200)[2])
            for r in (0.2, 0.4, 0.6, 0.8, 0.99)}
    if loud:
        print(f"\nAfter {sweeps} sweeps, column error {errors[-1]:.2e}")
        print(f"  f = {[round(float(z), 4) for z in f]}, "
              f"g = {[round(float(z), 4) for z in g]}")
        print(f"  pi* = {[[round(float(z), 3) for z in row] for row in plan]}")
        print(f"  mass that stays: bridge {held_mass(plan):.4f}, independent {independent:.2f}, "
              f"largest possible {largest:.1f}")
        print(f"  diagonal mass at r = 0.2, 0.4, 0.6, 0.8, 0.99: "
              f"{[round(held[r], 4) for r in (0.2, 0.4, 0.6, 0.8, 0.99)]}")

    pinned_time = half_survival_time(SURVIVAL)
    pinned_kernel = replacement_kernel(0.5, PRIOR)
    toward = pinned_rate(rates, pinned_kernel, 0, 2, 2)  # A move onto the terminal type.
    away = pinned_rate(rates, pinned_kernel, 2, 0, 2)  # A move off it.
    between = pinned_rate(rates, pinned_kernel, 0, 1, 2)  # Between two non-terminal types.
    one_edge = rate_error(toward, 1.0)  # The term one edge contributes at a network rate of 1.
    weights = [0.6, 0.4]  # The posterior over the two terminal types at this state.
    posterior = weights[0] * toward + weights[1] * between
    losses = {q: posterior_loss([toward, between], weights, q) for q in (posterior, 1.0, 0.5)}
    if loud:
        print(f"\nAt t = {pinned_time:.4f} the survival ratio to time 1 is 1/2, so K_t1 is "
              f"{float(pinned_kernel[0, 0]):.4f} on the diagonal and "
              f"{float(pinned_kernel[0, 1]):.4f} off it")
        print(f"  pinned rates: {toward:.4f} toward the terminal type, {away:.4f} away from it "
              f"and {between:.4f} between two others, a factor of {toward / away:.0f}")
        print(f"  one edge at a network rate of 1 contributes {one_edge:.4f}")
        print(f"  with two terminal types the network learns {posterior:.4f}, at a loss of "
              f"{losses[posterior]:.4f}, against {losses[1.0]:.4f} at 1 and {losses[0.5]:.4f} at 0.5")

    nodes, bonds, features = feature_count(ATOMS)
    graphs = state_space_size(ATOMS, ATOM_TYPES, BOND_TYPES)
    jump_rate = total_jump_rate(bonds, beta, BOND_TYPES)
    if loud:
        print(f"\n{ATOMS} heavy atoms: {nodes} node features and {bonds} bond features, "
              f"{features} in all")
        print(f"  {ATOM_TYPES}^{nodes} x {BOND_TYPES}^{bonds} = {graphs:.1e} graphs, "
              f"which no sum reaches")
        print(f"  the total jump rate out of a graph is {bonds} x {beta:.4f} x "
              f"(1 - 1/{BOND_TYPES}) = {jump_rate:.2f} per unit time")

    report = {
        "beta of alpha_bar(t) = 0.4^t": beta,
        "off-diagonal reference rate": float(rates[0, 1]),
        "K_01 on the diagonal": float(kernel[0, 0]),
        "K_01 off the diagonal": float(kernel[0, 1]),
        "rows of the generator sum to": float(rates.sum(dim=1).abs().max()),
        "column sums of diag(pi_0) K_01": [round(float(z), 2) for z in reference.sum(dim=0)],
        "the first g": [round(float(z), 4) for z in BOND_1 / reference.sum(dim=0)],
        "sweeps run": sweeps,
        "column error after the last sweep": errors[-1],
        "f": [round(float(z), 4) for z in f],
        "g": [round(float(z), 4) for z in g],
        "bond plan, no bond to double": float(plan[0, 2]),
        "bond mass that stays": held_mass(plan),
        "bond mass, independent coupling": independent,
        "bond mass, largest possible": largest,
        "diagonal mass at r = 0.2": held[0.2],
        "diagonal mass at r = 0.99": held[0.99],
        "time where the survival ratio is 1/2": pinned_time,
        "pinned rate toward the terminal": toward,
        "pinned rate away from it": away,
        "pinned rate between two others": between,
        "one edge at a network rate of 1": one_edge,
        "posterior average of the rates": posterior,
        "loss at the posterior average": losses[posterior],
        "loss at 1": losses[1.0],
        "loss at 0.5": losses[0.5],
        "node, bond and total features": [nodes, bonds, features],
        "graphs on nine heavy atoms": f"{graphs:.1e}",
        "total jump rate out of a graph": jump_rate,
    }
    print_report("Discrete Diffusion Schrodinger Bridge Matching, Chapter 7 Section 7.24", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()