#!/usr/bin/env python
"""Rebuild the M82 Chandra-to-SQUDE mock-tutorial figure set.

The pipeline is intentionally closed and physical:

1. Fit the extracted Chandra 4x4-arcmin PHA with the literature baseline
   XSPEC model (tbabs × [tbabs × (vapec + vapec + 2 gaussians) + tbabs × PL]).
2. Use that exact, fitted source model and a seeded Poisson draw to create a
   new 50 ks SQUDE fake_pha product.
3. Render three intentionally separated figures: Chandra data+total model+
   Chandra-response-folded components; SQUDE mock data+model only; and SQUDE
   data/model plus the Chandra-RMF control.
4. Plot the SQUDE fake spectrum in linear-linear coordinates and construct
   a comparison curve produced by

       source F(E) -> SQUDE ARF -> Chandra RMF -> rate-conserving rebin

   onto the SQUDE display bins.

No polynomial or empirical surrogate continuum is used.

Run after CIAO initialization:

    cd ~/program/M82
    export ASCDS_OVERRIDE=1
    source ~/software/ciao-4.18/bin/ciao.sh -q
    export PFILES="$PWD/pfiles;$ASCDS_INSTALL/param:"
    python scripts/make_m82_mock_tutorial_figure_set.py
"""

from __future__ import annotations

import argparse
import copy
import csv
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

os.environ.setdefault("MPLCONFIGDIR", "/tmp/m82-mock-tutorial-matplotlib")
os.environ.setdefault("XDG_CACHE_HOME", "/tmp/m82-mock-tutorial-cache")

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import scienceplots  # noqa: F401
from astropy.io import fits
from sherpa.astro.ui import (
    clean,
    fake_pha,
    fit,
    freeze,
    get_data,
    get_data_plot,
    get_delchi_plot,
    get_fit_results,
    get_model_plot,
    group_counts,
    link,
    load_pha,
    notice_id,
    save_pha,
    set_analysis,
    set_method,
    set_rng,
    set_source,
    set_stat,
    set_xsabund,
    set_xsxsect,
    thaw,
    ungroup,
)
from sherpa.astro.ui import powlaw1d, xsgaussian, xstbabs, xsvapec
from sherpa.utils.random import poisson_noise


ROOT = Path(__file__).resolve().parents[1]
CHANDRA_SPEC_DIR = ROOT / "data" / "spectra" / "m82_4x4arcmin_early100ks"
CHANDRA_PHA = CHANDRA_SPEC_DIR / "m82_4x4arcmin_early100ks_combined_src.pi"
CHANDRA_ARF = CHANDRA_SPEC_DIR / "m82_4x4arcmin_early100ks_combined_src.arf"
CHANDRA_RMF = CHANDRA_SPEC_DIR / "m82_4x4arcmin_early100ks_combined_src.rmf"
SQUDE_ARF = ROOT / "SQUDE_response" / "SQUDE_rsp_v1.1.1.arf"
SQUDE_RMF = ROOT / "SQUDE_response" / "SQUDE_rsp_v1.1.2.rmf"
OUT_DIR = ROOT / "data" / "fits" / "m82_mock_tutorial_response_comparison"
RNG_SEED = 20260719

COMPONENT_STYLE = {
    "cool_vapec": ("cool vapec", "#0072B2", "--"),
    "hot_vapec": ("hot vapec", "#E69F00", "--"),
    "cx_o": ("CX line proxy 0.777 keV", "#009E73", ":"),
    "cx_mg": ("CX line proxy 1.234 keV", "#CC79A7", ":"),
    "powerlaw": ("absorbed power law", "#7A3E00", "-."),
}


@dataclass(frozen=True)
class RmfKernel:
    """Dense OGIP response matrix and true/channel energy grids."""

    true_lo: np.ndarray
    true_hi: np.ndarray
    channel_lo: np.ndarray
    channel_hi: np.ndarray
    redistribution: np.ndarray  # shape: (output channel, true-energy bin)


@dataclass(frozen=True)
class ModelBundle:
    """Fitted model and response-folded physical components."""

    total: Any
    components: dict[str, Any]


def configure_plot_style() -> None:
    """Use the project-standard scientific plotting style without LaTeX."""
    plt.style.use(["science", "no-latex"])


def setup_literature_model(dataset_id: int) -> ModelBundle:
    """Create the documented VAPEC-based M82 baseline model and fit starts.

    Model:
        tbabs_Gal * [tbabs_wind * (vapec_cool + vapec_hot
                     + gaussian_0.777 + gaussian_1.234)
                     + tbabs_powerlaw * powerlaw]

    The thermal components are genuine XSPEC VAPEC components. The power-law is
    retained only for the documented absorbed unresolved hard-source continuum.
    """
    set_xsabund("wilm")
    set_xsxsect("vern")

    gal = xstbabs.gal
    absth = xstbabs.absth
    vcool = xsvapec.vcool
    vhot = xsvapec.vhot
    gce_o = xsgaussian.gce_o
    gce_mg = xsgaussian.gce_mg
    abspl = xstbabs.abspl
    pl = powlaw1d.pl

    total = gal * (absth * (vcool + vhot + gce_o + gce_mg) + abspl * pl)
    set_source(dataset_id, total)

    gal.nH = 0.04
    freeze(gal.nH)
    absth.nH = 0.45
    thaw(absth.nH)
    abspl.nH = 2.3
    thaw(abspl.nH)

    vcool.kT = 0.60
    thaw(vcool.kT)
    vcool.norm = 1.0e-2
    thaw(vcool.norm)
    vcool.Redshift = 0.000677
    freeze(vcool.Redshift)

    vhot.kT = 5.0
    freeze(vhot.kT)
    vhot.norm = 1.0e-3
    thaw(vhot.norm)
    vhot.Redshift = 0.000677
    freeze(vhot.Redshift)

    for element, value in (
        ("O", 0.20),
        ("Ne", 0.45),
        ("Mg", 1.00),
        ("Si", 1.00),
        ("S", 1.10),
        ("Fe", 0.30),
    ):
        getattr(vcool, element).val = value
        thaw(getattr(vcool, element))
        getattr(vhot, element).val = value
        link(getattr(vhot, element), getattr(vcool, element))

    for element in ("He", "C", "N", "Al", "Ar", "Ca", "Ni"):
        getattr(vcool, element).val = 1.0
        freeze(getattr(vcool, element))
        getattr(vhot, element).val = 1.0
        freeze(getattr(vhot, element))

    pl.gamma = 1.6
    freeze(pl.gamma)
    pl.ampl = 1.0e-3
    thaw(pl.ampl)

    for line, energy in ((gce_o, 0.777), (gce_mg, 1.234)):
        line.LineE = energy
        freeze(line.LineE)
        line.Sigma = 0.01
        freeze(line.Sigma)
        line.norm = 1.0e-5
        thaw(line.norm)

    return ModelBundle(
        total=total,
        components={
            "cool_vapec": gal * absth * vcool,
            "hot_vapec": gal * absth * vhot,
            "cx_o": gal * absth * gce_o,
            "cx_mg": gal * absth * gce_mg,
            "powerlaw": gal * abspl * pl,
        },
    )


def response_folded_components(bundle: ModelBundle, dataset_id: int) -> dict[str, Any]:
    """Return component plots after forcing a fresh response fold per component."""
    plots: dict[str, Any] = {}
    for name, component in bundle.components.items():
        component.cache_clear()
        set_source(dataset_id, component)
        plots[name] = copy.deepcopy(get_model_plot(dataset_id))
    bundle.total.cache_clear()
    set_source(dataset_id, bundle.total)
    return plots


def read_arf(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Read SPECRESP bin edges and effective area from an OGIP ARF."""
    with fits.open(path) as hdul:
        data = hdul["SPECRESP"].data
        return (
            np.asarray(data["ENERG_LO"], dtype=float),
            np.asarray(data["ENERG_HI"], dtype=float),
            np.asarray(data["SPECRESP"], dtype=float),
        )


def read_rmf_kernel(path: Path) -> RmfKernel:
    """Expand sparse OGIP MATRIX rows into R(output-channel, true-energy-bin)."""
    with fits.open(path) as hdul:
        matrix = hdul["MATRIX"].data
        ebounds = hdul["EBOUNDS"].data
        true_lo = np.asarray(matrix["ENERG_LO"], dtype=float)
        true_hi = np.asarray(matrix["ENERG_HI"], dtype=float)
        channel = np.asarray(ebounds["CHANNEL"], dtype=int)
        channel_lo = np.asarray(ebounds["E_MIN"], dtype=float)
        channel_hi = np.asarray(ebounds["E_MAX"], dtype=float)

        redistribution = np.zeros((len(channel), len(true_lo)), dtype=float)
        channel_offset = int(channel.min())
        for true_index, row in enumerate(matrix):
            values = np.asarray(row["MATRIX"], dtype=float)
            cursor = 0
            for first_channel, n_channel in zip(
                np.atleast_1d(row["F_CHAN"]),
                np.atleast_1d(row["N_CHAN"]),
                strict=True,
            ):
                start = int(first_channel) - channel_offset
                stop = start + int(n_channel)
                redistribution[start:stop, true_index] = values[cursor: cursor + int(n_channel)]
                cursor += int(n_channel)
            if cursor != len(values):
                raise RuntimeError(
                    f"RMF row {true_index}: used {cursor} values, expected {len(values)}"
                )
    return RmfKernel(true_lo, true_hi, channel_lo, channel_hi, redistribution)


def arf_on_rmf_grid(kernel: RmfKernel, arf_path: Path) -> np.ndarray:
    """Interpolate effective area, not the source model, onto RMF true-energy bins."""
    arf_lo, arf_hi, effective_area = read_arf(arf_path)
    return np.interp(
        0.5 * (kernel.true_lo + kernel.true_hi),
        0.5 * (arf_lo + arf_hi),
        effective_area,
        left=0.0,
        right=0.0,
    )


def fold_source_with_kernel(source_model: Any, kernel: RmfKernel, arf_path: Path) -> np.ndarray:
    """Forward fold F(E) through arbitrary ARF and RMF into count-rate density.

    The XSPEC source expression returns photon flux integrated over each true
    RMF energy bin, F_j [photons cm^-2 s^-1]. The exact discrete forward model:

        r_i = sum_j R_ij F_j A_j,

    is divided by the output-bin width to obtain counts s^-1 keV^-1.
    """
    photons_per_bin = np.asarray(source_model(kernel.true_lo, kernel.true_hi), dtype=float)
    count_rate = kernel.redistribution @ (photons_per_bin * arf_on_rmf_grid(kernel, arf_path))
    return count_rate / (kernel.channel_hi - kernel.channel_lo)


def rebin_density_conserving_rate(
    source_lo: np.ndarray,
    source_hi: np.ndarray,
    source_density: np.ndarray,
    target_lo: np.ndarray,
    target_hi: np.ndarray,
) -> np.ndarray:
    """Map density to a new grid while conserving its integrated count rate."""
    output = np.zeros_like(target_lo, dtype=float)
    for index, (lo, hi) in enumerate(zip(target_lo, target_hi, strict=True)):
        overlap = np.clip(np.minimum(source_hi, hi) - np.maximum(source_lo, lo), 0.0, None)
        output[index] = np.sum(source_density * overlap) / (hi - lo)
    return output


def validate_explicit_chandra_fold(source_model: Any, sherpa_chandra_model: Any) -> dict[str, float]:
    """Require the manually-expanded Chandra RMF to reproduce Sherpa's fold."""
    kernel = read_rmf_kernel(CHANDRA_RMF)
    explicit_density = fold_source_with_kernel(source_model, kernel, CHANDRA_ARF)
    explicit_grouped = rebin_density_conserving_rate(
        kernel.channel_lo,
        kernel.channel_hi,
        explicit_density,
        np.asarray(sherpa_chandra_model.xlo, dtype=float),
        np.asarray(sherpa_chandra_model.xhi, dtype=float),
    )
    native = np.asarray(sherpa_chandra_model.y, dtype=float)
    relative = np.abs(explicit_grouped - native) / np.maximum(np.abs(native), 1.0e-30)
    metrics = {
        "median_relative_error": float(np.median(relative)),
        "p95_relative_error": float(np.percentile(relative, 95)),
        "max_relative_error": float(np.max(relative)),
    }
    if metrics["p95_relative_error"] > 1.0e-5:
        raise RuntimeError(
            "Explicit RMF folding does not reproduce Sherpa's Chandra model: "
            f"p95={metrics['p95_relative_error']:.3e}"
        )
    return metrics


def validate_component_closure(total_model: Any, components: dict[str, Any]) -> dict[str, float]:
    """Require a response-folded total model to equal its component sum."""
    total = np.asarray(total_model.y, dtype=float)
    summed = np.zeros_like(total)
    for component in components.values():
        if not (np.array_equal(component.xlo, total_model.xlo) and np.array_equal(component.xhi, total_model.xhi)):
            raise RuntimeError("Component and total model grids differ")
        summed += np.asarray(component.y, dtype=float)
    relative = np.abs(summed - total) / np.maximum(np.abs(total), 1.0e-30)
    metrics = {
        "median_relative_error": float(np.median(relative)),
        "p95_relative_error": float(np.percentile(relative, 95)),
        "max_relative_error": float(np.max(relative)),
    }
    if metrics["p95_relative_error"] > 2.0e-5:
        raise RuntimeError(
            "Response-folded component closure failed: "
            f"p95={metrics['p95_relative_error']:.3e}"
        )
    return metrics


def validate_rebin_rate_conservation(
    source_lo: np.ndarray,
    source_hi: np.ndarray,
    source_density: np.ndarray,
    target_lo: np.ndarray,
    target_hi: np.ndarray,
    target_density: np.ndarray,
) -> dict[str, float]:
    """Validate integrated count rate over the common source/target support."""
    lower = max(float(np.min(source_lo)), float(np.min(target_lo)))
    upper = min(float(np.max(source_hi)), float(np.max(target_hi)))
    overlap = np.maximum(
        0.0,
        np.minimum(source_hi, upper) - np.maximum(source_lo, lower),
    )
    native_rate = float(np.sum(source_density * overlap))
    target_rate = float(np.sum(target_density * (target_hi - target_lo)))
    relative_error = abs(target_rate - native_rate) / max(abs(native_rate), 1.0e-30)
    if relative_error > 1.0e-10:
        raise RuntimeError(
            "Rate-conserving rebin validation failed: "
            f"relative_error={relative_error:.3e}"
        )
    return {
        "native_overlap_rate_counts_s": native_rate,
        "target_rate_counts_s": target_rate,
        "relative_error": relative_error,
    }


def model_on_data_grid(model: Any, data: Any) -> np.ndarray:
    """Rebin a response-folded model to exactly the displayed data bins."""
    return rebin_density_conserving_rate(
        np.asarray(model.xlo, dtype=float),
        np.asarray(model.xhi, dtype=float),
        np.asarray(model.y, dtype=float),
        np.asarray(data.xlo, dtype=float),
        np.asarray(data.xhi, dtype=float),
    )


def make_chandra_data_model_figure(
    data: Any,
    model: Any,
    components: dict[str, Any],
    path: Path,
    width: float,
) -> None:
    """Plot Chandra data, total model, and Chandra-response-folded components."""
    model_y = model_on_data_grid(model, data)
    figure, axis = plt.subplots(figsize=(width, width * 0.68))
    for name, component in components.items():
        label, color, linestyle = COMPONENT_STYLE[name]
        component_y = model_on_data_grid(component, data)
        axis.plot(
            data.x,
            component_y,
            color=color,
            ls=linestyle,
            lw=0.80 if width < 4 else 1.05,
            alpha=0.90,
            label=label,
        )
    axis.plot(data.x, model_y, color="#D55E00", lw=1.35, label="Chandra-RMF-folded total model")
    axis.errorbar(data.x, data.y, yerr=data.yerr, fmt=".", ms=2.5, color="0.10", label="Chandra ACIS data")
    axis.set_yscale("log")
    axis.set_xlim(0.3, 4.0)
    axis.set_ylim(1.0e-3, 10.0)
    axis.set_xlabel("Energy (keV)")
    axis.set_ylabel(r"Counts s$^{-1}$ keV$^{-1}$")
    axis.set_title(
        f"M82 4′×4′ Chandra ACIS ({len(data.x)} bins): total model and components",
        fontsize=8.2 if width < 4 else 11,
        pad=58 if width < 4 else 6,
    )
    if width < 4:
        axis.legend(
            loc="lower center",
            bbox_to_anchor=(0.5, 1.01),
            fontsize=4.6,
            ncol=2,
            frameon=True,
            columnspacing=0.8,
            handlelength=2.4,
        )
    else:
        axis.legend(loc="upper right", fontsize=7.1, ncol=2, frameon=True)
    axis.grid(alpha=0.20)
    figure.savefig(path, dpi=320, bbox_inches="tight")
    plt.close(figure)


def make_squde_data_model_figure(data: Any, model: Any, path: Path, width: float) -> None:
    """Plot only seeded SQUDE mock data and the same-model SQUDE prediction."""
    model_y = model_on_data_grid(model, data)
    mask = (data.x >= 0.50) & (data.x <= 1.70)
    figure, axis = plt.subplots(figsize=(width, width * 0.62))
    axis.errorbar(data.x[mask], data.y[mask], yerr=data.yerr[mask], fmt=".", ms=2.0 if width < 4 else 2.8, color="0.10", label="SQUDE mock data (50 ks)")
    axis.plot(data.x[mask], model_y[mask], color="#D55E00", lw=1.25, label="SQUDE-RMF-folded model")
    ymax = float(np.max(np.maximum(data.y[mask] + data.yerr[mask], model_y[mask])))
    axis.set_xlim(0.50, 1.70)
    axis.set_ylim(0.0, 1.10 * ymax)
    axis.set_xlabel("Energy (keV)")
    axis.set_ylabel(r"Counts s$^{-1}$ keV$^{-1}$")
    axis.set_title("M82 50 ks SQUDE mock: data and same-model prediction", fontsize=8.8 if width < 4 else 11)
    axis.legend(loc="upper right", fontsize=6.3 if width < 4 else 8.0, frameon=True)
    axis.grid(alpha=0.23)
    figure.savefig(path, dpi=320, bbox_inches="tight")
    plt.close(figure)


def resolution_control_ylim_top(
    data: Any,
    squde_model: Any,
    chandra_rmf_squde_arf: np.ndarray,
) -> float:
    """Return the plotted 0.5--1.7 keV upper limit for one grouping."""
    squde_y = model_on_data_grid(squde_model, data)
    chandra_y = rebin_density_conserving_rate(
        np.asarray(squde_model.xlo, dtype=float),
        np.asarray(squde_model.xhi, dtype=float),
        chandra_rmf_squde_arf,
        np.asarray(data.xlo, dtype=float),
        np.asarray(data.xhi, dtype=float),
    )
    mask = (data.x >= 0.50) & (data.x <= 1.70)
    ymax = float(np.max(np.maximum.reduce([
        data.y[mask] + data.yerr[mask],
        squde_y[mask],
        chandra_y[mask],
    ])))
    return 1.10 * ymax


def make_resolution_control_figure(
    data: Any,
    squde_model: Any,
    chandra_rmf_squde_arf: np.ndarray,
    path: Path,
    width: float,
    title: str | None = None,
    control_label: str = "Chandra RMF model (SQUDE ARF)",
    ylim_top: float | None = None,
) -> None:
    """Plot mock data, SQUDE model, and Chandra-RMF control on data bins."""
    squde_y = model_on_data_grid(squde_model, data)
    chandra_y = rebin_density_conserving_rate(
        np.asarray(squde_model.xlo, dtype=float),
        np.asarray(squde_model.xhi, dtype=float),
        chandra_rmf_squde_arf,
        np.asarray(data.xlo, dtype=float),
        np.asarray(data.xhi, dtype=float),
    )
    mask = (data.x >= 0.50) & (data.x <= 1.70)
    figure, axis = plt.subplots(figsize=(width, width * 0.62))
    axis.errorbar(data.x[mask], data.y[mask], yerr=data.yerr[mask], fmt=".", ms=2.0 if width < 4 else 2.8, color="0.10", label="SQUDE mock data (50 ks)", zorder=3)
    axis.plot(data.x[mask], squde_y[mask], color="#D55E00", lw=1.25, label="SQUDE RMF model", zorder=2)
    axis.plot(data.x[mask], chandra_y[mask], color="black", ls="-.", lw=1.15, label=control_label, zorder=4)
    auto_ylim_top = 1.10 * float(np.max(np.maximum.reduce([
        data.y[mask] + data.yerr[mask],
        squde_y[mask],
        chandra_y[mask],
    ])))
    axis.set_xlim(0.50, 1.70)
    axis.set_ylim(0.0, auto_ylim_top if ylim_top is None else ylim_top)
    axis.set_xlabel("Energy (keV)")
    axis.set_ylabel(r"Counts s$^{-1}$ keV$^{-1}$")
    axis.set_title(
        title or "M82 50 ks SQUDE mock: SQUDE versus Chandra RMF",
        fontsize=8.8 if width < 4 else 11,
        pad=72 if width < 4 else 6,
    )
    if width < 4:
        axis.legend(
            loc="lower center",
            bbox_to_anchor=(0.5, 1.01),
            fontsize=4.9,
            ncol=1,
            frameon=True,
            handlelength=2.6,
        )
    else:
        axis.legend(loc="upper right", fontsize=7.8, frameon=True)
    axis.grid(alpha=0.23)
    figure.savefig(path, dpi=320, bbox_inches="tight")
    plt.close(figure)


def write_curve_csv(
    model: Any,
    components: dict[str, Any],
    chandra_rmf_squde_arf: np.ndarray,
    path: Path,
) -> None:
    """Save the plotted forward-model curves as a machine-readable artifact."""
    fields = [
        "energy_lo_keV",
        "energy_hi_keV",
        "squde_rmf_counts_s_keV",
        "chandra_rmf_squde_arf_counts_s_keV",
        *[f"{name}_counts_s_keV" for name in components],
    ]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for index in range(len(model.x)):
            row = {
                "energy_lo_keV": f"{model.xlo[index]:.9g}",
                "energy_hi_keV": f"{model.xhi[index]:.9g}",
                "squde_rmf_counts_s_keV": f"{model.y[index]:.9g}",
                "chandra_rmf_squde_arf_counts_s_keV": f"{chandra_rmf_squde_arf[index]:.9g}",
            }
            row.update({f"{name}_counts_s_keV": f"{plot.y[index]:.9g}" for name, plot in components.items()})
            writer.writerow(row)


def main(output_dir: Path | None = None) -> None:
    """Run the pipeline, optionally writing all products outside the project."""
    output_dir = OUT_DIR if output_dir is None else output_dir.expanduser().resolve()
    squde_pha = output_dir / "m82_squde_50ks_from_current_chandra_bestfit.pha"
    configure_plot_style()
    output_dir.mkdir(parents=True, exist_ok=True)

    # Step 1: fresh Chandra fit using the literature VAPEC model.
    clean()
    load_pha(1, str(CHANDRA_PHA))
    set_analysis(1, "energy")
    notice_id(1, 0.3, 7.0)
    group_counts(1, 50)
    set_stat("chi2gehrels")
    set_method("levmar")
    bundle = setup_literature_model(1)
    fit(1)
    fit_result = get_fit_results()
    chandra_data = copy.deepcopy(get_data_plot(1))
    chandra_model = copy.deepcopy(get_model_plot(1))
    kernel_validation = validate_explicit_chandra_fold(bundle.total, chandra_model)
    chandra_components = response_folded_components(bundle, 1)
    chandra_model = copy.deepcopy(get_model_plot(1))
    chandra_component_validation = validate_component_closure(chandra_model, chandra_components)

    # Step 2: a seeded fake_pha realization from the exact fitted Chandra model.
    # fake_pha injects its UI-session RNG into ``method`` as a keyword. Set
    # that session RNG directly; a functools.partial(rng=...) is overridden.
    set_rng(np.random.default_rng(RNG_SEED))
    set_source(2, bundle.total)
    fake_pha(
        2,
        arf=str(SQUDE_ARF),
        rmf=str(SQUDE_RMF),
        exposure=50000.0,
        grouped=False,
        method=poisson_noise,
    )
    save_pha(2, str(squde_pha), clobber=True)
    set_analysis(2, "energy")
    notice_id(2, 0.3, 4.0)
    group_counts(2, 50)
    squde_data = copy.deepcopy(get_data_plot(2))
    squde_components = response_folded_components(bundle, 2)
    # Sherpa reuses a mutable model-plot object across datasets and calls.
    # Deep-copy after restoring the total so later plot requests cannot alias it.
    squde_model = copy.deepcopy(get_model_plot(2))
    component_validation = validate_component_closure(squde_model, squde_components)

    # Requested resolution-comparison operator: F(E) -> A_SQUDE -> R_Chandra,
    # then a count-rate-conserving projection to SQUDE display bins.
    chandra_kernel = read_rmf_kernel(CHANDRA_RMF)
    chandra_density_with_squde_arf = fold_source_with_kernel(bundle.total, chandra_kernel, SQUDE_ARF)
    chandra_rmf_squde_arf = rebin_density_conserving_rate(
        chandra_kernel.channel_lo,
        chandra_kernel.channel_hi,
        chandra_density_with_squde_arf,
        np.asarray(squde_model.xlo, dtype=float),
        np.asarray(squde_model.xhi, dtype=float),
    )
    control_conservation = validate_rebin_rate_conservation(
        chandra_kernel.channel_lo,
        chandra_kernel.channel_hi,
        chandra_density_with_squde_arf,
        np.asarray(squde_model.xlo, dtype=float),
        np.asarray(squde_model.xhi, dtype=float),
        chandra_rmf_squde_arf,
    )

    make_chandra_data_model_figure(
        chandra_data, chandra_model, chandra_components,
        output_dir / "m82_chandra_data_model_1col.png", 3.35,
    )
    make_chandra_data_model_figure(
        chandra_data, chandra_model, chandra_components,
        output_dir / "m82_chandra_data_model_2col.png", 7.00,
    )
    make_squde_data_model_figure(
        squde_data, squde_model,
        output_dir / "m82_squde_mock_data_model_1col.png", 3.35,
    )
    make_squde_data_model_figure(
        squde_data, squde_model,
        output_dir / "m82_squde_mock_data_model_2col.png", 7.00,
    )
    make_resolution_control_figure(
        squde_data, squde_model, chandra_rmf_squde_arf,
        output_dir / "m82_squde_mock_chandra_rmf_1col.png", 3.35,
    )
    make_resolution_control_figure(
        squde_data, squde_model, chandra_rmf_squde_arf,
        output_dir / "m82_squde_mock_chandra_rmf_2col.png", 7.00,
    )
    write_curve_csv(
        squde_model, squde_components, chandra_rmf_squde_arf,
        output_dir / "m82_squde_chandra_kernel_comparison.csv",
    )

    # Display-only min30 sensitivity branch from the exact same saved PHA.
    # Preserve the audited min50 assets above; regroup the in-memory dataset only.
    min50_ylim_top = resolution_control_ylim_top(
        squde_data, squde_model, chandra_rmf_squde_arf
    )
    ungroup(2)
    group_counts(2, 30)
    squde_data_min30 = copy.deepcopy(get_data_plot(2))
    squde_components_min30 = response_folded_components(bundle, 2)
    squde_model_min30 = copy.deepcopy(get_model_plot(2))
    component_validation_min30 = validate_component_closure(
        squde_model_min30, squde_components_min30
    )
    chandra_rmf_squde_arf_min30 = rebin_density_conserving_rate(
        chandra_kernel.channel_lo,
        chandra_kernel.channel_hi,
        chandra_density_with_squde_arf,
        np.asarray(squde_model_min30.xlo, dtype=float),
        np.asarray(squde_model_min30.xhi, dtype=float),
    )
    control_conservation_min30 = validate_rebin_rate_conservation(
        chandra_kernel.channel_lo,
        chandra_kernel.channel_hi,
        chandra_density_with_squde_arf,
        np.asarray(squde_model_min30.xlo, dtype=float),
        np.asarray(squde_model_min30.xhi, dtype=float),
        chandra_rmf_squde_arf_min30,
    )
    shared_ylim_top = max(
        min50_ylim_top,
        resolution_control_ylim_top(
            squde_data_min30,
            squde_model_min30,
            chandra_rmf_squde_arf_min30,
        ),
    )
    make_squde_data_model_figure(
        squde_data_min30,
        squde_model_min30,
        output_dir / "m82_squde_mock_data_model_min30_1col.png",
        3.35,
    )
    make_squde_data_model_figure(
        squde_data_min30,
        squde_model_min30,
        output_dir / "m82_squde_mock_data_model_min30_2col.png",
        7.00,
    )
    make_resolution_control_figure(
        squde_data_min30,
        squde_model_min30,
        chandra_rmf_squde_arf_min30,
        output_dir / "m82_squde_mock_chandra_rmf_min30_1col.png",
        3.35,
        ylim_top=shared_ylim_top,
    )
    make_resolution_control_figure(
        squde_data_min30,
        squde_model_min30,
        chandra_rmf_squde_arf_min30,
        output_dir / "m82_squde_mock_chandra_rmf_min30_2col.png",
        7.00,
        ylim_top=shared_ylim_top,
    )
    make_resolution_control_figure(
        squde_data,
        squde_model,
        chandra_rmf_squde_arf,
        output_dir / "m82_squde_mock_chandra_rmf_min50_slider_2col.png",
        7.00,
        ylim_top=shared_ylim_top,
    )
    write_curve_csv(
        squde_model_min30,
        squde_components_min30,
        chandra_rmf_squde_arf_min30,
        output_dir / "m82_squde_chandra_kernel_comparison_min30.csv",
    )

    def provenance_path(path: Path) -> str:
        """Avoid leaking host paths when an audit writes to a scratch directory."""
        try:
            return str(path.relative_to(ROOT))
        except ValueError:
            return path.name

    provenance = {
        "model": "tbabs_gal * (tbabs_wind * (vapec_cool + vapec_hot + gaussian_0.777 + gaussian_1.234) + tbabs_powerlaw * powerlaw)",
        "source_model_is_true_xspec_vapec": True,
        "chandra": {
            "pha": str(CHANDRA_PHA.relative_to(ROOT)),
            "arf": str(CHANDRA_ARF.relative_to(ROOT)),
            "rmf": str(CHANDRA_RMF.relative_to(ROOT)),
            "statistic": fit_result.statname,
            "method": fit_result.methodname,
            "statval": float(fit_result.statval),
            "dof": int(fit_result.dof),
            "rstat": float(fit_result.rstat),
            "parameter_names": list(fit_result.parnames),
            "parameter_values": [float(value) for value in fit_result.parvals],
        },
        "squde": {
            "pha": provenance_path(squde_pha),
            "arf": str(SQUDE_ARF.relative_to(ROOT)),
            "rmf": str(SQUDE_RMF.relative_to(ROOT)),
            "exposure_s": float(get_data(2).exposure),
            "raw_counts": int(np.sum(get_data(2).counts)),
            "poisson_rng_seed": RNG_SEED,
            "rng_control": "sherpa.astro.ui.set_rng(numpy.random.default_rng(seed))",
            "grouping_variants": {
                "min50": {
                    "method": "sherpa.astro.ui.group_counts",
                    "minimum_counts_per_bin": 50,
                    "notice_energy_keV": [0.3, 4.0],
                    "display_energy_keV": [0.5, 1.7],
                    "display_bins_0p3_4": int(len(squde_data.x)),
                    "plotted_bins_0p5_1p7": int(np.sum((squde_data.x >= 0.5) & (squde_data.x <= 1.7))),
                    "mock_figure_1col": "m82_squde_mock_data_model_1col.png",
                    "mock_figure_2col": "m82_squde_mock_data_model_2col.png",
                    "control_figure_1col": "m82_squde_mock_chandra_rmf_1col.png",
                    "control_figure_2col": "m82_squde_mock_chandra_rmf_2col.png",
                    "slider_control_figure_2col": "m82_squde_mock_chandra_rmf_min50_slider_2col.png",
                    "control_rebin_conservation": control_conservation,
                },
                "min30": {
                    "method": "sherpa.astro.ui.group_counts",
                    "minimum_counts_per_bin": 30,
                    "notice_energy_keV": [0.3, 4.0],
                    "display_energy_keV": [0.5, 1.7],
                    "display_bins_0p3_4": int(len(squde_data_min30.x)),
                    "plotted_bins_0p5_1p7": int(np.sum((squde_data_min30.x >= 0.5) & (squde_data_min30.x <= 1.7))),
                    "same_raw_pha": provenance_path(squde_pha),
                    "shared_control_ylim_top_counts_s_keV": shared_ylim_top,
                    "mock_figure_1col": "m82_squde_mock_data_model_min30_1col.png",
                    "mock_figure_2col": "m82_squde_mock_data_model_min30_2col.png",
                    "control_figure_1col": "m82_squde_mock_chandra_rmf_min30_1col.png",
                    "control_figure_2col": "m82_squde_mock_chandra_rmf_min30_2col.png",
                    "curve_csv": "m82_squde_chandra_kernel_comparison_min30.csv",
                    "control_rebin_conservation": control_conservation_min30,
                },
            },
        },
        "comparison_operator": "F(E) -> SQUDE ARF -> Chandra RMF -> rate-conserving rebin to SQUDE display bins",
        "validation": {
            "explicit_chandra_rmf_vs_sherpa": kernel_validation,
            "chandra_component_closure": chandra_component_validation,
            "squde_component_closure": component_validation,
            "squde_component_closure_min30": component_validation_min30,
        },
    }
    (output_dir / "m82_mock_tutorial_manifest.json").write_text(
        json.dumps(provenance, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )

    print(f"Wrote tutorial products to: {output_dir}")
    print(json.dumps(provenance["validation"], indent=2))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output-dir",
        type=Path,
        help="write products to this directory instead of the canonical data/fits location",
    )
    main(parser.parse_args().output_dir)