"""Compare the demo galaxy's Lyα emission components and dust runs. [AI-Codex]"""

import argparse
from pathlib import Path

import h5py
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import numpy as np

import thorplot as tp


COMPONENTS = ("recombination", "collisional", "stellar")
LABELS = ("Gas recombination", "Collisional excitation", "Stellar emission")
KPC_CM = 3.0856775814913673e21
C_KMS = 299792.458


def packet_budget(handle):
    """Return luminosity (erg/s) and summed photon-rate weights. [AI-Codex]"""
    if not handle.attrs.get("completed", False):
        raise ValueError(f"Incomplete run: {handle.filename}")
    if handle.attrs.get("mcrt_packet_weight") != "photon_rate":
        raise ValueError("These examples require current photon-rate output metadata")
    if "weight" not in handle:
        return 0.0, 0.0  # A completed stream can contain no escaping packets.
    weights = handle["weight"][:]
    wavelength = float(handle.attrs["line_rest_wavelength_angstrom"])
    packet_wavelength = wavelength + handle["dlambda"][:]
    if not np.all(np.isfinite(packet_wavelength)) or np.any(packet_wavelength <= 0):
        raise ValueError(f"Invalid packet wavelengths: {handle.filename}")
    energy_factor = wavelength / packet_wavelength
    if not np.all(np.isfinite(weights)) or np.any(weights < 0):
        raise ValueError(f"Invalid packet weights: {handle.filename}")
    luminosity = np.sum(weights * energy_factor, dtype=np.float64) * handle.attrs["mcrt_weight_unit_erg_s"]
    return float(luminosity), float(np.sum(weights, dtype=np.float64))


def read_run(directory):
    """Read physical maps, spectra, and packet luminosity budgets. [AI-Codex]"""
    result = {}
    for stream in ("input", "peel"):
        with h5py.File(directory / stream / "data.h5", "r") as handle:
            if not handle.attrs.get("completed", False):
                raise ValueError(f"Incomplete run: {handle.filename}")
            image = handle["sbmap/image"]
            spectrum = handle["aperturespectrum/dlambda"]
            if image.attrs["units"] != "erg/s/cm2/arcsec2":
                raise ValueError(f"Unexpected map units in {directory}/{stream}")
            if spectrum.attrs["units"] != "erg/s":
                raise ValueError(f"Unexpected spectrum units in {directory}/{stream}")
            # A single input observer can retain a leading length-one LOS axis.
            result[stream] = {
                "map": np.squeeze(image[:]).T,
                "spectrum": np.squeeze(spectrum[:]),
                "map_attrs": dict(image.attrs),
                "spec_attrs": dict(spectrum.attrs),
            }
            if stream == "input":
                result["emitted"], result["emitted_weight"] = packet_budget(handle)
                result["wavelength"] = float(handle.attrs["line_rest_wavelength_angstrom"])
        for name in ("map", "spectrum"):
            values = result[stream][name]
            if not np.all(np.isfinite(values)) or np.any(values < 0):
                raise ValueError(f"Invalid {name} values in {directory}/{stream}")
    with h5py.File(directory / "original/data.h5", "r") as handle:
        result["escaped"], result["escaped_weight"] = packet_budget(handle)
    if not np.isfinite(result["emitted"]) or result["emitted"] <= 0:
        raise ValueError(f"No emitted luminosity in {directory}")
    if not np.isfinite(result["escaped"]) or result["escaped"] < 0:
        raise ValueError(f"Invalid escaped luminosity in {directory}")
    return result


def check_geometry(runs):
    """Reject comparisons using different maps or spectral apertures. [AI-Codex]"""
    reference = runs[0]["peel"]
    for run in runs:
        if not np.isclose(run["wavelength"], runs[0]["wavelength"]):
            raise ValueError("Runs use different line wavelengths")
        for stream in ("input", "peel"):
            data = run[stream]
            for name in ("map", "spectrum"):
                if data[name].shape != reference[name].shape:
                    raise ValueError(f"Runs use different {name} shapes")
            for key in ("left_x", "right_x", "left_y", "right_y", "length_unit_cm",
                        "observer_redshift", "center", "north"):
                if not np.allclose(data["map_attrs"][key], reference["map_attrs"][key]):
                    raise ValueError(f"Inconsistent map {key}")
            for key in ("left", "right", "nbins", "aperture_radii", "center", "north"):
                if not np.allclose(data["spec_attrs"][key], reference["spec_attrs"][key]):
                    raise ValueError(f"Inconsistent spectrum {key}")


def report(runs, title):
    """Print absolute source luminosities and all-direction escape. [AI-Codex]"""
    print(f"\n{title}")
    print(f"{'Component':24s} {'Emitted [erg/s]':>16s} {'Escaped [erg/s]':>16s} {'Photon escape':>14s}")
    for label, run in zip(LABELS, runs):
        escape = run["escaped_weight"] / run["emitted_weight"]
        print(f"{label:24s} {run['emitted']:16.5e} {run['escaped']:16.5e} {escape:14.4f}")
    emitted = sum(run["emitted"] for run in runs)
    escaped = sum(run["escaped"] for run in runs)
    escape = sum(run["escaped_weight"] for run in runs) / sum(run["emitted_weight"] for run in runs)
    print(f"{'Total':24s} {emitted:16.5e} {escaped:16.5e} {escape:14.4f}")


def save(fig, stem, output_dir):
    """Save transparent documentation SVG and an opaque PNG preview. [AI-Codex]"""
    tp.save_doc_figure(fig, stem, output_dir=output_dir)
    tp.save_doc_figure(fig, stem, fmt="png", output_dir=output_dir, dpi=160,
                       transparent=False)
    plt.close(fig)


def plot_maps(baseline, dusty, args):
    """Compare intrinsic/escaping or dust-free/dusty component maps. [AI-Codex]"""
    rows = [(baseline, "peel"), (dusty, "peel")] if dusty else [(baseline, "input"), (baseline, "peel")]
    row_names = ("Without dust", "With dust") if dusty else ("Intrinsic", "Escaping toward +z")
    attrs = baseline[0]["peel"]["map_attrs"]
    box_kpc = attrs["length_unit_cm"] / KPC_CM
    extent = [(attrs[key] - 0.5) * box_kpc for key in ("left_x", "right_x", "left_y", "right_y")]
    maximum = max(run[stream]["map"].max() for runs, stream in rows for run in runs)
    if maximum <= 0:
        raise ValueError("All maps are empty")
    norm = LogNorm(vmin=maximum * 1e-5, vmax=maximum)
    cmap = plt.get_cmap("inferno").with_extremes(bad=plt.get_cmap("inferno")(0))
    fig, axes = plt.subplots(2, 3, figsize=(tp.DOC_FULL_WIDTH, 6.4), layout="constrained")
    for row, (runs, stream) in enumerate(rows):
        for column, (label, run) in enumerate(zip(LABELS, runs)):
            ax = axes[row, column]
            values = run[stream]["map"]
            im = ax.imshow(np.ma.masked_less_equal(values, 0), origin="lower", extent=extent,
                           cmap=cmap, norm=norm, interpolation="nearest")
            ax.set_facecolor(plt.get_cmap("inferno")(0))
            tp.imshow_panel(ax)
            if row == 0:
                ax.set_title(label, fontsize=12)
            if column == 0:
                ax.set_ylabel(row_names[row], fontsize=11)
            if row == 1 and column == 0:
                tp.scalebar(ax, 20, "20 kpc")
    tp.colorbar_minimal(fig, im, axes.ravel().tolist(), orientation="horizontal",
                        label=r"Surface brightness [erg s$^{-1}$ cm$^{-2}$ arcsec$^{-2}$]",
                        fraction=0.06, pad=0.03)
    stem = "galaxy-lya-dust-maps" if dusty else "galaxy-lya-components-maps"
    save(fig, stem + ("_dark" if args.dark else ""), args.output_dir)


def plot_spectra(baseline, dusty, args):
    """Plot absolute aperture luminosities per velocity interval. [AI-Codex]"""
    attrs = baseline[0]["peel"]["spec_attrs"]
    edges = np.linspace(attrs["left"], attrs["right"], int(attrs["nbins"]) + 1)
    edges *= C_KMS / baseline[0]["wavelength"]
    widths = np.diff(edges)
    radii = attrs["aperture_radii"] * baseline[0]["peel"]["map_attrs"]["length_unit_cm"] / KPC_CM
    indices = []
    for radius in (10, 30):
        matches = np.flatnonzero(np.isclose(radii, radius))
        if len(matches) != 1:
            raise ValueError(f"Expected one {radius} kpc spectral aperture")
        indices.append(int(matches[0]))
    radii = radii[indices]
    curves = [run["peel"]["spectrum"][indices] / widths for run in baseline]
    total = np.sum(curves, axis=0)
    dust_total = np.sum([run["peel"]["spectrum"][indices] for run in dusty], axis=0) / widths if dusty else None
    maximum = max(float(total.max()), float(dust_total.max()) if dusty else 0.0)
    colors = [tp.code_color(code, args.dark) for code in ("thor", "colt", "rascas")]
    fig, axes = plt.subplots(1, len(radii), figsize=(tp.DOC_FULL_WIDTH, 4.2),
                             sharex=True, sharey=True, layout="constrained")
    for index, ax in enumerate(axes):
        if dusty:
            ax.stairs(total[index] / 1e39, edges, color=colors[0], label="Without dust", linewidth=1.8)
            ax.stairs(dust_total[index] / 1e39, edges, color=colors[2], label="With dust", linewidth=1.8)
        else:
            for label, curve, color in zip(LABELS, curves, colors):
                ax.stairs(curve[index] / 1e39, edges, color=color, label=label, linewidth=1.6)
            ax.stairs(total[index] / 1e39, edges, color="white" if args.dark else "#222222",
                      label="Sum", linewidth=1.8, linestyle="--")
        ax.set_title(f"Within {radii[index]:g} kpc", fontsize=12)
        ax.set_xlabel(r"Velocity [km s$^{-1}$]")
        ax.set_xlim(-1000, 1000)
        ax.grid(axis="y", alpha=0.12)
        tp.strip_axes(ax, left=True)
    axes[0].set_ylim(0, 1.08 * maximum / 1e39)
    axes[0].set_ylabel(r"$L_v$ [$10^{39}$ erg s$^{-1}$ (km s$^{-1}$)$^{-1}$]")
    handles, labels = axes[0].get_legend_handles_labels()
    fig.legend(handles, labels, loc="outside upper center", ncol=2, frameon=False, fontsize=11)
    stem = "galaxy-lya-dust-spectra" if dusty else "galaxy-lya-components-spectra"
    save(fig, stem + ("_dark" if args.dark else ""), args.output_dir)


def main():
    """Read the documented runs, report budgets, and render comparisons. [AI-Codex]"""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dust", action="store_true", help="Compare matching dust-free and dusty runs")
    parser.add_argument("--dark", action="store_true")
    parser.add_argument("--data-dir", type=Path, default=Path("."))
    parser.add_argument("--output-dir", type=Path, default=Path("."))
    args = parser.parse_args()
    baseline = [read_run(args.data_dir / f"output-lya-{channel}") for channel in COMPONENTS]
    dusty = [read_run(args.data_dir / f"output-lya-{channel}-dust") for channel in COMPONENTS] if args.dust else None
    check_geometry(baseline + (dusty or []))
    report(baseline, "Without dust")
    if dusty:
        for channel, clean, attenuated in zip(COMPONENTS, baseline, dusty):
            if not np.isclose(clean["emitted"], attenuated["emitted"], rtol=1e-6, atol=0):
                raise ValueError(f"Source luminosity changed between dust runs: {channel}")
        report(dusty, "With dust")
    args.output_dir.mkdir(parents=True, exist_ok=True)
    with tp.style(dark=args.dark):
        plot_maps(baseline, dusty, args)
        plot_spectra(baseline, dusty, args)


if __name__ == "__main__":
    main()