"""Simplex diffusion on DNA, a belief at every position in place of a token.

Every model before this one draws a token at each reverse step and discards the rest of the
denoiser's distribution, so a position the model was unsure about looks the same as one it was
certain about. Here the state at a position is a probability vector over the four bases, and the
corruption is a Dirichlet around the uniform-replacement kernel. One reverse step thins the
current belief, draws an innovation and mixes the two.

Chapter 4, "From tokens to beliefs" and "Thinning, an innovation, and a mixture". Time runs from
t = 0 clean to t = 1 noise, the survival probability is alpha_t = 1 - t, the concentration is
c_t = gamma K / (1 - alpha_t), and the rates R_t(x, y) act on probability rows, p_dot = p R.
Trained on the synthetic DNA of cis6270.data. Prints the worked Dirichlet parameters, the worked
thinning weight and its Monte Carlo check, then the sample statistics.

    python lecture_4/simplex.py                    # train and generate
    python lecture_4/simplex.py --mixture-draws 10000  # a coarser check of the mixture
    python lecture_4/simplex.py --steps 20         # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import torch  # Tensors, gradients and the Dirichlet, Beta and categorical draws.
import torch.nn as nn  # The denoiser, where only the input layer changes.
import torch.nn.functional as F  # Cross-entropy at every position.

from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters  # The transformer the belief feeds.
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

GAMMA = 1.0  # The constant of the concentration schedule c_t = gamma K / (1 - alpha_t).
TIME_FLOOR = 1e-3  # The concentration grows like 1/(1 - alpha_t), so the grid stops above zero.


# %% 02. The forward marginal as a Dirichlet
#
# Fact 2 requires a forward marginal in closed form, and we want the belief to fluctuate around
# the marginal the token models already use,
# E[P_t | x_0] = alpha_t onehot(x_0) + (1 - alpha_t) pi. The concentration c_t sets the spread.
def concentration(alpha: float | torch.Tensor) -> float | torch.Tensor:
    """c_t = gamma K / (1 - alpha_t), the schedule that makes the increment nonnegative."""
    return GAMMA * DNA_SIZE / (1.0 - alpha)  # 8 at alpha_t = 0.5 and 20 at alpha_s = 0.8.


def belief_mean(alpha: float, x0: int) -> torch.Tensor:
    """E[P_t | x_0] over (A, C, G, T), the uniform-replacement kernel of the token models. [K]."""
    mean = torch.full((DNA_SIZE,), (1.0 - alpha) / DNA_SIZE)  # [K] the share from pi.
    mean[x0] = mean[x0] + alpha  # The clean token keeps the survival mass on top of its share.
    return mean  # [K] a distribution, since alpha + (1 - alpha) = 1.


def dirichlet_parameters(alpha: float, x0: int) -> torch.Tensor:
    """c_t E[P_t | x_0], the Dirichlet parameters at one position. Returns [K]."""
    return concentration(alpha) * belief_mean(alpha, x0)  # (5, 1, 1, 1) at alpha_t = 0.5.


def parameter_field(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
    """The same parameters for a whole batch. x0: [B, L], alpha: [B, 1] -> [B, L, K].

    Every coordinate other than the clean token has parameter gamma, constant in time, and the
    clean coordinate has gamma (1 + K alpha_t/(1 - alpha_t)), the factor of the uniform rates.
    """
    off = torch.full(x0.shape + (DNA_SIZE,), GAMMA, device=x0.device)  # [B, L, K] the pi share.
    clean = GAMMA * DNA_SIZE * alpha / (1.0 - alpha)  # [B, 1] the survival mass times c_t.
    return off.scatter_add(2, x0[:, :, None], clean[..., None].expand_as(x0[:, :, None]))


def corrupt(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
    """Draw P_t ~ Dir(c_t E[P_t | x_0]) at each position. x0: [B, L], alpha: [B, 1] -> [B, L, K]."""
    return torch.distributions.Dirichlet(parameter_field(x0, alpha)).sample()  # [B, L, K] beliefs.


# %% 03. Thinning, an innovation, and a mixture
#
# A Dirichlet with parameter a + b equals W V + (1 - W) U for V ~ Dir(a), U ~ Dir(b) and
# W ~ Beta(sum a, sum b). With this concentration schedule the increment b sits entirely on the
# clean coordinate, so the innovation U is the one-hot vector of the clean token.
def increment(alpha_t: float, alpha_s: float, x0: int) -> torch.Tensor:
    """b = c_s E[P_s | x_0] - c_t E[P_t | x_0], nonnegative for alpha_s > alpha_t. Returns [K]."""
    return dirichlet_parameters(alpha_s, x0) - dirichlet_parameters(alpha_t, x0)  # (12, 0, 0, 0).


def thinning_law(alpha_t: float, alpha_s: float) -> tuple[float, float]:
    """The two Beta parameters, sum a = c_t and sum b = c_s - c_t, or 8 and 12 here."""
    return float(concentration(alpha_t)), float(concentration(alpha_s) - concentration(alpha_t))


def reverse_step(belief: torch.Tensor, clean: torch.Tensor,
                 alpha_t: float, alpha_s: float) -> torch.Tensor:
    """P_s = W P_t + (1 - W) onehot(x_0), one Beta draw per position. belief: [B, L, K]."""
    first, second = thinning_law(alpha_t, alpha_s)  # The two parameters of the thinning weight.
    w = torch.distributions.Beta(  # [B, L] one weight per position, with mean c_t / c_s.
        torch.tensor(first, device=belief.device),
        torch.tensor(second, device=belief.device)).sample(belief.shape[:2])
    innovation = F.one_hot(clean, DNA_SIZE).to(belief.dtype)  # [B, L, K] the corner it moves to.
    return w[..., None] * belief + (1.0 - w[..., None]) * innovation  # [B, L, K] the mixture.


# %% 04. The denoiser that reads a belief
#
# Only the input layer changes. We look a token up in the embedding table, and a belief is the
# same K rows of that table combined in its own proportions, so a one-hot input recovers the
# denoiser of every other file in this lecture exactly.
class BeliefDenoiser(nn.Module):
    """d_theta(. | P_t, t), the shared sequence denoiser reading K numbers at every position."""

    def __init__(self, length: int, width: int = 128):
        super().__init__()
        self.inner = SequenceDenoiser(DNA_SIZE, length, width=width)  # Embeddings and transformer.

    def forward(self, beliefs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """beliefs: [B, L, K], t: [B] -> logits over the clean tokens [B, L, K]."""
        inner = self.inner
        positions = torch.arange(beliefs.shape[1], device=beliefs.device)  # [L].
        h = (beliefs @ inner.token_embedding.weight  # [B, L, D] the belief mixes the K embeddings.
             + inner.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
             + inner.time_embedding(t)[:, None, :])  # [B, 1, D] how far along the path we are.
        return inner.head(inner.transformer(h))  # [B, L, K] every position attends to every other.


# %% 05. The loss
#
# Fact 5 is the cross-entropy whose optimum is the posterior over clean tokens. The reverse step
# moves the denoiser's answer through the mixture of Section 03, and its weight already comes
# from the concentration schedule, so the time derivative of alpha_t multiplies nothing here.
def loss(model: nn.Module, x0: torch.Tensor) -> torch.Tensor:
    """One Monte Carlo sample of E[-log d_theta(x_0 | P_t, t)], summed over positions."""
    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
    alpha = 1.0 - t  # [B, 1] the linear schedule, alpha_0 = 1 and alpha_1 = 0.
    beliefs = corrupt(x0, alpha)  # [B, L, K] one Dirichlet draw per position.
    logits = model(beliefs, t[:, 0])  # [B, L, K] the denoiser reads the whole field of beliefs.
    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
    return ce.sum(-1).mean()  # Nats per sequence, with every position counted.


# %% 06. The sampler
#
# Generation starts from Dir(gamma, ..., gamma), the law the corruption ends at when alpha_1 = 0.
# Each step draws one clean token from the denoiser, thins the current belief and mixes the two,
# so the belief moves through the interior and reaches a corner at the last step.
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
    """Generate from beliefs scattered around pi down to clean DNA. Returns [B, L] over ACGT."""
    device = next(model.parameters()).device
    start = torch.full((batch, length, DNA_SIZE), GAMMA, device=device)  # [B, L, K] at alpha_1 = 0.
    belief = torch.distributions.Dirichlet(start).sample()  # [B, L, K] scattered around pi.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid, t down to 0.
    for t, s in zip(ts[:-1], ts[1:]):  # One interval [s, t] at a time.
        logits = model(belief, t.expand(batch))  # [B, L, K] the clean-token prediction.
        clean = torch.distributions.Categorical(logits=logits).sample()  # [B, L] the innovation.
        belief = reverse_step(belief, clean, float(1.0 - t), float(1.0 - s))  # [B, L, K].
    return belief.argmax(-1)  # [B, L] the corner each position has reached.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--mixture-draws", type=int, default=400_000,
                        help="draws in the check of the thinning representation")
    # The belief holds the denoiser's uncertainty between steps. Enough reverse steps have to run
    # for it to appear. A default run trains 1000 steps and measures 64 sequences, inside about a
    # minute.
    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)  # The same seed gives the same training batches and the same samples.
    device = pick_device(args.device)
    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.

    # Section 02, the worked Dirichlet at the two noise levels of the chapter, with a clean A.
    at_t = dirichlet_parameters(0.5, 0)  # [K] (5, 1, 1, 1) at the later of the two times.
    at_s = dirichlet_parameters(0.8, 0)  # [K] (17, 1, 1, 1) at the earlier one.
    total_t, total_s = at_t.sum(), at_s.sum()  # c_t and c_s, the totals that set every variance.
    spread_t = (at_t * (total_t - at_t) / (total_t ** 2 * (total_t + 1))).sqrt()  # [K] exact.
    spread_s = (at_s * (total_s - at_s) / (total_s ** 2 * (total_s + 1))).sqrt()  # [K] exact.
    talk(f"alpha_t = 0.5       c_t {concentration(0.5):.0f}  mean "
         f"{[round(float(v), 3) for v in belief_mean(0.5, 0)]}  parameters "
         f"{[round(float(v), 0) for v in at_t]}  std on A {float(spread_t[0]):.3f}")
    talk(f"alpha_s = 0.8       c_s {concentration(0.8):.0f}  mean "
         f"{[round(float(v), 2) for v in belief_mean(0.8, 0)]}  parameters "
         f"{[round(float(v), 0) for v in at_s]}")

    # Section 03, the increment, the thinning weight, and the mean the mixture lands on.
    step = increment(0.5, 0.8, 0)  # [K] all of the increase sits on the clean coordinate.
    first, second = thinning_law(0.5, 0.8)  # Beta(8, 12), with mean 0.4.
    mean_weight = first / (first + second)  # c_t / c_s, the average share the old belief keeps.
    mixed = mean_weight * belief_mean(0.5, 0) + (1.0 - mean_weight) * torch.eye(DNA_SIZE)[0]  # [K].
    talk(f"increment           {[round(float(v), 0) for v in step]}  W ~ Beta({first:.0f}, "
         f"{second:.0f})  mean {mean_weight:.1f}  mixture mean "
         f"{[round(float(v), 2) for v in mixed]}")

    # The same mixture by simulation, against the exact standard deviations of Dir(17, 1, 1, 1).
    thinned = torch.distributions.Dirichlet(at_t).sample((args.mixture_draws,))  # [N, K] draw V.
    weight = torch.distributions.Beta(  # [N] the Beta draw of the representation.
        torch.tensor(first), torch.tensor(second)).sample((args.mixture_draws,))
    drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(DNA_SIZE)[0]  # [N, K].
    talk(f"mixture check       {args.mixture_draws} draws  std "
         f"{[round(float(v), 3) for v in drawn.std(0)]}  against the exact "
         f"{[round(float(v), 3) for v in spread_s]}")

    # Train the denoiser on the shared synthetic DNA. We reseed here so the training batches stay
    # independent of how many draws the check above consumed.
    set_seed(args.seed)
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
    model = BeliefDenoiser(args.length, width=args.width).to(device)
    talk(f"device={device} parameters={count_parameters(model):,} "
         f"train={len(train_tokens)} length={args.length}")

    def step_loss(step_index: int) -> torch.Tensor:
        """One optimizer step draws its own batch of clean sequences."""
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        return loss(model, train_tokens[rows])  # [B, L] clean sequences into one scalar.

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
    with torch.no_grad():
        held_out = sum(float(loss(model, valid_tokens)) for _ in range(4)) / 4.0

    # Generate, then measure the structure the data was built with.
    torch.manual_seed(args.seed + 1)  # Fix the Beta and categorical draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
    report = {
        "c_t at alpha_t = 0.5": float(concentration(0.5)),
        "Dirichlet parameter on A at 0.5": float(at_t[0]),
        "standard deviation on A at 0.5": float(spread_t[0]),
        "Dirichlet parameter on A at 0.8": float(at_s[0]),
        "increment on A": float(step[0]),
        "mean thinning weight": mean_weight,
        "simulated standard deviation on A": float(drawn.std(0)[0]),
        "exact standard deviation on A": float(spread_s[0]),
        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation loss, nats/seq": held_out,
        "loss, nats per token": held_out / args.length,
        "GC fraction": gc_fraction(samples),
        "GC fraction in the data": gc_fraction(train_tokens),
        "motif fraction": motif_fraction(samples),
        "motif fraction in the data": motif_fraction(train_tokens),
        "distinct samples": unique_fraction(samples),
        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
    }
    print_report("simplex: beliefs in place of tokens", report, decode_dna(samples))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(samples))


# %% 08. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the Dirichlet numbers, train the belief denoiser, and generate.