"""Connect the metal absorption spectrum to gas along its sightline. [AI-Codex]"""

import argparse

import matplotlib.pyplot as plt
import numpy as np
import thorplot as tp
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"]
    first, last = sightline["cell_offsets"][:2]
    selection = slice(int(first), int(last))
    length_cm = sightline.attrs["length_unit_cm"]
    dl_cm = sightline["dl"][selection] * length_cm
    rho = sightline["density"][selection]  # Physical gas mass density, g/cm^3.
    gas_velocity = np.column_stack([
        sightline["velocity_" + component][selection] for component in "xyz"
    ])  # cm/s
    start, end = sightline["start"][0], sightline["end"][0]
    direction = (end - start) / np.linalg.norm(end - start)

    # Match the engine's wavelength Doppler shift with Hubble flow disabled.
    # Positive values mean redshift: gas moving toward the observer is negative.
    c_cms = 2.99792458e10
    beta_parallel = gas_velocity @ direction / c_cms
    beta_squared = np.sum((gas_velocity / c_cms) ** 2, axis=1)
    redshift = (1 - beta_parallel) / np.sqrt(1 - beta_squared) - 1
    velocity = redshift * c_cms / 1e5  # Wavelength-equivalent velocity, km/s.
    distance_edges = np.r_[0, np.cumsum(dl_cm)] / 3.0856775814913673e21  # kpc

    # Each velocity bin sums rho * dl, giving gas mass column per km/s.
    velocity_edges = np.linspace(-250, 250, 41)
    gas_column, _ = np.histogram(velocity, bins=velocity_edges, weights=rho * dl_cm)
    column_per_velocity = gas_column / np.diff(velocity_edges)

    fig, (profile, spectrum) = plt.subplots(
        2, 1, figsize=(tp.DOC_FULL_WIDTH, 8.4), layout="constrained"
    )
    gray = "#8a90b8" if args.dark else "#7e84a3"
    profile.stairs(rho, distance_edges, baseline=None, color=gray, label="Gas density")
    profile.set_yscale("log")
    profile.set_ylabel(r"Gas density [g cm$^{-3}$]")
    profile.set_xlabel("Distance from ray start [kpc]")
    profile.set_xlim(distance_edges[0], distance_edges[-1])
    profile.grid(axis="y", alpha=0.12)
    tp.strip_axes(profile, left=True)
    velocity_axis = profile.twinx()
    velocity_axis.stairs(velocity, distance_edges, baseline=None,
                         color=tp.code_color("thor", args.dark), label="Gas velocity")
    velocity_axis.set_ylabel(r"Velocity offset [km s$^{-1}$]")
    velocity_axis.spines["right"].set_visible(True)
    handles, labels = profile.get_legend_handles_labels()
    other_handles, other_labels = velocity_axis.get_legend_handles_labels()
    profile.legend(handles + other_handles, labels + other_labels, ncol=2,
                   loc="upper center", bbox_to_anchor=(0.5, 1.18))

    column_axis = spectrum.twinx()
    column_axis.stairs(column_per_velocity, velocity_edges, fill=True,
                       color=gray, alpha=0.18, label="Gas column / velocity bin")
    column_axis.set_ylabel(r"Gas column [g cm$^{-2}$ / (km s$^{-1}$)]")
    column_axis.ticklabel_format(axis="y", style="sci", scilimits=(0, 0), useMathText=True)
    column_axis.spines["right"].set_visible(True)
    column_axis.set_ylim(0, max(column_per_velocity.max() * 1.1, np.finfo(float).tiny))
    # Draw absorption lines above the translucent gas-column layer.
    spectrum.set_zorder(column_axis.get_zorder() + 1)
    spectrum.patch.set_visible(False)
    spectrum.axhline(1, color=gray, linewidth=0.8, alpha=0.5, zorder=0)
    wavelength = sightline["lam"][:]
    windows = sightline.attrs["per_species"]
    for ion, label, rest, code in [
        ("CIV", "C IV 1548", 1548.187, "thor"),
        ("SiIV", "Si IV 1394", 1393.755, "colt"),
        ("OVI", "O VI 1032", 1031.912, "rascas"),
    ]:
        key = "tau_" + ion
        window = windows[key]
        offset = window["bin_offset"]
        lam = wavelength[offset:offset + window["nbins"]]
        spectral_velocity = (lam / rest - 1) * 299792.458
        spectrum.plot(spectral_velocity, np.exp(-sightline[key][0, :]),
                      color=tp.code_color(code, args.dark), label=label)
    spectrum.set_xlim(-250, 250)
    spectrum.set_ylim(-0.03, 1.06)
    spectrum.set_xlabel(r"Velocity offset [km s$^{-1}$]")
    spectrum.set_ylabel("Normalized flux")
    spectrum.grid(axis="y", alpha=0.12)
    tp.strip_axes(spectrum, left=True)
    handles, labels = spectrum.get_legend_handles_labels()
    other_handles, other_labels = column_axis.get_legend_handles_labels()
    spectrum.legend(handles + other_handles, labels + other_labels, ncol=4,
                    loc="upper center", bbox_to_anchor=(0.5, 1.18))

    stem = "absorption-gas" + ("_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)