"""Plot THOR's Lyα surface-brightness map and aperture spectra. [AI-Codex]"""

import argparse

import h5py
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib.patches import Circle
import numpy as np
import thorplot as tp

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--dark", action="store_true")
args = parser.parse_args()
with tp.style(dark=args.dark):

    with h5py.File("output-tng-lya/peel/data.h5", "r") as f:
        image = f["sbmap/image"][:].T
        map_attrs = dict(f["sbmap/image"].attrs)
        spectra = f["aperturespectrum/dlambda"][:]
        spectrum_attrs = dict(f["aperturespectrum/dlambda"].attrs)

    # THOR stores map bounds and aperture radii as fractions of the domain length.
    length_cm = float(map_attrs["length_unit_cm"])
    if not np.isfinite(length_cm) or length_cm <= 0:
        raise ValueError("Map length_unit_cm must be finite and positive.")
    # The aperture spectrum may omit this shared length unit.
    if "length_unit_cm" in spectrum_attrs:
        spectrum_length_cm = float(spectrum_attrs["length_unit_cm"])
        if not np.isclose(length_cm, spectrum_length_cm, rtol=1e-6, atol=0):
            raise ValueError("Map and aperture spectrum length units disagree.")
    box_kpc = length_cm / 3.0856775814913673e21
    radii_kpc = spectrum_attrs["aperture_radii"] * box_kpc
    extent = [(map_attrs[k] - 0.5) * box_kpc
              for k in ("left_x", "right_x", "left_y", "right_y")]
    pixel_area = ((extent[1] - extent[0]) / image.shape[1]
                  * (extent[3] - extent[2]) / image.shape[0])
    # Both figures use the whole-cutout spectrum's total flux as their reference.
    # The map is relative flux per physical area, not observed flux per arcsec².
    reference_flux = spectra[-1].sum()
    if not reference_flux > 0:
        raise ValueError("The whole-cutout aperture contains no flux.")
    sb = image / (reference_flux * pixel_area)
    positive = sb[sb > 0]
    if not positive.size:
        raise ValueError("The surface-brightness map contains no flux.")
    colors = [tp.code_color("thor", args.dark),
              tp.code_color("colt", args.dark), tp.code_color("rascas", args.dark)]
    labels = [f"{r:g} kpc radius" for r in radii_kpc[:-1]] + ["Whole cutout"]
    suffix = "_dark" if args.dark else ""

    fig, ax = plt.subplots(figsize=(tp.DOC_FULL_WIDTH, 7.625), layout="constrained")
    im = ax.imshow(np.ma.masked_less_equal(sb, 0), origin="lower", extent=extent,
                   norm=LogNorm(vmin=positive.max() * 1e-4, vmax=positive.max()),
                   cmap="inferno", interpolation="nearest")
    ax.set_facecolor(plt.get_cmap("inferno")(0))
    for radius, color in zip(radii_kpc[:-1], colors):
        ax.add_patch(Circle((0, 0), radius, fill=False, color=color, linewidth=1.6))
    tp.imshow_panel(ax)
    tp.scalebar(ax, 20, "20 kpc")
    tp.colorbar_minimal(fig, im, ax,
                        label=r"Relative surface brightness [kpc$^{-2}$]")
    tp.save_doc_figure(fig, "lya-tng-sbmap" + suffix, output_dir=".")
    tp.save_doc_figure(fig, "lya-tng-sbmap" + suffix, fmt="png", output_dir=".", dpi=160)
    plt.close(fig)

    edges = np.linspace(spectrum_attrs["left"], spectrum_attrs["right"],
                        spectrum_attrs["nbins"] + 1) / 1215.67 * 299792.458
    flux = spectra / (reference_flux * np.diff(edges))
    fig, ax = plt.subplots(figsize=(tp.DOC_FULL_WIDTH, 4.8), layout="constrained")
    for values, color, label in zip(flux, colors, labels):
        ax.stairs(values, edges, color=color, label=label, linewidth=2.2)
    ax.set_xlabel(r"Velocity offset [km s$^{-1}$]")
    ax.set_ylabel(r"Relative flux [(km s$^{-1}$)$^{-1}$]")
    ax.margins(x=0.02)
    ax.grid(axis="y", alpha=0.12)
    tp.strip_axes(ax, left=True)
    tp.legend_top(ax, bbox_to_anchor=(0.5, 1.14))
    tp.save_doc_figure(fig, "lya-tng" + suffix, output_dir=".")
    tp.save_doc_figure(fig, "lya-tng" + suffix, fmt="png", output_dir=".", dpi=160)
    plt.close(fig)