#!/usr/bin/env python3
"""Add RGS-style reference line labels to the promoted connected-line contrast.

This is a new, non-overwriting visualization.  It reuses the exact V5 blue/red
arrays and adds only vertical reference guides and text labels.  The labels are
laboratory-energy guides, not independent line detections or fit assignments.
"""
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

import make_m82_line_resolution_contrast as base

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = ROOT / "data/fits/m82_line_resolution_contrast_labeled_v2"
PREFIX = "m82_squde_mock_min30_connected_vs_chandra_ccd_resolution_labeled_v2"
V5_MANIFEST = (
    ROOT
    / "data/fits/m82_line_resolution_contrast/"
    "m82_squde_mock_min30_connected_vs_chandra_ccd_resolution_manifest.json"
)
V5_MANIFEST_SHA256 = "ccbfdb17f5c60a7404b2096fbe205035beaad624eb9b87d909b4c705c0c65d5b"
V5_PRODUCER_SHA256 = "20269596175e950ee8c84636bd69307795212331b3f68424e02a599a619844d0"
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"

# Representative features are grouped from AtomDB 3.1.3 transitions and the
# project SQUDE Fe-L marker implementation, then checked against deterministic
# SQUDE-RMF model peaks.  Labels use keV to match the x axis.  They are
# reference guides, not detections or unique deblended identifications.
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 verify_inputs() -> dict[str, Any]:
    if file_sha256(V5_MANIFEST) != V5_MANIFEST_SHA256:
        raise RuntimeError("Promoted V5 manifest hash mismatch")
    if file_sha256(Path(base.__file__).resolve()) != V5_PRODUCER_SHA256:
        raise RuntimeError("Promoted V5 producer hash mismatch")
    if file_sha256(LINE_EVIDENCE) != LINE_EVIDENCE_SHA256:
        raise RuntimeError("Bounded AtomDB line-evidence hash mismatch")
    manifest = json.loads(V5_MANIFEST.read_text(encoding="utf-8"))
    if manifest.get("version") != "line-resolution-contrast-v5-symlink-safe":
        raise RuntimeError("Unexpected V5 manifest version")
    return manifest


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(0.50, 1.70)
    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
    is_narrow = width < 5
    height = 4.50 if is_narrow else 5.20
    fig = plt.figure(figsize=(width, height), dpi=dpi)
    grid = fig.add_gridspec(
        2,
        1,
        height_ratios=[1.90 if is_narrow else 1.70, 3.0],
        hspace=0.025,
        left=0.16 if is_narrow else 0.10,
        right=0.985,
        bottom=0.13 if is_narrow else 0.12,
        top=0.90,
    )
    label_ax = fig.add_subplot(grid[0])
    ax = fig.add_subplot(grid[1], sharex=label_ax)

    xmin, xmax = 0.50, 1.70
    mock_mask = (curves["mock_energy_keV"] >= xmin) & (
        curves["mock_energy_keV"] <= xmax
    )
    ccd_mask = (curves["ccd_energy_keV"] >= xmin) & (
        curves["ccd_energy_keV"] <= xmax
    )
    ax.plot(
        curves["mock_energy_keV"][mock_mask],
        curves["mock_rate_counts_s_keV"][mock_mask],
        color="#6FA8FF",
        lw=0.72 if is_narrow else 0.92,
        alpha=0.98,
        zorder=3,
        label=(
            "SQUDE mock (50 ks Poisson; connected data)"
            if is_narrow
            else "SQUDE mock, min30 (connected Poisson data; not a model)"
        ),
    )
    ax.plot(
        curves["ccd_energy_keV"][ccd_mask],
        curves["ccd_rate_counts_s_keV"][ccd_mask],
        color="#D62728",
        lw=1.25 if is_narrow else 1.65,
        alpha=1.0,
        zorder=4,
        label=(
            "Chandra ACIS-RMF control"
            if is_narrow
            else "Chandra ACIS RMF control (fixed SQUDE ARF)"
        ),
    )

    y_visible = np.concatenate(
        [
            curves["mock_rate_counts_s_keV"][mock_mask],
            curves["ccd_rate_counts_s_keV"][ccd_mask],
        ]
    )
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(0.0, float(np.nanmax(y_visible) * 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: SQUDE vs CCD · energy-line guide"
            if is_narrow
            else "M82: SQUDE mock vs CCD resolution · energy-line guide"
        ),
        fontsize=7.2 if is_narrow else 12.0,
        y=0.965,
    )
    ax.legend(
        loc="upper right",
        fontsize=4.5 if is_narrow else 7.2,
        title="Blue=data, not model · fixed SQUDE ARF" if is_narrow else None,
        title_fontsize=4.3 if is_narrow else None,
        frameon=True,
        framealpha=0.88 if is_narrow else 0.82,
        facecolor="white",
        edgecolor="none",
        handlelength=2.5,
    )
    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 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 path: {requested_output}")
    output_dir = requested_output.resolve(strict=False)
    if os.path.lexists(output_dir):
        raise FileExistsError(f"Refusing occupied resolved path: {output_dir}")
    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:
        v5_manifest = verify_inputs()
        curves, source_manifest = base.load_curves()
        outputs: dict[str, Any] = {}
        for role, width in (("1col", 3.35), ("2col", 7.00)):
            path = staging_dir / f"{PREFIX}_{role}.png"
            outputs[role] = {
                "path": path.name,
                "geometry": render(curves, path, width),
                "sha256": file_sha256(path),
            }

        plotted = v5_manifest["plotted_subset"]
        manifest = {
            "version": "line-resolution-contrast-labeled-v2-energy-band",
            "scientific_role": "annotation-only successor to promoted V5; blue/red scientific arrays are unchanged",
            "supersedes": {
                "version": "line-resolution-contrast-labeled-v1",
                "manifest_sha256": "7c9c3253ba47f78a5149fc00ceb5796f5c41bc3430cc29c45088afa7b1cbfd60",
                "reason": "V1 used wavelength text on a keV axis, placed labels over data, and omitted prominent deterministic-model peaks",
            },
            "scope_limitations": [
                "line labels are AtomDB/SQUDE catalog-pinned energy-reference guides, not independent 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",
                "blue spikes include Poisson, grouping, and unequal-bin rate-density effects",
                "red is a deterministic response control, not Chandra data",
                "figure is not fit-quality or physical-parameter evidence",
            ],
            "producer": {
                "path": "scripts/make_m82_line_resolution_contrast_labeled_v2.py",
                "sha256": file_sha256(Path(__file__).resolve()),
                "replay_command": "python -B scripts/make_m82_line_resolution_contrast_labeled_v2.py --output-dir <new-nonexistent-output-dir>",
            },
            "promoted_v5": {
                "manifest_path": str(V5_MANIFEST.relative_to(ROOT)),
                "manifest_sha256": V5_MANIFEST_SHA256,
                "producer_path": "scripts/make_m82_line_resolution_contrast.py",
                "producer_sha256": V5_PRODUCER_SHA256,
                "source_manifest_sha256": v5_manifest["source_manifest"]["observed_sha256"],
            },
            "label_catalog": {
                "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 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"),
            },
            "plotted_subset": plotted,
            "array_identity": {
                name: {
                    "length": int(len(values)),
                    "sha256_le_f8": array_sha256(values),
                }
                for name, values in curves.items()
            },
            "artist_counts": {
                "scientific_lines": 2,
                "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_dir / 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_output) or os.path.lexists(output_dir):
            raise FileExistsError("Output path appeared during build")
        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="new, nonexistent output directory",
    )
    main(parser.parse_args().output_dir)