"""Gumbel-Softmax Flow Matching: the path comes from the way a token is sampled.

Sampling a categorical is a deterministic function of noise. We add a Gumbel draw to each log
probability and take the largest. If we replace the argmax by a softmax, the winner becomes a point
of the simplex, and cooling the temperature slides that point from the center toward a corner. With
the noise held fixed the trajectory is a smooth curve, and its time derivative is the replicator
equation of evolutionary dynamics. At inference we set the noise to zero, and the velocity
collapses to a multiple of the direction toward the data corner.

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.22.1 to 5.22.6. Trained on the synthetic DNA of ``cis6270.data``. Prints the
Gumbel-max frequencies, the annealed path at four times, the replicator velocity, the marginal
velocity late in a trajectory, and the measurements of a generated batch.

    python lecture_5/gumbel.py                     # train, then 64 Euler steps on the simplex
    python lecture_5/gumbel.py --steps 20 --quiet  # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from cis6270.data import DNA_SIZE, decode_dna, dna_splits
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

TAU_MAX = 10.0  # The temperature at t = 0, hot enough that every letter holds nearly the same mass.
LAMBDA = 3.0  # The cooling rate, so tau(1) = 10 exp(-3) = 0.498.
BETA = 2.0  # The noise scale; a larger beta makes the paths from one data token spread less.


# %% 02. The Gumbel-max trick
#
# A Gumbel variable is the negative logarithm of an exponential clock, so the trick is Chapter 4's
# race between clocks written with logarithms, where the winner is an argmax and the rates are
# logits. All of the randomness sits in the noise, and the logits are ours to move.


def gumbel(shape, generator=None):
    """g ~ Gumbel(0, 1), as minus the log of minus the log of a uniform draw."""
    u = torch.rand(shape, generator=generator).clamp(min=1e-12, max=1 - 1e-12)
    return -torch.log(-torch.log(u))  # The distribution function is exp(-exp(-x)).


def gumbel_max(probabilities, draws, generator=None):
    """Frequencies of argmax_k [log pi_k + g_k] over ``draws`` samples. Returns [K]."""
    logits = probabilities.log()[None, :] + gumbel((draws, len(probabilities)), generator)
    counts = torch.bincount(logits.argmax(-1), minlength=len(probabilities)).double()
    return counts / counts.sum()  # Approaches pi itself.


# %% 03. The softmax relaxation and a cooling temperature
#
# A softmax of the noisy logits divided by a temperature is a point of the simplex. As the
# temperature falls the point approaches the one-hot of the argmax, and as it grows the point
# approaches the center. The temperature never reaches zero, so the path stays in the interior and
# its density stays defined.


def temperature(t, tau_max=TAU_MAX, rate=LAMBDA):
    """tau(t) = tau_max exp(-lambda t), or 10 at t = 0 and 0.498 at t = 1."""
    return tau_max * math.exp(-rate * t)  # 2.231 at t = 0.5 and 0.672 at t = 0.9.


# %% 04. The annealed path toward a data token
#
# We put the data token into the logits and let the cooling do the rest. That gives the
# conditional path. The noise is drawn once and held for the whole trajectory, so only the
# temperature changes with time and the point moves continuously toward the data corner.


def annealed_path(logits, t):
    """psi_t = softmax((e_k + g/beta)/tau(t)), with the logits already assembled. Returns [..., K]."""
    return torch.softmax(logits / temperature(t), dim=-1)  # A point of the simplex at every time.


def path_logits(targets, noise, K=DNA_SIZE, beta=BETA):
    """e_k + g/beta at every position. targets: [B, L], noise: [B, L, K] -> [B, L, K]."""
    return F.one_hot(targets, K) + noise / beta  # The data indicator plus the scaled Gumbel draw.


# %% 05. Differentiating the path: the replicator equation
#
# The only time dependence sits in one over tau. It grows at relative rate lambda. If we compose
# with the softmax Jacobian we get the current mass times the gap between a letter's logit and the
# mass-weighted average of all of them, the replicator equation with fitness lambda ell_i/tau(t).


def replicator_velocity(x, logits, t, rate=LAMBDA):
    """u_{t,i} = (lambda/tau(t)) x_i (ell_i - sum_j x_j ell_j). x, logits: [..., K] -> [..., K]."""
    average = (x * logits).sum(-1, keepdim=True)  # The mass-weighted average of the logits.
    return (rate / temperature(t)) * x * (logits - average)  # Coordinates summing to zero.


# %% 06. The inference-time velocity
#
# The sampler has only the current point, so we set the noise to zero. The logits become the data
# indicator, the bracket becomes the gap to the data coordinate, and the velocity collapses to a
# multiple of the direction toward the corner, with the factor x_k(1 - x_k) in the k-th component.


def inference_velocity(x, k, t, rate=LAMBDA):
    """u_t(x | x_1 = e_k) = (lambda/tau(t)) x_k (e_k - x). x: [..., K] -> [..., K]."""
    corner = F.one_hot(torch.tensor(k), x.shape[-1]).to(x)  # The data corner.
    return (rate / temperature(t)) * x[..., k, None] * (corner - x)  # Zero once x_k reaches zero.


def marginal_velocity(x, p_hat, t, rate=LAMBDA):
    """The posterior average over corners. x, p_hat: [..., K] -> [..., K]."""
    weight = (rate / temperature(t)) * x * p_hat  # [..., K] the coefficient of each corner.
    return weight - weight.sum(-1, keepdim=True) * x  # sum_k w_k (e_k - x), collected.


# %% 07. Training and the marginal velocity
#
# Cross-entropy on the data token trains the denoiser, as it did for the mixture paths and for the
# Dirichlet path, and the conditional velocity appears nowhere in the gradient. At sampling time the
# marginal velocity is the posterior average, the fourth appearance of the same rule in this chapter.


def relaxed_denoiser(length, width, K=DNA_SIZE, depth=2):
    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token."""
    model = SequenceDenoiser(K, length, width=width, depth=depth)
    model.token_embedding = nn.Linear(K, width, bias=False)  # It now takes a point of the simplex.
    return model


@torch.no_grad()
def generate(model, batch, length, steps, K=DNA_SIZE):
    """Integrate the marginal velocity from the center of the simplex to t = 1. Returns [B, L]."""
    x = annealed_path(gumbel((batch, length, K)) / BETA, 0.0)  # The source, with no data indicator.
    h = 1.0 / steps
    for k in range(steps):
        t = k * h
        p_hat = model(x, torch.full((batch,), t)).softmax(-1)  # [B, L, K] the denoiser.
        x = x + h * marginal_velocity(x, p_hat, t)  # One Euler step along the marginal velocity.
        x = x.clamp(min=1e-8)  # A large step can push a coordinate below zero.
        x = x / x.sum(-1, keepdim=True)  # Project back onto the simplex.
    return x.argmax(-1)  # Read the token off the largest coordinate.


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--draws", type=int, default=400_000, help="draws in the Gumbel-max check")
    parser.set_defaults(steps=1000, width=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)

    pi = torch.tensor([0.5, 0.2, 0.2, 0.1])  # The categorical of the chapter's quick check.
    report = {
        "gumbel_max_frequencies": [round(float(v), 4) for v in gumbel_max(pi, args.draws)],
        "target_probabilities": [round(float(v), 4) for v in pi],  # The frequencies approach these.
        "tau_at_0": round(temperature(0.0), 4),  # 10.
        "tau_at_1": round(temperature(1.0), 4),  # 0.498.
    }

    # Section 5.22.3: one base annealed toward G with the chapter's fixed noise.
    scaled = torch.tensor([0.0, -0.45, 0.40, -0.05])  # g/beta over A, C, G, T.
    logits = F.one_hot(torch.tensor(2), DNA_SIZE) + scaled  # The data indicator at G plus the noise.
    report["path_logits"] = [round(float(v), 2) for v in logits]  # (0, -0.45, 1.40, -0.05).
    for t in (0.0, 0.5, 0.9, 1.0):
        point = annealed_path(logits, t)
        report[f"path_at_{t}"] = [round(float(v), 4) for v in point]

    # Section 5.22.4: the replicator velocity at the midpoint.
    x = annealed_path(logits, 0.5)  # (0.2142, 0.1751, 0.4012, 0.2095).
    report["tau_at_0.5"] = round(temperature(0.5), 4)  # 2.231.
    report["rate_over_tau"] = round(LAMBDA / temperature(0.5), 3)  # 1.345.
    report["average_logit"] = round(float((x * logits).sum()), 4)  # 0.4724.
    noisy = replicator_velocity(x, logits, 0.5)
    report["replicator_velocity"] = [round(float(v), 3) for v in noisy]  # (-0.136, -0.217, 0.5, -0.147).
    report["replicator_sums_to"] = f"{float(noisy.sum()):.1e}"  # Zero, so the step stays on the simplex.

    # Section 5.22.5: the same point with the noise set to zero.
    clean = inference_velocity(x, 2, 0.5)
    report["inference_velocity"] = [round(float(v), 3) for v in clean]  # (-0.116, -0.094, 0.323, ...).

    # Section 5.22.6: one Euler step late in the trajectory.
    late = torch.tensor([0.138, 0.084, 0.647, 0.131])  # The base at t = 0.8.
    p_hat = torch.tensor([0.05, 0.10, 0.70, 0.15])  # The denoiser output there.
    report["tau_at_0.8"] = round(temperature(0.8), 4)  # 0.9072.
    coefficients = (LAMBDA / temperature(0.8)) * late * p_hat
    report["corner_coefficients"] = [round(float(v), 4) for v in coefficients]  # (0.0228, ..., 1.4976).
    v = marginal_velocity(late, p_hat, 0.8)
    report["marginal_velocity"] = [round(float(c), 3) for c in v]  # (-0.200, -0.108, 0.454, -0.146).
    report["euler_step_h_0.01"] = [round(float(c), 4) for c in late + 0.01 * v]  # (0.1360, ...).

    # Training on the shared DNA, then generation.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    data = splits["train_tokens"]  # [N, L].
    model = relaxed_denoiser(args.length, args.width)

    def loss_fn(step):
        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
        x1 = data[rows]  # [B, L] the letters the denoiser has to recover.
        t = torch.rand(args.batch_size)  # One time per sequence, uniform on [0, 1].
        noise = gumbel((args.batch_size, args.length, DNA_SIZE))  # Held fixed along one trajectory.
        cooled = path_logits(x1, noise) / torch.tensor(
            [temperature(float(s)) for s in t])[:, None, None]  # Each sequence at its own temperature.
        xt = cooled.softmax(-1)  # [B, L, K] a point on the conditional path.
        return F.cross_entropy(model(xt, t).transpose(1, 2), x1)  # -log <x_theta, x_1>.

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

    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])
    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("Gumbel-Softmax flow matching", report, decode_dna(tokens))
    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))


if __name__ == "__main__":
    main()