#
# Copyright 2024 Dan J. Bower
#
# This file is part of Atmodeller.
#
# Atmodeller is free software: you can redistribute it and/or modify it under the terms of the GNU
# General Public License as published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Atmodeller is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without
# even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with Atmodeller. If not,
# see <https://www.gnu.org/licenses/>.
#
"""Package level variables and initialises the package logger"""

__version__: str = "0.9.1"

import logging
import os

import jax
import jax.numpy as jnp
import numpy as np

# from beartype.claw import beartype_this_package
# beartype_this_package()

try:
    from typing import override as _override  # type: ignore valid for Python 3.12+
except ImportError:
    from typing_extensions import override as _override  # Python 3.11 and earlier

override = _override


jax.config.update("jax_enable_x64", True)
jnp.set_printoptions(precision=15)  # For better clarity in printed output
print("Atmodeller initialized with double precision (float64)")

# For debugging
# jax.config.update("jax_debug_nans", True)
# jax.config.update("jax_debug_infs", True)
# jax.config.update("jax_disable_jit", True)
# jax.config.update("jax_log_compiles", True)

# This prevents error_if from throwing an error when encountering NaN or Inf values. To actually
# find the root cause of NaN or Inf values, you should set this to "raise" or "breakpoint" as per
# https://docs.kidger.site/equinox/api/errors/
os.environ["EQX_ON_ERROR"] = "raise"

# Suppress warnings (notably from Equinox about static JAX arrays)
# if not sys.warnoptions:
#     import warnings

#     warnings.simplefilter("ignore")

# Thermodynamic standard state
TEMPERATURE_REFERENCE: float = 298.15
"""Enthalpy reference temperature in K (:math:`T_r` in the JANAF tables) :cite:p:`MZG02,Cha98`"""
PRESSURE_REFERENCE: float = 1.0
"""Standard state pressure in bar"""
GAS_STATE: str = "g"
"""Suffix to identify gases as per JANAF convention for the state of aggregation"""

# Initial solution guess
INITIAL_LOG_NUMBER_DENSITY: float = 50.0
"""Initial log number density

Empiricially determined. This value is mid-range for Earth-like planets.
"""
INITIAL_LOG_STABILITY: float = -30.0
"""Initial log stability.

Empirically determined.
"""

# Maximum x for which exp(x) is finite in 64-bit precision (to prevent overflow)
max_exp_input = np.log(np.finfo(np.float64).max)
# Minimum x for which exp(x) is non-zero in 64-bit precision
min_exp_input = np.log(jnp.finfo(np.float64).tiny)

# Lower and upper bounds on the hypercube which contains the root
LOG_NUMBER_DENSITY_LOWER: float = -170.0
"""Lower log number density for a species

For a gas species this corresponds to 3.17E-77 bar and 3.16E-78 bar at 3000 K and 298 K,
respectively.
"""
LOG_NUMBER_DENSITY_UPPER: float = 80.0
"""Upper log number density for a species

For a gas species this corresponds to 2294896 GPa and 227960 GPa at 3000 K and 298 K,
respectively. However, the choice of this upper limit is actually motivated by condensed species.
"""
LOG_STABILITY_LOWER: float = -700.0  # basically the same as min_exp_input
"""Lower stability for a species

Derived to ensure that the exponential function exp(x) does not underflow to zero
"""
LOG_STABILITY_UPPER: float = 35.0
"""Upper stability for a species

Empirically determined.
"""
TAU_MAX: float = 1.0e-3
"""Maximum tau scaling factor for species stability when using the tau cascade solver"""
TAU: float = 1.0e-25
"""Desired (i.e. final/minimium) tau scaling factor for species stability :cite:p:`LKK16`.

Tau effectively controls the minimum non-zero number density of unstable species. Formally, it
defines the number density of an unstable pure condensate with an activity of 1/e, which
corresponds to a log stability of zero.

This value is typically appropriate for condensate stability only, but if you additionally apply 
stability criteria to gas species you should reduce this value, maybe as low as 1e-60 to 1e-72 if
you want to ensure you do not truncated O2 at low temperatures. Hence you can override this default
using an argument to :class:`atmodeller.classes.InteriorAtmosphere`.
"""
TAU_NUM: int = 2
"""Number of tau values to solve between TAU_MAX and TAU (inclusive) for the tau cascade solver

Empirically determined. Basically, once a solution has been found for TAU_MAX the solver can
immediately proceed to TAU. This usually solves within a few steps on the first attempt.
"""

# Create the package logger.
# https://docs.python.org/3/howto/logging.html#library-config
logger: logging.Logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


def complex_formatter() -> logging.Formatter:
    """Complex formatter"""
    fmt: str = "[%(asctime)s - %(name)-30s - %(lineno)03d - %(levelname)-9s - %(funcName)s()]"
    fmt += " - %(message)s"
    datefmt: str = "%Y-%m-%d %H:%M:%S"
    formatter: logging.Formatter = logging.Formatter(fmt, datefmt=datefmt)

    return formatter


def simple_formatter() -> logging.Formatter:
    """Simple formatter for logging

    Returns:
        Formatter for logging
    """
    fmt: str = "[%(asctime)s - %(name)-30s - %(levelname)-9s] - %(message)s"
    datefmt: str = "%H:%M:%S"
    formatter: logging.Formatter = logging.Formatter(fmt, datefmt=datefmt)

    return formatter


def debug_logger() -> logging.Logger:
    """Sets up debug logging to the console.

    Returns:
        A logger
    """
    package_logger: logging.Logger = logging.getLogger(__name__)
    package_logger.setLevel(logging.DEBUG)
    package_logger.handlers = []
    console_handler: logging.Handler = logging.StreamHandler()
    console_formatter: logging.Formatter = simple_formatter()
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)

    return package_logger


def debug_file_logger() -> logging.Logger:
    """Sets up info logging to the console and debug logging to a file.

    Returns:
        A logger
    """
    # Console logger
    package_logger: logging.Logger = logging.getLogger(__name__)
    package_logger.setLevel(logging.DEBUG)
    package_logger.handlers = []
    console_handler: logging.Handler = logging.StreamHandler()
    console_formatter: logging.Formatter = simple_formatter()
    console_handler.setFormatter(console_formatter)
    console_handler.setLevel(logging.INFO)
    package_logger.addHandler(console_handler)
    # File logger
    file_handler: logging.Handler = logging.FileHandler(f"{__package__}.log")
    file_formatter: logging.Formatter = complex_formatter()
    file_handler.setFormatter(file_formatter)
    file_handler.setLevel(logging.DEBUG)
    package_logger.addHandler(file_handler)

    return package_logger


from atmodeller.classes import InteriorAtmosphere  # noqa: E402, F401
from atmodeller.containers import (  # noqa: E402, F401
    Planet,
    SolverParameters,
    Species,
    SpeciesCollection,
)
from atmodeller.utilities import (  # noqa: E402, F401
    bulk_silicate_earth_abundances,
    earth_oceans_to_hydrogen_mass,
)
