"""Calibrated compact-support anomaly score and numerical verification."""
import json
from pathlib import Path
import numpy as np
from scipy.special import gammaln
from scipy.integrate import quad
from sklearn.metrics import roc_auc_score


def parameters(d, R2):
    if R2 <= d + 2:
        raise ValueError("R2 must exceed d+2")
    gamma = 2.0 / (R2 - d - 2.0)
    a = 1.0 / gamma
    logC = (gammaln(d / 2 + 1 + a) - d / 2 * np.log(np.pi)
            - d / 2 * np.log(R2) - gammaln(1 + a))
    return gamma, a, logC


def compact_score(X, mu, cov, R2, barrier=False):
    """Negative log density; outside support is +inf (or finite training barrier)."""
    X = np.asarray(X)
    L = np.linalg.cholesky(cov)
    v = np.linalg.solve(L, (X - mu).T).T
    r2 = np.sum(v * v, axis=1)
    gamma, a, logC = parameters(X.shape[1], R2)
    u = 1.0 - r2 / R2
    if barrier:
        # Stable surrogate specified in the idea for optimization.
        return -logC - a * np.log(np.maximum(u, 1e-6)) + np.logaddexp(0., r2 - R2), r2
    score = np.full(len(X), np.inf)
    inside = r2 < R2
    score[inside] = -logC - a * np.log1p(-r2[inside] / R2)
    return score, r2


def gaussian_score(X, mu, cov):
    L = np.linalg.cholesky(cov)
    v = np.linalg.solve(L, (X - mu).T).T
    return np.sum(v * v, axis=1)


def sample_compact(rng, n, d, R2):
    """Exact sampler: r2/R2 ~ Beta(d/2, 1+1/gamma), direction uniform."""
    _, a, _ = parameters(d, R2)
    y = rng.beta(d / 2, a + 1, size=n)
    direction = rng.normal(size=(n, d))
    direction /= np.linalg.norm(direction, axis=1, keepdims=True)
    return direction * np.sqrt(R2 * y)[:, None]


def verify(d=5, seed=1013):
    rng = np.random.default_rng(seed)
    radii = [d + 2.5, d + 6, d + 20, d + 80]
    rows = []
    # Predictions: direct normalization is 1 and calibrated E[r^2] is d.
    for R2 in radii:
        gamma, a, logC = parameters(d, R2)
        z = sample_compact(rng, 250000, d, R2)
        r2 = np.sum(z*z, axis=1)
        surface = 2 * np.pi ** (d / 2) / np.exp(gammaln(d / 2))
        integral = quad(lambda t: surface * t ** (d - 1) * np.exp(logC)
                        * (1 - t * t / R2) ** a, 0, np.sqrt(R2),
                        epsabs=1e-10)[0]
        rows.append({"R2": R2, "gamma": gamma, "a": a,
                     "mean_r2": float(r2.mean()), "target_d": d,
                     "max_r2": float(r2.max()),
                     "normalization": float(integral),
                     "predicted_boundary_slope": -a})
    # Prediction: score near boundary has slope -1/gamma against log(epsilon).
    boundary_sweep = []
    eps = np.logspace(-2, -10, 9)
    for R2 in radii:
        _, a, logC = parameters(d, R2)
        x = np.sqrt(R2 * (1 - eps))[:, None] * np.array([[1.] + [0.] * (d - 1)])
        scores, _ = compact_score(x, np.zeros(d), np.eye(d), R2)
        observed = np.polyfit(np.log(eps), scores + logC, 1)[0]
        boundary_sweep.append({"R2": R2, "predicted_slope": -a,
                               "observed_slope": float(observed),
                               "abs_error": float(abs(observed + a))})
    R2 = d + 10
    outside = compact_score(np.array([[np.sqrt(R2) * 1.001] + [0.]*(d-1)]),
                            np.zeros(d), np.eye(d), R2)[0][0]
    rows_boundary = {"sweep": boundary_sweep,
                     "outside_is_infinite": bool(np.isinf(outside))}
    # Affine geometry prediction.
    A = np.array([[1.5, .2, 0, 0, 0], [.1, .8, .1, 0, 0],
                  [0, .1, 1.2, .1, 0], [0, 0, .1, .9, .2],
                  [0, 0, 0, .1, 1.1]])
    z = sample_compact(rng, 2000, d, d + 10)
    cov = A @ A.T
    _, r_orig = compact_score(z, np.zeros(d), np.eye(d), d + 10)
    _, r_aff = compact_score(z @ A.T, np.zeros(d), cov, d + 10)
    affine_err = float(np.max(np.abs(r_orig-r_aff)))
    # Secondary equal-feature OOD comparison.
    nominal = sample_compact(rng, 4000, d, d + 10)
    anomalies = rng.normal(size=(4000, d)) + np.array([2.2, 0, 0, 0, 0])
    X = np.vstack([nominal, anomalies])
    y = np.r_[np.zeros(len(nominal)), np.ones(len(anomalies))]
    gs = gaussian_score(X, np.zeros(d), np.eye(d))
    cs, _ = compact_score(X, np.zeros(d), np.eye(d), d + 10)
    cs_rank = np.where(np.isinf(cs), np.max(np.where(np.isfinite(cs), cs, 0)) + 1e6, cs)
    auc_g = roc_auc_score(y, gs); auc_c = roc_auc_score(y, cs_rank)
    qg = np.quantile(gs[:len(nominal)], .95); qc = np.quantile(cs_rank[:len(nominal)], .95)
    fpr_g = np.mean(gs[len(nominal):] <= qg); fpr_c = np.mean(cs_rank[len(nominal):] <= qc)
    return {"calibration_rows": rows, "boundary": rows_boundary,
            "affine_max_abs_radius_error": affine_err,
            "ood": {"gaussian_auc": float(auc_g), "compact_auc": float(auc_c),
                    "gaussian_anomaly_accept_rate_at_95pct_nominal": float(fpr_g),
                    "compact_anomaly_accept_rate_at_95pct_nominal": float(fpr_c)}}

if __name__ == "__main__":
    out = verify()
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))