#!/usr/bin/env python
"""Volume-limited sample of forming L>1e5 Lsun sources (true O stars) within 6 kpc.
RMS is flux-complete for L>2e4 Galaxy-wide (Davies+2011), so L>1e5 selection is
completeness-safe; binding issues are distances (fix with BeSSeL parallaxes),
contaminants (drop evolved types), the |l|<10 GC exclusion and RMS data defects
(fill from curated knowledge).

Steps: RMS L-selection at RMS distance with buffer -> BeSSeL Reid+2019 crossmatch
(position <3', |dvLSR|<10 km/s) -> adopt parallax distance + rescale L -> final cut
L>1e5, d<6 -> merge curated |l|<10 / defect sources -> ostar_forming_lt6kpc.{csv,md}."""
import csv
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
from astroquery.vizier import Vizier
import warnings
warnings.filterwarnings('ignore')

GOOD_TYPES = {'HII region', 'YSO', 'HII/YSO'}

rms = Table.read('rms_table1_raw.fits')
c_rms = SkyCoord(rms['RAJ2000'], rms['DEJ2000'], unit=(u.hourangle, u.deg))
L = np.array(rms['Lbol'], float)
D = np.array(rms['Dist'], float)
sel = (np.isfinite(L) & np.isfinite(D) &
       np.array([str(t) in GOOD_TYPES for t in rms['Type']]))
# buffer: keep anything that could scatter into (L>1e5, d<6) after distance fixes
sel &= (L * (6.5 / np.maximum(D, 0.1)) ** 2 > 1e5) & (D < 9.0)
cand = rms[sel]
c_cand = c_rms[sel]
print(f'RMS candidates (buffered): {len(cand)}')

v = Vizier(row_limit=-1)
bes = v.get_catalogs('J/ApJ/885/131')[0]
c_bes = SkyCoord(bes['RAJ2000'], bes['DEJ2000'], unit=(u.hourangle, u.deg))
plx = np.array(bes['plx'], float)
vb = np.array(bes['VLSR'], float)

# our extinction/literature distance revisions take precedence over raw RMS kinematics
master = {m['Name']: m for m in csv.DictReader(open('master_distances.csv'))}

rows = []
for i, r in enumerate(cand):
    sep = c_cand[i].separation(c_bes).arcmin
    dv = np.abs(float(r['vLSR']) - vb) if np.isfinite(float(r['vLSR'] or np.nan)) else np.full(len(bes), 99.)
    m = (sep < 3.0) & (dv < 10)
    d_rms = float(r['Dist'])
    nm = str(r['Name'])
    # priority: (1) our master revisions (incl. A/B-suffixed rows), (2) BeSSeL, (3) RMS
    mrow = master.get(nm) or master.get(nm + 'A') or master.get(nm + 'B')
    if mrow:
        d = float(mrow['d_best'])
        meth, ref = mrow['method'], mrow['ref']
    elif m.any():
        j = int(np.argmin(np.where(m, sep, np.inf)))
        d = 1.0 / plx[j]
        meth, ref = 'maser-parallax', f"BeSSeL:{bes['Name'][j]}"
    else:
        d, meth, ref = d_rms, 'kinematic-RMS', 'Lumsden2013/Urquhart2018'
    Lnew = float(r['Lbol']) * (d / d_rms) ** 2
    rows.append(dict(Name=str(r['Name']), Type=str(r['Type']),
                     RA=c_cand[i].ra.deg, Dec=c_cand[i].dec.deg,
                     glon=round(c_cand[i].galactic.l.deg, 4),
                     glat=round(c_cand[i].galactic.b.deg, 4),
                     vLSR=float(r['vLSR']) if np.isfinite(float(r['vLSR'] or np.nan)) else np.nan,
                     d_kpc=round(d, 2), method=meth, ref=ref,
                     L=round(Lnew, -2), origin='RMS'))

final = [r for r in rows if r['L'] > 1e5 and r['d_kpc'] < 6.0]
npar = sum(r['method'] == 'maser-parallax' for r in final)
print(f'RMS-side final: {len(final)} (parallax-anchored {npar})')

# curated |l|<10 & defect additions (region hosts forming/embedded O-star equivalent)
CURATED = [
    dict(Name='RCW38-IRS2', Type='HII/YSO', glon=267.94, glat=-1.07, d_kpc=1.70,
         method='Gaia-cluster', ref='Kuhn2019/Zucker2020', L=7.0e5, origin='curated',
         note='RMS blank-params defect'),
    dict(Name='NGC3576-S3M4', Type='HII/YSO', glon=291.27, glat=-0.70, d_kpc=2.40,
         method='kinematic+Gaia', ref='Persi1994/Andre2008', L=1.0e5, origin='curated',
         note='RMS corrupt-dist defect'),
    dict(Name='NGC6334-region', Type='HII/YSO', glon=351.42, glat=0.65, d_kpc=1.30,
         method='maser-parallax', ref='Wu2014/Chibueze2014', L=1.5e5, origin='curated',
         note='|l|<10; region L; Povich2026 Gaia 1.67 tension'),
    dict(Name='NGC6357-G353.2', Type='HII', glon=353.17, glat=0.90, d_kpc=1.70,
         method='maser-parallax', ref='Reid2014/Kuhn2019', L=3.5e5, origin='curated',
         note='|l|<10; Pis24 partially revealed - flag'),
    dict(Name='W31-G10.6', Type='HII', glon=10.62, glat=-0.38, d_kpc=4.95,
         method='maser-parallax', ref='Sanna2014', L=5.0e5, origin='curated',
         note='in RMS at l=10.6 but check inclusion'),
    dict(Name='M8-Hourglass', Type='HII', glon=5.97, glat=-1.17, d_kpc=1.34,
         method='Gaia-cluster', ref='Damiani2019', L=1.5e5, origin='curated',
         note='|l|<10; Her36 revealed O7.5 - flag as revealed'),
    dict(Name='W43-Main-complex', Type='HII/YSO', glon=30.75, glat=-0.06, d_kpc=5.49,
         method='maser-parallax', ref='Zhang2014', L=3.0e6, origin='curated',
         note='COMPLEX-LEVEL: RMS decomposes into pieces each <1e5 (2.8-6.9e4); '
              'W43-MM1 protocluster certainly forming O stars; scale-definition flag'),
]
names = {r['Name'] for r in final}
for cr in CURATED:
    if cr['Name'] not in names and cr['d_kpc'] < 6.0 and cr['L'] > 1e5:
        cr.setdefault('vLSR', np.nan)
        final.append(cr)

final.sort(key=lambda r: r['glon'])
fields = ['Name', 'Type', 'glon', 'glat', 'd_kpc', 'method', 'ref', 'L', 'origin', 'note']
with open('ostar_forming_lt6kpc.csv', 'w', newline='') as fh:
    w = csv.DictWriter(fh, fieldnames=fields, extrasaction='ignore')
    w.writeheader()
    for r in final:
        r.setdefault('note', '')
        w.writerow(r)

with open('ostar_forming_lt6kpc.md', 'w') as fh:
    fh.write('# Volume-limited forming O-star sample: L>1e5 Lsun, d<6 kpc\n\n')
    fh.write(f'**N = {len(final)}** ({sum(r["method"]=="maser-parallax" for r in final)} '
             'parallax-anchored). RMS flux-complete >2e4 Lsun Galaxy-wide (Davies+2011); '
             'contaminant types removed; |l|<10 + defect regions restored from curated '
             'knowledge.\n\n')
    fh.write('| # | Name | l | b | d (kpc) | L (1e5 Lsun) | type | dist | note |\n')
    fh.write('|---|------|---|---|---------|--------------|------|------|------|\n')
    for i, r in enumerate(final, 1):
        fh.write(f"| {i} | {r['Name']} | {r['glon']:.2f} | {r['glat']:+.2f} | "
                 f"{r['d_kpc']:.2f} | {r['L']/1e5:.1f} | {r['Type']} | "
                 f"{r['method']}/{r['ref']} | {r.get('note','')} |\n")
print(f'final N={len(final)} -> ostar_forming_lt6kpc.csv/md')
for r in final:
    print(f"  {r['Name']:24s} l={r['glon']:7.2f} d={r['d_kpc']:5.2f} L={r['L']:.1e} "
          f"{r['method']:16s} {r.get('note','')[:40]}")