"""Part 5: humanity's birth rank with a geophysical lifetime DISTRIBUTION in place of a fixed cap.

Same model as studies/mdwarf_early (snapshot in mdwarf_earliness_snapshot.py): stars of mass M form at
cosmic time b at rate psi(b) xi(M); with n hard steps the chance that a civilisation has appeared by
planet age a grows as a^n while the planet is habitable.  The habitable window of a planet is now
        L = min(tau_MS(M), L_geo),
with L_geo drawn from the Monte Carlo lifetime distribution of the planet's formation epoch, so

    N(<t) = int dM xi(M) w(M) int_0^t db psi(b) E[ min(t - b, tau_MS(M), L_geo)^n | b ] .

E[min(x, L_geo)^n] =: G_b(x) is tabulated from the lifetime samples at a few formation-epoch nodes and
interpolated linearly in stellar age between nodes.  Stars that form in the future use the node for
stars forming today.  Rank = N(<t0) / N(<infinity).
"""
import os

import numpy as np
import pandas as pd

import mdwarf_earliness_snapshot as E
from common import DATA_DIR


def tau_hz():
    """Stellar limit on the habitable window that is consistent with stellar.py: the smaller of the
    main-sequence lifetime used by the sister study and the median habitable-zone exit time from the
    Geneva tracks (0.5-1.0 Msun; log-log interpolation; above 1 Msun the 1 Msun ratio of exit time to
    main-sequence lifetime is kept; below 0.5 Msun the main-sequence lifetime is used).  Approximate:
    the exit time depends on the orbit and only its median is used."""
    hz = pd.read_csv(os.path.join(DATA_DIR, "stellar_hz_lifetimes.csv")).sort_values("mass_msun")
    m, t = hz["mass_msun"].values, hz["hz_exit_median_gyr"].values
    tau_ms = E.ms_lifetime_gyr(E.MASS_GRID)
    t_hz = np.exp(np.interp(np.log(E.MASS_GRID), np.log(m), np.log(t)))
    ratio_hi = t[-1] / float(E.ms_lifetime_gyr(m[-1]))
    t_hz = np.where(E.MASS_GRID > m[-1], ratio_hi * tau_ms, t_hz)
    t_hz = np.where(E.MASS_GRID < m[0], tau_ms, t_hz)
    return np.minimum(tau_ms, t_hz)


def g_table(life, n, xgrid, weights=None):
    """G(x) = E[min(x, L)^n] on xgrid from lifetime samples (optionally weighted)."""
    life = np.asarray(life, dtype=float)
    order = np.argsort(life)
    ls = life[order]
    w = np.full(len(ls), 1.0 / len(ls)) if weights is None else np.asarray(weights, dtype=float)[order] / np.sum(weights)
    cum_ln = np.concatenate([[0.0], np.cumsum(w * ls ** n)])
    cum_w = np.concatenate([[0.0], np.cumsum(w)])
    idx = np.searchsorted(ls, xgrid, side="right")           # samples with L <= x
    return cum_ln[idx] + xgrid ** n * (1.0 - cum_w[idx])


def rank_from_lifetimes(life_nodes, node_ages, n, sfh="MD14", imf="Kroupa", m_min_host=None, weights=None,
                        x_max=None, nx=3000, b_step=4, tau=None):
    """life_nodes: (K, N) lifetimes for formation-epoch nodes given as stellar age today, node_ages (K,)
    increasing.  Returns the rank of today's cosmic time.  `tau` (optional, on E.MASS_GRID) replaces the
    main-sequence lifetime as the stellar limit on the habitable window."""
    life_nodes = np.atleast_2d(life_nodes)
    node_ages = np.asarray(node_ages, dtype=float)
    t0 = E.age_today()
    b, psi, past, total = E._time_grids(sfh)
    if b_step > 1:                       # coarser birth-time grid for speed (end point kept); past/total unchanged
        idx = np.unique(np.concatenate([np.arange(0, len(b), b_step), [len(b) - 1]]))
        b, psi = b[idx], psi[idx]
    if tau is None:
        tau = E.ms_lifetime_gyr(E.MASS_GRID)      # as in the sister study; tau_hz() is the consistent alternative
    if x_max is None:
        x_max = float(life_nodes.max()) * 1.001 + 1e-6
    xgrid = np.linspace(0.0, x_max, nx)
    G = np.array([g_table(life_nodes[k], n, xgrid, weights) for k in range(len(node_ages))])    # (K, nx)
    age = t0 - b                                                                                # stellar age today
    # linear interpolation weights between nodes (clipped at the ends)
    a_c = np.clip(age, node_ages[0], node_ages[-1])
    hi = np.clip(np.searchsorted(node_ages, a_c, side="right"), 1, len(node_ages) - 1)
    lo = hi - 1
    wt = (a_c - node_ages[lo]) / (node_ages[hi] - node_ages[lo]) if len(node_ages) > 1 else np.zeros_like(a_c)

    def g_at(x, k_lo, k_hi, w_hi):
        xi_ = np.clip(x, 0.0, x_max) / x_max * (nx - 1)
        i0 = np.minimum(xi_.astype(int), nx - 2)
        fr = xi_ - i0
        g_lo = G[k_lo, i0] * (1 - fr) + G[k_lo, i0 + 1] * fr
        g_hi = G[k_hi, i0] * (1 - fr) + G[k_hi, i0 + 1] * fr
        return g_lo * (1 - w_hi) + g_hi * w_hi

    if len(node_ages) == 1:
        lo = hi = np.zeros_like(lo)
    x_num = np.minimum(age[None, :], tau[:, None])                    # (mass, b)
    kl, kh, ww = lo[None, :] + 0 * x_num.astype(int), hi[None, :] + 0 * x_num.astype(int), wt[None, :] + 0 * x_num
    inner_num = np.trapezoid(psi[None, :] * g_at(x_num, kl, kh, ww), b, axis=1)
    x_den = np.broadcast_to(tau[:, None], x_num.shape)
    inner_den = np.trapezoid(psi[None, :] * g_at(x_den, kl, kh, ww), b, axis=1)
    k0 = np.zeros(len(tau), dtype=int)                                # node for stars forming today / in the future
    fut = (total - past) * g_at(tau, k0, k0, np.zeros(len(tau)))
    xi = E.imf_number(E.MASS_GRID, imf) * E.MASS_GRID * E.DLNM
    w = np.ones_like(tau) if m_min_host is None else E.host_weight(m_min_host)
    return float((w * xi * inner_num).sum() / (w * xi * (inner_den + fut)).sum())


def rank_fixed_cap(n, cap, sfh="MD14", m_min_host=None):
    """Reference: the fixed-cap rank of the mdwarf_early study (all stars, weight 1)."""
    if m_min_host is None:
        return E.rank_step(n, cap, 1.0, model=sfh)
    return E.rank_host_cut(n, cap, m_min_host, model=sfh)


if __name__ == "__main__":
    # self-test: a delta-function lifetime distribution must reproduce the fixed-cap ranks
    for n in (1, 2, 4, 6):
        for cap in (10.0, 20.0):
            mine = rank_from_lifetimes(np.full((2, 50), cap), [0.0, 12.0], n)
            fine = rank_from_lifetimes(np.full((2, 50), cap), [0.0, 12.0], n, b_step=1)
            ref = rank_fixed_cap(n, cap)
            assert abs(mine / ref - 1) < 2e-3 and abs(fine / ref - 1) < 1e-4, (mine, fine, ref)
            print(f"n={n} cap={cap}: distribution code {mine:.5f}  mdwarf_early code {ref:.5f}")
    print("hosts > 0.6 Msun, cap 1e5:", [round(rank_fixed_cap(n, 1e5, m_min_host=0.6), 5) for n in (2, 4, 6)])