#!/usr/bin/env python3
"""Build MOS2 Fe-L high/low color maps with one shared spatial kernel."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import astropy.units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
from astropy.wcs.utils import proj_plane_pixel_area
from matplotlib.colors import LogNorm
from matplotlib.patches import Ellipse
from reproject import reproject_exact
from scipy.ndimage import gaussian_filter
from scipy.signal import fftconvolve

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"
)
QUICK = OUT / "quick_raw_exposure_rpc_authoritative"
SPECTRAL_CONTRACT = OUT / "r47_mos2_fel_background_spectral_contract.json"
ESAS_LOW = OUT / "esas_low_700_875_fullfov_nomask"
ESAS_HIGH = OUT / "esas_high_875_1050_fullfov_nomask"
POINT_SOURCE_REGIONS = (
    ROOT
    / "data/0900170101/rpc_hermes_v2/4.background/"
    "bkg_mos2_00700-01000_R19_nomask/mos2U005-bkg_region-sky.fits"
)
M104_RA_DEG = 189.9976
M104_DEC_DEG = -11.6231
PIXEL_SCALE_ARCSEC = 2.5


def gaussian_kernel_2d(sigma_pixels: float, truncate: float = 4.0) -> np.ndarray:
    radius = max(1, int(np.ceil(truncate * sigma_pixels)))
    axis = np.arange(-radius, radius + 1, dtype=float)
    xx, yy = np.meshgrid(axis, axis)
    kernel = np.exp(-0.5 * (xx**2 + yy**2) / sigma_pixels**2)
    return kernel / kernel.sum()


def compute_smoothed_ratio(
    *,
    counts_low: np.ndarray,
    counts_high: np.ndarray,
    exposure_low: np.ndarray,
    exposure_high: np.ndarray,
    background_low: np.ndarray | None = None,
    background_high: np.ndarray | None = None,
    sigma_pixels: float,
    min_native_exposure_fraction: float = 0.2,
    min_kernel_support_fraction: float = 0.6,
    max_fractional_error: float = 0.5,
) -> dict[str, np.ndarray]:
    """Subtract components, then apply the same Gaussian support to both bands."""

    arrays = [counts_low, counts_high, exposure_low, exposure_high]
    shape = np.asarray(counts_low).shape
    if any(np.asarray(array).shape != shape for array in arrays):
        raise ValueError("Counts and exposure maps must have identical shapes")
    if sigma_pixels <= 0:
        raise ValueError("sigma_pixels must be positive")

    counts_low = np.nan_to_num(np.asarray(counts_low, dtype=float), nan=0.0)
    counts_high = np.nan_to_num(np.asarray(counts_high, dtype=float), nan=0.0)
    exposure_low = np.nan_to_num(np.asarray(exposure_low, dtype=float), nan=0.0)
    exposure_high = np.nan_to_num(np.asarray(exposure_high, dtype=float), nan=0.0)
    if background_low is None:
        background_low = np.zeros(shape, dtype=float)
    if background_high is None:
        background_high = np.zeros(shape, dtype=float)
    background_low = np.nan_to_num(np.asarray(background_low, dtype=float), nan=0.0)
    background_high = np.nan_to_num(np.asarray(background_high, dtype=float), nan=0.0)

    net_low = counts_low - background_low
    net_high = counts_high - background_high
    smooth = lambda array: gaussian_filter(array, sigma_pixels, mode="constant", cval=0.0)
    smooth_net_low = smooth(net_low)
    smooth_net_high = smooth(net_high)
    smooth_exposure_low = smooth(exposure_low)
    smooth_exposure_high = smooth(exposure_high)

    with np.errstate(divide="ignore", invalid="ignore"):
        rate_low = smooth_net_low / smooth_exposure_low
        rate_high = smooth_net_high / smooth_exposure_high
        ratio_all = rate_high / rate_low

    kernel = gaussian_kernel_2d(sigma_pixels)
    variance_low = fftconvolve(np.clip(counts_low, 0.0, None), kernel**2, mode="same")
    variance_high = fftconvolve(np.clip(counts_high, 0.0, None), kernel**2, mode="same")
    with np.errstate(divide="ignore", invalid="ignore"):
        fractional_error = np.sqrt(
            variance_high / smooth_net_high**2 + variance_low / smooth_net_low**2
        )

    low_threshold = min_native_exposure_fraction * float(np.nanmax(exposure_low))
    high_threshold = min_native_exposure_fraction * float(np.nanmax(exposure_high))
    low_native_support = (exposure_low >= low_threshold).astype(float)
    high_native_support = (exposure_high >= high_threshold).astype(float)
    normalization = smooth(np.ones(shape, dtype=float))
    low_kernel_support = smooth(low_native_support) / normalization
    high_kernel_support = smooth(high_native_support) / normalization
    valid = (
        np.isfinite(ratio_all)
        & np.isfinite(fractional_error)
        & (low_kernel_support >= min_kernel_support_fraction)
        & (high_kernel_support >= min_kernel_support_fraction)
        & (smooth_net_low > 0)
        & (smooth_net_high > 0)
        & (fractional_error <= max_fractional_error)
    )
    ratio = np.where(valid, ratio_all, np.nan)
    fractional_error = np.where(valid, fractional_error, np.nan)
    rate_low = np.where(valid, rate_low, np.nan)
    rate_high = np.where(valid, rate_high, np.nan)

    return {
        "ratio": ratio,
        "fractional_error": fractional_error,
        "rate_low": rate_low,
        "rate_high": rate_high,
        "smooth_net_low": smooth_net_low,
        "smooth_net_high": smooth_net_high,
        "smooth_exposure_low": smooth_exposure_low,
        "smooth_exposure_high": smooth_exposure_high,
        "kernel_support": np.minimum(low_kernel_support, high_kernel_support),
        "valid": valid,
    }


def compute_adaptive_smoothed_ratio(
    *,
    counts_low: np.ndarray,
    counts_high: np.ndarray,
    exposure_low: np.ndarray,
    exposure_high: np.ndarray,
    background_low: np.ndarray | None = None,
    background_high: np.ndarray | None = None,
    sigma_grid_pixels: tuple[float, ...] = (4.0, 6.0, 8.0, 12.0, 16.0, 24.0, 32.0),
    target_fractional_error: float = 0.20,
    maximum_fractional_error: float = 0.35,
) -> dict[str, np.ndarray]:
    """Select one common local smoothing scale for both energy bands."""

    if tuple(sorted(sigma_grid_pixels)) != sigma_grid_pixels:
        raise ValueError("sigma_grid_pixels must be increasing")
    products = [
        compute_smoothed_ratio(
            counts_low=counts_low,
            counts_high=counts_high,
            exposure_low=exposure_low,
            exposure_high=exposure_high,
            background_low=background_low,
            background_high=background_high,
            sigma_pixels=sigma,
            max_fractional_error=np.inf,
        )
        for sigma in sigma_grid_pixels
    ]
    shape = np.asarray(counts_low).shape
    keys = (
        "ratio",
        "fractional_error",
        "rate_low",
        "rate_high",
        "smooth_net_low",
        "smooth_net_high",
    )
    selected = {key: np.full(shape, np.nan, dtype=float) for key in keys}
    sigma_map = np.full(shape, np.nan, dtype=float)
    chosen = np.zeros(shape, dtype=bool)
    for sigma, product in zip(sigma_grid_pixels, products):
        accept = (
            ~chosen
            & product["valid"]
            & (product["fractional_error"] <= target_fractional_error)
        )
        for key in keys:
            selected[key][accept] = product[key][accept]
        sigma_map[accept] = sigma
        chosen |= accept

    largest = products[-1]
    fallback = (
        ~chosen
        & largest["valid"]
        & (largest["fractional_error"] <= maximum_fractional_error)
    )
    for key in keys:
        selected[key][fallback] = largest[key][fallback]
    sigma_map[fallback] = sigma_grid_pixels[-1]
    chosen |= fallback
    selected["sigma_pixels"] = sigma_map
    selected["valid"] = chosen
    return selected


def load_image(path: Path, extension: int | str = 0) -> tuple[np.ndarray, fits.Header]:
    with fits.open(path) as hdul:
        return np.asarray(hdul[extension].data, dtype=float), hdul[extension].header.copy()


def assert_wcs_equal(reference_header: fits.Header, other_header: fits.Header) -> None:
    reference = WCS(reference_header)
    other = WCS(other_header)
    if reference.pixel_shape != other.pixel_shape:
        raise ValueError(f"WCS pixel shapes differ: {reference.pixel_shape} vs {other.pixel_shape}")
    probes = np.array([[0.0, 0.0], [449.5, 449.5], [899.0, 899.0]])
    ref_world = reference.all_pix2world(probes, 0)
    other_world = other.all_pix2world(probes, 0)
    if not np.allclose(ref_world, other_world, rtol=0.0, atol=1e-10):
        raise ValueError("WCS world-coordinate probes differ")


def project_uniform_sky_counts(
    exposure: np.ndarray,
    header: fits.Header,
    model_counts: float,
    skyarea_arcmin2: float,
    calibration_radius_arcmin: float = 15.0,
) -> tuple[np.ndarray, dict[str, float]]:
    """Project a uniform vignetted sky model and preserve its R0-15 normalization."""

    exposure = np.asarray(exposure, dtype=float)
    wcs = WCS(header)
    yy, xx = np.indices(exposure.shape, dtype=float)
    world = wcs.pixel_to_world(xx, yy)
    center = SkyCoord(M104_RA_DEG * u.deg, M104_DEC_DEG * u.deg, frame="fk5")
    separation = center.separation(world).to_value(u.arcmin)
    valid_circle = (
        (separation <= calibration_radius_arcmin)
        & (exposure >= 0.2 * float(np.nanmax(exposure)))
    )
    if not np.any(valid_circle):
        raise ValueError("No valid exposure pixels inside R0-15 calibration circle")
    mean_exposure = float(np.nanmean(exposure[valid_circle]))
    pixel_area_arcmin2 = float(proj_plane_pixel_area(wcs)) * 3600.0
    equivalent_pixels = skyarea_arcmin2 / pixel_area_arcmin2
    normalization = model_counts / (mean_exposure * equivalent_pixels)
    counts = normalization * exposure
    return counts, {
        "model_counts_r015": float(model_counts),
        "skyarea_arcmin2": float(skyarea_arcmin2),
        "pixel_area_arcmin2": pixel_area_arcmin2,
        "equivalent_pixels": float(equivalent_pixels),
        "mean_exposure_r015_s": mean_exposure,
        "counts_per_exposure_pixel": float(normalization),
        "full_map_counts": float(np.nansum(counts)),
    }


def equivalent_r015_component_counts(
    component: np.ndarray,
    exposure: np.ndarray,
    header: fits.Header,
    skyarea_arcmin2: float,
) -> float:
    """Estimate counts in an R0-15 extraction with the specified BACKSCAL area."""

    wcs = WCS(header)
    yy, xx = np.indices(component.shape, dtype=float)
    world = wcs.pixel_to_world(xx, yy)
    center = SkyCoord(M104_RA_DEG * u.deg, M104_DEC_DEG * u.deg, frame="fk5")
    valid = (
        (center.separation(world).to_value(u.arcmin) <= 15.0)
        & (exposure >= 0.2 * float(np.nanmax(exposure)))
    )
    pixel_area_arcmin2 = float(proj_plane_pixel_area(wcs)) * 3600.0
    equivalent_pixels = skyarea_arcmin2 / pixel_area_arcmin2
    return float(np.nanmean(component[valid]) * equivalent_pixels)


def reproject_counts_component(
    path: Path, target_header: fits.Header, target_shape: tuple[int, int]
) -> tuple[np.ndarray, dict[str, float]]:
    """Area-overlap reproject a counts/pixel component to the authoritative grid."""

    source, source_header = load_image(path)
    projected, footprint = reproject_exact(
        (source, WCS(source_header)),
        WCS(target_header),
        shape_out=target_shape,
        parallel=False,
    )
    projected = np.nan_to_num(projected, nan=0.0)
    return projected, {
        "source_sum": float(np.nansum(source)),
        "projected_sum": float(np.nansum(projected)),
        "footprint_pixels": int(np.count_nonzero(footprint)),
    }


def sky_circle_pixels(wcs: WCS, radius_arcmin: float, n: int = 361) -> tuple[np.ndarray, np.ndarray]:
    angle = np.linspace(0.0, 2.0 * np.pi, n)
    dec0 = np.deg2rad(M104_DEC_DEG)
    radius_deg = radius_arcmin / 60.0
    ra = M104_RA_DEG + radius_deg * np.sin(angle) / np.cos(dec0)
    dec = M104_DEC_DEG + radius_deg * np.cos(angle)
    return wcs.all_world2pix(ra, dec, 0)


def add_geometry(ax, wcs: WCS) -> None:
    for radius, linestyle in ((4.0, "--"), (7.0, "-")):
        x, y = sky_circle_pixels(wcs, radius)
        ax.plot(x, y, color="white", linestyle=linestyle, linewidth=0.7, alpha=0.9)
    cx, cy = wcs.all_world2pix(M104_RA_DEG, M104_DEC_DEG, 0)
    ax.plot(cx, cy, marker="+", color="cyan", markersize=6, markeredgewidth=1.0)


def point_source_exclusion_mask(shape: tuple[int, int]) -> np.ndarray:
    """Rasterize the ESAS sky-coordinate cheese ellipses on the 2.5-arcsec grid."""

    yy, xx = np.indices(shape, dtype=float)
    event_x = 3426.0 + 50.0 * xx
    event_y = 3426.0 + 50.0 * yy
    excluded = np.zeros(shape, dtype=bool)
    with fits.open(POINT_SOURCE_REGIONS) as hdul:
        for row in hdul["REGION"].data:
            if row["SHAPE"].strip().upper() != "!ELLIPSE":
                continue
            dx = event_x - float(row["X"][0])
            dy = event_y - float(row["Y"][0])
            angle = np.deg2rad(float(row["ROTANG"][0]))
            rotated_x = dx * np.cos(angle) + dy * np.sin(angle)
            rotated_y = -dx * np.sin(angle) + dy * np.cos(angle)
            radius_x = float(row["R"][0])
            radius_y = float(row["R"][1])
            excluded |= (rotated_x / radius_x) ** 2 + (rotated_y / radius_y) ** 2 <= 1.0
    return excluded


def add_point_source_overlay(ax) -> None:
    """Overlay ESAS cheese exclusion ellipses without removing their counts."""

    with fits.open(POINT_SOURCE_REGIONS) as hdul:
        regions = hdul["REGION"].data
        for row in regions:
            if row["SHAPE"].strip().upper() != "!ELLIPSE":
                continue
            x_pixel = (float(row["X"][0]) - 3426.0) / 50.0
            y_pixel = (float(row["Y"][0]) - 3426.0) / 50.0
            width = 2.0 * float(row["R"][0]) / 50.0
            height = 2.0 * float(row["R"][1]) / 50.0
            ax.add_patch(
                Ellipse(
                    (x_pixel, y_pixel),
                    width=width,
                    height=height,
                    angle=float(row["ROTANG"][0]),
                    fill=False,
                    edgecolor="#FF8C00",
                    linewidth=0.35,
                    linestyle="--",
                    alpha=0.55,
                )
            )


def configure_wcs_axis(ax) -> None:
    ax.coords[0].set_axislabel("RA", fontsize=7)
    ax.coords[1].set_axislabel("Dec", fontsize=7)
    ax.coords[0].set_ticklabel(size=5, exclude_overlapping=True)
    ax.coords[1].set_ticklabel(size=5, exclude_overlapping=True)


def plot_quicklook(product: 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)

    panels = [
        (product["rate_low"], "Fe-L low: 0.70–0.875 keV", "magma", "log"),
        (product["rate_high"], "Fe-L high: 0.875–1.05 keV", "magma", "log"),
        (product["ratio"], r"Observed $R_{\rm FeL}$", "viridis", "linear"),
        (product["fractional_error"], "Ratio statistical fractional error", "cividis", "linear"),
    ]
    for index, (data, title, cmap, scale) in enumerate(panels, start=1):
        ax = fig.add_subplot(*grid, index, projection=wcs)
        finite = data[np.isfinite(data)]
        if finite.size == 0:
            raise RuntimeError(f"No finite pixels for panel {title}")
        if scale == "log":
            positive = finite[finite > 0]
            vmin, vmax = np.nanpercentile(positive, [3, 99.5])
            image = ax.imshow(data, origin="lower", cmap=cmap, norm=LogNorm(vmin=vmin, vmax=vmax))
        else:
            if "Observed" in title:
                vmin, vmax = np.nanpercentile(finite, [5, 95])
            else:
                vmin, vmax = 0.0, min(0.5, float(np.nanpercentile(finite, 97)))
            image = ax.imshow(data, origin="lower", cmap=cmap, vmin=vmin, vmax=vmax)
        add_geometry(ax, wcs)
        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 color quick-look\npoint sources retained; background not subtracted; shared 60″ kernel",
        fontsize=8,
        y=0.995,
    )
    fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.965))
    stem = OUT / f"r47_mos2_fel_ratio_point_sources_included_raw_quicklook_{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 write_product(product: dict[str, np.ndarray], header: fits.Header) -> Path:
    primary = fits.PrimaryHDU(header=header)
    primary.header["PRODTYPE"] = "OBSERVED_COLOR_QUICKLOOK"
    primary.header["PTSRC"] = (True, "Point sources retained")
    primary.header["BKGSUB"] = (False, "No background subtraction")
    primary.header["ADAPTIVE"] = (False, "Fixed shared kernel used for both bands")
    primary.header["SIGMA"] = (24.0, "Shared Gaussian sigma in pixels")
    primary.header["SIGARC"] = (60.0, "Shared Gaussian sigma in arcsec")
    hdus = [primary]
    output_arrays = [
        ("RATIO", "ratio"),
        ("FRACERR", "fractional_error"),
        ("LOW_RATE", "rate_low"),
        ("HIGH_RATE", "rate_high"),
        ("LOW_NET_SMOOTH", "smooth_net_low"),
        ("HIGH_NET_SMOOTH", "smooth_net_high"),
        ("VALID", "valid"),
    ]
    if "sigma_pixels" in product:
        output_arrays.insert(-1, ("SIGMA_PIX", "sigma_pixels"))
    for name, key in output_arrays:
        array = product[key].astype(np.uint8 if key == "valid" else np.float32)
        hdus.append(fits.ImageHDU(array, header=header, name=name))
    path = OUT / "r47_mos2_fel_ratio_point_sources_included_raw_observed_color.fits"
    fits.HDUList(hdus).writeto(path, overwrite=True, checksum=True)
    return path


def run_quicklook() -> None:
    low_counts, low_header = load_image(QUICK / "mos2U005-obj-im-700-875.fits")
    high_counts, high_header = load_image(QUICK / "mos2U005-obj-im-875-1050.fits")
    low_exposure, low_exp_header = load_image(QUICK / "mos2U005-exp-im-700-875.fits")
    high_exposure, high_exp_header = load_image(QUICK / "mos2U005-exp-im-875-1050.fits")
    for header in (high_header, low_exp_header, high_exp_header):
        assert_wcs_equal(low_header, header)

    product = compute_smoothed_ratio(
        counts_low=low_counts,
        counts_high=high_counts,
        exposure_low=low_exposure,
        exposure_high=high_exposure,
        sigma_pixels=24.0,
        max_fractional_error=0.30,
    )
    fits_path = write_product(product, low_header)
    for width in ("1col", "2col"):
        plot_quicklook(product, low_header, width)

    valid_ratio = product["ratio"][product["valid"]]
    summary = {
        "product_class": "observed_color_quicklook",
        "point_sources_retained": True,
        "background_subtracted": False,
        "bands_keV": {"low": [0.70, 0.875], "high": [0.875, 1.05]},
        "shared_gaussian_sigma_pixels": 24.0,
        "shared_gaussian_sigma_arcsec": 60.0,
        "maximum_ratio_fractional_error": 0.30,
        "valid_pixels": int(product["valid"].sum()),
        "ratio_median": float(np.nanmedian(valid_ratio)),
        "ratio_percentiles_16_84": [float(x) for x in np.nanpercentile(valid_ratio, [16, 84])],
        "input_counts": {"low": int(low_counts.sum()), "high": int(high_counts.sum())},
        "fits_path": str(fits_path),
        "warning": "This is not a diffuse-gas temperature map: point sources and all backgrounds remain.",
    }
    (OUT / "r47_mos2_fel_ratio_point_sources_included_raw_quicklook_summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, indent=2))


def plot_background_product(
    product: 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)
    panels = [
        (product["rate_low"], "Net Fe-L low rate", "magma", "log"),
        (product["rate_high"], "Net Fe-L high rate", "magma", "log"),
        (product["ratio"], r"Background-corrected $R_{\rm FeL}$", "viridis", "linear"),
        (product["fractional_error"], "Statistical fractional error", "cividis", "error"),
    ]
    for index, (data, title, cmap, scale) in enumerate(panels, start=1):
        ax = fig.add_subplot(*grid, index, projection=wcs)
        finite = data[np.isfinite(data)]
        if scale == "log":
            positive = finite[finite > 0]
            vmin, vmax = np.nanpercentile(positive, [3, 99.5])
            image = ax.imshow(data, origin="lower", cmap=cmap, norm=LogNorm(vmin=vmin, vmax=vmax))
        elif scale == "error":
            image = ax.imshow(data, origin="lower", cmap=cmap, vmin=0.0, vmax=0.35)
        else:
            vmin, vmax = np.nanpercentile(finite, [5, 95])
            image = ax.imshow(data, origin="lower", cmap=cmap, vmin=vmin, vmax=vmax)
            add_point_source_overlay(ax)
        add_geometry(ax, wcs)
        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 color: QPB + native SP + sky subtracted\n"
        "point sources retained; orange ellipses mark ESAS source masks; shared 60″ kernel",
        fontsize=8,
        y=0.995,
    )
    fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.96))
    stem = OUT / f"r47_mos2_fel_ratio_point_sources_included_background_corrected_{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 plot_background_controls(
    controls: 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)
    titles = {
        "raw": "Raw observed color",
        "qpb": "QPB subtracted",
        "full": "QPB + native SP + sky subtracted",
        "stress": "MOS1 common-SP stress test",
    }
    finite_values = np.concatenate(
        [item["ratio"][np.isfinite(item["ratio"])] for item in controls.values()]
    )
    vmin, vmax = np.nanpercentile(finite_values, [5, 95])
    for index, key in enumerate(("raw", "qpb", "full", "stress"), start=1):
        ax = fig.add_subplot(*grid, index, projection=wcs)
        image = ax.imshow(
            controls[key]["ratio"], origin="lower", cmap="viridis", vmin=vmin, vmax=vmax
        )
        add_geometry(ax, wcs)
        add_point_source_overlay(ax)
        configure_wcs_axis(ax)
        ax.set_title(titles[key], 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 background-component controls; point sources retained", fontsize=8)
    fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.97))
    stem = OUT / f"r47_mos2_fel_ratio_background_component_controls_{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 run_background_corrected() -> None:
    low_counts, header = load_image(QUICK / "mos2U005-obj-im-700-875.fits")
    high_counts, high_header = load_image(QUICK / "mos2U005-obj-im-875-1050.fits")
    low_exposure, low_exp_header = load_image(QUICK / "mos2U005-exp-im-700-875.fits")
    high_exposure, high_exp_header = load_image(QUICK / "mos2U005-exp-im-875-1050.fits")
    for other in (high_header, low_exp_header, high_exp_header):
        assert_wcs_equal(header, other)
    shape = low_counts.shape
    contract = json.loads(SPECTRAL_CONTRACT.read_text())
    skyarea = float(contract["skyarea_arcmin2"])

    qpb_low, qpb_low_meta = reproject_counts_component(
        ESAS_LOW / "mos2U005-back-im-sky-700-875.fits", header, shape
    )
    qpb_high, qpb_high_meta = reproject_counts_component(
        ESAS_HIGH / "mos2U005-back-im-sky-875-1050.fits", header, shape
    )
    sp_low, sp_low_meta = reproject_counts_component(
        ESAS_LOW / "mos2U005-prot-im-sky-700-875.fits", header, shape
    )
    sp_high, sp_high_meta = reproject_counts_component(
        ESAS_HIGH / "mos2U005-prot-im-sky-875-1050.fits", header, shape
    )
    sp_low_estimate = equivalent_r015_component_counts(sp_low, low_exposure, header, skyarea)
    sp_high_estimate = equivalent_r015_component_counts(sp_high, high_exposure, header, skyarea)
    sp_low_scale = (
        contract["bands"]["low"]["soft_proton_native_mos2_model_counts"] / sp_low_estimate
    )
    sp_high_scale = (
        contract["bands"]["high"]["soft_proton_native_mos2_model_counts"] / sp_high_estimate
    )
    sp_low *= sp_low_scale
    sp_high *= sp_high_scale

    sky_low, sky_low_meta = project_uniform_sky_counts(
        low_exposure,
        header,
        contract["bands"]["low"]["sky_background_model_counts"],
        skyarea,
    )
    sky_high, sky_high_meta = project_uniform_sky_counts(
        high_exposure,
        header,
        contract["bands"]["high"]["sky_background_model_counts"],
        skyarea,
    )
    background_low = qpb_low + sp_low + sky_low
    background_high = qpb_high + sp_high + sky_high

    kwargs = {
        "counts_low": low_counts,
        "counts_high": high_counts,
        "exposure_low": low_exposure,
        "exposure_high": high_exposure,
        "sigma_pixels": 24.0,
        "max_fractional_error": 0.35,
    }
    raw = compute_smoothed_ratio(**kwargs)
    qpb_only = compute_smoothed_ratio(
        **kwargs, background_low=qpb_low, background_high=qpb_high
    )
    full = compute_smoothed_ratio(
        **kwargs, background_low=background_low, background_high=background_high
    )
    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"]
    )
    stress = compute_smoothed_ratio(
        **kwargs,
        background_low=qpb_low + sky_low + sp_low * common_low_factor,
        background_high=qpb_high + sky_high + sp_high * common_high_factor,
    )
    controls = {"raw": raw, "qpb": qpb_only, "full": full, "stress": stress}

    primary = fits.PrimaryHDU(header=header)
    primary.header["PRODTYPE"] = "BACKGROUND_CORRECTED_COLOR"
    primary.header["PTSRC"] = (True, "Point sources retained")
    primary.header["SIGARC"] = (60.0, "Shared Gaussian sigma in arcsec")
    hdus = [primary]
    arrays = {
        "LOW_COUNTS": low_counts,
        "HIGH_COUNTS": high_counts,
        "LOW_EXPOSURE": low_exposure,
        "HIGH_EXPOSURE": high_exposure,
        "LOW_QPB": qpb_low,
        "HIGH_QPB": qpb_high,
        "LOW_SP": sp_low,
        "HIGH_SP": sp_high,
        "LOW_SKY": sky_low,
        "HIGH_SKY": sky_high,
        "LOW_NET": low_counts - background_low,
        "HIGH_NET": high_counts - background_high,
        "LOW_RATE": full["rate_low"],
        "HIGH_RATE": full["rate_high"],
        "RATIO": full["ratio"],
        "FRACERR": full["fractional_error"],
        "VALID": full["valid"].astype(np.uint8),
        "RAW_RATIO": raw["ratio"],
        "QPB_RATIO": qpb_only["ratio"],
        "STRESS_RATIO": stress["ratio"],
    }
    for name, array in arrays.items():
        dtype = np.uint8 if name == "VALID" else np.float32
        hdus.append(fits.ImageHDU(np.asarray(array, dtype=dtype), header=header, name=name))
    fits_path = OUT / "r47_mos2_fel_ratio_point_sources_included_background_corrected.fits"
    fits.HDUList(hdus).writeto(fits_path, overwrite=True, checksum=True)

    for width in ("1col", "2col"):
        plot_background_product(full, header, width)
        plot_background_controls(controls, header, width)

    def ratio_stats(item: dict[str, np.ndarray]) -> dict[str, float | int]:
        values = item["ratio"][item["valid"]]
        return {
            "valid_pixels": int(item["valid"].sum()),
            "median": float(np.nanmedian(values)),
            "p16": float(np.nanpercentile(values, 16)),
            "p84": float(np.nanpercentile(values, 84)),
        }

    summary = {
        "product_class": "point_source_included_background_corrected_fe_l_color",
        "fits_path": str(fits_path),
        "input_counts": {"low": float(low_counts.sum()), "high": float(high_counts.sum())},
        "component_full_map_counts": {
            "qpb_low": float(qpb_low.sum()),
            "qpb_high": float(qpb_high.sum()),
            "sp_low": float(sp_low.sum()),
            "sp_high": float(sp_high.sum()),
            "sky_low": float(sky_low.sum()),
            "sky_high": float(sky_high.sum()),
        },
        "soft_proton_spectral_rescale": {"low": float(sp_low_scale), "high": float(sp_high_scale)},
        "qpb_reprojection": {"low": qpb_low_meta, "high": qpb_high_meta},
        "sp_reprojection": {"low": sp_low_meta, "high": sp_high_meta},
        "sky_projection": {"low": sky_low_meta, "high": sky_high_meta},
        "ratio_statistics": {key: ratio_stats(value) for key, value in controls.items()},
        "warning": "Point sources are retained; corrected ratio is not a diffuse-gas temperature map.",
    }
    (OUT / "r47_mos2_fel_ratio_background_corrected_summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, indent=2))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--quicklook", action="store_true", help="Build the raw observed-color quick-look")
    parser.add_argument(
        "--background-corrected",
        action="store_true",
        help="Build the QPB/native-SP/sky-subtracted point-source-included color map",
    )
    args = parser.parse_args()
    if not args.quicklook and not args.background_corrected:
        parser.error("Select --quicklook and/or --background-corrected")
    if args.quicklook:
        run_quicklook()
    if args.background_corrected:
        run_background_corrected()


if __name__ == "__main__":
    main()