"""Compare metal absorption along one sightline. [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):

    sightline = zarr.open("absorption-metals.zr", mode="r")["sightlines/default"]
    wavelength = sightline["lam"][:]  # Angstrom
    windows = sightline.attrs["per_species"]
    # Rest wavelengths match THOR's curated CIV, SiIV, and OVI transitions.
    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, ax = plt.subplots(figsize=(tp.DOC_FULL_WIDTH, 4.8), layout="constrained")
    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"]]
        tau = sightline[key][0, :]  # First (and only) ray
        velocity = (lam / rest - 1) * 299792.458  # km/s
        ax.plot(velocity, np.exp(-tau), label=label, color=color, linewidth=2)

    ax.set_xlabel(r"Velocity offset [km s$^{-1}$]")
    ax.set_ylabel("Normalized flux")
    ax.set_xlim(-250, 250)
    ax.set_ylim(-0.03, 1.06)
    tp.strip_axes(ax, left=True)
    ax.grid(axis="y", alpha=0.12)
    tp.legend_top(ax, bbox_to_anchor=(0.5, 1.14))
    stem = "absorption-metals" + ("_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)