#!/usr/bin/env python3
"""Export response-folded MOS2 R0-15 background counts for Fe-L image bands."""

from __future__ import annotations

import json
import os
from pathlib import Path

import numpy as np

CIAO_ROOT = "/home/huangrui/software/ciao-4.18"
os.environ["HEADAS"] = f"{CIAO_ROOT}/spectral"
os.environ["XSPEC_MDATA_DIR"] = f"{CIAO_ROOT}/spectral/modelData/"

from sherpa.astro import ui

ROOT = Path(__file__).resolve().parent
SESSION = (
    ROOT
    / "joint_spectrum_fitting_2T_basedon_region_v21"
    / "regionR0-15_spectrum.fitting"
)
OUT = (
    ROOT
    / "joint_spectrum_fitting_2T_basedon_region_v22"
    / "r47_mos2_fel_ratio_map_pilot_20260713"
    / "r47_mos2_fel_background_spectral_contract.json"
)
DATASET = "SRC_2"
BANDS = {"low": (0.70, 0.875), "high": (0.875, 1.05)}
COMMON_SP = {
    "PhoIndx1": 1.2965306521244002,
    "BreakE": 3.2,
    "PhoIndx2": 2.9106868928719,
    "norm": 6.24576708953368e-05,
}


def integrate_model(plot, values: np.ndarray, band: tuple[float, float]) -> float:
    centers = np.asarray(plot.x)
    widths = np.asarray(plot.xhi - plot.xlo)
    selected = (centers >= band[0]) & (centers < band[1])
    return float(np.sum(np.asarray(values)[selected] * widths[selected]))


def channel_counts(data, energy: np.ndarray, band: tuple[float, float]) -> float:
    selected = (energy >= band[0]) & (energy < band[1])
    return float(np.sum(np.asarray(data.counts)[selected]))


def model_values() -> tuple[object, np.ndarray]:
    plot = ui.get_model_plot(DATASET)
    return plot, np.asarray(plot.y).copy()


def main() -> None:
    ui.clean()
    ui.restore(str(SESSION))
    ui.ungroup(DATASET)
    data = ui.get_data(DATASET)
    qpb = ui.get_bkg(DATASET)
    _, rmf = data.get_response()
    channel_energy = 0.5 * (np.asarray(rmf.e_min) + np.asarray(rmf.e_max))
    exposure = float(data.exposure)
    qpb_scale = float(np.asarray(data.get_background_scale()).reshape(-1)[0])
    backscal = float(data.backscal)
    skyarea_arcmin2 = backscal * (1.0 / 20.0 / 60.0) ** 2

    # v21 ties all source/sky photon components to the first source dataset;
    # only the detector soft-proton component remains per dataset.
    source = ui.get_model_component("SRCspec_1_SrcApec")
    soft_proton = ui.get_model_component("SRCspec_2_SoftProton")
    line_names = (
        "SRCspec_2_mos1_InstLine1",
        "SRCspec_2_mos1_InstLine2",
        "SRCspec_2_mos_InstLine3",
    )
    lines = [ui.get_model_component(name) for name in line_names]

    plot, full_native = model_values()
    source_norm = float(source.norm.val)
    source.norm.val = 0.0
    _, background_native = model_values()
    source_only = full_native - background_native

    native_sp_parameters = {parameter.name: float(parameter.val) for parameter in soft_proton.pars}
    native_sp_norm = float(soft_proton.norm.val)
    soft_proton.norm.val = 0.0
    _, no_sp = model_values()
    native_sp_only = background_native - no_sp

    line_norms = [float(line.norm.val) for line in lines]
    for line in lines:
        line.norm.val = 0.0
    _, sky_only = model_values()
    instrumental_line_only = no_sp - sky_only

    for parameter in soft_proton.pars:
        common_value = COMMON_SP[parameter.name]
        if common_value > parameter.max:
            parameter.max = common_value
        if common_value < parameter.min:
            parameter.min = common_value
        parameter.val = common_value
    _, common_sp_plus_sky = model_values()
    common_sp_only = common_sp_plus_sky - sky_only

    rows = {}
    for label, band in BANDS.items():
        rows[label] = {
            "energy_keV": list(band),
            "raw_pha_counts": channel_counts(data, channel_energy, band),
            "qpb_pha_counts_unscaled": channel_counts(qpb, channel_energy, band),
            "qpb_pha_counts_scaled": qpb_scale * channel_counts(qpb, channel_energy, band),
            "source_apec_model_counts": exposure * integrate_model(plot, source_only, band),
            "sky_background_model_counts": exposure * integrate_model(plot, sky_only, band),
            "instrumental_line_model_counts": exposure
            * integrate_model(plot, instrumental_line_only, band),
            "soft_proton_native_mos2_model_counts": exposure
            * integrate_model(plot, native_sp_only, band),
            "soft_proton_common_r015_model_counts": exposure
            * integrate_model(plot, common_sp_only, band),
        }

    source.norm.val = source_norm
    soft_proton.norm.val = native_sp_norm
    for line, norm in zip(lines, line_norms):
        line.norm.val = norm

    payload = {
        "dataset": DATASET,
        "session": str(SESSION),
        "pha": str(data.name),
        "qpb_pha": str(qpb.name),
        "exposure_s": exposure,
        "backscal": backscal,
        "skyarea_arcmin2": skyarea_arcmin2,
        "qpb_scale": qpb_scale,
        "native_mos2_soft_proton_parameters": native_sp_parameters,
        "common_r015_soft_proton_parameters": COMMON_SP,
        "instrumental_line_norms_restored": dict(zip(line_names, line_norms)),
        "bands": rows,
        "interpretation": {
            "sky": "Vignetted photon background; project as exposure-map-proportional counts.",
            "qpb": "Unvignetted particle background; use mos_back detector map rotated to sky.",
            "soft_proton": "Use proton detector template; do not apply photon exposure-map vignetting.",
        },
    }
    OUT.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()