# coding: utf-8

# In[1]:

import copy

import numpy as np
from pymongo import MongoClient

from rnpp.hawkes_processes.datatypes import HawkesProcess

# Set to Data Base

client = MongoClient("129.26.78.40:27017")
db = client["PointProcesses"]
# clean database
collection = db['TrainingData']
# collection.remove({})

# Parameters

data_params = {'T': 100., 'beta': 0.4, 'alpha': 10., 'lambda_0': 1.}
data_params_list = []
lambda_0s = np.arange(1., 5., 1.)
for lambda_0 in lambda_0s:
    params = copy.deepcopy(data_params)
    params.update({'lambda_0': lambda_0})
    data_params_list.append(params)

for data_params in data_params_list:
    print 'generate data for ', data_params
    beta = data_params['beta']
    alpha = data_params['alpha']
    lambda_0 = data_params['lambda_0']
    T = data_params['T']

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

    hawkes_process_parameters = {"type_of_kernel": "exponential",
                                 "kernel_parameters": (alpha, beta),
                                 "base_intensity_parameters": {"type_of_base": "constant", "a": lambda_0},
                                 "excitations_kernel_parameters": None,
                                 "T": T,
                                 "inference": "variational"}

    hp_simulate_parameters = {"type_of_simulation": "stochastic_hawkes",
                              "T": T,
                              "epsilon": 1e-6,
                              "maximum_branching_level": 10}

    HP = HawkesProcess("Hawkes", hawkes_process_parameters)
    for i in range(300):
        arrivals = HP.simulate(hp_simulate_parameters)
        INTENSITY = HP.__intensity(WHERE_TO_EVALUATE_INTENSITY, arrivals)
        ROW_JSON = {"__intensity": INTENSITY.tolist(),
                    "support": support_tuple,
                    "arrivals": arrivals.tolist(),
                    "model_type": "HawkesProcess",
                    "model_parameters": hawkes_process_parameters,
                    "simulate_parameters": None,
                    "id_params": data_params}
        db["TrainingData"].insert_one(ROW_JSON)
