"""The flow map of an ordinary differential equation, with no learning anywhere.

A flow map X_{s,t}(x) returns the state at time t of the trajectory that passes through x at time
s, so one evaluation replaces the N Euler steps of Chapter 2. Here we integrate a velocity we
already know, the Gaussian flow of the chapter, and check against it every property the rest of
the lecture trains for: the three axioms, the inverse, the Jacobian and its differential equation,
Liouville's formula, the transport of a density, and the generator.

Chapter 6, Sections 6.3 to 6.12 of the notes. Runs on the scalar Gaussian flow and on a batch from
``cis6270.data.gaussian_mixture_2d``. Prints the four numbers of the chapter's listing, 4.100000
for the map, the cocycle agreeing with it, 0.700000 for the inverse and a Jacobian determinant of
3.000000, together with log 3, the density ratio and the log likelihood.

    python lecture_6/flow_map.py                 # 200 Runge-Kutta substeps per crossing
    python lecture_6/flow_map.py --steps 20      # 20 substeps, a second, and the same report
"""

# %% 01. Imports and configuration
import math

import torch
from torch.func import jacrev  # Differentiates the whole integration in its starting point.

from cis6270.data import GAUSSIAN_STD, gaussian_mixture_2d
from cis6270.runner import common_parser, print_report, save_run, set_seed

torch.set_num_threads(1)  # The integration is a scalar loop over a small batch, so one thread is
# the fastest here. ``tests/conftest.py`` sets the same thing for the test session, and this line
# acts when a reader runs the script directly.
torch.set_default_dtype(torch.float64)  # We check the axioms below the float32 noise floor.


SUBSTEPS = 200  # n, the Runge-Kutta substeps per crossing.
START = 0.7  # The state x at which the chapter evaluates every check.


# %% 02. The velocity field of the Gaussian flow
#
# Source N(0,1) and data N(2,9), paired by x_1 = 2 + 3 x_0 and joined by a straight line. The
# velocity is a translation at speed 2 plus an expansion whose strength 2/(1+2t) decays as the
# cloud widens.
def b(x, t):  # velocity field
    return 2.0 + 2.0 * (x - 2.0 * t) / (1.0 + 2.0 * t)  # Section 6.9, b(x,t).


def exact_map(x, s, t):
    """X_{s,t}(x) = 2t + (1+2t)(x - 2s)/(1 + 2s), the closed form the integrator has to reproduce."""
    return 2.0 * t + (1.0 + 2.0 * t) * (x - 2.0 * s) / (1.0 + 2.0 * s)  # [B] or scalar.


# %% 03. The flow map by fourth-order Runge-Kutta
#
# Four velocity evaluations per substep, at its start, twice at its midpoint and at its end,
# combined with the weights 1, 2, 2, 1 over six. The step h is negative when t is below s, so the
# same function integrates backward with no change.
def flow(x, s, t, n=SUBSTEPS):  # X_{s,t}(x) by RK4
    h = (t - s) / n
    for k in range(n):
        u = s + k * h
        k1 = b(x, u)
        k2 = b(x + 0.5 * h * k1, u + 0.5 * h)
        k3 = b(x + 0.5 * h * k2, u + 0.5 * h)
        k4 = b(x + h * k3, u + h)
        x = x + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
    return x


# %% 04. The three axioms, checked numerically
#
# Identity, cocycle and tangent characterize the map (Proposition 6.1). The first holds by
# definition, the second by uniqueness of the trajectory, the third by the defining equation.
def check_axioms(x, n):
    """Return the residual of each axiom at the state x, all three zero for the true map."""
    identity = flow(x, 0.5, 0.5, n) - x  # X_{t,t} = id.
    cocycle = flow(flow(x, 0.0, 0.5, n), 0.5, 1.0, n) - flow(x, 0.0, 1.0, n)  # X_{t,u} o X_{s,t}.
    h = 1e-6  # One-sided difference quotient at the diagonal.
    tangent = (flow(x, 0.5, 0.5 + h, n) - x) / h - b(x, 0.5)  # d/dt X at t = s equals b(x,s).
    return float(identity.abs()), float(cocycle.abs()), float(tangent.abs())


# %% 05. The Jacobian and the equation it satisfies
#
# dJ/dt = grad b(X_{s,t}(x), t) J, starting from J_{s,s} = I. In one dimension grad b is the single
# derivative 2/(1+2t), and the solution from zero is J_{0,t} = 1 + 2t.
def jacobian_by_ode(x, s, t, n):
    """Integrate the state and its Jacobian together, with one Euler step each on the same grid."""
    jacobian = torch.ones_like(x)  # J_{s,s} = I, the identity axiom differentiated in x.
    h = (t - s) / n
    for k in range(n):
        u = s + k * h
        grad_b = 2.0 / (1.0 + 2.0 * u)  # d b / d x at time u, the divergence in one dimension.
        jacobian = jacobian + h * grad_b * jacobian  # Linear in J, with the coefficient frozen.
        x = x + h * b(x, u)  # The trajectory the coefficient is read along.
    return jacobian


def liouville_log_det(s, t, n):
    """Integrate d/dt log det J = div b, on this flow the number log 3 over [0, 1]."""
    h = (t - s) / n
    grid = s + h * (torch.arange(n) + 0.5)  # Midpoints of the substeps.
    return float((2.0 / (1.0 + 2.0 * grid)).sum() * h)  # Riemann sum of the divergence.


# %% 06. Transporting a density along the flow
#
# p_t(X_{s,t}(x)) |det J_{s,t}(x)| = p_s(x). The flow stretches every length by 3 over the unit
# interval, so it divides every density by 3.
def normal_density(x, mean, std):
    """The scalar Gaussian density, used at both ends of the transport statement."""
    return math.exp(-((x - mean) ** 2) / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi))


def mixture_log_density_2d(points, scale, shift):
    """Log density of the four-Gaussian mixture after the affine map y = shift + scale * x.

    The pushforward of a Gaussian mixture under an affine map is a Gaussian mixture with the
    centers and the widths carried through. The transported density is then a closed form.
    """
    from cis6270.data import GAUSSIAN_CENTERS

    centers = shift + scale * GAUSSIAN_CENTERS.to(points)  # [4, 2] transported centers.
    std = scale * GAUSSIAN_STD  # Every component widens by the same factor.
    squared = ((points[:, None, :] - centers) ** 2).sum(dim=-1)  # [N, 4].
    component = -0.5 * squared / std ** 2 - math.log(2.0 * math.pi * std ** 2)  # [N, 4].
    return torch.logsumexp(component + math.log(0.25), dim=-1)  # [N] mixture log density.


# %% 07. The generator of the flow
#
# (L_t f)(x) = b(x,t) . grad f(x), and the rate of change of f along a trajectory is the generator
# applied to f and read at the moving point.
def generator(f, grad_f, x, t):
    """(L_t f)(x) for a scalar test function, in one dimension b times the derivative."""
    return b(x, t) * grad_f(x)  # Section 6.11, the generator.


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=SUBSTEPS, samples=512)
    args = parser.parse_args()  # --steps is the substep count n here, since nothing is trained.
    set_seed(args.seed)
    n = max(args.steps, 4)  # A floor of four substeps, since the weights are fourth order.

    # %% 09. The chapter's four lines
    x = torch.tensor(START)  # The single starting point every check below uses.
    one_crossing = flow(x, 0.0, 1.0, n)  # 2 + 3x = 4.1.
    two_crossings = flow(flow(x, 0.0, 0.5, n), 0.5, 1.0, n)  # The cocycle.
    there_and_back = flow(flow(x, 0.0, 1.0, n), 1.0, 0.0, n)  # The inverse.
    det_jacobian = jacrev(lambda z: flow(z, 0.0, 1.0, n))(x)  # det J = 3.
    if not args.quiet:
        print(f"X_0,1({START})      = {float(one_crossing):.6f}   closed form "
              f"{float(exact_map(x, 0.0, 1.0)):.6f}")
        print(f"cocycle through 0.5 = {float(two_crossings):.6f}")
        print(f"inverse X_1,0 o X_0,1 = {float(there_and_back):.6f}")
        print(f"det J_0,1           = {float(det_jacobian):.6f}")

    identity, cocycle, tangent = check_axioms(x, n)

    # %% 10. The Jacobian two ways, and Liouville's formula
    jacobian_ode = jacobian_by_ode(x, 0.0, 1.0, 4000)  # One Euler solve of dJ/dt = grad b . J.
    log_det = liouville_log_det(0.0, 1.0, 4000)  # Integral of the divergence over [0, 1].

    # %% 11. Transport, in one dimension and in two
    image = float(one_crossing)  # X_{0,1}(0.7) = 4.1.
    p_source = normal_density(START, 0.0, 1.0)  # p_0(0.7).
    p_data = normal_density(image, 2.0, 3.0)  # p_1(4.1).
    log_likelihood = math.log(p_source) - log_det  # Section 6.10, with the divergence integral.

    points = gaussian_mixture_2d(num_points=args.samples, seed=args.seed)  # [N, 2] source cloud.
    moved = exact_map(points, 0.0, 1.0)  # The same affine flow applied to both coordinates.
    log_p0 = mixture_log_density_2d(points, 1.0, 0.0)  # [N] density before transport.
    log_p1 = mixture_log_density_2d(moved, 3.0, 2.0)  # [N] density of the transported mixture.
    transport_gap = float((log_p1 + 2.0 * math.log(3.0) - log_p0).abs().max())  # det J = 3^2 = 9.

    # %% 12. The generator along one trajectory
    f, grad_f = (lambda z: z ** 2), (lambda z: 2.0 * z)  # A test function and its derivative.
    h = 1e-5
    along = (f(flow(x, 0.0, 0.5 + h, n)) - f(flow(x, 0.0, 0.5 - h, n))) / (2.0 * h)  # d/dt f(X).
    generator_gap = abs(float(along) - float(generator(f, grad_f, flow(x, 0.0, 0.5, n), 0.5)))

    report = {
        "X_0,1(0.7)": float(one_crossing),
        "cocycle through u=0.5": float(two_crossings),
        "inverse back to the start": float(there_and_back),
        "det J_0,1 by jacrev": float(det_jacobian),
        "det J_0,1 by its own ODE": float(jacobian_ode),
        "log det J_0,1 by Liouville": log_det,
        "identity axiom residual": identity,
        "cocycle axiom residual": cocycle,
        "tangent axiom residual": tangent,
        "p_0(0.7)": p_source,
        "p_1(4.1)": p_data,
        "density ratio p_0/p_1": p_source / p_data,
        "log p_1(4.1) by transport": log_likelihood,
        "log p_1(4.1) directly": math.log(p_data),
        "2-D transport gap": transport_gap,
        "generator gap": generator_gap,
        "substeps n": n,
    }
    print_report("Flow map of the Gaussian flow, Chapter 6", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()