#!/usr/bin/env python
"""AGORA figure using the published Cloudy C17.03 table. [AI-Codex]"""

import os
from pathlib import Path
import sys

# Use the repository helper, or thorplot.py downloaded alongside this script. [AI-Codex]
for parent in Path(__file__).resolve().parents:
    if (parent / "scripts/thorplot.py").is_file():
        sys.path.insert(0, str(parent / "scripts"))
        break
import thorplot as tp
import numpy as np
import zarr
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

DATA = os.getcwd()
PROJ = f"{DATA}/agora-projection.zr"
OUTDIR = DATA
C_KMS, KPC_CM, MP, XH = 299792.458, 3.0856775814913673e21, 1.6726e-24, 0.76
CMAP_NH, HW, ACC = "cividis", 48.0, "#ff5a36"  # accent for the sightline marker
# NIST centres used by the configured THOR transitions (uv_optical.tsv). [AI-Codex]
IONS = [
    ("O VI 1032", 1031.9270, "#ff5a36"),
    ("C IV 1548", 1548.2043, "#2fae74"),
    ("Si IV 1394", 1393.7546, "#a07be6"),
]
EDGE = "#dcdcdc"  # main frame/ticks: light, so they show on the dark cividis edges
# inset sits over the DARK colormap in BOTH page themes -> style it as translucent "dark glass"
# with LIGHT inner elements (independent of the page fg), so labels read in light mode too.
INSET_BG, INSET_FG, INSET_ALPHA = "#0a0a0a", "#f2f2f2", 0.5

pd = zarr.open(PROJ, mode="r")["projections/density"]
# Raw axes are camera right (-y), up (+z). imshow expects rows z, columns y. [AI-Codex]
col = np.asarray(pd)[::-1, :].T
Pp = dict(pd.attrs)["Parameters"]
half_p = 0.5 * Pp["widths"][0] * Pp["length_unit_cm"] / KPC_CM
logNH = np.log10(np.where(col > 0, col * XH / MP, np.nan))  # true H column [cm^-2]

ion_spectra = []
ray = zarr.open(f"{DATA}/agora-cloudy17-strong-lines.zr", mode="r")[
    "sightlines/default"
]
length_pkpc = Pp["length_unit_cm"] / KPC_CM
SL_Y, SL_Z = (
    np.asarray(ray["start"])[0, 1:] - np.asarray(Pp["position"])[1:]
) * length_pkpc
lam = ray["lam"][:]
tau = ray["tau"][0, :]
for name, lam0, color in IONS:
    velocity = (lam / lam0 - 1) * C_KMS
    selected = np.abs(velocity) < 500
    ion_spectra.append((name, velocity[selected], np.exp(-tau[selected]), color))

cy, cz = 0.0, 0.0  # galaxy / sightline centre
yy = ((np.arange(col.shape[1]) + 0.5) / col.shape[1] * 2 - 1) * half_p
zz = ((np.arange(col.shape[0]) + 0.5) / col.shape[0] * 2 - 1) * half_p
_crop = logNH[np.ix_(np.abs(zz - cz) <= HW, np.abs(yy - cy) <= HW)]
vmin = float(np.nanpercentile(_crop, 4))
vmax = 22.5  # fixed top: full range (24.6) washes out the CGM


def render(fg):
    """Render the original layout with the new C17 spectra. [AI-Codex]"""
    fig, axp = plt.subplots(figsize=(6.6, 6.0))
    im = axp.imshow(
        logNH,
        origin="lower",
        cmap=CMAP_NH,
        extent=(-half_p, half_p, -half_p, half_p),
        vmin=vmin,
        vmax=vmax,
        aspect="equal",
    )
    axp.set_xlim(cy - HW, cy + HW)
    axp.set_ylim(cz - HW, cz + HW)
    axp.set_xlabel("y [pkpc]")
    axp.set_ylabel("z [pkpc]")
    axp.set_title(
        "AGORA z = 4 GIZMO CGM",
        color=fg,
        fontsize=11,
    )
    cb = fig.colorbar(im, ax=axp, fraction=0.046, pad=0.02)
    cb.set_label(r"$\log_{10}\,N_{\rm H}\ [{\rm cm^{-2}}]$", color=fg)
    cb.ax.yaxis.set_tick_params(color=EDGE, labelcolor=fg)
    cb.outline.set_edgecolor(EDGE)
    axp.tick_params(color=EDGE, labelcolor=fg)
    axp.xaxis.label.set_color(fg)
    axp.yaxis.label.set_color(fg)
    for s in axp.spines.values():
        s.set_color(EDGE)
        s.set_linewidth(1.3)

    # Off-centre sightline along +x, into the page
    axp.plot(SL_Y, SL_Z, marker="o", ms=14, mfc="none", mec=ACC, mew=2.2, zorder=5)
    axp.plot(SL_Y, SL_Z, marker=".", ms=5, color=ACC, zorder=5)
    axp.annotate(
        "sightline (b ≈ 12 kpc, +x)",
        (SL_Y, SL_Z),
        xytext=(SL_Y, SL_Z + HW * 0.34),
        textcoords="data",
        color=ACC,
        fontsize=9.5,
        ha="center",
        fontweight="bold",
        arrowprops=dict(arrowstyle="->", color=ACC, lw=1.3),
    )

    # embedded spectrum: this one sightline in three ions, one mini-panel each
    # (stacked, shared velocity axis; lower-left, over the diffuse outskirts)
    fx0, fw, h, g = 0.105, 0.44, 0.082, 0.018
    n = len(ion_spectra)
    for k, (name, v, flux, color) in enumerate(ion_spectra):
        a = axp.inset_axes((fx0, 0.10 + (n - 1 - k) * (h + g), fw, h))  # k=0 on top
        a.add_patch(
            Rectangle(
                (0, 0),
                1,
                1,
                transform=a.transAxes,
                facecolor=INSET_BG,
                edgecolor="none",
                alpha=INSET_ALPHA,
                zorder=0,
            )
        )
        a.plot(v, flux, color=color, lw=1.4, zorder=3)
        a.axhline(1.0, color=INSET_FG, lw=0.4, ls=":", alpha=0.4, zorder=2)
        a.set_xlim(-500, 500)
        a.set_ylim(-0.08, 1.2)
        a.text(
            0.035,
            0.56,
            name,
            transform=a.transAxes,
            color=color,
            fontsize=6.8,
            fontweight="bold",
            zorder=4,
        )
        a.set_yticks([0, 1])
        a.tick_params(colors=INSET_FG, labelsize=5.8, length=2)
        for s in a.spines.values():
            s.set_color(INSET_FG)
        if k == 0:
            a.set_title(
                r"flux $e^{-\tau}$  (±500 km/s)", color=INSET_FG, fontsize=7.5, pad=2
            )
        if k < n - 1:
            a.set_xticklabels([])
        else:
            a.set_xlabel("v [km/s]", color=INSET_FG, fontsize=7)
    return fig


for suffix, fg in [("", "#222222"), ("_dark", "#e8e8e8")]:
    with tp.style(dark=bool(suffix)):
        fig = render(fg)
        out = tp.save_doc_figure(
            fig, f"result{suffix}", fmt="png", output_dir=OUTDIR, dpi=140
        )
        plt.close(fig)
        print("wrote", out)
for name, v, flux, color in ion_spectra:
    print(f"  {name}: peakτ={(-np.log(np.clip(flux,1e-30,1))).max():.2f}")
print(
    f"off-centre LOS, 3 strong lines; frame ±{HW} kpc; logN_H crop [{vmin:.1f},{vmax:.1f}]"
)