#!/usr/bin/env python
"""P&G2026 forward model in mass/L space, model-vs-model differential for the
1 O -> O + Kroupa-ODF B-star multiplicity.

Baseline (P&G): embedded LF and OB LF both drawn from the field IMF xi, 1:1.
  phi_OB(m)  = xi(m) t_MS(m)                         [final stars on MS]
  phi_emb(m) = INT_m^mmax xi(mf) t_form(mf)/mf dmf   [accreting through mass m]
  t_form(mf) = tau0 (mf/60)^alpha
Modified: embedded sources are cluster PRIMARIES psi = xi/B; OB = primaries+siblings = xi.
  phi_OB,mod  = xi t_MS                (unchanged: ALS sees all members)
  phi_emb,mod = INT_m^mmax (xi(mf)/B(mf)) t_form/mf dmf   (embedded traces primaries only)
B(m) = 1 + siblings(m)/primaries(m), Kroupa optimal sampling (closed form).
'Data' = baseline at P&G best-fit (m_k=18,s=2.45,s_f=3.86,alpha=0.48,tau0=4.15).
Refit (alpha,tau0) of the modified model to the data -> the shift.
"""
import numpy as np
from scipy.integrate import quad
from scipy.interpolate import interp1d
from mist_relations import t_MS
from multiplicity_correction import _KR, MMAX as MMAXK   # Kroupa pdf, cluster mmax

# --- P&G best-fit field IMF (broken power law) ---
MK, S, SF = 18.0, 2.45, 3.86
MMIN_F, MMAX_F = 1.0, 150.0
A0, T0 = 0.48, 4.15          # truth (Myr); m0=60
M0 = 60.0


def xi(m):
    m = np.asarray(m, float)
    lo = m ** (-S)
    hi = MK ** (SF - S) * m ** (-SF)
    return np.where(m < MK, lo, hi)


# --- Kroupa-ODF boost B(m): siblings of primaries (proxy primaries = xi) ---
_mfine = np.logspace(np.log10(8), np.log10(MMAXK), 4000)
_xik = np.array([_KR(m) for m in _mfine])
_cumk = np.concatenate([[0], np.cumsum(0.5 * (_xik[1:] + _xik[:-1]) * np.diff(_mfine))])
def cumXiK(m): return np.interp(m, _mfine, _cumk)
def XiK(a, b): return cumXiK(b) - cumXiK(a)
M8 = 8.0


def B_of_m(mgrid, prim):
    """boost B(m)=1+sib/prim on mgrid given a primary dN/dlogm 'prim' (>=8)."""
    B = np.ones_like(mgrid)
    # primary dN/dlogm -> treat as weights; sibling density at m from primaries m1>m:
    # sib(m) = sum_{m1>m} prim(m1) * [xiK(m)/XiK(m1,mmax)] * (per dlogm consistent)
    for i, m in enumerate(mgrid):
        if m < M8:
            continue
        sel = mgrid > m
        if not sel.any():
            continue
        kk = 1.0 / XiK(mgrid[sel], MMAXK)             # prefactor per primary mass
        # sibling dN/dlogm at m = xiK(m)*m*ln10 * sum_j prim_j * kk_j * dlogm
        dlogm = np.log(10) * (np.log10(mgrid[1]) - np.log10(mgrid[0]))
        sibdens = (_KR(m) * m * np.log(10)) * np.sum(prim[sel] * kk) * dlogm
        if prim[i] > 0:
            B[i] = 1 + sibdens / prim[i]
    return B


# --- LFs on a mass grid ---
MG = np.logspace(np.log10(8), np.log10(120), 60)
DLOGM = np.log(10) * (np.log10(MG[1]) - np.log10(MG[0]))


def phi_OB(mg):
    return xi(mg) * t_MS(mg) * mg * np.log(10)        # dN/dlogm


def phi_emb(mg, alpha, tau0, weight=None):
    """dN/dlogm of embedded (accreting) sources; weight(mf) multiplies xi (for psi=xi/B)."""
    out = np.zeros_like(mg)
    mf = np.logspace(np.log10(8), np.log10(MMAX_F), 400)
    w = np.ones_like(mf) if weight is None else weight(mf)
    tform = tau0 * (mf / M0) ** alpha
    integrand = xi(mf) * w * tform / mf               # per dmf
    for i, m in enumerate(mg):
        sel = mf >= m
        out[i] = np.trapezoid(integrand[sel], mf[sel]) * m * np.log(10)
    return out


def fit(data, model_emb_fn, agrid, tgrid):
    """Joint-A chi2 scan over (alpha,tau0). data=[OB..,emb..]; OB fixed."""
    dOB = phi_OB(MG)
    sig = np.sqrt(np.maximum(np.concatenate([dOB, data]), 1e-30))   # Poisson-like
    d = np.concatenate([dOB, data])
    best = (np.inf, None, None)
    chi = np.full((len(agrid), len(tgrid)), np.nan)
    for ia, a in enumerate(agrid):
        for it, tt in enumerate(tgrid):
            m = np.concatenate([dOB, model_emb_fn(a, tt)])
            A = np.sum(m * d / sig ** 2) / np.sum(m ** 2 / sig ** 2)
            c = np.sum((A * m - d) ** 2 / sig ** 2)
            chi[ia, it] = c
            if c < best[0]:
                best = (c, a, tt)
    return best, chi


if __name__ == '__main__':
    # data = baseline at truth
    emb_data = phi_emb(MG, A0, T0)
    OB_data = phi_OB(MG)

    agrid = np.linspace(0.2, 1.6, 29)
    tgrid = np.linspace(1.0, 60.0, 60)

    # --- baseline refit (validation: should recover 0.48, 4.15) ---
    b0, _ = fit(emb_data, lambda a, t: phi_emb(MG, a, t), agrid, tgrid)
    print(f"BASELINE refit:  alpha={b0[1]:.2f}  tau0={b0[2]:.1f} Myr  (truth 0.48, 4.15)  chi2={b0[0]:.1e}")

    # --- B(m) from ODF (primaries proxy = xi on MS) ---
    prim = phi_OB(MG)                          # primary dN/dlogm proxy
    B = B_of_m(MG, prim)
    Bint = interp1d(np.log10(MG), B, fill_value='extrapolate')
    def inv_B_weight(mf):
        return 1.0 / np.clip(Bint(np.log10(mf)), 1.0, None)
    print(f"B(m): min={B.min():.1f} max={B.max():.1f} median={np.median(B):.1f} "
          f"@16Msun={float(Bint(np.log10(16))):.1f} @60={float(Bint(np.log10(60))):.1f}")

    # --- modified refit: embedded traces primaries (xi/B) ---
    bm, chi = fit(emb_data, lambda a, t: phi_emb(MG, a, t, weight=inv_B_weight), agrid, tgrid)
    print(f"MODIFIED refit:  alpha={bm[1]:.2f}  tau0={bm[2]:.1f} Myr   chi2={bm[0]:.1e}")
    print(f"\n  => multiplicity shift: tau0 {T0:.1f} -> {bm[2]:.1f} Myr  (x{bm[2]/T0:.1f})")
    print(f"     alpha {A0:.2f} -> {bm[1]:.2f}  (d_alpha={bm[1]-A0:+.2f})")
    np.savez('forward_out.npz', MG=MG, B=B, emb_data=emb_data, OB_data=OB_data,
             baseline=b0[1:], modified=bm[1:])