#!/usr/bin/env python3
"""Compare frozen CGMsum measurements with explicitly conditional stack models.

No spectra are fitted. Run with MPLCONFIGDIR and TMPDIR on the scratch disk.
Input copies are pinned by stack_comparison_sources/input_manifest.json.
"""
from __future__ import annotations

import hashlib
import json
import math
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.ticker import FuncFormatter
import numpy as np
import pandas as pd
from scipy.integrate import quad
from scipy.interpolate import PchipInterpolator
import scienceplots  # noqa: F401

ROOT = Path(__file__).resolve().parent
SOURCES = ROOT / "stack_comparison_sources"
PRODUCTS = ROOT / "products"
FIGURES = ROOT / "figures"
ALPHA = 2.4
TARGET_MASS = 1.3e11  # Zhang M31-bin median, not a new measurement of M31.
DISTANCE = 780.0
MAJOR_PA = 37.3
CENTER = (10.68470833, 41.26875)
APERTURE = DISTANCE * math.tan(math.radians(15 / 60))
SIDE_COLORS = {"North/NW": "#195b9b", "South/SE": "#b13b44"}
MODEL_COLORS = ["#919191", "#4184b4", "#184d92", "#873e9b", "#d99819", "#17836e"]


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def beta_shape(radius, rc, beta):
    return (1 + (np.asarray(radius) / rc) ** 2) ** (0.5 - 3 * beta)


def integral(rc, beta, low=10.0, high=200.0):
    p = 1.5 - 3 * beta
    a, b = 1 + (low / rc) ** 2, 1 + (high / rc) ** 2
    return math.pi * rc**2 * (math.log(b/a) if abs(p) < 1e-10 else (b**p-a**p)/p)


def project(ra, dec):
    """TAN projection; positive minor coordinate is NW at PA + 270 deg."""
    ra, dec = np.deg2rad(ra), np.deg2rad(dec)
    ra0, dec0 = np.deg2rad(CENTER)
    dra = ra-ra0
    den = np.sin(dec0)*np.sin(dec)+np.cos(dec0)*np.cos(dec)*np.cos(dra)
    east = np.cos(dec)*np.sin(dra)/den
    north = (np.cos(dec0)*np.sin(dec)-np.sin(dec0)*np.cos(dec)*np.cos(dra))/den
    pa = np.deg2rad(MAJOR_PA)
    major = DISTANCE*(east*np.sin(pa)+north*np.cos(pa))
    minor = DISTANCE*(-east*np.cos(pa)+north*np.sin(pa))
    return major, minor


def load_inputs():
    manifest = json.loads((SOURCES / "input_manifest.json").read_text())
    for item in manifest["files"]:
        if sha256(SOURCES / item["file"]) != item["sha256"]:
            raise ValueError(f"Input hash mismatch: {item['file']}")
    original = pd.read_csv(SOURCES / "v19_fields.csv", dtype={"obsid": str})
    if len(original) != 14 or not original["quality_primary"].all():
        raise ValueError("Expected the frozen 14-field primary sample")
    fields = original[["obsid", "side", "ra_deg", "dec_deg", "rproj_kpc", "signed_rproj_kpc",
                       "nh_hi4pi_1e22_cm-2"]].copy()
    fields["observed_cgmsum_fluxunit"] = original["cgmsum_absorbed_0p5_2p0_figure_fluxunit"]
    fields["local_staterr_fluxunit"] = original["cgmsum_absorbed_0p5_2p0_staterr_figure_fluxunit"]
    fields["major_axis_kpc"], fields["minor_axis_kpc"] = project(fields.ra_deg, fields.dec_deg)
    fields["tan_radius_780kpc"] = np.hypot(fields.major_axis_kpc, fields.minor_axis_kpc)
    fields["nominal_aperture_extent_kpc"] = APERTURE
    sign = np.where(fields.side == "North/NW", 1, -1)
    if not np.array_equal(np.sign(fields.minor_axis_kpc), sign):
        raise ValueError("Minor-axis sign does not match frozen side labels")
    author = np.loadtxt(SOURCES / "author_mw_posterior.txt")
    if author.shape != (299, 4) or not np.all((author[:, 2] <= author[:, 1]) & (author[:, 1] <= author[:, 3])):
        raise ValueError("Unexpected author posterior table")
    catalogue = pd.read_csv(SOURCES / "he_li_catalogue.csv")
    logm = np.sort(catalogue.logMs.to_numpy())
    if len(logm) != 474 or not (np.all(logm[:237] < 10.38) and np.all(logm[237:] > 10.38)):
        raise ValueError("He & Li mass split is not 237/237")
    masses = {"full": 10**logm, "high": 10**logm[237:]}
    audit = json.loads((SOURCES / "spectral_audit.json").read_text())
    spectral = audit["spectral_convention"]
    transmission = spectral["absorbed_0p5_2p0_over_intrinsic_0p5_2p0"]
    geometric = (np.pi/10800)**2/(4*np.pi*(3.0856775814913673e21)**2)
    conversion = geometric * transmission / 1e-15
    return manifest, fields, author, masses, spectral, conversion


def make_models(author, masses):
    models = [
        dict(id="zhang_v1_mw", label="Zhang v1 MW (historical)", mass=5.5e10, rc=30., beta=.42, s0=3.1e35,
             source="https://arxiv.org/html/2401.17308v1", kind="historical intrinsic beta model"),
        dict(id="zhang_v3_mw", label="Zhang v3 MW", mass=5.5e10, rc=6., beta=.43, s0=10**36.7,
             source="https://arxiv.org/html/2401.17308v3", kind="published intrinsic beta model"),
        dict(id="zhang_author_mw", label="Zhang author MW (Oct 8)", mass=5.5e10,
             source="Author-supplied posterior summary, 2026-10-08", kind="intrinsic pointwise posterior summary"),
        dict(id="zhang_v3_m31", label="Zhang v3 M31 bin", mass=1.3e11, rc=4., beta=.37, s0=10**37.1,
             source="https://arxiv.org/html/2401.17308v3", kind="published intrinsic beta model"),
    ]
    for key, label, rc, beta, logl in (("full", "He & Li v2 full", 8.7, .52, 39.7),
                                      ("high", "He & Li v2 high mass", 9.1, .55, 39.8)):
        m = masses[key]
        models.append(dict(id="he_li_"+key, label=label, mass=float(np.median(m)), rc=rc, beta=beta,
                           s0=10**logl/integral(rc, beta), logL10_200=logl,
                           mass_moment=float(np.mean(m**ALPHA)), n=len(m),
                           source="https://arxiv.org/html/2601.16499v2#S4.T1",
                           kind="approximate intrinsic beta reconstruction from rounded Table 1 medians"))
    for idx, model in enumerate(models):
        model["color"] = MODEL_COLORS[idx]
        model["median_scale"] = (TARGET_MASS/model["mass"])**ALPHA
        model["moment_scale"] = (TARGET_MASS**ALPHA/model["mass_moment"] if "mass_moment" in model else None)
    return models


def evaluate(model, r, author, percentile=1):
    r = np.asarray(r)
    if model["id"] == "zhang_author_mw":
        if np.any((r < author[0, 0]) | (r > author[-1, 0])):
            raise ValueError("Author interpolation outside supplied radius grid")
        return 10**PchipInterpolator(np.log10(author[:, 0]), np.log10(author[:, percentile]))(np.log10(r))
    return model["s0"]*beta_shape(r, model["rc"], model["beta"])


def observed(ax, fields, xcol="signed_rproj_kpc"):
    for side, color in SIDE_COLORS.items():
        subset = fields[fields.side == side]
        ax.errorbar(subset[xcol], subset.observed_cgmsum_fluxunit,
                    yerr=subset.local_staterr_fluxunit, xerr=APERTURE,
                    fmt="o" if side == "North/NW" else "s", ms=3.5,
                    color=color, ecolor=color, elinewidth=.65, capsize=1.3,
                    label="CGMsum "+side, zorder=10)


def setup_axis(ax, title, ylim=(.025, 12), minor=False):
    ax.set_yscale("log")
    ax.set_ylim(*ylim)
    ax.set_xlim(-40, 40)
    ax.set_xticks([-40, -20, 0, 20, 40])
    ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g}"))
    ax.axvline(0, color=".8", lw=.7)
    ax.set_title(title, fontsize=9, loc="left")
    ax.set_ylabel("Absorbed 0.5-2 keV SB\n"+r"[$10^{-15}$ erg cm$^{-2}$ s$^{-1}$ arcmin$^{-2}$]")
    ax.set_xlabel("Projected minor-axis coordinate (kpc; NW > 0)" if minor else "Signed circular radius (kpc; NW > 0)")
    ax.grid(axis="y", alpha=.13)


def model_curve(ax, model, author, conversion, scale=1.0, label=None, style=None):
    style = style or (":" if model["id"] == "zhang_v1_mw" else "--" if model["id"] == "zhang_v3_mw" else "-")
    is_author = model["id"] == "zhang_author_mw"
    # The author table contains no values below 5 kpc; do not invent a continuation.
    segments = [(np.linspace(5, 40, 351), style)] if is_author else [
        (np.linspace(0, 10, 101), ":"),
        (np.linspace(10, 40, 301), style),
    ]
    for index, (r, segment_style) in enumerate(segments):
        y = evaluate(model, r, author)*conversion*scale
        for sign in (-1, 1):
            ax.plot(sign*r, y, color=model["color"], ls=segment_style, lw=1.35,
                    label=(label or model["label"]) if sign == 1 and index == len(segments)-1 else None)
        if is_author:
            for sign in (-1, 1):
                ax.fill_between(sign*r, evaluate(model, r, author, 2)*conversion*scale,
                                evaluate(model, r, author, 3)*conversion*scale,
                                color=model["color"], alpha=.12, lw=0)


def save(fig, stem, column):
    fig.tight_layout(pad=1.15, h_pad=2.0)
    for ext in ("png", "pdf"):
        fig.savefig(FIGURES / f"{stem}_{column}.{ext}", dpi=400,
                    metadata={"Creator": "M31 CGM conditional profile comparison"})
    plt.close(fig)


def make_figures(fields, author, models, conversion):
    plt.style.use(["science", "no-latex"])
    for column, width, height in (("1col", 3.35, 8.4), ("2col", 7., 9.0)):
        plt.rcParams.update({"font.size": 8 if column == "1col" else 10,
                             "axes.labelsize": 8 if column == "1col" else 10,
                             "legend.fontsize": 6.5 if column == "1col" else 8})
        fig, axes = plt.subplots(2, 1, figsize=(width, height))
        for ax, subset, title in zip(axes, (models[:3], models[3:]),
                                   ("(a) Zhang MW versions: native mass", "(b) M31 bin and He & Li: native masses")):
            setup_axis(ax, title)
            for model in subset:
                model_curve(ax, model, author, conversion)
            observed(ax, fields)
            ax.legend(loc="upper left", frameon=True, framealpha=.88)
            ax.text(.98, .02, "Analytic models: dotted below 10 kpc\nAuthor table starts at 5 kpc",
                    transform=ax.transAxes, ha="right", va="bottom", fontsize=6,
                    bbox=dict(facecolor="white", edgecolor="none", alpha=.8))
        save(fig, "stack_comparison_native", column)

        fig, axes = plt.subplots(2, 1, figsize=(width, height))
        for ax in axes:
            setup_axis(ax, "", ylim=(.04, 300))
        axes[0].set_title(r"(a) Conditional median scaling to $1.3\times10^{11}\ M_\odot$", loc="left", fontsize=8)
        for model in models[1:]:
            model_curve(axes[0], model, author, conversion, model["median_scale"])
        observed(axes[0], fields)
        axes[0].legend(loc="lower center", frameon=True, framealpha=.88, ncol=1 if column == "1col" else 2)
        axes[1].set_title("(b) He & Li weighting sensitivity; same target mass", loc="left", fontsize=8)
        model_curve(axes[1], models[3], author, conversion)
        for model in models[4:]:
            for key, ls, suffix in (("median_scale", "--", "median"), ("moment_scale", "-", "equal-number moment")):
                model_curve(axes[1], model, author, conversion, model[key],
                            label=model["label"]+": "+suffix, style=ls)
        observed(axes[1], fields)
        axes[1].legend(loc="lower center", frameon=True, framealpha=.88, ncol=1 if column == "1col" else 2)
        for ax in axes:
            ax.text(.98, .98, "R < 10 kpc: analytic continuation (dotted)",
                    transform=ax.transAxes, ha="right", va="top", fontsize=6)
        save(fig, "stack_comparison_mass", column)

        fig, axes = plt.subplots(2, 1, figsize=(width, height))
        for ax, subset, title in zip(axes, (models[:3], models[3:]),
                                   ("(a) Zhang MW versions: actual field locations", "(b) M31 bin and He & Li: actual field locations")):
            setup_axis(ax, title, ylim=(.025, 5.5), minor=True)
            # Every model is evaluated at the field's circular R, not at |minor|.
            for model in subset:
                for side in SIDE_COLORS:
                    sub = fields[fields.side == side].sort_values("minor_axis_kpc")
                    y = evaluate(model, sub.rproj_kpc.to_numpy(), author)*conversion
                    ax.plot(sub.minor_axis_kpc, y, marker="+", ls="none", ms=5,
                            color=model["color"], label=model["label"] if side == "North/NW" else None)
            observed(ax, fields, "minor_axis_kpc")
            ax.legend(loc="upper left", frameon=True, framealpha=.88)
        save(fig, "stack_comparison_vertical", column)

        fig, ax = plt.subplots(figsize=(width, width*1.05))
        for radius in (10, 20, 30):
            ax.add_patch(Circle((0, 0), radius, fill=False, color=".75", ls=":", lw=.7))
        for side, color in SIDE_COLORS.items():
            sub = fields[fields.side == side]
            ax.scatter(sub.major_axis_kpc, sub.minor_axis_kpc, s=17, color=color, label=side)
            for row in sub.itertuples():
                ax.add_patch(Circle((row.major_axis_kpc, row.minor_axis_kpc), APERTURE,
                                    fill=False, color=color, alpha=.4, lw=.7))
        ax.axvline(0, color=".55", lw=.7); ax.axhline(0, color=".55", lw=.7)
        ax.set(xlim=(-37,37), ylim=(-37,37), aspect="equal",
               xlabel="Projected major-axis coordinate (kpc)",
               ylabel="Projected minor-axis coordinate (kpc; NW > 0)")
        ax.set_title("14 CGMsum fields; PA = 37.3 deg, D = 780 kpc", fontsize=8)
        ax.legend(loc="lower left", frameon=True)
        save(fig, "stack_comparison_geometry", column)


def main():
    PRODUCTS.mkdir(exist_ok=True); FIGURES.mkdir(exist_ok=True)
    manifest, fields, author, masses, spectral, conversion = load_inputs()
    models = make_models(author, masses)
    rows, predictions = [], []
    radius_grid = np.unique(np.r_[np.linspace(0, 100, 1001), fields.rproj_kpc.to_numpy()])
    for model in models:
        r = radius_grid[radius_grid >= 5] if model["id"] == "zhang_author_mw" else radius_grid
        for scenario, scale in (("native", 1.), ("median_mass_1p3e11", model["median_scale"]),
                                ("equal_number_moment_1p3e11", model["moment_scale"])):
            if scale is None:
                continue
            median = evaluate(model, r, author)
            for radius, value in zip(r, median):
                row = dict(model_id=model["id"], scenario=scenario, radius_kpc=radius,
                           scale=scale, intrinsic_sb_erg_s_kpc2=value*scale,
                           absorbed_fluxunit=value*scale*conversion,
                           analytic_inner_continuation=bool(model["id"] != "zhang_author_mw" and radius < 10),
                           below_main_he_fit=bool(model["id"].startswith("he_") and radius < 10),
                           beyond_main_he_fit=bool(model["id"].startswith("he_") and radius > 50))
                if model["id"] == "zhang_author_mw":
                    row.update(p16_absorbed_fluxunit=float(evaluate(model, radius, author, 2)*scale*conversion),
                               p84_absorbed_fluxunit=float(evaluate(model, radius, author, 3)*scale*conversion))
                rows.append(row)
            for field in fields.itertuples():
                predictions.append(dict(obsid=field.obsid, model_id=model["id"], scenario=scenario,
                                        radius_kpc=field.rproj_kpc, minor_axis_kpc=field.minor_axis_kpc,
                                        prediction_fluxunit=float(evaluate(model, field.rproj_kpc, author)*scale*conversion)))
    curves = pd.DataFrame(rows)
    fields.to_csv(PRODUCTS / "stack_comparison_fields.csv", index=False)
    pd.DataFrame(models).to_csv(PRODUCTS / "stack_comparison_models.csv", index=False)
    curves.to_csv(PRODUCTS / "stack_comparison_curves.csv", index=False)
    pd.DataFrame(predictions).to_csv(PRODUCTS / "stack_comparison_field_predictions.csv", index=False)
    checks = {}
    for model in models[4:]:
        numeric = quad(lambda x: 2*np.pi*x*float(evaluate(model, x, author)), 10, 200)[0]
        checks[model["id"]+"_luminosity_relative_error"] = abs(numeric / 10**model["logL10_200"] - 1)
    zhang20 = float(evaluate(models[3], 20, author)*conversion)
    # Rounding in the published parameters produces a known 4.6 ppm offset.
    checks["zhang_m31_20kpc_relative_to_frozen_ledger"] = zhang20/0.8284536873547118-1
    checks["geometry_hypot_max_error_kpc"] = float(np.max(abs(np.hypot(fields.major_axis_kpc,fields.minor_axis_kpc)-fields.tan_radius_780kpc)))
    checks["frozen_radius_vs_tan_780_max_fractional_difference"] = float(np.max(abs(fields.rproj_kpc/fields.tan_radius_780kpc-1)))
    assert max(checks[m["id"]+"_luminosity_relative_error"] for m in models[4:]) < 1e-8
    assert abs(checks["zhang_m31_20kpc_relative_to_frozen_ledger"]) < 1e-5
    assert checks["frozen_radius_vs_tan_780_max_fractional_difference"] < .01
    summary = dict(schema_version=1, generated_date="2026-10-13", field_count=len(fields),
                   observed_quantity="CGMsum: MW + M31, with no component subtraction",
                   observed_errors="saved local-covariance statistical errors; no full systematic budget",
                   model_prediction="field-center or radial template; fixed NH; not masked-aperture/response likelihood",
                   spectral_transfer=spectral, intrinsic_to_absorbed_fluxunit_factor=conversion,
                   unit="1e-15 erg cm-2 s-1 arcmin-2, absorbed 0.5-2.0 keV",
                   target_mass_Msun=TARGET_MASS, target_definition="Zhang M31 stellar-mass-bin median; illustrative target",
                   mass_exponent=ALPHA, mass_scaling="fixed physical radius; no halo dilation; conditional luminosity-to-SB assumption",
                   moment_scaling="equal-number He catalogue moment, common shape; actual radial stacking weights unknown",
                   author_interval="pointwise 16th-84th posterior model percentiles, not galaxy scatter or full covariance",
                   plotted_radius_kpc=[0, 40],
                   analytic_inner_display="0 <= R < 10 kpc shown dotted as model continuation; He fitting mainly 10-50 kpc",
                   author_display="5-40 kpc only; no interpolation or extrapolation below the first supplied radius",
                   author_sha256=sha256(SOURCES/"author_mw_posterior.txt"),
                   geometry=dict(distance_kpc=DISTANCE, major_axis_pa_deg=MAJOR_PA,
                                 minor_axis_positive="NW", disk_deprojection=False,
                                 original_signed_radius="unchanged saved circular projected radius with NW/SE sign",
                                 aperture_extent_kpc=APERTURE),
                   models=models, input_hashes=manifest["files"], checks=checks,
                   generator_sha256=sha256(__file__))
    (PRODUCTS / "stack_comparison_summary.json").write_text(json.dumps(summary, indent=2)+"\n")
    make_figures(fields, author, models, conversion)
    print(json.dumps({"fields":len(fields),"models":len(models),"curves":len(curves),"checks":checks},indent=2))


if __name__ == "__main__":
    main()
