#!/usr/bin/env python
"""Referee item 1 (Mottram): quantitative completeness of the nearby embedded
census.

Tests:
 A. Compact (r<0.15 pc at catalogue distance) WISE-K HII regions within 3 kpc:
    are they in our census (positional match)? Flag implausible catalogue
    distances.
 B. CORNISH-South UC HII regions (radio-selected, independent of RMS mid-IR
    selection): what fraction have an RMS counterpart? -> direct measurement of
    RMS UC-HII incompleteness, distance-blind.
 C. Same for the broader CORNISH-S 'HII Region' class.
 D. Census members inside the CORNISH-S footprint: fraction with a CORNISH
    counterpart (sanity check in the reverse direction).
Honest caveat: CORNISH (2.5" interferometric) filters LARGE angular-size
sources, so this measures completeness for angularly compact (mostly distant)
UC HII; the truly nearby (d<1 kpc, >20") regime is covered by the
famous-nebula audit and test A.
"""
import csv
import re
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u

# ---- our census positions (l, b, d, L) --------------------------------------
def fnum(x, d=np.nan):
    try:
        return float(x)
    except (TypeError, ValueError):
        return d


cen_l, cen_b, cen_d, cen_L, cen_n = [], [], [], [], []
for o in csv.DictReader(open('master_distances.csv')):
    cen_l.append(fnum(o['GLON'])); cen_b.append(fnum(o['GLAT']))
    cen_d.append(fnum(o['d_best'])); cen_L.append(fnum(o['Lbol_best']))
    cen_n.append(o['Name'])
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
    c = cand[nm]
    db, dr, Lr = fnum(e['d_best']), fnum(c['Dist_RMS']), fnum(c['Lbol_RMS'])
    if np.isfinite(db) and np.isfinite(dr) and dr > 0:
        cen_l.append(fnum(c['GLON'])); cen_b.append(fnum(c['GLAT']))
        cen_d.append(db); cen_L.append(Lr * (db / dr) ** 2); cen_n.append(nm)
for o in csv.DictReader(open('curated_nearby_sample.csv')):
    note = str(o.get('note', ''))
    if 'REVEALED-EXCLUDE' in note:
        continue
    cen_l.append(fnum(o['GLON'])); cen_b.append(fnum(o['GLAT']))
    cen_d.append(fnum(o['Dist_kpc'])); cen_L.append(fnum(o['Lbol']))
    cen_n.append(o['CommonName'])
seen = set(); keep = []
for i, nm in enumerate(cen_n):
    if nm in seen:
        continue
    seen.add(nm); keep.append(i)
cen = SkyCoord(np.array(cen_l)[keep], np.array(cen_b)[keep], unit=u.deg,
               frame='galactic')
cend = np.array(cen_d)[keep]
in3 = cend < 3.0
print(f'census positions: {len(cen)} unique ({in3.sum()} at d<3)')

# ---- full RMS positions (for RMS-recovery test) -----------------------------
rms = Table.read('rms_table1_raw.fits')
rmsc = SkyCoord(rms['RAJ2000'], rms['DEJ2000'],
                unit=(u.hourangle, u.deg)).galactic

# ---- A. compact WISE-K regions ---------------------------------------------
NAME_RE = re.compile(r'G(\d+\.\d+)([+-]\d+\.\d+)')
wl, wb, wname, wd, wr, wf = [], [], [], [], [], []
for r in csv.DictReader(open('wise22_photometry.csv')):
    if r['set'] != 'wise':
        continue
    m = NAME_RE.match(r['name'])
    if not m:
        continue
    wl.append(float(m.group(1))); wb.append(float(m.group(2)))
    wname.append(r['name']); wd.append(float(r['d']))
    wr.append(float(r['rad_as'])); wf.append(fnum(r['f22rel']))
wc = SkyCoord(np.array(wl), np.array(wb), unit=u.deg, frame='galactic')
wd = np.array(wd); wr = np.array(wr)
rpc = wr / 206265 * wd * 1000
compact = rpc <= 0.15
print(f'\n=== A. compact WISE-K (r<0.15 pc at cat. distance, d<3): '
      f'{compact.sum()} of {len(wc)} ===')
idx, sep, _ = wc[compact].match_to_catalog_sky(cen)
idxr, sepr, _ = wc[compact].match_to_catalog_sky(rmsc)
for i, (j, s, jr, sr) in enumerate(zip(idx, sep.arcsec, idxr, sepr.arcsec)):
    nm = np.array(wname)[compact][i]
    dd = wd[compact][i]
    flag = 'SUSPECT-DIST' if dd < 0.6 else ''
    print(f'  {nm} d={dd:.2f} {flag}: census match {np.array(cen_n)[keep][j]} '
          f'@{s:.0f}\" | RMS match @{sr:.0f}\"')

# ---- B/C. CORNISH-S recovery by RMS ----------------------------------------
cs = Table.read('cornish_south_catalog.fits')
styp = np.array([str(x).strip() for x in cs['SourceType']])
csc = SkyCoord(np.array(cs['GLON'], float), np.array(cs['GLAT'], float),
               unit=u.deg, frame='galactic')
for cls in ('UCHII Region', 'HII Region'):
    sel = styp == cls
    sub = csc[sel]
    idx, sep, _ = sub.match_to_catalog_sky(rmsc)
    for rad in (12., 20.):
        rec = (sep.arcsec < rad).sum()
        print(f'\n=== {cls}: N={sel.sum()}; RMS counterpart within {rad:.0f}\": '
              f'{rec} ({100*rec/sel.sum():.0f}%) ===' if rad == 12 else
              f'    within {rad:.0f}\": {rec} ({100*rec/sel.sum():.0f}%)')
    # of the unmatched, how many are near OUR census members anyway (complexes)
    un = sep.arcsec >= 20
    idx2, sep2, _ = sub[un].match_to_catalog_sky(cen)
    near = (sep2.arcsec < 120).sum()
    print(f'    unmatched-by-RMS but within 2\' of a census member: {near}')

# ---- D. census members in CORNISH-S footprint -------------------------------
glc = cen.l.deg
infoot = in3 & (glc > 295) & (glc < 350) & (np.abs(cen.b.deg) < 0.5)
if infoot.sum():
    idx, sep, _ = cen[infoot].match_to_catalog_sky(csc)
    print(f'\n=== D. census d<3 members in CORNISH-S footprint: {infoot.sum()}; '
          f'with CORNISH source within 20\": {(sep.arcsec<20).sum()} ===')
    for nm, s, j in zip(np.array(cen_n)[keep][infoot], sep.arcsec, idx):
        print(f'  {nm}: nearest CORNISH {styp[j]} @{s:.0f}\"')