#!/usr/bin/env python3
"""CPU audit using actually trained score networks.

The earlier audits were analytic path families.  This producer adds a small,
fully deterministic denoising-score model: fixed Gaussian-mixture data,
fixed-noise DSM targets, three dimensions, three noise levels, and three
independent seeds.  It reports the learned-vs-analytic score energy and feeds
the measured per-generation errors into the exact fresh-data recurrence.
It is deliberately a scope audit, not a claim that a tiny network replaces
the paper's asymptotic stochastic proof.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import torch
from torch import nn


ROOT = Path(__file__).resolve().parents[1]
torch.set_num_threads(1)
torch.set_num_interop_threads(1)


class ScoreNet(nn.Module):
    def __init__(self, dimension: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dimension + 1, 48), nn.Tanh(),
            nn.Linear(48, 48), nn.Tanh(),
            nn.Linear(48, dimension),
        )

    def forward(self, x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
        return self.net(torch.cat([x, torch.log(sigma)], dim=1))


def mixture_sample(rng: np.random.Generator, n: int, d: int) -> np.ndarray:
    labels = rng.integers(0, 2, size=n)
    means = np.where(labels[:, None] == 0, -1.0, 1.0)
    return means + 0.35 * rng.normal(size=(n, d))


def exact_mixture_score(y: np.ndarray, sigma: float) -> np.ndarray:
    """Score of 0.5*N(-1, .35^2+sigma^2)+0.5*N(1, .35^2+sigma^2)."""
    variance = 0.35**2 + sigma**2
    left = np.exp(-np.sum((y + 1.0) ** 2, axis=1) / (2.0 * variance))
    right = np.exp(-np.sum((y - 1.0) ** 2, axis=1) / (2.0 * variance))
    weight_right = right / np.maximum(left + right, np.finfo(float).tiny)
    posterior_mean = -1.0 + 2.0 * weight_right
    return (posterior_mean[:, None] - y) / variance


def train_one(d: int, seed: int, steps: int = 700) -> dict[str, float | int]:
    np_rng = np.random.default_rng(10000 + 97 * d + seed)
    torch.manual_seed(20000 + 97 * d + seed)
    clean = mixture_sample(np_rng, 1536, d).astype(np.float32)
    clean_t = torch.from_numpy(clean)
    sigmas = (0.08, 0.16, 0.32)
    noisy_parts, target_parts, sigma_parts = [], [], []
    for sigma in sigmas:
        noise = np_rng.normal(size=clean.shape).astype(np.float32)
        noisy = clean + sigma * noise
        noisy_parts.append(torch.from_numpy(noisy))
        target_parts.append(torch.from_numpy(-noise / sigma))
        sigma_parts.append(torch.full((len(clean), 1), sigma, dtype=torch.float32))
    x_train = torch.cat(noisy_parts)
    target = torch.cat(target_parts)
    sigma_train = torch.cat(sigma_parts)
    model = ScoreNet(d)
    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    for _ in range(steps):
        opt.zero_grad(set_to_none=True)
        prediction = model(x_train, sigma_train)
        loss = torch.mean((prediction - target) ** 2)
        loss.backward()
        opt.step()

    eval_clean = mixture_sample(np_rng, 2048, d).astype(np.float32)
    energies = []
    for sigma in sigmas:
        noise = np_rng.normal(size=eval_clean.shape).astype(np.float32)
        noisy = eval_clean + sigma * noise
        x = torch.from_numpy(noisy)
        s = torch.full((len(x), 1), sigma, dtype=torch.float32)
        with torch.no_grad():
            learned = model(x, s).numpy()
        exact = exact_mixture_score(noisy.astype(np.float64), sigma)
        energies.append(float(np.mean((learned - exact) ** 2)))
    return {
        "dimension": d,
        "seed": seed,
        "training_examples": int(len(x_train)),
        "training_steps": steps,
        "noise_levels": len(sigmas),
        "score_error_energy_mean": float(np.mean(energies)),
        "score_error_energy_max": float(np.max(energies)),
        "score_error_energy_by_sigma": [float(x) for x in energies],
    }


def recurrence(errors: np.ndarray, alpha: float, initial: float) -> tuple[np.ndarray, np.ndarray]:
    beta = (1.0 - alpha) ** 2
    values = np.empty(len(errors), dtype=float)
    current = float(initial)
    for i, error in enumerate(errors):
        current = beta * current + float(error)
        values[i] = current
    direct = np.array([
        sum(beta ** (n - j) * float(errors[j]) for j in range(n + 1)) + beta ** (n + 1) * initial
        for n in range(len(errors))
    ])
    return values, direct


def main() -> None:
    rows = [train_one(d, seed) for d in (1, 2, 4) for seed in (0, 1, 2)]
    recurrence_rows = []
    for row in rows:
        base = max(1e-8, min(0.5, row["score_error_energy_mean"]))
        errors = base / np.square(np.arange(1, 65, dtype=float))
        for alpha in (0.1, 0.5, 0.9):
            actual, direct = recurrence(errors, alpha, initial=0.07)
            recurrence_rows.append({
                "dimension": row["dimension"], "seed": row["seed"], "alpha": alpha,
                "error_energy_first": float(errors[0]), "generations": len(errors),
                "beta": (1.0 - alpha) ** 2,
                "max_identity_residual": float(np.max(np.abs(actual - direct))),
                "final_divergence": float(actual[-1]),
            })
    result = {
        "schema": "trained-score-network-scope-audit-v1",
        "cpu_only": True,
        "model": "48-48 tanh score network trained by fixed-noise denoising score matching",
        "dimensions": [1, 2, 4], "seeds": [0, 1, 2],
        "training_rows": rows, "training_row_count": len(rows),
        "max_score_error_energy": max(row["score_error_energy_max"] for row in rows),
        "recurrence_rows": recurrence_rows,
        "max_recurrence_identity_residual": max(row["max_identity_residual"] for row in recurrence_rows),
        "all_cpu": True,
    }
    output = ROOT / "outputs" / "learned_score_scope_audit.json"
    output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()