#!/usr/bin/env python3
"""Exact CPU audit of time-inhomogeneous Gaussian path perturbations.

The original claims used identity-covariance endpoint shifts.  This audit
keeps the path calculation closed form but broadens the family: deterministic
time-varying drifts and linear state-dependent drifts are evaluated on several
time grids and independent dimensions.  All recurrences are Fraction-valued;
only the final logarithms/exponentials are converted to floats.
"""

from __future__ import annotations

import argparse
import json
import math
from fractions import Fraction
from pathlib import Path


def deterministic_rows(dim: int, steps: int) -> list[dict[str, object]]:
    dt = Fraction(1, steps)
    profiles = {
        "constant": [Fraction(1, 10)] * steps,
        "alternating": [Fraction(1, 10) if i % 2 == 0 else Fraction(-1, 10) for i in range(steps)],
        "ramp": [Fraction(2 * i - steps + 1, 10 * steps) for i in range(steps)],
        "canceling": [Fraction(1, 5) if i < steps // 2 else Fraction(-1, 5) for i in range(steps)],
    }
    rows = []
    for profile, drift in profiles.items():
        mean = sum(drift, Fraction(0)) * dt
        energy = Fraction(dim) * sum((u * u for u in drift), Fraction(0)) * dt
        endpoint_kl = Fraction(dim) * mean * mean / 2
        rows.append(
            {
                "family": "deterministic_drift",
                "profile": profile,
                "dimension": dim,
                "steps": steps,
                "energy": float(energy),
                "endpoint_kl": float(endpoint_kl),
                "kl_over_half_energy": float(endpoint_kl / (energy / 2)) if energy else 0.0,
                "girsanov_upper_pass": endpoint_kl <= energy / 2,
            }
        )
    return rows


def linear_rows(dim: int, steps: int) -> list[dict[str, object]]:
    dt = Fraction(1, steps)
    profiles = {
        "constant_linear": [Fraction(1, 10)] * steps,
        "alternating_linear": [Fraction(1, 10) if i % 2 == 0 else Fraction(-1, 10) for i in range(steps)],
        "ramped_linear": [Fraction(2 * i - steps + 1, 10 * steps) for i in range(steps)],
    }
    rows = []
    for profile, coefficients in profiles.items():
        variance = Fraction(0)
        path_energy = Fraction(0)
        for coefficient in coefficients:
            path_energy += coefficient * coefficient * variance * dt
            variance = (1 + coefficient * dt) ** 2 * variance + dt
        variance *= dim
        q_variance = Fraction(dim)
        # The scalar recurrence is independent per coordinate, so the exact
        # energy and KL scale linearly with dimension.
        exact_energy = path_energy * dim
        endpoint_kl = (variance / q_variance - 1 - math.log(float(variance / q_variance))) * dim / 2
        variance_ratio = float(variance / q_variance)
        chi_square = (1.0 / math.sqrt(variance_ratio * (2.0 - variance_ratio))) ** dim - 1.0
        rows.append(
            {
                "family": "linear_state_drift",
                "profile": profile,
                "dimension": dim,
                "steps": steps,
                "terminal_variance_ratio": variance_ratio,
                "energy": float(exact_energy),
                "endpoint_kl": float(endpoint_kl),
                "chi_square": chi_square,
                "chi2_over_energy": chi_square / float(exact_energy) if exact_energy else 0.0,
                "girsanov_upper_pass": endpoint_kl <= exact_energy / 2 + 1e-15,
                "chi_square_finite": variance_ratio < 2.0,
            }
        )
    return rows


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    rows = []
    for dim in (1, 2, 8, 32):
        for steps in (4, 8, 16, 32, 64, 128):
            rows.extend(deterministic_rows(dim, steps))
            rows.extend(linear_rows(dim, steps))
    linear = [row for row in rows if row["family"] == "linear_state_drift"]
    summary = {
        "cells": len(rows),
        "dimensions": [1, 2, 8, 32],
        "steps": [4, 8, 16, 32, 64, 128],
        "families": ["deterministic_drift", "linear_state_drift"],
        "max_kl_over_half_energy": max(float(row["kl_over_half_energy"]) for row in rows if "kl_over_half_energy" in row),
        "max_linear_endpoint_kl_over_half_energy": max(
            2 * float(row["endpoint_kl"]) / float(row["energy"]) for row in linear if row["energy"]
        ),
        "chi2_over_energy_range": [min(float(row["chi2_over_energy"]) for row in linear), max(float(row["chi2_over_energy"]) for row in linear)],
        "all_girsanov_upper_pass": all(bool(row["girsanov_upper_pass"]) for row in rows),
        "all_linear_chi_square_finite": all(bool(row["chi_square_finite"]) for row in linear),
        "exact_fraction_recurrence": True,
    }
    (args.output_dir / "path_gaussian_scope.json").write_text(
        json.dumps({"summary": summary, "rows": rows}, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0 if summary["all_girsanov_upper_pass"] and summary["all_linear_chi_square_finite"] else 2


if __name__ == "__main__":
    raise SystemExit(main())