"""Compare a shell-model exercise with the original static run. [AI-Codex]"""

import argparse

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

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("variant", choices=["denser", "outflow", "smoother"])
parser.add_argument("--dark", action="store_true")
args = parser.parse_args()
with tp.style(dark=args.dark):
    labels = {
        "denser": "10× density",
        "outflow": "200 km/s outflow",
        "smoother": "Static, 10,000 photons",
    }


    def read_spectrum(path):
        """Read escaped photon wavelengths and weights. [AI-Codex]"""
        with h5py.File(path, "r") as f:
            return f["dlambda"][:], f["weight"][:]


    baseline, baseline_weights = read_spectrum("output/original/data.h5")
    variant, variant_weights = read_spectrum(f"output-{args.variant}/original/data.h5")
    bins = np.linspace(min(baseline.min(), variant.min()),
                       max(baseline.max(), variant.max()), 51)

    fig, ax = plt.subplots(figsize=(tp.DOC_FULL_WIDTH, 4.8), layout="constrained")
    for values, weights, label, color, linestyle in [
        (baseline, baseline_weights, "Static, 1,000 photons", tp.code_color("analytic", dark=args.dark), "--"),
        (variant, variant_weights, labels[args.variant], tp.code_color("thor", dark=args.dark), "-"),
    ]:
        counts, _ = np.histogram(values, bins=bins, weights=weights)
        # Normalize independently so different photon counts remain comparable.
        density = counts / (weights.sum() * np.diff(bins))
        ax.stairs(density, bins, label=label, color=color, linestyle=linestyle,
                  linewidth=2.4)
    ax.axvline(0, color=tp.code_color("analytic", dark=args.dark), linewidth=0.8, alpha=0.4, zorder=0)
    ax.set_xlabel(r"$\Delta\lambda$ [$\AA$]")
    ax.set_ylabel(r"Normalized flux [$\AA^{-1}$]")
    tp.strip_axes(ax, left=True)
    ax.grid(axis="y", alpha=0.12)
    ax.margins(x=0.02)
    tp.legend_top(ax, bbox_to_anchor=(0.5, 1.14))
    stem = f"spectrum-{args.variant}" + ("_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)
    print(f"Saved spectrum-{args.variant}.png")