# !/usr/bin/python
# -*- coding: utf-8 -*-
# @time    : 2019/11/2 21:08
# @author  : Mo
# @function:


from keras_textclassification.data_preprocess.text_preprocess import load_json, save_json
from keras_textclassification.conf.path_config import path_model_dir
path_fast_text_model_vocab2index = path_model_dir + 'vocab2index.json'
path_fast_text_model_l2i_i2l = path_model_dir + 'l2i_i2l.json'

import numpy as np
import os


class PreprocessGenerator:
    """
        数据预处理, 输入为csv格式, [label,ques]
    """
    def __init__(self):
        self.l2i_i2l = None
        if os.path.exists(path_fast_text_model_l2i_i2l):
            self.l2i_i2l = load_json(path_fast_text_model_l2i_i2l)

    def prereocess_idx(self, pred):
        if os.path.exists(path_fast_text_model_l2i_i2l):
            pred_i2l = {}
            i2l = self.l2i_i2l['i2l']
            for i in range(len(pred)):
                pred_i2l[i2l[str(i)]] = pred[i]
            pred_i2l_rank = [sorted(pred_i2l.items(), key=lambda k: k[1], reverse=True)]
            return pred_i2l_rank
        else:
            raise RuntimeError("path_fast_text_model_label2index is None")

    def prereocess_pred_xid(self, pred):
        if os.path.exists(path_fast_text_model_l2i_i2l):
            pred_l2i = {}
            l2i = self.l2i_i2l['l2i']
            for i in range(len(pred)):
                pred_l2i[pred[i]] = l2i[pred[i]]
            pred_l2i_rank = [sorted(pred_l2i.items(), key=lambda k: k[1], reverse=True)]
            return pred_l2i_rank
        else:
            raise RuntimeError("path_fast_text_model_label2index is None")

    def preprocess_get_label_set(self, path):
        # 首先获取label,set,即存在的具体类
        label_set = set()
        len_all = 0
        file_csv = open(path, "r", encoding="utf-8")
        for line in file_csv:
            len_all += 1
            if len_all > 1:  # 第一条是标签'label,ques'，不选择
                line_sp = line.split(",")
                label_org = str(line_sp[0]).strip().upper()
                label_real = "NAN" if label_org=="" else label_org
                label_set.add(label_real)
        file_csv.close()
        return label_set, len_all

    def preprocess_label_ques_to_idx(self, embedding_type, batch_size, path, embed, rate=1):
        label_set, len_all = self.preprocess_get_label_set(path)
        # 获取label转index字典等, 如果label2index存在则不转换了, dev验证集合的时候用
        if not os.path.exists(path_fast_text_model_l2i_i2l):
            count = 0
            label2index = {}
            index2label = {}
            for label_one in label_set:
                label2index[label_one] = count
                index2label[count] = label_one
                count = count + 1

            l2i_i2l = {}
            l2i_i2l['l2i'] = label2index
            l2i_i2l['i2l'] = index2label
            save_json(l2i_i2l, path_fast_text_model_l2i_i2l)
        else:
            l2i_i2l = load_json(path_fast_text_model_l2i_i2l)

        # 读取数据的比例
        len_ql = int(rate * len_all)
        if len_ql <= 500:  # sample时候不生效,使得语料足够训练
            len_ql = len_all

        def process_line(line):
            # 对每一条数据操作，获取label和问句index
            line_sp = line.split(",")
            ques = str(line_sp[1]).strip().upper()
            label = str(line_sp[0]).strip().upper()
            label = "NAN" if label == "" else label
            que_embed = embed.sentence2idx(ques)
            label_zeros = [0] * len(l2i_i2l['l2i'])
            label_zeros[l2i_i2l['l2i'][label]] = 1
            return que_embed, label_zeros

        while True:
            file_csv = open(path, "r", encoding="utf-8")
            cout_all_line = 0
            cnt = 0
            x, y = [], []
            # 跳出循环
            if len_ql < cout_all_line:
                break
            for line in file_csv:
                cout_all_line += 1
                if cout_all_line > 1: # 第一条是标签'label,ques'，不选择
                    x_line, y_line = process_line(line)
                    x.append(x_line)
                    y.append(y_line)
                    cnt += 1
                    if cnt == batch_size:
                        if embedding_type in ['bert', 'albert']:
                            x_, y_ = np.array(x), np.array(y)
                            x_1 = np.array([x[0] for x in x_])
                            x_2 = np.array([x[1] for x in x_])
                            x_all = [x_1, x_2]
                        elif embedding_type == 'xlnet':
                            x_, y_ = x, np.array(y)
                            x_1 = np.array([x[0][0] for x in x_])
                            x_2 = np.array([x[1][0] for x in x_])
                            x_3 = np.array([x[2][0] for x in x_])
                            x_all = [x_1, x_2, x_3]
                        else:
                            x_all, y_ = np.array(x), np.array(y)

                        cnt = 0
                        yield (x_all, y_)
                        x, y =[], []
        file_csv.close()
        print("preprocess_label_ques_to_idx ok")





[((23, 11, 23, 55), 'os.path.exists', 'os.path.exists', ({(23, 26, 23, 54): 'path_fast_text_model_l2i_i2l'}, {}), '(path_fast_text_model_l2i_i2l)', False, 'import os\n'), ((27, 11, 27, 55), 'os.path.exists', 'os.path.exists', ({(27, 26, 27, 54): 'path_fast_text_model_l2i_i2l'}, {}), '(path_fast_text_model_l2i_i2l)', False, 'import os\n'), ((38, 11, 38, 55), 'os.path.exists', 'os.path.exists', ({(38, 26, 38, 54): 'path_fast_text_model_l2i_i2l'}, {}), '(path_fast_text_model_l2i_i2l)', False, 'import os\n'), ((24, 27, 24, 66), 'keras_textclassification.data_preprocess.text_preprocess.load_json', 'load_json', ({(24, 37, 24, 65): 'path_fast_text_model_l2i_i2l'}, {}), '(path_fast_text_model_l2i_i2l)', False, 'from keras_textclassification.data_preprocess.text_preprocess import load_json, save_json\n'), ((66, 15, 66, 59), 'os.path.exists', 'os.path.exists', ({(66, 30, 66, 58): 'path_fast_text_model_l2i_i2l'}, {}), '(path_fast_text_model_l2i_i2l)', False, 'import os\n'), ((78, 12, 78, 60), 'keras_textclassification.data_preprocess.text_preprocess.save_json', 'save_json', ({(78, 22, 78, 29): 'l2i_i2l', (78, 31, 78, 59): 'path_fast_text_model_l2i_i2l'}, {}), '(l2i_i2l, path_fast_text_model_l2i_i2l)', False, 'from keras_textclassification.data_preprocess.text_preprocess import load_json, save_json\n'), ((80, 22, 80, 61), 'keras_textclassification.data_preprocess.text_preprocess.load_json', 'load_json', ({(80, 32, 80, 60): 'path_fast_text_model_l2i_i2l'}, {}), '(path_fast_text_model_l2i_i2l)', False, 'from keras_textclassification.data_preprocess.text_preprocess import load_json, save_json\n'), ((116, 34, 116, 62), 'numpy.array', 'np.array', ({(116, 43, 116, 61): '[x[0] for x in x_]'}, {}), '([x[0] for x in x_])', True, 'import numpy as np\n'), ((117, 34, 117, 62), 'numpy.array', 'np.array', ({(117, 43, 117, 61): '[x[1] for x in x_]'}, {}), '([x[1] for x in x_])', True, 'import numpy as np\n'), ((115, 37, 115, 48), 'numpy.array', 'np.array', ({(115, 46, 115, 47): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((115, 50, 115, 61), 'numpy.array', 'np.array', ({(115, 59, 115, 60): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((121, 34, 121, 65), 'numpy.array', 'np.array', ({(121, 43, 121, 64): '[x[0][0] for x in x_]'}, {}), '([x[0][0] for x in x_])', True, 'import numpy as np\n'), ((122, 34, 122, 65), 'numpy.array', 'np.array', ({(122, 43, 122, 64): '[x[1][0] for x in x_]'}, {}), '([x[1][0] for x in x_])', True, 'import numpy as np\n'), ((123, 34, 123, 65), 'numpy.array', 'np.array', ({(123, 43, 123, 64): '[x[2][0] for x in x_]'}, {}), '([x[2][0] for x in x_])', True, 'import numpy as np\n'), ((120, 40, 120, 51), 'numpy.array', 'np.array', ({(120, 49, 120, 50): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((126, 40, 126, 51), 'numpy.array', 'np.array', ({(126, 49, 126, 50): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((126, 53, 126, 64), 'numpy.array', 'np.array', ({(126, 62, 126, 63): 'y'}, {}), '(y)', True, 'import numpy as np\n')]