#!/usr/bin/env python
"""ALS II catalog figures for the writeup:
(1) fig_als_topdown.png — Galactic-plane XY view, color-coded in 6 luminosity bins 1e3-1e6.
(2) fig_als_lf.png — fine-binned LF with Salpeter-IMF prediction locked to the lowest-L bin,
    plus a bootstrap-banded KDE.
SpType -> L as in pg_lf_final.py."""
import pickle
import re
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

SPT_L = {'O3': 5.8, 'O4': 5.68, 'O5': 5.55, 'O5.5': 5.5, 'O6': 5.4, 'O6.5': 5.32,
         'O7': 5.23, 'O7.5': 5.14, 'O8': 5.05, 'O8.5': 4.96, 'O9': 4.87, 'O9.5': 4.77,
         'B0': 4.6, 'B0.5': 4.4, 'B1': 4.15, 'B1.5': 3.95, 'B2': 3.75, 'B2.5': 3.6,
         'B3': 3.4, 'B4': 3.2, 'B5': 3.0, 'B6': 2.85, 'B7': 2.7, 'B8': 2.55, 'B9': 2.4}
CLASS_BOOST = {'I': 0.7, 'II': 0.45, 'III': 0.25, 'IV': 0.12, 'V': 0.0}
# ZAMS m<->L mapping for the Salpeter conversion
ZM = np.array([2, 3, 5, 7, 9, 12, 15, 20, 25, 32, 40, 60, 85])
ZL = np.array([1.26, 1.9, 2.75, 3.25, 3.62, 3.98, 4.25, 4.56, 4.79, 5.02, 5.2, 5.5, 5.72])


def gc_arrow(ax, color='crimson', L=1.55):
    """Arrow from the Sun (origin) toward the Galactic center (+X direction)."""
    ax.annotate('', xy=(0, L), xytext=(0, 0),
                arrowprops=dict(arrowstyle='-|>', color=color, lw=2.0, alpha=0.85,
                                shrinkA=0, shrinkB=0), zorder=6)
    ax.text(0.13, L, 'G.C.', color=color, fontsize=8.5, va='center', ha='left',
            fontweight='bold', zorder=6)


def load_als(dmax_pc=3000):
    als = pickle.load(open('als_combined.pkl', 'rb'))
    import astropy.units as u
    from astropy.coordinates import SkyCoord
    d = np.array(als['DistMode'], float)
    ok = np.isfinite(d) & (d > 0) & (d < dmax_pc)
    logL = np.full(len(als), np.nan)
    for i, s in enumerate(als['SpType']):
        s = str(s)
        m = re.match(r'^([OB])\s*(\d+(?:\.5)?)', s)
        if not m:
            continue
        key = m.group(1) + m.group(2)
        if key not in SPT_L:
            continue
        boost = 0.0
        mc = re.search(r'(I{1,3}|IV|V)\b', s)
        if mc:
            boost = CLASS_BOOST.get(mc.group(1), 0.0)
        logL[i] = SPT_L[key] + boost
    ok &= np.isfinite(logL)
    c = SkyCoord(als['RAJ2000'], als['DEJ2000'], unit=(u.hourangle, u.deg)).galactic
    gl = np.radians(c.l.deg[ok])
    dk = d[ok] / 1000
    x = dk * np.cos(gl)   # toward l=0
    y = dk * np.sin(gl)   # toward l=90
    return x, y, logL[ok]


def topdown(x, y, logL):
    fig, ax = plt.subplots(figsize=(7.2, 7))
    edges = [3, 3.5, 4, 4.5, 5, 5.5, 6]
    cmap = plt.cm.plasma(np.linspace(0.05, 0.95, 6))
    for i in range(6):
        m = (logL >= edges[i]) & (logL < edges[i+1])
        ax.scatter(y[m], x[m], s=3 + 5*i, alpha=0.5, color=cmap[i],
                   label=f'{edges[i]:.1f}$\\leq\\log L<${edges[i+1]:.1f} (N={m.sum()})')
    for r in (1, 2, 3):
        th = np.linspace(0, 2*np.pi, 200)
        ax.plot(r*np.sin(th), r*np.cos(th), 'k:', lw=0.7, alpha=0.5)
    ax.plot(0, 0, 'k*', ms=12)
    ax.annotate('Sun', (0, 0), xytext=(-16, 4), textcoords='offset points', fontsize=9)
    gc_arrow(ax)
    ax.set_xlabel('Y [kpc] (toward $l=90^\\circ$)')
    ax.set_ylabel('X [kpc] (toward Galactic center)')
    ax.set_aspect('equal')
    ax.set_xlim(-3.3, 3.3)
    ax.set_ylim(-3.3, 3.3)
    ax.legend(fontsize=7.5, loc='lower left')
    ax.set_title('ALS II OB stars, $d<3$ kpc (top-down)', fontsize=11)
    fig.tight_layout()
    fig.savefig('fig_als_topdown.png', dpi=140)
    plt.close(fig)


def topdown_hiL(x, y, logL, thresh=5.0):
    """Same top-down view but only sources above `thresh` in log L."""
    fig, ax = plt.subplots(figsize=(7.2, 7))
    edges = [5.0, 5.25, 5.5, 5.75, 6.5]
    lbls = ['5.00', '5.25', '5.50', r'$\geq$5.75']
    cmap = plt.cm.plasma(np.linspace(0.1, 0.9, len(edges) - 1))
    sel = logL >= thresh
    for i in range(len(edges) - 1):
        m = (logL >= edges[i]) & (logL < edges[i+1])
        top = f'{edges[i]:.2f}$\\leq\\log L<${edges[i+1]:.2f}' if i < len(edges) - 2 \
            else r'$\log L\geq5.75$'
        ax.scatter(y[m], x[m], s=18 + 14*i, alpha=0.6, color=cmap[i], edgecolor='k',
                   linewidth=0.3, label=f'{top} (N={m.sum()})')
    for r in (1, 2, 3):
        th = np.linspace(0, 2*np.pi, 200)
        ax.plot(r*np.sin(th), r*np.cos(th), 'k:', lw=0.7, alpha=0.5)
    ax.plot(0, 0, 'k*', ms=12)
    ax.annotate('Sun', (0, 0), xytext=(-16, 4), textcoords='offset points', fontsize=9)
    gc_arrow(ax)
    ax.set_xlabel('Y [kpc] (toward $l=90^\\circ$)')
    ax.set_ylabel('X [kpc] (toward Galactic center)')
    ax.set_aspect('equal')
    ax.set_xlim(-3.3, 3.3)
    ax.set_ylim(-3.3, 3.3)
    ax.legend(fontsize=8, loc='lower left')
    ax.set_title(f'ALS II OB stars, $\\log L\\geq{thresh:g}$, $d<3$ kpc (top-down); '
                 f'N={sel.sum()}', fontsize=10.5)
    fig.tight_layout()
    fig.savefig('fig_als_topdown_hiL.png', dpi=140)
    plt.close(fig)


def density_maps(x, y, logL, threshs=(4.0, 5.0)):
    """Surface-density (stars/kpc^2) maps of ALS OB stars above each threshold.

    Density from a Gaussian KDE of the (Y,X) positions, normalized so the integral
    over the plane equals the source count -> the colorbar is stars per kpc^2.
    """
    from scipy.stats import gaussian_kde
    gy = np.linspace(-3.3, 3.3, 220)
    gx = np.linspace(-3.3, 3.3, 220)
    GY, GX = np.meshgrid(gy, gx)
    pos = np.vstack([GY.ravel(), GX.ravel()])
    fig, axes = plt.subplots(1, len(threshs), figsize=(6.4 * len(threshs), 6.6))
    for ax, thr in zip(np.atleast_1d(axes), threshs):
        m = logL >= thr
        n = int(m.sum())
        kde = gaussian_kde(np.vstack([y[m], x[m]]), bw_method=0.14)
        dens = kde(pos).reshape(GX.shape) * n            # stars per kpc^2
        pcm = ax.pcolormesh(GY, GX, dens, cmap='magma', shading='auto')
        cb = fig.colorbar(pcm, ax=ax, fraction=0.046, pad=0.03)
        cb.set_label(r'surface density [stars kpc$^{-2}$]', fontsize=9)
        ax.scatter(y[m], x[m], s=4, color='w', alpha=0.35, linewidth=0)
        for r in (1, 2, 3):
            th = np.linspace(0, 2*np.pi, 200)
            ax.plot(r*np.sin(th), r*np.cos(th), 'w:', lw=0.7, alpha=0.6)
        ax.plot(0, 0, 'w*', ms=13)
        ax.annotate('Sun', (0, 0), xytext=(-16, 4), textcoords='offset points',
                    fontsize=9, color='w')
        gc_arrow(ax, color='w')
        ax.set_xlabel('Y [kpc] (toward $l=90^\\circ$)')
        ax.set_ylabel('X [kpc] (toward Galactic center)')
        ax.set_aspect('equal')
        ax.set_xlim(-3.3, 3.3)
        ax.set_ylim(-3.3, 3.3)
        ax.set_title(f'$\\log L\\geq{thr:g}$   (N={n})', fontsize=11)
    fig.suptitle('ALS II OB-star surface density, $d<3$ kpc', fontsize=12)
    fig.tight_layout()
    fig.savefig('fig_als_density.png', dpi=140)
    plt.close(fig)


def salpeter_logL_lf(logL_grid):
    """dN/dlogL for a Salpeter IMF (dN/dlogm ~ m^-1.35) mapped through ZAMS L(m)."""
    logm = np.interp(logL_grid, ZL, np.log10(ZM))
    dlogm_dlogL = np.gradient(logm, logL_grid)
    return 10 ** (-1.35 * logm) * dlogm_dlogL


def lf_fig(logL):
    bins = np.arange(3.0, 6.05, 0.1)
    cen = 0.5 * (bins[1:] + bins[:-1])
    h, _ = np.histogram(logL, bins=bins)
    fig, ax = plt.subplots(figsize=(7.5, 5))
    ax.step(cen, h, where='mid', lw=1.6, color='C3', label=f'ALS II <3 kpc (N={len(logL)})')
    # Salpeter locked to the lowest populated bin
    sal = salpeter_logL_lf(cen)
    i0 = np.argmax(h > 0)
    sal *= h[i0] / sal[i0]
    ax.plot(cen, sal, 'k--', lw=1.5, label='Salpeter IMF (via ZAMS $L(m)$), locked to lowest bin')
    # KDE + bootstrap band
    grid = np.linspace(3.0, 6.0, 300)
    kde = gaussian_kde(logL, bw_method=0.12)
    scale = len(logL) * 0.1     # convert density to counts per 0.1-dex bin
    ax.plot(grid, kde(grid) * scale, 'C0-', lw=2, label='KDE (bw=0.12 dex)')
    rng = np.random.default_rng(7)
    boots = np.array([gaussian_kde(rng.choice(logL, len(logL)), bw_method=0.12)(grid)
                      for _ in range(120)]) * scale
    lo, hi = np.percentile(boots, [16, 84], axis=0)
    ax.fill_between(grid, lo, hi, color='C0', alpha=0.25, label='KDE 68% bootstrap')
    ax.set_yscale('log')
    ax.set_ylim(0.5, 2 * h.max())
    ax.set_xlabel(r'$\log L$ [L$_\odot$] (from spectral type)')
    ax.set_ylabel('N per 0.1 dex')
    ax.legend(fontsize=8)
    ax.set_title('ALS II OB luminosity function, $d<3$ kpc', fontsize=11)
    fig.tight_layout()
    fig.savefig('fig_als_lf.png', dpi=140)
    plt.close(fig)


if __name__ == '__main__':
    x, y, logL = load_als()
    print(f'N={len(logL)}; logL>=4: {(logL>=4).sum()}, logL>=5: {(logL>=5).sum()}')
    topdown(x, y, logL)
    topdown_hiL(x, y, logL, thresh=5.0)
    density_maps(x, y, logL, threshs=(4.0, 5.0))
    lf_fig(logL)
    print('wrote fig_als_topdown.png, fig_als_topdown_hiL.png, '
          'fig_als_density.png, fig_als_lf.png')