# BSD 2-CLAUSE LICENSE

# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:

# Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# #ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# original author: Reza Hosseini
"""Functions to impute missing data using past data."""
import numpy as np

from greykite.common.features.timeseries_lags import build_agg_lag_df


def impute_with_lags(
        df,
        value_col,
        orders,
        agg_func=np.mean,
        iter_num=1):
    """A function to impute timeseries values (given in ``df``) and in ``value_col``
    with chosen lagged values or an aggregated of those.
    For example for daily data one could use the 7th lag to impute using the
    value of the same day of past week as opposed to the closest value available
    which can be inferior for business related timeseries.

    The imputation can be done multiple times by specifying ``iter_num``
    to decrease the number of missing in some cases.
    Note that there are no guarantees to impute all missing values with this
    method by design.
    However the original number of missing values and the final number of missing
    values are returned by the function along with the imputed dataframe.

    Parameters
    ----------
    df : `pandas.DataFrame`
        Input dataframe which must include `value_col` as a column.
    value_col : `str`
        The column name in ``df`` representing the values of the timeseries.
    orders : list of `int`
        The lag orders to be used for aggregation.
    agg_func : callable, default `np.mean`
        `pandas.Series` -> `float`
        An aggregation function to aggregate the chosen lags.
    iter_num : `int`, default `1`
        Maximum number of iterations to impute the series.
        Each iteration represent an imputation of the series using the provided
        lag orders (``orders``) and return an imputed dataframe.
        It might be the case that with one iterations some values
        are not imputed but with more iterations one can achieve more
        imputed values.

    Returns
    -------
    impute_info : `dict`
        A dictionary with following items:

        "df" : `pandas.DataFrame`
            A dataframe with the imputed values.
        "initial_missing_num" : `int`
            Initial number of missing values.
        "final_missing_num" : `int`
            Final number of missing values after imputations.
    """
    df = df.copy()
    nan_index = df[value_col].isnull()
    missing_num = sum(nan_index)
    initial_missing_num = missing_num

    def nan_agg_func(x):
        """To avoid warnings due to aggregation performed on sub-series
        with only NAs, we define an internal modified version of ``agg_func``
        """
        x = x.dropna()
        if len(x) == 0:
            return np.nan
        else:
            return agg_func(x)

    i = 0
    while i < iter_num and missing_num > 0:
        agg_lag_info = build_agg_lag_df(
            value_col=value_col,
            df=df,
            orders_list=[orders],
            interval_list=[],
            agg_func=nan_agg_func,
            agg_name="avglag")

        agg_lag_df = agg_lag_info["agg_lag_df"]
        # In this case the lagged dataframe will have only one column
        # which can be extracted from the list given in the ``"col_names"`` item
        agg_col_name = agg_lag_info["col_names"][0]
        df.loc[nan_index, value_col] = agg_lag_df.loc[nan_index, agg_col_name]
        nan_index = df[value_col].isnull()
        missing_num = sum(nan_index)
        i += 1

    return {
        "df": df,
        "initial_missing_num": initial_missing_num,
        "final_missing_num": missing_num}


def impute_with_lags_multi(
        df,
        orders,
        agg_func=np.mean,
        iter_num=1,
        cols=None):
    """Imputes every column of ``df`` using
    `~greykite.common.features.timeseries_impute.impute_with_lags`.

    Parameters
    ----------
    df : `pandas.DataFrame`
        Input dataframe which must include `value_col` as a column.
    orders : list of `int`
        The lag orders to be used for aggregation.
    agg_func : callable, default `np.mean`
        `pandas.Series` -> `float`
        An aggregation function to aggregate the chosen lags.
    iter_num : `int`, default `1`
        Maximum number of iterations to impute the series.
        Each iteration represent an imputation of the series using the provided
        lag orders (``orders``) and return an imputed dataframe.
        It might be the case that with one iterations some values
        are not imputed but with more iterations one can achieve more
        imputed values.
    cols : `list` [`str`] or None, default None
        Which columns to impute. If None, imputes all columns.

    Returns
    -------
    impute_info : `dict`
        A dictionary with following items:

        "df" : `pandas.DataFrame`
            A dataframe with the imputed values.
        "missing_info" : `dict`
            Dictionary with information about the missing info.

            Key = name of a column in ``df``
            Value = dictionary containing:

                "initial_missing_num" : `int`
                    Initial number of missing values.
                "final_missing_num" : `int`
                    Final number of missing values after imputation.
    """
    if cols is None:
        cols = df.columns
    impute_info = {
        "df": df.copy(),
        "missing_info": {}
    }
    for col in cols:
        # updates each column with the imputed values
        col_impute_info = impute_with_lags(
            df=df,
            value_col=col,
            orders=orders,
            agg_func=agg_func,
            iter_num=iter_num,
        )
        impute_info["df"][col] = col_impute_info["df"][col]
        impute_info["missing_info"][col] = {
            "initial_missing_num": col_impute_info["initial_missing_num"],
            "final_missing_num": col_impute_info["final_missing_num"],
        }
    return impute_info
