"""Compare metal absorption at fixed impact parameters across views. [AI-Codex]"""

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import thorplot as tp
import zarr


def main():
    """Plot individual rays and ring medians using THOR's sampling metadata. [AI-Codex]"""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="THOR Zarr output from the rings or views example")
    parser.add_argument("--dark", action="store_true")
    args = parser.parse_args()
    with tp.style(dark=args.dark):

        sightlines = zarr.open(str(args.input), mode="r")["sightlines/default"]
        radius = sightlines["impact_parameter"][:] * sightlines.attrs["length_unit_cm"] / 3.0856775814913673e21
        view_index = sightlines["view_index"][:]
        ring_index = sightlines["ring_index"][:]
        wavelength = sightlines["lam"][:]
        bin_width = wavelength[1] - wavelength[0]  # Uniform bins, Angstrom.
        colors = [tp.PALETTE[name] for name in ["indigo", "teal", "orange"]]
        if args.dark:
            colors = [tp.CODE_COLORS_DARK[name] for name in ["thor", "colt", "rascas"]]

        fig, axes = plt.subplots(1, 3, figsize=(tp.DOC_FULL_WIDTH, 3.8), layout="constrained")
        for ax, ion, title in zip(axes, ["CIV", "SiIV", "OVI"], ["C IV 1548", "Si IV 1394", "O VI 1032"]):
            # Compact per-ion windows match the grid example's equivalent-width calculation.
            equivalent_width = np.sum(-np.expm1(-sightlines["tau_" + ion][:]), axis=1) * bin_width
            for index, view in enumerate(sightlines.attrs["views"]):
                selected = view_index == index
                color = colors[index % len(colors)]
                label = (view["target_id"] + " / " + view["orientation_id"]
                         if "target_id" in view else "+z")
                ax.scatter(radius[selected], equivalent_width[selected], s=14, alpha=0.45, color=color)
                rings = np.unique(ring_index[selected])
                bins = [selected & (ring_index == ring) for ring in rings]
                ax.plot([radius[mask][0] for mask in bins],
                        [np.median(equivalent_width[mask]) for mask in bins],
                        "o-", markersize=4, color=color, label=label)
            ax.set(title=title, xlabel="Impact parameter [kpc]", ylim=(0, None))
            tp.strip_axes(ax, left=True)
            ax.grid(axis="y", alpha=0.12)
        axes[0].set_ylabel("Equivalent width [Å]")
        handles, labels = axes[0].get_legend_handles_labels()
        if len(labels) > 1:
            fig.legend(handles, labels, loc="outside upper center", ncol=min(len(labels), 3), frameon=False, fontsize=9)
        stem = args.input.name.removesuffix(".zr") + ("_dark" if args.dark else "")
        tp.save_doc_figure(fig, stem, output_dir=".")
        tp.save_doc_figure(fig, stem, fmt="png", output_dir=".", dpi=160)
        plt.close(fig)


if __name__ == "__main__":
    main()