import pickle, re
import numpy as np
from astroquery.vizier import Vizier
import time

Vizier.ROW_LIMIT = -1
def get(name):
    for a in range(6):
        try: return Vizier.get_catalogs(name)[0]
        except Exception as e:
            print("retry",a,e); time.sleep(8)
    raise RuntimeError(name)

# Use tablea1 = main clean sample (good astrometry). tablea2 = excluded (bad astrometry/colours)
t1 = get("J/MNRAS/504/2968/tablea1")
t2 = get("J/MNRAS/504/2968/tablea2")

def parse_class(spt):
    """Return ('O'|'B'|'other', subnum or None) from leading component of an OB SpType string."""
    if spt is None: return ('none', None)
    s = str(spt).strip()
    if s == '' or s == '--' or s.lower()=='nan': return ('none', None)
    # take first comma-separated component
    first = s.split(',')[0].strip()
    # strip leading non-letter junk
    m = re.match(r'^\(?([OB])\s*([0-9]+(?:\.[0-9]+)?)?', first)
    if not m:
        # try anywhere a leading O/B with number for whole string fallback
        m2 = re.match(r'^[^OB]*([OB])\s*([0-9]+(?:\.[0-9]+)?)?', first)
        if not m2: return ('other', None)
        m = m2
    letter = m.group(1)
    num = float(m.group(2)) if m.group(2) is not None else None
    return (letter, num)

def analyze(t, label):
    dist = np.array(t['DistMode'], dtype=float)  # pc
    plx  = np.array(t['plx'], dtype=float)
    spt  = t['SpType']
    n = len(t)
    cls = []; sub = []
    for s in spt:
        c, num = parse_class(s)
        cls.append(c); sub.append(num)
    cls = np.array(cls); sub = np.array([np.nan if x is None else x for x in sub], dtype=float)

    has_dist = np.isfinite(dist)
    print(f"\n===== {label}: N={n} =====")
    print(f"  rows with finite DistMode: {has_dist.sum()}  (missing: {(~has_dist).sum()})")
    neg_plx = np.isfinite(plx) & (plx<=0)
    print(f"  rows with negative/zero Gaia plx: {neg_plx.sum()}")
    print(f"  SpType: O={ (cls=='O').sum() }  B={ (cls=='B').sum() }  none/empty={ (cls=='none').sum() }  other(non-OB)={ (cls=='other').sum() }")

    for dmax_kpc in (2.0, 3.0):
        dmax = dmax_kpc*1000.0
        within = has_dist & (dist < dmax)
        isO = within & (cls=='O')
        isB = within & (cls=='B')
        # early-B: B0 - B3 (subnum <=3, or B with no number we treat separately)
        earlyB = isB & (sub <= 3.0)
        earlyB_strict2 = isB & (sub <= 2.0)
        lateB  = isB & (sub >= 4.0)
        B_nonum = isB & ~np.isfinite(sub)
        print(f"  d < {dmax_kpc} kpc:  O={isO.sum():5d}   B(all)={isB.sum():5d}   earlyB(B0-B3)={earlyB.sum():5d}   earlyB(B0-B2)={earlyB_strict2.sum():5d}   lateB(B4-B9)={lateB.sum():5d}   B_nonum={B_nonum.sum():4d}")
    return cls, sub, dist

analyze(t1, "tablea1 (MAIN, clean astrometry)")
analyze(t2, "tablea2 (EXCLUDED: bad astrometry/colours)")
print("\nDONE_ANALYZE")