"""Flow Map Language Models, the composition identity in denoiser coordinates.

We lift each token to a one-hot vector, the source is Gaussian over that space, and the state moves
through R^{L x K}, so no atom has to be split. If we step by the full remaining time 1 - s in place
of the elapsed time t - s, the average velocity becomes a two-time denoiser that lands on the
simplex. The map is then a convex combination of staying put and jumping to the prediction, and the
cocycle is a convex combination of two shorter predictions. A KL divergence measures the gap.

Chapter 6, Sections 6.25 to 6.27 of the notes. Trained on the synthetic DNA of
``cis6270.data.synthetic_dna``. Prints the convex weights, the composition weight gamma = 0.2 at
(0.2, 0.5, 0.8), the worked KL of 0.030747 nats, the two gradients that decide the loss, the
decoding clock of a 30,522-token vocabulary and the four steps it produces, and samples from the
trained map.

    python lecture_6/fmlm.py                 # 1500 steps, about half a minute
    python lecture_6/fmlm.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import TwoTimeMLP, count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train


torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
# than it saves. ``tests/conftest.py`` sets the same thing for the test session, and this line
# acts when a reader runs the script directly.

BOOK_VOCAB = 30522  # The vocabulary the chapter puts on the clock, a standard tokenizer.
ENDPOINT_FRACTION = 0.25  # Share of each batch drawn with t = 1, the slice the sampler reads.


# %% 02. The convex form and the composition weight
#
# X_{s,t}(x) = (1-t)/(1-s) x + (t-s)/(1-s) delta_{s,t}(x). Both weights are non-negative for
# s <= t <= 1 and they sum to one, so each step is a convex combination of staying put and jumping
# to the predicted endpoint. At t = s the pair is (1, 0) and the identity axiom holds.
def convex_weights(s, t):
    """The pair ((1-t)/(1-s), (t-s)/(1-s)), the weights on the state and on the prediction."""
    remaining = (1.0 - s).clamp(min=1e-6)
    return (1.0 - t) / remaining, (t - s) / remaining


def composition_weight(s, u, t):
    """gamma = (1-t)(u-s)/((1-u)(t-s)), the share of the long prediction made over the first leg."""
    return (1.0 - t) * (u - s) / (((1.0 - u) * (t - s)).clamp(min=1e-6))


# %% 03. The cocycle in denoiser coordinates, checked on the scalar flow
#
# delta_{s,t}(x) = gamma delta_{s,u}(x) + (1-gamma) delta_{u,t}(X_{s,u}(x)). On dx/dt = x the
# two-time denoiser is x[1 + (1-s)(e^{t-s}-1)/(t-s)], so every term is available in closed form.
def scalar_denoiser(x, s, t):
    """delta_{s,t}(x) on the flow dx/dt = x."""
    return x * (1.0 + (1.0 - s) * math.expm1(t - s) / (t - s))


def scalar_cocycle(x: float = 1.0, s: float = 0.2, u: float = 0.5, t: float = 0.8):
    """Every term of the denoiser cocycle at the chapter's triple of times."""
    gamma = (1.0 - t) * (u - s) / ((1.0 - u) * (t - s))
    first = scalar_denoiser(x, s, u)  # The prediction over the first leg.
    reached = x * math.exp(u - s)  # X_{s,u}(x), the state the first leg arrives at.
    second = scalar_denoiser(reached, u, t)  # The prediction over the second leg.
    long = scalar_denoiser(x, s, t)  # The prediction over the whole interval.
    keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
    return {"gamma": gamma, "delta_s,u": first, "delta_u,t": second, "delta_s,t": long,
            "two legs combined": gamma * first + (1.0 - gamma) * second,
            "convex weights": (round(keep, 4), round(move, 4)),
            "convex form": keep * x + move * long, "X_s,t(x)": x * math.exp(t - s)}


# %% 04. One term of the loss, on the DNA alphabet
#
# The composed target is a convex combination of two simplex points, so it is a distribution and
# the divergence against it needs no projection. Its gradient pushes mass onto the token with more
# weight in the two-leg decomposition than in the long prediction.
def worked_kl():
    """The chapter's four-token example, with the target composed from two student predictions."""
    first = torch.tensor([0.10, 0.20, 0.60, 0.10])  # delta_{0.2,0.5} over (A, C, G, T).
    second = torch.tensor([0.05, 0.10, 0.80, 0.05])  # delta_{0.5,0.8} at the state the leg reaches.
    gamma = (1.0 - 0.8) * (0.5 - 0.2) / ((1.0 - 0.5) * (0.8 - 0.2))
    target = gamma * first + (1.0 - gamma) * second  # [4] the composed target.
    student = torch.tensor([0.10, 0.15, 0.65, 0.10])  # The student's long prediction.
    terms = target * (target / student).log()  # [4] the four terms of the divergence.
    return {"worked gamma": gamma, "composed target": [round(v, 4) for v in target.tolist()],
            "target mass": float(target.sum()), "worked KL": float(terms.sum()),
            "the four terms": [round(v, 6) for v in terms.tolist()]}


# %% 05. The two gradients of a softmax head
#
# With a softmax head the KL gradient in the logits is q - p, so the size of the correction in
# coordinate i is the error in coordinate i. The squared error passes through the softmax Jacobian
# q_j(delta_ij - q_i) and carries a factor q_i in every term, so a token the model has learned to
# rule out gets almost no correction. Here we take both gradients by automatic differentiation.
# That checks the two expressions of Section 6.27 as well as evaluating them.
def logit_gradients(target, student):
    """The gradient of KL(p || q) and of half the squared error in the logits of q. Two [K]."""
    for_kl = student.log().clone().requires_grad_(True)  # Logits whose softmax is the student.
    divergence = (target * (target / torch.softmax(for_kl, -1)).log()).sum()
    kl_gradient, = torch.autograd.grad(divergence, for_kl)  # [K] equal to q - p.
    for_squared = student.log().clone().requires_grad_(True)  # A second graph on the same logits.
    squared = 0.5 * ((torch.softmax(for_squared, -1) - target) ** 2).sum()
    squared_gradient, = torch.autograd.grad(squared, for_squared)  # [K] with a q_i throughout.
    return kl_gradient, squared_gradient


def gradient_comparison(floor: float = 1e-6):
    """The two gradients at the chapter's student, and again once that student has saturated."""
    target = torch.tensor([0.06, 0.12, 0.76, 0.06])  # p, the composed target of Section 6.27.
    student = torch.tensor([0.10, 0.15, 0.65, 0.10])  # q, the student's long prediction.
    kl, squared = logit_gradients(target, student)
    ratio = (squared / kl).abs()  # Between a quarter and a half, coordinate by coordinate.
    rest = student[1:] * (1.0 - floor) / student[1:].sum()  # The other three keep their shares.
    saturated = torch.cat([torch.full_like(student[:1], floor), rest])  # [K], summing to one.
    kl_low, squared_low = logit_gradients(target, saturated)
    return {
        "KL gradient": [round(v, 5) for v in kl.tolist()],
        "squared-error gradient": [round(v, 5) for v in squared.tolist()],
        "squared over KL": [round(v, 4) for v in ratio.tolist()],
        "saturated student": [round(v, 6) for v in saturated.tolist()],
        "saturated KL gradient": float(kl_low[0]),
        # The surviving squared-error gradient is of order 1e-08, below what four printed
        # decimals show, so the report carries it in units of 1e-08.
        "saturated squared-error gradient, in 1e-08": 1e8 * float(squared_low[0]),
        "orders of magnitude between them": math.log10(
            abs(float(kl_low[0])) / abs(float(squared_low[0]))),
    }


# %% 06. The decoding clock
#
# A position decodes correctly when its data coordinate is the largest of the |V| coordinates,
# a shifted standard normal against the maximum of the rest. The reparameterized time tau is the
# fraction of the initial decoding error already removed, and the sampler runs on that clock.
def decoding_error(t: float, vocab_size: int, nodes: int = 4001) -> float:
    """P_e(t), the probability that one position decodes to the wrong token at time t."""
    if t >= 1.0:
        return 0.0
    shift = t / (1.0 - t)  # a = t/(1-t), the gap the data coordinate opens.
    z = torch.linspace(-12.0, 14.0, nodes, dtype=torch.float64)  # Quadrature grid over the draw.
    density = torch.exp(-0.5 * z * z) / math.sqrt(2.0 * math.pi)  # phi(z).
    cumulative = 0.5 * torch.erfc(-(z + shift) / math.sqrt(2.0))  # Phi(z + a).
    log_all = (vocab_size - 1) * cumulative.clamp(min=1e-300).log()  # All the others fall short.
    return 1.0 - float(torch.trapz(density * log_all.exp(), z))


def clock(vocab_size: int):
    """tau(t) = (P_e(0) - P_e(t))/P_e(0), as a callable, with the source error rate beside it."""
    start = decoding_error(0.0, vocab_size)
    return (lambda t: (start - decoding_error(t, vocab_size)) / start), start


def even_in_tau(vocab_size: int, pieces: int = 4):
    """Solve tau(t) = k/pieces by bisection and return the grid with its two endpoints."""
    tau, _ = clock(vocab_size)
    grid = [0.0]
    for k in range(1, pieces):
        low, high = 0.0, 1.0 - 1e-6
        for _ in range(60):  # Bisection; tau is increasing in t.
            middle = 0.5 * (low + high)
            low, high = (middle, high) if tau(middle) < k / pieces else (low, middle)
        grid.append(0.5 * (low + high))
    return grid + [1.0]


# %% 07. The model, a two-time denoiser with a softmax head
def predict(model, x, s, t, length, vocab):
    """delta_{s,t}(x) at every position, a point of the simplex produced by a softmax head."""
    logits = model(x.flatten(1), s, t).view(-1, length, vocab)  # [B, L, K] the network output.
    return torch.softmax(logits, dim=-1)  # [B, L, K] on the simplex at every position.


def step_map(x, prediction, s, t):
    """The convex form, the flow map written in denoiser coordinates."""
    keep, move = convex_weights(s, t)
    return keep[:, None, None] * x + move[:, None, None] * prediction  # [B, L, K].


def make_loss(model, data, length, vocab, batch_size, device):
    eye = torch.eye(vocab, device=device)  # The one-hot lift of the alphabet.

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1 = eye[data[index]]  # [B, L, K] the data, lifted.
        x0 = torch.randn_like(x1)  # [B, L, K] the Gaussian source.
        t_diag = torch.rand(batch_size, device=device)  # [B] a time for the diagonal term.
        xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1  # [B, L, K].
        diagonal = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log()
                     .gather(-1, data[index][..., None])).mean()  # Cross-entropy to the token.

        s = torch.rand(batch_size, device=device)  # [B] start time.
        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time.
        # The slice t = 1 is the one-step case the sampler reads, and a continuous draw never
        # lands on it, so we put a fixed share of the batch there on purpose (Section 6.26).
        endpoint = torch.rand(batch_size, device=device) < ENDPOINT_FRACTION  # [B] which rows.
        t = torch.where(endpoint, torch.ones_like(t), t)
        u = s + (t - s) * torch.rand(batch_size, device=device)  # [B] the split point.
        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1  # [B, L, K] the interpolant.
        with torch.no_grad():  # Both legs and the intermediate state are frozen targets.
            first = predict(model, xs, s, u, length, vocab)
            second = predict(model, step_map(xs, first, s, u), u, t, length, vocab)
            gamma = composition_weight(s, u, t)[:, None, None]
            target = gamma * first + (1.0 - gamma) * second  # [B, L, K], still on the simplex.
        student = predict(model, xs, s, t, length, vocab).clamp_min(1e-9)
        off_diagonal = (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean()
        return diagonal + off_diagonal

    return loss_fn


# %% 08. Sampling on a grid of times
@torch.no_grad()
def sample(model, count, length, vocab, grid, device):
    """Walk the grid with the convex form, and read the alphabet off the last prediction."""
    x = torch.randn(count, length, vocab, device=device)  # [N, L, K] the Gaussian source.
    for a, b in zip(grid[:-1], grid[1:]):
        s = torch.full((count,), a, device=device)
        t = torch.full((count,), b, device=device)
        x = step_map(x, predict(model, x, s, t, length, vocab), s, t)
    return x.argmax(-1)  # [N, L] token ids, read by an argmax at every position.


# %% 09. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
    data = tokens.to(device)  # [1024, L] the training sequences.

    report = scalar_cocycle()
    report.update(worked_kl())
    report.update(gradient_comparison())
    if not args.quiet:
        print(f"at the saturated student the KL gradient is "
              f"{report['saturated KL gradient']:.4f} and the squared-error gradient "
              f"{1e-8 * report['saturated squared-error gradient, in 1e-08']:.3e}")
    book_tau, book_start = clock(BOOK_VOCAB)
    for t in (0.60, 0.75, 0.80, 0.85):
        report[f"tau({t:.2f}) at |V|=30522"] = book_tau(t)
    book_grid = even_in_tau(BOOK_VOCAB)
    report["four steps on the clock"] = [round(t, 6) for t in book_grid]
    rounded = [round(t, 3) for t in book_grid]  # The chapter quotes the grid to three decimals.
    report["convex weights on that grid"] = [
        (round((1 - b) / (1 - a), 4), round((b - a) / (1 - a), 4))
        for a, b in zip(rounded[:-1], rounded[1:])]
    report["gamma on the last three"] = float(composition_weight(
        *(torch.tensor([v]) for v in rounded[1:4])))
    report["gamma over the long first leg"] = float(composition_weight(
        *(torch.tensor([v]) for v in (0.0, rounded[1], rounded[2]))))
    report["tau on a uniform grid"] = [round(book_tau(t), 5) for t in (0.25, 0.5, 0.75)]
    if not args.quiet:
        print(f"clock at |V|=30522: P_e(0) = {book_start:.6f}, four steps at "
              f"{[round(t, 3) for t in book_grid]}")

    model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining the two-time denoiser  parameters={count_parameters(model):,}")
    losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device),
                   steps=args.steps, lr=args.lr, quiet=args.quiet)

    dna_grid = even_in_tau(DNA_SIZE)  # The same clock, read at the four-letter alphabet.
    report["DNA clock grid"] = [round(t, 4) for t in dna_grid]
    samples: list[str] = []
    for name, grid in (("one step", [0.0, 1.0]), ("four, uniform", [0.0, 0.25, 0.5, 0.75, 1.0]),
                       ("four, on the clock", dna_grid)):
        torch.manual_seed(args.seed + 1)  # The same source noise for every schedule.
        drawn = sample(model, args.samples, args.length, DNA_SIZE, grid, device)
        report[f"{name}: GC fraction"] = gc_fraction(drawn)
        report[f"{name}: motif fraction"] = motif_fraction(drawn)
        report[f"{name}: token entropy"] = token_entropy(drawn, DNA_SIZE)
        report[f"{name}: unique fraction"] = unique_fraction(drawn)
        samples.append(f"{name:<20s} {decode_dna(drawn[:1])[0]}  GC {gc_fraction(drawn):.3f}")
    report["training GC fraction"] = gc_fraction(data)
    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

    print_report("Flow Map Language Models, Chapter 6", report, samples)
    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)


if __name__ == "__main__":
    main()