"""Compare the first Lyα simulation with the static-sphere solution. [AI-Codex]"""
import h5py
import numpy as np
import argparse

import thorplot as tp
import matplotlib.pyplot as plt

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

    with h5py.File("output/original/data.h5", "r") as f:
        dlambda = f["dlambda"][:]
        weight = f["weight"][:]

    # Parameters of my_first_sim.yaml; this comparison assumes a static sphere.
    temperature = 20000.0  # K
    tau0 = 1e6
    vthermal = 12.85 * np.sqrt(temperature / 10000.0)  # km/s
    lambda0 = 1215.668  # Å
    # x = (nu - nu0) / Doppler width; positive x is blueward of line center.
    x = -299792.458 / vthermal * dlambda / (lambda0 + dlambda)
    hist, bin_edges = np.histogram(x, bins=50, weights=weight)
    density = hist / (weight.sum() * np.diff(bin_edges))

    # Dijkstra et al. (2006) static-sphere solution, normalized to unit area.
    grid = np.linspace(-18, 18, 1200)
    a = 4.7e-4 * np.sqrt(10000.0 / temperature)
    profile = (np.sqrt(np.pi / 24) / (a * tau0) * grid**2
               / (1 + np.cosh(np.sqrt(2 * np.pi**3 / 27)
                              * np.abs(grid)**3 / (a * tau0))))
    profile /= np.trapezoid(profile, grid)

    fig, ax = plt.subplots(figsize=(tp.DOC_FULL_WIDTH, 4.8), layout="constrained")
    ax.stairs(density, bin_edges, color=tp.code_color("thor", dark=args.dark),
              label=f"THOR ({len(weight):,} photons)", linewidth=2.4)
    ax.plot(grid, profile, color=tp.code_color("analytic", dark=args.dark),
            linestyle="--", label="Static sphere (analytic)", linewidth=2.2)
    ax.set_xlabel("Frequency offset $x$")
    ax.set_ylabel("Normalized flux")
    ax.set_xlim(-16, 16)
    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 = "spectrum" + ("_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)