#!/usr/bin/env python3
"""Render a non-overwriting line-style SQUDE-vs-CCD resolution contrast.

The blue curve connects the centers of the existing min30 SQUDE mock display
bins.  It is a visualization of the seeded mock realization, not a model.
The red curve is the same-source, fixed-SQUDE-ARF Chandra-ACIS-RMF control on
the identical min30 display grid.  No fit, fake_pha draw, response fold,
resampling, or existing figure is changed by this script.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys
import uuid
from importlib.metadata import version as package_version
from pathlib import Path
from typing import Any

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import scienceplots  # noqa: F401

ROOT = Path(__file__).resolve().parents[1]
SOURCE_MANIFEST = (
    ROOT
    / "data/fits/m82_mock_tutorial_response_comparison/m82_mock_tutorial_manifest.json"
)
SOURCE_MANIFEST_SHA256 = "e9e0fa3c6005136227b4f29648c289ecca8d315c5220b20887028217e84d7f7f"
DEFAULT_OUTPUT = ROOT / "data/fits/m82_line_resolution_contrast"
PREFIX = "m82_squde_mock_min30_connected_vs_chandra_ccd_resolution"


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def array_sha256(values: np.ndarray) -> str:
    return hashlib.sha256(np.asarray(values, dtype="<f8").tobytes()).hexdigest()


def load_curves() -> tuple[dict[str, np.ndarray], dict[str, Any]]:
    observed_source_sha = file_sha256(SOURCE_MANIFEST)
    if observed_source_sha != SOURCE_MANIFEST_SHA256:
        raise RuntimeError(
            f"Source manifest pin mismatch: {observed_source_sha} != {SOURCE_MANIFEST_SHA256}"
        )
    manifest = json.loads(SOURCE_MANIFEST.read_text(encoding="utf-8"))
    variant = manifest["squde"]["grouping_variants"]["min30"]
    display = variant["display_grid"]
    mock_lo = np.asarray(display["energy_lo_keV"], dtype=float)
    mock_hi = np.asarray(display["energy_hi_keV"], dtype=float)
    mock_rate = np.asarray(display["data_rate_counts_s_keV"], dtype=float)
    ccd_rate = np.asarray(display["chandra_rmf_control_counts_s_keV"], dtype=float)

    if not (len(mock_lo) == len(mock_hi) == len(mock_rate) == len(ccd_rate)):
        raise RuntimeError("Mock and CCD-control display arrays do not share one grid")
    if not np.all(mock_hi > mock_lo):
        raise RuntimeError("Energy bins must have positive widths")

    expected = display["sha256_le_f8"]
    if array_sha256(mock_rate) != expected["data_rate_counts_s_keV"]:
        raise RuntimeError("Mock-rate array does not close to the source manifest")
    if array_sha256(ccd_rate) != expected["chandra_rmf_control_counts_s_keV"]:
        raise RuntimeError("CCD-control display array does not close to the source manifest")

    centers = 0.5 * (mock_lo + mock_hi)
    return {
        "mock_energy_keV": centers,
        "mock_rate_counts_s_keV": mock_rate,
        "ccd_energy_keV": centers.copy(),
        "ccd_rate_counts_s_keV": ccd_rate,
    }, manifest


def render(curves: dict[str, np.ndarray], path: Path, width: float) -> dict[str, Any]:
    is_narrow = width < 4.0
    height = 2.55 if is_narrow else 3.75
    dpi = 320
    mask_mock = (
        (curves["mock_energy_keV"] >= 0.50)
        & (curves["mock_energy_keV"] <= 1.70)
    )
    mask_ccd = (
        (curves["ccd_energy_keV"] >= 0.50)
        & (curves["ccd_energy_keV"] <= 1.70)
    )
    ymax = 1.08 * max(
        float(np.max(curves["mock_rate_counts_s_keV"][mask_mock])),
        float(np.max(curves["ccd_rate_counts_s_keV"][mask_ccd])),
    )

    with plt.style.context(["science", "no-latex"]):
        fig, ax = plt.subplots(figsize=(width, height))
        if is_narrow:
            fig.subplots_adjust(left=0.18, right=0.97, bottom=0.19, top=0.78)
        else:
            fig.subplots_adjust(left=0.105, right=0.985, bottom=0.16, top=0.84)

        ax.plot(
            curves["mock_energy_keV"][mask_mock],
            curves["mock_rate_counts_s_keV"][mask_mock],
            color="#6FA8FF",
            lw=0.85 if is_narrow else 1.15,
            alpha=0.98,
            label=(
                "SQUDE mock (50 ks Poisson; connected data)"
                if is_narrow
                else "SQUDE mock, min30 (connected Poisson data; not a model)"
            ),
            zorder=3,
        )
        ax.plot(
            curves["ccd_energy_keV"][mask_ccd],
            curves["ccd_rate_counts_s_keV"][mask_ccd],
            color="#D62728",
            lw=1.25 if is_narrow else 1.65,
            label=(
                "Chandra ACIS-RMF control"
                if is_narrow
                else "Chandra ACIS RMF control (fixed SQUDE ARF)"
            ),
            zorder=4,
        )

        ax.set_xlim(0.50, 1.70)
        ax.set_ylim(0.0, ymax)
        ax.set_xlabel("Energy (keV)")
        ax.set_ylabel(r"Counts s$^{-1}$ keV$^{-1}$")
        ax.set_title(
            "M82: SQUDE mock vs CCD resolution",
            fontsize=8.0 if is_narrow else 11.5,
            pad=4,
        )
        if not is_narrow:
            ax.text(
                0.02,
                0.965,
                "Expected models: same source + SQUDE ARF; RMF-only difference",
                transform=ax.transAxes,
                va="top",
                fontsize=6.8,
                color="0.18",
            )
            ax.text(
                0.02,
                0.905,
                "Blue: seeded 50 ks Poisson realization",
                transform=ax.transAxes,
                va="top",
                fontsize=6.8,
                color="0.18",
            )
        ax.legend(
            loc="upper right",
            fontsize=4.8 if is_narrow else 7.2,
            title="Blue=data, not model · fixed SQUDE ARF" if is_narrow else None,
            title_fontsize=4.6 if is_narrow else None,
            frameon=is_narrow,
            framealpha=0.84 if is_narrow else None,
            facecolor="white" if is_narrow else None,
            edgecolor="none" if is_narrow else None,
            handlelength=2.5,
        )
        ax.grid(alpha=0.16)
        ax.tick_params(which="both", direction="in", top=True, right=True)

        canvas_pixels = [int(width * dpi), int(height * dpi)]
        position = ax.get_position()
        geometry = {
            "canvas_pixels": canvas_pixels,
            "axes_bbox_fraction_lower_left": [
                float(position.x0),
                float(position.y0),
                float(position.width),
                float(position.height),
            ],
            "xlim_keV": [float(value) for value in ax.get_xlim()],
            "ylim_counts_s_keV": [float(value) for value in ax.get_ylim()],
        }
        with plt.rc_context({"savefig.bbox": None}):
            fig.savefig(path, dpi=dpi, bbox_inches=None)
        plt.close(fig)
    return geometry


def main(output_dir: Path) -> None:
    requested_output = output_dir.expanduser()
    if not requested_output.is_absolute():
        requested_output = Path.cwd() / requested_output
    if os.path.lexists(requested_output):
        raise FileExistsError(
            f"Refusing occupied lexical output path: {requested_output}. "
            "Choose a new path."
        )
    output_dir = requested_output.resolve(strict=False)
    if os.path.lexists(output_dir):
        raise FileExistsError(
            f"Refusing resolved occupied output path: {output_dir}. "
            "Choose a new path."
        )
    output_dir.parent.mkdir(parents=True, exist_ok=True)
    staging_dir = output_dir.parent / (
        f".{output_dir.name}.tmp-{os.getpid()}-{uuid.uuid4().hex}"
    )
    staging_dir.mkdir()

    try:
        curves, source = load_curves()
        outputs: dict[str, Any] = {}
        for label, width in (("1col", 3.35), ("2col", 7.00)):
            path = staging_dir / f"{PREFIX}_{label}.png"
            geometry = render(curves, path, width)
            outputs[label] = {
                "path": path.name,
                "sha256": file_sha256(path),
                "geometry": geometry,
            }

        plotted_mask = (
            (curves["mock_energy_keV"] >= 0.50)
            & (curves["mock_energy_keV"] <= 1.70)
        )
        plotted_arrays = {
            "energy_keV": curves["mock_energy_keV"][plotted_mask],
            "squde_mock_rate_counts_s_keV": curves["mock_rate_counts_s_keV"][plotted_mask],
            "chandra_rmf_control_counts_s_keV": curves["ccd_rate_counts_s_keV"][plotted_mask],
        }
        if len(plotted_arrays["energy_keV"]) != 971:
            raise RuntimeError("Expected exactly 971 plotted min30 bins")

        producer_path = Path(__file__).resolve()
        provenance = {
            "version": "line-resolution-contrast-v5-symlink-safe",
            "scientific_role": "alternative visualization only; no new fit, simulation, or response fold",
            "scope_limitations": [
                "response-resolution teaching control only; not fit-quality evidence or physical-parameter inference",
                "connected mock spikes include one Poisson realization, min30 grouping, and unequal-bin rate-density effects; do not identify individual physical lines from this display alone",
                "the Chandra-RMF curve is a deterministic control and is not a Chandra data spectrum",
            ],
            "producer": {
                "path": "scripts/make_m82_line_resolution_contrast.py",
                "sha256": file_sha256(producer_path),
                "replay_command": "python -B scripts/make_m82_line_resolution_contrast.py --output-dir <new-nonexistent-output-dir>",
                "output_policy": "default no-clobber; lexical output is checked with lexists before path resolution, complete products are built off-path, and the directory is renamed only after validation",
            },
            "runtime": {
                "python": sys.version.split()[0],
                "numpy": np.__version__,
                "matplotlib": matplotlib.__version__,
                "scienceplots": package_version("SciencePlots"),
            },
            "source_manifest": {
                "path": str(SOURCE_MANIFEST.relative_to(ROOT)),
                "expected_sha256": SOURCE_MANIFEST_SHA256,
                "observed_sha256": file_sha256(SOURCE_MANIFEST),
                "variant": "min30",
            },
            "line_contract": {
                "blue": "seeded SQUDE mock data-rate values connected at min30 display-bin centers; Poisson data, not a model, and no error bars are drawn",
                "red": "same-source Chandra ACIS RMF response control at fixed SQUDE ARF, read directly from the same min30 display grid and connected through the identical bin centers",
                "common_grid": "blue and red use identical min30 display-bin centers; this renderer adds no interpolation, smoothing, renormalization, or resampling",
                "inherited_upstream_operations": [
                    "R16 interpolates SQUDE ARF values at Chandra-RMF true-energy centers",
                    "R16 conservatively projects the Chandra-EBOUNDS control to the SQUDE model grid",
                    "R16 conservatively projects the model-grid control to the min30 display grid",
                ],
                "difference_isolation": "expected response models use the same source and SQUDE ARF and differ only by RMF; the blue displayed realization additionally contains seeded Poisson noise and grouping/rate-density effects",
                "energy_keV": [0.50, 1.70],
                "rate_unit": "counts s^-1 keV^-1",
            },
            "arrays": {
                name: {
                    "length": int(len(values)),
                    "sha256_le_f8": array_sha256(values),
                }
                for name, values in curves.items()
            },
            "plotted_subset": {
                "length": 971,
                "arrays": {
                    name: {
                        "values": values.tolist(),
                        "sha256_le_f8": array_sha256(values),
                    }
                    for name, values in plotted_arrays.items()
                },
            },
            "outputs": outputs,
        }
        manifest_path = staging_dir / f"{PREFIX}_manifest.json"
        manifest_path.write_text(
            json.dumps(provenance, indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )
        if os.path.lexists(requested_output) or os.path.lexists(output_dir):
            raise FileExistsError(
                f"Output path appeared during build: {requested_output} -> {output_dir}"
            )
        os.rename(staging_dir, output_dir)
    except BaseException:
        shutil.rmtree(staging_dir, ignore_errors=True)
        raise

    print(f"Wrote {outputs['1col']['path']}")
    print(f"Wrote {outputs['2col']['path']}")
    print(f"Wrote {manifest_path.name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT,
        help="independent output directory (default: %(default)s)",
    )
    main(parser.parse_args().output_dir)