import numpy as np
import torch
from exponential_smoothing.model import Smoothing_Model

"""
Exponential Smoothing models for time series forecasting.
All methods have been implemented from chapter 7 of the textbook as a reference.
'Hyndman, Rob J., and George Athanasopoulos. Forecasting: principles
and practice. OTexts, 2014.'
"""

class SES(Smoothing_Model):
    def __init__(self, dep_var, trend=None, seasonal=None, seasonal_periods=None, damped=False, indep_var=None, **kwargs):
        """Implementation of Simple Exponential Smoothing"""
        self.trend = None
        self.damped = False
        self.seasonal = None
        self.seasonal_periods = None

        super().set_allowed_kwargs(["alpha", "beta", "phi", "gamma"])
        
        super().__init__(dep_var, self.trend, self.seasonal, self.seasonal_periods, self.damped, **kwargs)
        
        initialize_params = ["alpha"]
        self.__initialize_params(initialize_params)

    def __initialize_params(self, initialize_params):
        """Initialize the model parameters"""
        super().initialize_params(initialize_params)

        self.initial_level = torch.tensor(self.init_components["level"])

        self.alpha, self.beta, self.gamma, self.phi = torch.tensor(list(self.params.values()), dtype=torch.float32)

    def __smooth_level(self, y, lprev):
        """Calculate level"""
        self.level =  torch.add(torch.mul(self.alpha, y),torch.mul((1 - self.alpha), lprev))
    
    def fit(self):
        """
        Fit the model to the data according to the equations:
        l_t =  alpha*y_t + (1 - alpha)*l_{t-1}
        """
        self.fitted = torch.zeros(self.dep_var.shape[0])
        for index,row in enumerate(self.dep_var):
            if index == 0:
                self.level = self.initial_level
                self.fitted[0] = row
            else:
                y_hat = self.level
                lprev = self.level
                self.__smooth_level(row, lprev)
                self.fitted[index] = y_hat  
    
    def predict(self,h):
        """
        Predict the next h values of the target variable.

        Arguments:

        *h (int): Future time-steps to forecast
        """
        self.forecast = torch.tensor([])
        for i in range(1,h+1):
            step_forecast = self.level
            self.forecast = torch.cat((self.forecast,step_forecast))
        return self.forecast


class HoltTrend(Smoothing_Model):
    def __init__(self, dep_var, trend="add", seasonal=None, seasonal_periods=None, damped=False, indep_var=None, **kwargs):
        """Implementation of Holt's Trend Method"""
        self.trend = "add"
        self.damped = damped
        self.seasonal = None
        self.seasonal_periods = None

        super().set_allowed_kwargs(["alpha", "beta", "phi", "gamma"])
        super().__init__(dep_var, self.trend, self.seasonal, self.seasonal_periods, self.damped, **kwargs)
        
        initialize_params = ["alpha", "beta"]
        if self.damped:
            initialize_params.append("phi")

        self.__initialize_params(initialize_params)


    def __initialize_params(self, initialize_params):
        """Initialize the model parameters"""

        super().initialize_params(initialize_params)

        self.initial_level = torch.tensor(self.init_components["level"])
        self.initial_trend = torch.tensor(self.init_components["trend"])

        self.alpha, self.beta, self.gamma, self.phi = torch.tensor(list(self.params.values()), dtype=torch.float32)

    def __smooth_level(self, y, lprev, bprev):
        """Calculate level"""
        self.level = torch.mul(torch.sub(1,self.alpha),torch.add(lprev, torch.mul(self.phi, bprev)))
        self.level = torch.add(torch.mul(self.alpha, y), self.level)

    def __smooth_trend(self, lprev, bprev):
        "Calculate trend"
        self.trend = torch.mul(self.beta, torch.sub(self.level, lprev))
        self.trend = torch.add(self.trend, torch.mul(torch.sub(1, self.beta), torch.mul(self.phi, bprev)))

    def fit(self):
        """
        Fit the model to the data according to the equations:
        l_t = alpha * y_t + (1 - alpha) * (l_{t-1} + b_{t-1})
        b_t = beta * (l_t - l_{t-1}) + (1 - beta)*phi*b_{t-1}
        """
        
        self.fitted = torch.zeros(size=self.dep_var.shape)
        for index, row in enumerate(self.dep_var):
            if index == 0:
                self.level = self.initial_level
                self.trend = self.initial_trend
                self.fitted[0] = row[0]
            else:
                y_hat = torch.add(self.level, torch.mul(self.phi, self.trend))
                lprev, bprev = self.level, self.trend
                self.__smooth_level(row, lprev, bprev)
                self.__smooth_trend(lprev, bprev)
                self.fitted[index] = y_hat
        

    def predict(self, h):
        """
        Predict the next h values of the target variable.

        Arguments:

        *h (int): Future time-steps to forecast
        """
        
        self.forecast = torch.tensor([])
        for i in range(1,h+1):

            damp_factor = torch.sum(torch.tensor([torch.pow(self.phi, x+1) for x in range(i)]))

            step_forecast = torch.add(self.level, torch.mul(damp_factor, self.trend))
            self.forecast = torch.cat((self.forecast, step_forecast))
            
        return self.forecast
    
class HoltWinters(Smoothing_Model):
    def __init__(self, dep_var, trend="add", seasonal="add", seasonal_periods=4, damped=False, indep_var=None, **kwargs):
        "Implementation of Holt-Winters' Seasonal Method"
        self.trend = trend
        self.damped = damped
        self.seasonal = seasonal
        self.seasonal_periods = seasonal_periods

        super().set_allowed_kwargs(["alpha", "beta", "phi", "gamma"])
        super().__init__(dep_var, self.trend, self.seasonal, self.seasonal_periods, self.damped, **kwargs)
        
        initialize_params = ["alpha", "beta", "gamma"]
        if self.damped:
            initialize_params.append("phi")

        self.__initialize_params(initialize_params)

    def __initialize_params(self, initialize_params):
        """Initialize the model parameters"""

        super().initialize_params(initialize_params)

        self.initial_level = torch.tensor(self.init_components["level"])
        self.initial_trend = torch.tensor(self.init_components["trend"])
        self.initial_seasonals = torch.tensor(self.init_components["seasonal"])

        self.alpha, self.beta, self.gamma, self.phi = torch.tensor(list(self.params.values()), dtype=torch.float32)


    def __smooth_level(self, y, lprev, bprev, seasonal):
        """Calculate level"""

        if self.seasonal == "add":

            self.level = torch.mul(torch.sub(1,self.alpha),torch.add(lprev, torch.mul(self.phi, bprev)))
            self.level = torch.add(torch.mul(self.alpha, torch.sub(y, seasonal)), self.level)
        
        elif self.seasonal == "mul":

            self.level = torch.mul(torch.sub(1,self.alpha),torch.add(lprev, torch.mul(self.phi, bprev)))
            self.level = torch.add(torch.mul(self.alpha, torch.divide(y, seasonal)), self.level)


    def __smooth_trend(self, lprev, bprev):
        """Calculate trend"""

        self.trend = torch.mul(self.beta, torch.sub(self.level, lprev))
        self.trend = torch.add(self.trend, torch.mul(torch.sub(1, self.beta), torch.mul(self.phi, bprev)))

    def __smooth_seasonal(self,  y, lprev, bprev, seasonal):
        '''This function calculates the smoothed trend for a given beta and level values'''
        
        if self.seasonal == "add":

            seasonal = torch.mul(torch.sub(1,self.gamma), seasonal)
            seasonal = torch.add(seasonal, torch.mul(self.gamma, torch.sub(torch.sub(y, lprev), torch.mul(self.phi, bprev))))
            self.seasonals = torch.cat((self.seasonals, seasonal))

        elif self.seasonal == "mul":

            seasonal = torch.mul(torch.sub(1,self.gamma), seasonal)
            seasonal = torch.add(seasonal, torch.mul(self.gamma, torch.divide(y, torch.add(lprev,torch.mul(self.phi, bprev)))))
            self.seasonals = torch.cat((self.seasonals, seasonal))


    def fit(self):
        """
        Fit the model to the data according to the equations:
        
        If seasonal is additive:
        l_t = alpha*(y_t - s_{t-m}) + (1 - alpha)*(l_{t-1} + b_{t-1})
        b_t = beta * (l_t - l_{t-1}) + (1 - beta)*phi*b_{t-1}
        s_t = gamma*(y_t - l_{t-1} - phi*b_{t-1}) + (1 - y)*s_{t-m}

        If seasonal is multiplicative:
        l_t = alpha*(y_t/s_{t-m}) + (1 - alpha)*(l_{t-1} + b_{t-1})
        b_t = beta * (l_t - l_{t-1}) + (1 - beta)*phi*b_{t-1}
        s_t = gamma*(y_t/(l_{t-1} + phi*b_{t-1})) + (1 - y)*s_{t-m}
        """
        
        self.fitted = torch.zeros(size=self.dep_var.shape)

        for index, row in enumerate(self.dep_var):

            if index == 0:
                self.level = self.initial_level
                self.trend = self.initial_trend
                seasonal = self.initial_seasonals[0]
                self.seasonals = torch.tensor([seasonal])
                self.fitted[0] = row
            
            elif index < self.seasonal_periods:

                seasonal = self.initial_seasonals[index]
                
                if self.seasonal == "add":
                    y_hat = torch.add(seasonal, torch.add(self.level, torch.mul(self.phi, self.trend)))
                elif self.seasonal == "mul":
                    y_hat = torch.mul(seasonal, torch.add(self.level, torch.mul(self.phi, self.trend)))

                lprev, bprev = self.level, self.trend
                self.__smooth_level(row, lprev, bprev, seasonal)
                self.__smooth_trend(lprev, bprev)

                self.seasonals = torch.cat((self.seasonals, seasonal))

                self.fitted[index] = y_hat
                
            
            else:

                seasonal = self.seasonals[-self.seasonal_periods]
                
                if self.seasonal == "add":
                    y_hat = torch.add(seasonal, torch.add(self.level, torch.mul(self.phi, self.trend)))
                elif self.seasonal == "mul":
                    y_hat = torch.mul(seasonal, torch.add(self.level, torch.mul(self.phi, self.trend)))

                lprev, bprev = self.level, self.trend
                self.__smooth_level(row, lprev, bprev, seasonal)
                self.__smooth_trend(lprev, bprev)
                self.__smooth_seasonal(row, lprev, bprev, seasonal)

                self.fitted[index] = y_hat
            

    def predict(self, h):
        """
        Predict the next h values of the target variable.

        Arguments:

        *h (int): Future time-steps to forecast
        """
        
        self.forecast = torch.tensor([])
        len_seasonal = self.seasonals.shape[0]
        for i in range(1,h+1):
       
            damp_factor = torch.sum(torch.tensor([torch.pow(self.phi, x+1) for x in range(i)]))
            k = int((i-1)/self.seasonal_periods)
   
            step_forecast = torch.add(self.level, torch.mul(damp_factor, self.trend))

            if self.seasonal == "mul":
                step_forecast = torch.mul(step_forecast, self.seasonals[len_seasonal+i-self.seasonal_periods*(k+1)])
            elif self.seasonal == "add":
                step_forecast = torch.add(step_forecast, self.seasonals[len_seasonal+i-self.seasonal_periods*(k+1)])
            
            self.forecast = torch.cat((self.forecast, step_forecast))
            
        return self.forecast



            
