"""The Schrodinger bridge over a finite set: rates in place of drifts.

A finite state space has no small displacement, so the control is a reweighting of jump rates and
relative entropy on paths is a time integral of a Bregman divergence between two rate matrices.
Everything else carries over. The chain rule splits the problem into an endpoint term and a bridge
term, and the two potentials are the Sinkhorn scalings of the reference endpoint law. Under Doob's
h-transform each reference rate is multiplied by the ratio of a backward potential at the
destination to its value at the current state.

Chapter 7, Sections 7.21 to 7.23.3 of the notes. Runs the two-state continuous-time example with
the generator Q, the source law (0.5, 0.5) and the requested terminal law (0.2, 0.8). Prints the
bridge rates at three times, the reverse rates, the Euler sampler's error as the step shrinks,
and the cost computed on the endpoints and on the paths.

    python lecture_7/discrete_bridge.py              # 10 Euler steps, then 20, 40 and 100
    python lecture_7/discrete_bridge.py --steps 20   # a finer sampler, 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 for the matrix exponentials and the time integral.

Q = torch.tensor([[-1.0, 1.0], [1.5, -1.5]], dtype=DTYPE)  # The reference generator.
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.
STATIONARY = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Proportional to (1.5, 1), the rows of Q.
GRID = 1001  # Points in the Simpson rule for the time integral of the rate error.


# %% 02. Relative entropy between two jump processes
#
# We collect the one-step relative entropies and let the step go to zero. What is left is a time
# integral of the Bregman divergence of u log u between the two rates. The divergence is
# nonnegative and vanishes only when the two rates agree. It is not symmetric in the ratio.
# Section 7.21.1.


def rate_error(controlled: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
    """Q^theta log(Q^theta / Q) - Q^theta + Q, the Bregman divergence of u log u, entry by entry."""
    safe = controlled.clamp(min=1e-300)  # A zero rate contributes the reference rate alone.
    return controlled * (safe / reference).log() - controlled + reference


def scaling_cost(c: float) -> float:
    """c log c - c + 1, the cost per unit of reference rate of scaling every rate by c."""
    return c * math.log(c) - c + 1.0


# %% 03. The Schrodinger problem for a chain
#
# The endpoint law of a chain with a time-independent generator is diag(pi_0^ref) e^Q. Sinkhorn
# scales that matrix to the two marginals, and those scalings are the two potentials.
# Section 7.21.2.


def endpoint_law(generator: torch.Tensor, start: torch.Tensor) -> torch.Tensor:
    """R_01 = diag(pi_0) e^Q, the reference probability of starting at x and ending at y."""
    return torch.diag(start) @ torch.matrix_exp(generator)  # [S, S].


def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int = 400
             ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Scale a positive matrix to the two marginals, returning f, g and the plan f R_01 g."""
    f = torch.ones_like(start)
    g = torch.ones_like(end)
    for _ in range(sweeps):
        g = end / (matrix.T @ f)  # Column sums become exactly pi_1.
        f = start / (matrix @ g)  # Row sums become exactly pi_0.
    return f, g, f[:, None] * matrix * g[None, :]


def relative_entropy(p: torch.Tensor, q: torch.Tensor) -> float:
    """Sum over entries of p log(p / q), for two positive matrices or vectors."""
    mask = p > 0
    return float((p[mask] * (p[mask] / q[mask]).log()).sum())


# %% 04. Rates reweighted by an h-transform
#
# The backward potential is g averaged over the reference from the current time to time 1. We
# multiply each off-diagonal rate by the ratio of that potential at the destination to its value at
# the current state, and then set the diagonal so that the row sums to zero. Section 7.22.


def backward_potential(generator: torch.Tensor, g: torch.Tensor, t: float) -> torch.Tensor:
    """phi_t = e^{Q(1-t)} g, the reference average of g at the final state seen from time t."""
    return torch.matrix_exp(generator * (1.0 - t)) @ g  # [S].


def tilted_generator(generator: torch.Tensor, phi: torch.Tensor) -> torch.Tensor:
    """Q*(x, y) = Q(x, y) phi(y) / phi(x) off the diagonal, with the diagonal closing each row."""
    tilted = generator * phi[None, :] / phi[:, None]  # [S, S] every entry reweighted.
    tilted.fill_diagonal_(0.0)  # Discard the reweighted diagonal, a quantity with no meaning.
    return tilted - torch.diag(tilted.sum(dim=1))  # Each row now sums to zero.


def bridge_marginal(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor,
                    g: torch.Tensor, t: float) -> torch.Tensor:
    """p*_t = psi_t phi_t, the forward potential carried under the reference times the backward one."""
    psi = (start * f) @ torch.matrix_exp(generator * t)  # [S] the forward factor.
    return psi * backward_potential(generator, g, t)  # [S] their entrywise product.


# %% 05. Sampling the bridge with Euler steps
#
# Over a step the chain moves to a different state with probability dt Q*(x, y) and stays with the
# remaining probability. This is the Euler sampler of Chapter 4 applied to a generator that
# changes with time. Section 7.23.1.


def euler_marginal(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor,
                   g: torch.Tensor, steps: int) -> tuple[torch.Tensor, list[float]]:
    """Propagate p <- p (I + dt Q*_t) and record the probability of leaving state 1 at each step."""
    dt = 1.0 / steps
    p = start.clone()
    leaving = []
    identity = torch.eye(len(start), dtype=DTYPE)
    for step in range(steps):
        tilted = tilted_generator(generator, backward_potential(generator, g, step * dt))
        leaving.append(float(tilted[0, 1] * dt))  # The jump probability out of state 1.
        p = p @ (identity + dt * tilted)  # One Euler step of the master equation.
    return p, leaving


# %% 06. The same bridge run backwards
#
# We match the probability flux on an edge in both directions. That fixes the reverse rate at
# the forward rate times the ratio of the marginal at the two ends of the edge. Section 7.23.2.


def reverse_rates(tilted: torch.Tensor, marginal: torch.Tensor) -> torch.Tensor:
    """Q_hat(y, x) = Q(x, y) p_t(x) / p_t(y), the rates under which the flux on an edge matches."""
    return (tilted * marginal[:, None] / marginal[None, :]).T  # [S, S], reversed in both indices.


# %% 07. The cost of the bridge, computed twice
#
# On the endpoints the cost is the relative entropy of the plan from the reference endpoint law. On
# the paths it is the time integral of the rate error, weighted by the marginal of the bridge at
# that time. The two agree. Section 7.23.3.


def path_cost(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor, g: torch.Tensor,
              points: int = GRID) -> tuple[float, float]:
    """Integrate the weighted rate error over [0, 1] by Simpson's rule.

    Returns the integral and the value of the integrand at t = 1/2, the one the chapter prints.
    """
    off = ~torch.eye(len(generator), dtype=torch.bool)  # [S, S] true away from the diagonal.
    times = torch.linspace(0.0, 1.0, points, dtype=DTYPE)
    values = []
    for t in times.tolist():
        tilted = tilted_generator(generator, backward_potential(generator, g, t))
        marginal = bridge_marginal(generator, start, f, g, t)
        error = torch.zeros_like(tilted)
        error[off] = rate_error(tilted[off], generator[off])  # Only the edges carry a cost.
        values.append(float((marginal * error.sum(dim=1)).sum()))
    weights = torch.ones(points, dtype=DTYPE)
    weights[1:-1:2] = 4.0  # Simpson's rule on an even number of intervals.
    weights[2:-1:2] = 2.0
    integral = float((weights * torch.tensor(values, dtype=DTYPE)).sum() / (3.0 * (points - 1)))
    return integral, values[(points - 1) // 2]


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=10)  # --steps is the coarsest number of Euler steps in the table.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet
    coarse = max(args.steps, 2)

    if loud:
        print("Scaling every reference rate by c costs c log c - c + 1 per unit of rate")
        print(f"  halving {scaling_cost(0.5):.4f}, doubling {scaling_cost(2.0):.4f}")

    R01 = endpoint_law(Q, PI_0)
    f, g, plan = sinkhorn(R01, PI_0, PI_1)
    phi_half = backward_potential(Q, g, 0.5)
    tilted_half = tilted_generator(Q, phi_half)
    tilted_zero = tilted_generator(Q, backward_potential(Q, g, 0.0))
    tilted_one = tilted_generator(Q, g)
    marginal_half = bridge_marginal(Q, PI_0, f, g, 0.5)
    marginal_one = bridge_marginal(Q, PI_0, f, g, 1.0)
    reverse = reverse_rates(tilted_half, marginal_half)
    if loud:
        print(f"\ne^Q = {[[round(float(z), 4) for z in row] for row in torch.matrix_exp(Q)]}")
        print(f"R_01 = {[[round(float(z), 4) for z in row] for row in R01]}, "
              f"column sums {[round(float(z), 4) for z in R01.sum(0)]}")
        print(f"  f = {[round(float(z), 4) for z in f]}, g = {[round(float(z), 4) for z in g]}, "
              f"g ratio {float(g[1] / g[0]):.3f}")
        print(f"  pi* = {[[round(float(z), 4) for z in row] for row in plan]}")
        print(f"  phi_1/2 = {[round(float(z), 4) for z in phi_half]}, "
              f"ratio {float(phi_half[1] / phi_half[0]):.2f}")
        print(f"  Q*_1/2 = {[[round(float(z), 4) for z in row] for row in tilted_half]}")
        print(f"  rates out of the two states: at t = 0 "
              f"{float(tilted_zero[0, 1]):.4f} and {float(tilted_zero[1, 0]):.4f}, "
              f"at t = 1 {float(tilted_one[0, 1]):.4f} and {float(tilted_one[1, 0]):.4f}")
        print(f"  marginal at t = 1/2 {[round(float(z), 4) for z in marginal_half]}, "
              f"at t = 1 {[round(float(z), 4) for z in marginal_one]}")
        print(f"  reverse rates at t = 1/2: {float(reverse[1, 0]):.4f} from 2 back to 1, "
              f"{float(reverse[0, 1]):.4f} from 1 back to 2")
        print(f"  the flux on the edge: {float(marginal_half[0] * tilted_half[0, 1]):.4f} "
              f"against {float(marginal_half[1] * reverse[1, 0]):.4f}")

    sampler = {}
    for steps in (coarse, 2 * coarse, 4 * coarse, 10 * coarse):
        terminal, leaving = euler_marginal(Q, PI_0, f, g, steps)
        sampler[steps] = (terminal, float((terminal - PI_1).abs().sum() / 2), leaving)
    if loud:
        print("\nsteps   terminal law                 error    first and last jump probability")
        for steps, (terminal, error, leaving) in sampler.items():
            print(f"  {steps:<5d} ({float(terminal[0]):.4f}, {float(terminal[1]):.4f})"
                  f"           {error:.4f}   {leaving[0]:.4f}, {leaving[-1]:.4f}")

    endpoint_cost = relative_entropy(plan, R01)
    integral, at_half = path_cost(Q, PI_0, f, g)
    speeds = {}
    for lam in (0.5, 1.0, 2.0, 4.0):
        scaled_law = endpoint_law(lam * Q, PI_0)
        _, _, scaled_plan = sinkhorn(scaled_law, PI_0, PI_1)
        speeds[lam] = relative_entropy(scaled_plan, scaled_law)
    limit = relative_entropy(PI_1, STATIONARY)
    if loud:
        print(f"\nKL(pi* || R_01) = {endpoint_cost:.4f}")
        print(f"  the integrand at t = 1/2 is {at_half:.4f}, and its integral is {integral:.4f}")
        print(f"  at lambda = 1/2, 1, 2, 4 the cost is "
              f"{[round(speeds[k], 4) for k in (0.5, 1.0, 2.0, 4.0)]}")
        print(f"  the limit KL(pi_1 || (0.6, 0.4)) is {limit:.4f}")

    report = {
        "cost of halving every rate": scaling_cost(0.5),
        "cost of doubling every rate": scaling_cost(2.0),
        "R_01 column sums": [round(float(z), 4) for z in R01.sum(dim=0)],
        "f": [round(float(z), 4) for z in f],
        "g": [round(float(z), 4) for z in g],
        "pi*_11": float(plan[0, 0]),
        "phi at t = 1/2": [round(float(z), 4) for z in phi_half],
        "bridge rate out of state 1": float(tilted_half[0, 1]),
        "bridge rate out of state 2": float(tilted_half[1, 0]),
        "rate out of state 1 at t = 1": float(tilted_one[0, 1]),
        "marginal at t = 1/2": [round(float(z), 4) for z in marginal_half],
        "marginal at t = 1": [round(float(z), 4) for z in marginal_one],
        "reverse rate, 2 back to 1": float(reverse[1, 0]),
        "Euler error at the coarsest step": sampler[coarse][1],
        "Euler error, step halved": sampler[2 * coarse][1],
        "Euler error, step quartered": sampler[4 * coarse][1],
        "Euler error at ten times the steps": sampler[10 * coarse][1],
        "KL(pi* || R_01)": endpoint_cost,
        "rate error integrand at t = 1/2": at_half,
        "its integral over [0, 1]": integral,
        "cost at lambda = 1/2": speeds[0.5],
        "cost at lambda = 4": speeds[4.0],
        "KL(pi_1 || the stationary law)": limit,
    }
    print_report("The bridge over a finite set, Chapter 7 Sections 7.21 to 7.23.3", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()