"""Correctors: extra motion that leaves the probability path where it is.

The master equation is linear in the velocity, so we can add a second velocity with no net flux into
any state without moving the marginals. Detailed balance is the easy way to arrange that, and the
mixture path already contains such a term, the difference between the forward velocity and the same
path run backward. If we add it with weight eta, a position unmasks more often and remasks in
between, so the sampler gets more chances to replace a token the rest of the sequence has ruled out.

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.13 and 5.14. The long simulation uses the letter frequency of the synthetic DNA
of ``cis6270.data`` as its denoiser. Prints the zero-flux check, the per-step unmasking table, the
jump counts at three corrector strengths, the parallel-step error at 64 steps, and the chance a
jointly revealed pair survives to t = 1 untouched.

    python lecture_5/corrector.py                     # 600 Euler steps on 40,000 positions
    python lecture_5/corrector.py --eta 10            # one corrector strength in the step table
    python lecture_5/corrector.py --steps 20 --quiet  # a coarse grid, visibly biased jump counts
"""

# %% 01. Imports and configuration
import math

import torch

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

EXACT = torch.double  # The worked arithmetic keeps full precision whatever the default dtype.
MASK = DNA_SIZE  # Index 4 is the mask symbol m, outside the four-letter alphabet.


# %% 02. Velocities that move samples without moving the path
#
# We can add a velocity v_t whose net flux into every state is zero to a velocity that generates
# p_t, and the sum generates the same path. Detailed balance, equal flux in both directions along
# every edge, is sufficient. If we reverse every term, the net flux becomes p_t(y) times a column of
# rates, and those rates sum to zero.


def net_flux(v, p):
    """sum_x v_t(y, x) p_t(x) at every state y. v: [K, K], p: [K] -> [K]."""
    return v @ p  # Zero at every entry when v moves no probability.


def detailed_balance_velocity(p, chosen):
    """Build v_t from chosen forward rates, with each reverse rate fixed by detailed balance.

    ``chosen`` maps an edge (x, y) to the rate v_t(y, x). The reverse rate is then
    v_t(x, y) = v_t(y, x) p_t(x) / p_t(y). This equalizes the flux along that edge.
    """
    v = torch.zeros(len(p), len(p), dtype=p.dtype)  # [K, K], destination indexing the row.
    for (x, y), rate in chosen.items():
        v[y, x] = rate  # The rate we pick, out of x toward y.
        v[x, y] = rate * p[x] / p[y]  # The reverse rate detailed balance forces.
    return v - torch.diag(v.sum(dim=0))  # Each column now sums to zero.


# %% 03. A zero-flux velocity on three states
#
# On the states A, C, G with p_t = (0.2, 0.3, 0.5) we can pick the rate 3 from A to C and the rate 1
# from A to G, and detailed balance returns 2 and 0.4. If we add this to a velocity that sends every
# state to G at rate 1, the total rate of leaving A goes from 1 to 5 while p_dot stays put.


def three_state_example():
    """Section 5.13.2, returning the two velocities, the path and the two rates of change."""
    p = torch.tensor([0.2, 0.3, 0.5], dtype=EXACT)  # [3] over A, C, G.
    v = detailed_balance_velocity(p, {(0, 1): 3.0, (0, 2): 1.0})  # A to C at 3, A to G at 1.
    u = torch.zeros(3, 3, dtype=EXACT)  # The velocity that does move the path.
    for x in range(2):
        u[2, x], u[x, x] = 1.0, -1.0  # A and C jump to G at rate 1; G never leaves.
    return p, u, v, net_flux(u, p), net_flux(u + v, p)  # The last two agree.


# %% 04. The corrector for the mask source
#
# Take the linear schedule, the weights alpha_t = 1 + eta t and beta_t = eta t, and the mask source.
# The backward velocity is then the point mass on the mask, and the combination collapses to two
# brackets.
# A masked position unmasks at total rate (1 + eta t)/(1 - t), split over tokens by the denoiser, and
# a visible position remasks at rate eta.


def corrector_velocity(p1, x_i, t, eta):
    """bar_u^i_t(., x) over the five symbols A, C, G, T, m. p1: [4] -> [5]."""
    current = torch.zeros(DNA_SIZE + 1, dtype=EXACT)
    current[x_i] = 1.0  # delta on the token the position holds, possibly the mask.
    denoiser = torch.cat([p1.to(EXACT), torch.zeros(1, dtype=EXACT)])  # The data token is never m.
    to_mask = torch.zeros(DNA_SIZE + 1, dtype=EXACT)
    to_mask[MASK] = 1.0  # delta_m, the destination of a remasking.
    forward = (1.0 + eta * t) / (1.0 - t) * (denoiser - current)  # The first bracket.
    return forward + eta * (to_mask - current)  # The second bracket, the circulation back to m.


def admissible_step(t, eta):
    """The largest Euler step with a redraw probability at most one, h <= (1 - t)/(1 + eta t)."""
    return (1.0 - t) / (1.0 + eta * t)  # 0.1, 0.0357 and 0.01 at t = 0.9 for eta = 0, 2, 10.


# %% 05. Counting the jumps
#
# Write M(t) for the probability that the position is masked. The two flows cancel, so M_dot(t) = -1
# for every eta and M(t) = 1 - t whatever the corrector strength. If we integrate the unmasking rate
# against M and the remasking rate against 1 - M, we get 1 + eta/2 unmaskings and eta/2 remaskings,
# for 1 + eta jumps per position.


def masked_fraction_derivative(t, eta, masked=None):
    """M_dot(t) = -M (1 + eta t)/(1 - t) + (1 - M) eta, equal to -1 at M = 1 - t."""
    masked = 1.0 - t if masked is None else masked  # The solution we expect.
    return -masked * (1.0 + eta * t) / (1.0 - t) + (1.0 - masked) * eta


def simulate_jumps(p1, eta, positions, steps, generator=None):
    """Run ``positions`` independent positions from the mask and count their jumps.

    Returns the fraction still masked at t = 0.5 and the mean number of jumps per position.
    """
    state = torch.full((positions,), MASK, dtype=torch.long)  # Every position masked at t = 0.
    jumps = torch.zeros(positions, dtype=torch.int32)  # One counter per position.
    h, halfway = 1.0 / steps, None  # The increment, and the masked fraction recorded at t = 0.5.
    for k in range(steps):
        t = k * h
        masked = state == MASK  # [P] which positions are still at the mask symbol.
        if halfway is None and t >= 0.5:
            halfway = float(masked.to(torch.get_default_dtype()).mean())  # M(0.5), put at 1 - t by the chapter.
        up = min(1.0, h * (1.0 + eta * t) / max(1.0 - t, 1e-12))  # Unmasking probability this step.
        down = h * eta  # Remasking probability this step.
        coin = torch.rand(positions, generator=generator)
        unmask = masked & (coin < up)  # Masked positions that reveal a token.
        remask = (~masked) & (coin < down)  # Visible positions that return to the mask.
        revealed = int(unmask.sum())  # Draw a token only where one is needed.
        if revealed:
            tokens = torch.multinomial(p1, revealed, replacement=True, generator=generator)
            state = state.masked_scatter(unmask, tokens)  # Split over tokens by the denoiser.
        state = state.masked_fill(remask, MASK)  # The circulation back to the mask symbol.
        jumps += (unmask | remask).to(torch.int32)  # Every reveal and remasking counts as a jump.
    return halfway, float(jumps.double().mean())  # Near 1 - t = 0.5, and near 1 + eta.


# %% 06. The parallel-step error
#
# Let the data be uniform over AA and TT and the source be all masks. With both positions masked the
# denoiser at each is (0.5, 0.5), and a step that reveals both at once draws them independently, so
# half of those draws land outside the data. With eta > 0 a pair revealed jointly is remasked at rate
# 2 eta and redrawn with the other position visible, where the denoiser is a point mass.


def joint_reveal_rate(steps):
    """The probability that the two reveal times land in the same cell of a uniform grid, h."""
    return 1.0 / steps  # Sum over cells of h squared, or 1/64 at 64 steps.


def pair_survival(eta, tau=0.5):
    """exp(-2 eta (1 - tau)), the chance a pair revealed at tau is untouched at t = 1."""
    return math.exp(-2.0 * eta * (1.0 - tau))  # 0.135 at eta = 2 and 4.54e-5 at eta = 10.


def simulate_survival(eta, samples, steps, tau=0.5, generator=None):
    """The fraction of pairs revealed at ``tau`` that neither position remasks before t = 1.

    Each of the two positions returns to the mask at rate eta, so the first remasking of either
    happens at rate 2 eta, and a pair that is remasked is redrawn with the other position visible.
    """
    remaining = max(1, int(steps * (1.0 - tau)))  # Grid cells left between tau and t = 1.
    h = (1.0 - tau) / remaining  # The increment, matching the grid the sampler runs on.
    untouched = torch.ones(samples, dtype=torch.bool)  # Pairs no remasking has reached.
    for _ in range(remaining):
        coin = torch.rand(samples, 2, generator=generator)  # One clock per position.
        untouched &= ~(coin < h * eta).any(dim=1)  # The survival ends as soon as either position remasks.
    return float(untouched.to(torch.get_default_dtype()).mean())  # Approaches exp(-2 eta (1 - tau)).


def simulate_pair(eta, samples, steps, generator=None):
    """The AA-or-TT two-position problem under the corrector. Returns the fraction outside the data."""
    state = torch.full((samples, 2), MASK, dtype=torch.long)  # Both positions masked at t = 0.
    h = 1.0 / steps
    for k in range(steps):
        t, last = k * h, k == steps - 1  # The last step reveals whatever is left.
        masked = state == MASK  # [S, 2] read before either update, so the two are simultaneous.
        other = state.flip(1)  # The token at the opposite position, read by the denoiser.
        guess = torch.randint(0, 2, state.shape, generator=generator)  # The (0.5, 0.5) denoiser.
        drawn = torch.where(other != MASK, other, guess)  # A point mass once the partner is visible.
        up = 1.0 if last else min(1.0, h * (1.0 + eta * t) / max(1.0 - t, 1e-12))
        down = 0.0 if last else h * eta  # No remasking on the step that ends the run.
        reveal = masked & (torch.rand(state.shape, generator=generator) < up)
        remask = (~masked) & (torch.rand(state.shape, generator=generator) < down)
        state = torch.where(reveal, drawn, state)  # Unmask, in parallel at both positions.
        state = torch.where(remask, torch.full_like(state, MASK), state)  # The circulation back to m.
    return float((state[:, 0] != state[:, 1]).double().mean())  # Outside {AA, TT}.


# %% 07. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--eta", type=float, default=2.0, help="corrector strength in the step table")
    parser.add_argument("--positions", type=int, default=40_000, help="positions in the simulation")
    parser.add_argument("--pairs", type=int, default=20_000, help="sequences in the pair simulation")
    args = parser.parse_args()
    set_seed(args.seed)

    p, u, v, with_u, with_both = three_state_example()
    report = {
        "zero_flux": [round(float(f), 6) for f in net_flux(v, p)],  # Zero at A, C and G.
        "p_dot_u": [round(float(f), 4) for f in with_u],  # (-0.2, -0.3, 0.5).
        "p_dot_u_plus_v": [round(float(f), 4) for f in with_both],  # The same three numbers.
        "rate_out_of_A_u": round(float(-u[0, 0]), 4),  # 1 under the velocity alone.
        "rate_out_of_A_both": round(float(-(u + v)[0, 0]), 4),  # 5 once the corrector is added.
    }

    # Section 5.14.3: one masked position at t = 0.4 with the chapter's denoiser and h = 0.01.
    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
    bar_u = corrector_velocity(p1, MASK, t=0.4, eta=args.eta)  # [5] rates out of the mask.
    seen = torch.zeros(DNA_SIZE)
    seen[1] = 1.0  # At a visible position the true denoiser is the point mass on the visible token.
    report["unmask_rate"] = round(float(-bar_u[MASK]), 4)  # (1 + eta t)/(1 - t); 3.0 at eta = 2.
    report["remask_rate"] = round(float(corrector_velocity(seen, 1, 0.4, args.eta)[MASK]), 4)
    report["step_probabilities"] = [round(float(0.01 * r), 5) for r in bar_u[:DNA_SIZE]]
    report["stay_masked"] = round(float(1.0 + 0.01 * bar_u[MASK]), 5)  # 0.97 at eta = 2.
    report["remask_probability"] = round(0.01 * args.eta, 4)  # eta h, or 0.02 at eta = 2.
    report["masked_derivative"] = round(masked_fraction_derivative(0.4, args.eta), 6)  # -1.
    report["admissible_step_at_0.9"] = round(admissible_step(0.9, args.eta), 4)  # 0.0357 at eta = 2.

    # The long simulation, with the letter frequency of the synthetic DNA as the denoiser.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
    frequency = counts / counts.sum()  # [4] the GC-rich background the lecture shares.
    for eta in (0.0, 2.0, 10.0):
        half, mean = simulate_jumps(frequency, eta, args.positions, args.steps)
        report[f"masked_at_0.5_eta_{eta:g}"] = round(half, 4)  # 1 - t = 0.5 at every eta.
        report[f"jumps_eta_{eta:g}"] = round(mean, 3)  # 1 + eta.

    # Section 5.14.4: the two positions that have to agree.
    report["joint_reveal_at_64"] = round(joint_reveal_rate(64), 4)  # 0.0156.
    report["error_at_64_predicted"] = round(0.5 * joint_reveal_rate(64), 4)  # Half of those are wrong.
    report["error_at_64_measured"] = f"{simulate_pair(0.0, args.pairs, 64):.5f}"  # Near 0.0078.
    report["error_at_grid_measured"] = f"{simulate_pair(0.0, args.pairs, args.steps):.5f}"
    for eta in (0.0, 2.0, 10.0):
        predicted = pair_survival(eta)  # exp(-2 eta (1 - tau)) at the midpoint tau = 0.5.
        report[f"survival_eta_{eta:g}"] = f"{predicted:.3e}"
        if predicted * args.pairs >= 20:  # Below that, the simulation has nothing to resolve.
            measured = simulate_survival(eta, args.pairs, args.steps)  # The same number, simulated.
            report[f"survival_eta_{eta:g}_measured"] = f"{measured:.3e}"
    print_report("correctors", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()