import os
import sys
import json
import datetime
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt


def covariates_from_line(trip_id, taxi_line, date0, datef):
    try:
        pickup_datetime = datetime.datetime.strptime(taxi_line[1],
                                                     '%Y-%m-%d %H:%M:%S')
        dropoff_datetime = datetime.datetime.strptime(taxi_line[2],
                                                      '%Y-%m-%d %H:%M:%S')
    except:
        return []

    if pickup_datetime > date0:
        if (pickup_datetime < datef):
            if (dropoff_datetime < datef):
                pickup_datetime = (pickup_datetime - date0).total_seconds()
                dropoff_datetime = (dropoff_datetime - date0).total_seconds()
            else:
                try:
                    dropoff_datetime = (pickup_datetime - datef).total_seconds()
                    pickup_datetime = (pickup_datetime - date0).total_seconds()
                except:
                    print sys.exc_info()
            passenger_count = int(taxi_line[3])
            trip_distance = float(taxi_line[4])
            PULocationID = int(taxi_line[7])
            DOLocationID = int(taxi_line[8])
            tolls_amount = float(taxi_line[14])
            return np.asarray([
                trip_id, pickup_datetime, passenger_count, tolls_amount, PULocationID,DOLocationID, trip_distance, dropoff_datetime
            ])
        else:
            return []
    else:
        return []

def create_arrivals_and_covariates_folder(folder_counter,
                                          results_dir,
                                          taxi_file,
                                          taxi_data,
                                          date0,
                                          datef,
                                          size_of_chunks=90000,
                                          size_from_file=100000):
    """

    :param results_dir:
    :param taxi_file:
    :param taxi_data:
    :param date0:
    :param datef:
    :param size_of_chunks:
    :param size_from_file:
    :return:

    trip_id, pickup_datetime, passenger_count, tolls_amount, PULocationID, DOLocationID, trip_distance, dropoff_datetime
    """
    columns_string = "trip_id,pickup_seconds,passenger_count,tolls_amount,PULocationID,DOLocationID,trip_distance,dropoff_seconds".split(
        ",")
    arrivals_header = dict(zip(range(len(columns_string)), columns_string))
    departures_header = {0: "trip_id", 1: "dropoff_seconds"}
    #==========================
    # LOAD ONE FULL FILE
    #==========================
    ARRIVALS_AND_COVARIATES = []
    for trip_id, taxi_line in enumerate(taxi_data[2:int(size_from_file + 2)]):
        covariates_line = covariates_from_line(trip_id, taxi_line.split(","), date0, datef)
        if len(covariates_line) > 0.:
            ARRIVALS_AND_COVARIATES.append(covariates_line)

    ARRIVALS_AND_COVARIATES = np.asarray(ARRIVALS_AND_COVARIATES)
    print ARRIVALS_AND_COVARIATES.shape
    # sort according to arrival
    pickup_sorted_index = ARRIVALS_AND_COVARIATES[:, 1].argsort()
    ARRIVALS_AND_COVARIATES_all = ARRIVALS_AND_COVARIATES[pickup_sorted_index, :]

    #=========================
    #DIVIDE IN TREE DAYS
    #=========================
    day_before = date0
    for i in range(3):
        day_next = day_before + datetime.timedelta(days=1)
        lower_border = (day_before-date0).total_seconds()
        upper_border = (day_next-date0).total_seconds()
        support_size = (day_next - day_before).total_seconds()

        day_block_index = np.where(np.logical_and(ARRIVALS_AND_COVARIATES_all[:,1] < upper_border, ARRIVALS_AND_COVARIATES_all[:,1] > lower_border))[0]
        ARRIVALS_AND_COVARIATES = ARRIVALS_AND_COVARIATES_all[day_block_index,:]
        ARRIVALS_AND_COVARIATES[:,1] = ARRIVALS_AND_COVARIATES[:,1] - lower_border
        ARRIVALS_AND_COVARIATES[:, -1] = ARRIVALS_AND_COVARIATES[:, -1] - lower_border
        on_day_index = np.where(ARRIVALS_AND_COVARIATES[:,-1] < support_size)[0]
        ARRIVALS_AND_COVARIATES = ARRIVALS_AND_COVARIATES[on_day_index ,:]

        try:
            print max(ARRIVALS_AND_COVARIATES[:,-1])
        except:
            print None


        day_before = day_next

        total_arrivals = 0.
        total_departures = 0.
        metadata = {
            "arrivals_files": {},
            "arrivals_file_header":
            arrivals_header,
            "departures_files": {},
            "departures_file_header":
            departures_header,
            "real_temporal_support": [(day_before.strftime("%Y-%m-%d %H:%M:%S"),
                                       day_next.strftime("%Y-%m-%d %H:%M:%S"))],
            "temporal_support": [(0., support_size)],
            "total_arrivals":
            0.,
            "total_departures":
            0.,
            "last_event":0.
        }

        current_folder = results_dir + str(folder_counter) + "/"
        os.makedirs(current_folder)
        folder_counter+=1
        #==========================
        # PADDINGS
        #==========================
        num_rows, num_cols = ARRIVALS_AND_COVARIATES.shape
        # SPLIT OF DATASETS
        number_of_splits = (num_rows / size_of_chunks)
        residual = num_rows % size_of_chunks

        if (num_rows % size_of_chunks) != 0:
            number_of_splits += 1
            padding_size = size_of_chunks - residual
        else:
            padding_size = 0

        # SPLIT FOR ARRIVALS
        ARRIVALS_AND_COVARIATES = np.pad(ARRIVALS_AND_COVARIATES,
                                         ((0, padding_size), (0, 0)), "constant")
        last_arrival = ARRIVALS_AND_COVARIATES[-1,1]
        ARRIVALS_AND_COVARIATES = np.split(
            ARRIVALS_AND_COVARIATES, number_of_splits, axis=0)
        for chunk_id, chunk in enumerate(ARRIVALS_AND_COVARIATES):
            np.savetxt(
                current_folder +
                "arrivals_covariates_{0}_.txt".format(chunk_id + 1), chunk)
            total_arrivals += chunk.shape[0]
            metadata["arrivals_files"][chunk_id + 1] = chunk.shape[0]

        # SORT FOR DEPARTURES
        ARRIVALS_AND_COVARIATES = np.concatenate(ARRIVALS_AND_COVARIATES, axis=0)
        pickup_sorted_index = ARRIVALS_AND_COVARIATES[:, -1].argsort()
        ARRIVALS_AND_COVARIATES = ARRIVALS_AND_COVARIATES[pickup_sorted_index, :]
        ARRIVALS_AND_COVARIATES = ARRIVALS_AND_COVARIATES[:, [0, -1]]
        last_departure = ARRIVALS_AND_COVARIATES[-1,1]

        ARRIVALS_AND_COVARIATES = np.split(
            ARRIVALS_AND_COVARIATES, number_of_splits, axis=0)
        for chunk_id, chunk in enumerate(ARRIVALS_AND_COVARIATES):
            np.savetxt(
                current_folder +
                "departures_covariates_{0}_.txt".format(chunk_id + 1), chunk)
            metadata["departures_files"][chunk_id + 1] = chunk.shape[0]
            total_departures += chunk.shape[0]

        metadata["total_arrivals"] = total_arrivals
        metadata["total_departures"] = total_departures
        metadata["last_event"] = max(last_arrival,last_departure)
        json.dump(metadata, open(current_folder + "metadata.json", "w"))
    return  folder_counter

if __name__ == "__main__":

    data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/raw_data/taxi/"
    #results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi_2/"
    results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/taxi_new/"


    #data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/RawData/Taxi/"
    #results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Taxi_Day/"

    folder_counter = 0
    all_taxi_files = os.listdir(data_dir)
    for taxi_file in all_taxi_files:
        taxi_file = taxi_file[16:-4]
        taxi_year = int(taxi_file.split("-")[0])
        taxi_month = int(taxi_file.split("-")[1])

        if taxi_month == 12:
            date0 = datetime.datetime(taxi_year, taxi_month, 1)
            datef = datetime.datetime(taxi_year, taxi_month, 4)
        else:
            date0 = datetime.datetime(taxi_year, taxi_month, 1)
            datef = datetime.datetime(taxi_year, taxi_month, 4)

        taxi_data = open(
            data_dir + 'yellow_tripdata_{0}.csv'.format(taxi_file),
            "r").read().split("\n")
        folder_counter = create_arrivals_and_covariates_folder(
            folder_counter,
            results_dir,
            taxi_file,
            taxi_data,
            date0,
            datef,
            size_of_chunks=1000,
            size_from_file=10000000)
