args = TrainArgs().get_parser()
args.gpu_ids = '0'
args.mode = "train"
args.raw_data_dir = './data/final/raw_data'
args.mid_data_dir = './data/final/mid_data'
args.aux_data_dir = "./data/final/preliminary_clean"
args.bert_dir = '/home/yuzhang/PycharmProjects/xf_event_extraction2020Top1/bert/torch_roberta_wwm'
args.output_dir = './out'
args.bert_type = 'roberta_wwm'
args.task_type = 'trigger'
args.max_seq_len = 320
args.train_epochs = 6
args.train_batch_size = 16
args.lr = 2e-5
args.other_lr = 2e-4
args.attack_train = "pgd"
args.swa_start = 4
args.eval_model = True
args.enhance_data = True
args.use_trigger_distance = True

raw_text = "学校公布成绩。"
trigger_word = "公布"
tokens = ['学', '校', '公', '布', '成', '绩', '。']
labels = [
    [0, 0], 
    [0, 0], 
    [1, 0], 
    [0, 1], 
    [0, 0], 
    [0, 0], 
    [0, 0]
]

# -*- coding: utf8 -*-
#
# -*- coding: utf8 -*-
#
import torch
from torch import nn

# shape = (3,5,2)
target = torch.tensor([
    [
        [0, 0],
        [0, 0],
        [1, 0],
        [0, 1],
        [0, 0],
    ],
    [
        [0, 0],
        [0, 0],
        [1, 0],
        [0, 1],
        [0, 0],
    ],
    [
        [0, 0],
        [0, 0],
        [1, 0],
        [0, 1],
        [0, 0],
    ],
], dtype=torch.float32)
pred = torch.randn(target.shape)

pred = nn.Sigmoid()(pred)

loss = nn.BCELoss()
print(loss(pred, target))

candidate_entities = []

start_ids = np.argwhere(logits[:, 0] > start_threshold)[:, 0]
end_ids = np.argwhere(logits[:, 1] > end_threshold)[:, 0]

# 选最短的
for _start in start_ids:
    for _end in end_ids:
        # 限定 trigger 长度不能超过 3
        if _end >= _start and _end - _start <= 2:
            # (start, end, start_logits + end_logits)
            candidate_entities.append((raw_text[_start: _end + 1], _start, logits[_start][0] + logits[_end][1]))
            break

tokenizer = BertTokenizer.from_pretrained(bert_dir)
bert = BertModel.from_pretrained(bert_dir)

# -*- coding: utf8 -*-
#
import math
from typing import List


def trace_path(vec, pathes):
    if not vec:
        return pathes
    valid_pathes = []
    for item in vec[0]:
        for path in pathes:
            if item > path[-1]:
                valid_pathes.append([*path, item])

    return trace_path(vec=vec[1:], pathes=valid_pathes)


def get_all_valid_path(vec: List[List[int]]) -> List[List[int]]:
    if not vec: return []
    for v in vec:
        if not v:
            return []
    return trace_path(vec=vec[1:], pathes=[[i] for i in vec[0]])


def _sub_sum(path: List[int]) -> int:
    _sum = 0
    for index, item in enumerate(path):
        try:
            next_item = path[index + 1]
            _sum += (next_item - item)
        except IndexError:
            pass
    return _sum


def get_shortest_path(vec: List[List[int]]) -> List[int]:
    """
    输出最短的那个路径
    :param vec:
    :return:
    """
    pathes = get_all_valid_path(vec=vec)
    short_path, short_sum = [], math.inf

    for path in pathes:
        _sum = _sub_sum(path=path)
        if short_sum > _sum:
            short_path = path
            short_sum = _sum
    return short_path

from unittest import TestCase

from path import get_all_valid_path, get_shortest_path


class TestPath(TestCase):
    def test_sample(self):
        self.assertEqual(get_all_valid_path([]), [])

        self.assertEqual(get_all_valid_path([[1, 3], [2, 5]]), [[1, 2], [1, 5], [3, 5]])

        self.assertEqual(get_shortest_path([[1, 3], [2, 5]]), [1, 2])

class ScalarMix(nn.Module):
    def __init__(self, n_layers, dropout=0.):
        super(ScalarMix, self).__init__()
        self.n_layers = n_layers

        self.weights = nn.Parameter(torch.zeros(n_layers))
        self.gamma = nn.Parameter(torch.tensor([1.0]))
        self.dropout = nn.Dropout(dropout)

    def __repr__(self):
        s = f"n_layers={self.n_layers}"
        if self.dropout.p > 0:
            s += f", dropout={self.dropout.p}"

        return f"{self.__class__.__name__}({s})"

    def forward(self, tensors):
        r"""
        Args:
            tensors (list[~torch.Tensor]):
                :math:`N` tensors to be mixed.
        Returns:
            The mixture of :math:`N` tensors.
        """

        normed_weights = self.dropout(self.weights.softmax(-1))
        weighted_sum = sum(w * h for w, h in zip(normed_weights, tensors))

        return self.gamma * weighted_sum

check_corrector_initialized()

def _initialize_corrector(self):
    # chinese common char 加载常用的汉字列表，这里大概3000多个常见的汉字。
    self.cn_char_set = self.load_set_file(self.common_char_path)
    # same pinyin 加载同音的列表，自定义
    self.same_pinyin = self.load_same_pinyin(self.same_pinyin_text_path)
    # same stroke 加载形似字，比如{坐:[座, ...]}， 自定义
    self.same_stroke = self.load_same_stroke(self.same_stroke_text_path)
    self.initialized_corrector = True

# 编码统一，utf-8 to unicode
text = convert_to_unicode(text)

def split_2_short_text(text, include_symbol=False):
    """
    长句切分为短句
    :param text: str
    :param include_symbol: bool
    :return: (sentence, idx)
    """
    result = []
    blocks = re_han.split(text)
    start_idx = 0
    for blk in blocks:
        if not blk:
            continue
        if include_symbol:
            result.append((blk, start_idx))
        else:
            if re_han.match(blk):
                result.append((blk, start_idx))
        start_idx += len(blk)
    return result

 # 词、频数dict, TODO: 这里虽然有，但是貌似没有用到
self.word_freq = self.load_word_freq_dict(self.word_freq_path)
# 自定义混淆集
self.custom_confusion = self._get_custom_confusion_dict(self.custom_confusion_path)
# 自定义切词词典
self.custom_word_freq = self.load_word_freq_dict(self.custom_word_freq_path)
self.person_names = self.load_word_freq_dict(self.person_name_path)
self.place_names = self.load_word_freq_dict(self.place_name_path)
self.stopwords = self.load_word_freq_dict(self.stopwords_path)
# 合并切词词典及自定义词典
self.custom_word_freq.update(self.person_names)
self.custom_word_freq.update(self.place_names)
self.custom_word_freq.update(self.stopwords)
self.word_freq.update(self.custom_word_freq) # TODO:这里
self.tokenizer = Tokenizer(dict_path=self.word_freq_path, custom_word_freq_dict=self.custom_word_freq,
                            custom_confusion_dict=self.custom_confusion)

('队员', 2, 4)
('少先队', 0, 3)
('少先队员', 0, 4)
('因该', 4, 6)
('为', 6, 7)
('老人', 7, 9)
('让', 9, 10)
('坐', 10, 11)

sent_scores = [-5.629326581954956, -6.566553155581156, -6.908517241477966, -7.255491574605306, -7.401519060134888, -7.489806890487671, -7.1438290278116865, -6.559153278668722, -6.858733296394348, -7.7903218269348145, -8.28114366531372]

def _get_maybe_error_index(scores, ratio=0.6745, threshold=2):
    """
    取疑似错字的位置，通过平均绝对离差（MAD）
    :param scores: np.array
    :param ratio: 正态分布表参数
    :param threshold: 阈值越小，得到疑似错别字越多
    :return: 全部疑似错误字的index: list
    """
    result = []
    scores = np.array(scores)
    if len(scores.shape) == 1:
        scores = scores[:, None]
    median = np.median(scores, axis=0)  # get median of all scores
    margin_median = np.abs(scores - median).flatten()  # deviation from the median
    # 平均绝对离差值
    med_abs_deviation = np.median(margin_median)
    if med_abs_deviation == 0:
        return result
    y_score = ratio * margin_median / med_abs_deviation
    # 打平
    scores = scores.flatten()
    maybe_error_indices = np.where((y_score > threshold) & (scores < median))
    # 取全部疑似错误字的index
    result = [int(i) for i in maybe_error_indices[0]]
    return result

[['因该', 4, 6, 'word'], ['坐', 10, 11, 'char']]

def get_lm_correct_item(self, cur_item, candidates, before_sent, after_sent, threshold=57, cut_type='char'):
    """
    通过语言模型纠正字词错误
    :param cur_item: 当前词
    :param candidates: 候选词
    :param before_sent: 前半部分句子
    :param after_sent: 后半部分句子
    :param threshold: ppl阈值, 原始字词替换后大于该ppl值则认为是错误
    :param cut_type: 切词方式, 字粒度
    :return: str, correct item, 正确的字词
    """
    result = cur_item
    if cur_item not in candidates:
        candidates.append(cur_item)
    # 对每个候选集进行拼接成新句子，然后进行计算ppl_score。
    ppl_scores = {i: self.ppl_score(segment(before_sent + i + after_sent, cut_type=cut_type)) for i in candidates}
    sorted_ppl_scores = sorted(ppl_scores.items(), key=lambda d: d[1])

    # 增加正确字词的修正范围，减少误纠
    top_items = []
    top_score = 0.0
    for i, v in enumerate(sorted_ppl_scores):
        v_word = v[0]
        v_score = v[1]
        if i == 0:
            top_score = v_score
            top_items.append(v_word)
        # 通过阈值修正范围
        elif v_score < top_score + threshold:
            top_items.append(v_word)
        else:
            break
    if cur_item not in top_items:
        result = top_items[0]
    return result

def ppl_score(self, words):

    
    """
    words比如：['少', '先', '队', '员', '应', '该', '为', '老', '人', '让', '坐']
    取语言模型困惑度得分，越小句子越通顺
    :param words: list, 以词或字切分
    :return:
    """
    self.check_detector_initialized()
    return self.lm.perplexity(' '.join(words))

pprint(sorted_ppl_scores)
[('应该', 144.39704182754554),
 ('因改', 236.80615502078768),
 ('因该', 284.14769660593794),
 ('听该', 357.8835799332408),
 ('因盖', 360.68106481988417),
 ('因核', 365.9438178618582),
 # 这里只截取一部分！！
]

#download from https://github.com/HIT-SCIR/ELMoForManyLangs
model_file = r'../auxiliary_data/zhs.model/'
# 初始化elmo
ELMO = word_emb_elmo.WordEmbeddings(model_file)
SIF = sent_emb_sif.SentEmbeddings(ELMO, lamda=1.0)

class SentEmbeddings():

    def __init__(self,
                 word_embeddor,
                 weightfile_pretrain='../auxiliary_data/dict.txt',
                 weightfile_finetune='../auxiliary_data/dict.txt',
                 weightpara_pretrain=2.7e-4,
                 weightpara_finetune=2.7e-4,
                 lamda=1.0,database="",embeddings_type="elmo"):

        self.word2weight_pretrain = get_word_weight(weightfile_pretrain, weightpara_pretrain)
        self.word2weight_finetune = get_word_weight(weightfile_finetune, weightpara_finetune)
        self.word_embeddor = word_embeddor
        self.lamda=lamda
        self.database=database
        self.embeddings_type=embeddings_type

def get_word_weight(weightfile="", weightpara=2.7e-4):
    """
    Get the weight of words by word_fre/sum_fre_words
    :param weightfile
    :param weightpara
    :return: word2weight[word]=weight : a dict of word weight
    """
    if weightpara <= 0:  # when the parameter makes no sense, use unweighted
        weightpara = 1.0
    word2weight = {}
    word2fre = {}
    with open(weightfile, encoding='UTF-8') as f:
        lines = f.readlines()
    # sum_num_words = 0
    sum_fre_words = 0
    for line in lines:
        word_fre = line.split()
        # sum_num_words += 1
        if (len(word_fre) >= 2):
            word2fre[word_fre[0]] = float(word_fre[1])
            sum_fre_words += float(word_fre[1])
        else:
            print(line)
    # 这地方效果：如果一个词出现的频率越高，那么它的权重就越低
    # 可看： https://github.com/sunyilgdx/SIFRank_zh/issues/14
    for key, value in word2fre.items():
        word2weight[key] = weightpara / (weightpara + value / sum_fre_words)
        # word2weight[key] = 1.0 #method of RVA
    return word2weight

zh_model = thulac.thulac(model_path=r'../auxiliary_data/thulac.models/',user_dict=r'../auxiliary_data/user_dict.txt')
elmo_layers_weight = [0.0, 1.0, 0.0]

class InputTextObj:
    """Represent the input text in which we want to extract keyphrases"""

    def __init__(self, zh_model, text=""):
        """
        :param is_sectioned: If we want to section the text.
        :param zh_model: the pipeline of Chinese tokenization and POS-tagger
        :param considered_tags: The POSs we want to keep
        """
        self.considered_tags = {'n', 'np', 'ns', 'ni', 'nz','a','d','i','j','x','g'}

        self.tokens = []
        self.tokens_tagged = []
        # self.tokens = zh_model.cut(text)
        word_pos = zh_model.cut(text)
        self.tokens = [word_pos[0] for word_pos in word_pos]
        self.tokens_tagged = [(word_pos[0],word_pos[1]) for word_pos in word_pos]
        assert len(self.tokens) == len(self.tokens_tagged)
        for i, token in enumerate(self.tokens): # 停用词处理
            if token.lower() in stopword_dict:
                self.tokens_tagged[i] = (token, "u")
            if token == '-':
                self.tokens_tagged[i] = (token, "-")
        self.keyphrase_candidate = extractor.extract_candidates(self.tokens_tagged, zh_model)

GRAMMAR_zh = """  NP:
        {<n.*|a|uw|i|j|x>*<n.*|uw|x>|<x|j><-><m|q>} # Adjective(s)(optional) + Noun(s)"""

def extract_candidates(tokens_tagged, no_subset=False):
    """
    Based on part of speech return a list of candidate phrases
    :param text_obj: Input text Representation see @InputTextObj
    :param no_subset: if true won't put a candidate which is the subset of an other candidate
    :return keyphrase_candidate: list of list of candidate phrases: [tuple(string,tuple(start_index,end_index))]
    """
    np_parser = nltk.RegexpParser(GRAMMAR_zh)  # Noun phrase parser
    keyphrase_candidate = []
    np_pos_tag_tokens = np_parser.parse(tokens_tagged)
    count = 0
    for token in np_pos_tag_tokens:
        if (isinstance(token, nltk.tree.Tree) and token._label == "NP"):
            np = ''.join(word for word, tag in token.leaves())
            length = len(token.leaves())
            start_end = (count, count + length)
            count += length
            keyphrase_candidate.append((np, start_end))

        else:
            count += 1

    return keyphrase_candidate

def get_sent_segmented(tokens):
    """
    按照.和。进行分段
    但是这个分法很个性，只有大于16才分割，否则不分
    :param tokens:
    :return:
    """
    min_seq_len = 16
    sents_sectioned = []
    if (len(tokens) <= min_seq_len):
        sents_sectioned.append(tokens)
    else:
        position = 0
        for i, token in enumerate(tokens):
            if (token == '.' or token =='。'):
                if (i - position >= min_seq_len):
                    sents_sectioned.append(tokens[position:i + 1])
                    position = i + 1
        if (len(tokens[position:]) > 0):
            sents_sectioned.append(tokens[position:])

    return sents_sectioned

def get_tokenized_words_embeddings(self, sents_tokened):
    """
    @see EmbeddingDistributor
    :param tokenized_sents: list of tokenized words string (sentences/phrases)
    :return: ndarray with shape (len(sents), dimension of embeddings)
    """
    max_len = max([len(sent) for sent in sents_tokened])
    elmo_embedding = self.elmo.sents2elmo(sents_tokened,output_layer=-2)
    elmo_embedding = [np.pad(emb, pad_width=((0,0),(0,max_len-emb.shape[1]),(0,0)) , mode='constant') for emb in elmo_embedding]
    elmo_embedding = torch.from_numpy(np.array(elmo_embedding))
    return elmo_embedding

def context_embeddings_alignment(elmo_embeddings, tokens_segmented):

    """
    Embeddings Alignment
    :param elmo_embeddings: The embeddings from elmo
    :param tokens_segmented: The list of tokens list
     <class 'list'>: [['今', '天', '天气', '真', '好', '啊'],['潮水', '退', '了', '就', '知道', '谁', '没', '穿', '裤子']]
    :return:
    """
    token_emb_map = {}
    n = 0
    for i in range(0, len(tokens_segmented)):
        # 一词多义，将相同词的emb append到一起
        for j, token in enumerate(tokens_segmented[i]):
            # 获取第一层的embedding, 1 for the first LSTM hidden layer
            emb = elmo_embeddings[i, 1, j, :]
            if token not in token_emb_map:
                token_emb_map[token] = [emb]
            else:
                token_emb_map[token].append(emb)
            n += 1
    # 求每个token的平均值
    anchor_emb_map = {}
    for token, emb_list in token_emb_map.items():
        average_emb = emb_list[0]
        for j in range(1, len(emb_list)):
            average_emb += emb_list[j]
        average_emb /= float(len(emb_list))
        anchor_emb_map[token] = average_emb
    # 替换掉elmo最后一层的词对应的词向量
    for i in range(0, elmo_embeddings.shape[0]):
        for j, token in enumerate(tokens_segmented[i]):
            emb = anchor_emb_map[token]
            elmo_embeddings[i, 2, j, :] = emb

    return elmo_embeddings

def splice_embeddings(elmo_embeddings,tokens_segmented):
    new_elmo_embeddings = elmo_embeddings[0:1, :, 0:len(tokens_segmented[0]), :]
    for i in range(1, len(tokens_segmented)):
        emb = elmo_embeddings[i:i + 1, :, 0:len(tokens_segmented[i]), :]
        new_elmo_embeddings = torch.cat((new_elmo_embeddings, emb), 2)
    return new_elmo_embeddings

def get_weight_list(word2weight_pretrain, word2weight_finetune, tokenized_sents, lamda, database=""):
    weight_list = []
    for word in tokenized_sents:
        word = word.lower()

        if(database==""):
            weight_pretrain = get_oov_weight(tokenized_sents, word2weight_pretrain, word, method="max_weight")
            weight=weight_pretrain
        else:
            weight_pretrain = get_oov_weight(tokenized_sents, word2weight_pretrain, word, method="max_weight")
            weight_finetune = get_oov_weight(tokenized_sents, word2weight_finetune, word, method="max_weight")
            weight = lamda * weight_pretrain + (1.0 - lamda) * weight_finetune
        weight_list.append(weight)

    return weight_list

def get_oov_weight(tokenized_sents,word2weight,word,method="max_weight"):
    try:
        word=wnl.lemmatize(word)
    except:
        nltk.download('wordnet')
        nltk.download('omw-1.4')

    if(word in word2weight):#
        return word2weight[word]

    if(word in stop_words):
        return 0.0

    if(word in english_punctuations or word in chinese_punctuations):#The oov_word is a punctuation
        return 0.0

    if(method=="max_weight"):#Return the max weight of word in the tokenized_sents
        max=0.0
        for w in tokenized_sents:
            if(w in word2weight and word2weight[w]>max):
                max=word2weight[w]
        return max
    return 0.0

def get_weighted_average(tokenized_sents, sents_tokened_tagged,weight_list, embeddings_list, embeddings_type="elmo"):
    # weight_list=get_normalized_weight(weight_list)
    assert len(tokenized_sents) == len(weight_list)
    num_words = len(tokenized_sents)
    e_test_list=[]
    if (embeddings_type == "elmo" or embeddings_type == "elmo_sectioned"):
        # assert num_words == embeddings_list.shape[1]
        sum = torch.zeros((3, 1024))
        for i in range(0, 3):
            for j in range(0, num_words):
                if(sents_tokened_tagged[j][1] in considered_tags):
                    e_test=embeddings_list[i][j]
                    e_test_list.append(e_test)
                    sum[i] += e_test * weight_list[j]

            sum[i] = sum[i] / float(num_words)
        return sum
    else:
        print('其他的忽略。')

def get_candidate_weighted_average(tokenized_sents, weight_list, embeddings_list, start,end,embeddings_type="elmo"):
    """获取候选词在每一层的平均向量"""
    # weight_list=get_normalized_weight(weight_list)
    assert len(tokenized_sents) == len(weight_list)
    # num_words = len(tokenized_sents)
    num_words =end - start
    e_test_list=[]
    if (embeddings_type == "elmo" or embeddings_type == "elmo_sectioned"):
        # assert num_words == embeddings_list.shape[1]
        sum = torch.zeros((3, 1024))
        for i in range(0, 3):
            for j in range(start, end):
                e_test=embeddings_list[i][j]
                e_test_list.append(e_test)
                sum[i] += e_test * weight_list[j]
            sum[i] = sum[i] / float(num_words)

        return sum

def get_dist_cosine(emb1, emb2, sent_emb_method="elmo",elmo_layers_weight=[0.0,1.0,0.0]):
    sum = 0.0
    assert emb1.shape == emb2.shape
    if(sent_emb_method=="elmo"):

        for i in range(0, 3):
            a = emb1[i]
            b = emb2[i]
            sum += cos_sim(a, b) * elmo_layers_weight[i]
        return sum

def get_all_dist(candidate_embeddings_list, text_obj, dist_list):
    '''
    :param candidate_embeddings_list:
    :param text_obj:
    :param dist_list:
    :return: dist_all
    '''

    dist_all={}
    for i, emb in enumerate(candidate_embeddings_list):
        phrase = text_obj.keyphrase_candidate[i][0]
        phrase = phrase.lower()
        phrase = wnl.lemmatize(phrase) # 词形还原
        if(phrase in dist_all):
            #store the No. and distance
            dist_all[phrase].append(dist_list[i])
        else:
            dist_all[phrase]=[]
            dist_all[phrase].append(dist_list[i])
    return dist_all

def get_final_dist(dist_all, method="average"):
    '''
    :param dist_all:
    :param method: "average"
    :return:
    '''

    final_dist={}

    if(method=="average"):

        for phrase, dist_list in dist_all.items():
            sum_dist = 0.0
            for dist in dist_list:
                sum_dist += dist
            if (phrase in stop_words):
                sum_dist = 0.0
            final_dist[phrase] = sum_dist/float(len(dist_list))
        return final_dist

def get_position_score(keyphrase_candidate_list, position_bias):
    length = len(keyphrase_candidate_list)
    position_score ={}
    for i,kc in enumerate(keyphrase_candidate_list):
        np = kc[0]
        p = kc[1][0]
        np = np.lower()
        np = wnl.lemmatize(np)
        if np in position_score:

            position_score[np] += 0.0
        else:
            position_score[np] = 1/(float(i)+1+position_bias)
    score_list=[]
    for np,score in position_score.items():
        score_list.append(score)
    score_list = softmax(score_list)

    i=0
    for np, score in position_score.items():
        position_score[np] = score_list[i]
        i+=1
    return position_score

[
    ('现在', 'time'),
    ('社区', 'n'),
    ('在', 'p'),
    ('进行', 'v'),
    ('网格化', 'n'),
    ('管理', 'n')
]

[('社区', 'n'), ('网格化', 'n'), ('管理', 'n')]

if before_token.pos == 'n' and cur_token.pos == 'n':
    print('那么就聚合')

def token_match(token, element):
    # ！开头表示词性
    if element.startswith('!'):
        return token.pos == element[1:]
    # @表示词
    elif element.startswith('@'):
        return token.word == element[1:]
    raise # 自行实现其他的

# 每个表达式判断
for element in express:
    for index，token in enumerate(tokens):
        if token_match(token, element):
            # 最终构成一个DAG

>>> sentence = [('the', 'DT'), ('little', 'JJ'), ('yellow', 'JJ'), ('dog', 'NN'), ('barked', 'VBD'), ('at', 'IN'), ('the', 'DT'), ('cat', 'NN')]
>>> grammar = "NP: {<DT>?<JJ>*<NN>}"
>>> cp = nltk.RegexpParser(grammar)
>>> result = cp.parse(sentence)

>>> result.pretty_print()
                            S                                       
     _______________________|______________________________          
    |        |              NP                             NP       
    |        |      ________|_________________        _____|____     
barked/VBD at/IN the/DT little/JJ yellow/JJ dog/NN the/DT     cat/NN