"""The stressmodels module contains all the stressmodels that available in
Pastas.

Supported Stressmodels
----------------------
The following stressmodels are supported and tested:

- StressModel
- StressModel2
- FactorModel
- StepModel
- WellModel

All other stressmodels are for research purposes only and are not (yet)
fully supported and tested.

TODO
----
- Test and support StepModel
- Test and support LinearTrend

"""

from importlib import import_module
from logging import getLogger

import numpy as np
import pandas as pd
from scipy.signal import fftconvolve

from .decorators import set_parameter
from .rfunc import One, Exponential, HantushWellModel
from .timeseries import TimeSeries
from .utils import validate_name

logger = getLogger(__name__)

__all__ = ["StressModel", "StressModel2", "Constant", "StepModel",
           "LinearTrend", "FactorModel", "RechargeModel"]


class StressModelBase:
    """StressModel Base class called by each StressModel object.

    Attributes
    ----------
    name : str
        Name of this stressmodel object. Used as prefix for the parameters.
    parameters : pandas.DataFrame
        Dataframe containing the parameters.

    """
    _name = "StressModelBase"

    def __init__(self, rfunc, name, tmin, tmax, up, meanstress, cutoff):
        self.rfunc = rfunc(up, meanstress, cutoff)
        self.parameters = pd.DataFrame(
            columns=['initial', 'pmin', 'pmax', 'vary', 'name'])
        self.name = validate_name(name)
        self.tmin = tmin
        self.tmax = tmax
        self.freq = None
        self.stress = []

    @property
    def nparam(self):
        return self.parameters.index.size

    def set_init_parameters(self):
        """Set the initial parameters (back) to their default values.

        """
        pass

    @set_parameter
    def set_initial(self, name, value):
        """Internal method to set the initial parameter value.

        Notes
        -----
        The preferred method for parameter setting is through the model.

        """
        self.parameters.loc[name, 'initial'] = value

    @set_parameter
    def set_pmin(self, name, value):
        """Internal method to set the lower bound of the parameter value.

        Notes
        -----
        The preferred method for parameter setting is through the model.

        """
        self.parameters.loc[name, 'pmin'] = value

    @set_parameter
    def set_pmax(self, name, value):
        """Internal method to set the upper bound of the
 parameter value.

        Notes
        -----
        The preferred method for parameter setting is through the model.

        """
        self.parameters.loc[name, 'pmax'] = value

    @set_parameter
    def set_vary(self, name, value):
        """I
nternal method to set if the parameter is varied during
        optimization.

        Notes
        -----
        The preferred method for parameter setting is through the model.

        """
        self.parameters.loc[name, 'vary'] = bool(value)

    def update_stress(self, **kwargs):
        """Method to update the settings of the individual TimeSeries.

        Notes
        -----
        For the individual options for the different settings please refer to
        the docstring from the TimeSeries.update_series() method.

        See Also
        --------
        ps.TimeSeries.update_series

        """
        for stress in self.stress:
            stress.update_series(**kwargs)

        if "freq" in kwargs:
            self.freq = kwargs["freq"]

    def handle_stress(self, stress, settings):
        """Method to handle user provided stress in init

        Parameters
        ----------
        stress: pandas.Series, pastas.TimeSeries or iterable
        settings: dict or iterable

        Returns
        -------
        stress: dict
            dictionary with strings

        """
        data = []

        if isinstance(stress, pd.Series):
            data.append(TimeSeries(stress, settings))
        elif isinstance(stress, dict):
            for i, value in enumerate(stress.values()):
                data.append(TimeSeries(value, settings=settings[i]))
        elif isinstance(stress, list):
            for i, value in enumerate(stress):
                data.append(TimeSeries(value, settings=settings[i]))
        else:
            logger.warning("provided stress format is unknown. Provide a"
                           "Series, dict or list.")
        return data

    def dump_stress(self, series=True):
        """Method to dump all stresses in the stresses list.

        Parameters
        ----------
        series: Boolean
            True if time series are to be exported, False if only the name
            of the time series are needed. Settings are always exported.

        Returns
        -------
        data: dict
            dictionary with the dump of the stresses.

        """
        data = []

        for stress in self.stress:
            data.append(stress.to_dict(series=series))

        return data

    def get_stress(self, p=None, tmin=None, tmax=None, freq=None,
                   istress=None, **kwargs):
        """Returns the stress or stresses of the time series object as a pandas
        DataFrame.

        If the time series object has multiple stresses each column
        represents a stress.

        Returns
        -------
        stress: pd.Dataframe
            Pandas dataframe of the stress(es)

        """
        return self.stress[0].series

    def to_dict(self, series=True):
        """Method to export the StressModel object.

        Returns
        -------
        data: dict
            dictionary with all necessary information to reconstruct the
            StressModel object.

        """
        data = {
            "stressmodel": self._name,
            "name": self.name,
            "stress": self.dump_stress(series)
        }
        return data

    def get_nsplit(self):
        """Determine in how many timeseries the contribution can be splitted"""
        if hasattr(self, 'nsplit'):
            return self.nsplit
        else:
            return len(self.stress)

    def get_block(self, p, dt, tmin, tmax):
        """Internal method to get the block-response from the respnse function"""
        if tmin is not None and tmax is not None:
            day = pd.to_timedelta(1, 'd')
            maxtmax = (pd.Timestamp(tmax) - pd.Timestamp(tmin)) / day
        else:
            maxtmax = None
        b = self.rfunc.block(p, dt, maxtmax=maxtmax)
        return b


class StressModel(StressModelBase):
    """Time series model consisting of the convolution of one stress with one
    response function.

    Parameters
    ----------
    stress: pandas.Series
        pandas Series object containing the stress.
    rfunc: rfunc class
        Response function used in the convolution with the stress.
    name: str
        Name of the stress.
    up: Boolean or None, optional
        True if response function is positive (default), False if negative.
        None if you don't want to define if response is positive or negative.
    cutoff: float, optional
        float between 0 and 1 to determine how long the response is (default
        is 99% of the actual response time). Used to reduce computation times.
    settings: dict or str, optional
        The settings of the stress. This can be a string referring to a
        predefined settings dict, or a dict with the settings to apply.
        Refer to the docstring of pastas.Timeseries for further information.
    metadata: dict, optional
        dictionary containing metadata about the stress. This is passed onto
        the TimeSeries object.
    meanstress: float, optional
        The mean stress determines the initial parameters of rfunc. The initial
        parameters are chosen in such a way that the gain of meanstress is 1.

    Examples
    --------
    >>> import pastas as ps
    >>> import pandas as pd
    >>> sm = ps.StressModel(stress=pd.Series(), rfunc=ps.Gamma, name="Prec", \
                            settings="prec")

    See Also
    --------
    pastas.rfunc
    pastas.timeseries.TimeSeries

    """
    _name = "StressMod