"""
The interface to load log datasets. The datasets currently supported include
HDFS and BGL.

Authors:
    LogPAI Team

"""

import pandas as pd
import os
import numpy as np
import re
import sys
from sklearn.utils import shuffle


def load_HDFS(log_file, label_file=None, window='session', train_ratio=0.8):
    """ Load HDFS structured log into train and test data

    Arguments
    ---------
        log_file: str, the file path of structured log.
        label_file: str, the file path of anomaly labels.
        window: str, the window options including `session` (default).
        train_ratio: float, the ratio of training data for train/test split.

    Returns
    -------
        (x_train, y_train): the training data
        (x_test, y_test): the testing data
    """

    assert window == 'session', "Only window=`session` is supported for HDFS dataset."

    if log_file.endswith('.npz'):
        pass

    elif log_file.endswith('.csv'):
        struct_log = pd.read_csv(log_file, engine='c', na_filter=False, memory_map=True)
        data_df = pd.DataFrame(columns=['EventSequence'])
        for idx, row in struct_log.iterrows():
            blkId_list = re.findall(r'(blk_-?\d+)', row['Content'])
            blkId_set = set(blkId_list)
            for blk_Id in blkId_set:
                if not blk_Id in data_df.index:
                    data_df.loc[blk_Id, 'EventSequence'] = []
                data_df.loc[blk_Id, 'EventSequence'].append(row['EventId'])

        label_data = pd.read_csv(label_file, engine='c', na_filter=False, memory_map=True)
        label_data = label_data.set_index('BlockId')
        label_dict = label_data['Label'].to_dict()
        data_df['Label'] = [1 if label_dict[x] == 'Anomaly' else 0 for x in data_df.index]
        # data_df['BLK'] = data_df.index
        # data_df.to_csv('out.csv', index=False)

        # Split train and test data
        pos_data = data_df[data_df['Label'] == 1].reset_index()
        neg_data = data_df[data_df['Label'] == 0].reset_index()
        num_train_pos = int(len(pos_data) * train_ratio)
        num_train_neg = int(len(neg_data) * train_ratio)
        train_data = pd.concat([pos_data.loc[:num_train_pos - 1, :], neg_data.loc[:num_train_neg - 1, :]])
        train_data = shuffle(train_data)
        test_data = pd.concat([pos_data.loc[num_train_pos:, :], neg_data.loc[num_train_neg:, :]])

        (x_train, y_train)  = (train_data['EventSequence'].values, train_data['Label'].values)
        (x_test, y_test) = (test_data['EventSequence'].values, test_data['Label'].values)

    else:
        raise NotImplementedError('load_HDFS only support csv and npz files!')

    print('====== Input data summary ======')
    print('Total: {} instances, {} anomaly, {} normal' \
          .format(len(data_df), len(pos_data), len(neg_data)))
    print('Train: {} instances, {} anomaly, {} normal' \
          .format(len(x_train), num_train_pos, num_train_neg))
    print('Test: {} instances, {} anomaly, {} normal\n' \
          .format(len(x_test), len(pos_data) - num_train_pos, len(neg_data) - num_train_neg))

    return (x_train, y_train), (x_test, y_test)


def load_BGL(log_file, label_file=None, window='sliding', time_interval=60, stepping_size=60, 
             train_ratio=0.8):
    """  TODO

    """


def bgl_preprocess_data(para, raw_data, event_mapping_data):
    """ split logs into sliding windows, built an event count matrix and get the corresponding label

    Args:
    --------
    para: the parameters dictionary
    raw_data: list of (label, time)
    event_mapping_data: a list of event index, where each row index indicates a corresponding log

    Returns:
    --------
    event_count_matrix: event count matrix, where each row is an instance (log sequence vector)
    labels: a list of labels, 1 represents anomaly
    """

    # create the directory for saving the sliding windows (start_index, end_index), which can be directly loaded in future running
    if not os.path.exists(para['save_path']):
        os.mkdir(para['save_path'])
    log_size = raw_data.shape[0]
    sliding_file_path = para['save_path']+'sliding_'+str(para['window_size'])+'h_'+str(para['step_size'])+'h.csv'

    #=============divide into sliding windows=========#
    start_end_index_list = [] # list of tuples, tuple contains two number, which represent the start and end of sliding time window
    label_data, time_data = raw_data[:,0], raw_data[:, 1]
    if not os.path.exists(sliding_file_path):
        # split into sliding window
        start_time = time_data[0]
        start_index = 0
        end_index = 0

        # get the first start, end index, end time
        for cur_time in time_data:
            if  cur_time < start_time + para['window_size']*3600:
                end_index += 1
                end_time = cur_time
            else:
                start_end_pair=tuple((start_index,end_index))
                start_end_index_list.append(start_end_pair)
                break
        # move the start and end index until next sliding window
        while end_index < log_size:
            start_time = start_time + para['step_size']*3600
            end_time = end_time + para['step_size']*3600
            for i in range(start_index,end_index):
                if time_data[i] < start_time:
                    i+=1
                else:
                    break
            for j in range(end_index, log_size):
                if time_data[j] < end_time:
                    j+=1
                else:
                    break
            start_index = i
            end_index = j
            start_end_pair = tuple((start_index, end_index))
            start_end_index_list.append(start_end_pair)
        inst_number = len(start_end_index_list)
        print('there are %d instances (sliding windows) in this dataset\n'%inst_number)
        np.savetxt(sliding_file_path,start_end_index_list,delimiter=',',fmt='%d')
    else:
        print('Loading start_end_index_list from file')
        start_end_index_list = pd.read_csv(sliding_file_path, header=None).values
        inst_number = len(start_end_index_list)
        print('there are %d instances (sliding windows) in this dataset' % inst_number)

    # get all the log indexes in each time window by ranging from start_index to end_index
    expanded_indexes_list=[]
    for t in range(inst_number):
        index_list = []
        expanded_indexes_list.append(index_list)
    for i in range(inst_number):
        start_index = start_end_index_list[i][0]
        end_index = start_end_index_list[i][1]
        for l in range(start_index, end_index):
            expanded_indexes_list[i].append(l)

    event_mapping_data = [row[0] for row in event_mapping_data]
    event_num = len(list(set(event_mapping_data)))
    print('There are %d log events'%event_num)

    #=============get labels and event count of each sliding window =========#
    labels = []
    event_count_matrix = np.zeros((inst_number,event_num))
    for j in range(inst_number):
        label = 0   #0 represent success, 1 represent failure
        for k in expanded_indexes_list[j]:
            event_index = event_mapping_data[k]
            event_count_matrix[j, event_index] += 1
            if label_data[k]:
                label = 1
                continue
        labels.append(label)
    assert inst_number == len(labels)
    print("Among all instances, %d are anomalies"%sum(labels))
    assert event_count_matrix.shape[0] == len(labels)
    return event_count_matrix, labels
