# coding=utf-8
import os, json
import numpy as np
from curLine_file import curLine, normal_transformer

def load_data(data_dir, split=None):
    data = []
    if split is None:
        files = [os.path.join(data_dir, file) for file in os.listdir(data_dir) if file.endswith('.txt')]
    else:
        if not split.endswith('.txt'):
            split += '.txt'
        files = [os.path.join(data_dir, f'{split}')]
    print("%s read files:" % curLine(), files)
    for file in files:
        with open(file) as f:
            for line in f:
                text1, text2, label = line.rstrip().split('\t')[:3]
                data.append({
                    'text1': text1,
                    'text2': text2,
                    'target': label,
                })
    return data


def load_embeddings_Chinese(folder, vocab, dim, lower, mode='freq'): #  中文字向量
    embedding_npy_file = os.path.join(folder, "char2vecChinese_DbqaSmpLog.npy")  # char2vecChinese_DbqaSmpLogZero.npy
    embedding_npy = np.load(embedding_npy_file)
    vocab_file = os.path.join(folder, "char2idChinese_DbqaSmpLog.json")
    with open(vocab_file, "r") as fr:
        vocab_map = json.load(fr)
    # 初始化字向量
    # embedding = np.zeros((len(vocab), dim))
    embedding = np.random.normal(loc=0.0, scale=1.0e-3, size=(len(vocab), dim))
    # init_boundary = 0.1
    # embedding = np.random.uniform(low=-init_boundary, high=init_boundary, size=(len(vocab), dim))

    count = np.zeros((len(vocab), 1))
    for token,index in vocab_map.items():
        if lower and mode != 'strict':
            token = token.lower()
        if token in vocab:
            index = vocab.index(token)
            vector = embedding_npy[index]
            if mode == 'freq' or mode == 'strict':
                if not count[index]:
                    embedding[index] = vector
                    count[index] = 1.
            elif mode == 'last':
                embedding[index] = vector
                count[index] = 1.
            elif mode == 'avg':
                embedding[index] += vector
                count[index] += 1.
            else:
                raise NotImplementedError('Unknown embedding loading mode: ' + mode)
    if mode == 'avg':
        inverse_mask = np.where(count == 0, 1., 0.)
        embedding /= count + inverse_mask
    embedding = embedding.tolist()
    return embedding

def load_embeddings_English(file, vocab, dim, lower, mode='freq'): #  glove的英文词向量
    embedding = np.zeros((len(vocab), dim))
    count = np.zeros((len(vocab), 1))
    with open(file) as f:
        for line in f:
            elems = line.rstrip().split()
            if len(elems) != dim + 1:
                continue
            token = elems[0]
            if lower and mode != 'strict':
                token = token.lower()
            if token in vocab:
                index = vocab.index(token)
                vector = [float(x) for x in elems[1:]]
                if mode == 'freq' or mode == 'strict':
                    if not count[index]:
                        embedding[index] = vector
                        count[index] = 1.
                elif mode == 'last':
                    embedding[index] = vector
                    count[index] = 1.
                elif mode == 'avg':
                    embedding[index] += vector
                    count[index] += 1.
                else:
                    raise NotImplementedError('Unknown embedding loading mode: ' + mode)
    if mode == 'avg':
        inverse_mask = np.where(count == 0, 1., 0.)
        embedding /= count + inverse_mask
    return embedding.tolist()
