import pandas as pd
from fbprophet import Prophet
from fbprophet.utilities import regressor_coefficients
from sklearn.base import BaseEstimator, RegressorMixin

from sklearn_ts.models.base import TimeSeriesModel


class ProphetModel(BaseEstimator, RegressorMixin, TimeSeriesModel):
    # https://facebook.github.io/prophet/docs/quick_start.html#python-api

    def __init__(self, target='new_cases', features=['date'], regressors=[], **kwargs):
        self.target = target
        self.features = features
        self.regressors = regressors
        self.kwargs = kwargs

        self.model = None
        self.predictions = None
        self.feature_importances_ = None

    def fit(self, X, y):
        df = pd.DataFrame(X, columns=self.features)

        # Necessary for prophet algo:
        df['ds'] = df['date']
        df['y'] = y.values

        m = Prophet(**self.kwargs)

        for regressor in self.regressors:
            m.add_regressor(regressor)

        m.fit(df)

        self.model = m
        return self

    def predict(self, X):
        df = pd.DataFrame(X, columns=self.features)
        df['ds'] = df['date']  # necessary for prophet
        predictions = self.model.predict(df)
        self.predictions = predictions[['ds', 'yhat_lower', 'yhat_upper']].\
            rename(columns={'yhat_lower': 'pi_lower', 'yhat_upper': 'pi_upper'})
        return predictions['yhat'].values

    def get_params(self, deep=True):
        return {
            **{"target": self.target, 'regressors': self.regressors, 'features': self.features},
            **self.kwargs
        }

    def set_params(self, **parameters):
        for parameter, value in parameters.items():
            setattr(self, parameter, value)
        return self

    def get_features(self):
        if len(self.regressors) > 1:
            return pd.DataFrame({
                'impact': regressor_coefficients(self.model)['coef'].tolist(),
                'feature': self.features[1:],  # first one is date
            })
