#!/usr/bin/env python3
"""Center/North connected mock spectra plus Center Chandra-RMF control.

Creates a new three-line, RGS-style labeled comparison without modifying any
Full-FoV or regional predecessor. Center and North retain their own native
min30 display grids and Poisson realizations; the red control belongs only to
Center and is read from the validated joint-center regional manifest.
"""
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_three_region_mock_tutorial/"
    "m82_three_region_joint_center_manifest.json"
)
SOURCE_MANIFEST_SHA256 = "1d6cbe29c1c0a47c6ced5a1fd9e65cd6d3e5d85207eece65a30777b9e8d7868b"
RGS_LABEL_SOURCE_SHA256 = "1fdfe12927c44ddca08e30f0425fead8c965c707960a3b33c5a2cbbe1007edca"
ATOMDB_LINE_LIST_SHA256 = "e5d8a76f571c8cea6facb3c7dee6d225291583f087c56231a72dd3695a9a88ee"
SQUDE_LINE_MARKER_SOURCE_SHA256 = "1d2250f3b774051d7fa2823f68d0595514e5fa997e7e92c3f8b4ba4239915e25"
LINE_EVIDENCE = ROOT / "data/fits/m82_line_label_v2_atomdb_evidence.json"
LINE_EVIDENCE_SHA256 = "3003f6831c4db4e4a7cbb54cb2a03c28c1e9a110d704aaa9c644a1c9aaa5cd20"
DEFAULT_OUTPUT = ROOT / "data/fits/m82_center_north_labeled_line_comparison_v2"
PREFIX = "m82_center_north_mock_center_chandra_rmf_labeled_v2"
ENERGY_RANGE_KEV = (0.5, 1.7)

# Representative features are grouped from AtomDB 3.1.3 transitions and the
# project SQUDE Fe-L marker implementation, then checked against deterministic
# Center/North SQUDE-RMF model peaks. Labels use keV to match the x axis.
LINE_LABELS = [
    {"energy_keV": 0.568, "label": "O VII triplet", "row": 0},
    {"energy_keV": 0.654, "label": "O VIII Lyα", "row": 1},
    {"energy_keV": 0.704, "label": "Fe XVIII", "row": 2},
    {"energy_keV": 0.727, "label": "Fe XVII blend", "row": 0},
    {"energy_keV": 0.774, "label": "Fe XVIII / O VIII", "row": 1},
    {"energy_keV": 0.794, "label": "Fe XVIII", "row": 4},
    {"energy_keV": 0.812, "label": "Fe XVII", "row": 2},
    {"energy_keV": 0.826, "label": "Fe XVII", "row": 0},
    {"energy_keV": 0.853, "label": "Fe XVIII", "row": 0},
    {"energy_keV": 0.873, "label": "Fe XVIII blend", "row": 3},
    {"energy_keV": 0.905, "label": "Fe XIX / Ne IX", "row": 2},
    {"energy_keV": 0.917, "label": "Fe XIX / Ne IX", "row": 0},
    {"energy_keV": 0.965, "label": "Fe XX", "row": 1},
    {"energy_keV": 0.997, "label": "Ni XIX", "row": 2},
    {"energy_keV": 1.009, "label": "Fe XXI", "row": 0},
    {"energy_keV": 1.022, "label": "Ne X Lyα", "row": 3},
    {"energy_keV": 1.048, "label": "Ni XX", "row": 0},
    {"energy_keV": 1.075, "label": "Fe XVIII / Ne IX", "row": 1},
    {"energy_keV": 1.102, "label": "Fe XVII", "row": 2},
    {"energy_keV": 1.168, "label": "Fe XXIV / Fe XIX", "row": 3},
    {"energy_keV": 1.211, "label": "Ne X Lyβ", "row": 0},
    {"energy_keV": 1.277, "label": "Ne X Lyγ", "row": 1},
    {"energy_keV": 1.343, "label": "Mg XI triplet", "row": 2},
    {"energy_keV": 1.472, "label": "Mg XII Lyα", "row": 3},
    {"energy_keV": 1.579, "label": "Mg XI", "row": 0},
]

plt.style.use(["science", "no-latex"])


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 extract_region(
    manifest: dict[str, Any], key: str
) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
    region = next(item for item in manifest["regions"] if item["key"] == key)
    variant = region["squde"]["grouping_variants"]["min30"]
    grid = variant["display_grid"]
    lo = np.asarray(grid["energy_lo_keV"], dtype=float)
    hi = np.asarray(grid["energy_hi_keV"], dtype=float)
    center = 0.5 * (lo + hi)
    rate = np.asarray(grid["data_rate_counts_s_keV"], dtype=float)
    control = np.asarray(grid["chandra_rmf_control_counts_s_keV"], dtype=float)
    if not (len(lo) == len(hi) == len(rate) == len(control)):
        raise RuntimeError(f"Incomplete min30 display arrays for {key}")
    if np.any(hi <= lo) or np.any(np.diff(lo) < 0):
        raise RuntimeError(f"Invalid display grid for {key}")
    mask = (center >= ENERGY_RANGE_KEV[0]) & (center <= ENERGY_RANGE_KEV[1])
    if int(mask.sum()) != int(variant["plotted_bins_0p5_1p7"]):
        raise RuntimeError(f"Plotted-bin mismatch for {key}")
    return {
        "energy_keV": center[mask].copy(),
        "data_rate_counts_s_keV": rate[mask].copy(),
        "chandra_rmf_control_counts_s_keV": control[mask].copy(),
    }, region


def load_curves() -> tuple[dict[str, np.ndarray], dict[str, Any], dict[str, Any]]:
    if file_sha256(SOURCE_MANIFEST) != SOURCE_MANIFEST_SHA256:
        raise RuntimeError("Regional source-manifest hash mismatch")
    if file_sha256(LINE_EVIDENCE) != LINE_EVIDENCE_SHA256:
        raise RuntimeError("Bounded AtomDB line-evidence hash mismatch")
    manifest = json.loads(SOURCE_MANIFEST.read_text(encoding="utf-8"))
    if manifest.get("version") != "joint-center-v6-complete-grouping-contract":
        raise RuntimeError("Unexpected regional manifest version")
    north, north_meta = extract_region(manifest, "v80_north")
    center, center_meta = extract_region(manifest, "v40_center")
    if array_sha256(center["energy_keV"]) == array_sha256(north["energy_keV"]):
        raise RuntimeError("Center and North unexpectedly share a display grid")
    curves = {
        "center_energy_keV": center["energy_keV"],
        "center_mock_rate_counts_s_keV": center["data_rate_counts_s_keV"],
        "north_energy_keV": north["energy_keV"],
        "north_mock_rate_counts_s_keV": north["data_rate_counts_s_keV"],
        "center_chandra_rmf_control_counts_s_keV": center[
            "chandra_rmf_control_counts_s_keV"
        ],
    }
    return curves, center_meta, north_meta


def annotate_reference_lines(
    ax: plt.Axes, label_ax: plt.Axes, width: float
) -> None:
    narrow = width < 5
    fontsize = 3.0 if narrow else 4.2
    leader_lw = 0.45 if narrow else 0.60
    for record in LINE_LABELS:
        energy = record["energy_keV"]
        ax.axvline(
            energy,
            color="#A0A0A0",
            linewidth=leader_lw,
            linestyle=":",
            alpha=0.42,
            zorder=1,
        )
        row_y = 0.02 + 0.19 * record["row"]
        label_ax.vlines(
            energy,
            0.0,
            max(0.0, row_y - 0.015),
            color="#8A8A8A",
            linewidth=leader_lw,
            linestyle=":",
            alpha=0.70,
            zorder=1,
        )
        label_text = f"{record['label']} {energy:.3f} keV"
        label_ax.text(
            energy,
            row_y,
            label_text,
            rotation=90,
            fontsize=fontsize,
            color="#303640",
            ha="left",
            va="bottom",
            alpha=0.96,
            clip_on=False,
            linespacing=0.92,
            zorder=2,
        )
    label_ax.set_xlim(*ENERGY_RANGE_KEV)
    label_ax.set_ylim(0.0, 1.0)
    label_ax.axis("off")


def render(curves: dict[str, np.ndarray], path: Path, width: float) -> dict[str, Any]:
    dpi = 320
    narrow = width < 5
    height = 4.50 if narrow else 5.20
    fig = plt.figure(figsize=(width, height), dpi=dpi)
    grid = fig.add_gridspec(
        2,
        1,
        height_ratios=[1.90 if narrow else 1.70, 3.0],
        hspace=0.025,
        left=0.16 if narrow else 0.10,
        right=0.985,
        bottom=0.13 if narrow else 0.12,
        top=0.90,
    )
    label_ax = fig.add_subplot(grid[0])
    ax = fig.add_subplot(grid[1], sharex=label_ax)
    ax.plot(
        curves["center_energy_keV"],
        curves["center_mock_rate_counts_s_keV"],
        color="royalblue",
        lw=0.76 if narrow else 1.00,
        alpha=0.98,
        zorder=4,
        label="Center SQUDE mock (50 ks Poisson; min30)",
    )
    ax.plot(
        curves["north_energy_keV"],
        curves["north_mock_rate_counts_s_keV"],
        color="black",
        lw=0.68 if narrow else 0.90,
        alpha=0.88,
        zorder=3,
        label="North SQUDE mock (50 ks Poisson; min30)",
    )
    ax.plot(
        curves["center_energy_keV"],
        curves["center_chandra_rmf_control_counts_s_keV"],
        color="#D62728",
        lw=1.25 if narrow else 1.65,
        alpha=1.0,
        zorder=5,
        label="Center Chandra-RMF control (fixed SQUDE ARF)",
    )
    y_values = np.concatenate(
        [
            curves["center_mock_rate_counts_s_keV"],
            curves["north_mock_rate_counts_s_keV"],
            curves["center_chandra_rmf_control_counts_s_keV"],
        ]
    )
    ax.set_xlim(*ENERGY_RANGE_KEV)
    ax.set_ylim(0.0, float(np.nanmax(y_values) * 1.08))
    annotate_reference_lines(ax, label_ax, width)
    ax.set_xlabel("Energy (keV)")
    ax.set_ylabel(r"Counts s$^{-1}$ keV$^{-1}$")
    fig.suptitle(
        (
            "M82: Center / North mocks + Center CCD"
            if narrow
            else "M82 Center and North SQUDE mocks vs Center CCD resolution"
        ),
        fontsize=7.0 if narrow else 11.8,
        y=0.965,
    )
    ax.legend(
        loc="upper right",
        fontsize=4.15 if narrow else 6.8,
        title=("connected data; red=Center control" if narrow else None),
        title_fontsize=4.0 if narrow else None,
        frameon=True,
        framealpha=0.89 if narrow else 0.83,
        facecolor="white",
        edgecolor="none",
        handlelength=2.4,
    )
    ax.grid(alpha=0.14)
    ax.minorticks_on()
    ax.tick_params(which="both", direction="in", top=True, right=True)
    position = ax.get_position()
    label_position = label_ax.get_position()
    geometry = {
        "figsize_inches": [float(width), float(height)],
        "dpi": dpi,
        "canvas_pixels": [int(width * dpi), int(height * dpi)],
        "axes_bbox_fraction": [
            float(position.x0),
            float(position.y0),
            float(position.width),
            float(position.height),
        ],
        "label_axes_bbox_fraction": [
            float(label_position.x0),
            float(label_position.y0),
            float(label_position.width),
            float(label_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 manifest_array(values: np.ndarray) -> dict[str, Any]:
    return {
        "length": int(len(values)),
        "sha256_le_f8": array_sha256(values),
        "values": [float(value) for value in values],
    }


def main(output_dir: Path) -> None:
    requested = output_dir.expanduser()
    if not requested.is_absolute():
        requested = Path.cwd() / requested
    if os.path.lexists(requested):
        raise FileExistsError(f"Refusing occupied lexical path: {requested}")
    resolved = requested.resolve(strict=False)
    if os.path.lexists(resolved):
        raise FileExistsError(f"Refusing occupied resolved path: {resolved}")
    resolved.parent.mkdir(parents=True, exist_ok=True)
    staging = resolved.parent / f".{resolved.name}.tmp-{os.getpid()}-{uuid.uuid4().hex}"
    staging.mkdir()
    try:
        curves, center_meta, north_meta = load_curves()
        outputs: dict[str, Any] = {}
        for role, width in (("1col", 3.35), ("2col", 7.00)):
            path = staging / f"{PREFIX}_{role}.png"
            outputs[role] = {
                "path": path.name,
                "geometry": render(curves, path, width),
                "sha256": file_sha256(path),
            }
        center_control = center_meta["squde"]["resolution_control"]
        manifest = {
            "version": "center-north-labeled-line-comparison-v2-energy-band",
            "scientific_role": "three-line regional response-resolution visualization",
            "supersedes": {
                "version": "center-north-labeled-line-comparison-v1",
                "manifest_sha256": "ce95210f84458edf606a464c6ec4701d0a8d162ab1a41b1430027ad23b7113cc",
                "reason": "V1 used wavelength text on a keV axis, placed labels over data, and omitted prominent deterministic-model peaks",
            },
            "curve_contract": {
                "center_blue": "Center 40x80 arcsec seeded 50 ks min30 SQUDE mock data on its native grouped display grid",
                "north_black": "North 80x80 arcsec seeded 50 ks min30 SQUDE mock data on its native grouped display grid",
                "center_red": "Center source folded through each epoch-specific Chandra RMF at fixed SQUDE ARF, then exposure-weighted and conservatively projected to the Center min30 display grid",
                "normalization": "absolute counts s^-1 keV^-1; no area, peak, or integral renormalization",
                "grid_policy": "Center blue/red share the Center min30 grid; North retains its distinct native min30 grid",
            },
            "scope_limitations": [
                "Center and North have different apertures and source populations; line heights are not same-area surface-brightness measurements",
                "blue and black are separate Poisson realizations with different grouped bins",
                "red is a deterministic Center response control, not Chandra data",
                "line labels are AtomDB/SQUDE catalog-pinned energy-reference guides, not detections or unique deblended fit assignments",
                "grouped blend labels name representative dominant transitions near deterministic SQUDE-RMF model peaks",
                "all label text is isolated in a dedicated annotation band and uses keV to match the x axis",
                "figure is not fit-quality or physical-parameter evidence",
            ],
            "producer": {
                "path": "scripts/make_m82_center_north_labeled_line_comparison_v2.py",
                "sha256": file_sha256(Path(__file__).resolve()),
                "replay_command": "python -B scripts/make_m82_center_north_labeled_line_comparison_v2.py --output-dir <new-nonexistent-output-dir>",
            },
            "source_manifest": {
                "path": str(SOURCE_MANIFEST.relative_to(ROOT)),
                "expected_sha256": SOURCE_MANIFEST_SHA256,
                "observed_sha256": file_sha256(SOURCE_MANIFEST),
                "version": "joint-center-v6-complete-grouping-contract",
            },
            "regions": {
                "center": {
                    "key": center_meta["key"],
                    "geometry": center_meta["geometry"],
                    "seed": center_meta["squde"]["seed"],
                    "exposure_s": center_meta["squde"]["exposure_s"],
                    "plotted_bins": int(len(curves["center_energy_keV"])),
                },
                "north": {
                    "key": north_meta["key"],
                    "geometry": north_meta["geometry"],
                    "seed": north_meta["squde"]["seed"],
                    "exposure_s": north_meta["squde"]["exposure_s"],
                    "plotted_bins": int(len(curves["north_energy_keV"])),
                },
            },
            "center_control_provenance": {
                "operator": center_control["operator"],
                "source_semantics": center_control["source_semantics"],
                "aggregation": center_control["aggregation"],
                "weight_basis": center_control["weight_basis"],
                "weight_sum": center_control["weight_sum"],
                "total_chandra_exposure_s": center_control[
                    "total_chandra_exposure_s"
                ],
                "fixed_arf_sha256": center_control["fixed_arf_sha256"],
                "epochs": len(center_control["epochs"]),
            },
            "line_labels": {
                "sources": [
                    {"path": "atomdb/apec_v3.1.3_linelist.fits", "sha256": ATOMDB_LINE_LIST_SHA256, "role": "AtomDB 3.1.3 transition energies and temperature-dependent emissivities"},
                    {"path": "SQUDE_propsoal/M82_data/mos1/1_fake_spectra.py", "sha256": SQUDE_LINE_MARKER_SOURCE_SHA256, "role": "project Fe-L/O/Ne/Mg representative marker implementation"},
                    {"path": "M31Center/rgs_emission_lines.py", "sha256": RGS_LABEL_SOURCE_SHA256, "role": "RGS-style annotation provenance"}
                ],
                "bounded_evidence": {"path": str(LINE_EVIDENCE.relative_to(ROOT)), "sha256": LINE_EVIDENCE_SHA256, "labels": 25},
                "selection_rule": "representative AtomDB/SQUDE transition groups cross-matched to deterministic Center/North SQUDE-RMF model peaks; not a complete emissivity table",
                "axis_and_label_unit": "keV",
                "labels": LINE_LABELS,
                "user_requested_coverage": [
                    {"requested_near_keV": 0.78, "anchor_keV": 0.774, "assignment": "Fe XVIII / O VIII blend"},
                    {"requested_near_keV": 0.87, "anchor_keV": 0.873, "assignment": "Fe XVIII blend"},
                    {"requested_near_keV": 0.99, "anchor_keV": 0.997, "assignment": "Ni XIX; adjacent Fe XXI is separately marked at 1.009 keV"},
                    {"requested_near_keV": 1.08, "anchor_keV": 1.075, "assignment": "Fe XVIII / Ne IX blend"},
                    {"requested_near_keV": 1.21, "anchor_keV": 1.211, "assignment": "Ne X Ly-beta"}
                ],
                "n_vii_boundary_exclusion": "0.500 keV anchor omitted because it lies on the displayed left boundary",
                "artist_contract": "25 light full-height guides in the spectrum panel plus 25 energy labels in a separate five-row top annotation band",
            },
            "runtime": {
                "python": sys.version.split()[0],
                "numpy": np.__version__,
                "matplotlib": matplotlib.__version__,
                "scienceplots": package_version("SciencePlots"),
            },
            "colors": {
                "center_mock": {"name": "royalblue", "hex": "#4169E1"},
                "north_mock": {"name": "black", "hex": "#000000"},
                "center_control": {"name": "red", "hex": "#D62728"},
            },
            "arrays": {
                name: manifest_array(values) for name, values in curves.items()
            },
            "artist_counts": {
                "scientific_lines": 3,
                "reference_guides": len(LINE_LABELS),
                "reference_text_labels": len(LINE_LABELS),
                "annotation_leader_collections": len(LINE_LABELS),
                "labels_over_data_panel": 0,
                "errorbars": 0,
                "markers": 0,
            },
            "outputs": outputs,
        }
        manifest_path = staging / f"{PREFIX}_manifest.json"
        manifest_path.write_text(
            json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        if os.path.lexists(requested) or os.path.lexists(resolved):
            raise FileExistsError("Output path appeared during build")
        os.rename(staging, resolved)
    except BaseException:
        shutil.rmtree(staging, 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="new, nonexistent output directory",
    )
    main(parser.parse_args().output_dir)