#!/usr/bin/env python
"""Per-source figure: extinction-density profiles from Edenhofer24 / Bayestar19 / DECaPS
overlaid, adopted distance marked. One PNG per target in extinction_profiles/figs/."""
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from extinction_profiles import load_targets, OUT

FIGDIR = os.path.join(OUT, 'figs')
os.makedirs(FIGDIR, exist_ok=True)


def load_npz(fname):
    p = os.path.join(OUT, fname)
    if not os.path.exists(p):
        return {}
    return dict(np.load(p))


def main():
    eden = load_npz('edenhofer_profiles.npz')
    bay = load_npz('bayestar19_profiles.npz')
    dec = load_npz('decaps_profiles.npz')
    targs = load_targets()
    for t in targs:
        n = t['name']
        fig, ax = plt.subplots(figsize=(7, 4))
        handles = 0
        if n in eden:
            d, dens, std = eden[n]
            norm = np.nanmax(dens) or 1
            ax.plot(d / 1000, dens / norm, color='C0', label='Edenhofer+24 (Gaia)')
            ax.fill_between(d / 1000, (dens - std) / norm, (dens + std) / norm,
                            color='C0', alpha=0.2)
            handles += 1
        for store, lab, col in ((bay, 'Bayestar19', 'C1'), (dec, 'DECaPS', 'C2')):
            if n in store:
                d, cum, dens = store[n]
                norm = np.nanmax(dens) or 1
                ax.plot(d / 1000, dens / norm, color=col, label=lab)
                handles += 1
        if not handles:
            plt.close(fig)
            continue
        ax.axvline(t['d_best'], color='k', ls='--', lw=1.5,
                   label=f"adopted {t['d_best']:.2f} kpc ({t['method']})")
        ax.set_xlabel('distance [kpc]')
        ax.set_ylabel('normalized extinction density dA/dd')
        ax.set_xlim(0, 6)
        ax.set_title(f"{n}  l={t['glon']:.2f} b={t['glat']:.2f}  conf={t['conf']}")
        ax.legend(fontsize=8, loc='upper right')
        fig.tight_layout()
        fig.savefig(os.path.join(FIGDIR, f'{n}.png'), dpi=110)
        plt.close(fig)
    print(f'figs -> {FIGDIR}')


if __name__ == '__main__':
    main()