#!/usr/bin/env python
"""All-sky DIFFUSE cold-dust backstop within 0.1 pc around each revealed ALS O
star within 3 kpc, from the SFD (Schlegel+1998) reddening map.

Where ATLASGAL does not cover (outer Galaxy l=60-300, or |b|>1.5), we still bound
the cold gas in the aperture from the SFD total-sightline reddening. This is a
CONSERVATIVE upper limit twice over: (1) SFD integrates the ENTIRE line of sight,
attributing all foreground+background dust to the 0.1 pc aperture; (2) it assumes
that whole column lies within the aperture. For revealed (low-A_V) stars the
sightline reddening is modest, so the implied aperture mass is small.

  N_H = 5.8e21 E(B-V)  cm^-2          (Bohlin+1978)
  M   = N_H * mu * m_H * pi (0.05 pc)^2   (mass in a 0.1-pc-wide sightline tube)

Output: circ_diffuse.ecsv
"""
import re
import pickle
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
from dustmaps.sfd import SFDQuery

MH, MSUN, PC = 1.6726e-24, 1.989e33, 3.086e18
MU = 1.4
NH_PER_EBV = 5.8e21
APER_D_PC = 0.1

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}

als = pickle.load(open('als_combined.pkl', 'rb'))
ra, dec, d, logL = [], [], [], []
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.)
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)

ebv = SFDQuery()(star)
NH = NH_PER_EBV * np.asarray(ebv, float)
A_aper = np.pi * (0.5 * APER_D_PC * PC) ** 2
Mdiff = NH * MU * MH * A_aper / MSUN

t = Table(dict(id=np.arange(len(d)), gl=gl, gb=gb, d=d, logL=logL,
               EBV_sfd=np.asarray(ebv, float), Mdiffuse_Msun=Mdiff))
t.write('circ_diffuse.ecsv', overwrite=True)
print(f'SFD diffuse backstop, N={len(d)}')
print(f'median M_diffuse: {np.nanmedian(Mdiff):.3f} Msun (all-sightline, conservative)')
print(f'|b|>1 median:  {np.nanmedian(Mdiff[np.abs(gb)>1]):.3f} Msun')
print(f'L>1e5 median:  {np.nanmedian(Mdiff[logL>5]):.3f} Msun')
print('wrote circ_diffuse.ecsv')