#!/usr/bin/env python
"""Quantify how the P&G2026 conclusion (tau0=4.2 Myr) changes when each RMS source
(=cluster primary) is accompanied by its Kroupa-ODF M>8 Msun siblings.

Optimal sampling places exactly N(>m) = Xi(m..mmax)/Xi(m1..mmax) stars above mass m
when the primary is m1 (k = 1/Xi(m1..mmax) ensures one star above m1). So the number
of M>8 OB stars implied by a primary of mass m1 is, in closed form,
    N_OB8(m1) = Integral_8^mmax xi dm / Integral_{m1}^mmax xi dm,
and N_sib = N_OB8 - 1.  (Validated below against the imf-package _opt_sample.)
"""
import sys
sys.path.insert(0, '/orange/adamginsburg/repos/imf')
import numpy as np
from scipy.integrate import quad
from scipy.interpolate import interp1d
import imf.imf as I

MMIN, MMAX, M8 = 0.03, 150.0, 8.0
_KR = I.get_massfunc('kroupa', mmin=MMIN, mmax=MMAX)


def xi_int(a, b):
    return quad(_KR, a, b, limit=100)[0]


def N_OB8(m1):
    """Optimal-sampling # of M>8 stars (incl. primary) given primary mass m1."""
    m1 = np.atleast_1d(m1).astype(float)
    out = np.where(m1 >= M8, [xi_int(M8, MMAX) / xi_int(mm, MMAX) for mm in m1], 0.0)
    return out


# ---------- validation against imf _opt_sample (small clusters; fast) ----------
class NumKroupa:
    def __init__(self, mmin=MMIN):
        self.mmin, self.mmax = mmin, MMAX
        self._kk = I.get_massfunc('kroupa', mmin=mmin, mmax=MMAX)
    def __call__(self, m): return self._kk(m)
    def integrate(self, a, b, **k): return quad(self._kk, a, b, limit=100)
    def m_integrate(self, a, b, **k): return quad(lambda m: m * self._kk(m), a, b, limit=100)


def validate():
    mfc = NumKroupa()
    print("validation: imf _opt_sample vs closed form (real Kroupa)")
    for Mcl in [40, 80]:
        masses, _ = I._opt_sample(Mcl, mfc, tolerance=0.0)
        masses = np.sort(np.atleast_1d(masses))[::-1]
        m1 = masses[0]
        n_samp = int((masses >= M8).sum())
        n_form = float(N_OB8(m1)[0])
        print(f"  Mcl={Mcl:5.0f}  m1={m1:6.1f}  N(>8) sampled={n_samp:3d}  closed-form={n_form:5.1f}")


# ---------- ZAMS mass<->luminosity (to map RMS L_bol to primary mass) ----------
# tabulated ZAMS (solar metallicity, e.g. Ekstrom+2012 / standard): m [Msun], log10 L/Lsun
_ZM = np.array([2, 3, 5, 7, 9, 12, 15, 20, 25, 32, 40, 60, 85, 120, 150])
_ZL = np.array([1.26, 2.05, 2.74, 3.26, 3.66, 4.07, 4.37, 4.74, 5.0,
                5.28, 5.48, 5.83, 6.08, 6.31, 6.43])
_logL_of_m = interp1d(np.log10(_ZM), _ZL, kind='linear', fill_value='extrapolate')
_m_of_logL = interp1d(_ZL, np.log10(_ZM), kind='linear', fill_value='extrapolate')


def mass_from_L(logL):
    return 10 ** _m_of_logL(np.clip(logL, _ZL.min(), _ZL.max()))


if __name__ == '__main__':
    # N_sib(m1) table
    print("\nM>8 siblings per primary (Kroupa ODF, closed form):")
    for m1 in [8, 10, 12, 16, 20, 25, 32, 40, 60, 85, 120]:
        print(f"  m1={m1:6.1f} Msun  N(M>8)={float(N_OB8(m1)[0]):6.2f}  N_sib={float(N_OB8(m1)[0])-1:6.2f}")

    # apply to the RMS embedded sample
    from astropy.table import Table
    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'])
    good = sf & np.isfinite(L) & (L > 0)
    logL = np.log10(L[good])
    # P&G LF range: 2.9 < logL < 6.6
    sel = (logL > 2.9) & (logL < 6.6)
    logL = logL[sel]
    m1 = mass_from_L(logL)
    nOB8 = N_OB8(m1)
    nsib = np.clip(nOB8 - 1, 0, None)
    prim_is_OB = m1 >= M8

    print(f"\nRMS embedded SF sample in LF range: {sel.sum()} sources")
    print(f"  primaries with m1>=8 Msun: {prim_is_OB.sum()} "
          f"({100*prim_is_OB.mean():.0f}%)")
    print(f"  median primary mass: {np.median(m1):.1f} Msun")
    # population mean siblings, over M>8 primaries (those that become OB stars)
    mean_sib = nsib[prim_is_OB].mean()
    mean_NOB = nOB8[prim_is_OB].mean()
    print(f"  <N_sib (M>8)> per OB primary = {mean_sib:.2f}")
    print(f"  <N_OB8 per OB primary>       = {mean_NOB:.2f}")
    print(f"\n  => total M>8 OB stars / primaries = {mean_NOB:.2f}")
    print(f"  => leading-order tau0 correction factor (1+<N_sib>) = {1+mean_sib:.2f}")
    print(f"     tau0: 4.2 Myr -> {4.2*(1+mean_sib):.1f} Myr  (leading order)")
    np.savez('multiplicity_out.npz', m1=m1, nsib=nsib, nOB8=nOB8, logL=logL)

    validate()   # slow sanity check, last so the result above always prints