"""Optimal transport, from a pairing of four points to the Gaussian closed form.

Two distributions fix how often each point appears and leave the assignment between them free. This
file prices every assignment with the squared cost, writes a pairing as a matrix, relaxes Monge's
map to Kantorovich's coupling, certifies the optimum with prices, and ends with the two Gaussians
whose optimal map is an affine function.

Chapter 7, Sections 7.1 to 7.6 of the notes. Runs on the chapter's own points, with no training and
no data. Prints the cost of each pairing, the plan and the prices of the two-by-two example, the
sorted pairing on three points, and the transport of N(0,1) to N(3,4).

    python lecture_7/optimal_transport.py              # every worked number of Sections 7.1 to 7.6
    python lecture_7/optimal_transport.py --steps 20   # a coarse quantile grid, for a first look
"""

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

import itertools
import math

import torch

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

DTYPE = torch.float64  # Double precision, because the chapter quotes six figures.


# %% 02. The cost of moving one unit of mass
#
# We price one unit carried from x to y by the squared Euclidean distance, so a move twice as long
# costs four times as much. Section 7.1.1.


def cost_matrix(x: torch.Tensor, y: torch.Tensor, p: float = 2.0) -> torch.Tensor:
    """C_ij = |x_i - y_j|^p on the line. x: [n], y: [m] -> [n, m]."""
    return (x[:, None] - y[None, :]).abs() ** p  # [n, m] one price per route.


def plan_cost(C: torch.Tensor, pi: torch.Tensor) -> float:
    """The Frobenius inner product <C, pi>, the mass on each route times the price of that route."""
    return float((C * pi).sum())


# %% 03. Four points and two pairings
#
# Sources at 0 and 2, targets at 1 and 3, each point with mass 1/2. Section 7.1.2 counts the two
# arrows of a pairing without weighting them by mass, and Section 7.1.3 weights them.


def two_pairings() -> dict:
    """The ordered and crossed pairings of the four points, by arrow count and by mass."""
    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Source points.
    y = torch.tensor([1.0, 3.0], dtype=DTYPE)  # Target points.
    C = cost_matrix(x, y)  # [[1, 9], [1, 1]] the four route prices.
    ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=DTYPE)  # 0 -> 1 and 2 -> 3.
    crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=DTYPE)  # 0 -> 3 and 2 -> 1.
    independent = torch.full((2, 2), 0.25, dtype=DTYPE)  # a_i b_j with a = b = (1/2, 1/2).
    return {
        "C": C,
        "arrows_ordered": float(C[0, 0] + C[1, 1]),  # 1 + 1, the two arrow lengths squared.
        "arrows_crossed": float(C[0, 1] + C[1, 0]),  # 9 + 1.
        "mass_ordered": plan_cost(C, ordered),  # Each arrow weighted by the 1/2 it carries.
        "mass_crossed": plan_cost(C, crossed),
        "mass_independent": plan_cost(C, independent),  # 1/4 on every route.
    }


# %% 04. A pairing written as a matrix
#
# A coupling is an n by m matrix with row sums a and column sums b, so pi 1_m = a and pi^T 1_n = b.
# The independent coupling a_i b_j is always in that set. Section 7.1.3.


def is_coupling(pi: torch.Tensor, a: torch.Tensor, b: torch.Tensor, tol: float = 1e-9) -> bool:
    """True when the row sums equal a, the column sums equal b, and no entry is negative."""
    return bool(
        (pi >= -tol).all()
        and (pi.sum(dim=1) - a).abs().max() < tol  # pi 1_m = a.
        and (pi.sum(dim=0) - b).abs().max() < tol  # pi^T 1_n = b.
    )


# %% 05. Monge's formulation, and the source that has to be split
#
# A map sends each source point to one destination, so for n equal masses it is a permutation and
# we minimize over n! of them. No map carries delta_0 to a law with mass at two places, because the
# pushforward of a point mass is a point mass. Sections 7.2 and 7.2.1.


def monge_permutations(x: torch.Tensor, y: torch.Tensor) -> list[tuple[float, tuple[int, ...]]]:
    """Every permutation of n equal point masses and its average cost, cheapest first."""
    n = len(x)
    C = cost_matrix(x, y)  # [n, n].
    costs = [(float(sum(C[i, p[i]] for i in range(n)) / n), p)
             for p in itertools.permutations(range(n))]  # Mass 1/n on each of the n routes.
    return sorted(costs)


def pushforward_of_a_point_mass(destination: float) -> dict[float, float]:
    """T_# delta_0 is the point mass at T(0), whatever the map T does elsewhere."""
    return {destination: 1.0}  # One destination, all of the mass, so nu is out of reach.


NU_ATOMS = {1.0: 0.5, 2.0: 0.5}  # nu = (1/2) delta_1 + (1/2) delta_2 of Section 7.2.1.


# %% 06. Kantorovich's relaxation, solved by hand on the two-by-two
#
# The feasible set of the chapter's example is the one-parameter family pi(r), and the cost along it
# is 5.8 - 8r. That cost decreases in r, so the optimum sits at the largest feasible r.
# Sections 7.3 to 7.3.2.

SOURCE_X = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Source points of the running example.
SOURCE_A = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Their masses.
TARGET_Y = torch.tensor([1.0, 3.0], dtype=DTYPE)  # Target points.
TARGET_B = torch.tensor([0.3, 0.7], dtype=DTYPE)  # Their required masses.


def feasible_plan(r: float) -> torch.Tensor:
    """The coupling with mass r from 0 to 1, the one degree of freedom the margins leave."""
    return torch.tensor([[r, 0.6 - r], [0.3 - r, 0.1 + r]], dtype=DTYPE)  # [2, 2].


def feasible_interval() -> tuple[float, float]:
    """We keep the four entries of pi(r) nonnegative, and the binding pair gives 0 <= r <= 0.3."""
    lower = max(0.0, -0.1)  # From r >= 0 and 0.1 + r >= 0.
    upper = min(0.6, 0.3)  # From 0.6 - r >= 0 and 0.3 - r >= 0.
    return lower, upper


def cost_along_the_segment(r: float, C: torch.Tensor) -> float:
    """<C, pi(r)> as a function of r alone, or 5.8 - 8r."""
    return plan_cost(C, feasible_plan(r))  # Linear in r with slope C11 - C12 - C21 + C22 = -8.


# %% 07. Prices and the dual
#
# A price f_i at each source point and g_j at each target point is feasible when f_i + g_j <= C_ij on
# every route. By complementary slackness we read the prices off the loaded routes of an optimal
# plan, one substitution at a time, once we pin one price to remove the shift.
# Sections 7.4 to 7.4.4.


def prices_from_loaded_routes(C: torch.Tensor, pi: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Solve f_i + g_j = C_ij on the routes carrying mass, with f_1 = 0 fixing the shift."""
    n, m = C.shape
    f = torch.full((n,), float("nan"), dtype=DTYPE)
    g = torch.full((m,), float("nan"), dtype=DTYPE)
    f[0] = 0.0  # Adding a constant to f and subtracting it from g changes no sum f_i + g_j.
    for _ in range(n + m):  # Each pass propagates along one more loaded route.
        for i in range(n):
            for j in range(m):
                if pi[i, j] <= 0:
                    continue  # An empty route imposes no equation.
                if not torch.isnan(f[i]) and torch.isnan(g[j]):
                    g[j] = C[i, j] - f[i]  # The route prices the target point.
                elif torch.isnan(f[i]) and not torch.isnan(g[j]):
                    f[i] = C[i, j] - g[j]  # The route prices the source point.
    return f, g


def slack(C: torch.Tensor, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
    """The matrix C_ij - f_i - g_j, nonnegative exactly when the prices are feasible."""
    return C - f[:, None] - g[None, :]  # [n, m].


def dual_value(a: torch.Tensor, f: torch.Tensor, b: torch.Tensor, g: torch.Tensor) -> float:
    """<a, f> + <b, g>, the value of a pair of prices."""
    return float(a @ f + b @ g)


# %% 08. Powers of a distance, and non-crossing on the line
#
# With the squared cost the difference between the ordered and crossed pairings is -2 times the
# product of the two gaps, so an optimal plan on the line contains no crossing. We then write W_p
# in the quantile form and compute it by sorting. Sections 7.5 and 7.5.1.


def crossing_difference(x1: float, x2: float, y1: float, y2: float) -> tuple[float, float]:
    """The ordered cost minus the crossed cost, expanded and in the collected form -2 dx dy."""
    expanded = (x1 - y1) ** 2 + (x2 - y2) ** 2 - (x1 - y2) ** 2 - (x2 - y1) ** 2
    collected = -2.0 * (x2 - x1) * (y2 - y1)  # Every squared term cancels, leaving the cross terms.
    return expanded, collected


def quantile_levels(levels: int) -> torch.Tensor:
    """Midpoints (i + 1/2) / levels of a uniform grid on [0, 1], the levels q the integral uses."""
    return (torch.arange(levels, dtype=DTYPE) + 0.5) / levels  # [levels].


def empirical_quantile(points: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
    """F^{-1}(q) = inf{x : F(x) >= q} for n equally weighted points, a step function of q."""
    ordered = points.sort().values  # [n] the sorted points, the steps of F^{-1}.
    index = torch.clamp((q * len(points)).floor().long(), max=len(points) - 1)  # [levels].
    return ordered[index]  # [levels] the quantile at each level.


def wasserstein_on_the_line(x: torch.Tensor, y: torch.Tensor, p: float, levels: int) -> float:
    """W_p^p as the integral over quantile levels q of |F_mu^{-1}(q) - F_nu^{-1}(q)|^p.

    With n equally weighted points the quantile function is a step function whose values are the
    sorted points, so the level-by-level pairing is the sorted pairing. We can therefore compute
    the whole integral by sorting each side once.
    """
    q = quantile_levels(levels)  # [levels].
    difference = empirical_quantile(x, q) - empirical_quantile(y, q)  # [levels].
    return float((difference.abs() ** p).mean())  # The integral of a step function over [0, 1].


# %% 09. Gaussians, and the case where the plan is a map
#
# For two Gaussians on the line the squared 2-Wasserstein distance is the squared difference of the
# means plus the squared difference of the standard deviations. The optimal map is then the
# increasing affine function matching quantiles. Sections 7.6 to 7.6.2.


def gaussian_w2_squared(m0: float, s0: float, m1: float, s1: float) -> float:
    """W_2^2 between N(m0, s0^2) and N(m1, s1^2)."""
    return (m0 - m1) ** 2 + (s0 - s1) ** 2


def brenier_map(x: torch.Tensor, m0: float, s0: float, m1: float, s1: float) -> torch.Tensor:
    """T(x) = m1 + (s1 / s0) (x - m0), the gradient of a convex function in one dimension."""
    return m1 + (s1 / s0) * (x - m0)  # Same shape as x.


# %% 10. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=3072)  # --steps is the number of quantile levels in Section 7.5.1.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet

    pairings = two_pairings()
    if loud:
        print("Four points: sources at 0 and 2, targets at 1 and 3, mass 1/2 each")
        print(f"  arrows   ordered {pairings['arrows_ordered']:.0f}   "
              f"crossed {pairings['arrows_crossed']:.0f}")
        print(f"  by mass  ordered {pairings['mass_ordered']:.0f}   "
              f"crossed {pairings['mass_crossed']:.0f}   "
              f"independent {pairings['mass_independent']:.0f}")

    C = cost_matrix(SOURCE_X, TARGET_Y)  # [[1, 9], [1, 1]] for the running example.
    low, high = feasible_interval()
    grid = torch.linspace(low, high, 7, dtype=DTYPE)  # Seven points along the feasible segment.
    best_r = float(grid[int(torch.tensor([cost_along_the_segment(float(r), C)
                                          for r in grid]).argmin())])
    pi_star = feasible_plan(high)  # The cost decreases in r, so the optimum is the right endpoint.
    cost_star = plan_cost(C, pi_star)
    if loud:
        print(f"\nTwo-by-two: a = {SOURCE_A.tolist()}, b = {TARGET_B.tolist()}, C = {C.tolist()}")
        print(f"  feasible r in [{low:.1f}, {high:.1f}],  cost(r) = "
              f"{cost_along_the_segment(0.0, C):.1f} - "
              f"{(cost_along_the_segment(0.0, C) - cost_along_the_segment(1.0, C)):.0f} r")
        print(f"  pi* = {pi_star.tolist()}   <C, pi*> = {cost_star:.1f}")

    f, g = prices_from_loaded_routes(C, pi_star)
    S = slack(C, f, g)
    dual = dual_value(SOURCE_A, f, TARGET_B, g)
    gap = cost_star - dual  # Equals sum_ij pi_ij (C_ij - f_i - g_j), term by term nonnegative.
    if loud:
        print(f"\nPrices f = {f.tolist()}, g = {g.tolist()}")
        print(f"  slack = {S.tolist()}   dual value {dual:.1f}   duality gap {gap:.1e}")
        print(f"  slackness: mass on a route of positive slack = "
              f"{float((pi_star * (S > 1e-12)).sum()):.1f}")

    expanded, collected = crossing_difference(0.0, 2.0, 1.0, 3.0)
    three_x = torch.tensor([4.0, 1.0, 6.0], dtype=DTYPE)  # Section 7.5.2.
    three_y = torch.tensor([9.0, 2.0, 3.0], dtype=DTYPE)
    permutations = monge_permutations(three_x.sort().values, three_y.sort().values)
    quantile_w22 = wasserstein_on_the_line(three_x, three_y, 2.0, args.steps)
    C3 = cost_matrix(three_x.sort().values, three_y.sort().values)
    f3 = torch.tensor([0.0, 0.0, 8.0], dtype=DTYPE)  # Three equations in six unknowns, one solution.
    g3 = torch.tensor([1.0, 1.0, 1.0], dtype=DTYPE)
    dual3 = dual_value(torch.full((3,), 1 / 3, dtype=DTYPE), f3,
                       torch.full((3,), 1 / 3, dtype=DTYPE), g3)
    if loud:
        print(f"\nNo crossing: expanded {expanded:+.0f} = collected {collected:+.0f}")
        print(f"  three points, six permutations: {[round(c, 3) for c, _ in permutations]}")
        print(f"  quantile integral on {args.steps} levels: {quantile_w22:.4f}")
        print(f"  slack of f = (0, 0, 8) and g = (1, 1, 1): {slack(C3, f3, g3).tolist()}")
        print(f"  dual value {dual3:.4f}")

    source = torch.special.ndtri(quantile_levels(args.steps))  # N(0,1) read at the same levels.
    mapped = brenier_map(source, 0.0, 1.0, 3.0, 2.0)  # T(x) = 3 + 2x.
    displacement_square = float(((mapped - source) ** 2).mean())  # E[(3 + X)^2] = 9 + 1.
    w22_gauss = gaussian_w2_squared(0.0, 1.0, 3.0, 2.0)
    if loud:
        probe = brenier_map(torch.tensor([-1.0, 0.0, 1.0], dtype=DTYPE), 0.0, 1.0, 3.0, 2.0)
        print(f"\nN(0,1) to N(3,4): T(x) = {float(probe[1]):.0f} + "
              f"{float(probe[2] - probe[1]):.0f}x, T(-1) = {float(probe[0]):.0f}, "
              f"T(0) = {float(probe[1]):.0f}, T(1) = {float(probe[2]):.0f}")
        print(f"  sampled mean square displacement {displacement_square:.3f}")

    report = {
        "ordered pairing, by arrows": pairings["arrows_ordered"],
        "crossed pairing, by arrows": pairings["arrows_crossed"],
        "ordered pairing, by mass": pairings["mass_ordered"],
        "crossed pairing, by mass": pairings["mass_crossed"],
        "independent coupling, by mass": pairings["mass_independent"],
        "atoms of T_# delta_0, of nu": f"{len(pushforward_of_a_point_mass(1.0))}, {len(NU_ATOMS)}",
        "permutations at n = 2, 3, 256": f"2, 6, a number with "
                                         f"{len(str(math.factorial(256)))} digits",
        "argmin of the cost over r": best_r,
        "<C, pi*>": cost_star,
        "dual value <a,f> + <b,g>": dual,
        "duality gap": gap,
        "pi* is a coupling": is_coupling(pi_star, SOURCE_A, TARGET_B),
        "W_2 of the two-by-two": cost_star ** 0.5,
        "sorted pairing on three points": permutations[0][0],
        "worst pairing on three points": permutations[-1][0],
        "W_2^2 from the quantile integral": quantile_w22,
        "dual value on three points": dual3,
        "W_2^2 of N(0,1) and N(3,4)": float(w22_gauss),
        "W_2 of N(0,1) and N(3,4)": float(w22_gauss) ** 0.5,
        "E[(T(X) - X)^2] on a sample": displacement_square,
    }
    print_report("Optimal transport, Chapter 7 Sections 7.1 to 7.6", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()