import os
import numpy as np
import json

def arrivals_departures_deep_service_file_data(ARRIVALS_DEPARTURES,time_support):
    """
    :parameter:
        ARRIVALS_DEPARTURES: dictionary obtained from queues class simulation
                             {arrival:departure}
    :return:
    """
    number_of_arrivals = len(ARRIVALS_DEPARTURES)
    ARRIVALS_DEPARTURES_ARRAY = np.zeros((number_of_arrivals,3))
    j = 0
    for arrival, departure in ARRIVALS_DEPARTURES.iteritems():
        ARRIVALS_DEPARTURES_ARRAY[j, 0] = j
        if arrival < time_support[0][1]:
            if departure < time_support[0][1]:
                ARRIVALS_DEPARTURES_ARRAY[j,1] = arrival
                ARRIVALS_DEPARTURES_ARRAY[j,2] = departure
            else:
                ARRIVALS_DEPARTURES_ARRAY[j, 1] = arrival
                ARRIVALS_DEPARTURES_ARRAY[j, 2] = arrival - time_support[0][1]
        j+=1

    arrivals_index_sort = np.argsort(ARRIVALS_DEPARTURES_ARRAY[:,1])
    departures_index_sort = np.argsort(ARRIVALS_DEPARTURES_ARRAY[:,2])
    ARRIVALS_DEPARTURES_ARRAY_D = ARRIVALS_DEPARTURES_ARRAY[departures_index_sort,:]
    ARRIVALS_DEPARTURES_ARRAY_D = ARRIVALS_DEPARTURES_ARRAY_D[:,np.asarray([0, 2])]
    return ARRIVALS_DEPARTURES_ARRAY[arrivals_index_sort,:],ARRIVALS_DEPARTURES_ARRAY_D

def create_deep_service_format(list_of_ARRIVALS_AND_DEPARTURES,temporal_support,results_dir):
    """
    create the files format for the deep service distribution
    :return:
    """
    for batch_index,ARRIVALS_AND_DEPARTURES in enumerate(list_of_ARRIVALS_AND_DEPARTURES):
        current_folder = results_dir + str(batch_index) + "/"
        os.makedirs(current_folder)
        arrivals_cov, departures_cov = arrivals_departures_deep_service_file_data(ARRIVALS_AND_DEPARTURES,temporal_support)
        np.savetxt(current_folder+"arrivals_covariates_1_.txt",arrivals_cov)
        np.savetxt(current_folder + "departures_covariates_1_.txt",departures_cov)
        file_metadata = {"files":{1:{"num_arrivals":arrivals_cov.shape[0],"num_departures":departures_cov.shape[0]}},
                         "temporal_support":temporal_support,
                         "arrivals_file_header":{0:"arrival_id",1:"arrivals_time",2:"departure_time"},
                         "departures_file_header":{0:"departure_id",1:"departure_time"},
                         "total_arrivals":arrivals_cov.shape[0],
                         "total_departures":departures_cov.shape[0]}
        json.dump(file_metadata,open(current_folder+"metadata.json","w"))

#===========================
# OLD COMPLICATED MODEL
#===========================


def interarrivals_waiting(data_folder):
    """
    Reads the file created by

    multivariate hawkes ensemble


    :return:
    QUEUES_0: list of experiments .- arrivals and departures
    QUEUES_1: list of experiments .- interarrival time, service time

    """
    QUEUES_0 = []
    QUEUES_1 = []
    f = open(data_folder + "_arrivals.txt")
    for line in f:
        arrivals_departures_txt = line.split(" ")
        arrivals = map(float, [arrival for j, arrival in enumerate(arrivals_departures_txt) if (j % 4 == 0)][:-1])
        departures = map(float, [arrivals_departures_txt[k] for k in range(2, len(arrivals_departures_txt), 4)])
        arrivals_departures = np.stack((arrivals, departures))

        arrivals_departures = arrivals_departures[:, arrivals_departures[0].argsort()]
        arrivals_ordered = arrivals_departures[0]
        arrivals_ordered = np.hstack((0, arrivals_ordered))

        interarrival_a = arrivals_ordered[1:] - arrivals_ordered[:-1]
        waiting = arrivals_departures[1] - arrivals_departures[0]
        interarrivals_waiting = np.stack((interarrival_a, waiting))

        QUEUES_0.append(arrivals_departures)
        QUEUES_1.append(interarrivals_waiting)

    return QUEUES_0 ,QUEUES_1

def service_events(arrivals_and_departures):
    number_of_arrivals = arrivals_and_departures.shape[1]
    print "number_of_arrivals " ,number_of_arrivals
    events = arrivals_and_departures.reshape(-1)
    events.sort()

    events_index = 0
    arrival_index = 0
    service_events_dict = {}
    while arrival_index < number_of_arrivals:
        print "arrival_index " ,arrival_index
        if events[events_index] == arrivals_and_departures[0][arrival_index]:
            service_events_dict[arrival_index] = {"a": arrivals_and_departures[0][arrival_index]}
            if arrivals_and_departures[1][arrival_index] is not None:
                k = 0
                E = []
                while events[events_index + k] < arrivals_and_departures[1][arrival_index]:
                    E.append(events[events_index + k])
                    k += 1
            else:
                E = [a for a in events[events_index:-1] if a is not None]
            service_events_dict[arrival_index]["service_events"] = E
            service_events_dict[arrival_index]["departure"] = arrivals_and_departures[1][arrival_index]
            arrival_index += 1
        else:
            events_index += 1

    return service_events_dict