"""Direct numerical contract for the resolution limit in Theorem 4.2.

This module uses only deterministic quadrature.  The discrete measures are
positive midpoint rules that weakly converge to a positive continuous measure
on [0, 1].  The continuous reference is computed independently with
Gauss--Legendre quadrature at two orders.
"""

from __future__ import annotations

import csv
import hashlib
import json
import os
import platform
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
from numpy.polynomial.legendre import leggauss
from scipy.spatial.distance import cdist


ROOT = Path(__file__).resolve().parents[2]
ARTIFACT_DIR = ROOT / ".openresearch" / "artifacts" / "claim_2"
FIXED_COMMAND = "uv run python repro/src/verify.py"


def _kernel(left: np.ndarray, right: np.ndarray, name: str, sigma: float) -> np.ndarray:
    distances = cdist(left[:, None], right[:, None], metric="euclidean")
    if name == "gaussian":
        return np.exp(-(distances**2) / (2.0 * sigma**2))
    if name == "exponential":
        return np.exp(-distances / sigma)
    raise ValueError(f"unsupported kernel: {name}")


def _density(points: np.ndarray, name: str) -> np.ndarray:
    if name == "uniform":
        return np.ones_like(points)
    if name == "affine":
        # Strictly positive on the closed domain and integrates to one.
        return 0.5 + points
    if name == "oscillatory":
        # Strict lower bound 0.55; integral is one.
        return 1.0 + 0.45 * np.cos(2.0 * np.pi * points)
    raise ValueError(f"unsupported density: {name}")


def _signal(points: np.ndarray, name: str) -> np.ndarray:
    if name == "constant":
        return np.ones_like(points)
    if name == "linear":
        return points
    if name == "quadratic":
        return points**2
    if name == "sin_2pi":
        return np.sin(2.0 * np.pi * points)
    if name == "cos_3pi":
        return np.cos(3.0 * np.pi * points)
    raise ValueError(f"unsupported signal: {name}")


def _sinkhorn_scaling(
    points: np.ndarray,
    weights: np.ndarray,
    kernel_name: str,
    sigma: float,
    tolerance: float = 2e-14,
    max_iterations: int = 2_000,
) -> tuple[np.ndarray, int, float]:
    kernel = _kernel(points, points, kernel_name, sigma)
    scaling = np.ones(points.size, dtype=np.float64)
    residual = float("inf")
    for iteration in range(1, max_iterations + 1):
        denominator = kernel @ (weights * scaling)
        scaling = np.sqrt(scaling / np.maximum(denominator, 1e-300))
        residual = float(
            np.max(np.abs(scaling * (kernel @ (weights * scaling)) - 1.0))
        )
        if residual < tolerance:
            return scaling, iteration, residual
    raise RuntimeError(
        f"Sinkhorn failed: kernel={kernel_name} n={points.size} residual={residual}"
    )


def _operator_on_probes(
    nodes: np.ndarray,
    weights: np.ndarray,
    scaling: np.ndarray,
    probes: np.ndarray,
    signal_name: str,
    kernel_name: str,
    sigma: float,
) -> tuple[np.ndarray, float]:
    cross_kernel = _kernel(probes, nodes, kernel_name, sigma)
    probe_denominator = cross_kernel @ (weights * scaling)
    probe_scaling = 1.0 / np.maximum(probe_denominator, 1e-300)
    values = probe_scaling * (
        cross_kernel @ (weights * scaling * _signal(nodes, signal_name))
    )
    constant_residual = float(
        np.max(
            np.abs(
                probe_scaling * (cross_kernel @ (weights * scaling))
                - np.ones(probes.size)
            )
        )
    )
    return values, constant_residual


def _midpoint_measure(n: int, density_name: str) -> tuple[np.ndarray, np.ndarray]:
    nodes = (np.arange(n, dtype=np.float64) + 0.5) / n
    weights = _density(nodes, density_name) / n
    # Normalization removes only finite quadrature error, preserves positivity,
    # and does not affect weak convergence.
    weights /= weights.sum()
    return nodes, weights


def _quantile_measure(n: int, density_name: str) -> tuple[np.ndarray, np.ndarray]:
    probabilities = (np.arange(n, dtype=np.float64) + 0.5) / n
    if density_name == "uniform":
        nodes = probabilities
    elif density_name == "affine":
        # Invert F(x) = (x + x^2) / 2.
        nodes = 0.5 * (-1.0 + np.sqrt(1.0 + 8.0 * probabilities))
    elif density_name == "oscillatory":
        # Invert F(x) = x + 0.45 sin(2 pi x)/(2 pi) by bisection.
        lower = np.zeros(n, dtype=np.float64)
        upper = np.ones(n, dtype=np.float64)
        for _ in range(60):
            middle = 0.5 * (lower + upper)
            cdf = middle + 0.45 * np.sin(2.0 * np.pi * middle) / (
                2.0 * np.pi
            )
            lower = np.where(cdf < probabilities, middle, lower)
            upper = np.where(cdf >= probabilities, middle, upper)
        nodes = 0.5 * (lower + upper)
    else:
        raise ValueError(f"unsupported density: {density_name}")
    return nodes, np.full(n, 1.0 / n, dtype=np.float64)


def _discrete_measure(
    n: int, density_name: str, discretization: str
) -> tuple[np.ndarray, np.ndarray]:
    if discretization == "midpoint":
        return _midpoint_measure(n, density_name)
    if discretization == "quantile":
        return _quantile_measure(n, density_name)
    raise ValueError(f"unsupported discretization: {discretization}")


def _gauss_measure(order: int, density_name: str) -> tuple[np.ndarray, np.ndarray]:
    canonical_nodes, canonical_weights = leggauss(order)
    nodes = 0.5 * (canonical_nodes + 1.0)
    weights = 0.5 * canonical_weights * _density(nodes, density_name)
    weights /= weights.sum()
    return nodes, weights


def _fit_slope(resolutions: list[int], errors: list[float]) -> float:
    safe = np.maximum(np.asarray(errors, dtype=np.float64), 1e-18)
    return float(np.polyfit(np.log(np.asarray(resolutions)), np.log(safe), 1)[0])


def _reference_outputs(
    order: int,
    density_name: str,
    kernel_name: str,
    sigma: float,
    signals: list[str],
    probes: np.ndarray,
) -> tuple[dict[str, np.ndarray], dict[str, float], int, float]:
    nodes, weights = _gauss_measure(order, density_name)
    scaling, iterations, residual = _sinkhorn_scaling(
        nodes, weights, kernel_name, sigma
    )
    values: dict[str, np.ndarray] = {}
    constant_residuals: dict[str, float] = {}
    for signal_name in signals:
        values[signal_name], constant_residuals[signal_name] = _operator_on_probes(
            nodes,
            weights,
            scaling,
            probes,
            signal_name,
            kernel_name,
            sigma,
        )
    return values, constant_residuals, iterations, residual


def _git_sha() -> str:
    result = subprocess.run(
        ["git", "rev-parse", "HEAD"],
        cwd=ROOT,
        check=True,
        capture_output=True,
        text=True,
    )
    return result.stdout.strip()


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def run_resolution_contract(config_path: Path | None = None) -> dict:
    start = time.perf_counter()
    if config_path is None:
        config_path = ROOT / "repro" / "config.json"
    config = json.loads(config_path.read_text(encoding="utf-8"))
    probes = np.linspace(0.0, 1.0, int(config["probe_count"]))
    resolutions = [int(value) for value in config["resolutions"]]
    reference_orders = [int(value) for value in config["reference_orders"]]
    signals = [str(value) for value in config["signals"]]
    cases: list[dict] = []

    for density_name in config["densities"]:
        for discretization in config["discretizations"]:
            for kernel_spec in config["kernels"]:
                kernel_name = str(kernel_spec["name"])
                sigma = float(kernel_spec["sigma"])
                coarse_reference, coarse_constant, coarse_iterations, coarse_residual = (
                    _reference_outputs(
                        reference_orders[0],
                        density_name,
                        kernel_name,
                        sigma,
                        signals,
                        probes,
                    )
                )
                fine_reference, fine_constant, fine_iterations, fine_residual = (
                    _reference_outputs(
                        reference_orders[1],
                        density_name,
                        kernel_name,
                        sigma,
                        signals,
                        probes,
                    )
                )
                reference_stability = {
                    signal_name: float(
                        np.max(
                            np.abs(
                                coarse_reference[signal_name]
                                - fine_reference[signal_name]
                            )
                        )
                    )
                    for signal_name in signals
                }
                error_curves = {signal_name: [] for signal_name in signals}
                constant_residuals: list[float] = []
                scaling_records: list[dict] = []
                for n in resolutions:
                    nodes, weights = _discrete_measure(
                        n, density_name, discretization
                    )
                    scaling, iterations, residual = _sinkhorn_scaling(
                        nodes, weights, kernel_name, sigma
                    )
                    scaling_records.append(
                        {
                            "n": n,
                            "iterations": iterations,
                            "node_residual": residual,
                            "min_weight": float(weights.min()),
                            "min_scaling": float(scaling.min()),
                        }
                    )
                    for signal_name in signals:
                        values, constant_residual = _operator_on_probes(
                            nodes,
                            weights,
                            scaling,
                            probes,
                            signal_name,
                            kernel_name,
                            sigma,
                        )
                        error_curves[signal_name].append(
                            float(
                                np.max(
                                    np.abs(
                                        values - fine_reference[signal_name]
                                    )
                                )
                            )
                        )
                        constant_residuals.append(constant_residual)

                signal_records = []
                for signal_name in signals:
                    errors = error_curves[signal_name]
                    signal_records.append(
                        {
                            "signal": signal_name,
                            "uniform_errors": errors,
                            "final_error": errors[-1],
                            "reduction_ratio": errors[-1]
                            / max(errors[0], 1e-300),
                            "loglog_slope": _fit_slope(resolutions, errors),
                            "reference_stability": reference_stability[
                                signal_name
                            ],
                        }
                    )
                cases.append(
                    {
                        "density": density_name,
                        "discretization": discretization,
                        "kernel": kernel_name,
                        "sigma": sigma,
                        "reference": {
                            "orders": reference_orders,
                            "coarse_iterations": coarse_iterations,
                            "fine_iterations": fine_iterations,
                            "coarse_node_residual": coarse_residual,
                            "fine_node_residual": fine_residual,
                            "max_constant_residual": max(
                                list(coarse_constant.values())
                                + list(fine_constant.values())
                            ),
                        },
                        "scalings": scaling_records,
                        "max_probe_constant_residual": max(
                            constant_residuals
                        ),
                        "signals": signal_records,
                    }
                )

    # This deliberately violates resolution increase: the same n=32 result is
    # repeated at every nominal t.  It must not satisfy the convergence gates.
    negative_cases = []
    for case in cases:
        negative_signals = []
        for record in case["signals"]:
            repeated = [record["uniform_errors"][0]] * len(resolutions)
            negative_signals.append(
                {
                    "signal": record["signal"],
                    "uniform_errors": repeated,
                    "final_error": repeated[-1],
                    "reduction_ratio": 1.0,
                    "loglog_slope": _fit_slope(resolutions, repeated),
                    "reference_stability": record["reference_stability"],
                }
            )
        negative_cases.append(
            {
                "density": case["density"],
                "discretization": case["discretization"],
                "kernel": case["kernel"],
                "sigma": case["sigma"],
                "reference": case["reference"],
                "max_probe_constant_residual": case[
                    "max_probe_constant_residual"
                ],
                "signals": negative_signals,
            }
        )

    result = {
        "schema_version": 1,
        "claim": "Theorem 4.2 resolution convergence",
        "variant": config["resolution_variant"],
        "domain": "[0,1]",
        "resolutions": resolutions,
        "probe_count": int(config["probe_count"]),
        "fixed_command": FIXED_COMMAND,
        "seed": int(config["seed"]),
        "cases": cases,
        "negative_control": {
            "name": "fixed_resolution_relabelled_as_increasing",
            "expected_to_pass": False,
            "cases": negative_cases,
        },
        "runtime_seconds": time.perf_counter() - start,
    }
    return result


def write_artifacts(result: dict) -> None:
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    raw_json = ARTIFACT_DIR / "raw_results.json"
    raw_json.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")

    with (ARTIFACT_DIR / "raw_results.csv").open(
        "w", newline="", encoding="utf-8"
    ) as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=[
                "variant",
                "density",
                "discretization",
                "kernel",
                "sigma",
                "signal",
                "n",
                "uniform_error",
            ],
        )
        writer.writeheader()
        for case in result["cases"]:
            for record in case["signals"]:
                for n, error in zip(result["resolutions"], record["uniform_errors"]):
                    writer.writerow(
                        {
                            "variant": result["variant"],
                            "density": case["density"],
                            "discretization": case["discretization"],
                            "kernel": case["kernel"],
                            "sigma": case["sigma"],
                            "signal": record["signal"],
                            "n": n,
                            "uniform_error": f"{error:.17g}",
                        }
                    )

    environment = {
        "git_sha": _git_sha(),
        "fixed_command": FIXED_COMMAND,
        "python": sys.version,
        "platform": platform.platform(),
        "processor": platform.processor(),
        "logical_cpu_count": os.cpu_count(),
        "numpy": np.__version__,
        "config_sha256": _sha256(ROOT / "repro" / "config.json"),
        "lock_sha256": _sha256(ROOT / "uv.lock"),
        "seed": result["seed"],
        "runtime_seconds": result["runtime_seconds"],
    }
    (ARTIFACT_DIR / "environment.json").write_text(
        json.dumps(environment, indent=2) + "\n", encoding="utf-8"
    )


def print_summary(result: dict) -> None:
    print("C2_RESOLUTION_VARIANT=" + str(result["variant"]))
    print("C2_FIXED_COMMAND=" + FIXED_COMMAND)
    print("C2_RESOLUTIONS=" + json.dumps(result["resolutions"]))
    for case in result["cases"]:
        for record in case["signals"]:
            print(
                "C2_RESULT "
                f"density={case['density']} discretization={case['discretization']} "
                f"kernel={case['kernel']} "
                f"sigma={case['sigma']:.6g} signal={record['signal']} "
                f"errors={json.dumps(record['uniform_errors'])} "
                f"ratio={record['reduction_ratio']:.8g} "
                f"slope={record['loglog_slope']:.8g} "
                f"ref_stability={record['reference_stability']:.8g}"
            )
    print(f"C2_RUNTIME_SECONDS={result['runtime_seconds']:.6f}")