"""The static Schrodinger bridge: a path law, the chain rule, and the scaling that solves it.

A path law over a finite time grid is an ordinary probability vector over paths, and its relative
entropy from a reference splits by the chain rule into a term for the two endpoints and a term for
the interior. Only the first term is constrained, so the bridge keeps the reference bridges exactly
and all of its freedom goes into the endpoint coupling. Iterative proportional fitting then computes
that coupling by alternate rescaling.

Chapter 7, Sections 7.12 to 7.15.3 of the notes. Runs on the two-state chain with transition matrix
A, source law (0.5, 0.5) and requested terminal law (0.2, 0.8). Prints the eight path probabilities,
the chain rule on two tables, the fitted potentials, the tilted matrices and the bridge path law.

    python lecture_7/static_bridge.py              # 4 rounds of proportional fitting
    python lecture_7/static_bridge.py --steps 1    # one round, the one worked by hand
"""

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

import itertools

import torch

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

DTYPE = torch.float64  # The chapter quotes six figures of the potentials and of the path law.

A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE)  # The reference transition matrix.
PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE)  # The source law.
PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE)  # The requested terminal law after two steps.
STEPS = 2  # N, the number of transitions, so a path is a triple and there are eight of them.


# %% 02. The probability of a whole path
#
# For a Markov process the probability of a list of states is the starting mass times the product of
# the transition probabilities along it. Section 7.12.1.


def path_law(start: torch.Tensor, matrices: list[torch.Tensor]) -> dict[tuple[int, ...], float]:
    """R(x_{0:N}) over every path of the chain, as a dictionary keyed by the list of states."""
    size = len(start)
    law = {}
    for path in itertools.product(range(size), repeat=len(matrices) + 1):
        weight = float(start[path[0]])  # The starting mass R_0(x_0).
        for k, matrix in enumerate(matrices):
            weight *= float(matrix[path[k], path[k + 1]])  # One transition probability per step.
        law[path] = weight
    return law  # K^(N+1) entries, summing to one.


def relative_entropy(p: dict, q: dict) -> float:
    """KL(P || R) = sum over paths of P log(P / R), the definition of Chapter 1 indexed by paths."""
    return sum(v * torch.log(torch.tensor(v / q[k], dtype=DTYPE)).item()
               for k, v in p.items() if v > 0)


# %% 03. The chain rule for relative entropy
#
# We factor each joint law into a marginal and a conditional and average the logarithm of the
# ratio. That splits the relative entropy into a marginal term and an averaged conditional term.
# Sections 7.13 and 7.13.1.


def kl_vectors(p: torch.Tensor, q: torch.Tensor) -> float:
    """Relative entropy between two probability vectors over the same states."""
    mask = p > 0
    return float((p[mask] * (p[mask] / q[mask]).log()).sum())


def chain_rule_terms(pi: torch.Tensor, other: torch.Tensor) -> tuple[float, float, float]:
    """The total, the marginal term and the averaged conditional term of KL(pi || other)."""
    total = kl_vectors(pi.flatten(), other.flatten())  # The four cells, computed directly.
    marginal = kl_vectors(pi.sum(dim=1), other.sum(dim=1))  # The disagreement about the row.
    rows = [kl_vectors(pi[i] / pi[i].sum(), other[i] / other[i].sum()) for i in range(len(pi))]
    conditional = float((pi.sum(dim=1) * torch.tensor(rows, dtype=DTYPE)).sum())
    return total, marginal, conditional


# %% 04. Splitting a path law at its two endpoints
#
# With the pair (x_0, x_N) in the place of the first coordinate and the interior in the place of the
# second, the chain rule separates the endpoint coupling from the bridges. The second term is free
# and nonnegative, so the bridge takes the reference bridges and the problem collapses onto the
# coupling. Section 7.13.2.


def endpoint_coupling(law: dict) -> torch.Tensor:
    """The law of the pair (x_0, x_N) under a path law, summed over the interior states."""
    coupling = torch.zeros(2, 2, dtype=DTYPE)
    for path, weight in law.items():
        coupling[path[0], path[-1]] += weight  # Collect every path with the same two ends.
    return coupling


def reference_bridge(x0: int, mid: int, xN: int, matrix: torch.Tensor,
                     two_step: torch.Tensor) -> float:
    """R(x_1 | x_0, x_2) = A(x_0, x_1) A(x_1, x_2) / A^2(x_0, x_2), the interior given the ends."""
    return float(matrix[x0, mid] * matrix[mid, xN] / two_step[x0, xN])


def fill_with_reference_bridges(coupling: torch.Tensor, matrix: torch.Tensor) -> dict:
    """P(x_{0:2}) = pi(x_0, x_2) R(x_1 | x_0, x_2), a coupling filled in with reference bridges."""
    two_step = matrix @ matrix
    return {(x0, mid, xN): float(coupling[x0, xN]) * reference_bridge(x0, mid, xN, matrix, two_step)
            for x0, mid, xN in itertools.product(range(2), repeat=3)}


# %% 05. The static problem and the two potentials
#
# The bridge coupling minimizes KL(pi || R_01) over the couplings of pi_0 and pi_1, the entropic
# problem with R_01 in place of the Gibbs kernel. The minimizer is then R_01 with one scaling
# per row and one per column, and we find those two scalings from the two marginals by
# alternation. Sections 7.14 and 7.14.2.


def reference_endpoint_law(start: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor:
    """R_01 = diag(pi_0) A^N, the reference probability of starting at x and ending at y."""
    return torch.diag(start) @ torch.linalg.matrix_power(matrix, steps)  # [K, K].


def proportional_fitting(
    R01: torch.Tensor, start: torch.Tensor, end: torch.Tensor, rounds: int,
) -> tuple[torch.Tensor, torch.Tensor, list[float]]:
    """Alternate g <- pi_1 / (R_01^T f) and f <- pi_0 / (R_01 g), returning the terminal errors.

    One round is a g update followed by an f update, started from f = (1, 1), so the row sums are
    exact at the end of a round and we measure the error on the columns.
    """
    f = torch.ones_like(start)
    g = torch.ones_like(end)
    errors: list[float] = []
    for _ in range(rounds):
        g = end / (R01.T @ f)  # Column sums become exactly pi_1.
        f = start / (R01 @ g)  # Row sums become exactly pi_0, and the columns move again.
        plan = f[:, None] * R01 * g[None, :]
        errors.append(float((plan.sum(dim=0) - end).abs().sum()))
    return f, g, errors


# %% 06. The squared cost from a Brownian reference
#
# For a Brownian motion of variance epsilon the transition density is a Gibbs kernel for the squared
# distance at strength 2 epsilon. When we minimize the relative entropy from its endpoint law we
# are solving the entropic transport problem, with the noise variance in the place of the
# regularization strength. Section 7.14.1.


def brownian_reference_plan(epsilon: float = 1.0) -> torch.Tensor:
    """The bridge coupling of the four points of Section 7.3.1 under a Brownian reference."""
    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Source points.
    y = torch.tensor([1.0, 3.0], dtype=DTYPE)  # Target points.
    a = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Their masses.
    b = torch.tensor([0.3, 0.7], dtype=DTYPE)
    endpoint_law = torch.exp(-(x[:, None] - y[None, :]) ** 2 / (2.0 * epsilon))  # exp(-C/2).
    f, g, _ = proportional_fitting(endpoint_law, a, b, 200)
    return f[:, None] * endpoint_law * g[None, :]  # [2, 2].


# %% 07. Making the bridge a Markov process
#
# The backward potential averages g over the reference, one matrix-vector product per step. We then
# reweight each transition by the ratio of the potential at the destination to the potential at the
# current state, and the rows of that transition matrix sum to one by the same recursion.
# Sections 7.15 and 7.15.1.


def backward_potential(matrix: torch.Tensor, g: torch.Tensor, steps: int) -> list[torch.Tensor]:
    """h_N = g and h_k = A h_{k+1}, returned as the list h_0, ..., h_N."""
    potentials = [g]
    for _ in range(steps):
        potentials.insert(0, matrix @ potentials[0])  # One step back, into the front of the list.
    return potentials


def tilted_matrix(matrix: torch.Tensor, h_now: torch.Tensor, h_next: torch.Tensor) -> torch.Tensor:
    """A^h_k(x, y) = A(x, y) h_{k+1}(y) / h_k(x), Doob's h-transform of one step."""
    return matrix * h_next[None, :] / h_now[:, None]  # [K, K], each row summing to one.


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=4)  # --steps is the number of proportional fitting rounds.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet
    rounds = max(args.steps, 1)

    reference = path_law(PI_0, [A] * STEPS)
    if loud:
        print("The eight paths of the two-state chain, with their reference probabilities")
        for path, weight in reference.items():
            print(f"  {','.join(str(s + 1) for s in path)}   R = {weight:.6f}")
        print(f"  they sum to {sum(reference.values()):.6f}")

    table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=DTYPE)  # pi of Section 7.13.1.
    uniform = torch.full((2, 2), 0.25, dtype=DTYPE)  # pi', the uniform law on the four cells.
    total, marginal, conditional = chain_rule_terms(table, uniform)
    if loud:
        print(f"\nChain rule on two tables: total {total:.4f} = marginal {marginal:.4f} "
              f"+ conditional {conditional:.4f}")

    R01 = reference_endpoint_law(PI_0, A, STEPS)
    reference_error = float((R01.sum(dim=0) - PI_1).abs().sum())
    f, g, errors = proportional_fitting(R01, PI_0, PI_1, rounds)
    plan = f[:, None] * R01 * g[None, :]
    independent = PI_0[:, None] * PI_1[None, :]
    if loud:
        print(f"\nA^2 = {[[round(float(z), 2) for z in row] for row in A @ A]}")
        print(f"R_01 = {[[round(float(z), 3) for z in row] for row in R01]}   "
              f"column sums {[round(float(z), 3) for z in R01.sum(0)]}   error {reference_error:.2f}")
        first_g = PI_1 / R01.sum(dim=0)  # The first update, with f = (1, 1).
        print(f"  first g = {[round(float(z), 4) for z in first_g]}   "
              f"R_01 g = {[round(float(z), 4) for z in R01 @ first_g]}")
        print(f"  first f = {[round(float(z), 4) for z in PI_0 / (R01 @ first_g)]}")
        print("  round   f                        g                        terminal error   ratio")
        fr = torch.ones_like(PI_0)
        for index in range(rounds):
            gr = PI_1 / (R01.T @ fr)
            fr = PI_0 / (R01 @ gr)
            previous = reference_error if index == 0 else errors[index - 1]
            ratio = f"{previous / errors[index]:.1f}" if errors[index] > 0 else "exact"
            print(f"  {index + 1:<7d} ({float(fr[0]):.6f}, {float(fr[1]):.6f})   "
                  f"({float(gr[0]):.6f}, {float(gr[1]):.6f})   {errors[index]:.3e}      {ratio}")
        print(f"  pi* = {[[round(float(z), 6) for z in row] for row in plan]}")
        print(f"  KL(pi* || R_01) = {kl_vectors(plan.flatten(), R01.flatten()):.4f}   "
              f"KL(independent || R_01) = {kl_vectors(independent.flatten(), R01.flatten()):.4f}")

    brownian = brownian_reference_plan()
    if loud:
        print("\nBrownian reference at epsilon = 1 on the four points of Section 7.3.1")
        print(f"  pi* = {[[round(float(z), 4) for z in row] for row in brownian]}, the plan the "
              f"Gibbs kernel at strength 2 produced")

    potentials = backward_potential(A, g, STEPS)
    tilted = [tilted_matrix(A, potentials[k], potentials[k + 1]) for k in range(STEPS)]
    middle = PI_0 @ tilted[0]
    terminal = middle @ tilted[1]
    bridge_law = fill_with_reference_bridges(plan, A)
    bridge_kl = relative_entropy(bridge_law, reference)
    middle_from_paths = sum(v for k, v in bridge_law.items() if k[1] == 0)
    if loud:
        print()
        for k, h in enumerate(potentials):
            print(f"h_{k} = ({float(h[0]):.6f}, {float(h[1]):.6f})   "
                  f"ratio {float(h[1] / h[0]):.2f}")
        for k, matrix in enumerate(tilted):
            print(f"A^h_{k} = {[[round(float(z), 6) for z in row] for row in matrix]}")
        print(f"  (0.5, 0.5) -> {[round(float(z), 6) for z in middle]} -> "
              f"{[round(float(z), 6) for z in terminal]}")
        print(f"  1 / h_0 = {[round(float(z), 6) for z in 1.0 / potentials[0]]}, "
              f"the f the scaling produced")
        product = torch.diag(PI_0) @ tilted[0] @ tilted[1]  # The endpoint law of the tilted chain.
        print(f"  diag(pi_0) A^h_0 A^h_1 = "
              f"{[[round(float(z), 6) for z in row] for row in product]}")
        print("\npath      R          P*         P*/R")
        for path in reference:
            name = ",".join(str(s + 1) for s in path)
            print(f"  {name}   {reference[path]:.6f}   {bridge_law[path]:.6f}   "
                  f"{bridge_law[path] / reference[path]:.6f}")

    report = {
        "R(1,1,2)": reference[(0, 0, 1)],
        "the eight paths sum to": sum(reference.values()),
        "chain rule, total": total,
        "chain rule, marginal term": marginal,
        "chain rule, conditional term": conditional,
        "reference terminal error": reference_error,
        "terminal error, round 1": errors[0],
        "terminal error, last round": errors[-1],
        "error ratio per round": reference_error / errors[0],
        "f": [round(float(z), 6) for z in f],
        "g": [round(float(z), 6) for z in g],
        "pi*_11": float(plan[0, 0]),
        "pi*_12": float(plan[0, 1]),
        "KL(pi* || R_01)": kl_vectors(plan.flatten(), R01.flatten()),
        "KL(pi_0 x pi_1 || R_01)": kl_vectors(independent.flatten(), R01.flatten()),
        "h_0": [round(float(z), 6) for z in potentials[0]],
        "h_2 = g": [round(float(z), 6) for z in potentials[-1]],
        "A^h_0(1,1)": float(tilted[0][0, 0]),
        "A^h_1(1,2)": float(tilted[1][0, 1]),
        "rows of A^h_0 sum to": float(tilted[0].sum(dim=1).max()),
        "marginal at the middle step": float(middle[0]),
        "marginal at the last step": [round(float(z), 6) for z in terminal],
        "f against 1 / h_0": float((f - 1.0 / potentials[0]).abs().max()),
        "KL(P* || R) over the paths": bridge_kl,
        "middle state from the paths": middle_from_paths,
    }
    print_report("The static Schrodinger bridge, Chapter 7 Sections 7.12 to 7.15.3", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()