'''
Created on Mar 21, 2017

@author: anon
'''

import numpy as np
from scipy import stats
from scipy.stats import poisson

# from rnpp.poisson_processes.datatypes import GaussianPoissonProcess


def sample_arrivals_from_poisson_sine():
    freq = .00075
    T = 4500
    dt = 0.01

    def f(t): return 0.1 * np.sin(2. * np.pi * freq * t) + 0.2

    x = Poisson_numberOfWork(T, dt, f, ())
    x = np.sort(x, axis=0)
    x = np.hstack((0, x))
    d = x[1:] - x[:-1]
    d = np.hstack((x[0], d))
    x.reshape((x.shape[0], 1))
    d.reshape((d.shape[0], 1))
    t = np.stack((x, d, np.zeros_like(d)), axis=1)

    max_arr = np.max(x)
    support = np.arange(0, max_arr, .01)
    true_intensity = f(support)
    return t, support, true_intensity


def generate_train_data_poisson(data_params, N_samples):
    """ This function generates arrivals from Hawkes process and stores them in the database """

    lenght_scale = data_params['length_scale']
    T = data_params['T']
    N = data_params['N']

    support_tuple = (0., T, 500)
    WHERE_TO_EVALUATE_INTENSITY = np.linspace(*support_tuple)

    inducing_points = np.linspace(0., T, N)
    gpp_parameters = {"intensity_kernel": {"type_of_kernel": "squared_exponential",
                                           "kernel_parameters": [1., 2., lenght_scale], },
                      "inducing_points": list(inducing_points),
                      "inducing_function": None,
                      "T": T,
                      "inference": "variational"}

    simulate_function = True
    GPP = GaussianPoissonProcess((gpp_parameters, simulate_function, 0))
    gpp_parameters["inducing_function"] = list(GPP.inducing_function)

    support_tuple = (0., T, 500)
    WHERE_TO_EVALUATE_INTENSITY = np.linspace(*support_tuple)
    INTENSITY = GPP.__intensity(WHERE_TO_EVALUATE_INTENSITY)
    all_arrivals = []
    for i in range(N_samples):
        arrivals = GPP.thinning(T)
        all_arrivals.append(arrivals)
    return INTENSITY, all_arrivals, support_tuple


def prepare_data(args):
    data = []
    for _ in range(args.batch_size):
        x, f = sample_arrivals_from_poisson_sine(args)
        data.append(x)
    min_seq = min([x.shape[0] for x in data])
    input_data = np.asarray([x[:min_seq] for x in data])
    return f, input_data


def normalize(low=0, high=1, A=None):
    mins = np.min(A, axis=0)
    maxs = np.max(A, axis=0)
    rng = maxs - mins

    return high - (((high - low) * (maxs - A)) / rng)


def compress_arrivals(arrivals, min_border=.5, max_border=8, step=.5):
    interarrivals = np.diff(arrivals)
    pvalues = []
    for i in np.arange(min_border, max_border, step):
        l = interarrivals[np.less_equal(interarrivals, i)]
        pvalues.append(stats.kstest(l, "expon")[1])

    t = np.arange(min_border, max_border, step)[np.argmax(pvalues)] / 24
    packages = np.split(arrivals, np.where(np.diff(arrivals) > t)[0] + 1)
    compressed = [package[0] for package in packages]
    return packages, compressed


def Poisson_numberOfWork(T, dT, function, functionParameters):
    """
    used for recurrent network

    Parameters:
    T: float
    dT: float
    function: function
    functionParametersT
    numberOfWorkDistribution: stast frozen rvs
    """
    rateBound = max(function(np.arange(0, T, dT), *functionParameters)
                    )  # DANGEROUS ASSUMES THAT THE CHANGES WITHIN DT ARE NEGLIGIBLE
    J = poisson.rvs(T * rateBound)
    datesInSeconds = np.random.uniform(0., T, J)
    intensities = function(datesInSeconds, *functionParameters) / rateBound
    r = np.random.uniform(0., 1., J)
    hidden_arrivals = np.take(datesInSeconds, np.where(r < intensities)[0])
    return hidden_arrivals


def doublePoisson_numberOfWork(T, dT, function, functionParameters, fineRate, numberOfWorkDistribution):
    """
    generates a set of arrivals from a functional form
    using the thinning process

    Parameters:
    T: float
    dT: float
    function: function
    functionParametersT
    numberOfWorkDistribution: stast frozen rvs
    """
    rateBound = max(function(np.arange(0, T, dT), *functionParameters)
                    )  # DANGEROUS ASSUMES THAT THE CHANGES WITHIN DT ARE NEGLIGIBLE
    J = poisson.rvs(T * rateBound)
    datesInSeconds = np.random.uniform(0., T, J)
    intensities = function(datesInSeconds, *functionParameters) / rateBound
    r = np.random.uniform(0., 1., J)
    hidden_arrivals = np.take(datesInSeconds, np.where(r < intensities)[0])
    number_of_hidden_arrivals = len(hidden_arrivals)
    number_of_arrivals = np.abs(np.ceil(
        numberOfWorkDistribution.rvs(size=number_of_hidden_arrivals)))
    observed_arrivals = []
    for k, N in enumerate(number_of_arrivals):
        datesInSeconds = np.random.exponential(1. / fineRate, size=int(N))
        observed_arrivals.extend(hidden_arrivals[k] + datesInSeconds.cumsum())
    return hidden_arrivals, number_of_arrivals, observed_arrivals
