# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "autograd>=1.8",
#   "matplotlib>=3.8",
#   "numpy>=1.26",
#   "pandas>=2.2",
#   "scipy>=1.12",
#   "scikit-learn>=1.5",
# ]
# ///
"""Audit Algorithm 1 on the authors' Boston M=N specialization.

The function definitions are executed directly from the immutable author-code
snapshot mounted at /official.  This wrapper records diagnostics that the
released experiment computes but does not persist: the sandwich target,
DQ+exact preconditioner equation residual, a Powell-hybrid root cross-check,
stability radius, and short sampled-path covariance errors.
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import numpy as np
from scipy.linalg import solve_sylvester
from scipy.optimize import root


OFFICIAL_COMMIT = "89194c6c49319f75cc0e46d35d2e6294bd7e8513"
OFFICIAL_ROOT = Path(os.environ.get("SGMCMC_OFFICIAL_ROOT", "/official"))
RESULTS_ROOT = Path(os.environ.get("SGMCMC_RESULTS_ROOT", "/results"))
SOURCE = (
    OFFICIAL_ROOT
    / "experiments/robust_linear_regression/real_data/robust_regression_boston.py"
)
DATA = Path(
    os.environ.get(
        "SGMCMC_BOSTON_CSV",
        str(RESULTS_ROOT / "boston-scaled/boston.csv"),
    )
)
OUTPUT = Path(
    os.environ.get(
        "SGMCMC_AUDIT_OUTPUT",
        str(RESULTS_ROOT / "boston-scaled/algorithm1_equation_audit.json"),
    )
)


def rel_frob(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.linalg.norm(a - b, "fro") / max(np.linalg.norm(b, "fro"), 1e-15))


def main() -> None:
    os.environ["MPLBACKEND"] = "Agg"
    source_text = SOURCE.read_text(encoding="utf-8")
    definitions = source_text.split("# MAIN", 1)[0]
    ns: dict[str, object] = {"__name__": "official_boston_definitions"}
    exec(compile(definitions, str(SOURCE), "exec"), ns)

    np.random.seed(100)
    X, y = ns["load_boston"](str(DATA), log_target=False, add_intercept=False)
    n, d = X.shape
    theta_ols = np.linalg.solve(X.T @ X, X.T @ y)
    residuals = y - X @ theta_ols
    sigma_working = float(np.std(residuals) + 1e-8)

    np.random.seed(999)
    z_samples = np.random.normal(0.0, sigma_working * 5.0, size=50)
    theta_hat, loss_hist = ns["rmsprop"](
        lambda th, xx, yy: ns["beta_loss_fixedZ"](
            th, xx, yy, z_samples, beta=1.5, sigma=sigma_working
        ),
        theta_ols,
        X,
        y,
        lr=0.05,
        n_iter=200,
        print_every=None,
    )
    fisher = ns["empirical_fisher_beta"](
        theta_hat, X, y, z_samples, beta=1.5, sigma=sigma_working
    )
    hessians, jacobian = ns["empirical_hessian_beta"](
        theta_hat, X, y, z_samples, beta=1.5, sigma=sigma_working
    )
    jacobian_inverse = ns["damped_inv"](jacobian, lam=1e-6)
    sandwich = jacobian_inverse @ fisher @ jacobian_inverse / n
    c_raw = ns["compute_C_raw"](hessians, fisher, sandwich)

    rows = []
    for batch_size in (16, int(0.1 * n)):
        c_bar = c_raw / batch_size
        denominator = c_bar + fisher / n
        lambda_official = (
            (fisher @ jacobian_inverse + jacobian_inverse @ fisher)
            @ np.linalg.inv(denominator + 1e-12 * np.eye(d))
            / n
        )

        def equation(preconditioner: np.ndarray) -> np.ndarray:
            left = (
                preconditioner @ jacobian @ sandwich
                + sandwich @ jacobian @ preconditioner
            )
            right = preconditioner @ (
                c_bar + jacobian @ sandwich @ jacobian
            ) @ preconditioner
            return left - right

        official_residual = equation(lambda_official)
        residual_scale = max(
            np.linalg.norm(lambda_official @ jacobian @ sandwich, "fro")
            + np.linalg.norm(sandwich @ jacobian @ lambda_official, "fro"),
            1e-15,
        )
        # For an invertible nonzero solution, X = Lambda^{-1} converts the
        # Riccati equation into the Sylvester equation
        #   (J S) X + X (S J) = Cbar + J S J.
        # Solving this linear equation avoids the trivial Lambda=0 root.
        q_matrix = c_bar + jacobian @ sandwich @ jacobian
        lambda_inverse = solve_sylvester(
            jacobian @ sandwich,
            sandwich @ jacobian,
            q_matrix,
        )
        lambda_sylvester = np.linalg.inv(lambda_inverse)
        sylvester_residual = equation(lambda_sylvester)
        sylvester_scale = max(
            np.linalg.norm(lambda_sylvester @ jacobian @ sandwich, "fro")
            + np.linalg.norm(sandwich @ jacobian @ lambda_sylvester, "fro"),
            1e-15,
        )
        solved = root(
            lambda flat: equation(flat.reshape(d, d)).ravel(),
            lambda_sylvester.ravel(),
            method="hybr",
            options={"maxfev": 20000},
        )
        lambda_root = solved.x.reshape(d, d)

        path_official = ns["sgd_beta_path"](
            theta_hat,
            X,
            y,
            z_samples,
            n_iters=50 * int(n / batch_size),
            batch_size=batch_size,
            lr=1.0,
            pre_matrix=lambda_official,
            beta=1.5,
            sigma=sigma_working,
            seed=260600293 + batch_size,
        )
        path_sylvester = ns["sgd_beta_path"](
            theta_hat,
            X,
            y,
            z_samples,
            n_iters=50 * int(n / batch_size),
            batch_size=batch_size,
            lr=1.0,
            pre_matrix=lambda_sylvester,
            beta=1.5,
            sigma=sigma_working,
            seed=260600293 + batch_size,
        )
        empirical_official = np.cov(
            path_official[len(path_official) // 2 :], rowvar=False, bias=True
        )
        empirical_sylvester = np.cov(
            path_sylvester[len(path_sylvester) // 2 :], rowvar=False, bias=True
        )
        sym_lambda = 0.5 * (lambda_official + lambda_official.T)
        sym_sylvester = 0.5 * (lambda_sylvester + lambda_sylvester.T)
        rows.append(
            {
                "batch_size": batch_size,
                "official_equation_relative_residual": float(
                    np.linalg.norm(official_residual, "fro") / residual_scale
                ),
                "official_equation_max_abs_residual": float(
                    np.max(np.abs(official_residual))
                ),
                "official_lambda_symmetry_relative_error": rel_frob(
                    lambda_official, lambda_official.T
                ),
                "official_lambda_symmetrized_min_eigenvalue": float(
                    np.linalg.eigvalsh(sym_lambda).min()
                ),
                "official_transition_spectral_radius": float(
                    max(abs(np.linalg.eigvals(np.eye(d) - lambda_official @ jacobian)))
                ),
                "sylvester_equation_relative_residual": float(
                    np.linalg.norm(sylvester_residual, "fro") / sylvester_scale
                ),
                "sylvester_equation_max_abs_residual": float(
                    np.max(np.abs(sylvester_residual))
                ),
                "sylvester_lambda_symmetry_relative_error": rel_frob(
                    lambda_sylvester, lambda_sylvester.T
                ),
                "sylvester_lambda_symmetrized_min_eigenvalue": float(
                    np.linalg.eigvalsh(sym_sylvester).min()
                ),
                "sylvester_transition_spectral_radius": float(
                    max(
                        abs(
                            np.linalg.eigvals(
                                np.eye(d) - lambda_sylvester @ jacobian
                            )
                        )
                    )
                ),
                "hybr_success": bool(solved.success),
                "hybr_message": str(solved.message),
                "hybr_nfev": int(solved.nfev),
                "hybr_equation_relative_residual": float(
                    np.linalg.norm(equation(lambda_root), "fro") / residual_scale
                ),
                "hybr_relative_distance_to_official": rel_frob(
                    lambda_root, lambda_official
                ),
                "hybr_relative_distance_to_sylvester": rel_frob(
                    lambda_root, lambda_sylvester
                ),
                "path_steps": int(len(path_official) - 1),
                "official_path_all_finite": bool(np.isfinite(path_official).all()),
                "official_path_covariance_relative_error": rel_frob(
                    empirical_official, sandwich
                ),
                "sylvester_path_all_finite": bool(
                    np.isfinite(path_sylvester).all()
                ),
                "sylvester_path_covariance_relative_error": rel_frob(
                    empirical_sylvester, sandwich
                ),
            }
        )

    payload = {
        "paper": "https://arxiv.org/abs/2606.00293",
        "official_code": (
            "https://github.com/wangyu1369/large-sample-sgmcmc-uq/tree/"
            + OFFICIAL_COMMIT
        ),
        "scope": (
            "Algorithm 1 Boston M=N specialization; numerical audit, not a "
            "generic root-convergence theorem"
        ),
        "dataset": {"rows": int(n), "features": int(d), "offline_M": int(n)},
        "stage_1": {
            "theta_hat_norm": float(np.linalg.norm(theta_hat)),
            "loss_initial": float(loss_hist[0]),
            "loss_final": float(loss_hist[-1]),
            "sandwich_min_eigenvalue": float(np.linalg.eigvalsh(sandwich).min()),
            "sandwich_max_eigenvalue": float(np.linalg.eigvalsh(sandwich).max()),
        },
        "stage_2": rows,
    }
    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()