import copy
import sys
if sys.version_info < (3,):
    range = xrange

import numpy as np
import pandas as pd
import scipy.stats as ss

from .. import families as fam
from .. import tsm as tsm
from .. import data_check as dc

from .kalman import *

class LLEV(tsm.TSM):
    """ Inherits time series methods from TSM class.

    **** LOCAL LEVEL MODEL ****

    Parameters
    ----------
    data : pd.DataFrame or np.array
        Field to specify the time series data that will be used.

    integ : int (default : 0)
        Specifies how many time to difference the time series.

    target : str (pd.DataFrame) or int (np.array)
        Specifies which column name or array index to use. By default, first
        column/array will be selected as the dependent variable.
    """

    def __init__(self,data,integ=0,target=None):

        # Initialize TSM object
        super(LLEV,self).__init__('LLEV')

        # Latent Variables
        self.integ = integ
        self.param_no = 2
        self.max_lag = 0
        self._z_hide = 0 # Whether to cutoff variance latent variables from results
        self.supported_methods = ["MLE","PML","Laplace","M-H","BBVI"]
        self.default_method = "MLE"
        self.model_name = "LLEV"
        self.multivariate_model = False

        # Format the data
        self.data, self.data_name, self.is_pandas, self.index = dc.data_check(data,target)
        self.data = self.data.astype(np.float)
        self.data_original = self.data

        # Difference data
        for order in range(self.integ):
            self.data = np.diff(self.data)
            self.data_name = "Differenced " + self.data_name

        self.data_length = self.data.shape[0]
        self._create_latent_variables()
        self.z_no = 2

    def _create_latent_variables(self):
        """ Creates model latent variables

        Returns
        ----------
        None (changes model attributes)
        """

        self.latent_variables.add_z('Sigma^2 irregular', fam.Flat(transform='exp'), fam.Normal(0,3))
        self.latent_variables.add_z('Sigma^2 level', fam.Flat(transform='exp'), fam.Normal(0,3))

    def _forecast_model(self,beta,h):
        """ Creates forecasted states and variances

        Parameters
        ----------
        beta : np.ndarray
            Contains untransformed starting values for latent variables

        Returns
        ----------
        a : np.ndarray
            Forecasted states

        P : np.ndarray
            Variance of forecasted states
        """     

        T, Z, R, Q, H = self._ss_matrices(beta)
        return univariate_kalman_fcst(self.data,Z,H,T,Q,R,0.0,h)

    def _model(self, data, beta):
        """ Creates the structure of the model

        Parameters
        ----------
        data : np.array
            Contains the time series

        beta : np.array
            Contains untransformed starting values for latent variables

        Returns
        ----------
        a,P,K,F,v : np.array
            Filted states, filtered variances, Kalman gains, F matrix, residuals
        """     

        T, Z, R, Q, H = self._ss_matrices(beta)
        return univariate_kalman(data,Z,H,T,Q,R,0.0)

    def _ss_matrices(self,beta):
        """ Creates the state space matrices required

        Parameters
        ----------
        beta : np.array
            Contains untransformed starting values for latent_variables

        Returns
        ----------
        T, Z, R, Q, H : np.array
            State space matrices used in KFS algorithm
        """     

        T = np.identity(1)
        R = np.identity(1)
        Z = np.identity(1)
        H = np.identity(1)*self.latent_variables.z_list[0].prior.transform(beta[0])
        Q = np.identity(1)*self.latent_variables.z_list[1].prior.transform(beta[1])

        return T, Z, R, Q, H

    def neg_loglik(self,beta):
        """ Creates the negative log likelihood of the model

        Parameters
        ----------
        beta : np.array
            Contains untransformed starting values for latent variables

        Returns
        ----------
        The negative log logliklihood of the model
        """         
        _, _, _, F, v = self._model(self.data,beta)
        loglik = 0.0
        for i in range(0,self.data.shape[0]):
            loglik += np.linalg.slogdet(F[:,:,i])[1] + np.dot(v[i],np.dot(np.linalg.pinv(F[:,:,i]),v[i]))
        return -(-((self.data.shape[0]/2)*np.log(2*np.pi))-0.5*loglik.T[0].sum())

    def mb_neg_loglik(self, beta, mini_batch):
        """ Creates the negative log likelihood of the model

        Parameters
        ----------
        beta : np.array
            Contains untransformed starting values for latent variables

        mini_batch : int
            Size of each mini batch of data

        Returns
        ----------
        The negative log logliklihood of the model
        """    

        rand_int =  np.random.randint(low=0, high=self.data.shape[0]-mini_batch-self.max_lag+1)
        sample = np.arange(start=rand_int, stop=rand_int+mini_batch)

        _, _, _, F, v = self._model(self.data[sample],beta)
        loglik = 0.0
        for i in range(0,len(sample)):
            loglik += np.linalg.slogdet(F[:,:,i])[1] + np.dot(v[i],np.dot(np.linalg.pinv(F[:,:,i]),v[i]))
        return -(-((len(sample)/2)*np.log(2*np.pi))-0.5*loglik.T[0].sum())

    def plot_predict(self, h=5, past_values=20, intervals=True, **kwargs):      
        """ Makes forecast with the estimated model

        Parameters
        ----------
        h : int (default : 5)
            How many steps ahead would you like to forecast?

        past_values : int (default : 20)
            How many past observations to show on the forecast graph?

        intervals : Boolean
            Would you like to show prediction intervals for the forecast?

        Returns
        ----------
        - Plot of the forecast
        """
        import matplotlib.pyplot as plt
        import seaborn as sns

        figsize = kwargs.get('figsize',(10,7))
        nsims = kwargs.get('nsims', 200)

        if self.latent_variables.estimated is False:
            raise Exception("No latent variables estimated!")
        else:
            # Retrieve data, dates and (transformed) latent variables   
            if self.latent_variables.estimation_method in ['M-H']:
                lower_final = 0
                upper_final = 0
                plot_values_final = 0
                date_index = self.shift_dates(h)
                plot_index = date_index[-h-past_values:]

                for i in range(nsims):

                    t_params = self.draw_latent_variables(nsims=1).T[0]
                    a, P = self._forecast_model(t_params, h)

                    plot_values = a[0][-h-past_values:]
                    forecasted_values = a[0][-h:]

                    lower = forecasted_values - 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(t_params[0]),0.5)
                    upper = forecasted_values + 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(t_params[0]),0.5)
                    lower_final += np.append(plot_values[-h-1], lower)
                    upper_final += np.append(plot_values[-h-1], upper)
                    plot_values_final += plot_values

                plot_values_final = plot_values_final / nsims
                lower_final = lower_final / nsims
                upper_final = upper_final / nsims

                plt.figure(figsize=figsize)
                if intervals == True:
                    plt.fill_between(date_index[-h-1:], lower_final, upper_final, alpha=0.2)            

                plt.plot(plot_index, plot_values_final)
                plt.title("Forecast for " + self.data_name)
                plt.xlabel("Time")
                plt.ylabel(self.data_name)
                plt.show()
            else:
                a, P = self._forecast_model(self.latent_variables.get_z_values(),h)
                date_index = self.shift_dates(h)
                plot_values = a[0][-h-past_values:]
                forecasted_values = a[0][-h:]

                lower = forecasted_values - 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(self.latent_variables.get_z_values()[0]),0.5)
                upper = forecasted_values + 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(self.latent_variables.get_z_values()[0]),0.5)
                lower = np.append(plot_values[-h-1],lower)
                upper = np.append(plot_values[-h-1],upper)

                plot_index = date_index[-h-past_values:]

                plt.figure(figsize=figsize)
                if intervals == True:
                    plt.fill_between(date_index[-h-1:], lower, upper, alpha=0.2)            

                plt.plot(plot_index,plot_values)
                plt.title("Forecast for " + self.data_name)
                plt.xlabel("Time")
                plt.ylabel(self.data_name)
                plt.show()

    def plot_fit(self, intervals=True, **kwargs):
        """ Plots the fit of the model

        Parameters
        ----------
        intervals : Boolean
            Whether to plot 95% confidence interval of states

        Returns
        ----------
        None (plots data and the fit)
        """
        import matplotlib.pyplot as plt
        import seaborn as sns

        figsize = kwargs.get('figsize',(10,7))
        series_type = kwargs.get('series_type','Smoothed')

        if self.latent_variables.estimated is False:
            raise Exception("No latent variables estimated!")
        else:
            date_index = copy.deepcopy(self.index)
            date_index = date_index[self.integ:self.data_original.shape[0]+1]

            if series_type == 'Smoothed':
                mu, V = self.smoothed_state(self.data,self.latent_variables.get_z_values())
            elif series_type == 'Filtered':
                mu, V, _, _, _ = self._model(self.data,self.latent_variables.get_z_values())
            else:
                mu, V = self.smoothed_state(self.data,self.latent_variables.get_z_values())

            mu = mu[0][:-1]
            V = V.ravel()

            plt.figure(figsize=figsize) 
            
            plt.subplot(3, 1, 1)
            plt.title(self.data_name + " Raw and " + series_type)   

            if intervals == True:
                alpha =[0.15*i/float(100) for i in range(50,12,-2)]
                plt.fill_between(date_index[2:], mu[2:] + 1.98*np.sqrt(V[:-1][2:]), mu[2:] - 1.98*np.sqrt(V[:-1][2:]), alpha=0.15,label='95% C.I.') 

            plt.plot(date_index,self.data,label='Data')
            plt.plot(date_index,mu,label=series_type,c='black')
            plt.legend(loc=2)
            
            plt.subplot(3, 1, 2)
            plt.title(self.data_name + " Local Level")  

            if intervals == True:
                alpha =[0.15*i/float(100) for i in range(50,12,-2)]
                plt.fill_between(date_index[2:], mu[2:] + 1.98*np.sqrt(V[:-1][2:]), mu[2:] - 1.98*np.sqrt(V[:-1][2:]), alpha=0.15,label='95% C.I.') 

            plt.plot(date_index,mu,label='Local Level')
            plt.legend(loc=2)
            plt.subplot(3, 1, 3)
            plt.title("Measurement Noise")  
            plt.plot(date_index,self.data-mu)
            plt.show()

    def predict(self, h=5, intervals=False, **kwargs):      
        """ Makes forecast with the estimated model

        Parameters
        ----------
        h : int (default : 5)
            How many steps ahead would you like to forecast?

        intervals : boolean (default: False)
            Whether to return prediction intervals

        Returns
        ----------
        - pd.DataFrame with predictions
        """     

        nsims = kwargs.get('nsims', 200)

        if self.latent_variables.estimated is False:
            raise Exception("No latent variables estimated!")
        else:
            # Retrieve data, dates and (transformed) latent variables   
            if self.latent_variables.estimation_method in ['M-H']:
                lower_1_final = 0
                upper_99_final = 0
                lower_5_final = 0
                upper_95_final = 0
                forecasted_values_final = 0
                date_index = self.shift_dates(h)

                for i in range(nsims):
                    t_params = self.draw_latent_variables(nsims=1).T[0]
                    a, P = self._forecast_model(t_params, h)

                    forecasted_values = a[0][-h:]
                    lower_5 = forecasted_values - 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(t_params[0]),0.5)
                    upper_95 = forecasted_values + 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(t_params[0]),0.5)
                    lower_5_final += lower_5
                    upper_95_final += upper_95
                    lower_1 = forecasted_values - 2.575*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(t_params[0]),0.5)
                    upper_99 = forecasted_values + 2.575*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(t_params[0]),0.5)
                    lower_1_final += lower_1
                    upper_99_final += upper_99
                    forecasted_values_final += forecasted_values

                forecasted_values_final = forecasted_values_final / nsims
                lower_1_final = lower_1_final / nsims
                lower_5_final = lower_5_final / nsims
                upper_95_final = upper_95_final / nsims
                upper_99_final = upper_99_final / nsims

                if intervals is False:
                    result = pd.DataFrame(forecasted_values_final)
                    result.rename(columns={0:self.data_name}, inplace=True)
                else:
                    prediction_05 = lower_5_final
                    prediction_95 = upper_95_final
                    prediction_01 = lower_1_final
                    prediction_99 = upper_99_final

                    result = pd.DataFrame([forecasted_values_final, prediction_01, prediction_05, 
                        prediction_95, prediction_99]).T
                    result.rename(columns={0:self.data_name, 1: "1% Prediction Interval", 
                        2: "5% Prediction Interval", 3: "95% Prediction Interval", 4: "99% Prediction Interval"}, 
                        inplace=True)

                result.index = date_index[-h:]

                return result
     
            else:
                # Retrieve data, dates and (transformed) latent variables         
                a, P = self._forecast_model(self.latent_variables.get_z_values(),h)
                date_index = self.shift_dates(h)
                forecasted_values = a[0][-h:]

                if intervals is False:
                    result = pd.DataFrame(forecasted_values)
                    result.rename(columns={0:self.data_name}, inplace=True)
                else:
                    prediction_05 = forecasted_values - 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(self.latent_variables.get_z_values()[0]),0.5)
                    prediction_95 = forecasted_values + 1.96*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(self.latent_variables.get_z_values()[0]),0.5)
                    prediction_01 = forecasted_values - 2.575*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(self.latent_variables.get_z_values()[0]),0.5)
                    prediction_99 = forecasted_values + 2.575*np.power(P[0][0][-h:] + self.latent_variables.z_list[0].prior.transform(self.latent_variables.get_z_values()[0]),0.5)

                    result = pd.DataFrame([forecasted_values, prediction_01, prediction_05, 
                        prediction_95, prediction_99]).T
                    result.rename(columns={0:self.data_name, 1: "1% Prediction Interval", 
                        2: "5% Prediction Interval", 3: "95% Prediction Interval", 4: "99% Prediction Interval"}, 
                        inplace=True)
     
                result.index = date_index[-h:]

                return result

    def predict_is(self, h=5, fit_once=True, fit_method='MLE', intervals=False):
        """ Makes dynamic in-sample predictions with the estimated model

        Parameters
        ----------
        h : int (default : 5)
            How many steps would you like to forecast?

        fit_once : boolean
            (default: True) Fits only once before the in-sample prediction; if False, fits after every new datapoint

        fit_method : string
            Which method to fit the model with

        intervals : boolean
            Whether to output prediction intervals or not

        Returns
        ----------
        - pd.DataFrame with predicted values
        """     

        predictions = []

        for t in range(0,h):
            x = LLEV(integ=self.integ,data=self.data_original[:(-h+t)])
            if fit_once is False:
                x.fit(fit_method=fit_method, printer=False)
            if t == 0:
                if fit_once is True:
                    x.fit(fit_method=fit_method, printer=False)
                    saved_lvs = x.latent_variables
                predictions = x.predict(1, intervals=intervals)
            else:
                if fit_once is True:
                    x.latent_variables = saved_lvs
                predictions = pd.concat([predictions,x.predict(1, intervals=intervals)])

        predictions.rename(columns={0:self.data_name}, inplace=True)
        predictions.index = self.index[-h:]

        return predictions

    def plot_predict_is(self, h=5, fit_once=True, fit_method='MLE', **kwargs):
        """ Plots forecasts with the estimated model against data
            (Simulated prediction with data)

        Parameters
        ----------
        h : int (default : 5)
            How many steps to forecast

        fit_once : boolean
            (default: True) Fits only once before the in-sample prediction; if False, fits after every new datapoint

        fit_method : string
            Which method to fit the model with

        Returns
        ----------
        - Plot of the forecast against data 
        """
        import matplotlib.pyplot as plt
        import seaborn as sns

        figsize = kwargs.get('figsize',(10,7))

        plt.figure(figsize=figsize)
        predictions = self.predict_is(h, fit_once=fit_once, fit_method=fit_method)
        data = self.data[-h:]
        plt.plot(predictions.index,data,label='Data')
        plt.plot(predictions.index,predictions,label='Predictions',c='black')
        plt.title(self.data_name)
        plt.legend(loc=2)   
        plt.show()          

    def simulation_smoother(self,beta):
        """ Koopman's simulation smoother - simulates from states given
        model latent variables and observations

        Parameters
        ----------

        beta : np.array
            Contains untransformed starting values for latent variables

        Returns
        ----------
        - A simulated state evolution
        """         

        T, Z, R, Q, H = self._ss_matrices(beta)

        # Generate e_t+ and n_t+
        rnd_h = np.random.normal(0,np.sqrt(H),self.data.shape[0]+1)
        q_dist = ss.multivariate_normal([0.0], Q)
        rnd_q = q_dist.rvs(self.data.shape[0]+1)

        # Generate a_t+ and y_t+
        a_plus = np.zeros((T.shape[0],self.data.shape[0]+1)) 
        a_plus[0,0] = np.mean(self.data[0:5])
        y_plus = np.zeros(self.data.shape[0])

        for t in range(0,self.data.shape[0]+1):
            if t == 0:
                a_plus[:,t] = np.dot(T,a_plus[:,t]) + rnd_q[t]
                y_plus[t] = np.dot(Z,a_plus[:,t]) + rnd_h[t]
            else:
                if t != self.data.shape[0]:
                    a_plus[:,t] = np.dot(T,a_plus[:,t-1]) + rnd_q[t]
                    y_plus[t] = np.dot(Z,a_plus[:,t]) + rnd_h[t]

        alpha_hat,_ = self.smoothed_state(self.data,beta)
        alpha_hat_plus,_ = self.smoothed_state(y_plus,beta)
        alpha_tilde = alpha_hat - alpha_hat_plus + a_plus
    
        return alpha_tilde

    def smoothed_state(self,data,beta):
        """ Creates the negative log marginal likelihood of the model

        Parameters
        ----------

        data : np.array
            Data to be smoothed

        beta : np.array
            Contains untransformed starting values for latent variables

        Returns
        ----------
        - Smoothed states
        """         

        T, Z, R, Q, H = self._ss_matrices(beta)
        alpha, V = univariate_KFS(data,Z,H,T,Q,R,0.0)
        return alpha, V

    def sample(self, nsims=1000):
        """ Samples from the posterior predictive distribution

        Parameters
        ----------
        nsims : int (default : 1000)
            How many draws from the posterior predictive distribution

        Returns
        ----------
        - np.ndarray of draws from the data
        """     
        if self.latent_variables.estimation_method not in ['BBVI', 'M-H']:
            raise Exception("No latent variables estimated!")
        else:
            lv_draws = self.draw_latent_variables(nsims=nsims)
            mus = [self.smoothed_state(self.data, lv_draws[:,i])[0][0][:-1] for i in range(nsims)]
            data_draws = np.array([np.random.normal(mus[i], np.sqrt(self.latent_variables.z_list[0].prior.transform(lv_draws[0,i])), mus[i].shape[0]) for i in range(nsims)])
            return data_draws

    def plot_sample(self, nsims=10, plot_data=True, **kwargs):
        """
        Plots draws from the posterior predictive density against the data

        Parameters
        ----------
        nsims : int (default : 1000)
            How many draws from the posterior predictive distribution

        plot_data boolean
            Whether to plot the data or not
        """

        if self.latent_variables.estimation_method not in ['BBVI', 'M-H']:
            raise Exception("No latent variables estimated!")
        else:
            import matplotlib.pyplot as plt
            import seaborn as sns

            figsize = kwargs.get('figsize',(10,7))
            plt.figure(figsize=figsize)
            date_index = self.index
            draws = self.sample(nsims).T
            plt.plot(date_index, draws, label='Posterior Draws', alpha=1.0)
            if plot_data is True:
                plt.plot(date_index, self.data, label='Data', c='black', alpha=0.5, linestyle='', marker='s')
            plt.title(self.data_name)
            plt.show()    

    def ppc(self, nsims=1000, T=np.mean):
        """ Computes posterior predictive p-value

        Parameters
        ----------
        nsims : int (default : 1000)
            How many draws for the PPC

        T : function
            A discrepancy measure - e.g. np.mean, np.std, np.max

        Returns
        ----------
        - float (posterior predictive p-value)
        """     
        if self.latent_variables.estimation_method not in ['BBVI', 'M-H']:
            raise Exception("No latent variables estimated!")
        else:
            lv_draws = self.draw_latent_variables(nsims=nsims)
            mus = [self.smoothed_state(self.data, lv_draws[:,i])[0][0][:-1] for i in range(nsims)]
            data_draws = np.array([np.random.normal(mus[i], np.sqrt(self.latent_variables.z_list[0].prior.transform(lv_draws[0,i])), mus[i].shape[0]) for i in range(nsims)])
            T_sims = T(self.sample(nsims=nsims), axis=1)
            T_actual = T(self.data)
            return len(T_sims[T_sims>T_actual])/nsims

    def plot_ppc(self, nsims=1000, T=np.mean, **kwargs):
        """ Plots histogram of the discrepancy from draws of the posterior

        Parameters
        ----------
        nsims : int (default : 1000)
            How many draws for the PPC

        T : function
            A discrepancy measure - e.g. np.mean, np.std, np.max
        """     
        if self.latent_variables.estimation_method not in ['BBVI', 'M-H']:
            raise Exception("No latent variables estimated!")
        else:
            import matplotlib.pyplot as plt
            import seaborn as sns

            figsize = kwargs.get('figsize',(10,7))

            lv_draws = self.draw_latent_variables(nsims=nsims)
            mus = [self.smoothed_state(self.data, lv_draws[:,i])[0][0][:-1] for i in range(nsims)]
            data_draws = np.array([np.random.normal(mus[i], np.sqrt(self.latent_variables.z_list[0].prior.transform(lv_draws[0,i])), mus[i].shape[0]) for i in range(nsims)])
            T_sim = T(self.sample(nsims=nsims), axis=1)
            T_actual = T(self.data)

            if T == np.mean:
                description = " of the mean"
            elif T == np.max:
                description = " of the maximum"
            elif T == np.min:
                description = " of the minimum"
            elif T == np.median:
                description = " of the median"
            else:
                description = ""

            plt.figure(figsize=figsize)
            ax = plt.subplot()
            ax.axvline(T_actual)
            sns.distplot(T_sim, kde=False, ax=ax)
            ax.set(title='Posterior predictive' + description, xlabel='T(x)', ylabel='Frequency');
            plt.show()