'''
Created on Aug 17, 2018

@author: anon
'''
import json
from math import ceil
import os
from glob import glob
from io import open

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import minmax_scale


def split_train_test_and_preprocess(data,
                                    batch_size,
                                    window_size,
                                    test_size=.1):
    train, test = train_test_split(
            data, test_size=test_size, shuffle=False, random_state=1)

    N = len(train)
    R = N % batch_size
    train_set = preprocess_data(train[:N - R], batch_size, window_size)
    test_set = preprocess_data(test, len(test), window_size)
    test_set = [t[0] for t in test_set]
    return train_set, test_set


def preprocess_data(x, batch_size, window_size):
    """
    Preprocess input sequences with variable length for training with BPTT.

    It is added the sequence length to each window and this information is used for the dynamic rnn.

    Args:
        x (list): N sequences, each sequence i has length L_i and each item j in the sequence i has dimension D.
        batch_size (int):
        window_size (int):

    Returns:
        ndarray: with shape [num_batches, batch_size, num_windows, window_size, D]
    """
    D = len(x[0][0])
    N = len(x)

    seqences_l = np.asarray(list(map(len, x)))

    max_s_l = np.max(seqences_l)
    reminder = max_s_l % window_size
    if reminder == 0:
        max_l = max_s_l
    else:
        max_l = max_s_l + (window_size - reminder)
    data = []
    for seq, s_len in zip(x, seqences_l):
        s_padded = np.pad(seq, ((0, max_l - s_len), (0, 0)), 'constant')
        data.append(s_padded)
    data = np.asarray(data)
    num_windows = ceil(max_l / window_size)
    f_w = seqences_l // window_size
    p_w = seqences_l % window_size
    z_w = num_windows - f_w - 1
    f_w_m = list(map(lambda x: np.ones(max_s_l // window_size) * window_size, f_w))
    if reminder != 0:
        f_w_m = list(map(lambda x, y: np.append(x, y), f_w_m, p_w))
        z_w_m = list(map(lambda x: np.zeros(x), z_w))
        f_w_m = list(map(lambda x, y: np.append(x, y), f_w_m, z_w_m))
    f_w_m = np.asarray(f_w_m)
    mask = np.zeros((N, max_l))
    for i in range(N):
        mask[i, :seqences_l[i]] = 1
    data = data.reshape(-1, batch_size, num_windows, window_size, D)
    w_l = f_w_m.reshape(-1, batch_size, num_windows)
    mask = mask.reshape(-1, batch_size, num_windows, window_size)
    return data, w_l, mask


def load_mempool_data(path):
    data = np.load(path)
    data = data.tolist()
    T = np.max([np.max([a[0] for a in m]) for m in data])
    d = []
    for m in data:
        m = np.asarray(m)
        m[:, :2] /= T
        d.append(m.tolist())
    return d, T


def load_github_arrivals(path):
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders[:180]:
        files = glob(path + '/' + d + '/arrivals_covariates_1.txt')
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 31536000.0
            data[:, 1] = data[:, 1] / T
            delta_a = data[1:, 1] - data[:-1, 1]
            data = np.concatenate(
                    (data[1:, 1, None], delta_a[:, None]), axis=1)
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)

    return batch, T


def load_github_service_time(path):
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders:
        files = glob(path + '/' + d + '/arrivals_covariates_1.txt')
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 31536000.0
            service_time = np.where(
                    data[1:, -1] > 0, data[1:, -1] - data[1:, 1], data[1:, -1])
            data = service_time[:, None] / T
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)

    return batch, T


def load_synthetic_nonlinear_arrivals(path):
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders[:10]:
        files = glob(path + '/' + d + '/arrivals_covariates_*_.txt')
        files_n = map(lambda x: int(x.split('_')[-2]), files)
        idx = np.argsort(files_n)
        files = np.asarray(files)[idx]
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 1000.0
            data[:, 1] = data[:, 1] / T
            delta_a = data[1:, 1] - data[:-1, 1]
            data = np.concatenate(
                    (data[1:, 1, None], delta_a[:, None]), axis=1)
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)

    return batch, T


def load_synthetic_nonlinear_service_time(path):
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders:
        files = glob(path + '/' + d + '/arrivals_covariates_*_.txt')
        files_n = map(lambda x: int(x.split('_')[-2]), files)
        idx = np.argsort(files_n)
        files = np.asarray(files)[idx]
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 1000.0
            service_time = np.where(
                    data[1:, -1] > 0, data[1:, -1] - data[1:, 1], data[1:, -1])
            data = service_time[:, None] / T
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)

    return batch, T


def load_stackoverflow_arrivals(path):
    # data = np.loadtxt(path)
    # T = -np.min(data[:, 12] - data[:, 2])
    # data[:, 2] = data[:, 2] / T
    # delta_a = data[1:, 2] - data[:-1, 2]
    # data = np.concatenate(
    #     (data[1:, 2, None], delta_a[:, None], data[1:, 7:12]), axis=1)
    # return data, T

    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders[:5]:
        files = glob(path + '/' + d + '/arrivals_covariates_*_.txt')
        files_n = map(lambda x: int(x.split('_')[-2]), files)
        idx = np.argsort(files_n)
        files = np.asarray(files)[idx]
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 604800.
            data[:, 2] = data[:, 2] / T
            delta_a = data[1:, 2] - data[:-1, 2]
            data = np.concatenate(
                    (data[1:, 2, None], delta_a[:, None], data[1:, 7:12]), axis=1)
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)
    return batch, T


def load_stackoverflow_service_time(path):
    # data = np.loadtxt(path)
    # T = -np.min(data[:, 12] - data[:, 2])
    # cov = data[1:, 7:12]
    # service_time = np.where(data[1:, 12] > 0, data[1:, 12] - data[1:, 2],
    #                         data[1:, 12])
    # data = np.concatenate((cov, service_time[:, None] / T), 1)
    # return data, T
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders:
        files = glob(path + '/' + d + '/arrivals_covariates_*_.txt')
        files_n = map(lambda x: int(x.split('_')[-2]), files)
        idx = np.argsort(files_n)
        files = np.asarray(files)[idx]
        row = []
        print("ljasdlkajsdlkajsdlkasjdlasjd",files)
        import pdb;pdb.set_trace()
        for f in files:
            data = np.loadtxt(f)
            T = 604800.
            service_time = np.where(
                    data[1:, -1] > 0, data[1:, -1] - data[1:, 2], data[1:, -1])
            cov = data[1:, 7:12]
            data = np.concatenate((cov, service_time[:, None] / T), 1)
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)

    return batch, T


def load_ny_arrivals(path):
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders:
        files = glob(path + '/' + d + '/arrivals_covariates_*_.txt')
        files_n = map(lambda x: int(x.split('_')[-2]), files)
        idx = np.argsort(files_n)
        files = np.asarray(files)[idx]
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 86400.
            data[:, [1, 7]] = data[:, [1, 7]] / T
            delta_a = data[1:, 1] - data[:-1, 1]
            covar = data[1:, 4:7]
            data = np.concatenate(
                    (data[1:, 1, None], delta_a[:, None], covar), axis=1)
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)
    return batch, T


def load_ny_service_time(path):
    folders = [f for f in os.listdir(
            path) if os.path.isdir(os.path.join(path, f))]
    folders = np.sort(folders)
    batch = []
    for d in folders:
        files = glob(path + '/' + d + '/arrivals_covariates_*_.txt')
        files_n = map(lambda x: int(x.split('_')[-2]), files)
        idx = np.argsort(files_n)
        files = np.asarray(files)[idx]
        row = []
        for f in files:
            data = np.loadtxt(f)
            T = 86400.
            service_time = np.where(
                    data[1:, -1] > 0, data[1:, 7] - data[1:, 1], data[1:, -1])
            cov = data[1:, 4:7]
            data = np.concatenate((cov, service_time[:, None] / T), 1)
            row.append(data)
        row = np.concatenate(row, axis=0)
        batch.append(row)

    return batch, T


def load_service_events(path):
    data = {
        'arrivals': [],
        'covariates': [],
        'service_events': [],
        'service_time': [],
        'ids': []
    }
    with open(path) as lines:
        for line in lines:
            try:
                o = json.loads(line)
                data['arrivals'].append(o['a'])
                data['ids'].append(o['i'])
                data['covariates'].append(o['covariates'])
                data['service_time'].append([o['s']] + o['covariates'])
                s_e = minmax_scale(o['service_events'])
                s_e_d = s_e[1:, 0] - s_e[:-1, 0]
                s_e_d = np.insert(s_e_d, 0, 0.)
                s_e = np.stack((s_e[:, 0], s_e_d, s_e[:, 1]), 1)
                elapsed = s_e[:, 0] - o['a']
                remain = (np.abs(o['s']) - elapsed) * np.sign(o['s'])
                s_e = minmax_scale(
                        np.concatenate(
                                (s_e, remain.reshape(-1, 1), elapsed.reshape(-1, 1)),
                                1))
                data['service_events'].append(
                        map(lambda x: x + o['covariates'], s_e.tolist()))
            except Exception as e:
                print(e)
    arrivals = minmax_scale(data['arrivals'])
    delta_arrivals = arrivals[1:] - arrivals[:-1]
    arrivals = np.stack((arrivals[1:], delta_arrivals), 1)
    data['arrivals'] = arrivals

    return data
