#!/usr/bin/env python3
"""Check 30/60/90 arcsec common-kernel stability of the Fe-L color morphology."""

from __future__ import annotations

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

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"
CENTER = SkyCoord(189.9976 * u.deg, -11.6231 * u.deg, frame="fk5")
SCALES = {"30arcsec": 12.0, "60arcsec": 24.0, "90arcsec": 36.0}


def analyze() -> tuple[dict[str, dict[str, np.ndarray]], fits.Header, dict[str, object]]:
    with fits.open(MAP) as hdul:
        header = hdul["LOW_COUNTS"].header.copy()
        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",
        )}
    low_background = arrays["LOW_QPB"] + arrays["LOW_SP"] + arrays["LOW_SKY"]
    high_background = arrays["HIGH_QPB"] + arrays["HIGH_SP"] + arrays["HIGH_SKY"]
    products = {}
    for label, sigma in SCALES.items():
        products[label] = felmap.compute_smoothed_ratio(
            counts_low=arrays["LOW_COUNTS"],
            counts_high=arrays["HIGH_COUNTS"],
            exposure_low=arrays["LOW_EXPOSURE"],
            exposure_high=arrays["HIGH_EXPOSURE"],
            background_low=low_background,
            background_high=high_background,
            sigma_pixels=sigma,
            max_fractional_error=0.35,
        )

    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)
    annulus = (radius >= 4.0) & (radius < 7.0)
    common = annulus.copy()
    for product in products.values():
        common &= product["valid"]
    pairs = {}
    labels = list(products)
    for index, first in enumerate(labels):
        for second in labels[index + 1 :]:
            left = products[first]["ratio"][common]
            right = products[second]["ratio"][common]
            result = spearmanr(left, right)
            pairs[f"{first}_vs_{second}"] = {
                "spearman_rho": float(result.statistic),
                "median_absolute_ratio_difference": float(np.median(np.abs(left - right))),
                "p95_absolute_ratio_difference": float(np.percentile(np.abs(left - right), 95)),
            }
    summary = {
        "point_sources_retained": True,
        "annulus": "R4-7 arcmin",
        "common_valid_pixels": int(common.sum()),
        "scales": {label: {"sigma_pixels": value, "sigma_arcsec": 2.5 * value} for label, value in SCALES.items()},
        "pairwise_morphology": pairs,
    }
    return products, header, summary


def plot(products: dict[str, dict[str, np.ndarray]], header: fits.Header, width: str) -> None:
    wcs = WCS(header)
    if width == "1col":
        fig = plt.figure(figsize=(3.35, 8.2))
        grid = (4, 1)
    else:
        fig = plt.figure(figsize=(7.0, 6.0))
        grid = (2, 2)
    finite = np.concatenate([product["ratio"][np.isfinite(product["ratio"])] for product in products.values()])
    vmin, vmax = np.percentile(finite, [5, 95])
    panels = [
        (products["30arcsec"]["ratio"], "shared sigma=30 arcsec", "ratio"),
        (products["60arcsec"]["ratio"], "shared sigma=60 arcsec", "ratio"),
        (products["90arcsec"]["ratio"], "shared sigma=90 arcsec", "ratio"),
        (products["60arcsec"]["fractional_error"], "60 arcsec statistical fractional error", "error"),
    ]
    for index, (data, title, kind) in enumerate(panels, start=1):
        ax = fig.add_subplot(*grid, index, projection=wcs)
        if kind == "ratio":
            image = ax.imshow(data, origin="lower", cmap="viridis", vmin=vmin, vmax=vmax)
            felmap.add_point_source_overlay(ax)
        else:
            image = ax.imshow(data, origin="lower", cmap="cividis", vmin=0.0, vmax=0.35)
        felmap.add_geometry(ax, wcs)
        felmap.configure_wcs_axis(ax)
        ax.set_title(title, fontsize=7)
        colorbar = fig.colorbar(image, ax=ax, pad=0.02, fraction=0.045)
        colorbar.ax.tick_params(labelsize=5)
    fig.suptitle("MOS2 Fe-L common-kernel stability; QPB+native-SP+sky subtracted", fontsize=8)
    fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.97))
    stem = OUT / f"r47_mos2_fel_ratio_common_kernel_stability_{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:
    products, header, summary = analyze()
    for width in ("1col", "2col"):
        plot(products, header, width)
    (OUT / "r47_mos2_fel_ratio_common_kernel_stability_summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()