#!/usr/bin/env python
"""Referee item 4 (Tan): apply the P&G Appendix-C correction recipe to the
embedded (RMS) and OB (ALS) catalogues and test whether the corrected
high-luminosity ratio recovers the volume-complete value.

Recipe as published:
  V(R) = 2 h pi R^2 (h = 39 pc), phi(L;R) from all sources with d <= R,
  g(R) = reference-range density trend normalized to g(3 kpc) = 1,
  phi_g(L;R) = phi(L;R)/g(R),
  C(L) = max_R phi_g(L;R) / phi_g(L;R0),  phi_final = C * phi_g(L;R0).
Reference ranges: embedded log L in [3.5,4.5]; OB log L in [4.0,4.6]
(well-populated, 'weakly incomplete'). We work in L for both (their OB LF is
in G_abs; the test is structural, not photometric).
"""
import re
import pickle
import numpy as np
from astropy.table import Table

H = 0.039  # kpc


def V(R):
    return 2 * H * np.pi * R ** 2


# ---- embedded (RMS, d<=18) --------------------------------------------------
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']])
ok = np.isin(typ, ['YSO', 'HII region', 'HII/YSO']) & np.isfinite(L) & (L > 0) \
    & np.isfinite(D) & (D > 0) & (D <= 18)
eL, eD = np.log10(L[ok]), D[ok]

# ---- OB (ALS, d<=12) --------------------------------------------------------
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}
sL, sD = [], []
for row in als:
    d = row['DistMode']
    if not np.isfinite(d) or d <= 0 or d > 12000:
        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.)
sL = np.array(sL); sD = np.array(sD)

Rgrid = np.array([2, 3, 4, 5, 6, 8, 10, 12, 15, 18], float)


def g_of_R(Darr, Larr, lref, R0ref=3.0, Rmax=18.):
    """reference-range cumulative density normalized at 3 kpc"""
    sel = (Larr >= lref[0]) & (Larr < lref[1])
    def dens(R):
        return np.sum(sel & (Darr <= R)) / V(R)
    d3 = dens(R0ref)
    return {R: dens(R) / d3 for R in Rgrid if R <= Rmax}


ge = g_of_R(eD, eL, (3.5, 4.5), Rmax=18)
go = g_of_R(sD, sL, (4.0, 4.6), Rmax=12)

print('P&G recipe applied, high-L (log L > 5) densities [kpc^-3]:')
print(f'{"R":>4} {"phi_e":>9} {"g_e":>6} {"phi_ge":>9} | '
      f'{"phi_o":>9} {"g_o":>6} {"phi_go":>9} | {"ratio_g":>8}')
pge, pgo = {}, {}
for R in Rgrid:
    fe = np.sum((eL >= 5) & (eD <= R)) / V(R)
    line = f'{R:4.0f} {fe:9.3f} {ge.get(R, np.nan):6.2f} '
    pe = fe / ge[R] if R in ge else np.nan
    pge[R] = pe
    line += f'{pe:9.3f} | '
    if R <= 12:
        fo = np.sum((sL >= 5) & (sD <= R)) / V(R)
        po = fo / go[R]
        pgo[R] = po
        line += f'{fo:9.2f} {go[R]:6.2f} {po:9.2f} | {pe/po:8.3f}'
    print(line)

# their completeness step
R0e, R0o = 18., 12.
Ce = max(pge.values()) / pge[R0e]
Co = max(v for v in pgo.values()) / pgo[R0o]
phi_e_final = Ce * pge[R0e]
phi_o_final = Co * pgo[R0o]
print(f'\nC_emb = {Ce:.2f} (peak at R='
      f'{[R for R,v in pge.items() if v==max(pge.values())][0]:.0f}), '
      f'C_OB = {Co:.2f} (peak at R='
      f'{[R for R,v in pgo.items() if v==max(pgo.values())][0]:.0f})')
print(f'final corrected densities: emb {phi_e_final:.3f}, OB {phi_o_final:.2f} '
      f'kpc^-3;  CORRECTED RATIO = {phi_e_final/phi_o_final:.3f}')
print(f'volume-complete (d<3) ratio = {np.sum((eL>=5)&(eD<=3))} local... '
      f'using census: 9/527 = {9/527:.3f}')
print(f'P&G best-fit model needs 0.185')