"""Shared constants, date distributions and plotting style for the timing study."""
from __future__ import annotations

import os

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
STUDY = os.path.abspath(os.path.join(HERE, ".."))
DATA = os.path.join(STUDY, "data")
FIGS = os.path.join(STUDY, "figures")
os.makedirs(DATA, exist_ok=True)
os.makedirs(FIGS, exist_ok=True)

SEED = 20261011

# ---------------------------------------------------------------------------------
# Baseline dates (Ga).  See data/events.csv for sources.
# ---------------------------------------------------------------------------------
T0_ZIRCON = 4.404          # Wilde et al. 2001; Kipping's start of habitability
T0_MOON = 4.510            # Moon-forming impact (Moody et al. calibration)
T0_LATE = 4.50             # latest plausible "oceans within ~10 Myr of the impact" start
LUCA_PAPER = (4.20, 4.09, 4.33)   # Moody et al. 2024 headline (centre, 2.5%, 97.5%)
LIFE_DATES = {             # label -> age (Ga) of earliest evidence used as upper bound on abiogenesis
    "LUCA 4.2 Ga (Moody+ 2024)": 4.20,
    "zircon graphite 4.1 Ga (contested)": 4.10,
    "Labrador 3.95 Ga (contested)": 3.95,
    "Isua 3.7 Ga (contested)": 3.70,
    "Dresser 3.48 Ga (accepted)": 3.48,
    "Apex 3.465 Ga (Kipping 2020 conservative)": 3.465,
}
FUTURE = {                 # label -> remaining habitable time (Gyr)
    "0.9 (Caldeira & Kasting 1992; Kipping)": 0.9,
    "1.08 (Ozaki & Reinhard 2021)": 1.08,
    "1.73 (Graham+ 2024 midpoint)": 1.73,
    "1.84 (Haqq-Misra & Wolf 2026)": 1.84,
}
F_SHORT = (1.08, 0.14)     # Ozaki & Reinhard 2021 mean, 1 sigma
F_LONG = (1.60, 1.87)      # Graham et al. 2024 / Haqq-Misra & Wolf 2026 range (uniform)

# transition dates used in the multi-step models (centre, lo, hi), Ga; see events.csv
OXYPHOTO = (2.99, 2.95, 3.09)
LECA = (1.735, 1.67, 1.80)
METAZOA = (0.60, 0.574, 0.80)
# environmental gates used by the pacing model (centre, lo, hi), Ga
GOE = (2.33, 2.22, 2.46)
NOE = (0.635, 0.539, 0.80)
POE = (0.41, 0.40, 0.42)


def split_normal(rng, centre, q025, q975, size):
    """Two-piece normal with the given median and 2.5%/97.5% quantiles."""
    z = rng.standard_normal(size)
    s_lo = (centre - q025) / 1.959964
    s_hi = (q975 - centre) / 1.959964
    return centre + np.where(z < 0, s_lo, s_hi) * z


def draw_future(rng, size, which="mix"):
    """Remaining habitable time (Gyr).  'short' = Ozaki & Reinhard; 'long' = Graham et al. /
    Haqq-Misra & Wolf; 'mix' = 50/50 mixture of the two (our baseline)."""
    short = np.clip(rng.normal(F_SHORT[0], F_SHORT[1], size), 0.3, None)
    long_ = rng.uniform(F_LONG[0], F_LONG[1], size)
    if which == "short":
        return short
    if which == "long":
        return long_
    pick = rng.random(size) < 0.5
    return np.where(pick, short, long_)


def draw_t0(rng, size):
    """Start of habitability (Ga): uniform between the zircon water evidence (4.404 Ga) and
    10 Myr after the Moon-forming impact (4.50 Ga)."""
    return rng.uniform(T0_ZIRCON, T0_LATE, size)


def draw_uniform(rng, triple, size):
    return rng.uniform(triple[1], triple[2], size)


# ---------------------------------------------------------------------------------
# Plot style (dataviz skill reference palette, light surface)
# ---------------------------------------------------------------------------------
C = dict(blue="#2a78d6", orange="#eb6834", aqua="#1baf7a", yellow="#eda100",
         magenta="#e87ba4", green="#008300", violet="#6250d6", red="#e34948")
SERIES = [C["blue"], C["orange"], C["aqua"], C["yellow"], C["magenta"], C["green"], C["violet"], C["red"]]
SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK2 = "#52514e"
GRID = "#e4e3df"


def style():
    import matplotlib as mpl
    mpl.use("Agg")
    mpl.rcParams.update({
        "figure.facecolor": SURFACE, "axes.facecolor": SURFACE, "savefig.facecolor": SURFACE,
        "axes.edgecolor": "#b9b8b2", "axes.labelcolor": INK2, "text.color": INK,
        "xtick.color": INK2, "ytick.color": INK2, "axes.grid": True, "grid.color": GRID,
        "grid.linewidth": 0.6, "axes.spines.top": False, "axes.spines.right": False,
        "font.size": 10, "axes.titlesize": 11, "axes.titleweight": "bold", "axes.titlelocation": "left",
        "legend.frameon": False, "lines.linewidth": 2.0, "figure.dpi": 130, "savefig.dpi": 160,
        "axes.axisbelow": True,
    })