#!/usr/bin/env python
"""Top-down disk map overlaying ALS existing O/early-B stars + our forming sample."""
import re, csv
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from astropy.table import Table

t = Table.read('als_tablea1.fits')
d = np.array(t['DistMode'], float) / 1000.0
l = np.radians(np.array(t['GLON'], float)); b = np.radians(np.array(t['GLAT'], float))
ax_ = d * np.cos(b) * np.sin(l); ay = d * np.cos(b) * np.cos(l)
def parse(s):
    m = re.match(r'\s*([OB])\s*(\d(?:\.\d)?)?', str(s))
    return (m.group(1), float(m.group(2)) if (m and m.group(2)) else np.nan) if m else ('x', np.nan)
cls = np.array([parse(s)[0] for s in t['SpType']]); sub = np.array([parse(s)[1] for s in t['SpType']])
O8plus = (cls == 'O') & np.isfinite(sub) & (sub <= 8)
O9 = (cls == 'O') & np.isfinite(sub) & (sub >= 8.5)
earlyB = (cls == 'B') & np.isfinite(sub) & (sub < 4)

def fnum(x, dd=np.nan):
    try: return float(x)
    except (ValueError, TypeError): return dd
fr = []
for o in csv.DictReader(open('master_distances.csv')):
    fr.append((float(o['GLON']), float(o['GLAT']), float(o['d_best']), float(o['Lbol_best']), o['Name']))
for o in csv.DictReader(open('curated_nearby_sample.csv')):
    if 'REVEALED-EXCLUDE' in str(o.get('note', '')):
        continue
    fr.append((float(o['GLON']), float(o['GLAT']), float(o['Dist_kpc']), fnum(o['Lbol']), o['CommonName']))
seen = set(); FR = []
for r in fr:
    if r[4] in seen: continue
    seen.add(r[4]); FR.append(r)
fl = np.radians([r[0] for r in FR]); fb = np.radians([r[1] for r in FR])
fd = np.array([r[2] for r in FR]); fL = np.array([r[3] for r in FR])
fx = fd * np.cos(fb) * np.sin(fl); fy = fd * np.cos(fb) * np.cos(fl)

fig, axes = plt.subplots(1, 2, figsize=(15, 7.7))
for ax, cut in zip(axes, (2.0, 3.0)):
    th = np.linspace(0, 2 * np.pi, 200)
    for rr in (1, 2, 3):
        ax.plot(rr * np.cos(th), rr * np.sin(th), color='0.85', lw=0.7, zorder=0)
    in_ = d < cut
    ax.scatter(ax_[earlyB & in_], ay[earlyB & in_], s=4, c='0.6', alpha=0.45,
               lw=0, label=f'ALS early-B ({(earlyB & in_).sum()})', zorder=1)
    ax.scatter(ax_[O9 & in_], ay[O9 & in_], s=9, c='#3b78b0', alpha=0.6, lw=0,
               label=f'ALS O9 ({(O9 & in_).sum()})', zorder=2)
    ax.scatter(ax_[O8plus & in_], ay[O8plus & in_], s=16, c='#1b3a66',
               label=f'ALS O8+ existing ({(O8plus & in_).sum()})', zorder=3)
    fin = (fd < cut)
    big = fin & (fL >= 1e5); mid = fin & (fL >= 1e4) & (fL < 1e5)
    ax.scatter(fx[mid], fy[mid], s=40, marker='^', facecolor='none', edgecolor='#d1495b',
               lw=1.0, label=f'forming 1e4-1e5 ({mid.sum()})', zorder=4)
    ax.scatter(fx[big], fy[big], s=130, marker='*', c='#d1495b', edgecolor='k', lw=0.5,
               label=f'forming O8+ L>1e5 ({big.sum()})', zorder=5)
    ax.plot(0, 0, marker='+', color='k', ms=10, zorder=6)
    ax.text(0, -0.05 * cut, 'Sun', ha='center', va='top', fontsize=8)
    ax.set_xlim(-cut * 1.15, cut * 1.15); ax.set_ylim(-cut * 1.15, cut * 1.15)
    ax.set_aspect('equal'); ax.set_xlabel('X (kpc)'); ax.set_ylabel('Y (kpc) [+ = GC]')
    ax.set_title(f'd < {cut:.0f} kpc')
    ax.legend(loc='upper right', fontsize=7, framealpha=0.9)
fig.suptitle('Existing massive stars (ALS, Gaia) vs forming sample (UC HII + L>1e4 YSO)', fontsize=13)
fig.tight_layout(rect=[0, 0, 1, 0.96])
fig.savefig('als_forming_overlay.png', dpi=140)
print('wrote als_forming_overlay.png')