"""Strong stochastic flow maps, where the driving path is an argument of the map.

A map that matches only the transition kernel draws fresh noise on every leg, so one full step and
two half steps follow different trajectories and their composition residual has a floor at twice
the variance of the interval. A strong map takes the path itself, projected onto the first few
shifted Legendre polynomials. The Chen relations build the coefficients of a long interval from the
coefficients of its halves, so one realization drives all three legs of the cocycle.

Chapter 6, Section 6.31 of the notes. Runs on the linear equation dX = X dt + sigma dW
and on a cloud of starting states from ``cis6270.data.gaussian_mixture_2d``. Prints the weak
map's floor of 2K_1 = 6.389, the Legendre variances 1, 1/3 and 1/5, the Chen fusion that turns
four half-interval draws into (0.50, -0.20), and the pathwise error against the coefficient count.

    python lecture_6/stochastic_map.py                 # 400,000 Monte Carlo draws, a few seconds
    python lecture_6/stochastic_map.py --steps 20      # the floor of 1,000 draws, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

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

torch.set_num_threads(1)  # The work here is elementwise, and one thread keeps the timing stable.
# ``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)  # The checks below sit well under the float32 noise floor.


SIGMA = 1.0  # The noise level of dX = X dt + sigma dW.
ORDER = 3  # Coefficients carried per interval, the number used in practice.
FINEST = 32  # Subintervals we draw the path on before the Chen relations fuse it upward.


# %% 02. The transition kernel, and the weak map built from it
#
# Start at x at time s. The state at t is then e^{t-s} x plus a centered Gaussian whose variance
# is K_{t-s}. A weak map reproduces that law on every interval with one fresh draw.
def kernel_variance(d):
    """K_d = sigma^2 (e^{2d} - 1)/2, the variance the equation accumulates over a length d."""
    return 0.5 * SIGMA ** 2 * (math.exp(2.0 * d) - 1.0)


def weak_step(x, d, noise):
    """The weak map, a function of a state and one fresh standard normal draw."""
    return math.exp(d) * x + math.sqrt(kernel_variance(d)) * noise


# %% 03. The cost of matching only the transition kernel
#
# The variances compose correctly, so the map is exact at the level of distributions. At the level
# of samples the three draws are independent, so the squared gap between one full step and two
# half steps averages to twice the full-interval variance however well the map is trained.
def weak_floor(x: float = 0.3, u: float = 0.5, t: float = 1.0, draws: int = 400_000, seed: int = 0):
    """The worked comparison of Section 6.31, and a Monte Carlo average of the floor."""
    full = weak_step(x, t, 0.8)  # One step with the chapter's third draw.
    half = weak_step(weak_step(x, u, 0.5), t - u, -1.2)  # Two legs with the first two draws.
    carried = math.exp(2.0 * (t - u)) * kernel_variance(u) + kernel_variance(t - u)
    matching = (math.exp(t - u) * math.sqrt(kernel_variance(u)) * 0.5
                - math.sqrt(kernel_variance(t - u)) * 1.2) / math.sqrt(kernel_variance(t))
    generator = torch.Generator().manual_seed(seed)
    e1, e2, e3 = (torch.randn(draws, generator=generator) for _ in range(3))
    gap = (math.sqrt(kernel_variance(t)) * e3
           - math.exp(t - u) * math.sqrt(kernel_variance(u)) * e1
           - math.sqrt(kernel_variance(t - u)) * e2)  # [draws] the difference of the two routes.
    return {
        "K_1": kernel_variance(t), "sqrt K_1": math.sqrt(kernel_variance(t)),
        "sqrt K_0.5": math.sqrt(kernel_variance(u)),
        "one full step": full, "two half steps": half, "gap": full - half,
        "squared gap": (full - half) ** 2,
        "carried variance equals K_1": carried, "floor 2 K_1": 2.0 * kernel_variance(t),
        "the draw that would close it": matching,
        "Monte Carlo floor": float((gap ** 2).mean()),
    }


# %% 04. The path as an input, through the shifted Legendre polynomials
#
# I^{(n)}_{s,t} integrates the n-th shifted Legendre polynomial of the normalized time against the
# increment. The zeroth coefficient is the plain increment, the first separates a path that did
# its travelling early from one that did it late, and distinct coefficients are independent.
def legendre(n: int, r):
    """The n-th shifted Legendre polynomial on [0, 1]."""
    return [torch.ones_like(r), 2.0 * r - 1.0, 6.0 * r * r - 6.0 * r + 1.0][n]


def draw_coefficients(shape, length: float, generator):
    """I^{(n)} ~ N(0, (t-s)/(2n+1)), drawn directly, with no path simulated."""
    scale = torch.tensor([math.sqrt(length / (2 * n + 1)) for n in range(ORDER)])  # [ORDER].
    return torch.randn(*shape, ORDER, generator=generator) * scale  # [..., ORDER].


def kernel_coefficients(length: float, nodes: int = 8001):
    """a_n(h) in e^{h(1-r)} = sum_n a_n P_n(r), and it turns the Ito integral into a finite sum."""
    r = torch.linspace(0.0, 1.0, nodes)
    weight = torch.exp(length * (1.0 - r))  # The kernel e^{t-u} written in normalized time.
    return torch.stack([(2 * n + 1) * torch.trapz(weight * legendre(n, r), r)
                        for n in range(ORDER)])  # [ORDER].


def strong_step(x, length: float, coefficients):
    """The solution map over one interval, driven by the coefficients of the path on it."""
    kernel = kernel_coefficients(length).to(x)  # [ORDER].
    return math.exp(length) * x + SIGMA * (coefficients * kernel).sum(-1)  # Same shape as x.


# %% 05. The Chen relations
#
# A shifted Legendre polynomial of the long interval expands in the shifted Legendre basis of each
# half, so the coefficients fuse by a fixed linear rule with signs alternating in the parity of
# n + m. In the first rule the increments add, and the minus sign in the second is the oddness
# of P_1 about the midpoint.
def chen(first, second):
    """Fuse the coefficients of two equal halves into the coefficients of their union."""
    return torch.stack([
        first[..., 0] + second[..., 0],  # The increments add.
        -0.5 * (first[..., 0] - second[..., 0]) + 0.5 * (first[..., 1] + second[..., 1]),
        0.25 * (first[..., 2] + second[..., 2]) - 0.75 * (first[..., 1] - second[..., 1]),
    ], dim=-1)


def fuse_to(pieces, target: int):
    """Halve the number of intervals by Chen fusion until ``target`` intervals remain."""
    while pieces.shape[-2] > target:
        pieces = chen(pieces[..., 0::2, :], pieces[..., 1::2, :])  # Fuse neighbouring pairs.
    return pieces  # [..., target, ORDER].


def worked_chen():
    """The chapter's four half-interval draws, fused into the coefficients of the unit interval."""
    halves = torch.tensor([[0.80, 0.20, 0.0], [-0.30, 0.50, 0.0]])  # [2, ORDER], the two legs.
    fused = chen(halves[0], halves[1])
    return {"I0 over [0,1]": float(fused[0]), "I1 over [0,1]": float(fused[1]),
            "variance of I1 at h=1": 1.0 / 3.0,
            "fused variance at h=1": 0.25 * 0.5 + 0.25 * 0.5 + 0.25 / 6.0 + 0.25 / 6.0}


# %% 06. The pathwise error against the number of coefficients
#
# A fine simulation gives both the exact Ito integral and the coefficients of the same path, so we
# can measure the truncation on the realization it was taken from.
def truncation_error(paths: int, grid: int, seed: int):
    """Mean absolute pathwise error of a one-step map reading 1, 2 and 3 coefficients."""
    generator = torch.Generator().manual_seed(seed)
    step = 1.0 / grid
    time = torch.arange(grid) * step  # [grid] left endpoints of the substeps.
    increments = torch.randn(paths, grid, generator=generator) * math.sqrt(step)  # [P, grid].
    exact = (torch.exp(1.0 - time)[None, :] * increments).sum(1)  # [P] the Ito integral.
    measured = torch.stack([(legendre(n, time)[None, :] * increments).sum(1)
                            for n in range(ORDER)], dim=1)  # [P, ORDER].
    kernel = kernel_coefficients(1.0)  # [ORDER].
    out = {f"coefficient variance I{n}": float(measured[:, n].var()) for n in range(ORDER)}
    for keep in range(1, ORDER + 1):
        approximate = (measured[:, :keep] * kernel[:keep]).sum(1)  # [P].
        out[f"pathwise error with {keep}"] = float((exact - approximate).abs().mean())
    return out


# %% 07. One path, several step counts
#
# A strong map crosses one trajectory of one realization in any number of hops, so its terminal
# state barely moves with the step count. A weak map redraws at every change of the step count,
# so its terminal state moves by the floor of Section 6.31.
def cross(x, coefficients, steps: int, weak: bool, generator=None):
    """Cross [0, 1] in ``steps`` equal hops, either along the fused path or with fresh draws."""
    pieces = fuse_to(coefficients, steps)  # [..., steps, ORDER] on this grid.
    length = 1.0 / steps
    for k in range(steps):
        if weak:
            noise = torch.randn(x.shape, generator=generator)  # A fresh draw for every leg.
            x = weak_step(x, length, noise)
        else:
            x = strong_step(x, length, pieces[..., k, :])
    return x


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=400_000, samples=512)
    args = parser.parse_args()  # --steps is the Monte Carlo draw count here, since nothing trains.
    set_seed(args.seed)
    pick_device(args.device)  # Every number here is double precision, and not every
    # accelerator carries it, so the work stays on the CPU whatever --device names.

    report = weak_floor(draws=max(args.steps, 1000), seed=args.seed)
    report.update(worked_chen())
    report.update(truncation_error(paths=1024, grid=4096, seed=args.seed))
    if not args.quiet:
        print(f"weak floor 2K_1 = {report['floor 2 K_1']:.3f}, Monte Carlo "
              f"{report['Monte Carlo floor']:.3f}")
        print(f"pathwise error {report['pathwise error with 1']:.4f} with the increment alone "
              f"against {report['pathwise error with 3']:.6f} with three coefficients")

    # %% 09. The same cloud at four step counts
    points = gaussian_mixture_2d(num_points=args.samples, seed=args.seed).double()  # [N, 2].
    generator = torch.Generator().manual_seed(args.seed + 1)
    path = draw_coefficients((args.samples, 2, FINEST), 1.0 / FINEST, generator)  # [N, 2, F, 3].
    reference = cross(points, path, FINEST, weak=False)  # The finest crossing, as the yardstick.
    samples: list[str] = []
    for steps in (2, 4, 8, 16):  # Short of FINEST, the yardstick against itself.
        strong = cross(points, path, steps, weak=False)  # One realization, several hop counts.
        weak_generator = torch.Generator().manual_seed(args.seed + 2)
        weak = cross(points, path, steps, weak=True, generator=weak_generator)
        report[f"{steps} hops: strong spread"] = float((strong - reference).norm(dim=1).mean())
        report[f"{steps} hops: weak spread"] = float((weak - reference).norm(dim=1).mean())
        samples.append(f"{steps:2d} hops  strong {float((strong - reference).norm(dim=1).mean()):.6f}"
                       f"  weak {float((weak - reference).norm(dim=1).mean()):.3f}")

    print_report("Strong stochastic flow maps, Chapter 6", report, samples)
    save_run(args.out, config=vars(args), report=report, samples=samples)


if __name__ == "__main__":
    main()