#!/usr/bin/env python
"""Mass-resolved OB-LF boost from ODF siblings, and its effect on the P&G slope
difference (which sets the growth exponent alpha) and normalization (tau0).

Baseline: each RMS primary -> 1 OB star.  OB mass function ~ phi_prim(m).
Modified: each primary of mass m1 also seeds siblings at mass m<m1 with optimal-
sampling density  k(m1)*xi(m) = xi(m)/Xi(m1,mmax).  So at OB mass m:
    siblings(m) = INT_{m}^{mmax} phi_prim(m1) * xi(m)/Xi(m1,mmax) dm1
    B(m) = 1 + siblings(m)/phi_prim(m)          (boost factor of the OB LF)
phi_prim(m) is the observed RMS primary mass function (from L_bol -> mass).
"""
import sys
sys.path.insert(0, '/orange/adamginsburg/repos/imf')
import numpy as np
from scipy.integrate import quad
from astropy.table import Table
from multiplicity_correction import _KR, mass_from_L, MMAX, M8

# ---- observed RMS primary masses ----
t = Table.read('../rms_table1_raw.fits')
L = np.array(t['Lbol'], float)
typ = np.array([str(x) for x in t['Type']])
sf = np.isin(typ, ['HII region', 'YSO', 'HII/YSO'])
logL = np.log10(L[sf & np.isfinite(L) & (L > 0)])
logL = logL[(logL > 2.9) & (logL < 6.6)]
m1samp = mass_from_L(logL)
m1samp = m1samp[m1samp >= M8]                  # OB-producing primaries

# primary mass function phi_prim(m) on a log grid (smooth via histogram in log m)
mgrid = np.logspace(np.log10(8), np.log10(150), 40)
lm_edges = np.linspace(np.log10(8), np.log10(150), 31)
lm_cen = 0.5 * (lm_edges[1:] + lm_edges[:-1])
mcen = 10 ** lm_cen
hist, _ = np.histogram(np.log10(m1samp), bins=lm_edges)
dlogm = lm_edges[1] - lm_edges[0]
phi_prim = hist / dlogm                         # dN/dlog m of primaries (counts)


# precompute xi CDF on a fine grid for O(1) interval integrals Xi(a,b)=cum(b)-cum(a)
_mfine = np.logspace(np.log10(M8), np.log10(MMAX), 4000)
_xifine = np.array([_KR(m) for m in _mfine])
_cum = np.concatenate([[0.0], np.cumsum(0.5 * (_xifine[1:] + _xifine[:-1]) * np.diff(_mfine))])


def cumXi(m):
    return np.interp(m, _mfine, _cum)


def Xi(a, b):
    return cumXi(b) - cumXi(a)


# k_i = 1/Xi(m1_i, mmax): optimal-sampling prefactor so each primary has 1 star above it
k_prim = 1.0 / Xi(m1samp, MMAX)


def siblings_per_bin(m_lo, m_hi):
    """Sum over observed primaries of their optimal-sampling siblings in [m_lo,m_hi)."""
    top = np.minimum(m_hi, m1samp)              # siblings only below each primary
    contrib = np.where(m1samp > m_lo, k_prim * (cumXi(top) - cumXi(m_lo)), 0.0)
    return contrib.sum()


def slope(mc, phi):
    ok = phi > 0
    if ok.sum() < 3:
        return np.nan
    return np.polyfit(np.log10(mc[ok]), np.log10(phi[ok]), 1)[0]


if __name__ == '__main__':
    # per-bin counts (ratio is what matters; uniform log bins => slope unaffected)
    prim_cnt = hist.astype(float)
    sib = np.array([siblings_per_bin(10**lm_edges[j], 10**lm_edges[j + 1])
                    for j in range(len(lm_cen))])
    phi_prim = prim_cnt
    B = np.where(prim_cnt > 0, 1 + sib / np.where(prim_cnt > 0, prim_cnt, np.nan), np.nan)
    phi_OB_mod = prim_cnt + sib

    print("mass-resolved OB-LF sibling boost B(m) = phi_OB_mod / phi_prim:")
    print(f"{'m':>7} {'phi_prim':>9} {'siblings':>9} {'B(m)':>6}")
    for m, pp, s, b in zip(mcen, phi_prim, sib, B):
        if pp > 0:
            print(f"{m:7.1f} {pp:9.1f} {s:9.1f} {b:6.2f}")

    # slope change of the OB mass function (proxy for beta_OB change)
    sl_base = slope(mcen, phi_prim)
    sl_mod = slope(mcen, phi_OB_mod)
    print(f"\nOB mass-function slope (dlog phi/dlog m):")
    print(f"  baseline (primaries=OB) : {sl_base:+.2f}")
    print(f"  modified (+ODF siblings): {sl_mod:+.2f}")
    print(f"  steepening d(slope)     : {sl_mod - sl_base:+.2f}")
    # P&G: beta_OB=(1-s-delta)/gamma, beta_HII=(1+alpha-s)/gamma; Dbeta=(alpha+delta)/gamma
    # OB MF slope ~ (1-s-delta) here (in mass space, per dlogm, before /gamma & L mapping).
    # A steeper OB LF (more negative) at fixed embedded LF increases Dbeta -> larger alpha.
    gamma_hi = 1.8
    dbeta_shift = -(sl_mod - sl_base) / gamma_hi  # Dbeta increases as OB LF steepens
    print(f"\n  implied increase in Dbeta=(alpha+delta)/gamma : {dbeta_shift:+.2f}  (gamma={gamma_hi})")
    print(f"  => increase in alpha ~ gamma*dDbeta = {gamma_hi*dbeta_shift:+.2f}")
    print(f"     (baseline alpha=0.5 -> ~{0.5 + gamma_hi*dbeta_shift:.2f}); tau0 rises with alpha along the P&G degeneracy")
    np.savez('boost_out.npz', mcen=mcen, phi_prim=phi_prim, sib=sib, B=B)