import copy

import numpy as np
import tensorflow as tf
import tqdm
from pymongo import MongoClient

from rnpp.hawkes_processes.datatypes import HawkesProcess
from rnpp.poisson_processes.datatypes import GaussianPoissonProcess


def reshape_input_for_training(sequences, batch_size, window_size):
    """
    
    :param sequences: 
    :param batch_size:
    :param window_size:
    :return:
    """
    B, L, D = sequences.shape
    NUM_BATCHES = B / batch_size
    NUM_WIN = L / window_size
    x = sequences[:NUM_BATCHES * batch_size, :NUM_WIN *
                                              window_size].reshape(NUM_BATCHES, batch_size, NUM_WIN, window_size, D)
    return x


def preprocess_sequence(args, x):
    """ This function takes a sequence, calculates the difference between
    two arrivals as additional input and reshapes the data according to 'num_steps'"""

    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), axis=1)
    num_arrivals = x.shape[0]
    pieces = num_arrivals / args.num_steps
    t = t[:pieces * args.num_steps].reshape(pieces, args.num_steps, 2)
    return t


def get_train_data(args, data_params, N, collection):
    """ This function load arrivals from the database and prepares it for training """
    rows = collection.find({"id_params": data_params})
    data = []
    for i, batch in enumerate(range(N)):
        arrivals = rows[i]['arrivals']
        x = preprocess_sequence(args, arrivals)
        data.append(x)

    # read __intensity from first dataset since it is the same for all
    f = rows[0]['__intensity']
    support = np.linspace(0., data_params['T'], 500)

    # cut to minimum sequence length
    min_seq = min([x.shape[0] for x in data])
    input_data = np.asarray([x[:min_seq] for x in data])

    return input_data, support, f


def get_train_data_du(args, data_file):
    """ This function load arrivals from Du et al """
    data = []
    data_raw = np.loadtxt(data_file)
    N = np.shape(data_raw)[0]
    for i, batch in enumerate(range(N)):
        arrivals = data_raw[i]
        x = preprocess_sequence(args, arrivals)
        data.append(x)

    # cut to minimum sequence length
    min_seq = min([x.shape[0] for x in data])
    input_data = np.asarray([x[:min_seq] for x in data])

    return input_data


def get_train_data_taxi(args, N):
    client = MongoClient("129.26.78.40:27017")
    db = client["NYCTaxi"]
    collection = db['trip_data_1_group_2']

    docs = collection.find({})
    A = []
    for i, doc in enumerate(docs):
        arrivals = np.sort(doc['arrivals'])
        A.append(arrivals)

    L = [len(a) for a in A]
    index_sort = np.argsort(L)[::-1]
    A_sort = [A[i] for i in index_sort]

    """ This function load arrivals from Du et al """
    data = []
    for i, batch in enumerate(range(N)):
        arrivals = A_sort[i]
        x = preprocess_sequence(args, arrivals)
        data.append(x)

    return data


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

    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(N):
        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}
        collection.insert_one(ROW_JSON)


def generate_train_data_poisson(data_params, N_samples, collection):
    """ 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)
    for i in range(N_samples):
        arrivals = GPP.thinning(T)
        ROW_JSON = {"__intensity": INTENSITY.tolist(),
                    "support": support_tuple,
                    "arrivals": arrivals.tolist(),
                    "model_type": "GaussianPoisson",
                    "model_parameters": gpp_parameters,
                    "simulate_parameters": None,
                    "id_params": data_params}
        collection.insert_one(ROW_JSON)


def train(input_data_test, input_data_train, args, directories):
    """ trains and tests the model """

    tf.reset_default_graph()
    with tf.Graph().as_default():
        model = RNNHPP(args)
        with tf.Session(config=tf.ConfigProto(
                intra_op_parallelism_threads=1)) as sess:
            cost, LH, ts, intensities = model.train(
                args, input_data_train, input_data_test, sess, 100, directories)

    return cost, LH, ts, intensities


def get_train_data_nfold(args, data_params, data_collection, fold):
    """ get train data for n fold cross validation """

    # get input data from collection
    input_data_train, support, f = get_train_data(
        args, data_params, fold * args.batch_size, data_collection)

    input_data = np.array_split(input_data_train, fold)

    return input_data, f


def get_train_data_taxi_nminibatch(args, N_batch):
    """ get train data for n fold cross validation """

    # get input data from collection
    input_data_train = get_train_data_taxi(
        args, N_batch * args.batch_size)

    input_data = np.array_split(input_data_train, N_batch)

    input_data_cut = []
    # cur per minibatch
    for i, data in enumerate(input_data):
        # cut to minimum sequence length
        min_seq = min([x.shape[0] for x in data])
        input_data_cut.append(np.asarray([x[:min_seq] for x in data]))

    return input_data_cut


def train_fold(input_data, data_params, model_params, args, fold, collection, n_trials):
    """ trains one fold and store results to the database """

    input_data_test = copy.deepcopy(input_data)[fold]
    input_data_train = copy.deepcopy(input_data)
    input_data_train.pop(fold)
    trials = range(n_trials)
    for trial in tqdm.tqdm(trials, total=n_trials):
        param_field = {'model_params': model_params, 'data_params': data_params,
                       'fold': fold, 'trial': trial}
        # to ensure that for one parameter combination we have only one unique document in the db
        collection.update(param_field, param_field, upsert=True)
        doc = collection.find_one({'model_params': model_params})
        model_id = str(doc['_id'])
        directories = {'model': args.store_dir + '/' + model_id,
                       'summary': args.log_dir}

        cost, LH, ts, intensities = train(
            input_data_test, input_data_train, args, directories)

        # import pdb
        # pdb.set_trace()

        # write result into database
        collection.update_one(
            param_field,
            {
                "$set": {
                    "cost": [float(cost_e) for cost_e in cost],
                    "likelihood": [float(LH_e) for LH_e in LH],
                    "intensities": [[float(int_e_t) for int_e_t in int_e] for int_e in intensities],
                    "ts": [[int(ts_e_t) for ts_e_t in ts_e] for ts_e in ts],
                }
            }
        )


def train_trial(input_data_train, input_data_test, model_params, args, data_params, collection, trial):
    """ trains one fold and store results to the database """

    param_field = {'model_params': model_params, 'data_params': data_params,
                   'trial': trial}
    # to ensure that for one parameter combination we have only one unique document in the db
    collection.update(param_field, param_field, upsert=True)
    doc = collection.find_one({'model_params': model_params})
    model_id = str(doc['_id'])
    directories = {'model': args.store_dir + '/' + model_id,
                   'summary': args.log_dir}

    cost, LH, ts, intensities = train(
        input_data_test, input_data_train, args, directories)

    # import pdb
    # pdb.set_trace()

    # write result into database
    collection.update_one(
        param_field,
        {
            "$set": {
                "cost": [float(cost_e) for cost_e in cost],
                "likelihood": [float(LH_e) for LH_e in LH],
                "intensities": [[float(int_e_t) for int_e_t in int_e] for int_e in intensities],
                "ts": [[int(ts_e_t) for ts_e_t in ts_e] for ts_e in ts],
            }
        }
    )