"""
Soil Moisture Accounting Procedure Monthly (SMAPm) model implementation in Python.
"""

import numpy as np
import pandas as pd

from scipy.optimize import differential_evolution
from spotpy.objectivefunctions import kge

CONVERSION_FACTOR = 2630  # 30.44 days/month * 86.4 s/mm/day


class SmapM:
    """
    Monthly SMAP (Soil Moisture Accounting Procedure) model.

    The monthly SMAP model is a simplified lumped rainfall-runoff model based
    on conceptual reservoirs. It consists of two reservoirs: soil and subsurface.

    Attributes
    ----------
    Str : float
        Soil saturation (mm). Range: 400-5000 mm. Default is 1000.
    Pes : float
        Exponent for soil evaporation. Range: 1-10. Default is 1.
    Crec : float
        Recharge coefficient. Range: 0-1. Default is 0.5.
    kkt : float
        Base flow recession coefficient (months). Range: 1-6. Default is 1.5.
    Tuin : float
        Initial soil moisture content (fraction). Range: 0-1. Default is 0.5.
    Ebin : float
        Initial base flow (m³/s). Default is 0.1.
    Ad : float
        Drainage area (km²). Default is 1.
    """

    def __init__(
        self,
        Str: float = 1000,  # 400 - 5000 mm
        Pes: float = 1,  # 1 - 10
        Crec: float = 0.5,  # 0 - 1
        kkt: float = 1.5,  # 1 - 6 meses
        Tuin: float = 0.5,  # 0 - 1
        Ebin: float = 0.1,  # m3/s
        Ad: float = 1,  # km2
    ):
        """
        Initialize an instance of the SmapM class.

        Parameters
        ----------
        Str : float, optional
            Soil saturation (mm). Default is 1000.
        Pes : float, optional
            Exponent for soil evaporation. Default is 1.
        Crec : float, optional
            Recharge coefficient. Default is 0.5.
        kkt : float, optional
            Base flow recession coefficient (months). Default is 1.5.
        Tuin : float, optional
            Initial soil moisture content (fraction). Default is 0.5.
        Ebin : float, optional
            Initial base flow (m³/s). Default is 0.1.
        Ad : float, optional
            Drainage area (km²). Default is 1.
        """
        # Inicialização das variáveis do modelo
        self.i = 0

        self.Str = Str
        self.Pes = Pes
        self.Crec = Crec
        self.kkt = kkt
        self.Tuin = Tuin
        self.Ebin = Ebin
        self.Ad = Ad

        self.Er = 0
        self.Es = 0
        self.Eb = self.Ebin / self.Ad * CONVERSION_FACTOR
        self.Rec = 0
        self.Rsolo = self.Rsolo0_calc(self.Tuin, self.Str)
        self.Rsub = self.Rsub0_calc(self.Ebin, self.kkt, self.Ad)
        self.Tu = self.Tu_calc(self.Rsolo, self.Str)

    def reset_variables(self):
        """
        Reset the model variables to their initial states.

        This method resets the step counter and recalculates the initial
        values for Eb, Rsolo, Rsub, and Tu based on the initial conditions
        (Tuin, Ebin, and other parameters).
        """
        self.i = 0
        self.Eb = self.Ebin / self.Ad * CONVERSION_FACTOR
        self.Rsolo = self.Rsolo0_calc(self.Tuin, self.Str)
        self.Rsub = self.Rsub0_calc(self.Ebin, self.kkt, self.Ad)
        self.Tu = self.Tu_calc(self.Rsolo, self.Str)

    def __str__(self) -> str:
        """
        Returns a string representation of the SmapM object.

        Returns
        -------
        str
            A string that includes the values of all parameters used in
            the model.
        """
        return (
            "SmapM Class Object\n"
            f"Str = {self.Str}\n"
            f"Pes = {self.Pes}\n"
            f"Crec = {self.Crec}\n"
            f"kkt = {self.kkt}\n"
            f"Ad = {self.Ad}\n"
            f"Tuin = {self.Tuin}\n"
            f"Ebin = {self.Ebin}"
        )

    def bounds(self) -> dict:
        """
        Returns the bounds for various hydrological parameters.

        Returns
        -------
        dict
            A dictionary containing the bounds for the following parameters:
            - "Str" : tuple of float
            Saturation (mm), range (400, 5000).
            - "Pes" : tuple of float
            Exponent for evaporation, range (1, 10).
            - "Crec" : tuple of float
            Recharge coefficient, range (0, 1).
            - "kkt" : tuple of float
            Base flow recession coefficient (months), range (1, 6).
        """
        return {
            "Str": (400, 5000),
            "Pes": (1, 10),
            "Crec": (0, 1),
            "kkt": (1, 6),
        }

    def check_bounds(self, params: dict) -> bool:
        """
        Check if parameters are within their acceptable bounds.

        Parameters
        ----------
        params : dict
            A dictionary containing parameter names as keys and their values.

        Returns
        -------
        bool
            True if all parameters are within bounds, False otherwise.
        """
        for param, value in params.items():
            if param in self.bounds():
                lower, upper = self.bounds()[param]
                if not (lower <= value <= upper):
                    return False
        return True

    # region: Funções de reservatórios
    def Rsolo_calc(self, Rsolo, P, Es, Er, Rec) -> float:
        """
        Calculate updated soil reservoir value.

        Parameters
        ----------
        Rsolo : float
            Current soil reservoir value (mm).
        P : float
            Precipitation input (mm).
        Es : float
            Soil evaporation (mm).
        Er : float
            Actual evapotranspiration (mm).
        Rec : float
            Recharge to subsurface reservoir (mm).

        Returns
        -------
        float
            Updated soil reservoir value (mm).
        """
        Rsolo = Rsolo + P - Es - Er - Rec
        return max(0, Rsolo)

    def Rsub_calc(self, Rsub, Rec, Eb) -> float:
        """
        Calculate updated subsurface reservoir value.

        Parameters
        ----------
        Rsub : float
            Current subsurface reservoir value (mm).
        Rec : float
            Recharge from soil reservoir (mm).
        Eb : float
            Base flow (mm).

        Returns
        -------
        float
            Updated subsurface reservoir value (mm).
        """
        Rsub = Rsub + Rec - Eb
        return Rsub

    def Rsolo0_calc(self, Tuin, Str) -> float:
        """
        Calculate initial soil reservoir value.

        Parameters
        ----------
        Tuin : float
            Initial soil moisture content (fraction, 0-1).
        Str : float
            Soil saturation capacity (mm).

        Returns
        -------
        float
            Initial soil reservoir value (mm).
        """
        Rsolo0 = Tuin * Str
        return Rsolo0

    def Rsub0_calc(self, Ebin, kkt, Ad) -> float:
        """
        Calculate initial subsurface reservoir value.

        Parameters
        ----------
        Ebin : float
            Initial base flow (m³/s).
        kkt : float
            Base flow recession coefficient (months).
        Ad : float
            Drainage area (km²).

        Returns
        -------
        float
            Initial subsurface reservoir value (mm).
        """
        Rsub0 = Ebin / (1 - 0.5 ** (1 / kkt)) / Ad * CONVERSION_FACTOR
        return Rsub0

    def Es_calc(self, P, Tu, Pes) -> float:
        """
        Calculate soil evaporation.

        Parameters
        ----------
        P : float
            Precipitation (mm).
        Tu : float
            Soil moisture content (fraction, 0-1).
        Pes : float
            Exponent for soil evaporation.

        Returns
        -------
        float
            Soil evaporation (mm).
        """
        return P * Tu**Pes

    def Er_calc(self, Tu, Ep) -> float:
        """
        Calculate actual evapotranspiration.

        Parameters
        ----------
        Tu : float
            Soil moisture content (fraction, 0-1).
        Ep : float
            Potential evapotranspiration (mm).

        Returns
        -------
        float
            Actual evapotranspiration (mm).
        """
        return Tu * Ep

    def Eb_calc(self, kkt, Rsub) -> float:
        """
        Calculate base flow from subsurface reservoir.

        Parameters
        ----------
        kkt : float
            Base flow recession coefficient (months).
        Rsub : float
            Subsurface reservoir value (mm).

        Returns
        -------
        float
            Base flow (mm).
        """
        return (1 - 0.5 ** (1 / kkt)) * Rsub

    def Rec_calc(self, Crec, Tu, Rsolo) -> float:
        """
        Calculate recharge to subsurface reservoir.

        Parameters
        ----------
        Crec : float
            Recharge coefficient (0-1).
        Tu : float
            Soil moisture content (fraction, 0-1).
        Rsolo : float
            Soil reservoir value (mm).

        Returns
        -------
        float
            Recharge to subsurface reservoir (mm).
        """
        Tu = max(0, min(1, Tu))
        return Crec * (Tu**4) * Rsolo

    def Tu_calc(self, Rsolo, Str) -> float:
        """
        Calculate soil moisture content.

        Parameters
        ----------
        Rsolo : float
            Soil reservoir value (mm).
        Str : float
            Soil saturation capacity (mm).

        Returns
        -------
        float
            Soil moisture content (fraction, 0-1).
        """
        return max(0, min(1, Rsolo / Str))

    def Q_bas(self, Eb, Ad) -> float:
        """
        Calculate base flow discharge.

        Parameters
        ----------
        Eb : float
            Base flow (mm).
        Ad : float
            Drainage area (km²).

        Returns
        -------
        float
            Base flow discharge (m³/s).
        """
        return (Eb) * Ad / CONVERSION_FACTOR

    def Q_calc(self, Es, Eb, Ad) -> float:
        """
        Calculate total discharge at the outlet.

        Parameters
        ----------
        Es : float
            Soil evaporation (mm).
        Eb : float
            Base flow (mm).
        Ad : float
            Drainage area (km²).

        Returns
        -------
        float
            Total discharge at the outlet (m³/s).
        """
        return (Es + Eb) * Ad / CONVERSION_FACTOR

    # endregion
    def run_step(self, prec, etp, reset=False) -> float:
        """
        Execute a single time step of the hydrological model.

        Parameters
        ----------
        prec : float
            Precipitation value for the current step (mm).
        etp : float
            Potential evapotranspiration value for the current step (mm).
        reset : bool, optional
            If True, resets the model state to initial conditions.
            Default is False.

        Returns
        -------
        float
            The calculated discharge value for the current step (m³/s).
        """
        if reset or self.i == 0:
            self.reset_variables()
            self.Tu = self.Tuin

        self.Er = self.Er_calc(self.Tu, etp)
        self.Es = self.Es_calc(prec, self.Tu, self.Pes)
        self.Eb = self.Eb_calc(self.kkt, self.Rsub)
        self.Rec = self.Rec_calc(self.Crec, self.Tu, self.Rsolo)
        self.Rsolo = self.Rsolo_calc(self.Rsolo, prec, self.Es, self.Er, self.Rec)
        self.Rsub = self.Rsub_calc(self.Rsub, self.Rec, self.Eb)
        self.Tu = self.Tu_calc(self.Rsolo, self.Str)

        self.i += 1

        return self.Q_calc(self.Es, self.Eb, self.Ad)

    def run(self, prec_arr, etp_arr, reset=True):
        """
        Execute the hydrological model with arrays of precipitation and evapotranspiration.

        Parameters
        ----------
        prec_arr : iterable
            An iterable of precipitation values (mm).
        etp_arr : iterable
            An iterable of potential evapotranspiration values (mm).
        reset : bool, optional
            If True, resets the model state before running. Default is True.

        Yields
        ------
        float
            The discharge value for each time step (m³/s).
        """
        if reset or self.i == 0:
            self.reset_variables()

        for prec, etp in zip(prec_arr, etp_arr):
            yield self.run_step(prec, etp, reset=False)

    def run_to_list(self, prec_arr, etp_arr, reset=True):
        """
        Execute the hydrological model and return results as a list.

        Parameters
        ----------
        prec_arr : iterable
            An iterable of precipitation values (mm).
        etp_arr : iterable
            An iterable of potential evapotranspiration values (mm).
        reset : bool, optional
            If True, resets the model state before running. Default is True.

        Returns
        -------
        list of float
            A list containing the discharge values for each time step (m³/s).
        """
        return list(self.run(prec_arr, etp_arr, reset))

    def run_to_dataframe(self, prec_arr, etp_arr, reset=True):
        """
        Executes the hydrological model and stores all intermediate variables
        and results into a Pandas DataFrame.

        Parameters
        ----------
        prec_arr : iterable
            An iterable of precipitation values.
        etp_arr : iterable
            An iterable of evapotranspiration values.
        reset : bool, optional
            If True, resets the model state before running. Default is True.

        Returns
        -------
        pd.DataFrame
            A DataFrame containing all intermediate variables and results for
            each step.
        """
        if reset or self.i == 0:
            self.reset_variables()

        data = []

        for prec, etp in zip(prec_arr, etp_arr):
            discharge = self.run_step(prec, etp, reset=False)

            # Store all variables for the current step
            data.append(
                {
                    "etp": etp,
                    "prec": prec,
                    "Er": self.Er,
                    "Es": self.Es,
                    "Eb": self.Eb,
                    "Rec": self.Rec,
                    "Rsolo": self.Rsolo,
                    "Rsub": self.Rsub,
                    "Tu": self.Tu,
                    "discharge": discharge,
                }
            )

        return pd.DataFrame(data)

    def calibrate(
        self,
        prec_arr,
        etp_arr,
        eval_arr,
        variables: list[str],
        obj_func=None,
        disp=False,
    ):
        """
        Calibrate the SMAPM model using the Differential Evolution algorithm.

        This method optimizes the specified model parameters to minimize the
        given objective function. Note that the objective function will be
        minimized, so metrics like KGE and NSE should be multiplied by -1.

        Parameters
        ----------
        prec_arr : iterable
            An iterable of precipitation values.
        etp_arr : iterable
            An iterable of evapotranspiration values.
        eval_arr : iterable
            An iterable of observed values.
        variables : list of str
            A list of variable names to be optimized.
        obj_func : callable, optional
            A callable function that takes observed and simulated values as
            input and returns a float to be minimized. If None, the KGE
            objective function will be used. The function signature should be
            `obj_func(observed, simulated)`.
        disp : bool, optional
            Whether to display optimization progress. Default is False.

        Returns
        -------
        scipy.optimize.OptimizeResult
            The result of the optimization process, as returned by
            `scipy.optimize.differential_evolution`.
        """

        invalid_vars = [var for var in variables if var not in self.__dict__]
        if invalid_vars:
            raise ValueError(f"Variables {', '.join(invalid_vars)} do not exist.")

        bounds = [self.bounds()[var] for var in variables]

        if obj_func is None:

            def default_obj_func(obs, sim):
                return -kge(obs, sim)

            obj_func = default_obj_func

        def objective(params):
            self.__dict__.update({var: val for var, val in zip(variables, params)})

            return obj_func(eval_arr, list(self.run(prec_arr, etp_arr)))

        result = differential_evolution(func=objective, bounds=bounds, disp=disp)

        return result

    def fast_calibrate(
        self,
        prec_arr,
        etp_arr,
        eval_arr,
        variables: list[str],
        obj_func=None,
        disp=False,
    ):
        """
        Fast calibration of the SMAPM model using integer-based Differential Evolution.

        This method optimizes the specified model parameters using integer values
        (multiplied by 10000) to achieve faster calibration with 4 decimal places
        of precision. The objective function will be minimized, so metrics like
        KGE and NSE should be multiplied by -1.

        Parameters
        ----------
        prec_arr : iterable
            An iterable of precipitation values.
        etp_arr : iterable
            An iterable of evapotranspiration values.
        eval_arr : iterable
            An iterable of observed values.
        variables : list of str
            A list of variable names to be optimized.
        obj_func : callable, optional
            A callable function that takes observed and simulated values as
            input and returns a float to be minimized. If None, the KGE
            objective function will be used. The function signature should be
            `obj_func(observed, simulated)`.
        disp : bool, optional
            Display optimization progress. Default is False.

        Returns
        -------
        scipy.optimize.OptimizeResult
            The result of the optimization process, as returned by
            `scipy.optimize.differential_evolution`. The optimized parameters
            are automatically scaled back to their original range.
        """

        invalid_vars = [var for var in variables if var not in self.__dict__]
        if invalid_vars:
            raise ValueError(f"Variables {', '.join(invalid_vars)} do not exist.")

        # Multiply bounds by 10000 to work with integers
        bounds = [(round(lb * 10000), round(ub * 10000)) for lb, ub in [self.bounds()[var] for var in variables]]

        if obj_func is None:

            def default_obj_func(obs, sim):
                return -kge(obs, sim)

            obj_func = default_obj_func

        def objective(params):
            # Divide by 10000 to get the actual parameter values
            actual_params = [p / 10000 for p in params]
            self.__dict__.update({var: val for var, val in zip(variables, actual_params)})

            return obj_func(eval_arr, list(self.run(prec_arr, etp_arr)))

        # Use integrality=True for all variables
        integrality = [True] * len(variables)

        result = differential_evolution(
            func=objective,
            bounds=bounds,
            integrality=integrality,
            disp=disp
        )

        # Scale the result back to original range
        result.x = result.x / 10000

        # Update model parameters with final calibrated values
        self.__dict__.update({var: val for var, val in zip(variables, result.x)})

        return result

    def pso_calibrate(
        self,
        prec_arr,
        etp_arr,
        eval_arr,
        variables: list[str],
        obj_func=None,
        options=None,
        n_particles=50,
        iters=100,
    ):
        """
        Calibrate the SMAPM model using Particle Swarm Optimization (PSO) BETA.

        This method optimizes the specified model parameters to minimize the
        given objective function. Note that the objective function will be
        minimized, so metrics like KGE and NSE should be multiplied by -1.

        Parameters
        ----------
        prec_arr : iterable
            An iterable of precipitation values.
        etp_arr : iterable
            An iterable of evapotranspiration values.
        eval_arr : iterable
            An iterable of observed values.
        variables : list of str
            A list of variable names to be optimized.
        obj_func : callable, optional
            A callable function that takes observed and simulated values as
            input and returns a float to be minimized. If None, the KGE
            objective function will be used. The function signature should be
            `obj_func(observed, simulated)`.
        options : dict, optional
            A dictionary of options for the PSO algorithm. Default is None.
            - c1: represents the cognitive component (default is 1.5).
            - c2: represents the social component (default is 1.5).
            - w: represents the inertia weight (default is 0.5).
        n_particles : int, optional
            Number of particles in the swarm. Default is 50.
        iters : int, optional
            Number of iterations for the PSO algorithm. Default is 100.

        Returns
        -------
        dict
            A dictionary containing the best parameters and the best objective
            function value.
        """
        try:
            import pyswarms as ps
        except ImportError:
            raise ImportError(
                "pyswarms is required for PSO calibration. "
                "Please install it via 'pip install pyswarms'."
            )

        invalid_vars = [var for var in variables if var not in self.__dict__]
        if invalid_vars:
            raise ValueError(f"Variables {', '.join(invalid_vars)} do not exist.")

        for var in variables:
            if var not in self.bounds():
                raise ValueError(f"Variable '{var}' is not defined in bounds.")

        bounds = [self.bounds()[var] for var in variables]
        lower_bounds, upper_bounds = zip(*bounds)
        print(f"Lower bounds: {lower_bounds}, Upper bounds: {upper_bounds}")

        if obj_func is None:

            def default_obj_func(obs, sim):
                result = kge(obs, sim)
                if hasattr(result, "item"):
                    return -float(result.item())
                return -float(result)

            obj_func = default_obj_func

        def objective(params):
            if params.ndim > 1:
                costs = []
                for p in params:
                    self.__dict__.update({var: val for var, val in zip(variables, p)})
                    simulated = self.run_to_list(prec_arr, etp_arr)
                    cost = obj_func(eval_arr, simulated)
                    costs.append(float(cost))
                return np.array(costs)
            else:
                self.__dict__.update({var: val for var, val in zip(variables, params)})
                simulated = self.run_to_list(prec_arr, etp_arr)
                cost = obj_func(eval_arr, simulated)
                return float(cost)

        if options is None:
            options = {"c1": 1.5, "c2": 1.5, "w": 0.5}

        optimizer = ps.single.GlobalBestPSO(
            n_particles=n_particles,
            dimensions=len(variables),
            options=options,
            bounds=(lower_bounds, upper_bounds),
        )

        best_cost, best_pos = optimizer.optimize(objective, iters=iters)

        for var, val in zip(variables, best_pos):
            setattr(self, var, val)

        return {
            "best_cost": best_cost,
            "best_params": {var: val for var, val in zip(variables, best_pos)},
        }
