#!/usr/bin/env python
"""Circumstellar COLD-dust mass (or upper limit) within a 0.1 pc (20,000 au)
aperture around each revealed ALS O star within 3 kpc.

Method: submillimeter continuum traces cold dust. We match each star against the
ATLASGAL 870um compact-source catalogue (Urquhart+2018 CSC; 19.2" beam ~ 0.1 pc
at d>~1 kpc, so a 0.1-pc reservoir would appear as a compact source). A matched
clump flux is converted to a gas mass at T=20 K (conservative -- lower T gives
higher mass) and treated as an UPPER LIMIT on truly circumstellar gas, since the
clump is almost certainly foreground/associated cloud rather than a bound
reservoir. Where the star lies in the ATLASGAL footprint with NO catalogued
source, the 5-sigma point-source sensitivity sets the mass upper limit.

Opacity: standard ATLASGAL/Ossenkopf-Henning value, per gram of GAS,
  kappa_gas(nu) = 0.0185 (nu/345GHz)^1.8 cm^2/g   (dust kappa 1.85, G2D=100)
so M_gas = F_nu d^2 / (kappa_gas(nu) B_nu(T)). (This differs from the fiducial
Hildebrand kappa used in the earlier beam-photometry figure, which over-states
gas masses ~100x; the standard opacity here gives literature-comparable,
tighter limits.)

Output: circ_cold.ecsv (one row per O star).
"""
import re
import pickle
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u

# ---- constants (cgs) -------------------------------------------------------
H, KB, CC = 6.626e-27, 1.381e-16, 2.998e10
MSUN, PC, JY = 1.989e33, 3.086e18, 1e-23
KAPPA0_GAS, NU0, BETA = 0.0185, 345e9, 1.8      # cm^2/g gas at 870um, per g gas
TCOLD = 20.0                                     # K, conservative
APER_PC = 0.1                                    # aperture radius-equivalent scale

# ATLASGAL: Urquhart+2018 CSC, typical rms ~50-70 mJy/beam; 5-sigma point limit.
ATLAS_FLIM_JY = 0.30
ATLAS_BEAM_AS = 19.2

SPT_L = {'O3': 5.8, 'O4': 5.68, 'O5': 5.55, 'O5.5': 5.5, 'O6': 5.4, 'O6.5': 5.32,
         'O7': 5.23, 'O7.5': 5.14, 'O8': 5.05, 'O8.5': 4.96, 'O9': 4.87, 'O9.5': 4.77}
BOOST = {'I': 0.7, 'II': 0.45, 'III': 0.25, 'IV': 0.12, 'V': 0.0}


def bnu(nu, T):
    return 2 * H * nu ** 3 / CC ** 2 / np.expm1(H * nu / (KB * T))


def kappa_gas(nu):
    return KAPPA0_GAS * (nu / NU0) ** BETA


def mgas(F_jy, lam_um, d_kpc, T=TCOLD):
    """gas mass [Msun] from flux (Jy) at lambda_um, distance kpc, temp T."""
    nu = CC / (lam_um * 1e-4)
    d = float(d_kpc) * 1000 * PC
    return float(F_jy) * JY * d ** 2 / (kappa_gas(nu) * bnu(nu, T)) / MSUN


def aper_as(d_kpc):
    return 206265.0 * APER_PC / (d_kpc * 1000.0)


# ---- O-star sample (identical loader to build_ostar_fir.py) ----------------
als = pickle.load(open('als_combined.pkl', 'rb'))
ra, dec, d, logL, name = [], [], [], [], []
for row in als:
    dd = row['DistMode']
    if not np.isfinite(dd) or dd <= 0 or dd > 3000:
        continue
    sp = str(row['SpType'])
    m = re.match(r'^O\s*(\d+(?:\.5)?)', sp)
    if not m or ('O' + m.group(1)) not in SPT_L:
        continue
    mc = re.search(r'(I{1,3}|IV|V)\b', sp)
    logL.append(SPT_L['O' + m.group(1)] + (BOOST.get(mc.group(1), 0.0) if mc else 0.0))
    c = SkyCoord(row['RAJ2000'], row['DEJ2000'], unit=(u.hourangle, u.deg))
    ra.append(c.ra.deg); dec.append(c.dec.deg); d.append(dd / 1000.)
    name.append(str(row.get('Name', '') if hasattr(row, 'get') else ''))
star = SkyCoord(ra, dec, unit=u.deg)
gl, gb = star.galactic.l.deg, star.galactic.b.deg
d = np.array(d); logL = np.array(logL)
print(f'O-star sample within 3 kpc: N={len(d)}')

# ---- ATLASGAL footprint + local catalogue match ---------------------------
atl = Table.read('atlasgal_csc_raw.fits')
atl_c = SkyCoord(np.array(atl['_RA'], float), np.array(atl['_DE'], float), unit=u.deg)
atl_fint = np.array(atl['Fint'], float)            # Jy, integrated 870um
# ATLASGAL survey footprint: |l|<60 (i.e. l<60 or l>300) and |b|<1.5
gl180 = np.where(gl > 180, gl - 360, gl)
in_atlas = (np.abs(gl180) < 60.0) & (np.abs(gb) < 1.5)

idx, sep2d, _ = star.match_to_catalog_sky(atl_c)
sep_as = sep2d.arcsec
# match radius: aperture size, but at least half the ATLASGAL beam
match_rad = np.maximum(aper_as(d), ATLAS_BEAM_AS / 2.0)
matched = sep_as < match_rad

rows = []
for i in range(len(d)):
    apas = aper_as(d[i])
    if matched[i]:
        F = atl_fint[idx[i]]
        M = mgas(F, 870., d[i])
        kind, note = 'detection', f'ATLASGAL clump {sep_as[i]:.1f}" away, Fint={F:.2f}Jy'
        flim = F
    elif in_atlas[i]:
        M = mgas(ATLAS_FLIM_JY, 870., d[i])
        kind, note = 'limit', f'ATLASGAL 5sigma ({ATLAS_FLIM_JY}Jy), no source'
        flim = ATLAS_FLIM_JY
    else:
        M = np.nan
        kind, note = 'nocov', 'outside ATLASGAL footprint'
        flim = np.nan
    rows.append(dict(id=i, gl=gl[i], gb=gb[i], d=float(d[i]), logL=float(logL[i]),
                     aper_as=apas, in_atlas=bool(in_atlas[i]),
                     Mcold_Msun=M, F870_Jy=flim, cold_kind=kind, cold_note=note))

t = Table(rows)
t.write('circ_cold.ecsv', overwrite=True)

det = t['cold_kind'] == 'detection'
lim = t['cold_kind'] == 'limit'
nocov = t['cold_kind'] == 'nocov'
print(f'ATLASGAL detections (clump within aperture): {det.sum()}')
print(f'ATLASGAL non-detections (5sigma limit):      {lim.sum()}')
print(f'outside ATLASGAL footprint:                  {nocov.sum()}')
mcov = t['Mcold_Msun'][~nocov]
print(f'median cold-gas limit (in footprint): {np.nanmedian(mcov):.3f} Msun')
print(f'  of L>1e5 stars: median {np.nanmedian(t["Mcold_Msun"][(~nocov)&(t["logL"]>5)]):.3f} Msun')
print('wrote circ_cold.ecsv')