"""Edit Flows: paths whose sequences change length.

Every path so far keeps a single length, because a substitution changes a token and never adds or
removes one. If we enlarge the set of moves to insertions, deletions and substitutions, that
constraint goes away. An alignment pads the two ends with a blank until they have the same number of
columns, and the mixture construction then applies column by column. The alignment is an auxiliary
variable we marginalize exactly as we marginalized the data sequence. The loss scores rates in place
of a distribution over tokens, and its minimizer is the rate at which an edit makes progress.

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, Section 5.17. Runs on the synthetic DNA of ``cis6270.data``, trimmed to varying lengths.
Prints the alignment of the chapter, the posterior over two alignments, the averaged edit rates, the
loss of one example, one sampler step, and the length statistics of a run on DNA.

    python lecture_5/edit_flows.py                     # 600 sampler steps on 256 DNA pairs
    python lecture_5/edit_flows.py --steps 20 --quiet  # a coarse grid, a visibly coarser length
"""

# %% 01. Imports and configuration
import math

import torch

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

BLANK = DNA_SIZE  # Index 4 is the blank symbol epsilon that padding writes into a column.


# %% 02. Three edits and the rate split
#
# Each position of the current sequence has three edits, and each rate splits into an intensity and
# a distribution over tokens. A deletion writes no token, so the network emits three intensities and
# two token distributions per position, or 3 + 2K numbers in place of the K of the denoiser.


def edit_rates(intensities, tokens=DNA_SIZE):
    """Count the numbers a network emits per position, and the total rate of leaving the sequence."""
    total = sum(sum(kind) for kind in intensities.values())  # -u_t(x, x), the total exit rate.
    return 3 + 2 * tokens, total  # Three intensities and two token distributions per position.


# %% 03. Alignments and the column-wise path
#
# An alignment puts the positions of the two sequences in correspondence by padding both with the
# blank. A blank in the target column is a deletion, a blank in the source column an insertion, and
# two different tokens a substitution. On the aligned pair each column follows the mixture path
# independently, and once we remove the blanks from z_t we have the sequence x_t the model reads.


def align(x0, x1):
    """A minimal alignment of two token lists, as the padded pair (z_0, z_1) of equal length."""
    n, m = len(x0), len(x1)
    cost = [[0] * (m + 1) for _ in range(n + 1)]  # Edit distance with unit costs.
    for i in range(n + 1):
        cost[i][0] = i  # Deleting the first i tokens of the source.
    for j in range(m + 1):
        cost[0][j] = j  # Inserting the first j tokens of the target.
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost[i][j] = min(cost[i - 1][j] + 1, cost[i][j - 1] + 1,
                             cost[i - 1][j - 1] + int(x0[i - 1] != x1[j - 1]))
    z0, z1, i, j = [], [], n, m
    while i or j:  # One backtrace through the table, taking the diagonal where it is minimal.
        if i and j and cost[i][j] == cost[i - 1][j - 1] + int(x0[i - 1] != x1[j - 1]):
            z0.append(x0[i - 1]), z1.append(x1[j - 1]), (i := i - 1), (j := j - 1)
        elif i and cost[i][j] == cost[i - 1][j] + 1:
            z0.append(x0[i - 1]), z1.append(BLANK), (i := i - 1)  # A deletion column.
        else:
            z0.append(BLANK), z1.append(x1[j - 1]), (j := j - 1)  # An insertion column.
    return z0[::-1], z1[::-1]  # [C] and [C] over the same number of columns.


def column_kinds(z0, z1):
    """Name the edit each column owes: ``del``, ``ins``, ``sub``, or ``none`` when the two agree."""
    kinds = []
    for a, b in zip(z0, z1):
        kinds.append("none" if a == b else "del" if b == BLANK else "ins" if a == BLANK else "sub")
    return kinds  # [C] one label per column.


def draw_state(z0, z1, finished):
    """z_t from a choice of which columns already hold their target, with the blanks removed."""
    held = [b if done else a for a, b, done in zip(z0, z1, finished)]  # [C] the aligned state.
    return [token for token in held if token != BLANK]  # [|x_t|] the sequence the model reads.


def expected_length(x0, x1, kappa_t):
    """E[|x_t|] = (1 - kappa_t)|x_0| + kappa_t|x_1|, the interpolation one position up."""
    return (1.0 - kappa_t) * len(x0) + kappa_t * len(x1)  # Linear from |x_0| to |x_1|.


# %% 04. Averaging over the alignments we cannot see
#
# Many alignments produce the same pair, and x_t does not record which one the path is on. Bayes'
# rule over the alignments weights each by the probability that its column draw produces the observed
# x_t, and the rates average against those weights.


def alignment_posterior(alignments, observed, kappa_t):
    """p_t(A | x_t) over a list of aligned pairs, by enumerating the column draws of each."""
    weights = []
    for z0, z1 in alignments:
        owed = [i for i, kind in enumerate(column_kinds(z0, z1)) if kind != "none"]  # Unfinished.
        total = 0.0
        for mask in range(1 << len(owed)):  # Every subset of the owed columns already finished.
            finished = [True] * len(z0)
            probability = 1.0
            for bit, column in enumerate(owed):
                done = bool(mask >> bit & 1)
                finished[column] = done
                probability *= kappa_t if done else 1.0 - kappa_t  # The column-wise mixture path.
            if draw_state(z0, z1, finished) == observed:
                total += probability  # This draw is consistent with what we see.
        weights.append(total)
    norm = sum(weights)
    return [w / norm for w in weights]  # One posterior weight per alignment.


def averaged_rates(alignments, observed, kappa_t, kappa_dot=1.0):
    """The marginal rate of each edit that finishes a column, averaged over the alignments.

    ``observed`` is the draw in which every column that owes an edit still holds its source symbol.
    The worked example of Section 5.17.2 reads the rates at that state.
    """
    rate = kappa_dot / (1.0 - kappa_t)  # The weight on every unfinished column.
    posterior, out = alignment_posterior(alignments, observed, kappa_t), {}
    for (z0, z1), weight in zip(alignments, posterior):
        held = [b if kind == "none" else a  # A column owing nothing is finished at every time.
                for a, b, kind in zip(z0, z1, column_kinds(z0, z1))]
        for i, (a, b, kind) in enumerate(zip(z0, z1, column_kinds(z0, z1))):
            seen = sum(1 for token in held[:i] if token != BLANK)  # Position inside x_t.
            if kind == "sub":
                name = f"sub({seen + 1},{DNA[b]})"  # Replace the token at that position.
            elif kind == "ins":
                name = f"ins({seen},{DNA[b]})"  # Place a token to the right of that position.
            elif kind == "del":
                name = f"del({seen + 1})"  # Remove the token at that position.
            else:
                continue  # Nothing is owed here.
            out[name] = out.get(name, 0.0) + weight * rate
    return posterior, out


# %% 05. The training loss
#
# The first term is the sum of the model's intensities and pushes all of them down. The second is a
# weighted cross-entropy on the edits that move x_t toward x_1 and pushes those up. If we collect
# the terms belonging to one edit we get lambda - w log lambda, with stationary point lambda = w.


def edit_loss(total_intensity, owed, kappa_t, kappa_dot=1.0):
    """L = sum of all intensities - sum over unfinished columns of the weight times log lambda."""
    weight = kappa_dot / (1.0 - kappa_t)  # kappa_dot/(1 - kappa), or 2 at t = 0.5.
    return total_intensity - sum(weight * math.log(lam) for lam in owed)  # The model's own rates.


def loss_minimizer(weight, grid):
    """The intensity minimizing lambda - w log lambda on a grid, put at w by the derivative."""
    values = [lam - weight * math.log(lam) for lam in grid]  # The contribution of one edit.
    return grid[min(range(len(grid)), key=values.__getitem__)]  # Near w, with w/lambda^2 > 0.


# %% 06. One step of the sampler
#
# Each position inserts with probability h lambda_ins, and deletes or substitutes with probability
# h(lambda_del + lambda_sub), the two sharing that draw in proportion to their intensities. We apply
# every chosen edit at once, so the step has the same second-order error as the parallel step on a
# fixed length.


def step_probabilities(intensities, h):
    """The per-position insertion and delete-or-substitute probabilities of one Euler step."""
    insert = [h * lam for lam in intensities["ins"]]  # One coin per position.
    change = [h * (d + s) for d, s in zip(intensities["del"], intensities["sub"])]  # A second coin.
    return insert, change


def untouched_probability(intensities, h):
    """The chance the step applies no edit anywhere, and the first-order estimate beside it."""
    insert, change = step_probabilities(intensities, h)
    exact = math.prod(1.0 - p for p in insert + change)  # Every coin coming up tails.
    _, total = edit_rates(intensities)
    return exact, 1.0 - h * total  # The gap is the second-order term from two edits at once.


def expected_length_change(intensities, h):
    """h times the insertion intensities minus the deletion intensities, in tokens."""
    return h * (sum(intensities["ins"]) - sum(intensities["del"]))  # Nothing fixes the length.


# %% 07. Edits on the synthetic DNA
#
# We align each pair of sequences once, and every column that owes an edit finishes at rate
# kappa_dot/(1 - kappa). The sequence the model reads is the aligned state with its blanks removed.
# The realized length wanders while its mean follows the interpolation.


def simulate_columns(pairs, steps, generator=None):
    """Run the column-wise path on a list of aligned pairs. Returns the lengths and the end states."""
    width = max(len(z0) for z0, _ in pairs)  # Pad every alignment to the widest one.
    src = torch.full((len(pairs), width), BLANK)
    tgt = torch.full((len(pairs), width), BLANK)
    valid = torch.zeros(len(pairs), width, dtype=torch.bool)
    for row, (z0, z1) in enumerate(pairs):
        src[row, :len(z0)] = torch.tensor(z0)
        tgt[row, :len(z1)] = torch.tensor(z1)
        valid[row, :len(z0)] = True  # Columns beyond the alignment take no part.
    done = (src == tgt) & valid  # A column whose two symbols agree owes no edit.
    h, halfway = 1.0 / steps, None
    for k in range(steps):
        t = k * h
        if halfway is None and t >= 0.5:
            held = torch.where(done, tgt, src)  # [P, C] the aligned state at this time.
            halfway = float(((held != BLANK) & valid).sum(dim=1).double().mean())
        lam = 1.0 if k == steps - 1 else min(1.0, h / max(1.0 - t, 1e-12))  # lambda_t per column.
        coin = torch.rand(src.shape, generator=generator) < lam
        done = done | (coin & valid)  # An unfinished column takes its target symbol.
    finished = float((done | ~valid).all(dim=1).double().mean())  # Fraction of pairs that arrived.
    return halfway, finished


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--pairs", type=int, default=256, help="DNA pairs aligned and run")
    args = parser.parse_args()
    set_seed(args.seed)

    # Section 5.17.1: the chapter's alignment of ACGT with AGAC over five columns.
    x0, x1 = encode_dna(["ACGT"])[0].tolist(), encode_dna(["AGAC"])[0].tolist()
    z0 = x0 + [BLANK]  # (A, C, G, T, eps).
    z1 = [x1[0], BLANK] + x1[1:]  # (A, eps, G, A, C).
    kinds = column_kinds(z0, z1)
    draw = [True, False, True, True, True]  # Column 2 behind, columns 4 and 5 already at the target.
    report = {
        "column_kinds": kinds,  # none, del, none, sub, ins.
        "draw_at_kappa_0.5": "".join(DNA[token] for token in draw_state(z0, z1, draw)),  # ACGAC.
        "draw_probability": round(0.5 ** sum(kind != "none" for kind in kinds), 4),  # 0.125.
        "expected_length": round(expected_length(x0, x1, 0.5), 4),  # 4, both ends having length 4.
        "numbers_per_position": edit_rates({"ins": [0.0], "del": [0.0], "sub": [0.0]})[0],  # 11.
    }

    # Section 5.17.2: the pair AC and AAC, with two alignments on three columns.
    short, long = encode_dna(["AC"])[0].tolist(), encode_dna(["AAC"])[0].tolist()
    first = ([short[0], short[1], BLANK], long)  # Substitute C with A, then insert C.
    second = ([short[0], BLANK, short[1]], long)  # Insert A, and leave the last column alone.
    observed = short  # The noisy sequence x_t = AC.
    for kappa_t in (0.25, 0.5, 0.75):
        weights = alignment_posterior([first, second], observed, kappa_t)
        report[f"posterior_A1_kappa_{kappa_t}"] = round(weights[0], 4)  # (1-kappa)/(2-kappa).
    posterior, rates = averaged_rates([first, second], observed, 0.5)
    for name in sorted(rates):
        report[f"rate_{name}"] = round(rates[name], 3)  # 0.667, 1.333 and 0.667.
    report["total_target_rate"] = round(sum(rates.values()), 3)  # 2.667.

    # Section 5.17.3: one unfinished column, a model intensity of 0.8, and 3.0 over all edits.
    report["example_loss"] = round(edit_loss(3.0, [0.8], kappa_t=0.5), 4)  # 3.446.
    grid = [0.01 * i for i in range(1, 801)]  # Intensities between 0.01 and 8.
    report["loss_minimizer"] = round(loss_minimizer(2.0, grid), 2)  # 2, the weight the edit receives.

    # Section 5.17.4: one whole step on ACG with the chapter's nine intensities.
    intensities = {"ins": [0.2, 0.3, 0.4], "del": [0.1, 0.9, 0.2], "sub": [0.4, 0.3, 0.5]}
    exact, first_order = untouched_probability(intensities, 0.05)
    insert, change = step_probabilities(intensities, 0.05)
    report["insert_probabilities"] = [round(p, 3) for p in insert]  # 0.010, 0.015, 0.020.
    report["change_probabilities"] = [round(p, 3) for p in change]  # 0.025, 0.060, 0.035.
    report["deletion_given_change"] = round(0.9 / (0.9 + 0.3), 4)  # 0.75, three times in four.
    report["untouched"] = round(exact, 4)  # 0.8452.
    report["untouched_first_order"] = round(first_order, 4)  # 0.835.
    report["second_order_gap"] = round(exact - first_order, 4)  # 0.0102.
    report["length_change"] = round(expected_length_change(intensities, 0.05), 4)  # -0.015.

    # The same path on the shared DNA, trimmed so that the two ends have different lengths.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    tokens = splits["train_tokens"]
    cuts = torch.randint(args.length - 4, args.length + 1, (2 * args.pairs,))  # Varying lengths.
    rows = [tokens[i % len(tokens), :cuts[i]].tolist() for i in range(2 * args.pairs)]
    pairs = [align(rows[2 * i], rows[2 * i + 1]) for i in range(args.pairs)]
    predicted = sum(expected_length(rows[2 * i], rows[2 * i + 1], 0.5) for i in range(args.pairs))
    measured, arrived = simulate_columns(pairs, args.steps)
    report["dna_columns_mean"] = round(sum(len(z0) for z0, _ in pairs) / len(pairs), 3)
    report["dna_length_at_0.5"] = round(measured, 3)  # Measured over the run.
    report["dna_length_predicted"] = round(predicted / args.pairs, 3)  # The interpolation.
    report["dna_pairs_arrived"] = round(arrived, 4)  # Every column finishes by t = 1.
    report["dna_example"] = " -> ".join("".join(DNA[t] for t in row) for row in rows[:2])
    print_report("edit flows", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()