"""Controlled numerical diagnostic for constrained-Bayesian MCA.

Data are generated exactly from the Gamma/Betaprime class-conditional model.
This checks numerical behavior under the model assumptions; it is not a
clinical validation or a simulation study of statistical performance.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats


EXAMPLE_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = EXAMPLE_DIR.parents[1]
if str(PROJECT_ROOT.parent) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT.parent))

from organized_uncertainty_analysis.bayesian.core import (
    P_with,
    estimate_ci_bundle,
    slope_at_x,
    theta_max,
    x_at_p,
)


SEED = 20260914
SAMPLE_SIZES = (1000, 500, 200, 60)
MCA_DRAWS = int(os.getenv("BAYES_MCA_DIAGNOSTIC_DRAWS", "5000"))
OUTPUT_DIR = Path(
    os.getenv("BAYES_MCA_DIAGNOSTIC_OUTPUT", str(EXAMPLE_DIR / "mca_diagnostic_outputs"))
).resolve()
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Feasible generating parameters: omega, a, b, s, k, vartheta.
OMEGA = 0.30
A = 8.0
B = 5.0
S = 4.0
K = 5.0
VARTTHETA = 1.0
CAP = theta_max(A, B, K, S)
if not (0.0 < VARTTHETA < CAP):
    raise RuntimeError("The selected generating parameters violate the model constraints.")
TRUE_THETA = (OMEGA, A, B, S, K, VARTTHETA, CAP)


def generate_exact_model_data(n: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
    """Draw y from prevalence and x from the matching class distribution."""
    rng = np.random.default_rng(seed)
    y = rng.binomial(1, OMEGA, size=n).astype(int)
    x = np.empty(n, dtype=float)
    nc = y == 0
    ae = ~nc
    x[nc] = rng.gamma(shape=K, scale=VARTTHETA, size=int(nc.sum()))
    x[ae] = stats.betaprime.rvs(A, B, scale=S, size=int(ae.sum()), random_state=rng)
    return np.clip(x, 1e-12, None), y


def interval(values: np.ndarray) -> tuple[float, float]:
    values = np.asarray(values, float)
    values = values[np.isfinite(values)]
    if not len(values):
        return np.nan, np.nan
    lo, hi = np.quantile(values, [0.025, 0.975])
    return float(lo), float(hi)


def run_one(n: int, seed: int):
    x, y = generate_exact_model_data(n, seed)
    x_upper = max(float(np.quantile(x, 0.995) * 1.10), 12.0)
    x_grid = np.linspace(max(1e-5, float(x.min()) * 0.75), x_upper, 400)
    result = estimate_ci_bundle(
        x,
        y,
        f"n={n}",
        x_grid,
        B_nonpar=0,
        B_param=0,
        M_mca=MCA_DRAWS,
        mca_seed=seed + 10000,
        seed=seed,
        use_prior_p=True,
        prior_r=(1.05, 1.05),
        tau=25.0,
        n_starts_clean=5,
        n_starts_boot=1,
        progress_every=max(MCA_DRAWS, 1),
        max_hessian_condition=1e6,
    )

    h = result["hessian_diagnostics"]
    m = result["diagnostics_mca"]
    x50_lo, x50_hi = interval(result["x50_mca"])
    s50_lo, s50_hi = interval(result["s50_mca"])
    band_width = np.asarray(result["mca_hi"]) - np.asarray(result["mca_lo"])

    summary = {
        "n": n,
        "events": int(y.sum()),
        "event_fraction": float(y.mean()),
        "objective": float(result["objective"]),
        "x50_hat": float(result["x50"]),
        "x50_mca_lo": x50_lo,
        "x50_mca_hi": x50_hi,
        "s50_hat": float(result["s50"]),
        "s50_mca_lo": s50_lo,
        "s50_mca_hi": s50_hi,
        "hessian_min_eigenvalue": float(h["minimum_eigenvalue"]),
        "hessian_original_condition": float(h["original_condition_number"]),
        "hessian_positive_definite": bool(h["original_positive_definite"]),
        "hessian_stabilized": bool(h["stabilized"]),
        "hessian_reliable": bool(h["reliable"]),
        "mca_attempted": int(m["attempted"]),
        "mca_accepted": int(m["successful_curve"]),
        "mca_acceptance_fraction": float(m["successful_curve"] / max(m["attempted"], 1)),
        "mca_band_mean_width": float(np.mean(band_width)),
        "mca_band_max_width": float(np.max(band_width)),
    }

    curve_table = pd.DataFrame(
        {
            "n": n,
            "x": x_grid,
            "true_probability": P_with(TRUE_THETA, x_grid),
            "fitted_probability": result["pmap"],
            "mca_lower": result["mca_lo"],
            "mca_upper": result["mca_hi"],
        }
    )
    return summary, curve_table


def main() -> None:
    true_x50 = x_at_p(TRUE_THETA, hi=20.0, hi_max=200.0)
    true_s50 = slope_at_x(TRUE_THETA, true_x50)
    summaries = []
    curves = []
    for index, n in enumerate(SAMPLE_SIZES):
        summary, curve = run_one(n, SEED + index)
        summaries.append(summary)
        curves.append(curve)

    summary_table = pd.DataFrame(summaries)
    summary_table.insert(3, "true_x50", true_x50)
    summary_table.insert(7, "true_s50", true_s50)
    curve_table = pd.concat(curves, ignore_index=True)
    summary_table.to_csv(OUTPUT_DIR / "mca_diagnostic_summary.csv", index=False)
    curve_table.to_csv(OUTPUT_DIR / "mca_diagnostic_curves.csv", index=False)

    fig, axes = plt.subplots(2, 2, figsize=(12, 8), sharey=True)
    for ax, n in zip(axes.flat, SAMPLE_SIZES):
        block = curve_table[curve_table["n"] == n]
        ax.fill_between(block["x"], block["mca_lower"], block["mca_upper"], alpha=0.22)
        ax.plot(block["x"], block["true_probability"], "--", color="tab:orange", label="generating risk")
        ax.plot(block["x"], block["fitted_probability"], color="black", label="MAP fit")
        ax.set_title(f"n={n}")
        ax.set_xlabel("x")
        ax.set_ylim(-0.02, 1.02)
    axes[0, 0].set_ylabel("P(y=1 | x)")
    axes[1, 0].set_ylabel("P(y=1 | x)")
    axes[0, 0].legend(frameon=True)
    fig.tight_layout()
    fig.savefig(OUTPUT_DIR / "mca_diagnostic.png", dpi=300, bbox_inches="tight")
    fig.savefig(OUTPUT_DIR / "mca_diagnostic.pdf", bbox_inches="tight")
    plt.close(fig)

    print(f"True x50={true_x50:.6f}; true s50={true_s50:.6f}")
    print(summary_table.to_string(index=False))
    print(f"\nOutputs: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()