"""Independent re-computation of one model (Neukum 1983 ages x main-belt shape) that shares no code
with model.py.  It exists to check the two numbers the conclusions lean on:
P(no impactor >= 440 km after 4.2 Ga) with the HSE mass used as an upper bound only, and with both
bounds.

Differences from model.py, on purpose:
  * basin ages are the Orgel et al. model ages as tabulated (no density errors, no re-dating);
  * the lunar sampling error is one Gamma(43) factor on the whole record instead of a weight per basin;
  * impactor sizes are resampled from the list of real main-belt bodies >= 150 km with Dirichlet
    weights (no interpolation between bodies);
  * the mass of the bodies below 150 km is the plain sum over real main-belt bodies.
Agreement with model.py to a few hundredths is expected; exact agreement is not.  The main
structural difference is the count at the threshold: three real bodies are >= 440 km (Ceres,
Vesta, Pallas) and the fourth (Hygiea) is 407 km, while model.py interpolates the cumulative count
between bodies and uses 3.6 at 440 km.  The check is therefore run for thresholds of 440 km (three
bodies) and 400 km (four bodies); model.py should fall between the two.

    python3 code/independent_check.py        # prints the comparison and writes data/independent_check.csv
"""
from __future__ import annotations

import os

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "data")
RHO = 3000.0
T_MOON, T_SPA, T_YOUNG = 4.51, 4.31, 3.81
HSE = (0.7e22, 3.0e22)
N_DRAW = 40000


def n1_neukum(T):
    return 5.44e-14 * np.expm1(6.93 * T) + 8.38e-4 * T


def earth_moon(v, a_re):
    ve, vm, re, rm = 11.186, 2.376, 6371.0, 1737.4
    ve2 = ve ** 2 / a_re
    vinf2 = v ** 2 - vm ** 2 - ve2
    return (re / rm) ** 2 * (1 + ve ** 2 / vinf2) / (1 + (vm ** 2 + ve2) / vinf2)


def run(seed=7):
    rng = np.random.default_rng(seed)
    neu = pd.read_csv(os.path.join(DATA, "neumann2015_basins.csv"))
    org = pd.read_csv(os.path.join(DATA, "orgel2018_lpsc_ages.csv"))
    sb = pd.read_csv(os.path.join(DATA, "sbdb_asteroids_d10.csv"))
    import unicodedata
    norm = lambda s: "".join(ch for ch in unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode().lower()
                             if ch.isalpha())
    age = dict(zip(org.name.map(norm), org.age_neukum_Ga))
    big = neu[neu.D_main_km >= 300]
    ages = np.array([age.get(norm(n), np.nan) for n in big.name])
    ages_d = ages[~np.isnan(ages)]
    n_und = int(np.isnan(ages).sum())
    pool = ages_d[ages_d >= 4.17]                 # undated basins: as old as Nectaris or older
    mb = np.sort(sb.loc[sb["class"].isin(["MBA", "IMB", "OMB"]), "diameter"].to_numpy())[::-1]
    mass = np.pi / 6 * RHO * (mb * 1e3) ** 3
    bigb = mb[mb >= 150.0]
    k = len(big) / n1_neukum(T_SPA)               # lunar basins per unit N1

    def c_lunar(T):                               # lunar basins >= 300 km younger than T
        return ((ages_d < T).sum() + n_und * (pool < T).mean() + k * n1_neukum(min(T, T_YOUNG))
                + k * max(n1_neukum(T) - n1_neukum(T_SPA), 0.0))

    c_tot, g42 = c_lunar(T_MOON), c_lunar(4.2) / c_lunar(T_MOON)
    out = dict(lunar_basins_since_4p51=c_tot, fraction_after_4p2=g42, n_draw=N_DRAW)
    R = earth_moon(rng.uniform(11, 22, N_DRAW), rng.uniform(40, 60.3, N_DRAW))
    d_ref = 0.0442 * 300 ** 1.062 * np.exp(rng.normal(0, 0.2, N_DRAW))
    lunar = rng.gamma(len(big), 1.0, N_DRAW) / len(big)          # sampling error of the lunar count
    n_ref = np.array([(mb >= d).sum() for d in d_ref])           # main-belt bodies >= d_ref
    m_small = np.array([mass[(mb >= d) & (mb < 150.0)].sum() for d in d_ref])
    n_earth_ref = R * c_tot * lunar                              # Earth impactors >= d_ref since 4.51 Ga
    p_none, m_tot, n440 = np.empty(N_DRAW), np.empty(N_DRAW), np.empty(N_DRAW)
    p_none4, n400 = np.empty(N_DRAW), np.empty(N_DRAW)
    for i in range(N_DRAW):
        w = rng.dirichlet(np.ones(len(bigb))) * len(bigb)        # sampling error of the main-belt tail
        nb = rng.poisson(n_earth_ref[i] * w.sum() / n_ref[i])
        d = rng.choice(bigb, size=nb, p=w / w.sum())
        m_tot[i] = n_earth_ref[i] * m_small[i] / n_ref[i] + (np.pi / 6 * RHO * (d * 1e3) ** 3).sum()
        n440[i] = (d >= 440.0).sum()
        p_none[i] = (1.0 - g42) ** n440[i]
        n400[i] = (d >= 400.0).sum()
        p_none4[i] = (1.0 - g42) ** n400[i]
    for lab, ok in (("upper", m_tot <= HSE[1]), ("two", (m_tot >= HSE[0]) & (m_tot <= HSE[1]))):
        out[f"P_none_ge440_after_4p2_{lab}"] = float(p_none[ok].mean())
        out[f"n_ge440_since_4p51_{lab}"] = float(n440[ok].mean())
        out[f"P_none_ge400_after_4p2_{lab}"] = float(p_none4[ok].mean())
        out[f"n_ge400_since_4p51_{lab}"] = float(n400[ok].mean())
        out[f"pass_fraction_{lab}"] = float(ok.mean())
    return out


def compare(write=True):
    out = run()
    rows = [dict(quantity=k, independent=v) for k, v in out.items()]
    path = os.path.join(DATA, "model_grid.csv")
    if os.path.exists(path):
        g = pd.read_csv(path)
        r = g[(g.chron == "neukum1983") & (g.sfd == "mainbelt")].iloc[0]
        ref = {"P_none_ge440_after_4p2_upper": r["upper_P_none_ge440_after_4.2"],
               "P_none_ge440_after_4p2_two": r["two_P_none_ge440_after_4.2"],
               "n_ge440_since_4p51_upper": r["upper_mean_n_ge440_since_4p51"],
               "n_ge440_since_4p51_two": r["two_mean_n_ge440_since_4p51"],
               "pass_fraction_upper": r["pass_fraction_upper"], "pass_fraction_two": r["pass_fraction_two"]}
        ref.update({k.replace("440", "400"): v for k, v in ref.items() if "440" in k})   # same model.py value: it must lie between
        for row in rows:
            row["model_py_neukum1983_mainbelt_ge440"] = ref.get(row["quantity"], np.nan)
    df = pd.DataFrame(rows)
    if write:
        df.round(4).to_csv(os.path.join(DATA, "independent_check.csv"), index=False)
    return df


if __name__ == "__main__":
    print(compare().to_string(index=False))