#!/usr/bin/env python
"""Write the <2 kpc and <3 kpc samples (UC HII + L>1e4 Lsun YSO + marginal) to
short markdown tables: sample_lt2kpc.md, sample_lt3kpc.md."""
import csv
import numpy as np

def fnum(x, d=np.nan):
    try:
        return float(x)
    except (ValueError, TypeError):
        return d

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='', ref=o['ref'], method=o['method']))
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 = fnum(e['d_best']); dr = fnum(c['Dist_RMS']); Lr = fnum(c['Lbol_RMS'])
    if not (np.isfinite(db) and np.isfinite(dr) and dr > 0):
        continue
    rows.append(dict(Name=nm, l=float(c['GLON']), b=float(c['GLAT']), d=db,
                     L=Lr * (db / dr) ** 2, type=c['Type'], origin='RMS', flag='',
                     ref=e['ref'], method=e['method']))
for o in csv.DictReader(open('curated_nearby_sample.csv')):
    if 'REVEALED-EXCLUDE' in str(o.get('note', '')):
        continue                                  # revealed OB star, not forming
    fl = 'marginal' if 'marginal' in str(o.get('note', '')).lower() else ''
    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=fl, ref=o['Dist_ref'], method=o['Dist_method']))

seen = set(); uniq = []
for r in rows:
    if r['Name'] in seen:
        continue
    seen.add(r['Name']); uniq.append(r)
rows = uniq

def sample(cut):
    s = [r for r in rows if r['d'] < cut and
         ((np.isfinite(r['L']) and r['L'] >= 1e4) or r['flag'] == 'marginal')]
    return sorted(s, key=lambda x: (-x['d'] if False else x['l']))

def is_hii(t): return 'HII' in t

for cut, fn in [(2.0, 'sample_lt2kpc.md'), (3.0, 'sample_lt3kpc.md')]:
    s = sample(cut)
    nh = sum(is_hii(r['type']) for r in s); ny = len(s) - nh
    nm = sum(r['flag'] == 'marginal' for r in s)
    with open(fn, 'w') as f:
        f.write(f"# (U)C HII + L>10$^4$ L$_\\odot$ YSO sample — d < {cut:.0f} kpc\n\n")
        f.write(f"**N = {len(s)}**  (HII ~{nh}, YSO ~{ny}" +
                (f", incl. {nm} marginal" if nm else "") + ")  \n")
        f.write("Best literature distances; L rescaled to best distance. "
                "Sorted by Galactic longitude.  \n")
        f.write("Sources: RMS (Lumsden+2013) + ext literature search + curated non-RMS "
                "(Orion/NGC6334/etc, off-plane or |l|<10 GC-excluded). "
                "Off-plane & on-plane IRAS/Akari completeness checked.\n\n")
        f.write("| # | Name | l | b | d (kpc) | L (L$_\\odot$) | type | dist ref | note |\n")
        f.write("|---|------|------|------|------|------|------|------|------|\n")
        for i, r in enumerate(s, 1):
            note = []
            if r['flag'] == 'marginal':
                note.append('**marginal**')
            if r['origin'] == 'curated':
                note.append('curated')
            Ls = f"{r['L']:.2e}" if np.isfinite(r['L']) else '--'
            f.write(f"| {i} | {r['Name']} | {r['l']:.2f} | {r['b']:+.2f} | {r['d']:.2f} | "
                    f"{Ls} | {r['type']} | {r['method']}/{r['ref']} | {', '.join(note)} |\n")
    print(f"wrote {fn}: N={len(s)} (HII {nh}, YSO {ny}, marginal {nm})")