#!/usr/bin/env python
"""Massive-star census + forming/existing timescale for the most massive stars.
(a) save ALS subset; (b) luminosity-matched O/B counts; the O8+ IMF+lifetime calc.

ALS II = Pantaleoni Gonzalez+2021 (2021MNRAS.504.2968P), J/MNRAS/504/2968/tablea1.
Forming sample = master_distances.csv + curated_nearby_sample.csv (UC HII + L>1e4 YSO).
"""
import os, re, csv
import numpy as np
from astropy.table import Table

# ---------- get ALS (cache to disk) ----------
ALS = 'als_tablea1.fits'
if os.path.exists(ALS):
    t = Table.read(ALS)
else:
    from astroquery.vizier import Vizier
    V = Vizier(columns=['ALS', 'SpType', 'DistMode', 'GLON', 'GLAT', 'RAJ2000', 'DEJ2000'],
               row_limit=-1)
    last = None
    for _ in range(4):
        try:
            t = V.get_catalogs('J/MNRAS/504/2968/tablea1')[0]
            break
        except Exception as e:                       # noqa: BLE001  (network retry)
            last = e; t = None
    if t is None:
        raise RuntimeError(f'ALS query failed: {last}')
    t.write(ALS, overwrite=True)
print('ALS tablea1 rows:', len(t))

d_kpc = np.array(t['DistMode'], float) / 1000.0      # pc -> kpc
spt = [str(s) for s in t['SpType']]

# parse leading O/B subtype
def parse(s):
    m = re.match(r'\s*([OB])\s*(\d(?:\.\d)?)?', s)
    if not m:
        return ('x', np.nan)
    return (m.group(1), float(m.group(2)) if m.group(2) else np.nan)
cls = np.array([parse(s)[0] for s in spt])
sub = np.array([parse(s)[1] for s in spt])

# subtype bins (dwarf ZAMS mass guide):
#   O3-O8 (>~22 Msun, L>~1e5)  | O9 (~18) | B0 (~15) | B1 (~11) | B2 (~9) | B3 (~7.5)
isO = cls == 'O'; isB = cls == 'B'
O8plus = isO & np.isfinite(sub) & (sub <= 8)          # O3..O8  = "O8+" most massive
O9     = isO & np.isfinite(sub) & (sub >= 8.5)         # O8.5/O9/O9.5
B0     = isB & (sub < 1)                                # B0, B0.5
B1     = isB & (sub >= 1) & (sub < 2)
B2     = isB & (sub >= 2) & (sub < 3)
B3     = isB & (sub >= 3) & (sub < 4)
earlyB = isB & np.isfinite(sub) & (sub < 4)            # B0-B3
allO   = isO

def cnt(mask, cut):
    return int(np.sum(mask & np.isfinite(d_kpc) & (d_kpc < cut)))

print('\n=== ALS counts by bin ===')
print(f"{'bin':18s}{'<2kpc':>8s}{'<3kpc':>8s}")
for nm, m in [('O8+ (O3-O8)', O8plus), ('O9 (O8.5-O9.5)', O9), ('all O', allO),
              ('B0', B0), ('B1', B1), ('B2', B2), ('B3', B3),
              ('early-B(B0-B3)', earlyB)]:
    print(f"{nm:18s}{cnt(m,2):>8d}{cnt(m,3):>8d}")

# (b) luminosity-matched existing massive = ZAMS L>1e4 ~ O + B0 + B1 (B1V logL~4.0)
lummatch = allO | B0 | B1
print(f"\n(b) Lum-matched existing (ZAMS L>1e4 ~ O+B0+B1):  <2kpc={cnt(lummatch,2)}  <3kpc={cnt(lummatch,3)}")
# B-clock bin for IMF calc: O9 + B0 + B1
Bclock = O9 | B0 | B1
NB2, NB3 = cnt(Bclock, 2), cnt(Bclock, 3)
print(f"B-clock (O9+B0+B1):                              <2kpc={NB2}  <3kpc={NB3}")
NO8_2, NO8_3 = cnt(O8plus, 2), cnt(O8plus, 3)
print(f"existing O8+ (O3-O8):                            <2kpc={NO8_2}  <3kpc={NO8_3}")

# ---------- forming sample (our catalog) ----------
def fnum(x, d=np.nan):
    try: return float(x)
    except (ValueError, TypeError): return d
frows = []
for o in csv.DictReader(open('master_distances.csv')):
    frows.append((float(o['GLON']), float(o['GLAT']), float(o['d_best']),
                  float(o['Lbol_best']), o['Type'], o['Name']))
for o in csv.DictReader(open('curated_nearby_sample.csv')):
    if 'REVEALED-EXCLUDE' in str(o.get('note', '')):
        continue
    frows.append((float(o['GLON']), float(o['GLAT']), float(o['Dist_kpc']),
                  fnum(o['Lbol']), o['Type'], o['CommonName']))
# dedup by name
seen = set(); fr = []
for r in frows:
    if r[5] in seen: continue
    seen.add(r[5]); fr.append(r)

def fcount(cut, Lmin, Lmax=np.inf):
    return [r for r in fr if r[2] < cut and np.isfinite(r[3]) and Lmin <= r[3] < Lmax]

print('\n=== FORMING sample (ours) by L bin ===')
for nm, lo, hi in [('L>1e5 (forming O8+)', 1e5, np.inf),
                   ('3e4-1e5 (forming O9-B0)', 3e4, 1e5),
                   ('1e4-3e4 (forming B0-B1)', 1e4, 3e4)]:
    n2 = len(fcount(2, lo, hi)); n3 = len(fcount(3, lo, hi))
    print(f"  {nm:26s} <2kpc={n2:3d}  <3kpc={n3:3d}")

print('\n*** Young (forming) O8+ (L>1e5) within 2 kpc: ***')
for r in sorted(fcount(2, 1e5), key=lambda x: -x[3]):
    print(f"    {r[5]:24s} l={r[0]:.2f} d={r[2]:.2f} L={r[3]:.2e} {r[4]}")
Nform_O8_2 = len(fcount(2, 1e5)); Nform_O8_3 = len(fcount(3, 1e5))

# ---------- IMF + lifetime calc ----------
# Salpeter dN/dM ~ M^-2.35; number ratio O8+(22-120) / Bclock(10.5-20)
def imf_int(a, b, g=2.35):
    return (a**(1-g) - b**(1-g)) / (g - 1)
ratio = imf_int(22, 120) / imf_int(10.5, 20)
# MS lifetimes (Schaller+1992 solar, count-weighted by IMF within each bin)
tau_B = 14.0      # O9(~8.7)+B0(~11.6)+B1(~18), IMF-weighted toward B1 -> ~14 Myr
tau_O8 = 5.0      # O8+ count-weighted toward O8/O7 (~23-28 Msun) -> ~5 Myr
print('\n=== IMF + lifetime calc (the most-massive-star timescale) ===')
print(f"Salpeter number ratio  N(O8+,22-120)/N(O9-B1,10.5-20) = {ratio:.3f}")
print(f"adopted MS lifetimes: tau_Bclock={tau_B} Myr, tau_O8+={tau_O8} Myr")
for cut, NB, NO8, Nf in [(2, NB2, NO8_2, Nform_O8_2), (3, NB3, NO8_3, Nform_O8_3)]:
    SFR_B = NB / tau_B                    # stars/Myr formed in B-clock bin (steady state)
    SFR_O8 = SFR_B * ratio               # predicted O8+ formation rate
    N_formed_over_tauB = NB * ratio      # O8+ formed during one B lifetime
    N_pred_exist = SFR_O8 * tau_O8       # predicted existing O8+ (MS)
    tau_form = Nf / SFR_O8 if SFR_O8 > 0 else np.nan   # forming timescale of O8+
    print(f"\n--- d < {cut} kpc ---")
    print(f"  N(B-clock O9+B0+B1) = {NB}")
    print(f"  O8+ formed over one tau_B (={tau_B}Myr) = NB*ratio = {N_formed_over_tauB:.0f}")
    print(f"  predicted EXISTING O8+ (SFR_O8 x tau_O8) = {N_pred_exist:.1f}   vs ALS observed = {NO8}")
    print(f"  forming O8+ (our sample L>1e5) = {Nf}")
    print(f"  => SFR(O8+) = {SFR_O8:.2f} /Myr")
    print(f"  => tau_form(O8+) = Nform/SFR_O8 = {tau_form:.2f} Myr")
    if NO8 > 0:
        print(f"  (direct route: tau_form = Nform/Nexist x tau_O8 = {Nf/NO8*tau_O8:.2f} Myr)")