import argparse
import time
import torch
from Models import get_model, get_model_token_classification
from Process import *
import torch.nn.functional as F
from Optim import CosineWithRestarts
from Batch import create_masks
import pdb
import dill as pickle
import argparse
from Models import get_model
from Beam import beam_search
# from nltk.corpus import wordnet
from torch.autograd import Variable
import re
import time
import random
import copy


def get_yunmus(file_path):
    with open(file_path, "r", encoding="utf-8") as f:
        yunmus = f.readlines()
        yunmus = [a.strip() for a in yunmus]

    return yunmus


def get_result(src, model, SRC, TRG, opt):
    src_mask = (src != SRC.vocab.stoi['<pad>']).unsqueeze(-2)
    output = model(src, src_mask)
    output = F.softmax(output, dim=-1)
    preds = torch.argmax(output, dim=-1)
    return ''.join([TRG.vocab.itos[tok] for tok in preds[0][:] if tok.item() != 0 ]).replace("_", "").replace(" ","")
    # return ' '.join([TRG.vocab.itos[tok] for tok in preds[0][:]])


def translate_sentence(sentence, model, opt, SRC, TRG):

    model.eval()
    indexed = []
    sentence = SRC.preprocess(sentence)

    for tok in sentence:
        # 保证完全等长，删除剔除UNK的代码
        # indexed.append(SRC.vocab.stoi[tok])
        if SRC.vocab.stoi[tok] != 0 or opt.floyd == True:
            indexed.append(SRC.vocab.stoi[tok])
        else:
            # indexed.append(get_synonym(tok, SRC))
            pass
    sentence = Variable(torch.LongTensor([indexed]))
    
    # if len(ori.split(" "))!=len(indexed):
    #     print("\n")
    #     print(ori)

    if opt.device == 0:
        sentence = sentence.cuda()

    sentence = get_result(sentence, model, SRC, TRG, opt)

    return sentence


def translate(i, opt, model, SRC, TRG):
    sentence = i

    sentences = adjust_len(sentence, opt)

    translated = []

    for sentence in sentences:
        sentence=sentence.split(" ")
        if len(sentence)==1 and sentence[0]=="":
            continue
        sentence=" ".join([a for a in sentence if a != " " and a!=""])
        # result=translate_sentence(sentence, model, opt, SRC, TRG)
        # if len(sentence.split(" "))!=len(result.split(" ")):
        #     print("\n")
        #     print(len(sentence.split(" ")))
        #     print(sentence)
        #     print(len(result.split(" ")))
        #     print(result)
        translated.append(translate_sentence(sentence, model, opt, SRC, TRG))

    return (' '.join(translated)[:])


def adjust_len(sentence, opt):
    sentence = sentence.split(" ")
    if len(sentence) < opt.max_len:
        return [" ".join(sentence)]
    tmp_list = re.sub(u"(?<=([。！？]))", r"|", " ".join(sentence)).split("|")
    result_list = []
    for line in tmp_list:
        if len(line.split(" ")) < opt.max_len:
            result_list.append(line)
        else:
            tmp = re.sub(u"(?<=([。，！？]))", r"|", line).split("|")
            for i in tmp:
                if(len(i.split(" "))) < opt.max_len:
                    result_list.append(i)
                else:
                    i = i.split(" ")
                    x = opt.max_len
                    last = 0
                    while x < len(i):
                        while i[x] not in opt.yunmus:
                            x -= 1
                        result_list.append(" ".join(i[last:x+1]))
                        last = x+1
                        x = last+opt.max_len
                    result_list.append(" ".join(i[last:]))
    return result_list


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('-load_weights', required=True)
    parser.add_argument('-pkl_dir', required=True)
    parser.add_argument('-k', type=int, default=3)
    parser.add_argument('-max_len', type=int, default=200)
    parser.add_argument('-d_model', type=int, default=512)
    parser.add_argument('-n_layers', type=int, default=6)
    # parser.add_argument('-src_lang', required=True)
    # parser.add_argument('-trg_lang', required=True)
    parser.add_argument('-heads', type=int, default=8)
    parser.add_argument('-dropout', type=int, default=0.1)
    parser.add_argument('-no_cuda', action='store_true')
    parser.add_argument('-floyd', action='store_true')
    parser.add_argument('-test_dir', type=str, required=True)
    parser.add_argument('-result_dir', type=str, required=True)
    parser.add_argument('-src_voc')
    parser.add_argument('-trg_voc')

    opt = parser.parse_args()

    opt.device = 0 if opt.no_cuda is False else -1

    assert opt.k > 0
    assert opt.max_len > 10

    SRC, TRG = create_fields(opt)
    model = get_model_token_classification(opt, len(SRC.vocab), len(TRG.vocab))
    opt.yunmus = get_yunmus("./data/voc/yunmus.txt")

    for file in os.listdir(opt.test_dir):
        print("filename:{}".format(file))
        contents = open(os.path.join(opt.test_dir, file)
                        ).read().strip().split('\n')
        # contents=random.sample(contents,10)

        start = time.time()
        translates = [translate(i, opt, model, SRC, TRG)
                      for i in tqdm(contents)]
        print("Average time: {}".format((time.time()-start)/len(contents)))

        with open(os.path.join(opt.result_dir, file), 'w', encoding='utf-8') as f:
            f.write("\n".join(translates))

    # while True:
    #     opt.text =input("Enter a sentence to translate (type 'f' to load from file, or 'q' to quit):\n")
    #     if opt.text=="q":
    #         break
    #     if opt.text=='f':
    #         fpath =input("Enter a sentence to translate (type 'f' to load from file, or 'q' to quit):\n")
    #         try:
    #             opt.text = ' '.join(open(opt.text, encoding='utf-8').read().split('\n'))
    #         except:
    #             print("error opening or reading text file")
    #             continue
    #     phrase = translate(opt, model, SRC, TRG)
    #     print('> '+ phrase + '\n')


if __name__ == '__main__':
    main()
