from __future__ import annotations

from contextlib import contextmanager
import hashlib
import json
from pathlib import Path
from typing import Iterator

import numpy as np
from numpy.polynomial.hermite import hermgauss
import pandas as pd
from sherpa.astro.xspec import XSapec, XSphabs, XSraymond, XSTBabs
from sherpa.astro.ui import (
    get_xsabund,
    get_xsxsect,
    set_xsabund,
    set_xsxsect,
)

HERE = Path(__file__).resolve().parent
HENLEY_SHELTON_INPUT_CSV = (
    HERE / "m31_cgmsum_henley_shelton2013_vizier_main_sample.csv"
)
M31_FOOTPRINT_INPUT_CSV = (
    HERE
    / "m31_cgmsum_locatelli2024_reference_beta0p5_m31_footprint_predictions.csv"
)
PRIMARY_MEASUREMENTS_INPUT_CSV = HERE / "m31_cgmsum_v19_primary_measurements_public.csv"
KELVIN_TO_KEV = 8.617333262e-8
KEV_TO_ERG = 1.602176634e-9
PARSEC_TO_CM = 3.0856775814913673e18
ARCMIN2_PER_STERADIAN = (180.0 * 60.0 / np.pi) ** 2


@contextmanager
def _xspec_settings(abundance: str, cross_section: str) -> Iterator[None]:
    original_abundance = get_xsabund()
    original_cross_section = get_xsxsect()
    set_xsabund(abundance)
    set_xsxsect(cross_section)
    try:
        yield
    finally:
        set_xsabund(original_abundance)
        set_xsxsect(original_cross_section)


def _efeds_disnht_band_fluxes(
    temperature_keV: float,
    abundance_solar: float,
) -> dict[str, float]:
    """Approximate the published disnht screen by log-NH quadrature."""
    edges = np.arange(0.3, 2.0 + 5.0e-5, 5.0e-5)
    midpoint = 0.5 * (edges[:-1] + edges[1:])
    plasma = XSapec()
    plasma.kT = temperature_keV
    plasma.Abundanc = abundance_solar
    plasma.Redshift = 0.0
    plasma.norm = 1.0
    photons = plasma(edges[:-1], edges[1:])

    nodes, raw_weights = hermgauss(40)
    log_nh = 20.51 + np.sqrt(2.0) * 0.117 * nodes
    weights = raw_weights / np.sqrt(np.pi)
    absorbed = np.zeros_like(photons)
    for weight, value in zip(weights, log_nh, strict=True):
        absorption = XSTBabs()
        absorption.nH = 10.0**value / 1.0e22
        absorbed += weight * absorption(edges[:-1], edges[1:]) * photons
    output: dict[str, float] = {}
    for name, low, high in (
        ("0p3_0p6", 0.3, 0.6),
        ("0p5_0p6", 0.5, 0.6),
    ):
        selected = (midpoint >= low) & (midpoint < high)
        output[name] = float(np.sum(absorbed[selected] * midpoint[selected]))
    return output


def convert_efeds_scenarios() -> pd.DataFrame:
    """Convert Ponti et al. (2023) Table 4 CGM fluxes to 0.5--2 keV."""
    scenarios = (
        ("negligible_swcx", 0.157, 0.068, 24.1, 5.6),
        ("high_swcx", 0.173, 0.058, 15.6, 4.9),
    )
    rows: list[dict[str, float | str]] = []
    with _xspec_settings("lodd", "vern"):
        for identifier, temperature, abundance, flux_03_06, flux_06_2 in scenarios:
            model_fluxes = _efeds_disnht_band_fluxes(temperature, abundance)
            soft_ratio = (
                model_fluxes["0p5_0p6"] / model_fluxes["0p3_0p6"]
            )
            central = (flux_06_2 + soft_ratio * flux_03_06) / 36.0
            rows.append(
                {
                    "scenario_id": identifier,
                    "temperature_keV": temperature,
                    "abundance_solar": abundance,
                    "mean_log10_nh_cm-2": 20.51,
                    "sigma_log10_nh": 0.117,
                    "published_0p3_0p6_1e13_deg-2": flux_03_06,
                    "published_0p6_2p0_1e13_deg-2": flux_06_2,
                    "model_ratio_0p5_0p6_to_0p3_0p6": soft_ratio,
                    "flux_0p5_2p0_plot_central_fluxunit": central,
                    "flux_0p5_2p0_reconstruction_low_fluxunit": central,
                    "flux_0p5_2p0_reconstruction_high_fluxunit": central,
                }
            )
    return pd.DataFrame(rows)


def convert_henley_shelton(
    source: pd.DataFrame,
) -> tuple[pd.DataFrame, dict[str, float | int]]:
    """Convert the published intrinsic Raymond-Smith brightnesses field by field."""
    detections = source.loc[
        source["surface_brightness_limit"].fillna("").ne("<")
    ].copy()
    missing_nh = detections["nhi_1e20_cm-2"].isna()
    excluded = detections.loc[missing_nh, "seq"].astype(int).tolist()
    converted = detections.loc[~missing_nh].copy()

    edges = np.linspace(0.5, 2.0, 30001)
    midpoint = 0.5 * (edges[:-1] + edges[1:])
    transmission: list[float] = []
    with _xspec_settings("angr", "bcmc"):
        temperatures = converted["temperature_1e6_K"].to_numpy()
        columns = converted["nhi_1e20_cm-2"].to_numpy()
        for temperature_1e6_k, nhi_1e20 in zip(
            temperatures, columns, strict=True
        ):
            plasma = XSraymond()
            plasma.kT = temperature_1e6_k * 1.0e6 * KELVIN_TO_KEV
            plasma.Abundanc = 1.0
            plasma.Redshift = 0.0
            plasma.norm = 1.0
            photons = plasma(edges[:-1], edges[1:])
            absorption = XSphabs()
            absorption.nH = nhi_1e20 / 100.0
            absorbed = absorption(edges[:-1], edges[1:]) * photons
            transmission.append(
                float(np.sum(absorbed * midpoint) / np.sum(photons * midpoint))
            )

    converted["absorbed_to_intrinsic_0p5_2p0_ratio"] = transmission
    converted["absorbed_0p5_2p0_fluxunit"] = (
        converted["surface_brightness_intrinsic_0p5_2p0_1e12"]
        * converted["absorbed_to_intrinsic_0p5_2p0_ratio"]
        / 3.6
    )
    values = converted["absorbed_0p5_2p0_fluxunit"].to_numpy()
    quantiles = np.quantile(values, [0.0, 0.16, 0.25, 0.5, 0.75, 0.84, 1.0])
    summary: dict[str, float | int] = {
        "published_main_sightlines": int(len(source)),
        "published_detections": int(len(detections)),
        "converted_detections": int(len(converted)),
        "excluded_without_unique_nh": int(missing_nh.sum()),
        "excluded_seq": excluded[0],
        "minimum_fluxunit": float(quantiles[0]),
        "p16_fluxunit": float(quantiles[1]),
        "q25_fluxunit": float(quantiles[2]),
        "median_fluxunit": float(quantiles[3]),
        "q75_fluxunit": float(quantiles[4]),
        "p84_fluxunit": float(quantiles[5]),
        "maximum_fluxunit": float(quantiles[6]),
    }
    return converted, summary


def _ueda_disk_emission_measure(
    galactic_l_deg: float,
    galactic_b_deg: float,
    central_electron_density_cm3: float = 3.4e-3,
) -> float:
    """Project the Ueda et al. (2022) exponential electron-density disk."""
    distance_kpc = np.linspace(0.0, 260.0, 100001)
    longitude = np.deg2rad(galactic_l_deg)
    latitude = np.deg2rad(galactic_b_deg)
    projected = distance_kpc * np.cos(latitude)
    radius_kpc = np.sqrt(
        8.2**2
        + projected**2
        - 2.0 * 8.2 * projected * np.cos(longitude)
    )
    height_kpc = distance_kpc * np.sin(latitude)
    electron_density = (
        central_electron_density_cm3
        * np.exp(-radius_kpc / 7.0)
        * np.exp(-np.abs(height_kpc) / 2.7)
    )
    # Ueda et al. define EM = integral n_e n_H ds and call n0 an electron density.
    return float(np.trapezoid(electron_density**2 / 1.2, distance_kpc) * 1000.0)


def _ueda_absorbed_flux_per_em(nh_1e22_cm2: float) -> float:
    """Return absorbed 0.5--2.0 keV Figure-3 units per cm^-6 pc."""
    edges = np.linspace(0.5, 2.0, 30001)
    midpoint = 0.5 * (edges[:-1] + edges[1:])
    plasma = XSapec()
    plasma.kT = 0.22
    plasma.Abundanc = 1.0
    plasma.Redshift = 0.0
    plasma.norm = 1.0
    absorption = XSphabs()
    absorption.nH = nh_1e22_cm2
    energy_flux = float(
        np.sum(
            absorption(edges[:-1], edges[1:])
            * plasma(edges[:-1], edges[1:])
            * midpoint
            * KEV_TO_ERG
        )
    )
    return (
        energy_flux
        * 1.0e-14
        * PARSEC_TO_CM
        / (4.0 * np.pi)
        / ARCMIN2_PER_STERADIAN
        / 1.0e-15
    )


def convert_ueda_disk(
    source: pd.DataFrame,
    measurement_source: pd.DataFrame | None = None,
) -> tuple[pd.DataFrame, dict[str, float | int]]:
    """Project the Ueda et al. (2022) disk through the 14-field footprint."""
    rows: list[dict[str, float | str]] = []
    error_by_obsid: dict[str, float] = {}
    if measurement_source is not None:
        selected = measurement_source.loc[
            measurement_source["quality_primary"].astype(str).str.lower().eq("true")
        ]
        error_by_obsid = {
            f"{int(row['obsid']):010d}": float(
                row[
                    "absorbed_flux_standard_0p5_2p0_staterr_erg_cm-2_s-1_arcmin-2"
                ]
                / 1.0e-15
            )
            for row in selected.to_dict(orient="records")
        }
    with _xspec_settings("lodd", "vern"):
        for record in source.to_dict(orient="records"):
            emission_measure = _ueda_disk_emission_measure(
                float(record["galactic_l_deg"]),
                float(record["galactic_b_deg"]),
            )
            flux_per_em = _ueda_absorbed_flux_per_em(
                float(record["nh_hi4pi_1e22_cm-2"])
            )
            nominal_flux = emission_measure * flux_per_em
            obsid = f"{int(record['obsid']):010d}"
            rows.append(
                {
                    "obsid": obsid,
                    "galactic_l_deg": float(record["galactic_l_deg"]),
                    "galactic_b_deg": float(record["galactic_b_deg"]),
                    "nh_hi4pi_1e22_cm-2": float(
                        record["nh_hi4pi_1e22_cm-2"]
                    ),
                    "nominal_em_cm-6_pc": emission_measure,
                    "absorbed_0p5_2p0_flux_per_em_fluxunit": flux_per_em,
                    "nominal_absorbed_0p5_2p0_fluxunit": nominal_flux,
                    "n0_minus_1sigma_absorbed_0p5_2p0_fluxunit": (
                        nominal_flux * (3.3 / 3.4) ** 2
                    ),
                    "n0_plus_1sigma_absorbed_0p5_2p0_fluxunit": (
                        nominal_flux * (3.5 / 3.4) ** 2
                    ),
                    "measurement_staterr_absorbed_0p5_2p0_fluxunit": (
                        error_by_obsid.get(obsid, np.nan)
                    ),
                }
            )
    converted = pd.DataFrame(rows)
    nominal = converted["nominal_absorbed_0p5_2p0_fluxunit"]
    summary: dict[str, float | int] = {
        "footprint_fields": int(len(converted)),
        "nominal_minimum_fluxunit": float(nominal.min()),
        "nominal_median_fluxunit": float(nominal.median()),
        "nominal_maximum_fluxunit": float(nominal.max()),
        "n0_sensitivity_low_fluxunit": float(
            converted["n0_minus_1sigma_absorbed_0p5_2p0_fluxunit"].min()
        ),
        "n0_sensitivity_high_fluxunit": float(
            converted["n0_plus_1sigma_absorbed_0p5_2p0_fluxunit"].max()
        ),
    }
    if error_by_obsid:
        errors = converted[
            "measurement_staterr_absorbed_0p5_2p0_fluxunit"
        ].to_numpy(dtype=float)
        if not np.isfinite(errors).all():
            raise ValueError("Missing measurement error for a Ueda footprint field")
        weights = errors**-2
        for column, key in (
            (
                "nominal_absorbed_0p5_2p0_fluxunit",
                "all_field_inverse_variance_fluxunit",
            ),
            (
                "n0_minus_1sigma_absorbed_0p5_2p0_fluxunit",
                "all_field_inverse_variance_n0_low_fluxunit",
            ),
            (
                "n0_plus_1sigma_absorbed_0p5_2p0_fluxunit",
                "all_field_inverse_variance_n0_high_fluxunit",
            ),
        ):
            values = converted[column].to_numpy(dtype=float)
            summary[key] = float(np.sum(values * weights) / np.sum(weights))
    return converted, summary


def main(output_dir: Path = HERE) -> None:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    efeds = convert_efeds_scenarios()
    hs_source = pd.read_csv(HENLEY_SHELTON_INPUT_CSV)
    hs_converted, hs_summary = convert_henley_shelton(hs_source)
    footprint_source = pd.read_csv(M31_FOOTPRINT_INPUT_CSV)
    measurement_source = pd.read_csv(PRIMARY_MEASUREMENTS_INPUT_CSV)
    ueda_converted, ueda_summary = convert_ueda_disk(
        footprint_source, measurement_source
    )

    efeds.to_csv(
        output_dir / "m31_cgmsum_efeds2023_absorbed_0p5_2p0.csv", index=False
    )
    hs_converted.to_csv(
        output_dir / "m31_cgmsum_henley_shelton2013_absorbed_0p5_2p0.csv",
        index=False,
    )
    ueda_converted.to_csv(
        output_dir
        / "m31_cgmsum_ueda2022_disk_m31_footprint_absorbed_0p5_2p0.csv",
        index=False,
    )
    audit = {
        "schema_version": 1,
        "figure_band": "absorbed 0.5-2.0 keV",
        "flux_unit": "1e-15 erg cm-2 s-1 arcmin-2",
        "efeds": {
            "reference": "Ponti et al. (2023)",
            "doi": "10.1051/0004-6361/202243992",
            "scope": "107.5 deg2 eFEDS field; l=220-235 deg, b=20-40 deg; western-sky comparator, not an M31-footprint prior",
            "native_quantity": "observed/best-fitting CGM-component surface brightness from Table 4; attenuated use follows the component-specific model rather than a verbatim table label",
            "disnht_sigma_primary_products": "PDF/arXiv=0.117; publisher HTML=0.177; executable uses PDF/arXiv value 0.117",
            "conversion": "Lodders APEC with the published kT/Z; derive F(0.5-0.6)/F(0.3-0.6) with 40-node lognormal-NH quadrature approximating disnht(logNH=20.51,sigma=0.117), then add the published 0.6-2.0 flux",
            "scenarios": efeds.to_dict(orient="records"),
            "envelope_low_fluxunit": float(
                efeds["flux_0p5_2p0_reconstruction_low_fluxunit"].min()
            ),
            "envelope_high_fluxunit": float(
                efeds["flux_0p5_2p0_reconstruction_high_fluxunit"].max()
            ),
        },
        "henley_shelton": {
            "reference": "Henley & Shelton (2013)",
            "doi": "10.1088/0004-637X/773/2/92",
            "vizier_catalog": "J/ApJ/773/92",
            "input_sha256": hashlib.sha256(
                HENLEY_SHELTON_INPUT_CSV.read_bytes()
            ).hexdigest(),
            "scope": "high-latitude XMM-Newton sightline population; not an M31-footprint prior",
            "native_quantity": "intrinsic 0.5-2.0 keV one-temperature Raymond-Smith halo surface brightness",
            "conversion": "field-specific XSraymond * XSphabs using published T and LAB NHI; angr abundances and bcmc cross-sections",
            "summary": hs_summary,
        },
        "ueda_disk": {
            "reference": "Ueda et al. (2022)",
            "doi": "10.1093/pasj/psac077",
            "scope": "exponential electron-density disk from the Table 3 2005-2009, |l|>105 deg, N=36 fit; projected through the 14-field M31 footprint",
            "parent_sample_domain": "75 < l < 285 deg; |b| > 15 deg",
            "adopted_fit_selection": "2005-2009; |l| > 105 deg; N=36",
            "native_quantity": "EM = integral n_e n_H ds for the approximately 0.22 keV MWH component",
            "published_disk_equation": "n = n0 exp(-R/7.0 kpc) exp(-z/2.7 kpc); Table 3 reports central electron density n_e0",
            "reprojection_augmentations": "midplane symmetry |z|; n_e/n_H=1.2; Rsun=8.2 kpc; explicit internal-observer line-of-sight transform",
            "implemented_model": "n_e = 3.4e-3 exp(-R/7.0 kpc) exp(-|z|/2.7 kpc) cm-3; n_H=n_e/1.2",
            "conversion": "field-specific Lodders APEC(kT=0.22 keV,Z=1) * phabs(HI4PI); absorbed 0.5-2.0 keV",
            "interpretation": "M31-footprint directional conditional model; nominal field spread and n0 marginal sensitivity are not a posterior interval",
            "input_sha256": hashlib.sha256(
                M31_FOOTPRINT_INPUT_CSV.read_bytes()
            ).hexdigest(),
            "summary": ueda_summary,
        },
    }
    (output_dir / "m31_cgmsum_external_mw_comparator_audit.json").write_text(
        json.dumps(audit, indent=2) + "\n", encoding="utf-8"
    )


if __name__ == "__main__":
    main()