#!/usr/bin/env python3
"""Generate study-specific speaker diagnostics from frozen Figure-3 authorities."""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scienceplots  # noqa: F401


PROJECT_DIR = Path(__file__).resolve().parent
PAPER_DIR = PROJECT_DIR.parent / "paper_apj_v19_cgmsum_draft"
UNIT = r"$10^{-15}$ erg cm$^{-2}$ s$^{-1}$ arcmin$^{-2}$"
COLORS = {
    "blue": "#2676b8",
    "orange": "#d47a2c",
    "green": "#398b5f",
    "purple": "#7553a6",
    "cyan": "#2a95a5",
    "gray": "#6f7782",
    "red": "#b74842",
}


def _read(name: str) -> pd.DataFrame:
    path = PAPER_DIR / name
    if not path.is_file():
        raise FileNotFoundError(path)
    return pd.read_csv(path)


def _weighted_mean(values: np.ndarray, errors: np.ndarray) -> float:
    weights = np.square(errors) ** -1
    return float(np.sum(values * weights) / np.sum(weights))


def _save(fig: plt.Figure, stem: str, column: str) -> None:
    for suffix in ("png", "pdf"):
        fig.savefig(
            PAPER_DIR / f"{stem}_{column}.{suffix}",
            dpi=300 if suffix == "png" else None,
            bbox_inches="tight",
            pad_inches=0.03,
        )
    plt.close(fig)


def _style_axes(axes: np.ndarray | list[plt.Axes]) -> None:
    for axis in np.asarray(axes).ravel():
        axis.grid(alpha=0.22, linewidth=0.5)


def plot_henley(column: str, width: float) -> None:
    data = _read("m31_cgmsum_henley_shelton2013_absorbed_0p5_2p0.csv")
    intrinsic = data["surface_brightness_intrinsic_0p5_2p0_1e12"].to_numpy() / 3.6
    absorbed = data["absorbed_0p5_2p0_fluxunit"].to_numpy()
    nh = data["nhi_1e20_cm-2"].to_numpy()
    quantiles = np.quantile(absorbed, [0.16, 0.25, 0.5, 0.75, 0.84])
    np.testing.assert_allclose(
        quantiles,
        [0.246097, 0.273548, 0.384610, 0.549082, 0.690362],
        rtol=0,
        atol=7e-7,
    )

    fig, axes = plt.subplots(2, 1, figsize=(width, width * 1.26), constrained_layout=True)
    scatter = axes[0].scatter(
        intrinsic,
        absorbed,
        c=nh,
        cmap="viridis",
        s=14 if column == "1col" else 20,
        edgecolor="none",
        alpha=0.85,
    )
    limit = max(float(intrinsic.max()), float(absorbed.max())) * 1.05
    axes[0].plot([0, limit], [0, limit], color=COLORS["gray"], linestyle="--", linewidth=0.8)
    axes[0].set(xlabel=f"Intrinsic 0.5–2.0 brightness [{UNIT}]", ylabel=f"Absorbed brightness [{UNIT}]")
    colorbar = fig.colorbar(scatter, ax=axes[0], pad=0.01)
    colorbar.set_label(r"LAB $N_{\rm H}$ [$10^{20}$ cm$^{-2}$]")

    ordered = np.sort(absorbed)
    ecdf = np.arange(1, len(ordered) + 1) / len(ordered)
    axes[1].step(ordered, ecdf, where="post", color=COLORS["blue"], linewidth=1.4)
    axes[1].axvspan(quantiles[1], quantiles[3], color=COLORS["blue"], alpha=0.16, label="Population IQR")
    axes[1].axvline(quantiles[2], color=COLORS["orange"], linewidth=1.4, label="Median = 0.384610")
    axes[1].axvline(quantiles[0], color=COLORS["gray"], linestyle=":", linewidth=0.9)
    axes[1].axvline(quantiles[4], color=COLORS["gray"], linestyle=":", linewidth=0.9, label="p16–p84")
    axes[1].set(xlabel=f"Absorbed 0.5–2.0 brightness [{UNIT}]", ylabel="ECDF", ylim=(0, 1.02))
    axes[1].legend(loc="lower right", fontsize=4.6 if column == "1col" else 7)
    axes[1].text(
        0.02,
        0.94,
        "N=86 population — not M31 direction",
        transform=axes[1].transAxes,
        va="top",
        weight="bold",
        fontsize=5 if column == "1col" else 8,
    )
    _style_axes(axes)
    fig.suptitle("H&S13 sightline-specific absorption conversion", fontsize="medium", weight="bold")
    _save(fig, "henley_population_diagnostic", column)


def plot_efeds(column: str, width: float) -> None:
    data = _read("m31_cgmsum_efeds2023_absorbed_0p5_2p0.csv")
    soft = data["published_0p3_0p6_1e13_deg-2"].to_numpy() * data[
        "model_ratio_0p5_0p6_to_0p3_0p6"
    ].to_numpy() * 100 / 3600
    hard = data["published_0p6_2p0_1e13_deg-2"].to_numpy() * 100 / 3600
    total = soft + hard
    published = data["flux_0p5_2p0_plot_central_fluxunit"].to_numpy()
    np.testing.assert_allclose(total, published, rtol=0, atol=8e-9)
    labels = ["Negligible\nSWCX", "High\nSWCX"]
    colors = [COLORS["purple"], "#aa8ac8"]

    fig, axes = plt.subplots(2, 1, figsize=(width, width * 1.12), constrained_layout=True, gridspec_kw={"height_ratios": [1.5, 1]})
    x = np.arange(2)
    axes[0].bar(x, hard, color=colors, alpha=0.95, label="Native 0.6–2.0")
    axes[0].bar(x, soft, bottom=hard, color=colors, alpha=0.42, hatch="///", label="Model-retained 0.5–0.6")
    for index, value in enumerate(total):
        axes[0].text(index, value + 0.012, f"{value:.6f}", ha="center", va="bottom", fontsize="small", weight="bold")
    axes[0].set(xticks=x, xticklabels=labels, ylabel=f"Reconstructed 0.5–2.0 [{UNIT}]", ylim=(0, max(total) * 1.28))
    axes[0].legend(
        loc="upper right",
        fontsize=4.2 if column == "1col" else 7,
        bbox_to_anchor=(0.99, 0.88 if column == "1col" else 1.0),
    )
    axes[0].text(
        0.02,
        0.97,
        "Two fixed scenarios — not a confidence interval",
        transform=axes[0].transAxes,
        va="top",
        weight="bold",
        fontsize=5 if column == "1col" else 8,
    )

    axes[1].axis("off")
    headers = ["Scenario", "kT", "Z", "soft fraction", "final [U]"]
    rows = [
        [
            labels[index].replace("\n", " "),
            f"{data.iloc[index]['temperature_keV']:.3f} keV",
            f"{data.iloc[index]['abundance_solar']:.3f}",
            f"{data.iloc[index]['model_ratio_0p5_0p6_to_0p3_0p6']:.6f}",
            f"{total[index]:.6f}",
        ]
        for index in range(2)
    ]
    table = axes[1].table(cellText=rows, colLabels=headers, loc="center", cellLoc="center")
    table.auto_set_font_size(False)
    table.set_fontsize(4.2 if column == "1col" else 8)
    table.scale(1, 1.45)
    fig.suptitle("eFEDS23 split-band scenario reconstruction", fontsize="medium", weight="bold")
    _save(fig, "efeds_scenario_diagnostic", column)


def _merge_signed_radius(data: pd.DataFrame) -> pd.DataFrame:
    primary = _read("m31_cgmsum_v19_primary_measurements_public.csv")[["obsid", "signed_rproj_kpc"]]
    merged = data.merge(primary, on="obsid", how="inner", validate="one_to_one")
    if len(merged) != 14:
        raise ValueError("directional diagnostic requires all fourteen M31 fields")
    return merged.sort_values("signed_rproj_kpc").reset_index(drop=True)


def plot_ueda(column: str, width: float) -> None:
    data = _merge_signed_radius(_read("m31_cgmsum_ueda2022_disk_m31_footprint_absorbed_0p5_2p0.csv"))
    radius = data["signed_rproj_kpc"].to_numpy()
    nominal = data["nominal_absorbed_0p5_2p0_fluxunit"].to_numpy()
    low = data["n0_minus_1sigma_absorbed_0p5_2p0_fluxunit"].to_numpy()
    high = data["n0_plus_1sigma_absorbed_0p5_2p0_fluxunit"].to_numpy()
    errors = data["measurement_staterr_absorbed_0p5_2p0_fluxunit"].to_numpy()
    weights = np.square(errors) ** -1
    weights /= weights.sum()
    mean = _weighted_mean(nominal, errors)
    np.testing.assert_allclose(mean, 0.311904286700407, rtol=0, atol=5e-13)

    fig, axes = plt.subplots(2, 1, sharex=True, figsize=(width, width * 1.18), constrained_layout=True, gridspec_kw={"height_ratios": [1.6, 1]})
    axes[0].vlines(radius, low, high, color=COLORS["green"], alpha=0.45, linewidth=2, label=r"Marginal $n_0$ sensitivity (not joint posterior)")
    axes[0].scatter(radius, nominal, color=COLORS["green"], s=18, zorder=3, label="Nominal field prediction")
    axes[0].axhspan(nominal.min(), nominal.max(), color=COLORS["green"], alpha=0.08, label="Deterministic footprint min–max")
    axes[0].axhline(mean, color=COLORS["orange"], linewidth=1.3, label=f"Matched weighted mean = {mean:.6f}")
    axes[0].set(ylabel=f"Predicted brightness [{UNIT}]")
    axes[0].legend(loc="upper left", fontsize=4.1 if column == "1col" else 6.5)
    axes[0].set_title(
        "14/14 fields inside adopted angular domain",
        loc="left",
        fontsize=5.8 if column == "1col" else 8,
        weight="bold",
        pad=2,
    )

    axes[1].bar(radius, weights, width=1.25, color=COLORS["green"], alpha=0.75)
    axes[1].set(xlabel="Signed projected M31 radius [kpc]", ylabel="Normalized\nstatistical weight")
    _style_axes(axes)
    fig.suptitle("Ueda22 directional footprint and estimator", fontsize="medium", weight="bold")
    _save(fig, "ueda_footprint_diagnostic", column)


def plot_halosat(column: str, width: float) -> None:
    data = _merge_signed_radius(_read("m31_cgmsum_halosat_kaaret2020_disk_adiabatic_halo_m31_extrapolation.csv"))
    if not bool(data["outside_halosat_fit_domain"].all()):
        raise ValueError("HaloSat diagnostic expects all M31 fields outside the fitted domain")
    radius = data["signed_rproj_kpc"].to_numpy()
    nominal = data["nominal_absorbed_0p5_2p0_fluxunit"].to_numpy()
    no_cross = data["no_cross_term_absorbed_0p5_2p0_fluxunit"].to_numpy()
    patchiness = data["patchiness_sigma_absorbed_0p5_2p0_fluxunit"].to_numpy()
    errors = data["measurement_staterr_absorbed_0p5_2p0_fluxunit"].to_numpy()
    mean = _weighted_mean(nominal, errors)
    no_cross_mean = _weighted_mean(no_cross, errors)
    np.testing.assert_allclose(mean, 0.5653859652455517, rtol=0, atol=5e-13)
    np.testing.assert_allclose(no_cross_mean, 0.438683, rtol=0, atol=8e-7)
    cross_share = (nominal - no_cross) / nominal * 100

    fig, axes = plt.subplots(2, 1, sharex=True, figsize=(width, width * 1.18), constrained_layout=True, gridspec_kw={"height_ratios": [1.6, 1]})
    axes[0].vlines(radius, nominal - patchiness, nominal + patchiness, color=COLORS["cyan"], alpha=0.28, linewidth=3, label="Converted predictor scatter (not CI)")
    axes[0].scatter(radius, nominal, color=COLORS["cyan"], s=18, zorder=3, label="Cross term retained")
    axes[0].scatter(radius, no_cross, color=COLORS["gray"], marker="x", s=22, zorder=3, label="No-cross structural check")
    axes[0].axhline(mean, color=COLORS["orange"], linewidth=1.3, label=f"Matched mean = {mean:.6f}")
    axes[0].set(ylabel=f"Predicted brightness [{UNIT}]")
    axes[0].legend(loc="upper left", fontsize=4.1 if column == "1col" else 6.5)
    axes[0].set_title(
        "14/14 M31 fields outside HaloSat fit domain",
        loc="left",
        fontsize=5.8 if column == "1col" else 8,
        weight="bold",
        pad=2,
    )

    axes[1].plot(radius, cross_share, color=COLORS["red"], marker="o", markersize=3, linewidth=1)
    axes[1].axhline(float(np.average(cross_share, weights=np.square(errors) ** -1)), color=COLORS["gray"], linestyle="--", linewidth=0.9)
    axes[1].set(xlabel="Signed projected M31 radius [kpc]", ylabel="Cross-term\nshare [%]")
    _style_axes(axes)
    fig.suptitle("HaloSat20 extrapolation and structural check", fontsize="medium", weight="bold")
    _save(fig, "halosat_footprint_diagnostic", column)


def main() -> None:
    PAPER_DIR.mkdir(parents=True, exist_ok=True)
    with plt.style.context(["science", "no-latex"]):
        for column, width in (("1col", 3.35), ("2col", 7.0)):
            scale = 5.2 if column == "1col" else 8.5
            with plt.rc_context(
                {
                    "font.size": scale,
                    "axes.labelsize": scale,
                    "axes.titlesize": scale + 1,
                    "xtick.labelsize": scale - 0.3,
                    "ytick.labelsize": scale - 0.3,
                    "legend.fontsize": scale - 0.5,
                    "figure.titlesize": scale + 1.5,
                }
            ):
                plot_henley(column, width)
                plot_efeds(column, width)
                plot_ueda(column, width)
                plot_halosat(column, width)
    print(f"Generated 16 reference diagnostic assets in {PAPER_DIR}")


if __name__ == "__main__":
    main()