#!/usr/bin/env python3
"""Integrate the point-source-included Fe-L component maps in R4-7 sectors."""

from __future__ import annotations

import csv
import json
from pathlib import Path

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
from scipy.stats import spearmanr

import build_r47_mos2_fel_ratio_map as felmap

try:
    import scienceplots  # noqa: F401

    plt.style.use(["science", "no-latex"])
except ImportError:
    pass

ROOT = Path(__file__).resolve().parent
OUT = (
    ROOT
    / "joint_spectrum_fitting_2T_basedon_region_v22"
    / "r47_mos2_fel_ratio_map_pilot_20260713"
)
MAP = OUT / "r47_mos2_fel_ratio_point_sources_included_background_corrected.fits"
CONTRACT = OUT / "r47_mos2_fel_background_spectral_contract.json"
FIT_TABLE = (
    ROOT
    / "joint_spectrum_fitting_2T_basedon_region_v22"
    / "r47_eight_axis_centered_v22_fit_diagnostics_spR015full_cxbNormFree_noMOSL3_mosIgn140-200_deep_semantic_hardened_20260712"
    / "r47_eight_axis_centered_v22_temperature_constraints.csv"
)
CALIBRATION = (
    ROOT
    / "joint_spectrum_fitting_2T_basedon_region_v22"
    / "r47_mos2_temperature_proxy_band_optimization_20260712"
    / "r47_mos2_temperature_proxy_calibration.csv"
)
CENTER = SkyCoord(189.9976 * u.deg, -11.6231 * u.deg, frame="fk5")
SECTORS = [
    ("S1", "disk W", 270.0),
    ("S2", "NW transition", 315.0),
    ("S3", "wind N", 0.0),
    ("S4", "NE transition", 45.0),
    ("S5", "disk E", 90.0),
    ("S6", "SE transition", 135.0),
    ("S7", "wind S", 180.0),
    ("S8", "SW transition", 225.0),
]


def circular_distance_deg(angle: np.ndarray, center: float) -> np.ndarray:
    return np.abs((angle - center + 180.0) % 360.0 - 180.0)


def load_fit_temperatures() -> dict[str, dict[str, object]]:
    with FIT_TABLE.open(newline="") as stream:
        return {row["sector_id"]: row for row in csv.DictReader(stream)}


def load_calibration() -> tuple[np.ndarray, np.ndarray]:
    temperatures, ratios = [], []
    with CALIBRATION.open(newline="") as stream:
        for row in csv.DictReader(stream):
            if row["ratio_name"] == "FeL_primary_high_over_low":
                temperatures.append(float(row["kT_keV"]))
                ratios.append(float(row["ratio_mean"]))
    order = np.argsort(ratios)
    return np.asarray(ratios)[order], np.asarray(temperatures)[order]


def ratio_and_error(
    mask: np.ndarray,
    low_counts: np.ndarray,
    high_counts: np.ndarray,
    low_exposure: np.ndarray,
    high_exposure: np.ndarray,
    low_background: np.ndarray,
    high_background: np.ndarray,
) -> tuple[float, float, float, float]:
    raw_low = float(low_counts[mask].sum())
    raw_high = float(high_counts[mask].sum())
    net_low = raw_low - float(low_background[mask].sum())
    net_high = raw_high - float(high_background[mask].sum())
    if net_low <= 0 or net_high <= 0:
        return np.nan, np.nan, net_low, net_high
    rate_low = net_low / float(low_exposure[mask].sum())
    rate_high = net_high / float(high_exposure[mask].sum())
    ratio = rate_high / rate_low
    fractional_error = np.sqrt(raw_high / net_high**2 + raw_low / net_low**2)
    return float(ratio), float(fractional_error), net_low, net_high


def analyze() -> tuple[list[dict[str, object]], dict[str, object]]:
    with fits.open(MAP) as hdul:
        header = hdul["LOW_COUNTS"].header
        arrays = {name: np.asarray(hdul[name].data, dtype=float) for name in (
            "LOW_COUNTS", "HIGH_COUNTS", "LOW_EXPOSURE", "HIGH_EXPOSURE",
            "LOW_QPB", "HIGH_QPB", "LOW_SP", "HIGH_SP", "LOW_SKY", "HIGH_SKY",
        )}
    yy, xx = np.indices(arrays["LOW_COUNTS"].shape, dtype=float)
    world = WCS(header).pixel_to_world(xx, yy)
    radius = CENTER.separation(world).to_value(u.arcmin)
    position_angle = CENTER.position_angle(world).to_value(u.deg) % 360.0
    annulus = (
        (radius >= 4.0)
        & (radius < 7.0)
        & (arrays["LOW_EXPOSURE"] >= 0.2 * arrays["LOW_EXPOSURE"].max())
        & (arrays["HIGH_EXPOSURE"] >= 0.2 * arrays["HIGH_EXPOSURE"].max())
    )
    point_source_excluded = felmap.point_source_exclusion_mask(arrays["LOW_COUNTS"].shape)
    contract = json.loads(CONTRACT.read_text())
    common_low_factor = (
        contract["bands"]["low"]["soft_proton_common_r015_model_counts"]
        / contract["bands"]["low"]["soft_proton_native_mos2_model_counts"]
    )
    common_high_factor = (
        contract["bands"]["high"]["soft_proton_common_r015_model_counts"]
        / contract["bands"]["high"]["soft_proton_native_mos2_model_counts"]
    )
    backgrounds = {
        "raw": (np.zeros_like(radius), np.zeros_like(radius)),
        "qpb": (arrays["LOW_QPB"], arrays["HIGH_QPB"]),
        "full": (
            arrays["LOW_QPB"] + arrays["LOW_SP"] + arrays["LOW_SKY"],
            arrays["HIGH_QPB"] + arrays["HIGH_SP"] + arrays["HIGH_SKY"],
        ),
        "stress": (
            arrays["LOW_QPB"] + arrays["LOW_SKY"] + common_low_factor * arrays["LOW_SP"],
            arrays["HIGH_QPB"] + arrays["HIGH_SKY"] + common_high_factor * arrays["HIGH_SP"],
        ),
    }
    fits_by_sector = load_fit_temperatures()
    calibration_ratio, calibration_kT = load_calibration()
    rows = []
    sector_pixel_total = 0
    for sector_id, role, center_pa in SECTORS:
        mask = annulus & (circular_distance_deg(position_angle, center_pa) < 22.5)
        masked = mask & ~point_source_excluded
        sector_pixel_total += int(mask.sum())
        row: dict[str, object] = {
            "sector_id": sector_id,
            "role": role,
            "sky_pa_deg": center_pa,
            "pixels": int(mask.sum()),
            "point_source_masked_pixels": int(masked.sum()),
            "raw_low_counts": float(arrays["LOW_COUNTS"][mask].sum()),
            "raw_high_counts": float(arrays["HIGH_COUNTS"][mask].sum()),
        }
        for label, (low_background, high_background) in backgrounds.items():
            ratio, fractional_error, net_low, net_high = ratio_and_error(
                mask,
                arrays["LOW_COUNTS"], arrays["HIGH_COUNTS"],
                arrays["LOW_EXPOSURE"], arrays["HIGH_EXPOSURE"],
                low_background, high_background,
            )
            row[f"{label}_ratio"] = ratio
            row[f"{label}_fractional_error"] = fractional_error
            row[f"{label}_net_low_counts"] = net_low
            row[f"{label}_net_high_counts"] = net_high
        masked_ratio, masked_error, masked_net_low, masked_net_high = ratio_and_error(
            masked,
            arrays["LOW_COUNTS"],
            arrays["HIGH_COUNTS"],
            arrays["LOW_EXPOSURE"],
            arrays["HIGH_EXPOSURE"],
            backgrounds["full"][0],
            backgrounds["full"][1],
        )
        row["point_source_masked_full_ratio"] = masked_ratio
        row["point_source_masked_full_fractional_error"] = masked_error
        row["point_source_masked_full_net_low_counts"] = masked_net_low
        row["point_source_masked_full_net_high_counts"] = masked_net_high
        full_ratio = float(row["full_ratio"])
        if calibration_ratio[0] <= full_ratio <= calibration_ratio[-1]:
            row["image_apec_equivalent_kT_keV"] = float(
                np.interp(full_ratio, calibration_ratio, calibration_kT)
            )
        else:
            row["image_apec_equivalent_kT_keV"] = np.nan
        if sector_id in fits_by_sector:
            row["spectral_fit_kT_keV"] = float(fits_by_sector[sector_id]["kT_keV"])
            row["constraint_class"] = fits_by_sector[sector_id]["constraint_class"]
        else:
            row["spectral_fit_kT_keV"] = np.nan
            row["constraint_class"] = "unavailable"
        rows.append(row)

    available = [row for row in rows if np.isfinite(row["spectral_fit_kT_keV"])]
    ratio_rho = spearmanr(
        [row["full_ratio"] for row in available],
        [row["spectral_fit_kT_keV"] for row in available],
    )
    masked_available = [
        row for row in available if np.isfinite(row["point_source_masked_full_ratio"])
    ]
    masked_ratio_rho = spearmanr(
        [row["point_source_masked_full_ratio"] for row in masked_available],
        [row["spectral_fit_kT_keV"] for row in masked_available],
    )
    kt_available = [row for row in available if np.isfinite(row["image_apec_equivalent_kT_keV"])]
    kt_rho = spearmanr(
        [row["image_apec_equivalent_kT_keV"] for row in kt_available],
        [row["spectral_fit_kT_keV"] for row in kt_available],
    )
    summary = {
        "annulus_pixels": int(annulus.sum()),
        "sector_union_pixels": sector_pixel_total,
        "geometry_closure": sector_pixel_total == int(annulus.sum()),
        "full_ratio_vs_spectral_kT_spearman_rho": float(ratio_rho.statistic),
        "full_ratio_vs_spectral_kT_spearman_pvalue": float(ratio_rho.pvalue),
        "point_source_masked_full_ratio_vs_spectral_kT_spearman_rho": float(masked_ratio_rho.statistic),
        "point_source_masked_full_ratio_vs_spectral_kT_spearman_pvalue": float(masked_ratio_rho.pvalue),
        "image_equivalent_kT_vs_spectral_kT_spearman_rho": float(kt_rho.statistic),
        "point_sources_retained": True,
        "s8_spectral_fit": "unavailable",
    }
    return rows, summary


def write_outputs(rows: list[dict[str, object]], summary: dict[str, object]) -> None:
    csv_path = OUT / "r47_mos2_fel_ratio_point_sources_included_sector_comparison.csv"
    with csv_path.open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    (OUT / "r47_mos2_fel_ratio_point_sources_included_sector_comparison_summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )


def plot(rows: list[dict[str, object]], width: str) -> None:
    labels = [row["sector_id"] for row in rows]
    x = np.arange(len(rows))
    if width == "1col":
        fig, axes = plt.subplots(2, 1, figsize=(3.35, 5.6))
    else:
        fig, axes = plt.subplots(1, 2, figsize=(7.0, 3.0))
    for key, label, color in (
        ("raw_ratio", "raw", "0.55"),
        ("qpb_ratio", "QPB-subtracted", "#4C78A8"),
        ("full_ratio", "QPB+SP+sky", "#E45756"),
        ("stress_ratio", "common-SP stress", "#72B7B2"),
        ("point_source_masked_full_ratio", "full + point mask", "#F2CF5B"),
    ):
        axes[0].plot(x, [row[key] for row in rows], marker="o", ms=3, lw=0.8, label=label, color=color)
    axes[0].set_xticks(x, labels)
    axes[0].set_ylabel(r"Integrated $R_{\rm FeL}$")
    axes[0].set_xlabel("R4–R7 sector")
    axes[0].legend(fontsize=6, ncol=2)

    fit_kT = np.asarray([row["spectral_fit_kT_keV"] for row in rows], dtype=float)
    image_kT = np.asarray([row["image_apec_equivalent_kT_keV"] for row in rows], dtype=float)
    valid = np.isfinite(fit_kT) & np.isfinite(image_kT)
    axes[1].scatter(fit_kT[valid], image_kT[valid], c=np.arange(len(rows))[valid], cmap="viridis", s=24)
    for index in np.where(valid)[0]:
        axes[1].annotate(labels[index], (fit_kT[index], image_kT[index]), fontsize=6, xytext=(2, 2), textcoords="offset points")
    limits = [0.55, 1.1]
    axes[1].plot(limits, limits, color="0.5", ls="--", lw=0.7)
    axes[1].set_xlim(limits)
    axes[1].set_ylim(limits)
    axes[1].set_xlabel("Full spectral-fit kT (keV)")
    axes[1].set_ylabel("Image APEC-equivalent kT (keV)")
    axes[1].text(0.03, 0.96, "point sources retained", transform=axes[1].transAxes, va="top", fontsize=6)
    fig.suptitle("MOS2 Fe-L image-sector cross-check", fontsize=8)
    fig.tight_layout()
    stem = OUT / f"r47_mos2_fel_ratio_point_sources_included_sector_comparison_{width}"
    fig.savefig(stem.with_suffix(".png"), dpi=320, bbox_inches="tight")
    fig.savefig(stem.with_suffix(".pdf"), bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    rows, summary = analyze()
    write_outputs(rows, summary)
    for width in ("1col", "2col"):
        plot(rows, width)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()