"""Discrete flow matching, the whole construction in one file.

We pick the probability path first, write down a velocity that generates it, and read the sampler off
the velocity. Each position follows the mixture path between a source token and a data token, and
its conditional velocity is a schedule factor times the difference of two point masses. The
posterior average of those conditional velocities is the same factor times a denoiser.
Cross-entropy on the data token recovers that denoiser, and one Euler step redraws a position from
it with probability lambda_t. Chapter 4 fixed the corruption and solved for its reverse; here the
corruption is a consequence of the path we chose.

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.5 to 5.11. Trained on the synthetic DNA of ``cis6270.data``. Prints the
chapter's worked numbers, the training loss and the measurements of a generated batch.

    python lecture_5/dfm.py                     # train and generate from the uniform source
    python lecture_5/dfm.py --source mask       # the all-mask source instead
    python lecture_5/dfm.py --steps 20 --quiet  # a few seconds, for a first look
"""

# %% 01. Imports and configuration

import torch
import torch.nn.functional as F
from torch.distributions import Categorical

from cis6270.data import DNA_SIZE, decode_dna, dna_splits, encode_dna
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters
from cis6270.runner import common_parser, print_report, save_run, set_seed, train


# %% 02. The class of Section 5.11, with both sources
#
# The constructor holds the vocabulary size V and the schedule kappa_t with its derivative
# kappa_dot_t, passed as functions of a time tensor. The defaults are the linear schedule
# kappa_t = t and kappa_dot_t = 1. A mask source returns the index of the mask symbol at every
# position, with the network's vocabulary raised by one to make room for it.


class DiscreteFlowMatching:
    """The source, the mixture path, the cross-entropy loss and one sampling step."""

    def __init__(self, V, kappa=lambda t: t, kappa_dot=lambda t: torch.ones_like(t), mask=None):
        self.V, self.k, self.kd = V, kappa, kappa_dot  # Alphabet size and the schedule.
        self.mask = mask  # Mask token id under the mask source, None under the uniform source.

    def source(self, shape, device="cpu"):
        if self.mask is None:
            return torch.randint(0, self.V, shape, device=device)  # Uniform source, [B, L].
        return torch.full(shape, self.mask, dtype=torch.long, device=device)  # All-mask source.

    def corrupt(self, x1, t):
        """One draw from p_t(. | x_0, x_1) by the coin flip of Section 5.5.5. x1: [B, L], t: [B]."""
        x0 = self.source(x1.shape, x1.device)  # The other end of the pair; [B, L].
        heads = torch.rand(x1.shape, device=x1.device) < self.k(t)[:, None]  # Heads with prob kappa_t.
        return torch.where(heads, x1, x0)  # The data token on heads and the source token on tails.


# %% 03. The conditional velocity
#
# If we differentiate the conditional path at one position we get kappa_dot_t times the difference
# of the two indicators. We then divide that flux by the probability 1 - kappa_t still sitting on
# the source token, so the flux becomes a rate out of the token the position currently holds.


def conditional_velocity(z, x1, t, V=DNA_SIZE):
    """u_t(a, z | x_1) = kappa_dot/(1 - kappa) [delta_{x_1}(a) - delta_z(a)]. Returns [V]."""
    factor = 1.0 / (1.0 - t)  # Schedule factor on the linear schedule, kappa_dot_t = 1.
    return factor * (F.one_hot(torch.tensor(x1), V) - F.one_hot(torch.tensor(z), V)).double()


# %% 04. The posterior average, collapsed onto the data token
#
# The conditional velocity depends on the pair only through the data token at this position, so every
# pair sharing that token contributes the same term and the sum over pairs collapses onto the four
# tokens. That leaves the denoiser p_{1|t}(. | x), the only quantity a network has to learn.


def marginal_velocity(p1, x_i, t, V=DNA_SIZE):
    """u^i_t(., x) = kappa_dot/(1 - kappa) [p_{1|t}(. | x) - delta_{x^i}]. p1: [V] -> [V]."""
    factor = 1.0 / (1.0 - t)  # Same schedule factor, now multiplying a distribution.
    return factor * (p1 - F.one_hot(torch.tensor(x_i), V).double())  # Signed, sums to zero.


def mask_posterior(data, weights, visible, kappa_t):
    """The posterior over the data token at the masked position, computed by enumeration.

    ``data`` is a list of candidate sequences and ``weights`` their probabilities under q. Position 0
    is visible and holds ``visible``; position 1 is masked. The time factor is common to every
    candidate and cancels, Proposition 6 of Discrete Flow Matching.
    """
    masked, seen = 1, 1  # One masked position and one visible position in this example.
    factor = (1.0 - kappa_t) ** masked * kappa_t ** seen  # Common to every surviving candidate.
    unnormalized = torch.zeros(DNA_SIZE, dtype=torch.double)  # One entry per base at position 1.
    for sequence, weight in zip(data, weights):
        if sequence[0] == visible:  # Drop any candidate that disagrees with the visible token.
            unnormalized[encode_dna([sequence])[0, 1]] += factor * weight
    return unnormalized / unnormalized.sum()  # [V] the denoiser at the masked position.


# %% 05. The cross-entropy loss and its minimizer
#
# The loss scores the true data token at every position, with no time weight and no rate in it, so we
# choose the sampler after training. If we average over the data token given X_t, the expected loss
# splits into the entropy of the true denoiser and a KL divergence from it to the network, and only
# the divergence depends on theta.


def dfm_loss(logits, x1):
    """Cross-entropy of Section 5.9, with the class dimension second. logits: [B, L, V]."""
    return F.cross_entropy(logits.transpose(1, 2), x1)  # Mean over the batch and over positions.


def entropy_and_gap(posterior, prediction):
    """The minimum of the loss at one position and the excess cost of a wrong prediction, in nats."""
    entropy = -(posterior * posterior.log()).sum()  # H(p_{1|t}), the floor no network beats.
    cross = -(posterior * prediction.log()).sum()  # The cross-entropy the prediction scores.
    return float(entropy), float(cross - entropy)  # The gap is the KL divergence to the prediction.


# %% 06. The Euler step as a two-part rule
#
# If we substitute the marginal velocity into the per-position Euler step, the two terms collect
# into a mixture of the current point mass and the denoiser. With probability lambda_t the position redraws
# from the denoiser, and with probability 1 - lambda_t it keeps the token it holds.


def redraw_probability(t, h):
    """lambda_t = h kappa_dot_t / (1 - kappa_t), the redraw probability of Section 5.8."""
    return h / (1.0 - t)  # Linear schedule; above one the sampler redraws with certainty.


@torch.no_grad()
def dfm_step(proc, model, xt, t, h):
    """One two-part step at every position at once. xt: [B, L], t: [B], h a scalar."""
    p1 = model(xt, t)[..., :proc.V].softmax(-1)  # [B, L, V] denoiser, the only network call.
    rate = proc.kd(t) / (1 - proc.k(t)).clamp(min=1e-4)  # kappa_dot/(1 - kappa), kept finite at t=1.
    redraw = torch.rand(xt.shape, device=xt.device) < (h * rate)[:, None]  # [B, L] Bernoulli.
    x1 = Categorical(probs=p1).sample()  # [B, L] one candidate token per position.
    return torch.where(redraw, x1, xt)  # Keep the candidate where the coin came up true.


# %% 07. Training and generation
#
# A training step draws one time per sequence, corrupts the batch at that time, and scores the true
# tokens. To generate we walk a uniform grid of ``sample_steps`` times from t = 0 to t = 1, and the
# grid is free to change because no step size enters the training loss.


@torch.no_grad()
def generate(proc, model, batch, length, steps, device="cpu"):
    """Run the sampler from the source to t = 1 and return the tokens. [B, L]."""
    x = proc.source((batch, length), device)  # Draw from the source distribution p.
    grid = torch.linspace(0.0, 1.0, steps + 1, device=device)  # The sampling times.
    for t, s in zip(grid[:-1], grid[1:]):
        x = dfm_step(proc, model, x, t.expand(batch), s - t)  # One forward pass per step.
    return x


# %% 08. The chapter's worked numbers
def worked_numbers():
    """The arithmetic Sections 5.5.5, 5.7.4, 5.8, 5.9.2 and 5.9.3 carry out by hand."""
    report = {}

    # Section 5.5.5: x_1 = ACGT, a uniform source draw x_0 = GTTA, kappa_t = 0.5, coins HTHT.
    x1, x0 = encode_dna(["ACGT"])[0], encode_dna(["GTTA"])[0]  # [4] each.
    heads = torch.tensor([True, False, True, False])  # The four coins of the example.
    xt = torch.where(heads, x1, x0)  # Data token on heads, source token on tails.
    report["mixture_draw"] = decode_dna(xt[None])[0]  # ATGA.

    # Section 5.7.4: one position holding G, denoiser (0.1, 0.6, 0.2, 0.1), t = 0.5.
    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1], dtype=torch.double)  # [V] over A, C, G, T.
    u = marginal_velocity(p1, x_i=2, t=0.5)  # Factor 1/0.5 = 2 times the bracket.
    report["velocity_at_G"] = [round(float(v), 4) for v in u]  # (0.2, 1.2, -1.6, 0.2).
    report["velocity_sums_to"] = f"{float(u.sum()):.1e}"  # Zero, the second rule.
    step = F.one_hot(torch.tensor(2), DNA_SIZE).double() + 0.1 * u  # One Euler step of size h = 0.1.
    report["euler_step_at_G"] = [round(float(v), 4) for v in step]  # (0.02, 0.12, 0.84, 0.02).
    report["redraw_probability"] = round(redraw_probability(0.5, 0.1), 4)  # lambda_t = 0.2.

    # Section 5.9.2: the noisy input ATGA scored against the target ACGT.
    scored = torch.tensor([0.7, 0.2, 0.6, 0.1], dtype=torch.double)  # One probability per position.
    report["example_loss_nats"] = round(float(-scored.log().sum()), 3)  # 4.780 nats.

    # Section 5.9.3: length two, the mask source, q(AC) = 0.5, q(AG) = 0.3, q(GC) = 0.2.
    candidates, weights = ["AC", "AG", "GC"], [0.5, 0.3, 0.2]
    half = mask_posterior(candidates, weights, visible="A", kappa_t=0.5)  # Denoiser at kappa = 0.5.
    late = mask_posterior(candidates, weights, visible="A", kappa_t=0.8)  # The same at kappa = 0.8.
    report["mask_posterior_C_G"] = [round(float(half[1]), 4), round(float(half[2]), 4)]
    report["mask_posterior_unchanged"] = bool(torch.allclose(half, late))  # The time cancels.
    live = half[half > 0]  # The two bases the posterior keeps.
    entropy, gap = entropy_and_gap(live, torch.tensor([0.5, 0.5], dtype=torch.double))
    report["posterior_entropy_nats"] = round(entropy, 4)  # 0.6616, the floor on the loss.
    report["kl_gap_to_uniform"] = round(gap, 4)  # 0.0316 above that floor.
    return report


# %% 09. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--source", choices=["uniform", "mask"], default="uniform")
    parser.set_defaults(steps=1000, width=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)

    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    data = splits["train_tokens"]  # [N, L] the shared synthetic DNA.
    mask = DNA_SIZE if args.source == "mask" else None  # Index 4 is the mask symbol m.
    vocab = DNA_SIZE + 1 if mask is not None else DNA_SIZE  # The network's vocabulary.
    proc = DiscreteFlowMatching(DNA_SIZE, mask=mask)  # Linear schedule by default.
    model = SequenceDenoiser(vocab, args.length, width=args.width, depth=2)

    def loss_fn(step):
        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
        x1 = data[rows]  # [B, L] the data end of the pair.
        t = torch.rand(args.batch_size)  # One time per sequence, uniform on [0, 1].
        xt = proc.corrupt(x1, t)  # [B, L] a draw from the conditional path.
        return dfm_loss(model(xt, t)[..., :DNA_SIZE], x1)  # Score the true token at every position.

    if not args.quiet:
        print(f"source={args.source} vocab={vocab} parameters={count_parameters(model):,}")
    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
    tokens = generate(proc, model, args.samples, args.length, args.sample_steps)

    report = worked_numbers()
    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])  # Averaged over the last steps.
    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
    print_report(f"discrete flow matching, {args.source} source", report, decode_dna(tokens))
    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))


if __name__ == "__main__":
    main()