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

STUDY = ''
RAW = 'fermi_raw/stellar_salting/'
CACHE = RAW + 'cache/'
DATA = STUDY + 'data/'
FIG = STUDY + 'figures/'
for _d in (RAW, CACHE, DATA, FIG, RAW + 'spec/'):
    os.makedirs(_d, exist_ok=True)

# neutron-capture elements reported by GALAH (lower-case = DR4 column stem)
NCAP = ['rb', 'sr', 'y', 'zr', 'mo', 'ru', 'ba', 'la', 'ce', 'nd', 'sm', 'eu']
# upper edge of the DR4 neural-network label range (flag bit 8 above it for the 1.5 group;
# the 1.0 group is clipped at 1.0 in the catalogue)
DR4_UPPER = {'y': 1.5, 'ba': 1.5, 'la': 1.5, 'ce': 1.5, 'nd': 1.5,
             'rb': 1.0, 'sr': 1.0, 'zr': 1.0, 'mo': 1.0, 'ru': 1.0, 'sm': 1.0, 'eu': 1.0}

# FGK main-sequence / subgiant selection used throughout
TEFF_MIN, TEFF_MAX, LOGG_MIN = 4500.0, 6500.0, 3.5

# solar-system s-process fractions: Bisterzo et al. 2014 (ApJ 787, 10), Table 1, column "TW+"
# (this-work GCE model, low- plus intermediate-mass AGB), VizieR J/ApJ/787/10/table1, values in
# per cent / 100.  The remainder (1 - s) is treated as the r-process (for Sr-Y-Zr it also contains
# the weak-s/LEPP contribution).  Used to define the natural (s + r mixture) locus.
S_FRAC = {'rb': 0.181, 'sr': 0.689, 'y': 0.719, 'zr': 0.663, 'mo': 0.387, 'ru': 0.289, 'ba': 0.852,
          'la': 0.755, 'ce': 0.835, 'pr': 0.499, 'nd': 0.575, 'sm': 0.314, 'eu': 0.060}
# alternative set for the robustness test: same table, column "T04+" (Travaglio et al. 2004)
S_FRAC_ALT = {'rb': 0.381, 'sr': 0.698, 'y': 0.683, 'zr': 0.649, 'mo': 0.379, 'ru': 0.241, 'ba': 0.807,
              'la': 0.625, 'ce': 0.752, 'pr': 0.477, 'nd': 0.545, 'sm': 0.296, 'eu': 0.058}

# Asplund, Amarsi & Grevesse 2021 (A&A 653, A141) photospheric log eps (meteoritic for Xe, Cs absent)
A_SUN = {'rb': 2.32, 'sr': 2.83, 'y': 2.21, 'zr': 2.59, 'nb': 1.47, 'mo': 1.88, 'ru': 1.75,
         'rh': 0.78, 'pd': 1.57, 'ba': 2.27, 'la': 1.11, 'ce': 1.58, 'pr': 0.75, 'nd': 1.42,
         'sm': 0.95, 'eu': 0.52, 'gd': 1.08, 'tc': None}
ATOMIC_MASS = {'rb': 85.47, 'sr': 87.62, 'y': 88.91, 'zr': 91.22, 'nb': 92.91, 'mo': 95.95,
               'tc': 98.0, 'ru': 101.07, 'rh': 102.91, 'pd': 106.42, 'ba': 137.33, 'la': 138.91,
               'ce': 140.12, 'pr': 140.91, 'nd': 144.24, 'sm': 150.36, 'eu': 151.96, 'gd': 157.25}

MSUN = 1.98841e30   # kg
M_H = 1.6735e-27    # kg
X_H = 0.7438        # present-day solar surface hydrogen mass fraction (Asplund+ 2021)


def binom_upper(k, n, cl=0.95):
    """Exact (Clopper-Pearson) one-sided upper limit on a binomial proportion."""
    from scipy.stats import beta
    if n <= 0:
        return np.nan
    return float(beta.ppf(cl, k + 1, max(n - k, 1e-9)))


def poisson_upper(k, cl=0.95):
    from scipy.stats import chi2
    return 0.5 * chi2.ppf(cl, 2 * (k + 1))