"""
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
from sklearn.utils import shuffle
from collections import OrderedDict
from sklearn.feature_extraction.text import CountVectorizer
from Model.tools.WordVector import *
from ast import literal_eval
def func(x):
    if len(x)<1:
        return
    else: return list(map(int, x.split()))

def extract(train_file,train_abnormal_file=None,test_file=None,test_abnormal_file=None,need_abnormal=0,abnormal_num=500):

    if train_file.endswith('.csv'):
        train_df=pd.read_csv(train_file,engine="c")
        test_df=pd.read_csv(test_file,engine="c")
        x_train=[]
        #x_train=np.array(train_df["Sequence"].map(lambda x:list(map(int,x.split()))))#把sequence分割成数字放回sequence，再列表化后赋值给x_train
        for i,row in enumerate(train_df["Sequence"]):
            print("row:",row)
            x_train.append(list(map(int,row.split())))
        x_train=np.array(x_train)
        y_train = np.array(train_df["label"].map(lambda x: int(x)).tolist())
        print(test_df["Sequence"])
        #x_test= np.array(test_df["Sequence"].map(lambda x: list(map(int, x.split()))))  # 把sequence分割成数字放回sequence，再列表化后赋值给x_train
        x_test=[]
        for i,row in enumerate(test_df["Sequence"]):
            print("row",i,":",row)
            x_test.append(list(map(int,row.split())))
        x_test=np.array(x_test)
        y_test = np.array(test_df["label"].map(lambda x: int(x)).tolist())
        print("x_train:",x_train.shape)
        print("y_train:", y_train.shape)
        print("x_test:", x_test.shape)
        print("y_test:", y_test.shape)
        return (x_train, y_train), (x_test, y_test)
    else:
        x_train=[]
        y_train=None
        x_test=[]
        y_test=[]

        with open(train_file,"r") as train_f:
            with open(test_file,"r") as test_f:
                with open(test_abnormal_file,"r") as test_abnormal_f:
                    for line in train_f.readlines():
                        x_train.append(np.array(list(map(int,line.split()))))

                    for line in test_f.readlines():
                        x_test.append(np.array(list(map(int,line.split()))))
                        y_test.append(0)
                    for line in test_abnormal_f.readlines():
                        x_test.append(np.array(list(map(int,line.split()))))
                        y_test.append(1)
        x_train=np.array(x_train)

        if train_abnormal_file!=None:
            y_train=[]
            with open(train_abnormal_file,"r") as wf:
                for line in wf.readlines():
                    y_train.append(np.array(list(map(int,line.split()))))

        x_test=np.array(x_test)
        y_test=np.array(y_test)
        if need_abnormal==1:
            x_train=np.concatenate((x_train,x_test[y_test>0][:abnormal_num]),axis=0)
        print("x_train:", x_train.shape)
        #print("y_train:", y_train)
        print("x_test:", x_test.shape)
        print("y_test:", y_test.shape)
        return (x_train, y_train), (x_test, y_test)

def _split_data(x_data, y_data=None, train_ratio=0, split_type='uniform'):
    if split_type == 'uniform' and y_data is not None:
        pos_idx = y_data > 0
        print(pos_idx,end="\n")
        x_pos = x_data[pos_idx] #取出label为正的，x_pos
        print('x_pos')
        print(x_pos[0])
        y_pos = y_data[pos_idx]#取出label为正的，y_pos
        print(~pos_idx,end="\n")
        x_neg = x_data[~pos_idx]#取出label为负的（0）,x_neg
        print('x_neg')
        print(x_neg[0])
        y_neg = y_data[~pos_idx]#取出label为负的（0）,y_neg
        train_pos = int(train_ratio * x_pos.shape[0])
        train_neg = int(train_ratio * x_neg.shape[0])
       # print("test")
        #print(np.hstack([x_pos[0:1], x_neg[0:1]]))
        x_train = np.hstack([x_pos[0:train_pos], x_neg[0:train_neg]])#水平方向叠加
        y_train = np.hstack([y_pos[0:train_pos], y_neg[0:train_neg]])
        x_test = np.hstack([x_pos[train_pos:], x_neg[train_neg:]])
        y_test = np.hstack([y_pos[train_pos:], y_neg[train_neg:]])
        print("train_pos,train_neg")
        print(train_pos,end="\n")
        print(train_neg)
        print("x_train")
        print(x_train.shape[0])
    elif split_type == 'sequential':
        num_train = int(train_ratio * x_data.shape[0])
        x_train = x_data[0:num_train]
        x_test = x_data[num_train:]
        if y_data is None:
            y_train = None
            y_test = None
        else:
            y_train = y_data[0:num_train]
            y_test = y_data[num_train:]
    # Random shuffle
    indexes = shuffle(np.arange(x_train.shape[0]))
    x_train = x_train[indexes]#x,y特征重新随机排列
    if y_train is not None:
        y_train = y_train[indexes]
        print("x_test:")
        print(x_test)
    return (x_train, y_train), (x_test, y_test)

def load_HDFS(log_file, label_file=None, window='session', parameter_feature=1,train_ratio=0.7, split_type='sequential', save_csv=True, window_size=0):
    """ 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, None for unlabeled data
        window: str, the window options including `session` (default).
        train_ratio: float, the ratio of training data for train/test split.
        split_type: `uniform` or `sequential`, which determines how to split dataset. `uniform` means
            to split positive samples and negative samples equally when setting label_file. `sequential`
            means to split the data sequentially without label_file. That is, the first part is for training,
            while the second part is for testing.

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

    print('====== Input data summary ======')

    if log_file.endswith('.npz'):
        # Split training and validation set in a class-uniform way
        data = np.load(log_file,allow_pickle=True)
        x_data = data['x_data']
        print('x_data:\n',x_data)
        y_data = data['y_data']

        (x_train, y_train), (x_test, y_test) = _split_data(x_data, y_data, train_ratio, split_type)

    elif log_file.endswith('.csv'):
        assert window == 'session', "Only window=session is supported for HDFS dataset."
        print("Loading", log_file)
        struct_log = pd.read_csv(log_file, engine='c',
                na_filter=False, memory_map=True,error_bad_lines=False)
        print("load over.start handling to Event Sequence.")
        data_dict = OrderedDict()#有序hashmap字典-
        for idx, row in struct_log.iterrows():  #行列循环  按行迭代，返还行的索引以及行本身的内容
            blkId_list = re.findall(r'(blk_-?\d+)', row['Content'])#匹配blk_或blk_-，数字1次或以上"
            blkId_set = set(blkId_list)#构建无序无重复blk集合"
            for blk_Id in blkId_set:          #为每个blkId匹配EventID，可能一个blk有多个EventId"
                if not blk_Id in data_dict:  #如果data_dict中没有这个blk，加上"
                    data_dict[blk_Id] = []
                data_dict[blk_Id].append(row['EventId'])#统计每个blk的事件
        data_df = pd.DataFrame(list(data_dict.items()), columns=['BlockId', 'EventSequence'])#将data_dict用items转化为数组后 矩阵化"
        print("data_df:")
        print(data_df)

        if label_file:
            # Split training and validation set in a class-uniform way
            label_data = pd.read_csv(label_file, engine='c', na_filter=False, memory_map=True)
            label_data = label_data.set_index('BlockId')#BlockId设置为索引"
            label_dict = label_data['Label'].to_dict()
            data_df['Label'] = data_df['BlockId'].apply(lambda x: 1 if label_dict[x] == 'Anomaly' else 0)#根据BlockID查找label字典，给矩阵添加label列"

            # Split train and test data
            #print(data_df['EventSequence'].values)
            (x_train, y_train), (x_test, y_test) = _split_data(data_df['EventSequence'].values,
                data_df['Label'].values, train_ratio, split_type)
            print("x_train",x_train)

            print(y_train.sum(), y_test.sum())

        if save_csv:
            data_df.to_csv('../data/HDFS/data_instances.csv', index=False)
            #X_train_fit()
        if window_size > 0:
            x_train, window_y_train, y_train = slice_hdfs(x_train, y_train, window_size)
            x_test, window_y_test, y_test = slice_hdfs(x_test, y_test, window_size)
            log = "{} {} windows ({}/{} anomaly), {}/{} normal"
            print(log.format("Train:", x_train.shape[0], y_train.sum(), y_train.shape[0], (1-y_train).sum(), y_train.shape[0]))
            print(log.format("Test:", x_test.shape[0], y_test.sum(), y_test.shape[0], (1-y_test).sum(), y_test.shape[0]))
            return (x_train, window_y_train, y_train), (x_test, window_y_test, y_test)

        if label_file is None:
            if split_type == 'uniform':
                split_type = 'sequential'
                print('Warning: Only split_type=sequential is supported \
                if label_file=None.'.format(split_type))
            # Split training and validation set sequentially
            x_data = data_df['EventSequence'].values
            (x_train, _), (x_test, _) = _split_data(x_data, train_ratio=train_ratio, split_type=split_type)
            print('Total: {} instances, train: {} instances, test: {} instances'.format(
                  x_data.shape[0], x_train.shape[0], x_test.shape[0]))
            return (x_train, None), (x_test, None), data_df
    else:
        raise NotImplementedError('load_HDFS() only support csv and npz files!')

    num_train = x_train.shape[0]
    num_test = x_test.shape[0]
    num_total = num_train + num_test
    num_train_pos = sum(y_train)
    num_test_pos = sum(y_test)
    num_pos = num_train_pos + num_test_pos

    print('Total: {} instances, {} anomaly, {} normal' \
          .format(num_total, num_pos, num_total - num_pos))
    print('Train: {} instances, {} anomaly, {} normal' \
          .format(num_train, num_train_pos, num_train - num_train_pos))
    print('Test: {} instances, {} anomaly, {} normal\n' \
          .format(num_test, num_test_pos, num_test - num_test_pos))

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

def slice_hdfs(x, y, window_size):#将X基于blk_id的序列每行进行window_size大小的窗口滑动，不足window_size大小的部分补#pad.
    results_data = []
    print("Slicing {} sessions, with window {}".format(x.shape[0], window_size))
    for idx, sequence in enumerate(x):
        seqlen = len(sequence)
        i = 0
        while (i + window_size) < seqlen:
            slice = sequence[i: i + window_size]
            results_data.append([idx, slice, sequence[i + window_size], y[idx]])
            i += 1
        else:
            slice = sequence[i: i + window_size]
            slice += ["#Pad"] * (window_size - len(slice))
            results_data.append([idx, slice, "#Pad", y[idx]])

    results_df = pd.DataFrame(results_data, columns=["SessionId", "EventSequence", "Label", "SessionLabel"])
    results_df.to_csv("日志序列滑动窗口.txt")
    print("Slicing done, {} windows generated".format(results_df.shape[0]))
    return results_df[["SessionId", "EventSequence"]], results_df["Label"], results_df["SessionLabel"]

#固定窗口大小，步长根据索引，提取特征
def load_BGL_feature(raw_data,event_data,window_size=40 ,stepping_size=10,train_ratio=0.8):
    """  TODO

    """
    label_data, time_data = raw_data[:, 0], raw_data[:, 1]
    start_index = 0
    end_index = 0

#提取bgl数据特征，窗口大小不固定，根据时间戳
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]
    print("log_size:",log_size)
    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
   # print("raw_data:",raw_data)
    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))
                #print("pair",start_end_pair)
                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
            if j-i<5:
                #print("j-i ",j-i)
                continue
            else:
                start_end_pair = tuple((start_index, end_index))
                start_end_index_list.append(start_end_pair)
        #inst_number为一共的窗口数，也是样本数
        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):
            #print("l:",l)
            expanded_indexes_list[i].append(l)

#    event_mapping_data = event_mapping_data.to_numpy().tolist()
    #print("event_mapping_data",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 = []
    #*****事件计数矩阵 并给出事件序列data_instance*******
    data_instance=pd.DataFrame(columns=["EventSequence","Label"])#事件序列矩阵

    event_count_matrix = np.zeros((inst_number,event_num+1))
    print("event_mapping",len(event_mapping_data))
    print("expend_list:",len(expanded_indexes_list))
    pd.DataFrame(expanded_indexes_list).to_csv("../data/BGL/expend_list",columns=None)
    "generating data_instances and CountMatrix"
    for j in tqdm(range(inst_number)):
        #print(j)
        label = 0   #0 represent success, 1 represent failure
        EventSequence=[]
        #print("len:", len(expanded_indexes_list[j]))
        #print(expanded_indexes_list[inst_number-1])
        for k in expanded_indexes_list[j]:
            #print(k)
            event_index = event_mapping_data[k]
            #print("event_index",event_index)
            EventSequence.append(event_index)
            event_count_matrix[j, event_index] += 1
            if label_data[k]:
                label = 1
                continue
        data_instance.loc[j]=[EventSequence,label]
        #print("EventSequence:",EventSequence)

        labels.append(label)
    #print("data_instances:",data_instance)
    data_instance.to_csv("../data/BGL/data_instances.csv",index=False)
    assert inst_number == len(labels)
    print("Among all instances, %d are anomalies"%sum(labels))
    assert event_count_matrix.shape[0] == len(labels)
    print("event_count:",event_count_matrix.shape)
    print("labels:",len(labels))
    return event_count_matrix, labels
#将hdfs以块block为单位，块整体都在训练集里，按data_instance进行划分
def X_train_extract(data_dir="../data/HDFS",block_num=100000,train_ratio=0.5,normal_num=5000,normal_abnormal_ratio=1,file="../data/HDFS/data_instances.csv",template="../data/HDFS/HDFS.log_templates.csv",data_type='HDFS'):
    if data_type == 'HDFS':
        Event_seq = pd.read_csv(file, engine='c',
                                na_filter=False, memory_map=True)
        #根据日志分析的事件模板排序template记录Events种类和序号，用Events序号代替解析后的日志模板ID
        Events_list=pd.read_csv(template, engine='c',
                                na_filter=False, memory_map=True)
        Events = Events_list["EventId"].tolist()
        Events.insert(0,"#")
        #Events.append("#")
        length = block_num
        max_length=len(Event_seq)
        if length>len(Event_seq):
            length=len(Event_seq)
            print("max blocks:",length)
        mid = int(length * train_ratio)
        # Event_seq将原事件ID按mapping替换成数字
        # 数据前一半的正常序列作为训练，后一半的正常+异常作测试
        for i in range(max_length):
            strs = Event_seq.at[i, "EventSequence"].replace('[', '').replace(']', '').replace('\'', '').replace(',',
                                                                                ' ').split()
            # print("strs:",strs)
            #用template生成的Events序号代替解析后的日志模板ID
            for j, ch in enumerate(strs):
                if ch in Events:
                    strs[j] = str(Events.index(ch))
                else:
                    print("unknown event!")
            Event_seq.at[i, "EventSequence"] = ' '.join(strs)
        print("事件总数:", len(Events))
        print("事件ID顺序", Events)

        Event_seq.to_csv("../data/HDFS/hdfs.csv")
        #前length条随机重排！！！
        Event_last=Event_seq[length:]
        pd.DataFrame(Event_last[['EventSequence', 'Label']]).to_csv("../data/HDFS/MLog_log_last.csv", index=None)
        print("last length:",len(Event_last))
        Event_seq=Event_seq[:length]
        Event_seq = Event_seq.sample(frac=1).reset_index(drop=True)
        Event_seq=Event_seq.reset_index(drop=True)
        Event_seq.to_csv("../data/HDFS/hdfs_shuffled.csv")

        df = Event_seq[:block_num][['EventSequence', 'Label']]
        normal_df = pd.DataFrame(pd.DataFrame(df).loc[df['Label'] == 0])
        abnormal_df = pd.DataFrame(pd.DataFrame(df).loc[df['Label'] == 1])

        #去除重复的
        normal_df1=normal_df[:normal_num]
        normal_df2=normal_df.drop(normal_df1.index)

        test_df=pd.DataFrame(pd.concat([pd.DataFrame(normal_df2),pd.DataFrame(abnormal_df)]))

        print("无监督的HDFS：")
        print("train_sequence:", len(normal_df1), " train_normal_df:", len(normal_df1), " train_abnormal_df:",
              0)
        print("test_sequence:", len(test_df), " test_normal_df:", len(normal_df2), " test_abnormal_df:",
              len(abnormal_df))

        with open("../data/HDFS/hdfs_train", 'w') as wf:
            for i,row in normal_df1.iterrows():
                    wf.write(row["EventSequence"] + "\n")

        with open("../data/HDFS/hdfs_test_normal", 'w') as wf1:
            with open("../data/HDFS/hdfs_test_abnormal", 'w') as wf2:

                  for i,row in normal_df2.iterrows():
                        wf1.write(row["EventSequence"] + "\n")
                  for i,row in abnormal_df.iterrows():
                        wf2.write(row["EventSequence"] + "\n")
        with open("../data/HDFS/Events_mapping", 'w') as wf3:
            wf3.write(str(Events))
        print("keys:", Event_seq.columns)

        abnormal_num=int(normal_num*(1.0/normal_abnormal_ratio))
        df = Event_seq[:][['EventSequence', 'Label']]
        normal_df= pd.DataFrame(df).loc[df['Label'] == 0]
        abnormal_df = pd.DataFrame(df).loc[df['Label'] == 1]
        abnormal_df1=abnormal_df[:abnormal_num]
        normal_df1=normal_df[:normal_num]
        abnormal_df2=abnormal_df[abnormal_num:]
        normal_df2 = normal_df[abnormal_num * normal_abnormal_ratio:]
        train_df = pd.concat([pd.DataFrame(normal_df1), pd.DataFrame(abnormal_df1)])
        predict_df=pd.concat([pd.DataFrame(normal_df2), pd.DataFrame(abnormal_df2)])
        print("有监督的HDFS：")
        print("train_sequence:", len(train_df), " train_normal_df:", len(normal_df1), " train_abnormal_df:",len(abnormal_df1))
        print("test_sequence:", len(predict_df), " test_normal_df:", len(normal_df2)," test_abnormal_df:", len(abnormal_df2))

        train_df.columns=["Sequence", 'label']
        predict_df.columns=["Sequence", 'label']
        pd.DataFrame(train_df).to_csv("../data/HDFS/MLog_log_train.csv", index=None)
        pd.DataFrame(train_df).to_csv("../data/HDFS/MLog_log_valid.csv", index=None)
        pd.DataFrame(predict_df).to_csv("../data/HDFS/MLog_log_test.csv", index=None)
        """
        df1 = Event_seq[0:mid][['EventSequence', 'Label']]
        df1.columns = ["Sequence", 'label']
        normal_df = pd.DataFrame(df1).loc[df1['label'] == 0]

        abnormal_df = pd.DataFrame(df1).loc[df1['label'] == 1]
        abnormal_len = len(abnormal_df)

        #normal_df = normal_df.sample(n=int(abnormal_len * (normal_ratio/(1-normal_ratio))), random_state=1)
        train_df = pd.concat([pd.DataFrame(normal_df), pd.DataFrame(abnormal_df)])
        print("train_sequence:",len(train_df)," train_normal_df:", len(normal_df)," train_abnormal_df:", len(abnormal_df))

        pd.DataFrame(train_df).to_csv("../data/HDFS/MLog_log_train.csv", index=None)
        pd.DataFrame(train_df).to_csv("../data/HDFS/MLog_log_valid.csv", index=None)
        df2 = Event_seq[mid:length][['EventSequence', 'Label']]
        df2.columns = ["Sequence", 'label']
        print("test_sequence:", len(df2), " test_normal_df:", len(pd.DataFrame(df2).loc[df2['label'] == 0]), " test_abnormal_df:",len(pd.DataFrame(df2).loc[df2['label'] == 1]))
        pd.DataFrame(df2).to_csv("../data/HDFS/MLog_log_test.csv", index=None)"""
        # pd.DataFrame(df2).to_csv("../data/HDFS/MLog_log_valid.csv", index=None)
        return Events

    elif data_type == 'BGL'or data_type=='Thunderbird':
        """
        if data_type=='BGL':
            file="../data/BGL/data_instances.csv"
        elif data_type=='Thunderbird':
            file="../data/Thunderbird/data_instances.csv"
        """

        Event_seq = pd.read_csv(file, engine='c',
                                na_filter=False, memory_map=True)
        # 根据日志分析的事件模板排序template记录Events种类和序号，用Events序号代替解析后的日志模板ID
        Events_list = pd.read_csv(template, engine='c',
                                  na_filter=False, memory_map=True)
        Events = Events_list["EventId"].tolist()
        Events.insert(0, "#")
        #print("mapping2:",Events)
        # Events.append("#")
        length = block_num
        max_length = len(Event_seq)
        if length > len(Event_seq):
            length = len(Event_seq)
            print("max blocks:", length)
        mid = int(length * train_ratio)
        # Event_seq将原事件ID按mapping替换成数字
        # 数据前一半的正常序列作为训练，后一半的正常+异常作测试
        for i in range(max_length):
            strs = Event_seq.at[i, "EventSequence"].replace('[', '').replace(']', '').replace('\'', '').replace(',', ' ').split()
            # print("strs:",strs)
            """
            # 用template生成的Events序号代替解析后的日志模板ID
            for j, ch in enumerate(strs):
                if ch in Events:
                    strs[j] = str(Events.index(ch))
                else:
                    print("unknown event!")
                    """
            Event_seq.at[i, "EventSequence"] = ' '.join(strs)

        print("事件总数:", len(Events))
        print("事件ID顺序", Events)

        Event_seq.to_csv("../data/"+data_type+"/"+data_type+".csv")
        Event_last = Event_seq[length:]
        pd.DataFrame(Event_last[['EventSequence', 'Label']]).to_csv("../data/"+data_type+"/MLog_log_last.csv", index=None)
        print("last length:", len(Event_last))
        Event_seq = Event_seq[:length]
        # 前length条随机重排！！！
        Event_seq = Event_seq.sample(frac=1).reset_index(drop=True)
        Event_seq = Event_seq.reset_index(drop=True)
        Event_seq.to_csv("../data/"+data_type+"/"+data_type+"_shuffled.csv")

        df = Event_seq[:][['EventSequence', 'Label']]
        total_num=len(df)

        normal_df = pd.DataFrame(pd.DataFrame(df).loc[df['Label'] == 0])
        abnormal_df = pd.DataFrame(pd.DataFrame(df).loc[df['Label'] == 1])

        # 去除重复的
        normal_df1 = normal_df[:int(len(normal_df)*train_ratio)]
        normal_df2 = normal_df.drop(normal_df1.index)
        abnormal_df1=abnormal_df[:int(len(abnormal_df)*train_ratio)]
        abnormal_df2=abnormal_df.drop(abnormal_df1.index)

        test_df = pd.DataFrame(pd.concat([pd.DataFrame(normal_df2), pd.DataFrame(abnormal_df2)]))

        print("无监督的"+data_type+"：")
        print("train_sequence:", len(normal_df1)+len(abnormal_df1), " train_normal_df:", len(normal_df1), " train_abnormal_df:",
              len(abnormal_df1))
        print("test_sequence:", len(test_df), " test_normal_df:", len(normal_df2), " test_abnormal_df:",
              len(abnormal_df2))

        with open("../data/"+data_type+"/"+data_type+"_train", 'w') as wf:
            for i, row in normal_df1.iterrows():
                wf.write(row["EventSequence"] + "\n")
        with open("../data/"+data_type+"/"+data_type+"_train_abnormal", 'w') as wf:
            for i, row in abnormal_df1.iterrows():
                print(row["EventSequence"])
                wf.write(row["EventSequence"] + "\n")
        with open("../data/"+data_type+"/"+data_type+"_test_normal", 'w') as wf1:
            with open("../data/"+data_type+"/"+data_type+"_test_abnormal", 'w') as wf2:

                for i, row in normal_df2.iterrows():
                    wf1.write(row["EventSequence"] + "\n")
                for i, row in abnormal_df2.iterrows():
                    wf2.write(row["EventSequence"] + "\n")
        with open("../data/"+data_type+"/Events_mapping", 'w') as wf3:
            wf3.write(str(Events))
        print("keys:", Event_seq.columns)

        abnormal_num = int(normal_num * (1.0 / normal_abnormal_ratio))
        df = Event_seq[:][['EventSequence', 'Label']]
        normal_df = pd.DataFrame(df).loc[df['Label'] == 0]
        abnormal_df = pd.DataFrame(df).loc[df['Label'] == 1]
        abnormal_df1 = abnormal_df[:abnormal_num]
        normal_df1 = normal_df[:normal_num]
        abnormal_df2 = abnormal_df[abnormal_num:]
        normal_df2 = normal_df[abnormal_num * normal_abnormal_ratio:]
        train_df = pd.concat([pd.DataFrame(normal_df1), pd.DataFrame(abnormal_df1)])
        predict_df = pd.concat([pd.DataFrame(normal_df2), pd.DataFrame(abnormal_df2)])
        print("有监督的"+data_type+"：")
        print("train_sequence:", len(train_df), " train_normal_df:", len(normal_df1), " train_abnormal_df:",
              len(abnormal_df1))
        print("test_sequence:", len(predict_df), " test_normal_df:", len(normal_df2), " test_abnormal_df:",
              len(abnormal_df2))

        train_df.columns = ["Sequence", 'label']
        predict_df.columns = ["Sequence", 'label']
        pd.DataFrame(train_df).to_csv("../data/"+data_type+"/MLog_log_train.csv", index=None)
        pd.DataFrame(train_df).to_csv("../data/"+data_type+"/MLog_log_valid.csv", index=None)
        pd.DataFrame(predict_df).to_csv("../data/"+data_type+"/MLog_log_test.csv", index=None)

        return Events

#将初步事件序列data_instance 按需要比例划分整理到 hdfs_train、hdfs_test,robust_train,bgl_train等

def BGL_data(file="../data/BGL/1000000_BGL.log_structured.csv",template="../data/BGL/1000000_BGL.log_templates.csv"):
    read_file=pd.read_csv(file)
    mapping=getMapping(template)
    mapping.insert(0, "#")
    #替换标签为0 1 ,EventID替换为数字
    read_file.Label=read_file.Label.map(lambda x: 0 if x=='-' else 1)
    read_file.EventId=read_file.EventId.map(lambda x: mapping.index(x))

    data_out=pd.concat([read_file["Label"],read_file["Timestamp"]],axis=1)
    print("data_out:",data_out)
    Events_mapping=read_file["EventId"].values.tolist()
    #print("Events:",Events_mapping)
    return data_out.values,Events_mapping

def get_BGL_template(template="../data/BGL/1000000_BGL.log_templates.csv"):
    read_file = pd.read_csv(template)
    templates=dict()
    for i,row in read_file.iterrows():
        templates[i]=row["EventTemplate"]
    return templates

def getMapping(file):
    read_file = pd.read_csv(file)
    Events_mapping=read_file["EventId"].tolist()
    return  Events_mapping

if __name__=='__main__':
    #HDFS:
    """
    (x_train, y_train), (x_test, y_test) = load_HDFS("../data/HDFS/HDFS.log_structured.csv",
                                                                label_file="../data/HDFS/anomaly_label.csv",
                                                                window='session',
                                                                train_ratio=0.7,
                                                                split_type='uniform',
                                                                save_csv=True)#加载数据，提取每个blkID的事件集，保存至/HDFS/data_instance
    """

    #Events_mapping=X_train_extract(block_num=100000,train_ratio=0.5,normal_num=6000,normal_abnormal_ratio=3)#与X_train_fit不同，该函数仅仅作用于hdfs,用于完整提取一个blockid所有日
    #Events_mapping = X_train_extract(block_num=1000000, train_ratio=0.5, normal_num=6000,normal_abnormal_ratio=1)  # 与X_train_fit不同，该函数仅仅作用于hdfs,用于完整提取一个blockid所有日

    Events_Bert_vectors("../data/HDFS/HDFS.log_templates.csv",tf_idf=0,IDF=1,scalermin=0.7,scalermax=1,events_semantic="../data/HDFS/events_semantic.json")

    #BGL:
    """
    para = dict()
    para["save_path"] = "../data/BGL/"
    para["step_size"] = 0.0001
    para["window_size"] = 0.0001

    #Events_mapping指事件ID序列
    BGL_data, Events_mapping = BGL_data(file="../data/BGL/BGL.log_structured.csv",template="../data/BGL/BGL.log_templates.csv")
    print("len",BGL_data.size)
    print("Events_mapping:",len(Events_mapping))
    Events_count, labels = bgl_preprocess_data(para, BGL_data, Events_mapping)
    print("mapping",Events_mapping)
    """
    #X_train_extract(block_num=100000,data_type='BGL',train_ratio=0.5,file="../data/BGL/data_instances.csv",template="../data/BGL/BGL.log_templates.csv",normal_num=5000,normal_abnormal_ratio=1)

    ##事件ID映射根据 templates文件

    #templates=Detect_Events("../data/BGL/BGL.log_templates.csv")
    #Events_vectors(templates,events_semantic="../data/BGL/events_semantic.json",tf_idf=1,template="../data/BGL/BGL.log_templates.csv")
    #Events_Bert_vectors("../data/BGL/BGL.log_templates.csv",tf_idf=0,IDF=1,scalermin=0.7 ,scalermax=1,events_semantic="../data/BGL/events_semantic.json")


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Dec 22 22:56:58 2021

@author: stepianno
"""
from itertools import combinations_with_replacement
from collections import Counter
p1 = p2 = 0
s1 = 2
s2 = 1
die = 0
first = True
while p1 < 1000 and p2 < 1000:
    die += 3
    move = (die*3-3)%10
    if first:
        s1 = (s1+move)%10
        if not s1:
            p1 += 10
        else:
            p1 += s1
    else:
        s2 = (s2+move)%10
        if not s2:
            p2 += 10
        else:
            p2 += s2
    first = not first
print('part a:', min(p1,p2)*die)

ways = combinations_with_replacement([1,2,3], 3)
paths = Counter([sum(way) for way in ways])
start = (2,1,True,0,0)
dim_cache = {}
def dirac(s1, s2, first=True, p=0):
    turn = (s1,s2,first,p)
    if turn in dim_cache:
        return dim_cache[turn]
    if p.real >= 21:
        return 1
    elif p.imag >= 21:
        return 1j
    w = 0
    for a in range(1,4):
        for b in range(1,4):
            for c in range(1,4):
                k = a+b+c
                if first:
                    s1n = (s1 + k)%10
                    if not s1n:
                        pn = p + 10
                    else:
                        pn = p + s1n
                    w += dirac(s1n,s2,False,pn)
                else:
                    s2n = (s2 + k)%10
                    if not s2n:
                        pn = p + 10j
                    else:
                        pn = p + (s2n*1j)
                    w += dirac(s1,s2n,True,pn)
    dim_cache[turn] = w
    return dim_cache[turn]

dirac(2,1)
p1 = int(dim_cache[(2,1,True,0)].real)
p2 = int(dim_cache[(2,1,True,0)].imag)
print('part b:', max(p1,p2))


import turtle as t

f = open("C:/haedal/5/turtle.txt", 'r')
lines = f.readlines()
t.shape('turtle')

for line in lines:
    val = float(line)
    i = lines.index(line)

    if i%2==0 :
        t.forward(val)
    else :
        t.right(val)
f.close()


import numpy as np
from scipy.interpolate import lagrange

# every calculation is based in mod
mod = 3017
inverseModArr = [1, 6, 4, 3, 9, 2, 8, 7, 5, 10]
#note:
# y = pow(3, -1, 11)
# y is the inverse of 3 mod 11

# functions---------------------------------

# plug in an x to calculate polynomials in terms of mod n in vector
def computePolyVec(polyVec, x, size):
    rarr = [0 for i in range (size)]
    for i, val in enumerate(polyVec):
        rarr[i] = polyVec[i](x) % mod
    return rarr

# input x value to batch index of queries in terms of mod n
def computeBatchInput(batch, x):
    rarr = [[0 for col in range(14)] for row in range(4)]
    for i, val in enumerate(batch):
        for j, val2 in enumerate(batch[i]):
            rarr[i][j] = batch[i][j](x) % mod
    return rarr

# converts vector to mod n
def convertMod(resArr):
    rarr = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for i, val in enumerate(resArr):
        rarr[i] = resArr[i] % mod
    return rarr

# converts result vector to mod n
def convertMod2(resArr):
    rarr = [0, 0, 0, 0, 0, 0, 0]
    for i, val in enumerate(resArr):
        rarr[i] = resArr[i] % mod
    return rarr

# user functions:

# 1 degree lagrange interpolating polynomial
def create1DPoly(sbv1, sbv2, xes):
    res=[0, 0, 0, 0]
    for i in range(4):
        mul = [0, 0]
        mul[0] = pow((xes[0] - xes[1])%mod, -1, mod)
        mul[1] = pow((xes[1] - xes[0])%mod, -1, mod)
        tempPoly1 = [(sbv1[i] * mul[0])%mod, (-1*(sbv1[i] *((xes[1] * mul[0])%mod))%mod)%mod]
        tempPoly2 = [(sbv2[i] * mul[1])%mod, (-1*(sbv2[i] *((xes[0] * mul[1])%mod))%mod)%mod]
        res[i] = np.poly1d([(tempPoly1[0] + tempPoly2[0])%mod, (tempPoly1[1] + tempPoly2[1])%mod])
    return res

# 2 degree lagrange interpolating polynomial
# this is to create batch index of queries
def create2DPoly(ioq1, ioq2, ioq3, xCoor):
    res = [[0 for col in range(14)] for row in range(4)]
    for j in range(0,4):
        for i in range(0,14):
            mul=[0, 0, 0]
            mul[0] = pow((((xCoor[0]-xCoor[1])%mod)*((xCoor[0]-xCoor[2])%mod))% mod, -1, mod)
            mul[1] = pow((((xCoor[1]-xCoor[0])%mod)*((xCoor[1]-xCoor[2])%mod))% mod, -1, mod)
            mul[2] = pow((((xCoor[2]-xCoor[0])%mod)*((xCoor[2]-xCoor[1])%mod))% mod, -1, mod)
            tempPoly1 = [(ioq1[j][i]*mul[0])% mod, (ioq1[j][i] * ((mul[0] * ((-1 * (xCoor[2]+xCoor[1]))%mod))%mod))%mod, (ioq1[j][i] * ((mul[0] * ((xCoor[1]*xCoor[2])%mod))%mod))%mod]
            tempPoly2 = [(ioq2[j][i]*mul[1])% mod, (ioq2[j][i] * ((mul[1] * ((-1 * (xCoor[2]+xCoor[0]))%mod))%mod))%mod, (ioq2[j][i] * ((mul[1] * ((xCoor[0]*xCoor[2])%mod))%mod))%mod]
            tempPoly3 = [(ioq3[j][i]*mul[2])% mod, (ioq3[j][i] * ((mul[2] * ((-1 * (xCoor[1]+xCoor[0]))%mod))%mod))%mod, (ioq3[j][i] * ((mul[2] * ((xCoor[0]*xCoor[1])%mod))%mod))%mod]
            res[j][i] = np.poly1d([(((tempPoly1[0]+tempPoly2[0])%mod) + tempPoly3[0])%mod, (((tempPoly1[1]+tempPoly2[1])%mod) + tempPoly3[1])%mod, (((tempPoly1[2]+tempPoly2[2])%mod) + tempPoly3[2])%mod])
    return res

# 3 degree lagrange interpolating polynomial
def create3DPoly(s1, s2, s3, s4, x):
    res=[0, 0, 0, 0, 0, 0, 0]
    for i in range(7):
        mul = [0, 0, 0, 0]
        mul[0] = pow(((((x[0] - x[1])%mod * ((x[0] - x[2])%mod))%mod) * ((x[0] - x[3])%mod))% mod, -1, mod)
        mul[1] = pow(((((x[1] - x[0])%mod * ((x[1] - x[2])%mod))%mod) * ((x[1] - x[3])%mod))% mod, -1, mod)
        mul[2] = pow(((((x[2] - x[0])%mod * ((x[2] - x[1])%mod))%mod) * ((x[2] - x[3])%mod))% mod, -1, mod)
        mul[3] = pow(((((x[3] - x[0])%mod * ((x[3] - x[1])%mod))%mod) * ((x[3] - x[2])%mod))% mod, -1, mod)
        tempPoly1 = [(s1[i] * mul[0])%mod, (s1[i] * ((mul[0] * (((-1 * x[3])-x[2]-x[1])%mod))%mod))%mod, (s1[i] * ((mul[0] * ((((x[2]*x[3])%mod + (x[1]*x[3])%mod + (x[1]*x[2])%mod)%mod))%mod)))%mod, ((-1*(s1[i] * ((mul[0] * (x[1]*x[2]*x[3])%mod))%mod))%mod)%mod]
        tempPoly2 = [(s2[i] * mul[1])%mod, (s2[i] * ((mul[1] * (((-1 * x[3])-x[2]-x[0])%mod))%mod))%mod, (s2[i] * ((mul[1] * ((((x[2]*x[3])%mod + (x[0]*x[3])%mod + (x[0]*x[2])%mod)%mod))%mod)))%mod, ((-1*(s2[i] * ((mul[1] * (x[0]*x[2]*x[3])%mod))%mod))%mod)%mod]
        tempPoly3 = [(s3[i] * mul[2])%mod, (s3[i] * ((mul[2] * (((-1 * x[3])-x[1]-x[0])%mod))%mod))%mod, (s3[i] * ((mul[2] * ((((x[1]*x[3])%mod + (x[0]*x[3])%mod + (x[0]*x[1])%mod)%mod))%mod)))%mod, ((-1*(s3[i] * ((mul[2] * (x[0]*x[1]*x[3])%mod))%mod))%mod)%mod]
        tempPoly4 = [(s4[i] * mul[3])%mod, (s4[i] * ((mul[3] * (((-1 * x[2])-x[1]-x[0])%mod))%mod))%mod, (s4[i] * ((mul[3] * ((((x[1]*x[2])%mod + (x[0]*x[2])%mod + (x[0]*x[1])%mod)%mod))%mod)))%mod, ((-1*(s4[i] * ((mul[3] * (x[0]*x[1]*x[2])%mod))%mod))%mod)%mod]
        res[i] = np.poly1d([(tempPoly1[0]+tempPoly2[0]+tempPoly3[0]+tempPoly4[0])%mod, (tempPoly1[1]+tempPoly2[1]+tempPoly3[1]+tempPoly4[1])%mod, (tempPoly1[2]+tempPoly2[2]+tempPoly3[2]+tempPoly4[2])%mod, (tempPoly1[3]+tempPoly2[3]+tempPoly3[3]+tempPoly4[3])%mod])
    return res

# end functions-----------------------------

# What each server contains-----------------
# database
sData = [[2010, 830, 126, 446, 285, 669, 1],
        [2009, 109, 876, 234, 294, 345, 2],
        [2010, 928, 135, 458, 288, 572, 3],
        [2013, 311, 281, 391, 286, 570, 4],
        [2007, 116,	732, 445, 972, 110, 3],
        [2009, 105, 111, 448, 287, 522, 6],
        [2010, 830, 123, 450, 245, 563, 2],
        [2002, 311, 281, 494, 288, 511, 5],
        [2009, 600, 332, 20, 222, 883, 1],
        [2007, 178, 778, 655, 100, 901, 6],
        [2007, 65, 853, 733, 888, 300, 1],
        [2009, 9, 654, 554, 123, 44, 3],
        [2010, 732, 771, 555, 432, 12, 6],
        [2007, 900, 900, 221, 555, 309, 2]]

# index of queries
covid = [[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0],
        [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
        [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0]]

hypertension = [[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0],
        [0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]]

diabetes = [[0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0],
        [0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]]

# index of queries labels
ioq_labels = [0, 1, 2]

# batch index of queries for covid, hypertension, and diabetes respecitely
batchIndServer = create2DPoly(covid, hypertension, diabetes, ioq_labels)

# end of server contains--------------------

# Start of protocol-------------------------
# printing database
print("Database:")
for i in range(0,14):
    print(sData[i])
print("\n")

# userside
# Alice wants the top 2 years of covid admissions and hypertension admissions of yolo count
# covid
#alice = [0, 1, 0, 0]
alice = [334, 821, 105, 234]
# hypertension
alice2 = [0, 1, 0, 0]

# labels for vectors above
xCoor = [0,1]
print("Alice's Vectors:")
print(alice)
print(xCoor[0])
print(alice2)
print(xCoor[1])
print("\n")

# generate polynomial vector out of alice's vector and random point vector
# this is where you use approx lagrange mod
polyVec = create1DPoly(alice, alice2, xCoor)

print("Alice's generate polynomials:")
print(polyVec)
print("\n")

# plug x=3 and x=4 into Alice's polynomial vector to send to the 2 servers
aliceQ1 = computePolyVec(polyVec, 4, 4)
aliceQ2 = computePolyVec(polyVec, 5, 4)
aliceQ3 = computePolyVec(polyVec, 6, 4)
aliceQ4 = computePolyVec(polyVec, 7, 4)

print("Query1 sent to Server 1 with x = 4")
print(aliceQ1)
print("Query2 sent to Server 2 with x = 5")
print(aliceQ2)
print("Query3 sent to Server 3 with x = 6")
print(aliceQ3)
print("Query4 sent to Server 4 with x = 7")
print(aliceQ3)
print("\n")

# severside
print("Covid index of queries:")
for i in range(0,4):
    print(covid[i])
print("\n")

print("Hypertension index of queries:")
for i in range(0,4):
    print(hypertension[i])
print("\n")

print("Diabetes index of queries:")
for i in range(0,4):
    print(diabetes[i])
print("\n")

print("Batch Index of queries:")
for i in range(0,4):
    print(batchIndServer[i])
print("\n")

# server plugs in x to batch index of queries
# server 1
s1_idxq = computeBatchInput(batchIndServer, 4)
# server 2
s2_idxq = computeBatchInput(batchIndServer, 5)
# server 3
s3_idxq = computeBatchInput(batchIndServer, 6)
# server 3
s4_idxq = computeBatchInput(batchIndServer, 7)

print("Server 1 batch index of queries x = 4")
for i in range(0,4):
    print(s1_idxq[i])
print("\n")
print("Server 2 batch index of quereis x = 5")
for i in range(0,4):
    print(s2_idxq[i])
print("\n")
print("Server 3 batch index of quereis x = 6")
for i in range(0,4):
    print(s3_idxq[i])
print("\n")
print("Server 4 batch index of quereis x = 7")
for i in range(0,4):
    print(s4_idxq[i])
print("\n")
print("\n")

# Matrix multiply Alice's respective server query and the batch index of queries calculated then convert to mod n
# server 1
s1_Q1 = np.dot( aliceQ1, s1_idxq)
s2_Q2 = np.dot( aliceQ2, s2_idxq)
s3_Q3 = np.dot( aliceQ3, s3_idxq)
s4_Q4 = np.dot( aliceQ4, s4_idxq)
s1_Q1 = convertMod(s1_Q1)
s2_Q2 = convertMod(s2_Q2)
s3_Q3 = convertMod(s3_Q3)
s4_Q4 = convertMod(s4_Q4)

print("Matrix Multiply of Alice's Query 1 and server1")
print(s1_Q1)
print("Matrix Multiply of Alice's Query 2 and server2")
print(s2_Q2)
print("Matrix Multiply of Alice's Query 3 and server3")
print(s3_Q3)
print("Matrix Multiply of Alice's Query 4 and server4")
print(s4_Q4)
print("\n")

# resulting query(x)*indexofqueries(x) matrix multiply to database and converted to mod n
# server1
s1_Q1_data = np.dot(s1_Q1, sData)
# server2
s2_Q2_data = np.dot(s2_Q2, sData)
# server3
s3_Q3_data = np.dot(s3_Q3, sData)
# server3
s4_Q4_data = np.dot(s4_Q4, sData)

s1_Q1_data = convertMod2(s1_Q1_data)
s2_Q2_data = convertMod2(s2_Q2_data)
s3_Q3_data = convertMod2(s3_Q3_data)
s4_Q4_data = convertMod2(s4_Q4_data)

print("Matrix Multiply of Alice's Query 1, server1, and database")
print(s1_Q1_data)
print("Matrix Multiply of Alice's Query 2, server2, and database")
print(s2_Q2_data)
print("Matrix Multiply of Alice's Query 3, server3, and database")
print(s3_Q3_data)
print("Matrix Multiply of Alice's Query 4, server4, and database")
print(s4_Q4_data)
print("\n")

# Alice reconstructs vectors to polynomial vectors
# this is where you use approximate lagrange mod
serverLab = [4, 5, 6, 7]
resPoly = create3DPoly(s1_Q1_data, s2_Q2_data, s3_Q3_data, s4_Q4_data, serverLab)
print("Reconstructed polynomials from server results:")
print(resPoly)
print("\n")

# Alice inputs x to resultant polynomial vectors to get data
data = computePolyVec(resPoly, 0, 7)
print("Final Result Query 1:")
print(data)
print("Actual Result Query 1:")
print(np.dot(np.dot(alice,covid), sData))

data2 = computePolyVec(resPoly, 1, 7)
print("Final Result: Query 2:")
print(data2)
print("Actual Result Query 2:")
print(np.dot(np.dot(alice2, hypertension), sData))
print("\n")

# Total number of 2 years of the highest admissions who have hypertension at Yolo County
print("Total number of 2 years of the highest admissions who have hypertension at Yolo County:")
print(data2[4])
print("\n")

# Total number of 2 years of the highest admissions who have covid at Yolo County
print("Total number of 2 years of the highest admissions who have covid at Yolo County:")
print(data[2])
print("\n")


# coding: utf-8

# 必要ModuleをImport

# In[1]:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from torch.utils.data import DataLoader
import math

import pickle
import tqdm
from collections import Counter

from torch.utils.data import Dataset
import random
import numpy as np

from utils import GELU, PositionwiseFeedForward, LayerNorm, SublayerConnection, LayerNorm

import matplotlib
get_ipython().magic('matplotlib inline')
import matplotlib.pyplot as plt
from ipywidgets import FloatProgress
from IPython.display import display, clear_output

# In[2]:

input_train_txt = './data/chatdata/splitted_1.txt'
input_valid_txt = './data/chatdata/splitted_2.txt'
processed_train_txt = './data/chatdata/train_X.txt'
processed_valid_txt = './data/chatdata/valid_X.txt'

# Next Sentence Predictionのために, 意味的に連続する文章をtab区切りで並べる前処理をデータセットに対して行います.

# In[3]:

# 偶数行の文章を奇数行の文章と接続するメソッド
def load_data(path):
    with open(path, encoding='utf-8') as f:
        even_rows = []
        odd_rows = []
        all_f = f.readlines()
        for row in all_f[2::2]:
            even_rows.append(row.strip().replace('\n', ''))
        for row in all_f[1::2]:
            odd_rows.append(row.strip().replace('\n', ''))
    min_rows_len = int(min(len(even_rows), len(odd_rows)))
    even_rows = even_rows[:min_rows_len]
    odd_rows = odd_rows[:min_rows_len]

    concat_rows = []
    for even_r, odd_r in zip(even_rows, odd_rows):
        concat_r = '\t'.join([even_r, odd_r])
        concat_rows.append(concat_r)
    return concat_rows

# In[4]:

train_data = load_data(input_train_txt)
valid_data = load_data(input_valid_txt)

# ランダムに並び替える
random.shuffle(train_data)
random.shuffle(valid_data)

# In[5]:

with open(processed_train_txt, 'w') as f:
    f.write('\n'.join(train_data))

# In[6]:

with open(processed_valid_txt, 'w') as f:
    f.write('\n'.join(valid_data))

# Attentionセルを定義する

# In[7]:

class Attention(nn.Module):
    """
    Scaled Dot Product Attention(縮小付き内積注意)
    (参考) http://deeplearning.hatenablog.com/entry/transformer

    """

    def forward(self, query, key, value, mask=None, dropout=None):
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1))

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        p_attn = F.softmax(scores, dim=-1)

        if dropout is not None:
            p_attn = dropout(p_attn)

        """
        Product Attentionとvalueの内積を取る(内積注意)
        """

        return torch.matmul(p_attn, value), p_attn

# Multi Head Attentionを定義する

# In[8]:

class MultiHeadedAttention(nn.Module):

    def __init__(self, h, d_model, dropout=0.1):
        super().__init__()
        assert d_model % h == 0

        """
        d_modelはモデル全体の次元数(largeだと512).head(=h)の数で等分される
        次元数はvalueとkeyで常に同数
        # We assume d_v always equals d_k
        """

        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h

        """
        nn.Linear(入力する特徴量のsize,出力〜)
        https://pytorch.org/docs/stable/nn.html#torch.nn.Linear
        """
        self.linear_layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(3)])
        self.output_linear = nn.Linear(d_model, d_model)

        """
        Attentionクラスの読み込み
        """
        self.attention = Attention()

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        """
        queryがバッチサイズになる。マスク無し。
        """
        batch_size = query.size(0)

        """
        l(x)の表記の仕方がよくわからない。
        .viewでtensorのサイズを調整。-1ですべての要素が横に並ぶ。
        """
        query, key, value = [l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linear_layers, (query, key, value))]

        x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)

        """
        contiguousはメモリの別々の部分に格納されている情報を統合するためにつかう。
        """
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)

        return self.output_linear(x)

# Transformerを定義する

# In[9]:

class TransformerBlock(nn.Module):
    """
    Bidirectional Encoder = Transformer (self-attention)
    Transformer = MultiHead_Attention + Feed_Forward with sublayer connection
    """

    def __init__(self, hidden, attn_heads, feed_forward_hidden, dropout):
        """
        :param hidden: hidden size of transformer
        :param attn_heads: head sizes of multi-head attention
        :param feed_forward_hidden: feed_forward_hidden, usually 4*hidden_size
        :param dropout: dropout rate
        """

        super().__init__()
        """
        MultiHeadedAttentionクラスを自己注意としてインスタンス化
        """
        self.attention = MultiHeadedAttention(h=attn_heads, d_model=hidden, dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(d_model=hidden, d_ff=feed_forward_hidden, dropout=dropout)
        self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout)
        self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, mask):
        """
        公開されていないメソッドやインスタンス変数にだけ、アンダースコアを先頭に付けてください。
        https://teratail.com/questions/41277
        query,key,valueへの3つの自己注意に対応しているのだろうか>_x, _x, _x,
        """
        x = self.input_sublayer(x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask))
        x = self.output_sublayer(x, self.feed_forward)
        return self.dropout(x)

# BERTクラスを定義する

# In[10]:

"""
googleのスクリプトの、BertPretrainModel(nn.module)にあたる
"""

class BERT(nn.Module):

    def __init__(self, vocab_size, hidden=768, n_layers=12, attn_heads=12, dropout=0.1):
        """
        :param vocab_size: vocab_size of total words
        :param hidden: BERT model hidden size
        :param n_layers: numbers of Transformer blocks(layers)
        :param attn_heads: number of attention heads
        :param dropout: dropout rate
        """

        super().__init__()
        self.hidden = hidden
        self.n_layers = n_layers
        self.attn_heads = attn_heads

        """
        通常feed forwardの隠れ層は4倍
        """
        self.feed_forward_hidden = hidden * 4

        # embedding for BERT
        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=hidden, dropout=dropout)

        """
        ・レイヤーをリストで保持してしまうと、、リストで保持しているレイヤーのパラメータは
        パラメータとして認識されず呼び出されない。そういうときはModulelistを使う。
        https://qiita.com/perrying/items/857df46bb6cdc3047bd8
        TransformerBlockクラスをModulelistを使ってインスタンス化
        """
        self.transformer_blocks = nn.ModuleList([TransformerBlock(hidden, attn_heads, hidden * 4, dropout) for _ in range(n_layers)])

    def forward(self, x, segment_info):
        # xの中で0以上は1, 0未満は0として, maskテンソルを作る
        """
        dim (int) – the index at which to insert the singleton dimension
        https://pytorch.org/docs/stable/torch.html#torch.unsqueeze
        （予想）x>0のとき、値が１のシングルトンの次元をまず生成。
        次に1度だけ、最初に作った値が１のシングルトンの次元の数だけ横に並べる。
        もういちどそれを値１のシングルトンの次元を生成。
        （よくわからない）
        """

        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)

        x = self.embedding(x, segment_info)

        for transformer in self.transformer_blocks:
            x = transformer.forward(x, mask)
        return x

    """
    1.できればapexのインストール
    2.FusedLayerNormをBertLayerFormとしてimport
    3.from apex.normaliztion.used_layer_norm import FusedLayerNorm as BertLayerNorm
    4.(497)BerPreTrainModel以下にある、init_bert_weights メソッドを、BERTクラス(ココ)に追加
    """
    #from apex.normalization.fused_layer_norm import FusedLayerNorm as BertLayerNorm

    def init_bert_weights(self,module):
      #initialize the weights.

      if isinstance(module,(nn.Linear,nn.Embedding)):
          #cf https://github.com/pytorch/pytorch/pull/5617
          """
          (148) initializer_range=0.02):
          module.weight.data.normal_(mean=0.0, std=0.02に→self.config.initializer_range)
          """

          module.weight.data.normal_(mean=0.0, std=0.02)

       #elif isinstance(module, BertLayerNorm)
      elif isinstance(module, LayerNorm):
          module.bias.data.zero_()
          module.weight.data.fill_(1.0)
      if isinstance(module, nn.Linear)and module.bias is not None:
          module.bias.data.zero_()

# In[11]:

"""
ここに、
Class BertForQuestionAnswering(BERT):
を入れるのがいいかもしれない。
"""

class BertForQuestionAnswering(BERT):

   #def __init__(selfトル→,config):
   def __init__(self):
     #super(BertForQuestionAnswering, self).__init__(トル→config)
     super(BertForQuestionAnswering, self).__init__()
     #self.bert = BERT→BertModel(トル→config)
     self.bert = BERT()
     #self.qa_outputs = nn.Linear(hidden→config.hidden_size, 2)
     self.qa_outputs = nn.Linear(hidden, 2)
     self.apply(self.init_bert_weights)

   def forward(self,input_ids,token_type_ids=None,attention_mask=None,start_positions=None,end_positions=None):
     sequence_output,_ = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
     logits = self.qa_outputs(sequence_output)
     start_logits, end_logits = logits.split(1, dim=-1)
     start_logits = start_logits.squeeze(-1)
     end_logits = end_logits.squeeze(-1)

     if start_positions is not None and end_positions is not None:
         #if we are on multi-GPU, split add a dimension
         if len(start_positions.size()) > 1:
             start_positions = start_positions.squeeze(-1)
         if len(end_positions.size()) > 1:
             end_positions = end_positions.squeeze(-1)
         #sometimes the start/end posisions are outside our model inputs, we ignore these terms
         ignored_index = start_logits.size(1)
         start_positions.clamp_(0, ignored_index)
         end_positions.clamp_(0,ignored_index)

         loss_fct = CrossEntropyLoss(ignore_index = ignored_index)
         start_loss = loss_fct(start_logits, start_positions)
         end_loss = loss_fct(end_logits, end_positions)
         total = (start_loss + end_loss) / 2
         return total_loss

     else:
         return start_logits, end_logits

# BERTのEmbedding層を定義する

# In[12]:

class TokenEmbedding(nn.Embedding):
    def __init__(self, vocab_size, embed_size=512):
        super().__init__(vocab_size, embed_size, padding_idx=0)

class PositionalEmbedding(nn.Module):

    def __init__(self, d_model, max_len=512):
        super().__init__()

        pe = torch.zeros(max_len, d_model).float()
        pe.require_grad = False

        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)).float().exp()

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return self.pe[:, :x.size(1)]

class SegmentEmbedding(nn.Embedding):
    def __init__(self, embed_size=512):
        super().__init__(3, embed_size, padding_idx=0)

class BERTEmbedding(nn.Module):
    """
    BERT Embedding which is consisted with under features
        1. TokenEmbedding : 通常のEMbedding
        2. PositionalEmbedding : sin, cosを用いた位置情報付きEmbedding
        2. SegmentEmbedding : Sentenceのセグメント情報 (sent_A:1, sent_B:2)
    """
    def __init__(self, vocab_size, embed_size, dropout=0.1):
        super().__init__()
        self.token = TokenEmbedding(vocab_size=vocab_size, embed_size=embed_size)
        self.position = PositionalEmbedding(d_model=self.token.embedding_dim)
        self.segment = SegmentEmbedding(embed_size=self.token.embedding_dim)
        self.dropout = nn.Dropout(p=dropout)
        self.embed_size = embed_size

    def forward(self, sequence, segment_label):
        x = self.token(sequence) + self.position(sequence) + self.segment(segment_label)
        return self.dropout(x)

# 学習用にマスク予測・隣接文予測の層を追加する

# In[13]:

class BERTLM(nn.Module):
    """
    BERT Language Model
    Next Sentence Prediction Model + Masked Language Model
    """

    def __init__(self, bert: BERT, vocab_size):
        """
        :param bert: BERT model which should be trained
        :param vocab_size: total vocab size for masked_lm
        """

        super().__init__()
        self.bert = bert
        self.next_sentence = NextSentencePrediction(self.bert.hidden)
        self.mask_lm = MaskedLanguageModel(self.bert.hidden, vocab_size)

    def forward(self, x, segment_label):
        x = self.bert(x, segment_label)
        return self.next_sentence(x), self.mask_lm(x)

class NextSentencePrediction(nn.Module):
    """
    2クラス分類問題 : is_next, is_not_next
    """

    def __init__(self, hidden):
        """
        :param hidden: BERT model output size
        """
        super().__init__()
        self.linear = nn.Linear(hidden, 2)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, x):
        return self.softmax(self.linear(x[:, 0]))

class MaskedLanguageModel(nn.Module):
    """
    入力系列のMASKトークンから元の単語を予測する
    nクラス分類問題, nクラス : vocab_size
    """

    def __init__(self, hidden, vocab_size):
        """
        :param hidden: output size of BERT model
        :param vocab_size: total vocab size
        """
        super().__init__()
        self.linear = nn.Linear(hidden, vocab_size)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, x):
        return self.softmax(self.linear(x))

# BERT用のVocabを生成するクラスを定義する

# In[14]:

import pickle
import tqdm
from collections import Counter

class TorchVocab(object):
    """
    :property freqs: collections.Counter, コーパス中の単語の出現頻度を保持するオブジェクト
    :property stoi: collections.defaultdict, string → id の対応を示す辞書
    :property itos: collections.defaultdict, id → string の対応を示す辞書
    """
    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],
                 vectors=None, unk_init=None, vectors_cache=None):
        """
        :param coutenr: collections.Counter, データ中に含まれる単語の頻度を計測するためのcounter
        :param max_size: int, vocabularyの最大のサイズ. Noneの場合は最大値なし. defaultはNone
        :param min_freq: int, vocabulary中の単語の最低出現頻度. この数以下の出現回数の単語はvocabularyに加えられない.
        :param specials: list of str, vocabularyにあらかじめ登録するtoken
        :param vecors: list of vectors, 事前学習済みのベクトル. ex)Vocab.load_vectors
        """
        self.freqs = counter
        counter = counter.copy()
        min_freq = max(min_freq, 1)

        self.itos = list(specials)
        # special tokensの出現頻度はvocabulary作成の際にカウントされない
        for tok in specials:
            del counter[tok]

        max_size = None if max_size is None else max_size + len(self.itos)

        # まず頻度でソートし、次に文字順で並び替える
        words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])
        words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)

        # 出現頻度がmin_freq未満のものはvocabに加えない
        for word, freq in words_and_frequencies:
            if freq < min_freq or len(self.itos) == max_size:
                break
            self.itos.append(word)

        # dictのk,vをいれかえてstoiを作成する
        self.stoi = {tok: i for i, tok in enumerate(self.itos)}

        self.vectors = None
        if vectors is not None:
            self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache)
        else:
            assert unk_init is None and vectors_cache is None

    def __eq__(self, other):
        if self.freqs != other.freqs:
            return False
        if self.stoi != other.stoi:
            return False
        if self.itos != other.itos:
            return False
        if self.vectors != other.vectors:
            return False
        return True

    def __len__(self):
        return len(self.itos)

    def vocab_rerank(self):
        self.stoi = {word: i for i, word in enumerate(self.itos)}

    def extend(self, v, sort=False):
        words = sorted(v.itos) if sort else v.itos
        for w in words:
            if w not in self.stoi:
                self.itos.append(w)
                self.stoi[w] = len(self.itos) - 1

class Vocab(TorchVocab):
    def __init__(self, counter, max_size=None, min_freq=1):
        self.pad_index = 0
        self.unk_index = 1
        self.eos_index = 2
        self.sos_index = 3
        self.mask_index = 4
        super().__init__(counter, specials=["<pad>", "<unk>", "<eos>", "<sos>", "<mask>"], max_size=max_size, min_freq=min_freq)

    # override用
    def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:
        pass

    # override用
    def from_seq(self, seq, join=False, with_pad=False):
        pass

    @staticmethod
    def load_vocab(vocab_path: str) -> 'Vocab':
        with open(vocab_path, "rb") as f:
            return pickle.load(f)

    def save_vocab(self, vocab_path):
        with open(vocab_path, "wb") as f:
            pickle.dump(self, f)

# テキストファイルからvocabを作成する
class WordVocab(Vocab):
    def __init__(self, texts, max_size=None, min_freq=1):
        print("Building Vocab")
        counter = Counter()
        for line in texts:
            if isinstance(line, list):
                words = line
            else:
                words = line.replace("\n", "").replace("\t", "").split()

            for word in words:
                counter[word] += 1
        super().__init__(counter, max_size=max_size, min_freq=min_freq)

    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):
        if isinstance(sentence, str):
            sentence = sentence.split()

        seq = [self.stoi.get(word, self.unk_index) for word in sentence]

        if with_eos:
            seq += [self.eos_index]  # this would be index 1
        if with_sos:
            seq = [self.sos_index] + seq

        origin_seq_len = len(seq)

        if seq_len is None:
            pass
        elif len(seq) <= seq_len:
            seq += [self.pad_index for _ in range(seq_len - len(seq))]
        else:
            seq = seq[:seq_len]

        return (seq, origin_seq_len) if with_len else seq

    def from_seq(self, seq, join=False, with_pad=False):
        words = [self.itos[idx]
                 if idx < len(self.itos)
                 else "<%d>" % idx
                 for idx in seq
                 if not with_pad or idx != self.pad_index]

        return " ".join(words) if join else words

    @staticmethod
    def load_vocab(vocab_path: str) -> 'WordVocab':
        with open(vocab_path, "rb") as f:
            return pickle.load(f)

def build(corpus_path, output_path, vocab_size=None, encoding='utf-8', min_freq=1):
    with open(corpus_path, "r", encoding=encoding) as f:
        vocab = WordVocab(f, max_size=vocab_size, min_freq=min_freq)

    print("VOCAB SIZE:", len(vocab))
    vocab.save_vocab(output_path)

# Dataloaderを定義する.
# ここで文章中の単語をMASKする処理と,隣り合う文章を一定確率でシャッフルする処理を同時に行う

# In[15]:

class BERTDataset(Dataset):
    """
    このクラスで訓練モード化どうかを指定する項目を初期化
    torch.utils.data.Dataset （抽象クラス）
    __len__,__getitem__ ともにDatasetクラスでは必要なメソッド
    """
    def __init__(self, corpus_path, vocab, seq_len, label_path='None', encoding="utf-8", corpus_lines=None, is_train=True):
        self.vocab = vocab
        self.seq_len = seq_len
        self.is_train = is_train

        with open(corpus_path, "r", encoding=encoding) as f:
            self.datas = [line[:-1].split("\t") for line in f]
        if label_path:
            self.labels_data = torch.LongTensor(np.loadtxt(label_path))
        else:
            # ラベル不要の時はダミーデータを埋め込む
            self.labels_data = [0 for _ in range(len(self.datas))]

    def __len__(self):
        return len(self.datas)

    def __getitem__(self, item):
        t1, (t2, is_next_label) = self.datas[item][0], self.random_sent(item)
        t1_random, t1_label = self.random_word(t1)
        t2_random, t2_label = self.random_word(t2)
        labels = self.labels_data[item]

        # [CLS] tag = SOS tag, [SEP] tag = EOS tag
        t1 = [self.vocab.sos_index] + t1_random + [self.vocab.eos_index]
        t2 = t2_random + [self.vocab.eos_index]

        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]
        t2_label = t2_label + [self.vocab.pad_index]

        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]
        bert_input = (t1 + t2)[:self.seq_len]
        bert_label = (t1_label + t2_label)[:self.seq_len]

        padding = [self.vocab.pad_index for _ in range(self.seq_len - len(bert_input))]
        bert_input.extend(padding), bert_label.extend(padding), segment_label.extend(padding)

        output = {"bert_input": bert_input,
                  "bert_label": bert_label,
                  "segment_label": segment_label,
                  "is_next": is_next_label,
                  "labels": labels}

        return {key: torch.tensor(value) for key, value in output.items()}

    def random_word(self, sentence):
        tokens = sentence.split()
        output_label = []

        for i, token in enumerate(tokens):
            if self.is_train: # Trainingの時は確率的にMASKする
                prob = random.random()
            else:  # Predictionの時はMASKをしない
                prob = 1.0
            if prob < 0.15:
                prob /= 0.15

                # 80% randomly change token to mask token
                if prob < 0.8:
                    tokens[i] = self.vocab.mask_index

                # 10% randomly change token to random token
                elif prob < 0.9:
                    tokens[i] = random.randrange(len(self.vocab))

                # 10% randomly change token to current token
                else:
                    tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)

                output_label.append(self.vocab.stoi.get(token, self.vocab.unk_index))

            else:
                tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)
                output_label.append(0)

        return tokens, output_label

    def random_sent(self, index):
        # output_text, label(isNotNext:0, isNext:1)
        if random.random() > 0.5:
            return self.datas[index][1], 1
        else:
            return self.datas[random.randrange(len(self.datas))][1], 0

# Trainerクラスを定義する.
# BERTの事前学習ではふたつの言語モデル学習を行う.
# 1. Masked Language Model : 文章中の一部の単語をマスクして,予測を行うタスク.
# 2. Next Sentence prediction : ある文章の次に来る文章を予測するタスク.

# In[16]:

class BERTTrainer:
    #def __init__(self, bert: BERT, vocab_size: int,
    def __init__(self, bert: BertForQuestionAnswering, vocab_size: int,
                 train_dataloader: DataLoader, test_dataloader: DataLoader = None,
                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01,
                 with_cuda: bool = True, log_freq: int = 10):
        """
        :param bert: BERT model
        :param vocab_size: vocabに含まれるトータルの単語数
        :param train_dataloader: train dataset data loader
        :param test_dataloader: test dataset data loader [can be None]
        :param lr: 学習率
        :param betas: Adam optimizer betas
        :param weight_decay: Adam optimizer weight decay param
        :param with_cuda: traning with cuda
        :param log_freq: logを表示するiterationの頻度
        """

        # GPU環境において、GPUを指定しているかのフラグ
        cuda_condition = torch.cuda.is_available() and with_cuda
        self.device = torch.device("cuda:0" if cuda_condition else "cpu")

        self.bert = bert
        self.model = BERTLM(bert, vocab_size).to(self.device)

        if torch.cuda.device_count() > 1:
            print("Using %d GPUS for BERT" % torch.cuda.device_count())
            self.model = nn.DataParallel(self.model)

        self.train_data = train_dataloader
        self.test_data = test_dataloader

        self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)

        # masked_token予測のためのLoss関数を設定
        self.criterion = nn.NLLLoss()
        self.log_freq = log_freq
        print("Total Parameters:", sum([p.nelement() for p in self.model.parameters()]))

        self.train_lossses = []
        self.train_accs = []

    def train(self, epoch):
        self.iteration(epoch, self.train_data)

    def test(self, epoch):
        self.iteration(epoch, self.test_data, train=False)

    def iteration(self, epoch, data_loader, train=True):
        """
        :param epoch: 現在のepoch
        :param data_loader: torch.utils.data.DataLoader
        :param train: trainかtestかのbool値
        """
        str_code = "train" if train else "test"

        data_iter = tqdm.tqdm(enumerate(data_loader), desc="EP_%s:%d" % (str_code, epoch), total=len(data_loader), bar_format="{l_bar}{r_bar}")

        avg_loss = 0.0
        total_correct = 0
        total_element = 0

        for i, data in data_iter:
            # 0. batch_dataはGPU or CPUに載せる
            data = {key: value.to(self.device) for key, value in data.items()}

            # 1. forward the next_sentence_prediction and masked_lm model
            next_sent_output, mask_lm_output = self.model.forward(data["bert_input"], data["segment_label"])

            # 2-1. NLLLoss(negative log likelihood) : next_sentence_predictionのLoss
            next_loss = self.criterion(next_sent_output, data["is_next"])

            # 2-2. NLLLoss(negative log likelihood) : predicting masked token word
            mask_loss = self.criterion(mask_lm_output.transpose(1, 2), data["bert_label"])

            # 2-3. next_lossとmask_lossの合計をlossとする
            loss = next_loss + mask_loss

            # 3. training時のみ,backwardとoptimizer更新を行う
            if train:
                self.optim.zero_grad()
                loss.backward()
                self.optim.step()

            # next sentence prediction accuracy
            correct = next_sent_output.argmax(dim=-1).eq(data["is_next"]).sum().item()
            avg_loss += loss.item()
            total_correct += correct
            total_element += data["is_next"].nelement()

            post_fix = {
                "epoch": epoch,
                "iter": i,
                "avg_loss": avg_loss / (i + 1),
                "avg_acc": total_correct / total_element * 100,
                "loss": loss.item()
            }

            if i % self.log_freq == 0:
                data_iter.write(str(post_fix))

        print("EP%d_%s, avg_loss=" % (epoch, str_code), avg_loss / len(data_iter), "total_acc=", total_correct * 100.0 / total_element)
        self.train_lossses.append(avg_loss / len(data_iter))
        self.train_accs.append(total_correct * 100.0 / total_element)

    def save(self, epoch, file_path="output/bert_trained.model"):
        """
        Saving the current BERT model on file_path

        :param epoch: current epoch number
        :param file_path: model output path which gonna be file_path+"ep%d" % epoch
        :return: final_output_path
        """
        output_path = file_path + ".ep%d" % epoch
        torch.save(self.bert.cpu(), output_path)
        self.bert.to(self.device)
        print("EP:%d Model Saved on:" % epoch, output_path)
        return output_path

    """
    loadをここに書く場合、こんな感じかもしれない

    def load(self, file_path="output/bert_trained.model"):
        pretrained_path = file_path
        torch.load(self.bert.cpu(), pretrained_path)
        self.bert.to(self.device)
        print("Model loaded on:" ,pretrained_path)
        return pretrained_path
    # ココまでで終わり
    """

# In[17]:

import datetime
dt_now = str(datetime.datetime.now()).replace(' ', '')

# In[18]:

# 訓練用パラメタを定義する
train_dataset=processed_train_txt
test_dataset=processed_valid_txt
vocab_path='./data/vocab'+ dt_now +'.txt'
output_model_path='./output/bertmodel'+ dt_now

hidden=256 #768
layers=8 #12
attn_heads=8 #12
seq_len=60

batch_size=4
epochs=10
num_workers=5
with_cuda=True
log_freq=20
corpus_lines=None

lr=1e-3
adam_weight_decay=0.00
adam_beta1=0.9
adam_beta2=0.999

dropout=0.0

min_freq=7

corpus_path=processed_train_txt
label_path=None

# In[19]:

build(corpus_path, vocab_path, min_freq=min_freq)

print("Loading Vocab", vocab_path)
vocab = WordVocab.load_vocab(vocab_path)

print("Loading Train Dataset", train_dataset)
train_dataset = BERTDataset(train_dataset, vocab, seq_len=seq_len, label_path=label_path, corpus_lines=corpus_lines)

print("Loading Test Dataset", test_dataset)
test_dataset = BERTDataset(test_dataset, vocab, seq_len=seq_len, label_path=label_path) if test_dataset is not None else None

print("Creating Dataloader")
train_data_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers)
test_data_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers) if test_dataset is not None else None

print("Building BERT model")
bert = BERT(len(vocab), hidden=hidden, n_layers=layers, attn_heads=attn_heads, dropout=dropout)

# In[20]:

print("Creating BERT Trainer")
trainer = BERTTrainer(bert, len(vocab), train_dataloader=train_data_loader, test_dataloader=test_data_loader,
                      lr=lr, betas=(adam_beta1, adam_beta2), weight_decay=adam_weight_decay,
                      with_cuda=with_cuda, log_freq=log_freq)

# In[25]:

print("Training Start")
for epoch in range(epochs):
    trainer.train(epoch)
    # Model Save
    trainer.save(epoch, output_model_path)
    trainer.test(epoch)

# In[73]:

"""
     (参考) QAの初期化
     def __init__(self):
      super(BertForQuestionAnswering, self).__init__()
      self.bert = BERT()
      self.qa_outputs = nn.Linear(hidden, 2)
      self.apply(self.init_bert_weights)
      #def prediction(self, "dummyQ")
"""

class QAprediction(BERTTrainer):

    #trainerの初期化
    def __init__(self, bert: BertForQuestionAnswering, vocab_size: int,
                 #train_dataloader: DataLoader, test_dataloader: DataLoader = None,
                 #train_dataloaderも使わないのでNoneに、逆にtest_dataloaderのひとつだけを使う
                 test_dataloader: DataLoader, train_dataloader: DataLoader = None,
                 #model.eval()にしてあるので関係なさそう
                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01,
                 with_cuda: bool = True, log_freq: int = 10):
        """
        :param bert: BERT model
        :param vocab_size: vocabに含まれるトータルの単語数
        :param train_dataloader: train dataset data loader
        :param test_dataloader: test dataset data loader [can be None]
        :param lr: 学習率
        :param betas: Adam optimizer betas
        :param weight_decay: Adam optimizer weight decay param
        :param with_cuda: traning with cuda
        :param log_freq: logを表示するiterationの頻度
        """

    #def test(self, epoch):
    #test_dataにindexを渡す
    #学習したモデルのロード
    def prediction(self, file_path="output/bert_trained.model"):
        """
        test_data のインデックスは何だろう？

        """
        self.test_data = test_dataloader
        # test_data (=dataloader)の中にあるデータのひとつだけを指定して推論してみる。下は１０番目。
        index = 10

        predict_data = self.test_data[index]
        pretrained_path = file_path
        torch.load(self.bert.cpu(), pretrained_path)

        model.eval()

        qa_output = model(predict_data)

        print("予測結果は{}".format(qa_output))

        #イテレーションはいらない。
        #self.iteration(epoch, self.test_data, train=False)


from numpy import *

'''
Superposition of two matrices using the algorithm
develpoed by Kabsch and Wolfgang

'''

class align_to_substrate:

    # Transform coordinates into numpy array
    # Requires coordinates in pdb format
    # Returns numpy array with coordinates
    def get_data(self,filename):
        # Need to add try/catch
        fname = open(filename,'r')
        # Coordinates
        dt = []
        # atom names
        a_n = []
        # Looping over line in file
        for line in fname:
            sl = line.split()
            # Bug fix 28-12-2009
            x = float(line[31:38])
            y = float(line[38:46])
            z = float(line[46:54])
            tmp = x,y,z
            dt.append(tmp)
            a_n.append(sl[2])
        fname.close()
        return array(dt),a_n

    # Writes the transformed coordinates to file
    # Requires coordinates translated and rotated
    # Returns xyz file of rotated and translated
    def write_data(self,atomnames,coordinates):
        al = len(atomnames)
        #    assert al == len(coordinates),'Lengths of lists do not match'
        wname = open('superimposed.xyz','w')
        wname.write(str(al)+'\n Generated to Enzyme Design\n')
        for i in range(al):
            wname.write(atomnames[i]+'\t'+str(coordinates[i][0])+'   '+str(coordinates[i][1])+'   '+str(coordinates[i][2])+'\n')

    def write_pdb(self,atomnames,coordinates,fname='superimposed.pdb'):
        '''Write a pdb file of the aligned coordinates '''
        al = len(atomnames)
        # Name of chain
        ch = '  SUB'
        wname = open(fname,'w')
        for i in range(al):
            wname.write('ATOM   '+str(i).rjust(3)+'   '+atomnames[i]+' '+ch+' X '+str(i).rjust(3)+'    '+str(coordinates[i][0])[0:6]+'  '+str(coordinates[i][1])[0:6]+'  '+str(coordinates[i][2])[0:6]+'\n')

    # Kabsch Algorithm for superimposing
    # Superimpose matrix1 on matrix2
    def get_rotate_translate(self,matrix1,matrix2):
        assert shape(matrix1) == shape(matrix2), 'Matrices not of same dimensions'

        # Store number of rows
        nrows = shape(matrix1)[0]

        # Getting centroid position for each selection
        avg_pos1 = matrix1.sum(axis=0)/nrows
        avg_pos2 = matrix2.sum(axis=0)/nrows

        # Translation of matrices
        avg_matrix1 = matrix1-avg_pos1
        avg_matrix2 = matrix2-avg_pos2

        # Covariance matrix
        covar = dot(avg_matrix1.T,avg_matrix2)

        # Do the SVD in order to get rotation matrix
        u,s,wt = linalg.svd(covar)

        # Rotation matrix
        # Transposition of u,wt
        rot_matrix = wt.T*u.T

        # Insure a right-handed coordinate system
        # need to fix this!!
        # if linalg.det(rot_matrix) > 0:
        wt[2] = -wt[2]
        rot_matrix = transpose(dot(transpose(wt),transpose(u)))

        trans_matrix = avg_pos2-dot(avg_pos1,rot_matrix)
        return trans_matrix, rot_matrix

    # Returning the superimposed coordinates
    # Requires two arrays same length
    # Returns array of transformed coordinates
    def get_transformed_coor(self,a,b):
        tr, rt = self.get_rotate_translate(a,b)
        nw_coor = dot(a,rt) + tr
        return nw_coor


binarytree = [None]*20
def root(data):
    if binarytree[0] != None:
        print("Tree already has a root")
    else:
        binarytree[0] = data

def set_left(data,parent):
    if binarytree[parent] == None:
        print("Cant set a child at",(parent*2)+1,", no parent found")
    else:
        binarytree[2*parent+1] = data

def set_right(data,parent):
    if binarytree[parent] == None:
        print("Cant set a child at",(parent*2)+2,", no parent found")
    else:
        binarytree[2*parent+2] = data

def print_binarytree():
    for i in range(20):
        if binarytree[i] != None:
            print(binarytree[i], end=" ")
        else:
            print("-",end=" ")
    print()

root(8)
set_left(3,0)
set_right(10,0)
set_left(1,1)
set_right(6,1)
set_right(14,2)
set_left(4,4)
set_right(7,4)
set_left(13,6)
print_binarytree()


import paddle
import paddle.nn as nn

from layers import BBB_Linear, BBB_Conv2d
from layers import BBB_LRT_Linear, BBB_LRT_Conv2d
from layers import FlattenLayer, ModuleWrapper

class BBBLeNet(ModuleWrapper):
    '''The architecture of LeNet with Bayesian Layers'''

    def __init__(self, outputs, inputs, priors, layer_type='lrt', activation_type='relu'):
        super(BBBLeNet, self).__init__()

        self.num_classes = outputs
        self.layer_type = layer_type
        self.priors = priors

        if layer_type == 'lrt':
            BBBLinear = BBB_LRT_Linear
            BBBConv2d = BBB_LRT_Conv2d
        elif layer_type == 'bbb':
            BBBLinear = BBB_Linear
            BBBConv2d = BBB_Conv2d
        else:
            raise ValueError("Undefined layer_type")

        if activation_type == 'softplus':
            self.act = nn.Softplus()
        elif activation_type == 'relu':
            self.act = nn.ReLU()
        elif activation_type == 'leakyrelu':
            self.act = nn.LeakyReLU()
        else:
            raise ValueError("Only softplus or relu supported")

        self.conv1 = BBBConv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2,priors=priors)
        self.act1 = self.act
        self.max_pool1 = nn.MaxPool2D(kernel_size=2,  stride=2)         # 16

        self.conv3 = BBBConv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1,priors=priors)#12
        self.act1 = self.act
        self.max_pool2 = nn.MaxPool2D(kernel_size=2, stride=2)   # 6

        self.flatten = FlattenLayer(6 * 6 * 16)

        self.linear1 = BBBLinear(in_features=16*6*6, out_features=128,priors=priors)
        # self.drop2 = nn.Dropout(0.25)
        self.act3 = self.act

        self.linear2 = BBBLinear(in_features=128, out_features=64,priors=priors)
        # self.drop2 = nn.Dropout(0.25)
        self.act4 = self.act

        self.linear3 = BBBLinear(in_features=64, out_features=10,priors=priors)

    # def forward(self,x):
    #     y = self.conv1(x)
    #     y = self.max_pool1(y)
    #     y = self.conv2(y)
    #     y = self.max_pool2(y)
    #     y = self.flatten(y)
    #     y = self.linear1(y)
    #     y = self.linear2(y)
    #     y = self.linear3(y)

    #     kl = 0.0
    #     for module in self.sublayers():
    #         if hasattr(module, 'kl_loss'):
    #             kl = kl + module.kl_loss()

    #     return y,kl


import random
import numpy as np
import math
import statistics

def schwef():
    repeticiones = 100
    iteraciones = 10000

    d = 10
    sigma = 0.1
    mejor = 10000
    rep = 10000

    # BUsqueda local iterada
    for i in range(repeticiones):
        # Busqueda local
        for j in range(iteraciones):
            # Vecindario
            x = []
            for i in range(d):
                x.append(random.uniform(-500, 500))

            suma = 0
            for k in range(d):
                xi = x[k]
                suma = suma + xi*math.sin(math.sqrt(abs(xi)))

            y = 418.9829 * d - suma
            s = np.random.normal(y, sigma)
            if s < rep:
                rep = s

        if rep < mejor:
            mejor = rep

    return mejor

if __name__ == "__main__":
    lista_mejor = []
    for i in range(30):
        resultado = schwef()
        print(resultado)
        lista_mejor.append(resultado)

    prom = statistics.mean(lista_mejor)
    st_dev = statistics.pstdev(lista_mejor)

    print("Promedio mejor: ", prom)
    print("Desviación estanadar mejor: ", st_dev)


# code was heavily based on https://github.com/hollance/BlazeFace-PyTorch
# This work is licensed under the same terms as MediaPipe (Apache License 2.0)
# https://github.com/google/mediapipe/blob/master/LICENSE

import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F

class BlazeBlock(nn.Layer):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
        super(BlazeBlock, self).__init__()

        self.stride = stride
        self.channel_pad = out_channels - in_channels

        if stride == 2:
            self.max_pool = nn.MaxPool2D(kernel_size=stride, stride=stride)
            padding = 0
        else:
            padding = (kernel_size - 1) // 2

        self.convs = nn.Sequential(
            nn.Conv2D(in_channels=in_channels,
                      out_channels=in_channels,
                      kernel_size=kernel_size,
                      stride=stride,
                      padding=padding,
                      groups=in_channels),
            nn.Conv2D(in_channels=in_channels,
                      out_channels=out_channels,
                      kernel_size=1,
                      stride=1,
                      padding=0),
        )

        self.act = nn.ReLU()

    def forward(self, x):
        if self.stride == 2:
            h = F.pad(x, [0, 2, 0, 2], "constant", 0)
            x = self.max_pool(x)
        else:
            h = x
        if self.channel_pad > 0:
            x = F.pad(x, [0, 0, 0, self.channel_pad, 0, 0, 0, 0], "constant", 0)

        return self.act(self.convs(h) + x)

class BlazeFace(nn.Layer):
    """The BlazeFace face detection model.
    """
    def __init__(self):
        super(BlazeFace, self).__init__()

        self.num_classes = 1
        self.num_anchors = 896
        self.num_coords = 16
        self.score_clipping_thresh = 100.0
        self.x_scale = 128.0
        self.y_scale = 128.0
        self.h_scale = 128.0
        self.w_scale = 128.0
        self.min_score_thresh = 0.75
        self.min_suppression_threshold = 0.3

        self._define_layers()

    def _define_layers(self):
        self.backbone1 = nn.Sequential(
            nn.Conv2D(in_channels=3,
                      out_channels=24,
                      kernel_size=5,
                      stride=2,
                      padding=0),
            nn.ReLU(),
            BlazeBlock(24, 24),
            BlazeBlock(24, 28),
            BlazeBlock(28, 32, stride=2),
            BlazeBlock(32, 36),
            BlazeBlock(36, 42),
            BlazeBlock(42, 48, stride=2),
            BlazeBlock(48, 56),
            BlazeBlock(56, 64),
            BlazeBlock(64, 72),
            BlazeBlock(72, 80),
            BlazeBlock(80, 88),
        )

        self.backbone2 = nn.Sequential(
            BlazeBlock(88, 96, stride=2),
            BlazeBlock(96, 96),
            BlazeBlock(96, 96),
            BlazeBlock(96, 96),
            BlazeBlock(96, 96),
        )

        self.classifier_8 = nn.Conv2D(88, 2, 1)
        self.classifier_16 = nn.Conv2D(96, 6, 1)

        self.regressor_8 = nn.Conv2D(88, 32, 1)
        self.regressor_16 = nn.Conv2D(96, 96, 1)

    def forward(self, x):
        x = F.pad(x, [1, 2, 1, 2], "constant", 0)

        b = x.shape[0]

        x = self.backbone1(x)  # (b, 88, 16, 16)
        h = self.backbone2(x)  # (b, 96, 8, 8)

        c1 = self.classifier_8(x)  # (b, 2, 16, 16)
        c1 = c1.transpose([0, 2, 3, 1])  # (b, 16, 16, 2)
        c1 = c1.reshape([b, -1, 1])  # (b, 512, 1)

        c2 = self.classifier_16(h)  # (b, 6, 8, 8)
        c2 = c2.transpose([0, 2, 3, 1])  # (b, 8, 8, 6)
        c2 = c2.reshape([b, -1, 1])  # (b, 384, 1)

        c = paddle.concat((c1, c2), axis=1)  # (b, 896, 1)

        r1 = self.regressor_8(x)  # (b, 32, 16, 16)
        r1 = r1.transpose([0, 2, 3, 1])  # (b, 16, 16, 32)
        r1 = r1.reshape([b, -1, 16])  # (b, 512, 16)

        r2 = self.regressor_16(h)  # (b, 96, 8, 8)
        r2 = r2.transpose([0, 2, 3, 1])  # (b, 8, 8, 96)
        r2 = r2.reshape([b, -1, 16])  # (b, 384, 16)

        r = paddle.concat((r1, r2), axis=1)  # (b, 896, 16)
        return [r, c]

    def load_weights(self, path):
        paddle.load_dict(paddle.load(path))
        self.eval()

    def load_anchors(self, path):
        self.anchors = paddle.to_tensor(np.load(path), dtype='float32')
        assert (self.anchors.shape == 2)
        assert (self.anchors.shape[0] == self.num_anchors)
        assert (self.anchors.shape[1] == 4)

    def load_anchors_from_npy(self, arr):
        self.anchors = paddle.to_tensor(arr, dtype='float32')
        assert (len(self.anchors.shape) == 2)
        assert (self.anchors.shape[0] == self.num_anchors)
        assert (self.anchors.shape[1] == 4)

    def _preprocess(self, x):
        """Converts the image pixels to the range [-1, 1]."""
        return x.astype('float32') / 127.5 - 1.0

    def predict_on_image(self, img):
        """Makes a prediction on a single image.

        Arguments:
            img: a NumPy array of shape (H, W, 3) or a Paddle tensor of
                 shape (3, H, W). The image's height and width should be
                 128 pixels.

        Returns:
            A tensor with face detections.
        """
        if isinstance(img, np.ndarray):
            img = paddle.to_tensor(img).transpose((2, 0, 1))

        return self.predict_on_batch(img.unsqueeze(0))[0]

    def predict_on_batch(self, x):
        """Makes a prediction on a batch of images.

        Arguments:
            x: a NumPy array of shape (b, H, W, 3) or a Paddle tensor of
               shape (b, 3, H, W). The height and width should be 128 pixels.

        Returns:
            A list containing a tensor of face detections for each image in
            the batch. If no faces are found for an image, returns a tensor
            of shape (0, 17).

        Each face detection is a Paddle tensor consisting of 17 numbers:
            - ymin, xmin, ymax, xmax
            - x,y-coordinates for the 6 keypoints
            - confidence score
        """
        if isinstance(x, np.ndarray):
            x = paddle.to_tensor(x).transpose((0, 3, 1, 2))

        assert x.shape[1] == 3
        assert x.shape[2] == 128
        assert x.shape[3] == 128

        x = self._preprocess(x)

        with paddle.no_grad():
            out = self.__call__(x)

        detections = self._tensors_to_detections(out[0], out[1], self.anchors)

        filtered_detections = []
        for i in range(len(detections)):
            faces = self._weighted_non_max_suppression(detections[i])
            faces = paddle.stack(faces) if len(faces) > 0 else paddle.zeros(
                (0, 17))
            filtered_detections.append(faces)

        return filtered_detections

    def _tensors_to_detections(self, raw_box_tensor, raw_score_tensor, anchors):
        """The output of the neural network is a tensor of shape (b, 896, 16)
        containing the bounding box regressor predictions, as well as a tensor
        of shape (b, 896, 1) with the classification confidences.

        Returns a list of (num_detections, 17) tensors, one for each image in
        the batch.
        """
        assert len(raw_box_tensor.shape) == 3
        assert raw_box_tensor.shape[1] == self.num_anchors
        assert raw_box_tensor.shape[2] == self.num_coords

        assert len(raw_score_tensor.shape) == 3
        assert raw_score_tensor.shape[1] == self.num_anchors
        assert raw_score_tensor.shape[2] == self.num_classes

        assert raw_box_tensor.shape[0] == raw_score_tensor.shape[0]

        detection_boxes = self._decode_boxes(raw_box_tensor, anchors)

        thresh = self.score_clipping_thresh
        raw_score_tensor = raw_score_tensor.clip(-thresh, thresh)
        detection_scores = F.sigmoid(raw_score_tensor).squeeze(axis=-1)

        mask = detection_scores >= self.min_score_thresh
        mask = mask.numpy()
        detection_boxes = detection_boxes.numpy()
        detection_scores = detection_scores.numpy()

        output_detections = []
        for i in range(raw_box_tensor.shape[0]):
            boxes = paddle.to_tensor(detection_boxes[i, mask[i]])
            scores = paddle.to_tensor(
                detection_scores[i, mask[i]]).unsqueeze(axis=-1)
            output_detections.append(paddle.concat((boxes, scores), axis=-1))

        return output_detections

    def _decode_boxes(self, raw_boxes, anchors):
        """Converts the predictions into actual coordinates using
        the anchor boxes. Processes the entire batch at once.
        """
        boxes = paddle.zeros_like(raw_boxes)

        x_center = raw_boxes[:,:, 0] / self.x_scale * \
            anchors[:, 2] + anchors[:, 0]
        y_center = raw_boxes[:,:, 1] / self.y_scale * \
            anchors[:, 3] + anchors[:, 1]

        w = raw_boxes[:, :, 2] / self.w_scale * anchors[:, 2]
        h = raw_boxes[:, :, 3] / self.h_scale * anchors[:, 3]

        boxes[:, :, 0] = y_center - h / 2.  # ymin
        boxes[:, :, 1] = x_center - w / 2.  # xmin
        boxes[:, :, 2] = y_center + h / 2.  # ymax
        boxes[:, :, 3] = x_center + w / 2.  # xmax

        for k in range(6):
            offset = 4 + k * 2
            keypoint_x = raw_boxes[:,:, offset] / \
                self.x_scale * anchors[:, 2] + anchors[:, 0]
            keypoint_y = raw_boxes[:,:, offset + 1] / \
                self.y_scale * anchors[:, 3] + anchors[:, 1]
            boxes[:, :, offset] = keypoint_x
            boxes[:, :, offset + 1] = keypoint_y

        return boxes

    def _weighted_non_max_suppression(self, detections):
        """The alternative NMS method as mentioned in the BlazeFace paper:
        The input detections should be a Tensor of shape (count, 17).
        Returns a list of Paddle tensors, one for each detected face.

        """
        if len(detections) == 0:
            return []

        output_detections = []

        # Sort the detections from highest to lowest score.
        remaining = paddle.argsort(detections[:, 16], descending=True).numpy()
        detections = detections.numpy()

        while len(remaining) > 0:
            detection = detections[remaining[0]]

            first_box = detection[:4]
            other_boxes = detections[remaining, :4]
            ious = overlap_similarity(paddle.to_tensor(first_box),
                                      paddle.to_tensor(other_boxes))

            mask = ious > self.min_suppression_threshold
            mask = mask.numpy()

            overlapping = remaining[mask]
            remaining = remaining[~mask]

            weighted_detection = detection.copy()
            if len(overlapping) > 1:
                coordinates = detections[overlapping, :16]
                scores = detections[overlapping, 16:17]
                total_score = scores.sum()
                weighted = (coordinates * scores).sum(axis=0) / total_score
                weighted_detection[:16] = weighted
                weighted_detection[16] = total_score / len(overlapping)

            output_detections.append(paddle.to_tensor(weighted_detection))

        return output_detections

def intersect(box_a, box_b):
    """Compute the area of intersect between box_a and box_b.
    Args:
      box_a: (tensor) bounding boxes, Shape: [A,4].
      box_b: (tensor) bounding boxes, Shape: [B,4].
    Return:
      (tensor) intersection area, Shape: [A,B].
    """
    A = box_a.shape[0]
    B = box_b.shape[0]
    max_xy = paddle.minimum(box_a[:, 2:].unsqueeze(1).expand((A, B, 2)),
                            box_b[:, 2:].unsqueeze(0).expand((A, B, 2)))
    min_xy = paddle.maximum(box_a[:, :2].unsqueeze(1).expand((A, B, 2)),
                            box_b[:, :2].unsqueeze(0).expand((A, B, 2)))
    inter = paddle.clip((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]

def jaccard(box_a, box_b):
    """Compute the jaccard overlap of two sets of boxes.
    Args:
        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
    Return:
        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
    """
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2] - box_a[:, 0]) *
              (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)
    area_b = ((box_b[:, 2] - box_b[:, 0]) *
              (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)
    union = area_a + area_b - inter
    return inter / union

def overlap_similarity(box, other_boxes):
    """Computes the IOU between a bounding box and set of other boxes."""
    return jaccard(box.unsqueeze(0), other_boxes).squeeze(0)

def init_model():
    net = BlazeFace()
    net.load_weights("blazeface.pdparams")
    net.load_anchors("anchors.npy")

    net.min_score_thresh = 0.75
    net.min_suppression_threshold = 0.3

    return net