"""Diffusion Schrodinger Bridge: iterative proportional fitting run on path laws.

Each half-step of proportional fitting is a relative-entropy projection onto the set of path laws
with one marginal correct, and the chain rule solves it in closed form by keeping the conditional
law of the previous iterate and resetting that one marginal. The conditional law is the current
process run backward in time, so each half-step reverses the current process and restarts it from
the requested law at the other end.

Chapter 7, Section 7.19.1 of the notes. Runs the half-steps exactly on the two-state chain with
transition matrix A, source law (0.5, 0.5) and requested terminal law (0.2, 0.8), and checks the
diffusion time reversal in closed form on a Gaussian. Prints the backward matrices, the law after
each half-step, the forward matrices of the first round and the error at the free end.

    python lecture_7/dsb.py              # 4 rounds, the first one worked by hand
    python lecture_7/dsb.py --steps 1    # the single round the chapter writes out
"""

# %% 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  # The chapter quotes six figures of the backward and forward matrices.

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.
TILTED_0 = torch.tensor([[0.599160, 0.400840], [0.299266, 0.700734]], dtype=DTYPE)  # A^h_0.
TILTED_1 = torch.tensor([[0.307297, 0.692703], [0.112490, 0.887510]], dtype=DTYPE)  # A^h_1.


# %% 02. A path law as a joint table
#
# With two states and two steps a path law is a table over the triple (x_0, x_1, x_2), so every
# projection below is an exact operation on eight numbers. Section 7.12.1.


def reference_path_law(start: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor:
    """R(x_{0:N}) as a tensor with one axis per time, from the starting law and one matrix per step."""
    law = start.clone()  # [K] the law at time 0.
    for _ in range(steps):
        law = law.unsqueeze(-1) * matrix  # Multiply by the transition matrix on a new axis.
    return law  # [K] * (N + 1).


def marginal(law: torch.Tensor, axis: int) -> torch.Tensor:
    """The law of the state at one time, summed over every other time."""
    other = [d for d in range(law.dim()) if d != axis]
    return law.sum(dim=other)  # [K].


def one_step_conditional(law: torch.Tensor, step: int) -> torch.Tensor:
    """P(x_{k+1} | x_k), the forward transition matrix the path law carries at that step."""
    pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)])  # [K, K].
    return pair / pair.sum(dim=1, keepdim=True)  # Normalize each row.


def reverse_step(law: torch.Tensor, step: int) -> torch.Tensor:
    """P(x_k | x_{k+1}), Bayes' rule on one step, the reversal a network has to learn."""
    pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)])  # [K, K].
    return (pair / pair.sum(dim=0, keepdim=True)).T  # [K, K] indexed by the later state first.


# %% 03. The two half-steps
#
# The odd iterate keeps the conditional law of the previous iterate given the final state and
# replaces the marginal at time 1, and the even iterate does the same at time 0. Section 7.19.1.


def fix_terminal(law: torch.Tensor, end: torch.Tensor) -> torch.Tensor:
    """P'(x_{0:N}) = pi_1(x_N) P(x_{0:N-1} | x_N), the projection onto the laws ending at pi_1."""
    last = marginal(law, law.dim() - 1)  # [K] the terminal marginal of the current law.
    return law * (end / last)  # Reweight every path by the ratio at its final state.


def fix_initial(law: torch.Tensor, start: torch.Tensor) -> torch.Tensor:
    """P''(x_{0:N}) = pi_0(x_0) P'(x_{1:N} | x_0), the projection onto the laws starting at pi_0."""
    first = marginal(law, 0)  # [K] the initial marginal of the current law.
    shape = [-1] + [1] * (law.dim() - 1)  # Broadcast the ratio along the first axis.
    return law * (start / first).reshape(shape)


def proportional_fitting(law: torch.Tensor, start: torch.Tensor, end: torch.Tensor,
                         rounds: int) -> tuple[torch.Tensor, list[float], list[torch.Tensor]]:
    """Alternate the two half-steps, returning the law, the terminal error and the law each round."""
    errors: list[float] = []
    iterates: list[torch.Tensor] = []
    for _ in range(rounds):
        law = fix_terminal(law, end)  # One half-step, and the terminal marginal is then exact.
        law = fix_initial(law, start)  # The other, and the initial marginal is then exact.
        errors.append(float((marginal(law, law.dim() - 1) - end).abs().sum()))
        iterates.append(law.clone())
    return law, errors, iterates


# %% 04. The reversal that a network learns
#
# On a finite state space the reversal is Bayes' rule on a matrix. On R^d it is the time reversal of
# a diffusion, with drift minus the forward drift plus epsilon times the score of the current
# marginal. The score of the current iterate has no closed form. Section 7.19.1.


def reverse_drift(x: torch.Tensor, t: torch.Tensor, epsilon: float, variance_0: float
                  ) -> torch.Tensor:
    """beta(x, t) = -b + epsilon grad log p_t for a driftless diffusion started from N(0, v_0)."""
    variance = variance_0 + epsilon * t  # The marginal of sqrt(epsilon) dB started from N(0, v_0).
    return epsilon * (-x / variance)  # The score of a centered Gaussian is -x over its variance.


def reversal_check(epsilon: float, variance_0: float, steps: int, paths: int,
                   generator: torch.Generator) -> tuple[float, float]:
    """Run the reversal backward from p_1 and compare its variance at time 0 with v_0."""
    dt = 1.0 / steps
    variance_1 = variance_0 + epsilon * 1.0  # The forward marginal at time 1.
    x = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(variance_1)
    for step in range(steps):
        t = torch.tensor(1.0 - step * dt, dtype=DTYPE)  # Time runs down from 1 to 0.
        noise = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(dt)
        x = x + reverse_drift(x, t, epsilon, variance_0) * dt + math.sqrt(epsilon) * noise
    return float(x.var()), variance_0


# %% 05. 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)
    generator = torch.Generator().manual_seed(args.seed)

    reference = reference_path_law(PI_0, A, 2)  # [2, 2, 2] the law P^0 = R.
    marginals = [marginal(reference, k) for k in range(3)]
    B2 = reverse_step(reference, 1)  # P(x_1 | x_2), the reversal of the last step.
    B1 = reverse_step(reference, 0)  # P(x_0 | x_1), the reversal of the earlier step.
    if loud:
        print("The reference started from pi_0, read at the three time points")
        print(f"  {[[round(float(z), 3) for z in m] for m in marginals]}")
        print(f"  B_2 = {[[round(float(z), 6) for z in row] for row in B2]}")
        print(f"  B_1(1,1) = {float(B1[0, 0]):.6f}")

    half = fix_terminal(reference, PI_1)  # The first half-step, a reset of the terminal marginal.
    backward_laws = [marginal(half, 2), marginal(half, 1), marginal(half, 0)]
    drift_at_zero = float((backward_laws[2] - PI_0).abs().sum())
    full = fix_initial(half, PI_0)  # The second half-step, a reset of the initial marginal.
    F0 = one_step_conditional(full, 0)
    F1 = one_step_conditional(full, 1)
    terminal = marginal(full, 2)
    terminal_error = float((terminal - PI_1).abs().sum())
    tilt_gap = max(float((F0 - TILTED_0).abs().max()), float((F1 - TILTED_1).abs().max()))
    if loud:
        print("\nThe backward chain started from the requested terminal law")
        written = ["(" + ", ".join(f"{float(z):.6f}" for z in m) + ")" for m in backward_laws]
        print(f"  {' -> '.join(written)}")
        print(f"  total absolute error at step 0: {drift_at_zero:.4f}")
        print(f"\nThe forward matrices of the second half-step")
        print(f"  F_0 = {[[round(float(z), 6) for z in row] for row in F0]}")
        print(f"  F_1 = {[[round(float(z), 6) for z in row] for row in F1]}")
        print(f"  terminal law {[round(float(z), 6) for z in terminal]}, "
              f"off by {terminal_error:.4f}")
        print(f"  largest gap against the tilted matrices A^h_0 and A^h_1: {tilt_gap:.4f}")

    _, errors, _ = proportional_fitting(reference, PI_0, PI_1, rounds)
    ratios = [errors[k] / errors[k + 1] for k in range(len(errors) - 1) if errors[k + 1] > 0]
    if loud:
        print(f"\nterminal error by round {['%.3e' % e for e in errors]}")
        print(f"  ratios {[round(r, 1) for r in ratios]}")

    simulated, exact = reversal_check(1.0, 1.0, 200, 8192, generator)
    if loud:
        print(f"\nThe diffusion reversal on a Gaussian, run backward from p_1")
        print(f"  variance at time 0: simulated {simulated:.3f}, exact {exact:.3f}")

    report = {
        "reference marginal at step 1": [round(float(z), 3) for z in marginals[1]],
        "reference marginal at step 2": [round(float(z), 3) for z in marginals[2]],
        "B_2(1,1)": float(B2[0, 0]),
        "B_2(2,1)": float(B2[1, 0]),
        "B_1(1,1)": float(B1[0, 0]),
        "backward law after one step": [round(float(z), 6) for z in backward_laws[1]],
        "backward law at step 0": [round(float(z), 6) for z in backward_laws[2]],
        "error at the free end, half-step": drift_at_zero,
        "F_0(1,1)": float(F0[0, 0]),
        "F_1(1,2)": float(F1[0, 1]),
        "terminal law after one round": [round(float(z), 6) for z in terminal],
        "terminal error after one round": terminal_error,
        "gap against the tilted matrices": tilt_gap,
        "terminal error, last round": errors[-1],
        "error ratio per round": ratios[0] if ratios else float("nan"),
        "rounds run": rounds,
        "reversal variance, simulated": simulated,
        "reversal variance, exact": exact,
    }
    print_report("Diffusion Schrodinger Bridge, Chapter 7 Section 7.19.1", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()