import os
import copy
import datetime
import numpy as np
import multiprocessing as mp

# =============================
# SYNTHETIC /GITHUB
# =============================

"""

# parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/ProcessorSharing/"
# data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/ProcessorSharing/0/"

#parallel_data_dir = "data/Synthetic/nonlinear_phase_type/"
#data_dir = "data/Synthetic/nonlinear_phase_type/0/"

parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Nonlinear_PS/"
data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Nonlinear_PS/0/"

DEPARTURES_TIME_INDEX = 1
ARRIVALS_TIME_INDEX = 1
ARRIVALS_ID_INDEX = 0
DEPARTURES_ID_INDEX = 0
REAL_DEPARTURE_TIME_INDEX = 2
U_SELECTION = np.asarray([1,2,3,4])
U_DIMENSION = len(U_SELECTION)

{"0": "arrival_id", "1": "arrivals_time", "2": "departure_time"}
"""

# =============================
# STACK EXCHANGE
# =============================

# """
# data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/bitcoin_exchange_debug/"
# data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/Service_Data/bitcoin_exchange/2011/"
# parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/Service_Data/bitcoin_exchange/"

parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Stackoverflow_weeks_small/"
data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Stackoverflow_weeks_small/0/"

DEPARTURES_TIME_INDEX = 1
ARRIVALS_TIME_INDEX = 2
ARRIVALS_ID_INDEX = 1
DEPARTURES_ID_INDEX = 0
MAXIMUM_SERVICE_LOAD = 2
REAL_DEPARTURE_TIME_INDEX = 12
U_SELECTION = np.asarray([2, 7, 8, 9, 10, 11, 12, 13, 14])
U_DIMENSION = len(U_SELECTION)

# {"0": "arrivals_id",
# "1": "accepted_answer_id",
# "2": "arrival_time",
# "3": "owner_id",
# "4": "score",
# "5": "favorite_count",
# "6": "view_count",
# "7": "tag0",
# "8": "tags1",
# "9": "tags2",
# "10": "tags3",
# "11": "tags4",
# "12": "departure_time}

# """

# =============================
# TAXI
# =============================

"""
# np.asarray([trip_id, pickup_datetime, passenger_count, tolls_amount, PULocationID, DOLocationID, trip_distance, service_time, dropoff_datetime])

# parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi/"
# data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi/2017-01/"

parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Taxi_Small/"
data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Taxi_Small/2017-01/"

#parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi_min/"
#data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi_min/2017-03/"

# parallel_data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi_small/"
# data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi_small/2017-02/"


DEPARTURES_TIME_INDEX = 1
ARRIVALS_TIME_INDEX = 1
REAL_DEPARTURE_TIME_INDEX = 7
ARRIVALS_ID_INDEX = 0
DEPARTURES_ID_INDEX = 0
U_SELECTION = np.asarray([1,4,5,6,7,8,9])
U_DIMENSION = len(U_SELECTION)

{"0": "trip_id",
 "1": "pickup_seconds",
 "2": "passenger_count",
 "3": "tolls_amount",
 "4": "PULocationID",
 "5": "DOLocationID",
 "6": "trip_distance",
 "7": "dropoff_seconds"}
"""

# ====GENERAL=====
RESAMPLE_WAIT = 100
MAXIMUM_SERVICE_LOAD = 3
SEQ_LENGHT = 3
ONE_BATCH_SIZE = SEQ_LENGHT * MAXIMUM_SERVICE_LOAD * U_DIMENSION
ONE_SEQUENCE_STEP_SIZE = MAXIMUM_SERVICE_LOAD * U_DIMENSION


def initialize_system_index(data_dir):
    """
    :param data_dir:
    :return:
    """
    all_arrivals_files = [a for a in os.listdir(
        data_dir) if "arrivals_covariates" in a]
    arrivals_files_numbers = [int(arrivals_file_name.split(
        "_")[-2]) for arrivals_file_name in all_arrivals_files]
    arrivals_files_numbers.sort()
    all_arrivals_files = ["arrivals_covariates_{0}_.txt".format(
        num) for num in arrivals_files_numbers]
    all_departures_files = [a for a in os.listdir(
        data_dir) if "departures_covariates" in a]
    departures_files_numbers = [int(arrivals_file_name.split(
        "_")[-2]) for arrivals_file_name in all_departures_files]
    departures_files_numbers.sort()
    all_departures_files = ["departures_covariates_{0}_.txt".format(
        num) for num in departures_files_numbers]

    arrivals_covariates_full = np.loadtxt(data_dir + all_arrivals_files[0])
    departures_covariates_full = np.loadtxt(data_dir + all_departures_files[0])

    arrival_time = arrivals_covariates_full[0][ARRIVALS_TIME_INDEX]
    departure_time = departures_covariates_full[0][DEPARTURES_TIME_INDEX]
    if arrival_time < departure_time:
        m_j = 0
        t_j = arrival_time
        arrivals_on_service = 1
    else:
        m_j = 1
        t_j = departure_time
        arrivals_on_service = 0

    total_num_arrivals_file = len(all_arrivals_files)
    total_num_departures_file = len(all_departures_files)

    system_index = {"data_dir": data_dir,
                    "current_arrival_file_index": 1,
                    "arrivals_index_in_file": 0,
                    "current_departure_file_index": 1,
                    "departure_index_in_file": 0,
                    "arrivals_time": arrival_time,
                    "arrivals_covariate": arrivals_covariates_full[0],
                    "departure_time": departure_time,
                    "departures_covariate": departures_covariates_full[0],
                    "m_j": m_j,
                    "t_j": t_j,
                    "j": 0,
                    "arrivals_on_service": arrivals_on_service,
                    "arrivals_covariates_full": arrivals_covariates_full,
                    "departures_covariates_full": departures_covariates_full,
                    "arrivals_covariates_num": arrivals_covariates_full.shape[0],
                    "departures_covariates_num": departures_covariates_full.shape[0],
                    "total_num_arrivals_file": total_num_arrivals_file,
                    "total_num_departures_file": total_num_departures_file}
    return system_index, None



def update_system_index(system_index_old):
    """
    obtain the next event by comparing all next arrivals and departures

    load new files if index are exceeded, and updates system index

    :return:
    """
    system_index = copy.deepcopy(system_index_old)

    data_dir = system_index["data_dir"]
    current_arrival_file_index = system_index["current_arrival_file_index"]
    arrivals_index_in_file = system_index["arrivals_index_in_file"]
    current_departure_file_index = system_index["current_departure_file_index"]
    departure_index_in_file = system_index["departure_index_in_file"]
    j = system_index["j"]
    arrivals_on_service = system_index["arrivals_on_service"]
    # if last event arrivals, move in arrivals
    if system_index["m_j"] == 0:
        # next arrival inside current file
        if arrivals_index_in_file < system_index["arrivals_covariates_num"] - 1:
            arrivals_index_in_file += 1
            system_index["arrivals_index_in_file"] = arrivals_index_in_file
            arrival_time = system_index["arrivals_covariates_full"][arrivals_index_in_file][ARRIVALS_TIME_INDEX]
            system_index["arrivals_covariate"] = system_index["arrivals_covariates_full"][arrivals_index_in_file]
        else:  # next arrival in new file
            if current_arrival_file_index < system_index["total_num_arrivals_file"]:
                current_arrival_file_index += 1
                system_index["current_arrival_file_index"] = current_arrival_file_index

                arrivals_index_in_file = 0  # start of new file
                system_index["arrivals_index_in_file"] = arrivals_index_in_file
                arrivals_covariates_full = np.loadtxt(
                    data_dir + 'arrivals_covariates_{0}_.txt'.format(current_arrival_file_index))
                system_index["arrivals_covariates_full"] = arrivals_covariates_full
                system_index["arrivals_covariates_num"] = arrivals_covariates_full.shape[0]
                arrival_time = system_index["arrivals_covariates_full"][arrivals_index_in_file][ARRIVALS_TIME_INDEX]
                system_index["arrivals_covariate"] = system_index["arrivals_covariates_full"][arrivals_index_in_file]
            else:
                arrival_time = np.infty
        departure_time = system_index["departure_time"]
        system_index["arrivals_time"] = arrival_time
    # if last event departure move in departures
    if system_index["m_j"] == 1:
        # next departure inside current file
        if departure_index_in_file < system_index["departures_covariates_num"] - 1:
            departure_index_in_file += 1
            system_index["departure_index_in_file"] = departure_index_in_file
            departure_time = system_index["departures_covariates_full"][departure_index_in_file][DEPARTURES_TIME_INDEX]
            system_index["departures_covariate"] = system_index["departures_covariates_full"][arrivals_index_in_file]
        else:  # next arrival in new file
            if current_departure_file_index < system_index["total_num_departures_file"]:
                current_departure_file_index += 1
                system_index["current_departure_file_index"] = current_departure_file_index

                departure_index_in_file = 0  # start of new file
                system_index["departure_index_in_file"] = departure_index_in_file
                departures_covariates_full = np.loadtxt(
                    data_dir + 'departures_covariates_{0}_.txt'.format(current_departure_file_index))
                system_index["departures_covariates_full"] = departures_covariates_full
                system_index["departures_covariates_num"] = departures_covariates_full.shape[0]
                departure_time = system_index["departures_covariates_full"][departure_index_in_file][
                    DEPARTURES_TIME_INDEX]
                system_index["departures_covariate"] = system_index["departures_covariates_full"][
                    departure_index_in_file]
            else:
                departure_time = np.infty
        system_index["departure_time"] = departure_time
        arrival_time = system_index["arrivals_time"]
    # ========================================
    # UPDATE CURRENT TIME
    # ========================================
    if arrival_time < departure_time:
        m_j = 0
        t_j = arrival_time
        arrivals_on_service = arrivals_on_service + 1
    else:
        m_j = 1
        t_j = departure_time
        if arrivals_on_service > 0:
            arrivals_on_service = arrivals_on_service - 1

    j += 1
    system_index["m_j"] = m_j
    system_index["t_j"] = t_j
    system_index["j"] = j
    system_index["arrivals_on_service"] = arrivals_on_service

    return system_index


def next_data_input(input):
    """
    arrival_and_covariates = [arrivals_id (0), accepted_answer_id (1), arrival_time (2), owner_id (3),
                              score (4), favorite_count (5), view_count (6), departure_time(7), tags(5) (7+5)]

    departure_and_covariates = [arrivals_id,arrival_time,parent_question,owner_id,
                                score]

    :param data_dir:
    :param system_index:
    :param U_truncated:
    :return: [1,BACKPROP_SEQ_LENGHT,MAX_SERVICE_LOAD,input_to_training]
    *input_to_training: elapsed_time,
                        covariates,
                        true_remaining_service_time,
                        time_to_next_event (float),
                        leaving_system (bool)
    """

    batch_index = input[0]
    past_t_j = input[1]
    events_batch = input[2]
    batch = input[3]
    system_index = input[4]
    U_truncated = input[5]

    if system_index["t_j"] != np.inf:
        EVENTS_BATCH = []
        BATCH = []
        LEAVING_SYSTEM = []
        IN_SERVICE = []
        wait_index = 0

        for j in range(SEQ_LENGHT):
            system_index_next = update_system_index(system_index)

            #==========================================
            # CURRENT BATCH ENDED
            #==========================================
            if np.isnan(system_index["t_j"] - past_t_j):
                left_to_padded = SEQ_LENGHT - len(EVENTS_BATCH)
                U = np.zeros((left_to_padded, MAXIMUM_SERVICE_LOAD, U_DIMENSION))

                BATCH = np.asarray(BATCH)
                BATCH = np.concatenate([BATCH, U], axis=0)

                EVENTS_BATCH = np.asarray(EVENTS_BATCH)
                non_zero = EVENTS_BATCH.shape[0]
                EVENTS_BATCH = np.concatenate([EVENTS_BATCH, np.zeros((left_to_padded, 4))], axis=0)
                return batch_index, past_t_j, EVENTS_BATCH, BATCH, system_index, U_truncated, non_zero

            wait_index += 1
            EVENTS_BATCH.append(np.asarray([system_index["t_j"], system_index["t_j"] - past_t_j, system_index["m_j"],system_index["arrivals_on_service"]]))

            # =========================================================
            # SERVICE DATA UPDATE ACCORDING TO EVENTS
            # =========================================================
            if system_index["m_j"] == 1:  # departure occurred
                if U_truncated is not None:
                    LEAVING_SYSTEM.append(system_index["departures_covariate"][DEPARTURES_ID_INDEX])
            elif system_index["m_j"] == 0:  # arrival ocurred
                IN_SERVICE.append(system_index["arrivals_covariate"])
            if wait_index == RESAMPLE_WAIT:
                pass
            # =========================================================
            # TIME TO NEXT EVENT
            # =========================================================
            "elapsed_time,covariates,true_remaining_service_time,time_to_next_event(float),leaving_system(bool)"
            if U_truncated is not None:
                pass
            else:
                system_index = system_index_next
                past_t_j = system_index["t_j"]
                U = np.zeros((MAXIMUM_SERVICE_LOAD, U_DIMENSION))
                BATCH.append(U)

        BATCH = np.asarray(BATCH)
        return batch_index, past_t_j, EVENTS_BATCH, BATCH, system_index, U_truncated, len(EVENTS_BATCH)
    #=====================================================
    # CURRENT SEQUENCE
    #=====================================================
    else:
        EVENTS_BATCH = np.zeros((SEQ_LENGHT, 4))
        BATCH = np.zeros((SEQ_LENGHT, MAXIMUM_SERVICE_LOAD, U_DIMENSION))
        return batch_index, system_index["t_j"], EVENTS_BATCH, BATCH, system_index, U_truncated, 0


def data_load(data_dir, maximum_time_of_study=1e12):
    past_t_j = 0.
    system_index, U_truncated = initialize_system_index(data_dir)
    for j in range(10):
        batch_index, past_t_j, EVENTS_BATCH, BATCH, system_index, U_truncated, non_zero = next_data_input(
            (0, past_t_j, None, None, system_index, U_truncated))

        print BATCH.shape

        print "Non zero ", non_zero
        print "\n"


# =========================================================
# PARALLEL CREATION OF BATCHES
# =========================================================


def parallel_data_load(data_dir, NUMBER_OF_EPOCHS=100, number_of_process=5):
    """
    [BATCH,BACKPROP_SEQ_LENGHT,MAX_SERVICE_LOAD,input_to_training]

    *input_to_training: elapsed_time,
                        covariates,
                        true_remaining_service_time,
                        time_to_next_event (float),
                        leaving_system (bool)

    :param data_dir:
    :param maximum_time_of_study:
    :return:
    """
    pool = mp.Pool(number_of_process)
    batch_files = [a for a in os.listdir(data_dir) if "tags" not in a]
    NUMBER_OF_BATCHES = len(batch_files)
    batch_directories = [data_dir + batch_file +
                         "/" for batch_file in batch_files]

    for epoch_index in range(NUMBER_OF_EPOCHS):
        SYSTEM_INDEX_PER_BATCH = [
            (j, 0., None, None, initialize_system_index(data_dir)[0], initialize_system_index(data_dir)[1], None) for
            j, data_dir
            in enumerate(batch_directories)]
        sequence_section_index = 0
        d0 = datetime.datetime.now()
        while True:
            SYSTEM_INDEX_PER_BATCH = pool.map(
                next_data_input, SYSTEM_INDEX_PER_BATCH)
            SYSTEM_INDEX_PER_BATCH.sort()
            EVENTS_BATCH = np.array([batch_line[2]
                                     for batch_line in SYSTEM_INDEX_PER_BATCH])
            BATCH = np.array([batch_line[3]
                              for batch_line in SYSTEM_INDEX_PER_BATCH])
            NON_ZERO = np.asarray([batch_line[-1]
                                   for batch_line in SYSTEM_INDEX_PER_BATCH])
            # ======================
            # HERE GOES TRAINING
            # ======================
            # print NON_ZERO
            # print EVENTS_BATCH.shape
            # print "\n"
            # =======================
            # HERE CHECK TO FINISH
            # =======================
            if sequence_section_index % 100 == 0:
                print "Steps into the sequence: ", sequence_section_index * SEQ_LENGHT
                df = datetime.datetime.now()
                print "steps time: ", (df - d0).total_seconds()
            sequence_section_index += 1
            min_t_j = min(
                [system_index["t_j"] for j, past_t_j, EVENTS_BATCH, BATCH, system_index, U, non_zero in
                 SYSTEM_INDEX_PER_BATCH])
            if min_t_j == np.infty:
                print "ALL SEQUENCES DONE"
                df = datetime.datetime.now()
                print "epoch time: ", (df - d0).total_seconds()
                break


if __name__ == "__main__":
    parallel_data_load(parallel_data_dir, NUMBER_OF_EPOCHS=2, number_of_process=10)
    # data_load(data_dir)