#!/usr/bin/env python
"""Labeled clone of the top-down ALS OB-star surface-density map
(fig_als_density.png): same KDE density in stars/kpc^2 for logL>=4 and >=5,
d<3 kpc, with the major nearby star-forming complexes / OB associations marked
at their (l, d). Confirms the density peaks are the known complexes.

Frame (matching make_als_figs.py): X = d cos l (toward GC), Y = d sin l
(toward l=90); plotted as (Y horizontal, X vertical), Sun at origin.
Output: fig_als_density_labeled.png
"""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
from circ_mass_ob_common import load_ob

# nearby complexes: (name, l_deg, d_kpc, text-offset dx, dy in kpc)
COMPLEXES = [
    ('Cygnus X',   80.0, 1.4,  0.55,  0.15),
    ('Carina',    287.5, 2.3, -0.10, -0.35),
    ('M17',        15.0, 2.0,  0.45,  0.45),
    ('NGC 6357',  353.2, 1.7, -0.55,  0.75),
    ('NGC 6334',  351.2, 1.3, -0.85,  0.10),
    ('Sco OB1',   343.5, 1.6, -0.75, -0.45),
    ('M8/M20',      6.0, 1.3,  0.60, -0.35),
    ('Gem OB1',   189.0, 1.8, -0.55,  0.15),
    ('Orion',     209.0, 0.4, -0.30, -0.30),
    ('Vela',      263.0, 0.9, -0.35, -0.20),
    ('Rosette',   207.0, 1.4, -0.60, -0.05),
    ('W3/4/5',    134.5, 2.0, -0.45,  0.25),
    ('Cep OB2',   102.5, 0.9,  0.45,  0.20),
    ('Per OB1',   135.0, 2.3, -0.20, -0.40),
]


def gc_arrow(ax, color='w', L=1.1):
    ax.annotate('', xy=(0, L), xytext=(0, 0),
                arrowprops=dict(arrowstyle='-|>', color=color, lw=2.0, alpha=0.9,
                                shrinkA=0, shrinkB=0), zorder=7)
    ax.text(0.13, L, 'G.C.', color=color, fontsize=8.5, va='center', ha='left',
            fontweight='bold', zorder=7)


def mark_complexes(ax):
    for name, l, d, dx, dy in COMPLEXES:
        lr = np.radians(l)
        X, Y = d * np.cos(lr), d * np.sin(lr)          # X toward GC, Y toward l=90
        ax.plot(Y, X, 'o', mfc='none', mec='cyan', mew=1.4, ms=11, zorder=8)
        ax.annotate(name, (Y, X), xytext=(Y + dx, X + dy), color='cyan',
                    fontsize=8.5, fontweight='bold', zorder=9,
                    ha='center', va='center',
                    arrowprops=dict(arrowstyle='-', color='cyan', lw=0.7, alpha=0.7))


ra, dec, gl, gb, d, logL, isO = load_ob()
lr = np.radians(gl)
x = d * np.cos(lr)      # toward GC
y = d * np.sin(lr)      # toward l=90

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, 2, figsize=(13.2, 6.8))
for ax, thr in zip(axes, (4.0, 5.0)):
    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
    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=3, color='w', alpha=0.25, 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.5)
    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)
    mark_complexes(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.set_title(f'$\\log L\\geq{thr:g}$   (N={n})', fontsize=11)
fig.suptitle('ALS II OB-star surface density with SF complexes labelled, '
             '$d<3$ kpc', fontsize=12)
fig.tight_layout()
fig.savefig('fig_als_density_labeled.png', dpi=140)
print('wrote fig_als_density_labeled.png')