import sys, pickle, os, random
import numpy as np

## tags, BIO
tag2label = {"O": 0,
             "B-P": 1, "I-P": 2,
             "B-S": 3, "I-S": 4,
             }
# tag2label = {"O": 0,
#              "B-PER": 1, "I-PER": 2,
#              "B-LOC": 3, "I-LOC": 4,
#              "B-ORG": 5, "I-ORG": 6
#              }

data_list = ['实现祖国的完全统一，是海内外全体/o 中华/ns 儿女的共同心愿，也是历史赋予我们的重任。',
             '我们热诚希望/o 致公党/nt 在今后的工作中，充分发挥与海外联系广泛的优势，多渠道、多层次、多形式地开展海外联络工作，积极宣传/o 邓小平/nr “一国两制”的科学构想和/o 江泽民/nr 同志关于/o 台湾/ns 问题的八项主张，为保持/o 香港/ns 的繁荣稳定，为/o 澳门/ns 的顺利回归和促进祖国完全统一作出新的更大的贡献。',
             '致公党第十一次全国代表大会/nt 是/o 致公党/nt 历史上一次重要的会议。',
             '大会将完成 致公党中央/nt 领导集体跨世纪的新老交替，确定今后一个时期的工作任务。',
             '我们相信，面临新的形势任务，这次大会即将选出的新一届/o 致公党中央/nt 领导集体，一定能够在实现新老交替的基础上完成政治交接，把老一辈领导人同/o 中国共产党/nt 亲密合作的优良传统继承下来并发扬光大，切实加强自身建设，不断提高/o 致公党/nt 的整体素质，充分发挥参政党作用，团结和带领广大成员和所联系的归侨、侨眷，埋头苦干，扎实工作，在建设有/o 中国/ns 特色社会主义伟大事业中创造出新的业绩。',
             '让我们在/o 邓小平/nr 理论的指引下，更加紧密地团结在以/o 江泽民/nr 同志为核心的/o 中共中央/nt 周围，以高度的历史责任感和时代紧迫感，积极投身于改革开放和社会主义现代化建设的伟大实践，共同谱写我国社会主义现代化建设的新篇章，迎接辉煌的２１世纪！',
             '祝中国致公党第十一次全国代表大会/nt 圆满成功！',
             ]


def read_corpus(corpus_path):
    """
    read corpus and return the list of samples
    :param corpus_path:
    :return: data
    """
    data = []
    with open(corpus_path, encoding='utf-8') as fr:
        lines = fr.readlines()
    sent_, tag_ = [], []
    for line in lines:
        if line != '\n':
            [char, label] = line.strip().split()
            sent_.append(char)
            tag_.append(label)
        else:
            data.append((sent_, tag_))
            sent_, tag_ = [], []

    return data


def vocab_build(vocab_path, corpus_path, min_count):
    """
    :param vocab_path:
    :param corpus_path:
    :param min_count:
    :return:
    """
    data = read_corpus(corpus_path)
    word2id = {}
    for sent_, tag_ in data:
        for word in sent_:
            if word.isdigit():
                word = '<NUM>'
            elif ('\u0041' <= word <='\u005a') or ('\u0061' <= word <='\u007a'):
                word = '<ENG>'
            if word not in word2id:
                word2id[word] = [len(word2id)+1, 1]
            else:
                word2id[word][1] += 1
    low_freq_words = []
    for word, [word_id, word_freq] in word2id.items():
        if word_freq < min_count and word != '<NUM>' and word != '<ENG>':
            low_freq_words.append(word)
    for word in low_freq_words:
        del word2id[word]

    new_id = 1
    for word in word2id.keys():
        word2id[word] = new_id
        new_id += 1
    word2id['<UNK>'] = new_id
    word2id['<PAD>'] = 0

    print(len(word2id))
    with open(vocab_path, 'wb') as fw:
        pickle.dump(word2id, fw)


def sentence2id(sent, word2id):
    """

    :param sent:
    :param word2id:
    :return:
    """
    sentence_id = []
    for word in sent:
        if word.isdigit():
            word = '<NUM>'
        elif ('\u0041' <= word <= '\u005a') or ('\u0061' <= word <= '\u007a'):
            word = '<ENG>'
        if word not in word2id:
            word = '<UNK>'
        sentence_id.append(word2id[word])
    return sentence_id


def read_dictionary(vocab_path):
    """

    :param vocab_path:
    :return:
    """
    vocab_path = os.path.join(vocab_path)
    with open(vocab_path, 'rb') as fr:
        word2id = pickle.load(fr)
    print('vocab_size:', len(word2id))
    return word2id


def random_embedding(vocab, embedding_dim):
    """

    :param vocab:
    :param embedding_dim:
    :return:
    """
    embedding_mat = np.random.uniform(-0.25, 0.25, (len(vocab), embedding_dim))
    embedding_mat = np.float32(embedding_mat)
    return embedding_mat


def pad_sequences(sequences, pad_mark=0):
    """

    :param sequences:
    :param pad_mark:
    :return:
    """
    max_len = max(map(lambda x: len(x), sequences))
    seq_list, seq_len_list = [], []
    for seq in sequences:
        seq = list(seq)
        seq_ = seq[:max_len] + [pad_mark] * max(max_len - len(seq), 0)
        seq_list.append(seq_)
        seq_len_list.append(min(len(seq), max_len))
    return seq_list, seq_len_list


def batch_yield(data, batch_size, vocab, tag2label, shuffle=False):
    """

    :param data:
    :param batch_size:
    :param vocab:
    :param tag2label:
    :param shuffle:
    :return:
    """
    if shuffle:
        random.shuffle(data)

    seqs, labels = [], []
    for (sent_, tag_) in data:
        sent_ = sentence2id(sent_, vocab)
        label_ = [tag2label[tag] for tag in tag_]

        if len(seqs) == batch_size:
            yield seqs, labels
            seqs, labels = [], []

        seqs.append(sent_)
        labels.append(label_)

    if len(seqs) != 0:
        yield seqs, labels


if __name__ == '__main__':
    word2id = read_dictionary('data_path/word2id.pkl')
    print(word2id)
    # vocab_build('data_path/word2id.pkl', 'data_path/train_data', min_count=1)