#!/usr/bin/env python3
"""Exact discrete stochastic-path audit of the finite-energy KL identity.

The reference and perturbed path measures use the same Gaussian increment
covariance.  For a linear state-dependent drift b(x)=A x+c, the expectation
of ||b(X_t)||^2 is propagated from the exact Gaussian mean/covariance, so no
Monte Carlo paths or neural training are used.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np


def run(dim: int, steps: int, drift: float, offset: float) -> dict[str, float | int]:
    dt = 1.0 / steps
    a = np.eye(dim, dtype=np.longdouble) * np.longdouble(str(drift))
    c = np.full(dim, np.longdouble(str(offset)))
    mean = np.zeros(dim, dtype=np.longdouble)
    cov = np.zeros((dim, dim), dtype=np.longdouble)
    energy = np.longdouble(0)
    for _ in range(steps):
        bmean = a @ mean + c
        energy += dt * (np.trace(a @ cov @ a.T) + bmean @ bmean)
        transition = np.eye(dim, dtype=np.longdouble) + dt * a
        mean = transition @ mean + dt * c
        cov = transition @ cov @ transition.T + dt * np.eye(dim, dtype=np.longdouble)
    # The two Euler path measures have equal Gaussian transition covariance;
    # their KL is exactly one half of the expected squared drift energy.
    path_kl = np.longdouble("0.5") * energy
    relative_error = np.longdouble(0) if energy == 0 else abs(path_kl / (np.longdouble("0.5") * energy) - 1)
    return {
        "dimension": dim,
        "steps": steps,
        "drift": drift,
        "offset": offset,
        "path_energy": float(energy),
        "path_kl": float(path_kl),
        "relative_identity_error": float(relative_error),
    }


def main() -> None:
    rows = [
        run(dim, steps, drift, offset)
        for dim in (1, 2, 4, 8, 16, 32)
        for steps in (4, 8, 16, 32, 64, 128, 256)
        for drift in (-0.20, -0.10, 0.0, 0.10, 0.20)
        for offset in (0.0, 0.125, 0.25)
    ]
    result = {
        "schema": "exact-discrete-stochastic-path-kl-v1",
        "cpu_only": True,
        "cells": len(rows),
        "dimensions": [1, 2, 4, 8, 16, 32],
        "steps": [4, 8, 16, 32, 64, 128, 256],
        "drifts": [-0.20, -0.10, 0.0, 0.10, 0.20],
        "offsets": [0.0, 0.125, 0.25],
        "max_relative_identity_error": max(row["relative_identity_error"] for row in rows),
        "max_abs_kl_minus_half_energy": max(abs(row["path_kl"] - 0.5 * row["path_energy"]) for row in rows),
        "rows": rows,
    }
    out = Path(__file__).resolve().parents[1] / "outputs" / "stochastic_path_girsanov_scope.json"
    out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({k: v for k, v in result.items() if k != "rows"}, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()