"""MeanFlow, the Eulerian identity written in the average-velocity parameterization.

If we substitute X_{s,t}(x) = x + (t-s) v_{s,t}(x) into the Eulerian identity, we are left with one
equation, v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with the correction term a total derivative
along the trajectory. One forward-mode pass along the tangent (b, 1, 0) in the arguments (x, s, t)
computes that derivative without forming a Jacobian, and the target uses the realized velocity
x_1 - x_0 of the interpolant in place of b, so a step costs one forward pass and one derivative.

Chapter 6, Sections 6.13, 6.15 and 6.18 of the notes. Trained on the four Gaussian blobs of
``cis6270.data.gaussian_mixture_2d``. Prints the identity worked by hand on dx/dt = x, where the
average velocity is 1.718x, the instantaneous velocity is x and the correction is 0.718x, and then
one-step, two-step and four-step samples from the trained map.

    python lecture_6/meanflow.py                 # 2000 steps, about twenty seconds
    python lecture_6/meanflow.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import gaussian_mixture_2d
from cis6270.metrics import mean_log_density, mode_report
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.

DIAGONAL_FRACTION = 0.25  # Share of each batch drawn with t = s, where the identity is flow matching.


# %% 02. The average velocity of the scalar flow, in closed form
#
# The exact map of dx/dt = x multiplies the state by e^{t-s}, so the displacement over the elapsed
# time is x (e^{t-s} - 1)/(t-s). On the diagonal the limit is x, or b(x,t).
def scalar_average_velocity(x, s, t):
    """v_{s,t}(x) = x (e^{t-s} - 1)/(t - s) on the flow dx/dt = x."""
    elapsed = t - s
    return x * torch.expm1(elapsed) / elapsed  # expm1 keeps the limit accurate near the diagonal.


# %% 03. The identity, checked on the scalar flow by hand
#
# v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with d/ds = partial_s + b(x,s) . grad_x. At s = 0,
# t = 1 the partial derivative is -x, the spatial gradient is e - 1, and the transport term is
# therefore (e-1)x, so the correction is x(e-2) = 0.718x.
def scalar_check(x_value: float = 1.0):
    """Return every term of the identity at s = 0, t = 1, each one computed here."""
    x = torch.tensor(x_value)
    s, t = torch.tensor(0.0), torch.tensor(1.0)
    average = scalar_average_velocity(x, s, t)  # v_{0,1}(x) = x(e-1).
    instant = x  # b(x, 0) = x on this flow.
    partial_s = torch.func.jvp(  # Derivative in the start time alone, with x and t held fixed.
        lambda ss: scalar_average_velocity(x, ss, t), (s,), (torch.ones_like(s),))[1]
    gradient_x = torch.func.jvp(  # Derivative in the state alone, the other half of d/ds.
        lambda xx: scalar_average_velocity(xx, s, t), (x,), (torch.ones_like(x),))[1]
    transport = instant * gradient_x  # b(x,s) . grad_x v, the second piece of the total derivative.
    total = partial_s + transport  # d/ds v_{s,t}(x) = x(e - 2).
    return {
        "v_0,1(1)": float(average),
        "b(1, 0)": float(instant),
        "partial_s v": float(partial_s),
        "grad_x v_0,1": float(gradient_x),
        "d/ds v_0,1": float(total),
        "identity gap": float((average - (instant + (t - s) * total)).abs()),
        "one map evaluation": float(x + (t - s) * average),
        "one Euler step": float(x + (t - s) * instant),
        "e": math.e,
    }


# %% 04. The Gaussian flow, where the correction vanishes
#
# A linear interpolant between paired Gaussians produces straight trajectories, and on a straight
# trajectory the average velocity equals the initial velocity, so the total derivative is zero.
def gaussian_check(x_value: float = 0.7):
    """v_{0,t}(x) = 2 + 2x for every end time, and b(x,0) is the same number."""
    x = torch.tensor(x_value)
    average = torch.stack([  # Three end times, all giving the same average velocity.
        (2.0 * t + (1.0 + 2.0 * t) * x - x) / t for t in (0.25, 0.5, 1.0)])
    instant = 2.0 + 2.0 * (x - 0.0) / 1.0  # b(x, 0) on the Gaussian flow.
    return {"v_0,t(0.7) spread": float(average.max() - average.min()),
            "v_0,1(0.7)": float(average[-1]), "b(0.7, 0)": float(instant)}


# %% 05. The MeanFlow target, by one forward-mode pass
#
# The tangent (x_1 - x_0, 1, 0) in the arguments (x, s, t) contracts the full Jacobian of v into
# the total derivative along the trajectory: the realized velocity in the state slot, one in the
# start-time slot, and zero in the end-time slot because the end time is held fixed.
def mean_flow_target(model, xs, s, t, realized):
    """b(x,s) + (t-s) d/ds v_{s,t}(x), detached, the right-hand side of the identity."""
    f = lambda xx, ss: model(xx, ss, t)
    _, total = torch.func.jvp(f, (xs, s), (realized, torch.ones_like(s)))  # [B, d] d/ds v.
    return (realized + (t - s)[:, None] * total).detach()  # [B, d] the target, held fixed.


def make_loss(model, data, batch_size, device):
    """Draw a coupling, draw times from the triangle, and regress v onto its own identity."""

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1 = data[index]  # [B, 2] the data draw.
        x0 = torch.randn_like(x1)  # [B, 2] the source draw.
        s = torch.rand(batch_size, device=device)  # [B] start time.
        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time, s <= t <= 1.
        diagonal = torch.rand(batch_size, device=device) < DIAGONAL_FRACTION  # [B] which rows anchor.
        t = torch.where(diagonal, s, t)  # On those rows the correction term has a factor of zero.
        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant I_s.
        realized = x1 - x0  # [B, 2] the derivative of the interpolant, available in closed form.
        target = mean_flow_target(model, xs, s, t, realized)
        return ((model(xs, s, t) - target) ** 2).mean()

    return loss_fn


# %% 06. Sampling, at any step count the reader chooses
@torch.no_grad()
def sample(model, x0, steps: int):
    """Cross [0, 1] in equal pieces, one evaluation of the map per piece."""
    x = x0
    grid = torch.linspace(0.0, 1.0, steps + 1, device=x0.device)  # [steps+1] the time grid.
    for s, t in zip(grid[:-1], grid[1:]):
        x = model.flow_map(x, s.expand(len(x)), t.expand(len(x)))  # X_{s,t}(x).
    return x  # [N, 2] generated points.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=2000, batch_size=128, samples=512)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device)  # [2048, 2].

    report = scalar_check()
    report.update(gaussian_check())
    if not args.quiet:
        print(f"on dx/dt = x at s=0, t=1: v = {report['v_0,1(1)']:.6f} = b {report['b(1, 0)']:.6f}"
              f" + correction {report['d/ds v_0,1']:.6f}")
        print(f"one map evaluation {report['one map evaluation']:.6f} against one Euler step "
              f"{report['one Euler step']:.6f}, exact e = {math.e:.6f}")

    model = TwoTimeMLP(dim=2, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining MeanFlow  parameters={count_parameters(model):,}")
    losses = train(model, make_loss(model, data, args.batch_size, device),
                   steps=args.steps, lr=args.lr, quiet=args.quiet)

    torch.manual_seed(args.seed + 1)
    x0 = torch.randn(args.samples, 2, device=device)  # [N, 2] the source points every run shares.
    samples: list[str] = []
    for steps in (1, 2, 4):
        points = sample(model, x0, steps)
        stats = mode_report(points)
        report[f"{steps}-step nearest distance"] = stats["nearest_distance"]
        report[f"{steps}-step mode entropy"] = stats["mode_entropy"]
        report[f"{steps}-step log density"] = mean_log_density(points.cpu())
        samples.append(f"{steps} step(s): nearest {stats['nearest_distance']:.3f}  "
                       f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
    one, four = sample(model, x0, 1), sample(model, x0, 4)
    report["one step against four"] = float((one - four).norm(dim=1).mean())
    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

    print_report("MeanFlow, Chapter 6", report, samples)
    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)


if __name__ == "__main__":
    main()