import jieba
import re
from config import *
import torch
from model import get_padding_mask,get_subsequent_mask
import sacrebleu

# 中文分词
def divided_zh(sentence):
    return jieba.lcut(sentence)

# 英文分词
def divided_en(sentence):
    # 使用正则表达式匹配单词和标点符号
    pattern = r'\w+|[^\w\s]'
    return re.findall(pattern, sentence)

def get_vocab(lang='en'):
    if lang == 'en':
        file_path = EN_VOCAB_PATH
    elif lang == 'zh':
        file_path = ZH_VOCAB_PATH
    
    with open(file_path, encoding='utf-8') as file:
        lines = file.read()

    id2vocab = lines.split('\n')
    vocab2id = {v:k for k,v in enumerate(id2vocab)}
    return id2vocab, vocab2id

#逐字生成预测值
def batch_greedy_decode(model, src_x, src_mask, max_len=MAX_LEN):
    src_x= src_x.to(device)
    src_mask=src_mask.to(device)
    #获取中文词表
    zh_id2vocab, _ = get_vocab('zh')
    #encoder 输出
    memory = model.encoder(src_x, src_mask)
    # 初始化decoder 输入
    prob_x = torch.tensor([[SOS_ID]] * src_x.size(0))
    prob_x = prob_x.to(device)

    for _ in range(max_len):
        #推理时mask
        # prob_mask = get_padding_mask(prob_x, PAD_ID).to(device)
        prob_pad_mask = get_padding_mask(prob_x, PAD_ID).to(device)
        prob_subsqueent_mask = get_subsequent_mask(prob_x.size(1)).to(device)
        prob_mask = prob_pad_mask | prob_subsqueent_mask
        prob_mask = prob_mask != 0
        
        decoder_output = model.decoder(prob_x, prob_mask, memory, src_mask) # shape 3,1,512
        output = model.generator(decoder_output[:, -1, :]) #shape 3,19
        output = torch.softmax(output,dim=-1)
        predict = torch.argmax(output, dim=-1, keepdim=True)
        #把之前的值拼到一起
        prob_x = torch.concat([prob_x, predict], dim=-1)
        #如果预测的当前值是eos，则表示整个序列结束
        if torch.all(predict==EOS_ID).item():
            break
   #把预测出的索引值，转成对应的词   
    batch_prob_text = []
    for prob in prob_x:
        prob_text = []
        for prob_id in prob:
            if prob_id == SOS_ID:
                continue
            if prob_id == EOS_ID:
                break
            prob_text.append(zh_id2vocab[prob_id])
        batch_prob_text.append(''.join(prob_text))
    return batch_prob_text


def bleu_score(hyp, refs):
    bleu = sacrebleu.corpus_bleu(hyp, refs, tokenize='zh')
    return round(bleu.score, 2)


if __name__ == '__main__':
    target = '我喜欢读书'
    vocabs = divided_zh(target)
    _, zh_vocab2id = get_vocab('zh')
    tokens = [zh_vocab2id.get(v, UNK_ID) for v in vocabs]
    print(tokens)