"""The forward corruption as a stochastic differential equation.

A deterministic flow is fixed by its velocity at every point, and a random path needs one more
piece, the perturbation applied over each small interval. A Brownian increment of an interval of
length h has size sqrt(h), and one Euler-Maruyama step adds that increment to the Euler step of
Chapter 2. A population started at one point spreads at a rate the noise scale sets, and a
restoring drift keeps that spread bounded. Ito's lemma and the Fokker-Planck equation then turn a
single drift and noise scale into an equation for the density. Every quantity here is a simulation
or a closed form.

In this file, as in Chapter 3, t = 0 holds a clean sample and corruption grows with t.

Chapter 3, Sections 3.2 through 3.10 of the notes. Prints the chapter's worked step from x = 2, the
Ornstein-Uhlenbeck mean and variance at t = log 2, the Ito check on E[W_t^2], and the
Fokker-Planck estimate of the Gaussian peak at t = 0.1.

    python lecture_3/sde.py                 # 400 Euler-Maruyama steps, 16384 paths
    python lecture_3/sde.py --steps 20      # a coarse grid, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

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

# `--steps` is the number of Euler-Maruyama steps on the grid and `--samples` the number of paths,
# and together they set the whole cost of a run.
TERMINAL_TIME = 1.0  # T, the end of the simulated interval.


# %% 02. Drift and Brownian increments
#
# dX_t = f(X_t, t) dt + g(t) dW_t. The drift f fixes the average displacement over a short interval
# and the Brownian increment dW_t spreads the possible next states around that average. If we set
# g = 0 we recover the ordinary differential equation of Chapter 2, with f as the velocity field.


def brownian_increment(h: float, shape: tuple[int, ...], generator=None) -> torch.Tensor:
    """One Brownian increment over an interval of length h, drawn as sqrt(h) * epsilon."""
    epsilon = torch.randn(shape, generator=generator)  # [*shape] standard Gaussian draw.
    return math.sqrt(h) * epsilon  # Covariance h I, so each coordinate has variance h.


# %% 03. The size of a Brownian increment
#
# When we multiply a standard Gaussian by sqrt(h) its covariance scales by h, so a coordinate of
# the increment has standard deviation sqrt(h). If we halve the interval, the typical displacement
# drops by sqrt(2) = 1.414. The ratio sqrt(h)/h = 1/sqrt(h) grows without bound as the grid
# refines. That is why dW_t appears in the equation and never a velocity dW_t/dt.


def increment_scale(h: float, f_value: float = 1.0, g_value: float = 1.0) -> dict:
    """The standard deviation of one coordinate of the increment, and two ratios built from it."""
    std = math.sqrt(h)  # Typical displacement of the noise over an interval of length h.
    return {
        "std": std,
        "std_over_h": std / h,  # Grows without bound as the grid refines.
        "noise_over_drift": g_value * std / (abs(f_value) * h),  # g / (|f| sqrt(h)).
    }


# %% 04. The Euler-Maruyama step
#
# x_{k+1} = x_k + f(x_k, t_k) h + g(t_k) sqrt(h) epsilon_k. The drift term is the Euler step of
# Chapter 2 unchanged. The noise term is the only addition, carrying the exact law of a Brownian
# increment over an interval of length h.


def euler_maruyama_step(x_k, f_value, g_value, h: float, epsilon):
    """One step of the scheme, with the Gaussian draw epsilon_k supplied by the caller."""
    drift_part = f_value * h  # Displacement the drift alone would produce.
    noise_part = g_value * math.sqrt(h) * epsilon  # Displacement from this draw.
    return x_k + drift_part + noise_part, drift_part, noise_part


# %% 05. Fresh randomness and the spreading cloud
#
# Each step draws a new epsilon_k, so two runs from the same state give two trajectories. If we
# start from a common point x, the next state is Gaussian with mean x + f h and covariance
# g^2 h I, so the drift moves the center and the noise sets the width. Without a drift the
# variances of independent increments add, so after N steps the variance is N g^2 h = g^2 T and is
# unbounded.


def driftless_cloud(num_paths: int, steps: int, g_value: float, terminal: float, generator=None):
    """Simulate f = 0 on a grid of ``steps`` intervals and return the terminal states. [N]."""
    h = terminal / steps  # Step size, a time.
    x = torch.zeros(num_paths)  # [N] every path starts at the origin.
    for _ in range(steps):  # Independent increments accumulate additively in variance.
        x = x + g_value * brownian_increment(h, (num_paths,), generator)
    return x  # [N] terminal states, with variance g^2 T in the limit of many paths.


# %% 06. A restoring drift
#
# dX_t = -kappa X_t dt + g dW_t. We multiply by the integrating factor e^{kappa t} to clear the
# drift. Under the Ito isometry the variance of the remaining stochastic integral is an ordinary
# integral, so X_t given x_0 is Gaussian with mean e^{-kappa t} x_0 and variance
# g^2 (1 - e^{-2 kappa t})/(2 kappa). The stationary variance g^2/(2 kappa) is the level at
# which the inward pull of the drift offsets the outward spread of the noise.


def ornstein_uhlenbeck_law(x_0: float, kappa: float, g_value: float, t: float) -> dict:
    """The exact conditional mean and variance of the Ornstein-Uhlenbeck process at time t."""
    decay = math.exp(-kappa * t)  # Fraction of the initial signal that remains.
    stationary = g_value ** 2 / (2.0 * kappa)  # Variance the process settles on.
    variance = stationary * (1.0 - decay ** 2)  # Rises from zero toward the stationary value.
    return {"mean": decay * x_0, "variance": variance, "std": math.sqrt(variance),
            "stationary_variance": stationary}


def simulate_ornstein_uhlenbeck(num_paths, steps, x_0, kappa, g_value, t, generator=None):
    """Euler-Maruyama on the same process, for comparison with the closed form. Returns [N]."""
    h = t / steps
    x = torch.full((num_paths,), float(x_0))  # [N] all paths start at x_0.
    for _ in range(steps):
        x = x - kappa * x * h + g_value * brownian_increment(h, (num_paths,), generator)
    return x  # [N] states at time t.


# %% 07. Ito's lemma
#
# d phi(X_t) = (phi_t + f phi_x + (1/2) g^2 phi_xx) dt + g phi_x dW_t. The squared increment
# (Delta X)^2 tends to g^2 Delta t and does not vanish, so the second-order Taylor term is as
# large as the first-order drift term and belongs in the dt part. The chapter's test function is
# phi(x) = x^2 on X_t = W_t, where phi_x = 2x and phi_xx = 2, giving d(W_t^2) = dt + 2 W_t dW_t.


def ito_drift(f_value, g_value, phi_x, phi_xx):
    """The dt bracket of Ito's lemma for a probe with no explicit time dependence."""
    return f_value * phi_x + 0.5 * g_value ** 2 * phi_xx  # Drift term plus curvature correction.


def check_ito_on_square(num_paths: int, steps: int, terminal: float, generator=None) -> dict:
    """Integrate d(W_t^2) = dt + 2 W_t dW_t on a grid and compare with the exact E[W_t^2] = t."""
    h = terminal / steps
    w = torch.zeros(num_paths)  # [N] Brownian position.
    accumulated = torch.zeros(num_paths)  # [N] right-hand side of Ito's lemma, integrated.
    without_curvature = torch.zeros(num_paths)  # [N] the same sum with the dt term dropped.
    for _ in range(steps):
        dw = brownian_increment(h, (num_paths,), generator)  # [N] one increment.
        accumulated = accumulated + ito_drift(0.0, 1.0, 2.0 * w, 2.0) * h + 2.0 * w * dw
        without_curvature = without_curvature + 2.0 * w * dw  # Ordinary chain rule.
        w = w + dw  # Advance the path itself.
    return {
        "realized_mean_square": float((w ** 2).mean()),  # Direct average of W_T^2.
        "ito_mean": float(accumulated.mean()),  # Same quantity through Ito's lemma.
        "no_curvature_mean": float(without_curvature.mean()),  # Near zero, the wrong answer.
        "exact_mean_square": terminal,  # E[W_T^2] = T.
    }


# %% 08. The Fokker-Planck equation
#
# d_t p_t = -div(f p_t) + (1/2) g^2 laplacian(p_t). With f = 0 and g = 1 and a start at N(0, 1),
# the solution stays Gaussian with variance 1 + t. Both sides are then in closed form. Autograd
# supplies the two derivatives, so we can check the identity directly.


def spreading_gaussian_log_density(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """log N(x; 0, 1 + t), the exact solution of d_t p = (1/2) d_x^2 p from N(0, 1)."""
    variance = 1.0 + t  # Variance grows at the rate the diffusion injects it.
    return -0.5 * torch.log(2.0 * math.pi * variance) - x ** 2 / (2.0 * variance)


def fokker_planck_residual(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """d_t p - (1/2) d_x^2 p on the closed-form solution, zero at every (x, t)."""
    x = x.clone().requires_grad_(True)  # [M] evaluation points.
    t = t.clone().requires_grad_(True)  # [M] evaluation times.
    p = spreading_gaussian_log_density(x, t).exp()  # [M] the density itself.
    ones = torch.ones_like(p)
    dp_dt, dp_dx = torch.autograd.grad(p, (t, x), ones, create_graph=True)  # First derivatives.
    d2p_dx2, = torch.autograd.grad(dp_dx, x, torch.ones_like(dp_dx))  # Second derivative in x.
    return (dp_dt - 0.5 * d2p_dx2).detach()  # The equation with f = 0 and g = 1.


def peak_extrapolation(step: float) -> dict:
    """The chapter's first-order estimate of the peak density after a short interval."""
    p_0 = 1.0 / math.sqrt(2.0 * math.pi)  # p_0(0), the height of the standard Gaussian.
    curvature = -p_0  # p_0''(x) = (x^2 - 1) p_0(x), so p_0''(0) = -p_0(0).
    rate = 0.5 * 1.0 * curvature  # d_t p_t(0) at t = 0, with g = 1.
    linear = p_0 + step * rate  # One-step extrapolation from the rate.
    exact = 1.0 / math.sqrt(2.0 * math.pi * (1.0 + step))  # Peak of N(0, 1 + step).
    return {"p_0": p_0, "rate": rate, "linear": linear, "exact": exact, "gap": exact - linear}


# %% 09. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=400, samples=16384)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)  # Printed for consistency; this file runs on the CPU.
    generator = torch.Generator().manual_seed(args.seed)  # Local generator, so the report repeats.

    # The chapter's step: x_k = 2, f = -1, g = 0.5, h = 0.04 and the draw epsilon_k = 1.5.
    h = 0.04
    landed, drift_part, noise_part = euler_maruyama_step(2.0, -1.0, 0.5, h, 1.5)
    scale = increment_scale(h, f_value=-1.0, g_value=0.5)
    three_draws = [euler_maruyama_step(2.0, -1.0, 0.5, h, e)[0] for e in (1.5, -0.5, 0.2)]

    # The cloud without a drift: variance g^2 T after the whole interval.
    terminal_states = driftless_cloud(args.samples, args.steps, 0.5, TERMINAL_TIME, generator)

    # The Ornstein-Uhlenbeck example: x_0 = 2, kappa = 1, g = sqrt(2), evaluated at t = log 2.
    ou_time = math.log(2.0)
    ou = ornstein_uhlenbeck_law(2.0, 1.0, math.sqrt(2.0), ou_time)
    ou_paths = simulate_ornstein_uhlenbeck(
        args.samples, args.steps, 2.0, 1.0, math.sqrt(2.0), ou_time, generator,
    )

    ito = check_ito_on_square(args.samples, args.steps, 0.25, generator)  # t = 0.25.
    grid = torch.linspace(-3.0, 3.0, 61)  # [61] evaluation points for the density equation.
    residual = fokker_planck_residual(grid, torch.full_like(grid, 0.5))  # At t = 0.5.
    peak = peak_extrapolation(0.1)

    report = {
        "euler_maruyama x_{k+1}": landed,  # 2.11 in the notes.
        "drift displacement f h": drift_part,  # -0.04.
        "noise displacement g sqrt(h) eps": noise_part,  # 0.15.
        "typical noise over drift": scale["noise_over_drift"],  # 2.5 at this step size.
        "increment std at h=0.04": scale["std"],  # 0.2.
        "three endpoints, mean": sum(three_draws) / 3.0,  # 2.00 against the drift's 1.96.
        "driftless terminal variance": float(terminal_states.var(unbiased=False)),  # g^2 T = 0.25.
        "OU mean at t=log 2": ou["mean"],  # 1, half the initial signal.
        "OU variance at t=log 2": ou["variance"],  # 0.75 of the stationary 1.
        "OU simulated mean": float(ou_paths.mean()),
        "OU simulated variance": float(ou_paths.var(unbiased=False)),
        "E[W_t^2] at t=0.25, exact": ito["exact_mean_square"],  # 0.25.
        "E[W_t^2] simulated directly": ito["realized_mean_square"],
        "E[W_t^2] through Ito's lemma": ito["ito_mean"],
        "E[W_t^2] without the curvature": ito["no_curvature_mean"],  # Near zero.
        "Fokker-Planck residual, largest": float(residual.abs().max()),  # Zero up to float error.
        "peak p_0(0)": peak["p_0"],  # 0.3989.
        "peak rate at t=0": peak["rate"],  # -0.1995.
        "peak linear estimate at t=0.1": peak["linear"],  # 0.3790.
        "peak exact at t=0.1": peak["exact"],  # 0.3804.
        "second-order gap": peak["gap"],  # 0.0014.
    }
    if not args.quiet:
        print(f"device={device} steps={args.steps} paths={args.samples} seed={args.seed}")
    print_report("Forward corruption without a network", report)
    save_run(args.out, config=vars(args), report=report)


# %% 10. Run the script
if __name__ == "__main__":
    main()