#!/usr/bin/env python3
"""CPU-only finite-support audit of the Theorem 3.7 coverage regime."""

from __future__ import annotations

import itertools
import json
import math

import numpy as np
from scipy.optimize import minimize


def sigmoid(z: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-np.clip(z, -40.0, 40.0)))


def loss(logit: np.ndarray, target: np.ndarray, weights: np.ndarray) -> float:
    return float(np.sum(weights * (np.logaddexp(0.0, logit) - target * logit)))


def fit(design: np.ndarray, target: np.ndarray, weights: np.ndarray) -> np.ndarray:
    def objective(theta: np.ndarray) -> tuple[float, np.ndarray]:
        logit = np.sum(design * theta[None, :], axis=1)
        probability = sigmoid(logit)
        value = loss(logit, target, weights)
        gradient = np.sum(design * (weights * (probability - target))[:, None], axis=0)
        return value, gradient

    result = minimize(
        lambda theta: objective(theta),
        np.zeros(design.shape[1], dtype=float),
        jac=True,
        method="L-BFGS-B",
        bounds=[(-20.0, 20.0)] * design.shape[1],
        options={"ftol": 1e-15, "gtol": 1e-12, "maxiter": 500, "maxls": 80},
    )
    if not result.success and np.linalg.norm(result.jac, ord=np.inf) > 2e-8:
        raise RuntimeError(f"logistic fit failed: {result.message}")
    return np.asarray(result.x, dtype=float)


def coverage(assignments: list[int], dimension: int, window: int) -> bool:
    target = set(range(dimension))
    return all(
        set(assignments[start : start + window]) == target
        for start in range(len(assignments) - window + 1)
    )


def audit_cell(dimension: int, trial: int, depth: int, window: int) -> dict:
    base = np.asarray(list(itertools.product((-1.0, 0.0, 1.0), repeat=dimension)))
    latent = base[:, 0]
    raw = np.column_stack([0.8 * latent + 0.2 * base[:, index] for index in range(dimension)])
    weights = np.full(len(raw), 1.0 / len(raw))
    rng = np.random.default_rng(731_000 + 97 * dimension + trial)
    linear = rng.uniform(0.18, 0.65, dimension) * rng.choice((-1.0, 1.0), dimension)
    # Interactions make the full-information linear optimum distinct from the
    # true conditional logit, so one complete feature pass is not exact.
    interaction = 0.8 * raw[:, 0] * raw[:, 1] * raw[:, 2]
    if dimension >= 5:
        interaction += 0.55 * raw[:, 2] * raw[:, 3] * raw[:, 4]
    target = sigmoid(np.sum(raw * linear[None, :], axis=1) + interaction)
    global_design = np.column_stack([np.ones(len(raw)), raw])
    global_theta = fit(global_design, target, weights)
    global_logit = np.sum(global_design * global_theta[None, :], axis=1)
    global_loss = loss(global_logit, target, weights)
    assignments = [index % dimension for index in range(depth)]
    parent_logit = np.zeros(len(raw))
    losses = [loss(parent_logit, target, weights)]
    for feature in assignments:
        design = np.column_stack([np.ones(len(raw)), raw[:, feature], parent_logit])
        theta = fit(design, target, weights)
        parent_logit = np.sum(design * theta[None, :], axis=1)
        losses.append(loss(parent_logit, target, weights))
    excess = losses[-1] - global_loss
    bp_star = float(np.sum(np.abs(global_theta[1:])))
    return {
        "dimension": dimension,
        "trial": trial,
        "D": depth,
        "M": window,
        "M_coverage": coverage(assignments, dimension, window),
        "support_atoms": len(raw),
        "B_p_star": bp_star,
        "B_X": float(np.sqrt(np.max(np.sum(weights[:, None] * raw**2, axis=0)))),
        "global_loss": global_loss,
        "final_excess_bce": excess,
        "upper_bound": bp_star * float(np.sqrt(np.max(np.sum(weights[:, None] * raw**2, axis=0)))) * window / math.sqrt(depth),
        "scaled_excess_sqrt_D_over_M": excess * math.sqrt(depth) / window,
        "all_losses_nonincreasing": all(
            losses[index + 1] <= losses[index] + 3e-10
            for index in range(len(losses) - 1)
        ),
    }


def main() -> None:
    rows = [
        audit_cell(dimension, trial, depth, dimension)
        for dimension in (3, 4, 5, 6, 7)
        for trial in range(4)
        for depth in (dimension, 2 * dimension, 4 * dimension, 8 * dimension, 16 * dimension)
    ]
    summary = {
        "cells": len(rows),
        "support_atom_range": [min(row["support_atoms"] for row in rows), max(row["support_atoms"] for row in rows)],
        "all_m_coverage": all(row["M_coverage"] for row in rows),
        "all_losses_nonincreasing": all(row["all_losses_nonincreasing"] for row in rows),
        "all_upper_bound_certificates": all(row["final_excess_bce"] <= row["upper_bound"] + 1e-10 for row in rows),
        "all_excess_nonnegative": all(row["final_excess_bce"] >= -1e-10 for row in rows),
        "positive_excess_cells": sum(row["final_excess_bce"] > 1e-8 for row in rows),
        "scaled_excess_range": [min(row["scaled_excess_sqrt_D_over_M"] for row in rows), max(row["scaled_excess_sqrt_D_over_M"] for row in rows)],
    }
    print(json.dumps({"schema": "nia-limited-coverage-bce-v1", "summary": summary}, indent=2, sort_keys=True))
    if not all(summary[key] for key in ("all_m_coverage", "all_losses_nonincreasing", "all_upper_bound_certificates")):
        raise SystemExit("limited-coverage audit gate failed")
    if summary["positive_excess_cells"] < 75:
        raise SystemExit("misspecified full-information control was not exercised")


if __name__ == "__main__":
    main()