"""Shared paths, constants and helpers for the planet_lifetime study."""
import os
import numpy as np

CODE_DIR = os.path.dirname(os.path.abspath(__file__))
STUDY_DIR = os.path.dirname(CODE_DIR)
DATA_DIR = os.path.join(STUDY_DIR, "data")
FIG_DIR = os.path.join(STUDY_DIR, "figures")
RAW_DIR = os.environ.get("FERMI_RAW", "fermi_raw/planet_lifetime")
for _d in (DATA_DIR, FIG_DIR, RAW_DIR):
    os.makedirs(_d, exist_ok=True)

SEED = 20261012
LN2 = np.log(2.0)

# ---------------------------------------------------------------------------------------------
# Long-lived heat-producing nuclides.
# Half-lives [Gyr] and specific heat production [W/kg of nuclide]: Turcotte & Schubert,
# Geodynamics (3rd ed., 2014), table 4.2.  Isotopic fractions today are the terrestrial values.
# ---------------------------------------------------------------------------------------------
NUCLIDES = ("U238", "U235", "Th232", "K40")
HALF_LIFE = {"U238": 4.47, "U235": 0.704, "Th232": 14.0, "K40": 1.25}          # Gyr
SPEC_HEAT = {"U238": 9.46e-5, "U235": 5.69e-4, "Th232": 2.64e-5, "K40": 2.92e-5}  # W/kg
ISO_FRAC = {"U238": 0.9928, "U235": 0.0071, "Th232": 1.0, "K40": 1.19e-4}       # of the element, today

# Bulk silicate Earth today (McDonough & Sun 1995, Chem. Geol. 120, 223): U 20.3 ppb, Th 79.5 ppb,
# K 240 ppm.  These give 20 TW for the 4.0e24 kg silicate Earth.
BSE_PPB = {"U": 20.3, "Th": 79.5, "K": 240.0e3}
M_BSE = 4.0e24            # kg, mantle + crust
AGE_EARTH = 4.5           # Gyr, age of the Earth used for calibration
AGE_SUN = 4.57            # Gyr, age of the Sun (meteoritic) used for decay corrections
T_UNIVERSE = 13.79        # Gyr (Planck 2018), same as studies/mdwarf_early
SEC_PER_GYR = 3.15576e16


def earth_heat_now_TW():
    """Present-day BSE heat production by nuclide [TW]."""
    out = {}
    for nuc in NUCLIDES:
        el = "U" if nuc.startswith("U") else ("Th" if nuc.startswith("Th") else "K")
        out[nuc] = BSE_PPB[el] * 1e-9 * ISO_FRAC[nuc] * SPEC_HEAT[nuc] * M_BSE / 1e12
    return out


def earth_heat_initial_TW():
    """BSE heat production by nuclide at Earth's formation [TW] (today's value undone for decay)."""
    now = earth_heat_now_TW()
    return {k: v * 2.0 ** (AGE_EARTH / HALF_LIFE[k]) for k, v in now.items()}


def heat_TW(age_gyr, f_th=1.0, f_u=1.0, f_k=1.0, mass_scale=1.0):
    """Radiogenic heat production [TW] of a planet of age `age_gyr` whose initial Th, U and K
    concentrations are f_th, f_u, f_k times the Earth's initial concentrations.

    f_u multiplies both U isotopes (the 235U/238U production ratio is assumed universal)."""
    h0 = earth_heat_initial_TW()
    age = np.asarray(age_gyr, dtype=float)
    out = 0.0
    for nuc, f in (("U238", f_u), ("U235", f_u), ("Th232", f_th), ("K40", f_k)):
        out = out + f * h0[nuc] * 2.0 ** (-age / HALF_LIFE[nuc])
    return mass_scale * out