"""Subgiant 'tracer' weights from MIST isochrones.

For a population formed at look-back time tau with metallicity [M/H], the
number of stars that sit inside a subgiant selection box today is

    N_box = psi(tau, [M/H]) * G(tau, [M/H]),   G = int_box xi(M) dM

where psi is the number of stars formed per unit time (per unit IMF
normalisation) and xi(M) ~ M^-2.3 is the IMF.  G is the IMF-weighted width of
the initial-mass interval that an isochrone of age tau spends inside the box.
Each observed subgiant therefore stands for I0/G Sun-like stars formed at the
same epoch, with I0 = int_0.8^1.2 xi dM.

Two boxes are implemented:
  * the Xiang & Rix (2022) box in (Teff, M_Ks), whose boundaries depend on
    [Fe/H] (their Extended Data Table 1), plus M_K > 0.5 and Teff < 6800 K;
  * a Gaia colour-magnitude box used for the volume-complete local sample.
"""
import numpy as np
from common import load_mist, IMF_SLOPE

# Xiang & Rix (2022) Extended Data Table 1:  [Fe/H], slope1, zpt1, slope2, zpt2
# box:  M_K > slope1*Teff + zpt1  (hot side of the RGB line)
#       M_K < slope2*Teff + zpt2  (bright side of the turn-off line)
XR_BOX = np.array([
    [0.4, -0.005, 26.00, -0.0014, 10.00],
    [0.2, -0.005, 26.25, -0.0014, 10.10],
    [0.0, -0.005, 26.50, -0.0014, 10.20],
    [-0.2, -0.0045, 24.25, -0.0014, 10.30],
    [-0.4, -0.004, 22.00, -0.0014, 10.50],
    [-0.6, -0.004, 22.20, -0.0014, 10.70],
    [-0.8, -0.004, 22.60, -0.00125, 9.90],
    [-1.0, -0.004, 23.00, -0.00125, 10.00],
    [-1.2, -0.004, 23.20, -0.001, 8.65],
    [-1.4, -0.004, 23.40, -0.001, 8.70],
    [-1.6, -0.004, 23.60, -0.001, 8.75],
    [-1.8, -0.004, 23.80, -0.001, 8.80],
    [-2.0, -0.004, 24.00, -0.001, 8.85],
    [-2.2, -0.004, 24.20, -0.001, 8.90],
    [-2.5, -0.004, 24.20, -0.001, 8.95]])
XR_BOX = XR_BOX[np.argsort(XR_BOX[:, 0])]


def xr_box_margins(teff, mk, feh):
    """Signed distances to the two box lines (d1 > 0 and d2 < 0 inside)."""
    f = np.clip(feh, -2.5, 0.4)
    s1 = np.interp(f, XR_BOX[:, 0], XR_BOX[:, 1])
    z1 = np.interp(f, XR_BOX[:, 0], XR_BOX[:, 2])
    s2 = np.interp(f, XR_BOX[:, 0], XR_BOX[:, 3])
    z2 = np.interp(f, XR_BOX[:, 0], XR_BOX[:, 4])
    return mk - (s1 * teff + z1), mk - (s2 * teff + z2)


def in_xr_box(teff, mk, feh, dteff=0.0, dmk=0.0):
    """dteff, dmk: systematic model-minus-data offsets used in robustness
    tests (shift the isochrones before applying the box)."""
    d1, d2 = xr_box_margins(teff + dteff, mk + dmk, feh)
    return (d1 > 0) & (d2 < 0) & (mk + dmk > 0.5) & (teff + dteff < 6800.0)


def imf_integral(m_lo, m_hi, slope=IMF_SLOPE):
    a = 1.0 - slope
    return (m_hi ** a - m_lo ** a) / a


def _g_one_iso(mini, inside, slope):
    """IMF integral over the parts of an isochrone flagged `inside`
    (linear interpolation of the box edge between adjacent points)."""
    g = 0.0
    for a in range(len(mini) - 1):
        if mini[a + 1] <= mini[a]:
            continue
        if inside[a] and inside[a + 1]:
            g += imf_integral(mini[a], mini[a + 1], slope)
        elif inside[a] != inside[a + 1]:
            mid = 0.5 * (mini[a] + mini[a + 1])
            lo, hi = (mini[a], mid) if inside[a] else (mid, mini[a + 1])
            g += imf_integral(lo, hi, slope)
    return g


class TracerWeights:
    """G(tau, [M/H]) on the MIST grid for a box function
    box(teff, mk_or_dict, feh) -> bool array."""

    def __init__(self, box="xr", slope=IMF_SLOPE, dteff=0.0, dmk=0.0,
                 min_eep=353, max_eep=605):
        mist = load_mist()
        self.fehs = np.array(sorted(mist.keys()))
        self.logages = np.unique(mist[self.fehs[0]][:, 0])
        self.G = np.zeros((len(self.logages), len(self.fehs)))
        self.mto = np.zeros_like(self.G)     # TAMS (EEP 454) initial mass
        self.msg = np.full_like(self.G, np.nan)   # mean mass in box
        for j, feh in enumerate(self.fehs):
            arr = mist[feh]
            for i, la in enumerate(self.logages):
                iso = arr[np.isclose(arr[:, 0], la)]
                iso = iso[(iso[:, 1] >= min_eep) & (iso[:, 1] <= max_eep)]
                if len(iso) < 5:
                    continue
                eep, mini = iso[:, 1], iso[:, 2].astype(float)
                teff = 10.0 ** iso[:, 3].astype(float)
                mk = iso[:, 6].astype(float)
                k = np.searchsorted(eep, 454)
                self.mto[i, j] = mini[min(k, len(mini) - 1)]
                if box == "xr":
                    inside = in_xr_box(teff, mk, feh, dteff, dmk)
                else:
                    inside = box(iso, feh)
                self.G[i, j] = _g_one_iso(mini, inside, slope)
                if inside.any():
                    self.msg[i, j] = mini[inside].mean()

    def _interp(self, grid, age_gyr, mh):
        la = np.log10(np.clip(np.asarray(age_gyr, float), 0.8, 19.9)) + 9.0
        mh = np.clip(np.asarray(mh, float), self.fehs[0], self.fehs[-1])
        x = np.interp(la, self.logages, np.arange(len(self.logages)))
        y = np.interp(mh, self.fehs, np.arange(len(self.fehs)))
        i0 = np.clip(np.floor(x).astype(int), 0, len(self.logages) - 2)
        j0 = np.clip(np.floor(y).astype(int), 0, len(self.fehs) - 2)
        fx, fy = x - i0, y - j0
        return ((1 - fx) * (1 - fy) * grid[i0, j0] + fx * (1 - fy) * grid[i0 + 1, j0]
                + (1 - fx) * fy * grid[i0, j0 + 1] + fx * fy * grid[i0 + 1, j0 + 1])

    def g(self, age_gyr, mh):
        return self._interp(self.G, age_gyr, mh)

    def m_turnoff(self, age_gyr, mh):
        return self._interp(self.mto, age_gyr, mh)


def alive_fraction(tw, age_gyr, mh, m_lo=0.8, m_hi=1.2, slope=IMF_SLOPE):
    """IMF-weighted fraction of stars born with m_lo < M < m_hi that are
    still on the main sequence (M < M_TAMS) at the given age."""
    mto = tw.m_turnoff(age_gyr, mh)
    top = np.clip(mto, m_lo, m_hi)
    return imf_integral(m_lo, top, slope) / imf_integral(m_lo, m_hi, slope)