#!/usr/bin/env python3
"""Plot frozen CGMsum and MW comparators in observer-centred Galactic coordinates.

No spectra, emissivities or model parameters are refitted. MW model values are
consumed from the prior audited per-sightline tables. Run with scratch TMPDIR
and MPLCONFIGDIR. Input copies are pinned by mw_coordinate_sources/manifest.
"""
from __future__ import annotations
import hashlib
import json
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
import numpy as np
import pandas as pd
from astropy.coordinates import SkyCoord
import astropy.units as u
import scienceplots  # noqa: F401

ROOT = Path(__file__).resolve().parent
SOURCES = ROOT / 'mw_coordinate_sources'
PRODUCTS = ROOT / 'products'
FIGURES = ROOT / 'figures'
R_SUN = 8.2
APERTURE_DEG = 15 / 60
UNIT = '1e-15 erg cm-2 s-1 arcmin-2; absorbed 0.5-2.0 keV'
SIDE_COLORS = {'North/NW': '#195b9b', 'South/SE': '#b13b44'}
MODELS = [
    dict(id='locatelli_reference', label='Locatelli 2024 reference (extrap.)', color='#873e9b',
         kT_keV=.15, Z_solar=.3, coordinate='disk n(R,z) + spherical halo n(r); integrate each (l,b)',
         transfer='audited response-smoothed O VIII normalization; Z=0.3 spectral bridge',
         domain='M31 is outside western-sky fitted longitude domain',
         source='https://arxiv.org/abs/2310.10715'),
    dict(id='ueda_disk', label='Ueda 2022 disk', color='#d99819',
         kT_keV=.22, Z_solar=1., coordinate='exponential disk n(R,z); integrate each (l,b)',
         transfer='audited source-temperature APEC, Lodders, phabs/vern and field HI4PI',
         domain='all 14 fields inside adopted |l| > 105 deg disk-fit angular selection',
         source='https://arxiv.org/abs/2209.01698'),
    dict(id='halosat_combined', label='HaloSat 2020 combined (extrap.)', color='#17836e',
         kT_keV=.225, Z_solar=.3, coordinate='disk n(R,z) + adiabatic halo n(r); integrate each (l,b)',
         transfer='audited sample-median temperature, Wilms/Verner TBabs and field HI4PI',
         domain='all 14 fields outside original b < -30 deg fit domain',
         source='https://arxiv.org/abs/2011.00126'),
]


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


def gc_angle(l_deg, b_deg):
    """Angular separation from the ideal Galactic centre (l,b)=(0,0)."""
    cosine = np.cos(np.deg2rad(b_deg)) * np.cos(np.deg2rad(l_deg))
    return np.rad2deg(np.arccos(np.clip(cosine, -1, 1)))


def load_inputs():
    manifest = json.loads((SOURCES / 'input_manifest.json').read_text())
    for item in manifest['files']:
        p = SOURCES / item['file']
        if p.stat().st_size != item['bytes'] or sha256(p) != item['sha256']:
            raise ValueError('Input mismatch: ' + item['file'])
    frozen = pd.read_csv(SOURCES / 'fields.csv', dtype={'obsid': str})
    if len(frozen) != 14 or not frozen.quality_primary.all():
        raise ValueError('Expected frozen 14-field primary sample')
    columns = ['obsid', 'side', 'ra_deg', 'dec_deg', 'rproj_kpc', 'galactic_l_deg',
               'galactic_b_deg', 'nh_hi4pi_1e22_cm-2']
    fields = frozen[columns].copy()
    fields['observed_cgmsum_fluxunit'] = frozen.cgmsum_absorbed_0p5_2p0_figure_fluxunit
    fields['local_staterr_fluxunit'] = frozen.cgmsum_absorbed_0p5_2p0_staterr_figure_fluxunit
    fields['psi_gc_deg'] = gc_angle(fields.galactic_l_deg, fields.galactic_b_deg)
    cosine = np.cos(np.deg2rad(fields.psi_gc_deg))
    fields['closest_point_signed_s_kpc'] = R_SUN * cosine
    # Minimize r(s) only over the observable ray s>=0, not the infinite line.
    fields['ray_min_gc_radius_kpc'] = np.where(cosine > 0, R_SUN*np.sqrt(1-cosine**2), R_SUN)
    fields['infinite_line_impact_parameter_kpc'] = R_SUN*np.sqrt(1-cosine**2)
    fields['nominal_aperture_radius_deg'] = APERTURE_DEG
    sky = SkyCoord(fields.ra_deg.to_numpy()*u.deg, fields.dec_deg.to_numpy()*u.deg).galactic
    coordinate_error = max(np.max(abs(sky.l.deg-fields.galactic_l_deg)),
                           np.max(abs(sky.b.deg-fields.galactic_b_deg)))
    psi_error = np.max(abs(sky.separation(SkyCoord(l=0*u.deg,b=0*u.deg,frame='galactic')).deg-fields.psi_gc_deg))
    assert coordinate_error < 1e-8 and psi_error < 1e-8
    assert (fields.closest_point_signed_s_kpc < 0).all()
    assert (fields.ray_min_gc_radius_kpc == R_SUN).all()
    values = {'locatelli_reference': frozen.mw_reference_actual_nh_figure_fluxunit.to_numpy()}
    for name, model_id in [('ueda.csv', 'ueda_disk'), ('halosat.csv', 'halosat_combined')]:
        table = pd.read_csv(SOURCES/name, dtype={'obsid': str}).set_index('obsid')
        assert table.index.is_unique and set(table.index) == set(fields.obsid)
        ordered = table.loc[fields.obsid]
        assert np.allclose(ordered.galactic_l_deg, fields.galactic_l_deg, atol=1e-8, rtol=0)
        assert np.allclose(ordered.galactic_b_deg, fields.galactic_b_deg, atol=1e-8, rtol=0)
        assert np.allclose(ordered['nh_hi4pi_1e22_cm-2'], fields['nh_hi4pi_1e22_cm-2'], atol=1e-10, rtol=0)
        values[model_id] = ordered.nominal_absorbed_0p5_2p0_fluxunit.to_numpy()
        if model_id == 'halosat_combined':
            assert ordered.outside_halosat_fit_domain.all()
    rows = []
    for model in MODELS:
        for row, value in zip(fields.itertuples(), values[model['id']]):
            rows.append(dict(obsid=row.obsid, model_id=model['id'], galactic_l_deg=row.galactic_l_deg,
                             galactic_b_deg=row.galactic_b_deg, psi_gc_deg=row.psi_gc_deg,
                             absorbed_prediction_fluxunit=float(value), domain=model['domain'],
                             transfer=model['transfer']))
    predictions = pd.DataFrame(rows)
    population = pd.read_csv(SOURCES/'henley.csv', dtype={'obsid': str})
    assert len(population) == 86 and (abs(population.glat_deg) > 30).all()
    population['psi_gc_deg'] = gc_angle(population.glon_deg, population.glat_deg)
    checks = dict(independent_coordinate_max_error_deg=float(coordinate_error),
                  independent_gc_angle_max_error_deg=float(psi_error),
                  outward_sightlines=int((cosine < 0).sum()), ray_min_radius_kpc=R_SUN,
                  model_prediction_rows=len(predictions), population_detection_rows=len(population))
    return manifest, fields, predictions, population, checks


def observed(ax, fields, xcolumn, *, nominal_extent=True):
    for side, color in SIDE_COLORS.items():
        sub = fields[fields.side == side]
        ax.errorbar(sub[xcolumn], sub.observed_cgmsum_fluxunit,
                    yerr=sub.local_staterr_fluxunit, xerr=APERTURE_DEG if nominal_extent else None,
                    fmt='o' if side == 'North/NW' else 's', color=color,
                    ms=3.5, lw=.65, capsize=1.3, label='CGMsum '+side, zorder=5)


def save(fig, stem, column):
    fig.tight_layout(pad=1.2, h_pad=2)
    for ext in ['png', 'pdf']:
        fig.savefig(FIGURES/f'{stem}_{column}.{ext}', dpi=400)
    plt.close(fig)


def make_figures(fields, predictions, population):
    plt.style.use(['science', 'no-latex'])
    for column, width in [('1col', 3.35), ('2col', 7.)]:
        plt.rcParams.update({'font.size': 8 if column == '1col' else 10,
                             'legend.fontsize': 6 if column == '1col' else 8})
        fig, axes = plt.subplots(2, 1, figsize=(width, 8.5 if column == '1col' else 8.7))
        for ax, xcolumn, title, xlabel in [
            (axes[0], 'galactic_b_deg', '(a) Galactic latitude: disk geometry', 'Galactic latitude b (deg)'),
            (axes[1], 'psi_gc_deg', '(b) Angular separation from the Galactic centre', r'Angle to Galactic centre $\psi$ (deg)'),
        ]:
            observed(ax, fields, xcolumn)
            for model in MODELS:
                sub = predictions[predictions.model_id == model['id']]
                ax.scatter(sub[xcolumn], sub.absorbed_prediction_fluxunit, marker='+', s=28,
                           lw=.9, color=model['color'], label=model['label'], zorder=4)
            ax.set(xlim=(fields[xcolumn].min()-.55, fields[xcolumn].max()+.55), ylim=(0, 1.75),
                   xlabel=xlabel, ylabel='Absorbed 0.5-2 keV SB\n'+r'[$10^{-15}$ erg cm$^{-2}$ s$^{-1}$ arcmin$^{-2}$]')
            ax.set_title(title, loc='left', fontsize=8)
            ax.grid(axis='y', alpha=.15)
            ax.legend(loc='lower right', frameon=True, framealpha=.9)
            ax.text(.02, .98, 'MW predictions at actual (l, b, NH)\nNo continuous one-coordinate fit',
                    ha='left', va='top', transform=ax.transAxes, fontsize=6)
        save(fig, 'stack_comparison_mw_coordinates', column)
        fig, ax = plt.subplots(figsize=(width, 4.1 if column == '1col' else 4.5))
        ax.scatter(population.glat_deg, population.absorbed_0p5_2p0_fluxunit,
                   s=10, alpha=.5, color='.45', label='H&S 2013: 86 high-latitude detections')
        observed(ax, fields, 'galactic_b_deg', nominal_extent=False)
        ax.axvspan(fields.galactic_b_deg.min(), fields.galactic_b_deg.max(), color='#195b9b', alpha=.10)
        ax.set(xlim=(-90,90), ylim=(.06,3), yscale='log', xlabel='Galactic latitude b (deg)',
               ylabel='Absorbed 0.5-2 keV SB\n'+r'[$10^{-15}$ erg cm$^{-2}$ s$^{-1}$ arcmin$^{-2}$]')
        ax.yaxis.set_major_formatter(FuncFormatter(lambda y, _: f'{y:g}'))
        ax.set_title('Population context; H&S values are not M31 foregrounds', loc='left', fontsize=7)
        ax.legend(loc='upper left', frameon=True)
        ax.grid(axis='y', alpha=.15)
        save(fig, 'stack_comparison_mw_population', column)


def main():
    manifest, fields, predictions, population, checks = load_inputs()
    fields.to_csv(PRODUCTS/'stack_comparison_mw_fields.csv', index=False)
    predictions.to_csv(PRODUCTS/'stack_comparison_mw_predictions.csv', index=False)
    population.to_csv(PRODUCTS/'stack_comparison_mw_population.csv', index=False)
    summary = dict(schema_version=1, observed_quantity='unchanged approximately 0.175 keV MW+M31 CGMsum',
                   unit=UNIT, field_count=14, no_refit=True, solar_radius_kpc=R_SUN, solar_height_kpc=0,
                   galactic_centre='ideal (l,b)=(0,0)',
                   psi_definition='acos(cos(b)*cos(l)); observer-centred angle, not gas distance',
                   line_of_sight_radius='r(s)^2=Rsun^2+s^2-2*Rsun*s*cos(b)*cos(l), s>=0',
                   intrinsic_radius_is_not_a_field_coordinate=True,
                   aperture='nominal 15 arcmin angular radius, not exact extraction masks',
                   predictions='frozen audited per-field LOS predictions with field-specific HI4PI NH; no continuous b-only or psi-only model',
                   spectral_boundary='same absorbed band; study-specific temperature/abundance and audit transfer retained, no common-temperature refit',
                   population_boundary='86 convertible detections only, |b|>30; detection-selected population context, not local foreground prediction',
                   models=MODELS, input_hashes=manifest['files'], checks=checks,
                   coordinate_ranges={col:[float(fields[col].min()),float(fields[col].max())]
                                      for col in ['galactic_l_deg','galactic_b_deg','psi_gc_deg']},
                   generator_sha256=sha256(__file__))
    (PRODUCTS/'stack_comparison_mw_summary.json').write_text(json.dumps(summary,indent=2)+'\n')
    make_figures(fields,predictions,population)
    print(json.dumps(summary['checks'],indent=2))


if __name__ == '__main__':
    main()
