"""Map equivalent width and absorption-weighted velocity across metal sightlines. [AI-Codex]"""

import argparse

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


def absorption_moments(tau, wavelength, rest_wavelength, bin_width):
    """Return equivalent width and the flux-deficit-weighted velocity centroid. [AI-Codex]"""
    absorption = -np.expm1(-tau)  # 1 - exp(-tau), accurate for weak absorption.
    weight = np.sum(absorption, axis=1)
    velocity = (wavelength / rest_wavelength - 1) * 299792.458  # km/s
    centroid = np.full(weight.shape, np.nan)
    np.divide(np.sum(absorption * velocity, axis=1), weight,
              out=centroid, where=weight > 0)
    return weight * bin_width, centroid


def add_scalebar(ax, extent):
    """Mark 20 physical kpc with an outlined bar and label. [AI-Codex]"""
    bar_x = extent[0] + 0.07 * (extent[1] - extent[0])
    bar_y = extent[2] + 0.09 * (extent[3] - extent[2])
    bar, = ax.plot([bar_x, bar_x + 20], [bar_y, bar_y], color="white",
                   linewidth=3.5, solid_capstyle="projecting")
    bar.set_path_effects([
        patheffects.Stroke(linewidth=5, foreground="black", capstyle="projecting"),
        patheffects.Normal(),
    ])
    text = ax.annotate("20 kpc", (bar_x + 10, bar_y), xytext=(0, 6),
                       textcoords="offset points", ha="center", color="white")
    text.set_path_effects([patheffects.withStroke(linewidth=2, foreground="black")])


def main(dark=False):
    """Plot three ions with a shared, zero-centred velocity color scale. [AI-Codex]"""
    sightlines = zarr.open("absorption-metals-map.zr", mode="r")["sightlines/default"]
    start = sightlines["start"][:]
    unit_kpc = sightlines.attrs["length_unit_cm"] / 3.0856775814913673e21
    # Match the projection camera: image-right is -x, image-up is +y.
    x = (0.5 - start[:, 0]) * unit_kpc
    y = (start[:, 1] - 0.5) * unit_kpc
    xs, x_index = np.unique(x, return_inverse=True)
    ys, y_index = np.unique(y, return_inverse=True)
    assert len(xs) * len(ys) == len(start)
    assert len(np.unique(y_index * len(xs) + x_index)) == len(start)
    dx, dy = xs[1] - xs[0], ys[1] - ys[0]
    extent = [xs[0] - dx / 2, xs[-1] + dx / 2,
              ys[0] - dy / 2, ys[-1] + dy / 2]
    lam = sightlines["lam"][:]
    bin_width = lam[1] - lam[0]  # Angstrom; shared uniform wavelength bins.
    lines = [("CIV", "C IV 1548", 1548.187),
             ("SiIV", "Si IV 1394", 1393.755),
             ("OVI", "O VI 1032", 1031.912)]
    maps = []
    for ion, label, rest in lines:
        key = "tau_" + ion
        window = sightlines.attrs["per_species"][key]
        offset = window["bin_offset"]
        wavelength = lam[offset:offset + window["nbins"]]
        equivalent_width, centroid = absorption_moments(
            sightlines[key][:], wavelength, rest, bin_width)
        values = np.empty((2, len(ys), len(xs)))
        values[0, y_index, x_index] = equivalent_width
        values[1, y_index, x_index] = centroid
        maps.append(values)

    finite_velocities = np.concatenate([values[1].ravel() for values in maps])
    finite_velocities = finite_velocities[np.isfinite(finite_velocities)]
    max_velocity = np.max(np.abs(finite_velocities)) if finite_velocities.size else 0
    velocity_limit = max(10, np.ceil(max_velocity / 10) * 10)
    velocity_cmap = plt.get_cmap("RdBu_r").copy()
    velocity_cmap.set_bad("0.85")  # No absorption: the velocity centroid is undefined.

    fig, axes = plt.subplots(2, 3, figsize=(tp.DOC_FULL_WIDTH, 8.7), layout="constrained")
    for column, ((ion, label, rest), values) in enumerate(zip(lines, maps)):
        top, bottom = axes[:, column]
        top.set_title(label)
        im = top.imshow(values[0], origin="lower", extent=extent, cmap="viridis",
                        vmin=0, interpolation="nearest")
        fig.colorbar(im, ax=top, orientation="horizontal", fraction=0.06, pad=0.03,
                     label="Equivalent width [Å]")
        velocity_im = bottom.imshow(values[1], origin="lower", extent=extent,
                                    cmap=velocity_cmap, vmin=-velocity_limit,
                                    vmax=velocity_limit, interpolation="nearest")
        for ax in [top, bottom]:
            ax.set_axis_off()
            add_scalebar(ax, extent)
    fig.colorbar(velocity_im, ax=axes[1, :].tolist(), orientation="horizontal",
                 fraction=0.06, pad=0.03, shrink=0.7,
                 label=r"Absorption-weighted velocity offset [km s$^{-1}$]")
    stem = "absorption-metals-map" + ("_dark" if dark else "")
    tp.save_doc_figure(fig, stem, fmt="png", output_dir=".", dpi=200)
    plt.close(fig)


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