#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Copyright (C) 2025 Tristan Hoellinger
# Distributed under the GNU General Public License v3.0 (GPLv3).
# See the LICENSE file in the root directory for details.
# SPDX-License-Identifier: GPL-3.0-or-later
# ----------------------------------------------------------------------

__author__ = "Tristan Hoellinger"
__version__ = "0.1.0"
__date__ = "2025"
__license__ = "GPLv3"

"""Global parameters for this project."""

import os
from pathlib import Path
from typing import cast
import numpy as np

WHICH_SPECTRUM = "eh"  # available options are "eh" and "class"

ROOT_PATH = cast(str, os.getenv("WIP3M_ROOT_PATH"))
if ROOT_PATH is None:
    raise EnvironmentError("Please set the 'WIP3M_ROOT_PATH' environment variable.")

OUTPUT_PATH = cast(str, os.getenv("WIP3M_OUTPUT_PATH"))
if OUTPUT_PATH is None:
    raise EnvironmentError("Please set the 'WIP3M_OUTPUT_PATH' environment variable.")

# Default verbose level
# 0: errors only, 1: info, 2: warnings+, 3: all diagnostics, 4+: debug
DEFAULT_VERBOSE_LEVEL = 2

# Baseline seeds for reproducibility
BASELINE_SEEDNORM = 100050599
BASELINE_SEEDNOISE = 200050599
BASELINE_SEEDPHASE = 300050599

# Fiducial cosmological parameters
h_planck = 0.6766
Omega_b_planck = 0.02242 / h_planck**2
Omega_m_planck = 0.3111
nS_planck = 0.9665
sigma8_planck = 0.8102

TAU_REIO = 0.066

planck_mean = np.array([h_planck, Omega_b_planck, Omega_m_planck, nS_planck, sigma8_planck])
planck_cov = np.diag(np.array([0.0042, 0.00030, 0.0056, 0.0038, 0.0060]) ** 2)

# Mapping from cosmological parameter names to corresponding indices
cosmo_params_names = [r"$h$", r"$\Omega_b$", r"$\Omega_m$", r"$n_S$", r"$\sigma_8$"]
cosmo_params_name_to_idx = {"h": 0, "Omega_b": 1, "Omega_m": 2, "n_s": 3, "sigma8": 4}

# Minimum k value used in the normalisation of the summaries
MIN_K_NORMALISATION = 4e-2

params_planck_kmax_missing = {
    "h": h_planck,
    "Omega_r": 0.0,
    "Omega_q": 1.0 - Omega_m_planck,
    "Omega_b": Omega_b_planck,
    "Omega_m": Omega_m_planck,
    "m_ncdm": 0.0,
    "Omega_k": 0.0,
    "tau_reio": TAU_REIO,
    "n_s": nS_planck,
    "sigma8": sigma8_planck,
    "w0_fld": -1.0,
    "wa_fld": 0.0,
    "WhichSpectrum": WHICH_SPECTRUM,
}


def z2a(z):
    return 1.0 / (1 + z)


def cosmo_small_to_full_dict(cosmo_min):
    """Return a full cosmology dictionary from a minimal one.

    Parameters
    ----------
    cosmo_min : dict
        Minimal cosmology dictionary.

    Returns
    -------
    cosmo_full : dict
        Full cosmology dictionary.

    """
    cosmo_full = {
        "h": cosmo_min["h"],
        "Omega_r": 0.0,
        "Omega_q": 1 - cosmo_min["Omega_m"],
        "Omega_b": cosmo_min["Omega_b"],
        "Omega_m": cosmo_min["Omega_m"],
        "m_ncdm": 0.0,
        "Omega_k": 0.0,
        "tau_reio": TAU_REIO,
        "n_s": cosmo_min["n_s"],
        "sigma8": cosmo_min["sigma8"],
        "w0_fld": -1.0,
        "wa_fld": 0.0,
        "k_max": cosmo_min["k_max"],
        "WhichSpectrum": cosmo_min["WhichSpectrum"],
    }
    return cosmo_full

# pyright: reportUnsupportedDunderAll=false
__all__ = [k for k in globals() if not k.startswith("_")]
