"""
Cosmological tables for the integrated model (v3).

For a population of planets that form following the cosmic star-formation history and on which a
civilisation arrives after n sequential hard steps inside a habitable window of length L (Gyr),
this module tabulates, on a common cosmic-time grid:

  b[n, L, t]     birth-date density of civilisations (per Gyr), normalised so that the integral over
                 all time is 1 if every planet eventually forms
  F[n, L, t]     cumulative fraction born by t
  kern[n, L, t]  (4 pi / 3) * integral b(t') (eta(t) - eta(t'))^3 dt'   [Gly^3]
                 = comoving volume swept by light-speed spheres centred on each birth; multiply by
                 v^3 * (expander origins per civilisation) * (civilisations per galaxy) * (galaxies per Gly^3)
                 to get -ln P(a random point is still unreached at t)

Because every quantity is linear in the birth law, a world with several host classes (G, K, M) is
just a weighted sum of table rows. Cosmology and star-formation history as in grabby.py.

Run: python3 cosmo_tables.py  -> cosmo_tables.npz
"""
import os
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
H0 = 67.7 / 977.8          # 1/Gyr
OM, OL = 0.31, 0.69
T0 = 13.79
RHO_GAL = 2.88e5           # Milky-Way equivalents per Gly^3 (0.01 per Mpc^3)
T_DELAY0 = 4.5             # Gyr: fixed delay to intelligence when there are no hard steps

N_STEPS = np.arange(0, 9)                                    # 0..8 hard steps
L_GRID = np.array([5.6, 8, 10, 12, 15, 20, 30, 45, 70, 100, 200, 400, 1000.0])

_a = np.logspace(-3.5, 80, 60000)
_Ha = H0 * np.sqrt(OM / _a ** 3 + OL)
_dt = 1.0 / (_a * _Ha)
_de = 1.0 / (_a ** 2 * _Ha)
_t = np.concatenate([[0], np.cumsum(0.5 * (_dt[1:] + _dt[:-1]) * np.diff(_a))]) + 2.0 / (3 * H0 * np.sqrt(OM)) * _a[0] ** 1.5
_eta = np.concatenate([[0], np.cumsum(0.5 * (_de[1:] + _de[:-1]) * np.diff(_a))]) + 2.0 / (H0 * np.sqrt(OM)) * _a[0] ** 0.5


def a_of_t(t):
    return np.interp(t, _t, _a)


def eta_of_t(t):
    return np.interp(t, _t, _eta)


def sfr(t):
    """Madau & Dickinson (2014) cosmic star-formation rate density, extrapolated into the future with the same form."""
    zp1 = 1.0 / a_of_t(np.maximum(t, 1e-3))
    return np.where(t > 0.05, 0.015 * zp1 ** 2.7 / (1 + (zp1 / 2.9) ** 5.6), 0.0)


def build(n_t=280, t_max=3000.0):
    t = np.unique(np.concatenate([np.logspace(np.log10(0.3), np.log10(t_max), n_t), [T0]]))
    eta = eta_of_t(t)
    tf = np.logspace(np.log10(0.05), np.log10(t_max), 6000)
    sfr_total = np.trapezoid(sfr(tf), tf)
    nt = len(t)
    b = np.zeros((len(N_STEPS), len(L_GRID), nt))
    u = (np.arange(600) + 0.5) / 600                          # midpoints in tau / tau_max
    for i, n in enumerate(N_STEPS):
        for j, L in enumerate(L_GRID):
            if n == 0:
                b[i, j] = sfr(t - T_DELAY0) / sfr_total if L >= T_DELAY0 else 0.0
                continue
            tau_max = np.minimum(L, t)[:, None]               # (nt, 1)
            tau = u[None, :] * tau_max                        # (nt, 600)
            dens = n * tau ** (n - 1) / L ** n                # arrival-time density inside the window
            psi = sfr(t[:, None] - tau)
            b[i, j] = (psi * dens).mean(axis=1) * tau_max[:, 0] / sfr_total
    F = np.zeros_like(b)
    F[..., 1:] = np.cumsum(0.5 * (b[..., 1:] + b[..., :-1]) * np.diff(t), axis=-1)
    kern = np.zeros_like(b)
    for k in range(1, nt):
        d3 = (eta[k] - eta[: k + 1]) ** 3
        kern[..., k] = 4 * np.pi / 3 * np.trapezoid(b[..., : k + 1] * d3, t[: k + 1], axis=-1)
    return dict(t=t, eta=eta, b=b, F=F, kern=kern, n_steps=N_STEPS, L_grid=L_GRID, i0=int(np.searchsorted(t, T0)),
                sfr_total=sfr_total, frac_sf_done=float(np.trapezoid(sfr(tf[tf <= T0]), tf[tf <= T0]) / sfr_total))


if __name__ == "__main__":
    tab = build()
    np.savez_compressed(os.path.join(HERE, "cosmo_tables.npz"), **tab)
    t, i0 = tab["t"], tab["i0"]
    print("grid points", len(t), " t[i0] =", t[i0], " fraction of all star formation already done:", round(tab["frac_sf_done"], 3))
    print("total births integral (should be <= 1):")
    for i, n in enumerate(N_STEPS[:7]):
        print(f"  n={n}: " + " ".join(f"L={L:g}:{tab['F'][i, j, -1]:.2f}" for j, L in enumerate(L_GRID[[0, 3, 5, 9, 12]])))
    print("fraction of each class's civilisations born by now, F(T0):")
    for i, n in enumerate(N_STEPS[:7]):
        print(f"  n={n}: " + " ".join(f"L={L:g}:{tab['F'][i, j, i0]:.3g}" for j, L in zip([0, 3, 5, 9, 12], L_GRID[[0, 3, 5, 9, 12]])))
    # sanity: galaxies whose expanders would have reached us, per unit (origins per galaxy), v = 0.1c and 0.5c
    for i in (1, 3, 6):
        k0 = tab["kern"][i, 0, i0] / max(tab["F"][i, 0, i0], 1e-30)
        print(f"  n={N_STEPS[i]}: V_eff = kern/F at T0 = {k0:.0f} Gly^3 -> N_reach(0.1c) = {RHO_GAL * 1e-3 * k0:.3g}, N_reach(0.5c) = {RHO_GAL * 0.125 * k0:.3g}")