#!/usr/bin/env python
"""Referee item 2 (Padoan): is the shell-ratio rise a real Galactic gradient or
selection? Decomposition: compute the raw embedded/OB ratio vs heliocentric
distance RESTRICTED to a fixed galactocentric annulus (7.5 < R_gal < 9.5 kpc,
the solar annulus). At fixed R_gal the ambient environment (surface density,
arm mix in the azimuthal average) cannot produce a heliocentric-distance trend;
any residual rise is selection. R0 = 8.2 kpc.
"""
import re
import pickle
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u

R0 = 8.2


def rgal(d, glon_deg):
    gl = np.radians(glon_deg)
    return np.sqrt(R0 ** 2 + d ** 2 - 2 * R0 * d * np.cos(gl))


# RMS embedded
t = Table.read('rms_table1_raw.fits')
L = np.array(t['Lbol'], float)
D = np.array(t['Dist'], float)
typ = np.array([str(x) for x in t['Type']])
emb = np.isin(typ, ['YSO', 'HII region', 'HII/YSO'])
g = SkyCoord(t['RAJ2000'], t['DEJ2000'], unit=(u.hourangle, u.deg)).galactic
Rg_e = rgal(D, g.l.deg)
sele = emb & np.isfinite(L) & (L >= 1e4) & np.isfinite(D) & (D > 0)

# ALS OB
als = pickle.load(open('als_combined.pkl', 'rb'))
SPT = {'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,
       'B0': 4.6, 'B0.5': 4.4, 'B1': 4.15, 'B1.5': 3.95, 'B2': 3.75, 'B2.5': 3.6,
       'B3': 3.4}
BOOST = {'I': 0.7, 'II': 0.45, 'III': 0.25, 'IV': 0.12, 'V': 0.0}
sd, sl, sglon = [], [], []
c = SkyCoord(als['RAJ2000'], als['DEJ2000'], unit=(u.hourangle, u.deg)).galactic
for i, row in enumerate(als):
    d = row['DistMode']
    if not np.isfinite(d) or d <= 0:
        continue
    m = re.match(r'^([OB])\s*(\d+(?:\.5)?)', str(row['SpType']))
    if not m or (m.group(1) + m.group(2)) not in SPT:
        continue
    mc = re.search(r'(I{1,3}|IV|V)\b', str(row['SpType']))
    sl.append(SPT[m.group(1) + m.group(2)] +
              (BOOST.get(mc.group(1), 0.0) if mc else 0.0))
    sd.append(float(d) / 1000.)
    sglon.append(c.l.deg[i])
sd = np.array(sd); sl = np.array(sl); sglon = np.array(sglon)
Rg_s = rgal(sd, sglon)

shells = [(0, 2), (2, 3), (3, 4), (4, 6), (6, 9)]
print('raw embedded/OB ratio (L>1e4) vs heliocentric shell:')
print(f'{"shell":>10} {"all R_gal":>18} {"solar annulus 7.5-9.5":>24}')
for a, b in shells:
    e_all = int(np.sum(sele & (D > a) & (D <= b)))
    o_all = int(np.sum((sd > a) & (sd <= b) & (sl >= 4.0)))
    ann_e = sele & (D > a) & (D <= b) & (Rg_e > 7.5) & (Rg_e < 9.5)
    ann_o = (sd > a) & (sd <= b) & (sl >= 4.0) & (Rg_s > 7.5) & (Rg_s < 9.5)
    ea, oa = int(ann_e.sum()), int(ann_o.sum())
    r_all = e_all / o_all if o_all else np.inf
    r_ann = ea / oa if oa else np.inf
    print(f'{a}-{b} kpc: {e_all:4d}/{o_all:5d} = {r_all:6.3f}   '
          f'{ea:4d}/{oa:5d} = {r_ann:6.3f}')
print('\nIf the rise were a real galactocentric gradient, the solar-annulus')
print('column would be flat; a rise there is pure selection.')