# import libs
import logging
from typing import Dict, Literal, Optional, Any
from pyThermoDB.core import TableEquation
import pycuc
from scipy import integrate
from pyThermoLinkDB.models.component_models import ComponentEquationSource
# locals
from ..configs.thermo_props import Cp_IG_INTEGRATION_UNIT

# NOTE: Logger
logger = logging.getLogger(__name__)


def Cp_integral(
    eq_src: ComponentEquationSource,
    T_ref: float,
    T: float,
    output_unit: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """
    Integrate the heat capacity equation from T_ref to T (J/mol).

    Parameters
    ----------
    eq_src : ComponentEquationSource
        The heat capacity equation to integrate.
    T_ref : float
        The reference temperature (lower limit of integration).
    T : float
        The target temperature (upper limit of integration).
    output_unit : Optional[str], optional
        The desired output unit for the integrated value, by default None. If None, the unit from the equation will be used.

    Returns
    -------
    Optional[Dict[str, Any]]
        The integrated heat capacity value, or None if an error occurs as:
        {
            'value': float,
            'unit': str
        }

    Notes
    -----
    output_unit is important because:
    - the heat capacity unit is defined, for example, as energy per mole per Kelvin (e.g., J/mol.K), so: the integrated value will be in energy per mole (e.g., J/mol), as a result.
    - If output_unit is None, the result will be in the unit defined by the heat capacity equation.
    """
    try:
        # SECTION: equation expression
        # NOTE: extract Cp equation
        # ! args
        eq_args = eq_src.args
        # ! arg symbols
        eq_arg_symbols = eq_src.arg_symbols
        # ! returns
        eq_returns = eq_src.returns
        # >> get return units
        returns_outer_key, returns_inner = next(
            iter(eq_returns.items()))
        eq_return_unit: str = returns_inner['unit']

        # ! return symbols
        eq_return_symbols = eq_src.return_symbols
        # ! equation
        Cp_eq = eq_src.source

        # NOTE: Cp integral
        # unit: [J/mol]
        # init unit_
        unit_ = None

        # scipy integrate method
        def integrand_1(T):
            '''
            Integral of Cp equation, unit: [J/mol]
            '''
            nonlocal unit_
            res_ = Cp_eq.cal(T=T)
            cal_ = res_.get('value', None)
            unit_ = res_.get('unit', None)  # ! get unit from equation

            # >> check
            if cal_ is None:
                logger.error(
                    f"Failed to calculate Cp for equation at T={T} K.")
                raise

            # >> check
            if (
                not isinstance(cal_, str) and
                not isinstance(cal_, float)
            ):
                logger.error(
                    f"Invalid Cp value at T={T} K: {cal_}")
                raise

            # >> check
            if unit_ is None:
                logger.error(
                    f"Equation does not have a unit for Cp at T={T} K.")
                raise

            return cal_

        # ! check integral
        function_integral = Cp_eq.body_integral

        # NOTE: calculate integral
        if function_integral:
            # calc
            # method 1 (still takes return unit from equation)
            _eq_Cp_integral = Cp_eq.cal_integral(
                T1=T_ref,
                T2=T
            )
        else:
            # calc
            # method 2 (still takes return unit from equation)
            _eq_Cp_integral, _ = integrate.quad(
                integrand_1,
                T_ref,
                T
            )

            # check
            if unit_ is None:
                logger.error(
                    f"Equation does not have a unit for Cp integral.")
                raise

            # >> check
            if _eq_Cp_integral is None:
                logger.error(
                    f"Failed to calculate Cp integral from T={T_ref} K to T={T} K.")
                return None

            # TODO: FINALLY convert Cp equation to [J/mol.K] (after integration means the unit is [J/mol])
            # ! check output unit
            if output_unit is not None:
                if (
                    output_unit == 'J/mol' or
                    output_unit == 'j/mol'
                ):
                    # >>> use `J/mol.K`` as default only for integration unit, the actual unit after integration should be `J/mol`
                    _eq_Cp_integral = pycuc.to(
                        _eq_Cp_integral,
                        f"{unit_} => J/mol.K"
                    )
                else:
                    # >>> convert to desired unit
                    _eq_Cp_integral = pycuc.to(
                        _eq_Cp_integral,
                        f"{unit_} => {output_unit}"
                    )

        return {
            'value': _eq_Cp_integral,
            'unit': output_unit if output_unit is not None else unit_
        }
    except Exception as e:
        logger.error(
            f"Error integrating Cp equation from T={T_ref} K to T={T} K: {e}")
        return None


def Cp__RT_integral(
    eq_src: ComponentEquationSource,
    T_ref: float,
    T: float,
    R: float,
    R_unit: str = 'J/mol.K',
) -> Optional[Dict[str, Any]]:
    """
    Integrate the dimensionless heat capacity equation Cp/RT from T_ref to T.

    Parameters
    ----------
    eq_src : ComponentEquationSource
        The heat capacity equation to integrate.
    T_ref : float
        The reference temperature (lower limit of integration).
    T : float
        The target temperature (upper limit of integration).
    R : float
        The universal gas constant.
    R_unit : str, optional
        The unit of the universal gas constant R, by default 'J/mol.K'.

    Returns
    -------
    Optional[Dict[str, Any]]
        The integrated dimensionless heat capacity value, or None if an error occurs.

    Notes
    -----
    output_unit is important because:
    - the heat capacity unit is defined, for example, as energy per mole per Kelvin (e.g., J/mol.K), so: the integrated value will be dimensionless when divided by RT.
    - as Cp and R normally have the same unit basis (energy per mole per Kelvin), the result of the integration Cp/RT is dimensionless.
    - so R_unit is needed to ensure compatibility between Cp and R units during calculation.
    - for better calculation performance, R is converted to the Cp unit before calculation.
    - the term Cp/RT is dimensionless, after integration, the result remains identical even if different units are used for Cp and R, as long as they are compatible.
    """
    try:
        # SECTION: equation expression
        # NOTE: extract Cp equation
        # ! args
        eq_args = eq_src.args
        # ! arg symbols
        eq_arg_symbols = eq_src.arg_symbols
        # ! returns
        eq_returns = eq_src.returns
        # >> get return units
        returns_outer_key, returns_inner = next(
            iter(eq_returns.items())
        )
        # ! return unit
        eq_return_unit: str = returns_inner['unit']
        # ! return symbols
        eq_return_symbols = eq_src.return_symbols
        # ! equation
        Cp_eq = eq_src.source

        # SECTION: Cp and R unit compatibility check
        # ! compatible unit check between Cp and R is done in the pycuc.to function
        # NOTE: compatibility unit
        # >> check Cp unit compatibility with R unit
        if eq_return_unit.lower() != R_unit.lower():
            # set R unit compatible with Cp unit
            R = pycuc.to(
                value=R,
                unit_conversion_block=f"{R_unit} => {eq_return_unit}"
            )

        # SECTION: integral [Cp/RT]
        # unit: [dimensionless]
        # init unit
        unit_ = None
        # scipy integrate method

        def integrand_0(T):
            '''
            Dimensionless integrand for Cp/RT calculation.
            '''
            # modify the nonlocal variable
            nonlocal unit_
            res_ = Cp_eq.cal(T=T)
            cal_ = res_.get('value', None)
            unit_ = res_.get('unit', None)

            # >> check
            if cal_ is None:
                logger.error(
                    f"Failed to calculate Cp for equation at T={T} K.")
                raise

            # >> check
            if not isinstance(cal_, str) and not isinstance(cal_, float):
                logger.error(
                    f"Invalid Cp value at T={T} K: {cal_}")
                raise

            # >> check
            if unit_ is None:
                logger.error(
                    f"Equation does not have a unit for Cp at T={T} K.")
                raise

            # NOTE: calculate Cp/RT
            res = cal_/(T*R)
            return res

        # ! check custom
        custom_integral = Cp_eq.custom_integral

        # SECTION: Calculate integral
        if custom_integral.get('Cp/RT', None) is not None:
            # calc
            # method 1 (custom)
            # >>> already takes return unit from equation and make a dimensionless unit
            _eq_Cp_integral_Cp__RT = Cp_eq.cal_custom_integral(
                'Cp/RT',
                T1=T_ref,
                T2=T
            )
        else:
            # calc
            # method 2 (scipy integrate)
            # >>> raise error if Cp unit is not compatible with R unit.
            _eq_Cp_integral_Cp__RT, _ = integrate.quad(
                integrand_0,
                T_ref,
                T
            )

            # >> check
            if _eq_Cp_integral_Cp__RT is None:
                logger.error(
                    f"Failed to calculate Cp/RT integral from T={T_ref} K to T={T} K.")
                return None

            # check
            if unit_ is None:
                logger.error(
                    f"Equation does not have a unit for Cp/RT integral.")
                raise

            # TODO: FINALLY convert Cp equation to [J/mol.K]
            # NOTE: the unit after integration is dimensionless, so no need to convert
            # ! very important check output unit as Cp and R have different units

        return {
            'value': _eq_Cp_integral_Cp__RT,
            'unit': 'dimensionless'
        }
    except Exception as e:
        logger.error(
            f"Error integrating Cp/RT equation from T={T_ref} K to T={T} K: {e}")
        return None
