#!/usr/bin/env python
"""Extend the embedded-vs-revealed + ALMA-coverage check to the full <3 kpc sample.
Per source: nearest 2MASS PSC star (Ks, sep) = 'revealed' indicator; ALMA bands.
Flag candidate-revealed = a BRIGHT NIR point (Ks below the revealed-OB threshold at
that distance) coincident (<2") with the source -> like M8/M20 (already-formed O star).
Embedded forming sources are NIR-faint/blank at mm/MSX position."""
import csv, warnings
import numpy as np
warnings.filterwarnings('ignore')
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
from astroquery.vizier import Vizier

# ---- build <3 kpc sample (same logic as write_sample_lists) ----
def fnum(x):
    try: return float(x)
    except (ValueError, TypeError): return np.nan
rows = []
for o in csv.DictReader(open('master_distances.csv')):
    rows.append(dict(Name=o['Name'], l=float(o['GLON']), b=float(o['GLAT']),
                     d=float(o['d_best']), L=float(o['Lbol_best']), type=o['Type'],
                     origin='RMS', flag=''))
ext = {o['Name']: o for o in csv.DictReader(open('distances_ext.csv'))}
cand = {o['Name']: o for o in csv.DictReader(open('candidate_distances_ext.csv'))}
for nm, e in ext.items():
    if nm not in cand: continue
    db, dr, Lr = fnum(e['d_best']), fnum(cand[nm]['Dist_RMS']), fnum(cand[nm]['Lbol_RMS'])
    if not (np.isfinite(db) and np.isfinite(dr) and dr > 0): continue
    rows.append(dict(Name=nm, l=float(cand[nm]['GLON']), b=float(cand[nm]['GLAT']),
                     d=db, L=Lr*(db/dr)**2, type=cand[nm]['Type'], origin='RMS', flag=''))
for o in csv.DictReader(open('curated_nearby_sample.csv')):
    rev = 'REVEALED-EXCLUDE' in str(o.get('note', ''))
    rows.append(dict(Name=o['CommonName'], l=float(o['GLON']), b=float(o['GLAT']),
                     d=float(o['Dist_kpc']), L=fnum(o['Lbol']), type=o['Type'],
                     origin='curated', flag='revealed-excl' if rev else 'curated'))
seen = set(); uniq = []
for r in rows:
    if r['Name'] in seen: continue
    seen.add(r['Name']); uniq.append(r)
# <3kpc members: L>=1e4 (or marginal/revealed-excl kept for reference) AND d<3
samp = [r for r in uniq if r['d'] < 3 and (r['L'] >= 1e4 or r['flag'] in ('revealed-excl',)
        or 'marginal' in r['Name'].lower())]
print(f'<3 kpc sources to check: {len(samp)}')

# need RA/Dec: RMS from rms table; curated from their l,b
rms = Table.read('rms_table1_raw.fits'); rn = [str(x) for x in rms['Name']]
def coord(r):
    if r['origin'] == 'RMS' and r['Name'] in rn:
        i = rn.index(r['Name'])
        return SkyCoord(rms['RAJ2000'][i], rms['DEJ2000'][i], unit=(u.hourangle, u.deg))
    return SkyCoord(r['l'], r['b'], unit='deg', frame='galactic').icrs
coords = SkyCoord([coord(r) for r in samp])

# ---- 2MASS PSC ----
Vizier.ROW_LIMIT = 5
v = Vizier(columns=['RAJ2000', 'DEJ2000', 'Jmag', 'Kmag'], catalog='II/246/out')
res = v.query_region(coords, radius=4*u.arcsec)
tm = res[0] if res else None
# match nearest per source
def nir(c):
    if tm is None: return (np.nan, np.nan)
    mc = SkyCoord(tm['RAJ2000'], tm['DEJ2000'], unit='deg')
    sep = c.separation(mc).arcsec
    if sep.min() > 4: return (np.nan, np.nan)
    k = int(np.argmin(sep)); return (float(np.ma.filled(tm['Kmag'][k], np.nan)), sep[k])

# ---- ALMA index ----
ai = Table.read('alma_footprint_index.fits')
aic = SkyCoord(np.array(ai['crval1'], float), np.array(ai['crval2'], float), unit='deg')
ah = np.minimum(np.array(ai['naxis1'], float), np.array(ai['naxis2'], float))*np.array(ai['cdelt_as'], float)/2
af = np.array(ai['freq_ghz'], float)
def band(f):
    for lo, hi, b in [(84,116,3),(125,163,4),(163,211,5),(211,275,6),(275,373,7),(385,500,8),(602,720,9),(787,950,10)]:
        if lo <= f <= hi: return b
    return 0
def alma(c):
    cov = c.separation(aic).arcsec < ah
    bs = sorted(set(band(af[j]) for j in np.where(cov)[0]) - {0})
    return bs

# revealed threshold: apparent Ks of a B0V (M_Ks~-2.5) at distance d, no extinction (bright limit)
def ks_thresh(d):
    return -2.5 + 5*np.log10(d*1000) - 5 + 3   # +3 mag slack for Av/early-B
out = []
print(f"\n{'Name':22s}{'type':10s}{'d':>4s}{'L':>9s}{'Ks':>6s}{'sep':>5s}{'ALMA':>10s}  class")
nrev = nemb = nalma = 0
for r, c in zip(samp, coords):
    ks, sep = nir(c); bs = alma(c)
    if bs: nalma += 1
    revealed = np.isfinite(ks) and sep < 2.0 and ks < ks_thresh(r['d'])
    cls = 'curated-REVEALED' if r['flag'] == 'revealed-excl' else \
          ('cand-REVEALED' if revealed else 'embedded')
    if 'REVEALED' in cls: nrev += 1
    else: nemb += 1
    out.append((r['Name'], r['type'], r['d'], r['L'], ks, sep, bs, cls))
    bstr = ','.join(f'B{b}' for b in bs) or '-'
    kstr = f"{ks:.1f}" if np.isfinite(ks) else '--'
    sstr = f"{sep:.1f}" if np.isfinite(sep) else '--'
    if 'REVEALED' in cls or r['L'] > 5e4:   # print revealed + luminous
        print(f"{r['Name']:22s}{r['type'][:9]:10s}{r['d']:4.1f}{r['L']:9.1e}{kstr:>6s}{sstr:>5s}{bstr:>10s}  {cls}")
print(f"\nTOTAL <3kpc: {len(samp)} | embedded: {nemb} | revealed(cand+curated): {nrev} | with ALMA: {nalma}")
# save
T = Table(rows=[(o[0], o[1], o[2], o[3], o[4], o[5], ','.join(f'B{b}' for b in o[6]), o[7]) for o in out],
          names=['Name', 'type', 'd_kpc', 'Lbol', 'Ks', 'sep_as', 'ALMA_bands', 'class'])
T.write('embedded_check_3kpc.csv', overwrite=True)
print('wrote embedded_check_3kpc.csv')