"""Categorical Flow Maps, the Lagrangian identity on the simplex and the bound that replaces it.

A softmax head makes the two-time prediction a distribution at every pair of times, and the
Lagrangian identity rewritten in those coordinates reads pi_{s,t} = pi_{t,t}(X_{s,t}) minus a
coefficient times the end-time derivative of the prediction. The derivative has no sign
constraint, so the residual splits into a teacher gap, a difference of two distributions, and a
drift, the scaled derivative. A parallelogram bound and Pinsker's inequality turn the squared
residual into four KL divergences plus two squared drifts, and at the exact map the two pieces
are antiparallel, so the bound is strictly positive where the residual is zero.

Chapter 6, Section 6.28 of the notes. Trained on the four-token phrases of
``cis6270.data.phrases``. Prints the chapter's worked cancellation on three tokens, where the gap
is the exact negative of the drift and the bound reads 0.199 around a residual of zero, and then
two models trained on the residual and on its bound.

    python lecture_6/categorical_map.py                 # two models, 1000 steps each
    python lecture_6/categorical_map.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import torch

from cis6270.data import phrases
from cis6270.metrics import 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_PRIOR = (0.5, 0.3, 0.2)  # The three data probabilities of the chapter's worked example.
BOOK_SOURCE = (-0.5739, -0.8344, -0.0584)  # The source draw I_0 the chapter starts from.
BOOK_STATES = 3000  # States the chapter averages the two terms of the bound over.


# %% 02. The exact model on three tokens
#
# With a Gaussian source and the linear interpolant, the posterior over the data token is a
# softmax whose only state dependence is the coordinate of that token. The velocity follows from
# D_t(x) = x + (1-t) b(x,t), and the flow integrates to twelve digits with Runge-Kutta.
def exact_denoiser(x, t, prior):
    """pi_{t,t}(x), the posterior over the data token given the state. [B, K]."""
    return torch.softmax(prior.log() + t * x / (1.0 - t) ** 2, dim=-1)


def exact_velocity(x, t, prior):
    """b(x,t) = (pi_{t,t}(x) - x)/(1 - t) for the linear interpolant. [B, K]."""
    return (exact_denoiser(x, t, prior) - x) / (1.0 - t)


def exact_flow(x, s, t, prior, n: int = 2000):
    """X_{s,t}(x) by fourth-order Runge-Kutta on the exact velocity."""
    h = (t - s) / n
    for k in range(n):
        u = s + k * h
        k1 = exact_velocity(x, u, prior)
        k2 = exact_velocity(x + 0.5 * h * k1, u + 0.5 * h, prior)
        k3 = exact_velocity(x + 0.5 * h * k2, u + 0.5 * h, prior)
        k4 = exact_velocity(x + h * k3, u + h, prior)
        x = x + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
    return x


def exact_two_time(x, s, t, prior):
    """pi_{s,t}(x) read off the convex parameterization, [(1-s) X_{s,t}(x) - (1-t) x]/(t-s)."""
    return ((1.0 - s) * exact_flow(x, s, t, prior) - (1.0 - t) * x) / (t - s)


# %% 03. The coefficient, the gap and the drift
#
# gamma_{s,t} = (t-s)(1-t)/(1-s) vanishes on the diagonal and at t = 1 and never exceeds 1/4 on
# the square. The residual scaled by the remaining time splits as (1-t) r = a + b, with a the gap
# between the student's prediction and the prediction at the transported state.
def gamma_coefficient(s, t):
    """gamma_{s,t} = (t-s)(1-t)/(1-s)."""
    return (t - s) * (1.0 - t) / (1.0 - s).clamp(min=1e-6)


def end_time_derivative(x, s, t, prior, n: int = 2000):
    """d pi_{s,t}/dt from the integral representation, (1-s)/((t-s)(1-t)) times the bracket."""
    transported = exact_flow(x, s, t, prior, n)
    prediction = ((1.0 - s) * transported - (1.0 - t) * x) / (t - s)
    at_output = exact_denoiser(transported, t, prior)
    return (1.0 - s) / ((t - s) * (1.0 - t)) * (at_output - prediction), prediction, at_output


# %% 04. The worked cancellation of Section 6.28
#
# The derivative above comes from the closed form of Section 6.26, and that form already carries
# the cancellation, so the central difference of pi_{s,t} in t below checks the two pieces
# against each other without assuming the identity.
def worked_cancellation(s: float = 0.2, t: float = 0.6):
    """The three-token example, with every number integrated or differentiated here."""
    prior = torch.tensor(BOOK_PRIOR, dtype=torch.float64)
    source = torch.tensor(BOOK_SOURCE, dtype=torch.float64)
    x = (1.0 - s) * source + s * torch.eye(3, dtype=torch.float64)[0]  # Data token is the first.
    derivative, prediction, at_output = end_time_derivative(x, s, t, prior)
    gamma = (t - s) * (1.0 - t) / (1.0 - s)
    gap = prediction - at_output  # a_{s,t}, a difference of two simplex points.
    drift = gamma * derivative  # b_{s,t}, the scaled derivative of the prediction.
    divergence = (at_output * (at_output / prediction).log()).sum()  # KL(teacher || student).
    difference = (exact_two_time(x, s, t + 1e-5, prior)
                  - exact_two_time(x, s, t - 1e-5, prior)) / 2e-5  # A central check.
    return {
        "state x": [round(v, 4) for v in x.tolist()],
        "X_s,t(x)": [round(v, 4) for v in exact_flow(x, s, t, prior).tolist()],
        "pi_s,t(x)": [round(v, 4) for v in prediction.tolist()],
        "pi_t,t(X_s,t(x))": [round(v, 4) for v in at_output.tolist()],
        "gamma_s,t": gamma,
        "d pi / dt": [round(v, 4) for v in derivative.tolist()],
        "central difference": [round(v, 4) for v in difference.tolist()],
        "gap a": [round(v, 3) for v in gap.tolist()],
        "drift b": [round(v, 3) for v in drift.tolist()],
        "residual |a + b|^2": float(((gap + drift) ** 2).sum()),
        "4 KL": 4.0 * float(divergence),
        "2 |b|^2": 2.0 * float((drift ** 2).sum()),
        "the bound": 4.0 * float(divergence) + 2.0 * float((drift ** 2).sum()),
        "teacher entropy": -float((at_output * at_output.log()).sum()),
        "cross-entropy": -float((at_output * prediction.log()).sum()),
    }


def averaged_bound(count: int = BOOK_STATES, s: float = 0.2, t: float = 0.6, seed: int = 0,
                   n: int = 400):
    """The same two terms averaged over states drawn from the interpolant at the same times."""
    prior = torch.tensor(BOOK_PRIOR, dtype=torch.float64)
    generator = torch.Generator().manual_seed(seed)
    which = torch.multinomial(prior.expand(count, 3), 1, generator=generator).squeeze(1)
    x1 = torch.eye(3, dtype=torch.float64)[which]  # [N, 3] the lifted data token.
    x0 = torch.randn(count, 3, generator=generator, dtype=torch.float64)  # [N, 3] the source.
    x = (1.0 - s) * x0 + s * x1  # [N, 3] the interpolant at the start time.
    derivative, prediction, at_output = end_time_derivative(x, s, t, prior, n)
    drift = ((t - s) * (1.0 - t) / (1.0 - s)) * derivative
    divergence = (at_output * (at_output / prediction).log()).sum(-1).mean()
    return {"averaged 4 KL": 4.0 * float(divergence),
            "averaged 2 |b|^2": 2.0 * float((drift ** 2).sum(-1).mean()),
            "averaged residual": float(((prediction - at_output + drift) ** 2).sum(-1).mean())}


# %% 05. The student, and the two losses it can be trained on
def predict(model, x, s, t, length, vocab):
    """pi_{s,t}(x) at every position, a softmax head over the vocabulary."""
    return torch.softmax(model(x.flatten(1), s, t).view(-1, length, vocab), dim=-1)


def pieces(model, xs, s, t, length, vocab):
    """The gap a and the drift b of the split, with the teacher detached."""
    f = lambda tt: predict(model, xs, s, tt, length, vocab)
    prediction, derivative = torch.func.jvp(f, (t,), (torch.ones_like(t),))  # One forward mode pass.
    keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
    transported = keep[:, None, None] * xs + move[:, None, None] * prediction  # X_{s,t}(x).
    teacher = predict(model, transported.detach(), t, t, length, vocab).detach()  # pi_{t,t}.
    drift = gamma_coefficient(s, t)[:, None, None] * derivative
    return prediction - teacher, drift, teacher, prediction


def make_loss(model, data, length, vocab, batch_size, device, bounded: bool):
    eye = torch.eye(vocab, device=device)

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

        s = 0.9 * torch.rand(batch_size, device=device)  # [B] start time, kept off t = 1.
        t = s + (0.95 - s) * torch.rand(batch_size, device=device)  # [B] end time below one.
        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1
        gap, drift, teacher, prediction = pieces(model, xs, s, t, length, vocab)
        if bounded:  # Four KL divergences plus two squared drifts, the bound of Section 6.28.
            divergence = (teacher * (teacher.clamp_min(1e-9)
                                     / prediction.clamp_min(1e-9)).log()).sum(-1).mean()
            return anchor + 4.0 * divergence + 2.0 * (drift ** 2).sum(-1).mean()
        return anchor + ((gap + drift) ** 2).sum(-1).mean()  # The residual itself.

    return loss_fn


# %% 06. Sampling, and the test for a valid phrase
@torch.no_grad()
def sample(model, count, length, vocab, steps, device):
    """Walk a uniform grid with the convex form and read the phrase off the last prediction."""
    x = torch.randn(count, length, vocab, device=device)
    grid = torch.linspace(0.0, 1.0, steps + 1)
    for a, b in zip(grid[:-1], grid[1:]):
        s, t = a.expand(count).to(device), b.expand(count).to(device)
        prediction = predict(model, x, s, t, length, vocab)
        keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
        x = keep[:, None, None] * x + move[:, None, None] * prediction
    return x.argmax(-1)  # [N, L].


def valid_fraction(drawn, group: int = 4):
    """A phrase is in the language when position i holds a token from the i-th quarter."""
    slots = torch.arange(drawn.shape[1], device=drawn.device)  # [L].
    return float(((drawn // group) == slots[None, :]).all(dim=1).to(torch.get_default_dtype()).mean())


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1000, batch_size=128, samples=256, sample_steps=4)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    tokens, vocabulary = phrases(num_phrases=1024, seed=args.seed)
    data, length, vocab = tokens.to(device), tokens.shape[1], len(vocabulary)

    report = worked_cancellation()
    report.update(averaged_bound(seed=args.seed))  # The chapter's three thousand states.
    if not args.quiet:
        print(f"gap {report['gap a']} and drift {report['drift b']} cancel to "
              f"{report['residual |a + b|^2']:.2e}, bound {report['the bound']:.3f}")

    samples: list[str] = []
    for name, bounded in (("the residual", False), ("the bound", True)):
        set_seed(args.seed)  # Both students start from the same weights and see the same batches.
        model = TwoTimeMLP(dim=length * vocab, width=args.width).to(device)
        if not args.quiet:
            print(f"\ntraining on {name}  parameters={count_parameters(model):,}")
        losses = train(model, make_loss(model, data, length, vocab, args.batch_size, device,
                                        bounded), steps=args.steps, lr=args.lr, quiet=args.quiet)
        for steps in (1, args.sample_steps):
            torch.manual_seed(args.seed + 1)
            drawn = sample(model, args.samples, length, vocab, steps, device)
            report[f"{name}, {steps} step(s): valid"] = valid_fraction(drawn)
            report[f"{name}, {steps} step(s): distinct"] = unique_fraction(drawn)
            report[f"{name}, {steps} step(s): entropy"] = token_entropy(drawn, vocab)
        report[f"{name}: final loss"] = sum(losses[-20:]) / len(losses[-20:])
        drawn = sample(model, 4, length, vocab, args.sample_steps, device)
        samples.append(f"{name:<14s} " + " | ".join(
            " ".join(vocabulary[int(i)] for i in row) for row in drawn[:2]))

    print_report("Categorical Flow Maps, Chapter 6", report, samples)
    save_run(args.out, config=vars(args), report=report, samples=samples)


if __name__ == "__main__":
    main()