"""Transport as motion in time: the kinetic action, straight lines, and displacement interpolation.

A plan records which source point ends where and nothing about the trajectory between them. The
Benamou-Brenier formulation replaces the plan by a density and a velocity field and measures the
transport by the energy the particles expend. Its minimizer moves every particle in a straight line
at constant speed, with the endpoints paired by the optimal plan.

Chapter 7, Sections 7.7 to 7.7.5 of the notes. Runs on the one-particle example of Section 7.7.1
and on the Gaussian pair N(0,1) to N(3,4), both in closed form. Prints the action of a steady and a
dawdling particle, the two sides of the averaged bound, the interpolated density against the mixture
of the endpoints, and the velocity field with its action density.

    python lecture_7/displacement.py              # 1024 times on the grid
    python lecture_7/displacement.py --steps 20   # a coarse time grid, for a first look
"""

# %% 01. Imports and configuration
from __future__ import annotations

import math

import torch

from cis6270.runner import common_parser, print_report, save_run, set_seed

DTYPE = torch.float64  # Double precision, so the closed forms match the chapter digit for digit.

MEAN_0, STD_0 = 0.0, 1.0  # The source N(0, 1).
MEAN_1, STD_1 = 3.0, 2.0  # The target N(3, 4), with standard deviation 2.


# %% 02. One particle
#
# The squared displacement of a particle is at most the integral of its squared speed, by
# Cauchy-Schwarz against the constant function 1 on an interval of length 1, with equality when the
# velocity is constant in time. Section 7.7.1.


def action_of_a_path(positions: torch.Tensor, times: torch.Tensor) -> float:
    """The integral of the squared speed along a trajectory sampled at the given times."""
    velocity = (positions[1:] - positions[:-1]) / (times[1:] - times[:-1])  # [n-1] piecewise speed.
    return float((velocity ** 2 * (times[1:] - times[:-1])).sum())  # Integral of ||x_dot||^2.


def steady_particle(times: torch.Tensor) -> torch.Tensor:
    """A particle covering the distance 2 at the constant speed 2, so its action is 4."""
    return 2.0 * times  # x(t) = 2t.


def dawdling_particle(times: torch.Tensor) -> torch.Tensor:
    """A particle that waits until t = 1/2 and then covers the same distance at speed 4."""
    return torch.clamp(4.0 * (times - 0.5), min=0.0)  # Zero on [0, 1/2], then 4(t - 1/2).


# %% 03. Averaging the bound over the particles
#
# The joint law of the two endpoints of a particle is a coupling of mu and nu, so the action is at
# least the average squared displacement under that coupling, and that is at least W_2^2.
# Section 7.7.2.


def average_squared_displacement(x: torch.Tensor, y: torch.Tensor, pi: torch.Tensor) -> float:
    """The average of ||x_1 - x_0||^2 over endpoint pairs drawn from the coupling pi."""
    return float((pi * (x[:, None] - y[None, :]) ** 2).sum())  # Sum over routes of mass times cost.


# %% 04. Straight lines
#
# We draw the endpoint pair from the optimal coupling and move along the segment between them.
# The velocity is constant along it. Cauchy-Schwarz is then tight for every particle, and the
# coupling is optimal, so the second inequality is tight as well. Section 7.7.3.


def straight_path(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """x(t) = (1 - t) x_0 + t x_1, with the constant velocity x_1 - x_0."""
    return (1.0 - t) * x0 + t * x1  # Same shape as x0 and x1 broadcast against t.


# %% 05. Displacement interpolation
#
# We read off the law of the particles at an intermediate time and get the displacement
# interpolation. Under the optimal map X_1 = 3 + 2 X_0 of the Gaussian example, X_t collects into
# 3t + (1 + t) X_0, so rho_t is Gaussian with mean 3t and standard deviation 1 + t. Section 7.7.4.


def displacement_law(t: float) -> tuple[float, float]:
    """The mean and standard deviation of rho_t for the Gaussian pair, both linear in t."""
    return 3.0 * t, 1.0 + t  # Mean slides from 0 to 3, standard deviation from 1 to 2.


def normal_density(x: torch.Tensor, mean: float, std: float) -> torch.Tensor:
    """The density of N(mean, std^2) evaluated at every entry of x."""
    return torch.exp(-0.5 * ((x - mean) / std) ** 2) / (std * math.sqrt(2.0 * math.pi))


def mixture_density(x: torch.Tensor) -> torch.Tensor:
    """Half the source density plus half the target density, the other path between the two ends."""
    return 0.5 * normal_density(x, MEAN_0, STD_0) + 0.5 * normal_density(x, MEAN_1, STD_1)


def gaussian_w2(m0: float, s0: float, m1: float, s1: float) -> float:
    """W_2 between two Gaussians on the line, the closed form of Section 7.6."""
    return math.sqrt((m0 - m1) ** 2 + (s0 - s1) ** 2)


# %% 06. The velocity field behind the displacement path
#
# Each particle obeys X_t = 3t + (1 + t) X_0 with the constant velocity 3 + X_0. We solve the
# trajectory for X_0 and substitute. The velocity is then a function of the current position
# alone. Section 7.7.5.


def velocity_field(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """v_t(x) = 3 + (x - 3t) / (1 + t), the mean speed plus a stretch about the mean."""
    return 3.0 + (x - 3.0 * t) / (1.0 + t)  # Broadcasts over a grid of states and times.


def continuity_residual(x: float, t: float) -> float:
    """The left side of d_t rho_t + d_x (rho_t v_t) = 0, computed by automatic differentiation."""
    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
    ts = torch.tensor(t, dtype=DTYPE, requires_grad=True)
    rho = normal_density(xs, 3.0 * ts, 1.0 + ts)  # rho_t written with the two linear schedules.
    flux = rho * velocity_field(xs, ts)  # The probability flux rho_t v_t.
    d_rho_dt = torch.autograd.grad(rho, ts, retain_graph=True)[0]  # Keep the graph for the flux.
    d_flux_dx = torch.autograd.grad(flux, xs)[0]
    return float((d_rho_dt + d_flux_dx).detach())


def action_density(t: float, levels: int) -> float:
    """The integral of rho_t(x) v_t(x)^2 over x, or E[(3 + X_0)^2] at every time."""
    mean, std = displacement_law(t)
    q = (torch.arange(levels, dtype=DTYPE) + 0.5) / levels  # Quantile levels of rho_t.
    x = mean + std * torch.special.ndtri(q)  # [levels] a stratified sample of rho_t.
    speed = velocity_field(x, torch.tensor(t, dtype=DTYPE))  # [levels].
    return float((speed ** 2).mean())  # The average of the squared speed under rho_t.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1024)  # --steps is the number of times on the grid.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet
    n = max(args.steps, 2)
    times = torch.linspace(0.0, 1.0, n + 1, dtype=DTYPE)  # [n+1] the time grid.

    steady = action_of_a_path(steady_particle(times), times)
    dawdle = action_of_a_path(dawdling_particle(times), times)
    if loud:
        print(f"One particle from 0 to 2 on a grid of {n} steps")
        print(f"  constant speed 2: action {steady:.3f}, squared displacement {2.0 ** 2:.0f}")
        print(f"  waiting then speed 4: action {dawdle:.3f}, the same two endpoints")

    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Four points, mass 1/2 each.
    y = torch.tensor([1.0, 3.0], dtype=DTYPE)
    ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=DTYPE)
    crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=DTYPE)
    independent = torch.full((2, 2), 0.25, dtype=DTYPE)
    w22_points = average_squared_displacement(x, y, ordered)  # The optimal coupling on the line.
    if loud:
        print(f"\nAveraging over particles, four points with mass 1/2")
        for name, pi in (("ordered", ordered), ("independent", independent), ("crossed", crossed)):
            print(f"  {name:<12s} E||x_1 - x_0||^2 = "
                  f"{average_squared_displacement(x, y, pi):.1f}  >= W_2^2 = {w22_points:.1f}")

    # Straight lines under the optimal coupling, with the endpoints of the Gaussian example.
    q = (torch.arange(n, dtype=DTYPE) + 0.5) / n
    x0 = MEAN_0 + STD_0 * torch.special.ndtri(q)  # [n] a stratified sample of N(0, 1).
    x1 = MEAN_1 + (STD_1 / STD_0) * (x0 - MEAN_0)  # The Brenier map 3 + 2x, the optimal pairing.
    paths = straight_path(x0[None, :], x1[None, :], times[:, None])  # [n+1, n] every trajectory.
    speeds = (paths[1:] - paths[:-1]) / (times[1:] - times[:-1])[:, None]  # [n, n] piecewise speed.
    action = float((speeds ** 2 * (times[1:] - times[:-1])[:, None]).sum(dim=0).mean())
    squared_displacement = float(((x1 - x0) ** 2).mean())
    if loud:
        print("\nStraight lines under the optimal coupling")
        print(f"  action averaged over the particles {action:.4f}  "
              f"average squared displacement {squared_displacement:.4f}")

    mean_half, std_half = displacement_law(0.5)
    grid = torch.linspace(-6.0, 9.0, 15001, dtype=DTYPE)  # A dense grid for the two densities.
    mixture = mixture_density(grid)
    peak = float(grid[int(mixture.argmax())])
    interpolated = normal_density(grid, mean_half, std_half)
    separation = float(0.5 * (interpolated - mixture).abs().sum() * (grid[1] - grid[0]))
    half_distance = gaussian_w2(MEAN_0, STD_0, mean_half, std_half)
    whole_distance = gaussian_w2(MEAN_0, STD_0, MEAN_1, STD_1)
    if loud:
        print(f"\nAt t = 1/2 the displacement interpolation is N({mean_half:.1f}, "
              f"{std_half ** 2:.2f}) with standard deviation {std_half:.1f}")
        print(f"  the mixture of the two ends peaks at x = {peak:.2f} and is not a Gaussian")
        print(f"  total variation between the two candidate densities {separation:.3f}")
        print(f"  W_2(rho_0, rho_1/2) = {half_distance:.3f} = "
              f"{half_distance / whole_distance:.1f} x W_2(mu, nu) = {whole_distance:.3f}")

    mean_speed = float(velocity_field(torch.tensor(3.0 * 0.5, dtype=DTYPE),
                                      torch.tensor(0.5, dtype=DTYPE)))  # A particle at the mean.
    above_speed = float(velocity_field(torch.tensor(3.0 * 0.5 + 1.5, dtype=DTYPE),
                                       torch.tensor(0.5, dtype=DTYPE)))  # One standard deviation up.
    residual = max(abs(continuity_residual(xx, tt))
                   for xx in (-1.0, 0.0, 1.5, 4.0) for tt in (0.1, 0.5, 0.9))
    densities = [action_density(float(t), 2048) for t in torch.linspace(0.0, 1.0, 5, dtype=DTYPE)]
    landing = x0 + 1.0 * velocity_field(x0, torch.zeros((), dtype=DTYPE))  # One Euler step of size 1.
    one_step = float((landing - x1).abs().max())  # The step lands on the Brenier map 3 + 2 x_0.
    if loud:
        print(f"\nv_t(x) = 3 + (x - 3t) / (1 + t)")
        print(f"  at the mean, speed {mean_speed:.0f}; one standard deviation above, "
              f"speed {above_speed:.0f}")
        print(f"  action density at five times {[round(d, 4) for d in densities]}")
        print(f"  largest continuity residual on the test states {residual:.2e}")

    report = {
        "action of the steady particle": steady,
        "action of the dawdling particle": dawdle,
        "squared displacement of both": 4.0,
        "E||x1-x0||^2, ordered": average_squared_displacement(x, y, ordered),
        "E||x1-x0||^2, independent": average_squared_displacement(x, y, independent),
        "E||x1-x0||^2, crossed": average_squared_displacement(x, y, crossed),
        "action of the straight paths": action,
        "W_2^2 of the Gaussian pair": whole_distance ** 2,
        "rho_1/2 mean, standard deviation": f"{mean_half:.1f}, {std_half:.1f}",
        "peak of the mixture": peak,
        "total variation of the two paths": separation,
        "W_2(rho_0, rho_1/2)": half_distance,
        "W_2(mu, nu)": whole_distance,
        "speed at the mean": mean_speed,
        "speed one deviation above": above_speed,
        "action density, every time": max(densities),
        "continuity equation residual": residual,
        "one Euler step of size 1 minus T": one_step,
    }
    print_report("Transport as motion, Chapter 7 Sections 7.7 to 7.7.5", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()