

import os
import json
import numpy as np
from shutil import copyfile

def concatenate_hidden_states(covariates_data,hidden_data,concatenated_data,batch_size=500,synthetic_file=True):
    """"
    """
    bath_folders = [a for a in os.listdir(covariates_data)]
    bath_folders.sort()
    HIDDEN_TRAIN = np.load(hidden_data + "train_hidden_states.npy")
    HIDDEN_TRAIN = np.concatenate(HIDDEN_TRAIN, axis=0)
    HIDDEN_TEST = np.load(hidden_data + "test_hidden_states.npy")
    HIDDEN_TEST = np.concatenate(HIDDEN_TEST, axis=0)

    number_of_arrivals_in_test = 25
    for folder_id in range(batch_size):
        ready = False
        bath_folder = covariates_data + bath_folders[folder_id]
        metadata = json.loads(open(covariates_data + bath_folders[folder_id] + "/metadata.json", "r").read())
        total_arrivals = metadata["total_arrivals"]
        number_of_arrivals_in_train = total_arrivals - number_of_arrivals_in_test
        if synthetic_file:
            files_number_in_batch = [int(a[-6]) for a in os.listdir(bath_folder) if "arrivals_covariates" in a]
        else:
            files_number_in_batch = [int(a[-5]) for a in os.listdir(bath_folder) if "arrivals_covariates" in a]

        files_number_in_batch.sort()
        # undo seq lenght slides
        hidden_train = np.concatenate(HIDDEN_TRAIN[folder_id, :, :], axis=0)

        number_of_arrivals_left_in_train = number_of_arrivals_in_train
        number_of_arrivals_to_file = 0
        os.makedirs(concatenated_data + bath_folders[folder_id])
        #================
        # copy departures
        #================
        departures_files = [a for a in os.listdir(bath_folder) if "departures" in a]
        for d_file in departures_files:
            copyfile(bath_folder+"/"+d_file,concatenated_data + bath_folders[folder_id]+"/"+d_file)
        for file_num in files_number_in_batch:
            # ==============================
            # read files
            if not synthetic_file:
                arrivals_covariates = np.loadtxt(bath_folder + "/arrivals_covariates_{0}.txt".format(file_num))
            else:
                arrivals_covariates = np.loadtxt(bath_folder + "/arrivals_covariates_{0}_.txt".format(file_num))

            number_of_arrivals_in_file = arrivals_covariates.shape[0]
            # ==============================
            # HOW MUCH TO CONCATENATE
            # ==============================
            if number_of_arrivals_left_in_train > number_of_arrivals_in_file:
                hidden_for_file = hidden_train[
                                  number_of_arrivals_to_file:number_of_arrivals_to_file + number_of_arrivals_in_file]
            else:
                hidden_for_file = hidden_train[
                                  number_of_arrivals_to_file:number_of_arrivals_to_file + number_of_arrivals_left_in_train]
                hidden_for_file = np.concatenate([hidden_for_file, HIDDEN_TEST[folder_id]])
                ready = True
            concatenated = np.concatenate([arrivals_covariates, hidden_for_file], axis=1)
            np.savetxt(
                concatenated_data + bath_folders[folder_id] + "/arrivals_covariates_hidden_{0}.txt".format(file_num),
                concatenated)
            if ready:
                break

            number_of_arrivals_to_file += number_of_arrivals_in_file
            number_of_arrivals_left_in_train = number_of_arrivals_in_train - number_of_arrivals_to_file
            # =================
            number_of_arrivals_to_file += number_of_arrivals_in_file


if __name__=="__main__":
    #covariates_data = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/Concatenate/github_covariates/"
    #hidden_data = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/Concatenate/github_hidden/"
    #concatenated_data = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/Concatenate/github_concatenated/"


    covariates_data = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/concatenate_synthetic/nonlinear_phase_type/"
    hidden_data = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/concatenate_synthetic/hidden_states/synthetic/phase/"
    concatenated_data = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/concatenate_synthetic/nonlinear_phase_hidden_covariates/"

    #covariates_data = "/parallel/stackoverflow_covariates/"
    #hidden_data = "/BigData/anon/hidden_states/stackoverflow/"
    #concatenated_data "/BigData/anon/covariates_hidden/stackoverflow/"

    #covariates_data = "/BigData/anon/taxi_covariates/"
    #hidden_data = "/BigData/anon/hidden_states/ny/"
    #concatenated_data = "/BigData/anon/covariates_hidden/ny/"

    #covariates_data = "/parallel/github_covariates/"
    #hidden_data = "/BigData/anon/hidden_states/github/"
    #concatenated_data = "/BigData/anon/covariates_hidden/github/"

    concatenate_hidden_states(covariates_data, hidden_data, concatenated_data,500,True)

