"""Environmental curves on a common 1-Myr grid (age in Ma, index 0 = present).

Phanerozoic oxygen (three independent families):
  gcso  : GEOCARBSULFOR Monte Carlo, isotope mass balance (Krause et al. 2018; outputs of Mills et al. 2023),
          1000 realisations, 550-0 Ma.  Columns of the saved table: median, 16th and 84th percentile.
  copse : COPSE forward model (Tostevin & Mills 2020 baseline 'standard' plus nine feedback variants of
          Vitali et al.), 900-0 Ma (the first ~100 Myr are spin-up from an imposed 0.1 PAL and are not used).
  char  : charcoal (inertinite) envelope of Glasspool & Scott 2010 as bundled with SCION / pySCION
          (upper and lower bound every 10 Myr, 400-0 Ma); the midpoint is used.
Temperature: PhanDA global mean surface temperature (Judd et al. 2024; stage means, 485-0 Ma) and the
  Scotese et al. 2021 global average temperature (540-0 Ma, 1-Myr steps).
CO2: PhanDA median (485-0 Ma).

Conversions: O2 mixing ratio m (fraction) <-> mass relative to present R (PAL): m = R / (R + 3.762).
"""
import glob
import os

import numpy as np
import pandas as pd
import scipy.io as sio

from config import RAW, DATA

K_N2 = 3.762  # (1 - 0.2095) / 0.2095: moles of non-O2 gas per mole of O2 today


def mr_to_pal(m):
    m = np.asarray(m, float)
    return K_N2 * m / (1.0 - m)


def pal_to_mr(r):
    r = np.asarray(r, float)
    return r / (r + K_N2)


def load_gcso():
    """Return ages 0..550 and an array (551, n_ok) of O2 in per cent (n_ok = 947 of 1000), row index = age in Ma."""
    m = sio.loadmat(os.path.join(RAW, "geocarbsulfor", "x", "outputs_14_7_2022.mat"),
                    squeeze_me=True, struct_as_record=False)["mc_outputs"]
    t = np.asarray(m.time)[:, 0]          # -550 ... 0
    o2 = np.asarray(m.O2_percent)         # (551, 1000)
    order = np.argsort(-t)                # age 0 first
    ages = (-t[order]).astype(float)
    assert np.allclose(ages, np.arange(551))
    o2 = o2[order, :]
    ok = ~np.isnan(o2).any(axis=0)       # 53 of the 1000 published realisations contain NaN (failed runs)
    return ages, o2[:, ok]


def load_copse():
    """Return ages 0..900 and dict name -> O2 per cent on that grid."""
    out = {}
    ages = np.arange(0, 901, dtype=float)
    pat = os.path.join(RAW, "copse", "x", "*", "output", "COPSE_*.txt")
    for f in sorted(glob.glob(pat)):
        name = os.path.basename(f)[6:-4]
        d = pd.read_csv(f, usecols=["time_myr", "mrO2"])
        a = -d["time_myr"].to_numpy()
        v = 100.0 * d["mrO2"].to_numpy()
        o = np.argsort(a)
        a, v = a[o], v[o]
        keep = np.concatenate([[True], np.diff(a) > 0])
        out[name] = np.interp(ages, a[keep], v[keep])
    return ages, out


def load_charcoal():
    m = sio.loadmat(os.path.join(RAW, "pyscion", "x", "pySCION", "data", "geochem_data_2024.mat"),
                    squeeze_me=True)
    x = -np.asarray(m["O2_x"], float)
    y = np.asarray(m["O2_y"], float)
    b = np.round(x / 10.0) * 10.0
    df = pd.DataFrame(dict(age=b, o2=y)).groupby("age")["o2"].agg(["min", "max"]).reset_index()
    ages = np.arange(0, 401, dtype=float)
    lo = np.interp(ages, df["age"], df["min"])
    hi = np.interp(ages, df["age"], df["max"])
    return ages, lo, hi, df


def load_temperature():
    p = pd.read_csv(os.path.join(RAW, "pyscion", "x", "pySCION", "data", "PhanDA_GMSTandCO2_percentiles.csv"))
    ages = np.arange(0, 486, dtype=float)
    o = np.argsort(p["AverageAge"].to_numpy())
    a = p["AverageAge"].to_numpy()[o]
    phanda = np.interp(ages, a, p["GMST_50"].to_numpy()[o])
    co2 = np.interp(ages, a, p["CO2_50"].to_numpy()[o])
    s = sio.loadmat(os.path.join(RAW, "pyscion", "x", "pySCION", "data", "Scotese_GAT_2021.mat"), squeeze_me=True)
    sa = np.abs(np.asarray(s["Scotese_2021_age"], float))
    sv = np.asarray(s["Scotese_2021_GAT"], float)
    o = np.argsort(sa)
    ages_s = np.arange(0, 541, dtype=float)
    scot = np.interp(ages_s, sa[o], sv[o])
    return ages, phanda, co2, ages_s, scot


def build():
    ages_g, g = load_gcso()
    ages_c, c = load_copse()
    ages_ch, lo, hi, chtab = load_charcoal()
    ages_t, phanda, co2, ages_s, scot = load_temperature()
    grid = np.arange(0, 901, dtype=float)
    tab = pd.DataFrame({"age_Ma": grid})

    def put(name, ages, vals):
        col = np.full(len(grid), np.nan)
        col[: len(ages)] = vals
        tab[name] = col

    put("gcso_median", ages_g, np.median(g, axis=1))
    put("gcso_p16", ages_g, np.quantile(g, 0.16, axis=1))
    put("gcso_p84", ages_g, np.quantile(g, 0.84, axis=1))
    for k, v in c.items():
        put("copse_" + k, ages_c, v)
    put("charcoal_lo", ages_ch, lo)
    put("charcoal_hi", ages_ch, hi)
    put("charcoal_mid", ages_ch, 0.5 * (lo + hi))
    put("phanda_gmst", ages_t, phanda)
    put("phanda_co2", ages_t, co2)
    put("scotese_gat", ages_s, scot)
    tab.to_csv(os.path.join(DATA, "curves_1myr.csv"), index=False, float_format="%.4f")
    return dict(grid=grid, table=tab, gcso_ages=ages_g, gcso=g, copse_ages=ages_c, copse=c)


if __name__ == "__main__":
    b = build()
    t = b["table"]
    pd.set_option("display.width", 220)
    print(t.iloc[::25, :].round(2).to_string())