"""Compare a 3 by 3 grid of metal absorption sightlines. [AI-Codex]"""

import argparse

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

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

    sightlines = zarr.open("absorption-metals-grid.zr", mode="r")["sightlines/default"]
    wavelength = sightlines["lam"][:]
    windows = sightlines.attrs["per_species"]
    rows = {int(ray_id): row for row, ray_id in enumerate(sightlines["ray_id"][:])}
    lines = [
        ("CIV", "C IV 1548", 1548.187, tp.CODE_COLORS_DARK["thor"] if args.dark else tp.PALETTE["indigo"]),
        ("SiIV", "Si IV 1394", 1393.755, tp.CODE_COLORS_DARK["colt"] if args.dark else tp.PALETTE["teal"]),
        ("OVI", "O VI 1032", 1031.912, tp.CODE_COLORS_DARK["rascas"] if args.dark else tp.PALETTE["orange"]),
    ]

    fig, axes = plt.subplots(3, 3, figsize=(tp.DOC_FULL_WIDTH, 7.4), sharex=True, sharey=True,
                             layout="constrained")
    # IDs increase left to right, top to bottom, matching the locator map.
    for ray_id, ax in enumerate(axes.flat, start=1):
        ax.axhline(1, color=tp.code_color("analytic", dark=args.dark), linewidth=0.8, zorder=0)
        for ion, label, rest, color in lines:
            key = "tau_" + ion
            window = windows[key]
            offset = window["bin_offset"]
            lam = wavelength[offset:offset + window["nbins"]]
            velocity = (lam / rest - 1) * 299792.458
            tau = sightlines[key][rows[ray_id], :]
            ax.plot(velocity, np.exp(-tau), color=color, linewidth=1.6, label=label)
        ax.text(0.05, 0.08, str(ray_id), transform=ax.transAxes, fontweight="bold")
        tp.strip_axes(ax, left=True)
        ax.grid(axis="y", alpha=0.12)
        ax.set_xlim(-250, 250)
        ax.set_ylim(-0.03, 1.06)
        ax.set_xticks([-200, 0, 200])
        ax.set_yticks([0, 0.5, 1])
    fig.supxlabel(r"Velocity offset [km s$^{-1}$]")
    fig.supylabel("Normalized flux")
    fig.legend(*axes[0, 0].get_legend_handles_labels(), loc="outside upper center",
               ncol=3, frameon=False)
    stem = "absorption-metals-grid" + ("_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)