 
import copy
from rnpp.GP import helper
from rnpp.GP.kernels.stationary import SquaredExponential
from rnpp.poisson_processes.intensity_functions import FUNCTIONS

from matplotlib import  pyplot as plt
from matplotlib import gridspec
from mpmath import mp
import mpmath
from scipy.optimize import minimize
from scipy.special import erf
from scipy.stats import poisson

import numpy as np

EULER_MASCHERONI = 0.57721566

mp.dps = 25
mp.pretty = True


class PoissonProcess(object):
    
    def __init__(self, parameters):
        self.type = parameters["type"]
        self.function = parameters["function"]
        self.parameters = parameters["function_parameters"]
    
    def __intensity(self, t):
        return FUNCTIONS[self.function](t, *self.parameters)
        
    def thinning(self, T, arrival=0.):
        """
        Here we follow:
            Tractable Nonparametric Bayesian Inference in Poisson Processes with Gaussian Process Intensities
        """
        time_support = np.arange(arrival, T, (T - arrival) * 0.001)
        FUNCTION = self.__intensity(time_support)        
        rateBound = max(FUNCTION)
        if min(FUNCTION) < 0:
            print ("Negativ Intensity function in Poisson Process Thinning")
            raise Exception
        J = poisson.rvs((T - arrival) * rateBound)
        datesInSeconds = np.random.uniform(arrival, T, J)
        intensities = self.__intensity(datesInSeconds) / rateBound
        r = np.random.uniform(0., 1., J)
        ARRIVALS = np.take(datesInSeconds, np.where(r < intensities)[0])
        ARRIVALS.sort()
        
        return ARRIVALS
    
    def homogeneous(self, rate, T):
        if rate < 0:
            print ("Negative Intensity function in Poisson Process Thinning")
            raise Exception
        numberOfArrivals = np.random.poisson(T * rate)
        ARRIVALS = np.random.uniform(0., T, size=numberOfArrivals)
        ARRIVALS.sort()
        return ARRIVALS
        
    def simulate(self, simulate_parameters): 
        if self.type == "homogeneous":
            return self.homogeneous(self.parameters["rate"], simulate_parameters["T"])
        else:
            return self.thinning(simulate_parameters["T"])
    
    #============================
    #  PLOTS
    #============================
    
    def plot_arrivals(self, arrivals, save_plot=None):
        fig, ax = plt.subplots(1, 1, figsize=(12, 1.5))
        ax.vlines(arrivals, ymin=-0.1, ymax=1.1, color="#F88017", alpha=.9, label="Packed Arrivals")
        plot_limit = (np.min(arrivals), np.max(arrivals))
        ax.set_xlim(plot_limit)
        ax.set_ylim(0., 1.)
        ax.get_yaxis().set_visible(False)
        ax.set_xlabel(r'Time ($t$)', fontsize=12)
        fig.tight_layout()
        if save_plot == None:
            plt.show()
        plt.close()
        
        
class GaussianPoissonProcess(PoissonProcess):
    """
    This is a class to specify a point process with __intensity function given by:
    
    :math:`\lambda(t) = \lambda_{0}(t) +\sum_{t:t > T_i}Y(T_i)\mu(t - T_i)`
    
    """

    def __init__(self, values):
        gaussian_poisson_process_parameters, generate_function, id = values
        self.id = id
        self.T = gaussian_poisson_process_parameters["T"]
        self.intensity_kernel_parameters = gaussian_poisson_process_parameters["intensity_kernel"]
        self.type_of_intensity_kernel = self.intensity_kernel_parameters["type_of_kernel"] 
        self.excitations_kernel_parameters = self.intensity_kernel_parameters["kernel_parameters"]
        
        self.inducing_points = np.asarray(gaussian_poisson_process_parameters["inducing_points"])
        self.inducing_function = np.asarray(gaussian_poisson_process_parameters["inducing_function"])
        self.exact_inducing_points = copy.copy(self.inducing_points)
        self.exact_inducing_function = copy.copy(self.inducing_function)
        # self.exact_square_functon = np.sqrt(self.inducing_function)
        if self.type_of_intensity_kernel == "squared_exponential":
            self.intensity_kernel = SquaredExponential(*self.excitations_kernel_parameters)
        
        if generate_function:
            self.inducing_function = helper.generate_function(self.intensity_kernel, self.inducing_points) 
            self.number_of_inducing_points = len(self.inducing_points)
            self.set_inducing_variables(self.excitations_kernel_parameters)
            self.exact_inducing_points = copy.copy(self.inducing_points)
            self.exact_inducing_function = copy.copy(self.inducing_function)
        else:
            self.number_of_inducing_points = len(self.inducing_points)
            self.set_inducing_variables(self.excitations_kernel_parameters)
        
        self.number_of_lower_diagonal = int(0.5 * self.number_of_inducing_points * (self.number_of_inducing_points - 1) + self.number_of_inducing_points)
        
    #==================================
    # SIMULATIONS
    #==================================
    
    def __intensity(self, t, pure_function=False, points=None, function=None, kernel_parameter=None):
        if np.all(points == None):
            if not pure_function:
                self.K_t_induced = self.intensity_kernel.k(t, self.inducing_points)
                return np.nan_to_num((np.dot(self.K_t_induced, np.dot(self.K_induced_inverse, self.inducing_function))) ** 2)
            else:
                self.K_t_induced = self.intensity_kernel.k(t, self.inducing_points)
                return np.nan_to_num(np.dot(self.K_t_induced, np.dot(self.K_induced_inverse, self.inducing_function)))
        else:
            points = np.array(points)
            function = np.array(function)
            self.set_inducing_variables(self.excitations_kernel_parameters, points)
            if not pure_function:
                self.K_t_induced = self.intensity_kernel.k(t, points)
                return np.nan_to_num((np.dot(self.K_t_induced, np.dot(self.K_induced_inverse, function))) ** 2)
            else:
                self.K_t_induced = self.intensity_kernel.k(t, points)
                return np.nan_to_num(np.dot(self.K_t_induced, np.dot(self.K_induced_inverse, function)))
                
    #===================
    #  INFERENCE
    #===================
    def initialize_inference(self, arrivals, mean_ratio, variance_ratio):
        # SET DATA
        self.arrivals = arrivals[:]
        self.arrivals.sort()
        #=================================================================
        # ESTIMATE INTENSITIES AND PRIORS
        #=================================================================
        self.empirical_intensity = len(self.arrivals) / self.T
        sqrt_empirical = np.sqrt(self.empirical_intensity)
        self.m_prior = np.repeat(sqrt_empirical, self.number_of_inducing_points)
        # self.m_prior = np.random.normal(0.,sqrt_empirical,self.number_of_inducing_points)
        #=================================================================
        # DEFINE BOUNDS FOR MINIMIZATION PROCEDURE
        #=================================================================
        m_bounds = [(-mean_ratio * sqrt_empirical, mean_ratio * sqrt_empirical) for i in range(self.number_of_inducing_points)]
        SL_bounds = [(-sqrt_empirical * variance_ratio, sqrt_empirical * variance_ratio) for a in range(self.number_of_lower_diagonal)]
        inducing_omega_bounds = [(-np.pi * 0.5, np.pi * 0.5) for i in range(self.number_of_inducing_points)]
        self.mimization_bounds = m_bounds + SL_bounds + inducing_omega_bounds
        #=================================================================
        # AUXILIARY VARIABLES FOR THE CALCULATION OF THE BOUND
        #=================================================================
        self.set_psi(self.inducing_points)
        self.precompute_arrivals_average_matrices(self.inducing_points)
        #=================================================================
        # VERBATIM
        #=================================================================
        print ("Gaussian Poisson Process ")
        print ("Number of arrivals ", len(self.arrivals))
        print ("Constant density ", self.empirical_intensity)
        print ("Empirical Inducing values ", sqrt_empirical)
        print ("Prior defined from constant __intensity ")
        print (self.m_prior)
        print ("Bounds on the __intensity training")
        print ([a for a in self.mimization_bounds])
        #================================================
        # FIRST ESTIMATERS OF VARIATIONAL APPROXIMATION
        #================================================
        SL = self.KL  # + np.tril(np.random.uniform(0.1,0.2,(self.number_of_inducing_points,self.number_of_inducing_points)))
        self.SL = self.KL
        initial_SL = copy.copy(self.SL) 
        np.random.seed(self.id)
        initial_m = np.random.normal(0., self.empirical_intensity, size=self.number_of_inducing_points)
        initial_w = self.inducing_points_to_omega(self.inducing_points)
        #================================================
        
        return initial_m, initial_SL, initial_w
    
    def inference_worker(self, something):
        #==============================================
        # INFERENCE PARAMETERS
        #==============================================
        arrivals, inference_parameters = something
        max_iterations = inference_parameters["max_iterations"]
        data_dir = inference_parameters["data_dir"]
        inference_name = inference_parameters["inference_name"]
        method = inference_parameters["method"]
        mean_ratio = inference_parameters["mean_ratio"]
        variance_ratio = inference_parameters["variance_ratio"] 
        initial_m, initial_SL, initial_w = self.initialize_inference(arrivals, mean_ratio, variance_ratio)

        print ("#==============================================================")
        print ("# Initial guess")
        print ("# Initial guess for mean of approximate inducing functions (mean)")
        print (list(initial_m))
        print ("# Initial guess for inducing points positions ")
        print (list(self.omega_to_inducing_points(initial_w)))
        print ("#==============================================================")
        #================================================
        # BOUND OPTIMIZATION
        #================================================
        x0 = self.parameters_to_vector(initial_m, initial_SL, initial_w)
        result = minimize(self.bound,
                          x0,
                          method=method,
                          options={'disp': True, 'maxiter':max_iterations},
                          bounds=self.mimization_bounds)
        #================================================
        # RESULTS
        #================================================
        final_x = result.x
        final_m, final_SL, final_w = self.vector_to_parameters(final_x)
        print ("#=======================================================")
        print ("Optimization Ready")
        print ("Result ")
        print ("Square of approximation (__intensity) ")
        print (final_m ** 2)
        print ("Inducing points position ")
        print (self.omega_to_inducing_points(final_w))
        self.final_inducing_points = self.omega_to_inducing_points(final_w)
        print ("#=======================================================")
        
        return {"process_id":self.id,
                "initial_m":initial_m,
                "initial_SL":initial_SL,
                "initial_w":initial_w,
                "initial_bound":self.bound(x0),
                "final_m":final_m,
                "final_SL":final_SL,
                "final_w":final_w,
                "final_bound":self.bound(final_x)}
    
    def bound(self, x):
        m, SL, w = self.vector_to_parameters(x)
        
        #SETTING NEW VALUES FOR ALL THE K_INDUCING FUNCTIONS=============================
        points = np.asarray(self.omega_to_inducing_points(w))
        self.inducing_points = points
        self.set_inducing_variables(self.excitations_kernel_parameters, points=points)
        self.set_psi(points)
        self.precompute_arrivals_average_matrices(points)
        #=================================================================================
        S = np.dot(SL, SL.T)
        
        # PRIOR
        KL_q_p = self.KL_prior(m, S)
        
        # ARRIVALS
        LogAverageArrivals = []
        for constant_variance, a, b in zip(self.constant_trace_variance, self.AZ, self.BZ):
            SquareAverage = np.dot(np.dot(m, a), m)
            Variance = constant_variance + np.trace(np.dot(self.K_induced_inverse, np.dot(S, b)))
            hyper_geometric_argument = -0.5 * (SquareAverage / Variance)
            LogAverageArrivals.append(-self.G(hyper_geometric_argument) + np.log(0.5 * Variance) - EULER_MASCHERONI)
        LogAverageArrivals = np.asarray(LogAverageArrivals)
        
        # INTEGRALS
        variance, gamma, lenght_scale = self.excitations_kernel_parameters
        gamma = np.sqrt(gamma)
        SquaredAverageIntegral = np.dot(m, np.dot(self.K_induced_inverse, np.dot(np.dot(self.PSI, self.K_induced_inverse), m)))
        A = np.trace(np.dot(self.K_induced_inverse, self.PSI))
        B = np.trace(np.dot(self.K_induced_inverse, np.dot(S, np.dot(self.K_induced_inverse, self.PSI))))
        VarianceIntegral = gamma * self.T - A + B
        
        LogAverageArrivals = LogAverageArrivals.sum()
        return -(-(SquaredAverageIntegral + VarianceIntegral) + LogAverageArrivals - KL_q_p) 
    
    def set_inducing_variables(self, kernel_parameters, points=None):
        points = np.array(points)
        # self.inducing_points = points
        # # print "inverting induced points kernel"
        if np.all(points == None):
            self.excitations_kernel_parameters = kernel_parameters
            self.K_induced = self.intensity_kernel.k(self.inducing_points) + np.diag(np.repeat(0.000001, self.number_of_inducing_points))
            self.KL = helper.noisechol(self.K_induced)
            L = np.linalg.inv(self.KL)
            self.K_induced_inverse = np.dot(L.transpose(), L)
            self.K_determinant = np.linalg.det(self.K_induced)
        else:
            self.excitations_kernel_parameters = kernel_parameters
            number_of_points = len(points)
            self.K_induced = self.intensity_kernel.k(points) + np.diag(np.repeat(0.000001, number_of_points))
            self.KL = helper.noisechol(self.K_induced)
            L = np.linalg.inv(self.KL)
            self.K_induced_inverse = np.dot(L.transpose(), L)
            self.K_determinant = np.linalg.det(self.K_induced)
            
    def precompute_arrivals_average_matrices(self, points=None):
        """
        Here we define a list with the matrices evaluated at the induced points 
        found in Eq (13) and Eq (14) of
        
        Variational Inference for Gaussian Process Modulated Poisson Processes
        
        AXZ = K^{-1}_zz k_zx k_xz K^{-1}_zz
        BXZ = K^{-1}_zz k_zx k_xz
        """
        points = np.array(points)
        Trace_Variance = []
        AZ = []
        BZ = []
        if np.all(points == None):
            # EVALUATE AT ARRIVALS 
            for arrival in self.arrivals:
                K_z_arrivals = self.intensity_kernel.k(self.inducing_points, np.asarray([arrival]))
                K_z_arrivals_z = np.dot(K_z_arrivals, K_z_arrivals.T)
                A = np.dot(self.K_induced_inverse, K_z_arrivals_z)
                AZ.append(np.dot(A, self.K_induced_inverse))
                BZ.append(A)
                Trace_Variance.append((self.intensity_kernel.k(np.asarray([arrival]), np.asarray([arrival])) - np.trace(A))[0, 0])
        else:
            # EVALUATE AT ARRIVALS 
            for arrival in self.arrivals:
                K_z_arrivals = self.intensity_kernel.k(points, np.asarray([arrival]))
                K_z_arrivals_z = np.dot(K_z_arrivals, K_z_arrivals.T)
                A = np.dot(self.K_induced_inverse, K_z_arrivals_z)
                AZ.append(np.dot(A, self.K_induced_inverse))
                BZ.append(A)
                Trace_Variance.append((self.intensity_kernel.k(np.asarray([arrival]), np.asarray([arrival])) - np.trace(A))[0, 0])
                
        self.AZ = AZ 
        self.BZ = BZ
        self.constant_trace_variance = Trace_Variance[:]

    def set_psi(self, points=None):
        """
        Integral over the square of the kernel c.f. equation (17)
        """
        if np.all(points == None):
            PSI = []
            for z in self.inducing_points:
                row = []
                for zp in self.inducing_points:
                    row.append(self.IntegratedSquareKernel(z, zp))
                PSI.append(np.asarray(row))
            self.PSI = np.asarray(PSI)
        else:
            PSI = []
            for z in points:
                row = []
                for zp in points:
                    row.append(self.IntegratedSquareKernel(z, zp))
                PSI.append(np.asarray(row))
            self.PSI = np.asarray(PSI)
        
    def IntegratedSquareKernel(self, z, zp):
        variance, gamma, lenght_scale = self.excitations_kernel_parameters
        gamma = np.sqrt(gamma)
        zh = (z + zp) * 0.5
        E1 = erf((zh - self.T) / np.sqrt(lenght_scale))
        E2 = erf((zh) / np.sqrt(lenght_scale))
        return -(gamma * 2) * np.sqrt(np.pi * lenght_scale) * 0.5 * np.exp(-((z - zp) ** 2) / (4 * lenght_scale)) * (E1 - E2)
    
    def KL_prior(self, m, S):
        A = np.trace(np.dot(self.K_induced_inverse, S))
        m_difference = self.m_prior - m
        B = np.dot(m_difference, np.dot(self.K_induced_inverse, m_difference)) - len(self.m_prior)
        b = np.linalg.det(S)
        C = np.log(self.K_determinant / b)
        return 0.5 * (A + B + C)
    
    def G(self, z, dt=10e-7):
        return (mpmath.hyp1f1(dt, 0.5, z) - mpmath.hyp1f1(0., 0.5, z)) / dt
 
    #==========================================================
    #                         PLOTS
    #==========================================================
    
    def plot_intensity_and_arrivals(self, arrivals, inference_plot=False, final_w=None, final_m=None, save_to_file=None):
        time_support = np.linspace(0., self.T, 1000)
        true_intensity = self.__intensity(time_support,
                                        False,
                                        self.exact_inducing_points,
                                        self.exact_inducing_function)
        if inference_plot:
            final_inducing_points = self.omega_to_inducing_points(final_w)
            inferred_intensity = self.__intensity(time_support, False, points=final_inducing_points, function=final_m)
            # inferred_intensity_pure = self.__intensity(time_support,pure_function=True,points=self.final_inducing_points,function=self.final_m)
        plot_limit = (0., self.T)
        
        fig = plt.figure(figsize=(12, 6))
        if inference_plot:
            gs = gridspec.GridSpec(3, 1, height_ratios=[7, 7, 1])
        else:
            gs = gridspec.GridSpec(2, 1, height_ratios=[7, 1])
            
        grid_index = 0
        # INFERRED INTENSITY
        if inference_plot:
            ax = plt.subplot(gs[grid_index])
            ax.plot(time_support, inferred_intensity, "-", label="Inference")
            ax.plot(final_inducing_points, final_m ** 2., "o")
            ax.set_ylabel(r'Intensities', fontsize=18)
            ax.get_xaxis().set_visible(False)
            ax.set_xlim(plot_limit)
            # axp = ax.twinx()
            # axp.plot(self.inducing_points,self.final_m,"o")
            # axp.plot(time_support,inferred_intensity_pure,"-")
            grid_index += 1
            
        # TRUE INTENSITY
        ax = plt.subplot(gs[grid_index])
        ax.plot(time_support, true_intensity, "k-", label="Intensities", color="#168EF7")
        ax.set_ylabel(r'Intensities', fontsize=18)
        ax.get_xaxis().set_visible(False)
        ax.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=2, mode="expand", borderaxespad=0.)
        ax.set_xlim(plot_limit)
        ax.get_xaxis().set_visible(False)
        grid_index += 1
        
        # ARRIVALS
        ax = plt.subplot(gs[grid_index])
        ax.vlines(arrivals, ymin=0, ymax=1, color="#F88017", alpha=.9, label="Arrivals")
        ax.set_xlim(plot_limit)
        ax.set_xlabel(r'Time ($t$)', fontsize=18)
        
        if save_to_file is None:
            plt.show()
        else:
            plt.savefig(save_to_file)
        plt.close()
        
    #=======================================
    # UTILS
    #=======================================
    def vector_to_parameters(self, x):
        m = x[:self.number_of_inducing_points] 
        SL_reshaped = x[self.number_of_inducing_points:self.number_of_inducing_points + self.number_of_lower_diagonal]
        SL = np.zeros((self.number_of_inducing_points, self.number_of_inducing_points))
        SL[np.tril_indices(self.number_of_inducing_points)] = SL_reshaped
        w = x[self.number_of_inducing_points + self.number_of_lower_diagonal:]
        return m, SL, w
    
    def parameters_to_vector(self, m, SL, w):
        SL = SL[np.tril_indices(self.number_of_inducing_points)]
        return np.concatenate([m, SL, w])

    def omega_to_inducing_points(self, inducing_w):
        return  self.T * 0.5 + self.T * np.sin(inducing_w) * 0.5 
    
    def inducing_points_to_omega(self, points):
        return np.arcsin((points * 2. - self.T) / (self.T))

