from __future__ import annotations

import json
from pathlib import Path

import pandas as pd
from sherpa.astro.ui import set_xsabund, set_xsxsect

from audit_apj_v19_cgmsum_conditional_priors import energy_flux_per_norm

HERE = Path(__file__).resolve().parent
OUTPUT_CSV = HERE / "m31_cgmsum_apec_bandpass_matrix.csv"
OUTPUT_JSON = HERE / "m31_cgmsum_apec_bandpass_audit.json"

TEMPERATURES_KEV = (0.17531, 0.2)
ABUNDANCES_SOLAR = (0.1, 0.3)
NH_VALUES_1E22 = (None, 0.04, 0.056, 0.08)
BANDS_KEV = (
    (0.2, 0.4),
    (0.3, 0.4),
    (0.4, 1.25),
    (0.4, 2.0),
    (0.5, 2.0),
    (0.3, 2.0),
    (0.2, 2.0),
)


def band_label(low_keV: float, high_keV: float) -> str:
    return f"{low_keV:g}-{high_keV:g}"


def calculate_band_matrix() -> pd.DataFrame:
    set_xsabund("angr")
    set_xsxsect("vern")
    rows: list[dict[str, float | str]] = []
    for temperature in TEMPERATURES_KEV:
        for abundance in ABUNDANCES_SOLAR:
            intrinsic_total = energy_flux_per_norm(
                temperature, abundance, None, 0.2, 2.0
            )
            for nh in NH_VALUES_1E22:
                fluxes = {
                    band_label(low, high): energy_flux_per_norm(
                        temperature, abundance, nh, low, high
                    )
                    for low, high in BANDS_KEV
                }
                state_total = fluxes["0.2-2"]
                standard_flux = fluxes["0.5-2"]
                for low, high in BANDS_KEV:
                    label = band_label(low, high)
                    flux = fluxes[label]
                    rows.append(
                        {
                            "kT_keV": temperature,
                            "abundance_solar": abundance,
                            "nh_1e22_cm-2": -1.0 if nh is None else nh,
                            "state": "unabsorbed" if nh is None else "absorbed",
                            "band_keV": label,
                            "flux_per_apec_norm": flux,
                            "fraction_of_intrinsic_0p2_2": flux / intrinsic_total,
                            "fraction_of_same_state_0p2_2": flux / state_total,
                            "ratio_to_same_state_0p5_2": flux / standard_flux,
                        }
                    )
    return pd.DataFrame(rows)


def summarize(matrix: pd.DataFrame) -> dict[str, object]:
    flux = matrix.pivot_table(
        index=["kT_keV", "abundance_solar", "nh_1e22_cm-2", "state"],
        columns="band_keV",
        values="flux_per_apec_norm",
    )
    ratios = {
        "flux_0p4_1p25_over_0p4_2": flux["0.4-1.25"] / flux["0.4-2"],
        "flux_0p4_1p25_over_0p5_2": flux["0.4-1.25"] / flux["0.5-2"],
        "flux_0p3_2_over_0p4_1p25": flux["0.3-2"] / flux["0.4-1.25"],
        "flux_0p2_0p4_over_0p2_2": flux["0.2-0.4"] / flux["0.2-2"],
    }
    ranges = {
        name: {"minimum": float(values.min()), "maximum": float(values.max())}
        for name, values in ratios.items()
    }
    representative = matrix.loc[
        (matrix["kT_keV"] == 0.17531)
        & (matrix["abundance_solar"] == 0.3)
        & (matrix["nh_1e22_cm-2"] == 0.056)
    ].set_index("band_keV")
    return {
        "model": "xsphabs*xsapec",
        "abundance_table": "angr",
        "cross_sections": "vern",
        "parameter_grid": {
            "kT_keV": list(TEMPERATURES_KEV),
            "abundance_solar": list(ABUNDANCES_SOLAR),
            "nh_1e22_cm-2": ["unabsorbed", 0.04, 0.056, 0.08],
        },
        "ratio_ranges_over_full_grid": ranges,
        "representative_kT0p17531_Z0p3_NH0p056": {
            band: {
                "fraction_of_intrinsic_0p2_2": float(
                    representative.loc[band, "fraction_of_intrinsic_0p2_2"]
                ),
                "fraction_of_same_state_0p2_2": float(
                    representative.loc[band, "fraction_of_same_state_0p2_2"]
                ),
                "ratio_to_same_state_0p5_2": float(
                    representative.loc[band, "ratio_to_same_state_0p5_2"]
                ),
            }
            for band in representative.index
        },
    }


def main() -> None:
    matrix = calculate_band_matrix()
    audit = summarize(matrix)
    matrix.to_csv(OUTPUT_CSV, index=False)
    OUTPUT_JSON.write_text(
        json.dumps(audit, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(f"Wrote {OUTPUT_CSV}")
    print(f"Wrote {OUTPUT_JSON}")


if __name__ == "__main__":
    main()