#!/usr/bin/env python3
"""CPU-only correlated non-Gaussian path audit for Claims 1--3.

The earlier audit used independent product mixtures.  This producer uses
non-separable two-dimensional Gaussian mixtures whose component covariances
have correlations +/-0.7.  Endpoint densities are closed-form mixtures and
the path energy is a closed-form second-moment integral; only deterministic
fixed-grid quadrature is used for KL and chi-squared integrals.
"""

from __future__ import annotations

import hashlib
import json
import math
from pathlib import Path

import numpy as np


PAPER_SHA256 = "fe979c798cd48a5af02f6c647ecc29b2d6a937841adfa8ceedb492b1c2d81583"
FAMILIES = (
    (
        "rho-positive",
        np.array([[-2.0, 0.0], [0.0, 2.0], [2.0, -1.0]]),
        np.array([0.25, 0.50, 0.25]),
        np.array([[[0.45, 0.315], [0.315, 0.45]], [[0.35, -0.245], [-0.245, 0.35]], [[0.55, 0.385], [0.385, 0.55]]]),
    ),
    (
        "rho-negative",
        np.array([[-3.0, 1.0], [0.0, -2.0], [2.0, 2.0]]),
        np.array([0.20, 0.30, 0.50]),
        np.array([[[0.50, -0.35], [-0.35, 0.50]], [[0.40, 0.28], [0.28, 0.40]], [[0.30, -0.21], [-0.21, 0.30]]]),
    ),
    (
        "mixed-correlation",
        np.array([[-2.0, -2.0], [1.0, 0.0], [2.0, 3.0]]),
        np.array([0.30, 0.40, 0.30]),
        np.array([[[0.35, 0.245], [0.245, 0.35]], [[0.60, -0.42], [-0.42, 0.60]], [[0.45, 0.315], [0.315, 0.45]]]),
    ),
)


def paper_sha() -> str:
    return hashlib.sha256((Path(__file__).resolve().parents[1] / "source" / "paper_v1.pdf").read_bytes()).hexdigest()


def density(points: np.ndarray, means: np.ndarray, weights: np.ndarray, covariances: np.ndarray, drift: float) -> np.ndarray:
    result = np.zeros(len(points), dtype=float)
    scale = math.exp(drift)
    variance_add = (math.exp(2.0 * drift) - 1.0) / (2.0 * drift) if abs(drift) > 1e-12 else 1.0
    for mean, weight, covariance in zip(means, weights, covariances):
        transformed_mean = scale * mean
        transformed_covariance = scale * scale * covariance + variance_add * np.eye(2)
        inverse = np.linalg.inv(transformed_covariance)
        determinant = np.linalg.det(transformed_covariance)
        delta = points - transformed_mean
        result += weight * np.exp(-0.5 * np.einsum("ni,ij,nj->n", delta, inverse, delta)) / (2.0 * np.pi * np.sqrt(determinant))
    return result


def audit() -> dict[str, object]:
    if paper_sha() != PAPER_SHA256:
        raise RuntimeError("paper source hash changed")
    axis = np.linspace(-12.0, 12.0, 801)
    grid_x, grid_y = np.meshgrid(axis, axis, indexing="ij")
    points = np.stack([grid_x.ravel(), grid_y.ravel()], axis=1)
    spacing = float(axis[1] - axis[0])
    rows: list[dict[str, float | str]] = []
    for family, means, weights, covariances in FAMILIES:
        reference = np.zeros(len(points), dtype=float)
        for mean, weight, covariance in zip(means, weights, covariances):
            inverse = np.linalg.inv(covariance + np.eye(2))
            determinant = np.linalg.det(covariance + np.eye(2))
            delta = points - mean
            reference += weight * np.exp(-0.5 * np.einsum("ni,ij,nj->n", delta, inverse, delta)) / (2.0 * np.pi * np.sqrt(determinant))
        for drift in (-0.15, -0.10, -0.05, 0.05, 0.10, 0.15):
            perturbed = density(points, means, weights, covariances, drift)
            times = np.linspace(0.0, 1.0, 2001)
            scale = np.exp(drift * times)
            variance_add = (np.exp(2.0 * drift * times) - 1.0) / (2.0 * drift)
            second_moment = []
            for scalar, add in zip(scale, variance_add):
                second_moment.append(sum(weight * (np.trace(scalar * scalar * covariance + add * np.eye(2)) + np.dot(scalar * mean, scalar * mean)) for mean, weight, covariance in zip(means, weights, covariances)))
            energy = float(np.trapezoid(drift * drift * np.asarray(second_moment), times))
            safe_reference = np.maximum(reference, np.finfo(float).tiny)
            safe_perturbed = np.maximum(perturbed, np.finfo(float).tiny)
            mass_reference = float(np.sum(reference) * spacing * spacing)
            mass_perturbed = float(np.sum(perturbed) * spacing * spacing)
            kl = float(np.sum(perturbed * np.log(safe_perturbed / safe_reference)) * spacing * spacing)
            chi2 = float(np.sum(perturbed * perturbed / safe_reference) * spacing * spacing - 1.0)
            rows.append({"family": family, "drift": drift, "endpoint_kl": kl, "endpoint_chi2": chi2, "path_energy": energy, "kl_over_half_energy": kl / (0.5 * energy), "chi2_over_energy": chi2 / energy, "mass_error": max(abs(mass_reference - 1.0), abs(mass_perturbed - 1.0))})
    assert len(rows) == 18
    assert max(row["kl_over_half_energy"] for row in rows) < 1.0
    assert min(row["endpoint_chi2"] for row in rows) > 0.0
    return {
        "schema": "correlated-nonseparable-mixture-path-v1",
        "cpu_only": True,
        "paper_sha256": PAPER_SHA256,
        "cells": len(rows),
        "families": [family[0] for family in FAMILIES],
        "drifts": [-0.15, -0.10, -0.05, 0.05, 0.10, 0.15],
        "max_kl_over_half_energy": max(row["kl_over_half_energy"] for row in rows),
        "min_chi2_over_energy": min(row["chi2_over_energy"] for row in rows),
        "max_chi2_over_energy": max(row["chi2_over_energy"] for row in rows),
        "max_mass_error": max(row["mass_error"] for row in rows),
        "rows": rows,
    }


if __name__ == "__main__":
    print(json.dumps(audit(), indent=2, sort_keys=True))