"""A jump rate written as a velocity, and the master equation it satisfies.

Nothing here is learned. A velocity on a finite set is a rate matrix read from the destination's
side, and two rules fix which tables of numbers qualify. The master equation is the condition under
which such a table generates a probability path. The schedule factor of the mixture path turns out
to be a hazard rate whose survival curve is the path itself. The discrete velocity has the same
shape as the continuous one of Chapter 2, with a posterior distribution in place of a posterior
mean.

Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).

Chapter 5, Sections 5.2, 5.3, 5.6.3 and 5.8.1. The long integration starts from the letter frequency
of the synthetic DNA of ``cis6270.data``. Prints the chapter's worked velocity, its Euler step, the
hazard at two times, and the two velocities side by side.

    python lecture_5/velocity.py                     # 600 Euler steps in the long integration
    python lecture_5/velocity.py --steps 20 --quiet  # a coarser grid, visibly less accurate
"""

# %% 01. Imports and configuration
import math

import torch
import torch.nn.functional as F

from cis6270.data import DNA, DNA_SIZE, dna_splits
from cis6270.runner import common_parser, print_report, save_run, set_seed

torch.set_default_dtype(torch.double)  # Full precision, because the report prints six decimal places.



# %% 02. A jump rate written as a velocity
#
# Chapter 4 wrote R_t(x, y) for the rate of jumping from x to y. Discrete flow matching writes the
# same number with the destination first, u_t(y, x) = R_t(x, y). If we fix x and let y vary, we get
# one vector attached to the state x, with one component per destination.


def velocity_from_rates(rates):
    """Transpose a Chapter 4 rate matrix R_t(x, y) into a Chapter 5 velocity u_t(y, x)."""
    return rates.transpose(0, 1)  # [K, K] with the destination now indexing the row.


def jump_to_one_state(target, size=DNA_SIZE, rate=1.0):
    """Every state other than ``target`` leaves toward it at ``rate``, and ``target`` never leaves."""
    u = torch.zeros(size, size)  # [K, K], row = destination, column = current state.
    for x in range(size):
        if x == target:
            continue
        u[target, x] = rate  # The one positive off-diagonal entry in column x.
        u[x, x] = -rate  # The diagonal is minus the total rate of leaving x.
    return u


# %% 03. The two rules a velocity obeys
#
# Off the diagonal the entries are nonnegative, since probability cannot flow out of a state at a
# negative rate. Every column sums to zero, since probability leaving one state has to arrive at
# another. Once we choose the off-diagonal entries, the two rules fix the diagonal.


def obeys_the_two_rules(u, tolerance=1e-9):
    """True when every off-diagonal entry is nonnegative and every column sums to zero."""
    off = u - torch.diag(torch.diagonal(u))  # The matrix with the diagonal removed.
    return bool((off >= -tolerance).all() and (u.sum(dim=0).abs() < tolerance).all())


def set_diagonal(u):
    """u_t(x, x) = -sum_{y != x} u_t(y, x), the total rate of leaving x."""
    u = u.clone()
    u.fill_diagonal_(0.0)  # Keep the off-diagonal entries and discard whatever the diagonal held.
    return u - torch.diag(u.sum(dim=0))  # Each column now sums to zero.


# %% 04. The master equation in velocity form
#
# The rate of change of p_t(y) is the velocity into y times the probability of each current state,
# summed over states. As a matrix equation this is p_dot = u p, where p is a column vector. This
# chapter writes the velocity in the column convention, so the product needs no transpose.


def master_equation(u, p):
    """p_dot_t(y) = sum_x u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K]."""
    return u @ p  # The matrix multiplies the probability vector from the left.


def arrivals_and_departures(u, p, y):
    """The flux arriving at y from every other state, and the flux leaving y, as two numbers."""
    others = [x for x in range(len(p)) if x != y]  # Every state except y itself.
    arriving = sum(float(u[y, x] * p[x]) for x in others)  # u_t(y, x) p_t(x) along each edge in.
    leaving = sum(float(u[x, y] * p[y]) for x in others)  # u_t(x, y) p_t(y) along each edge out.
    return arriving, leaving  # Their difference is p_dot_t(y).


def integrate(u, p, steps):
    """Carry p forward by ``steps`` Euler steps over one unit of time under a constant velocity."""
    h = 1.0 / steps  # The time increment of one step.
    for _ in range(steps):
        p = p + h * master_equation(u, p)  # p_{t+h} = p_t + h p_dot_t.
    return p  # [K] after one unit of time.


# %% 05. One velocity on DNA, with numbers
#
# Let the token be a single DNA base, let every base other than G jump to G at rate 1, and let G
# never leave. Because the two rules hold, probability moves into G and the total stays at one.


def worked_master_equation():
    """Section 5.3.2 with p_t = (0.1, 0.2, 0.3, 0.4) and the all-to-G velocity."""
    u = jump_to_one_state(DNA.index("G"))  # [4, 4] the matrix printed in the chapter.
    p = torch.tensor([0.1, 0.2, 0.3, 0.4])  # [4] over A, C, G, T.
    p_dot = master_equation(u, p)  # (-0.1, -0.2, 0.7, -0.4).
    arriving, leaving = arrivals_and_departures(u, p, DNA.index("G"))  # 0.7 in and 0 out.
    return u, p, p_dot, arriving, leaving


# %% 06. Survival and hazard
#
# The schedule factor kappa_dot_t/(1 - kappa_t) is a hazard rate, and the probability that a position
# still holds its source token is its survival curve. If we separate variables we get
# S(t) = 1 - kappa_t, the same weight the mixture path puts on the source token.


def hazard(t):
    """kappa_dot_t/(1 - kappa_t) on the linear schedule, or 1/(1 - t)."""
    return 1.0 / (1.0 - t)  # 1 at t = 0, 2 at t = 0.5, 10 at t = 0.9.


def survival(steps, upper=0.9):
    """Integrate dS/dt = -hazard(t) S(t) from S(0) = 1 and return S at ``upper``."""
    h, s = upper / steps, 1.0  # One increment and the initial survival probability.
    for k in range(steps):
        s = s - h * hazard(k * h) * s  # Euler on the survival equation.
    return s  # Exactly 1 - kappa_t = 1 - upper, since the two factors telescope at every grid.


# %% 07. The continuous and discrete velocities side by side
#
# In continuous space the velocity is the schedule factor times the posterior mean minus the current
# point. On tokens it is the same factor times the posterior distribution minus the point mass on the
# current token. One step of size h displaces the point by h u_t(x). On tokens the same step
# redraws from the denoiser with probability lambda_t, and the two agree in expectation on the
# one-hot embedding.


def continuous_velocity(x, x1_hat, t):
    """u_t(x) = kappa_dot/(1 - kappa) [x_hat_{1|t}(x) - x] at a point of R^d. x: [d] -> [d]."""
    return hazard(t) * (x1_hat - x)  # The displacement that remains, divided by the time remaining.


def discrete_velocity(x_i, p1, t, size=DNA_SIZE):
    """u^i_t(., x) = kappa_dot/(1 - kappa) [p_{1|t}(. | x) - delta_{x^i}]. Returns [K]."""
    return hazard(t) * (p1 - F.one_hot(torch.tensor(x_i), size).double())  # Signed, sums to zero.


def side_by_side(x_i, p1, t, h):
    """The continuous Euler step at the one-hot point and the discrete redraw distribution."""
    x = F.one_hot(torch.tensor(x_i), len(p1)).double()  # The current token as a corner of R^K.
    moved = x + h * continuous_velocity(x, p1, t)  # Displace by h times the velocity.
    lam = h * hazard(t)  # lambda_t, the redraw probability of the two-part rule.
    drawn = (1.0 - lam) * x + lam * p1  # Keep the current token, or redraw from the denoiser.
    return moved, drawn, lam  # The two coincide, the last row of Table 5.1.


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    args = parser.parse_args()
    set_seed(args.seed)

    u, p, p_dot, arriving, leaving = worked_master_equation()
    report = {
        "two_rules_hold": obeys_the_two_rules(u),  # Nonnegative off the diagonal, columns sum to zero.
        "diagonal_recovered": bool(torch.allclose(u, set_diagonal(u))),  # The rules fix the diagonal.
        "p_dot": [round(float(v), 4) for v in p_dot],  # (-0.1, -0.2, 0.7, -0.4).
        "p_dot_sums_to": round(float(p_dot.sum()), 6),  # Zero, so probability is conserved.
        "flux_into_G": round(arriving, 4),  # 0.7, the three incoming rates against their probabilities.
        "flux_out_of_G": round(leaving, 4),  # 0, since G has no outgoing rate.
        "euler_step_h_0.1": [round(float(v), 4) for v in p + 0.1 * p_dot],  # (0.09, 0.18, 0.37, 0.36).
        "hazard_at_0.5": round(hazard(0.5), 4),  # 2.
        "hazard_at_0.9": round(hazard(0.9), 4),  # 10.
        "survival_at_0.9": f"{survival(args.steps):.6f}",  # 1 - 0.9 = 0.1 at any grid.
        "survival_exact": 0.1,  # S(t) = 1 - kappa_t.
    }

    # Section 5.8.1: one position holding G with the denoiser of the chapter.
    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
    moved, drawn, lam = side_by_side(DNA.index("G"), p1, t=0.5, h=0.1)
    report["discrete_velocity_at_G"] = [round(float(v), 4) for v in discrete_velocity(2, p1, 0.5)]
    report["continuous_step"] = [round(float(v), 4) for v in moved]  # (0.02, 0.12, 0.84, 0.02).
    report["discrete_redraw"] = [round(float(v), 4) for v in drawn]  # The identical vector.
    report["redraw_probability"] = round(lam, 4)  # lambda_t = 0.2.

    # The same velocity run for a whole unit of time, starting from the data's letter frequency.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
    start = counts / counts.sum()  # [4] the empirical frequency of A, C, G, T in the training DNA.
    final = integrate(u, start, args.steps)  # Euler on the master equation, args.steps increments.
    exact = start.clone()
    exact[:] = start * math.exp(-1.0)  # Every base other than G decays as exp(-t).
    exact[DNA.index("G")] = 1.0 - (start.sum() - start[DNA.index("G")]) * math.exp(-1.0)
    report["dna_start"] = [round(float(v), 4) for v in start]  # The GC-rich background of the data.
    report["dna_after_one_unit"] = [round(float(v), 4) for v in final]  # Euler at args.steps.
    report["dna_exact"] = [round(float(v), 4) for v in exact]  # The closed-form solution.
    report["euler_error"] = round(float((final - exact).abs().max()), 6)  # Shrinks with args.steps.
    print_report("a jump rate written as a velocity", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()