import torch

import argparse

from datasets import QQRDataset,QQR_data,BertClassificationDataset

from tqdm import tqdm

from gensim.models import KeyedVectors

import time

from torch.utils.data import DataLoader

from models import SemNN,SemLSTM,SemAttention

from transformers import AutoTokenizer

import os

import torch.nn as nn

import json

from transformers import BertForSequenceClassification

model_type1_list = ['SemNN','SemAttention','SemLSTM']

model_type2_list = ['Bert']

def inference(args):

    batch_size = args.batch_size

    save_path = args.savepath

    data_dir = args.datadir

    w2v_path = args.w2v_path

    max_length = args.max_length

    model_path = args.model_path

    model_name = args.model_name

    in_feat = args.in_feat

    dropout_prob = args.dropout_prob

    if model_name in model_type1_list:

        begin_time = time.perf_counter()

        w2v_model = KeyedVectors.load_word2vec_format(w2v_path,binary=False)

        end_time = time.perf_counter()

        print("Load {} cost {:.2f}s".format(w2v_path,end_time-begin_time))

        w2v_map = w2v_model.key_to_index

    elif model_name in model_type2_list:

        tokenizer = AutoTokenizer.from_pretrained(w2v_path)

    device = torch.device("cuda:{}".format(args.gpu) if torch.cuda.is_available() else 'cpu')

    if not os.path.exists(save_path):

        os.makedirs(save_path)

    data = QQR_data(data_dir)

    if model_name in model_type1_list:

        test_dataset = QQRDataset(data.get_test_data(),data.get_labels(),w2v_map=w2v_map,max_length=max_length)

    elif model_name in model_type2_list:

        test_dataset = BertClassificationDataset(data.get_test_data(),tokenizer=tokenizer,label_list=data.get_labels(),max_length=max_length)

    id2label = test_dataset.id2label

    dataloader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=4)

    if model_name == "SemNN":

        model = SemNN(

            in_feat=100,

            num_labels=len(data.get_labels()),

            dropout_prob=dropout_prob,

            w2v_mapping=w2v_model

        )

    elif model_name == "SemLSTM":

        model = SemLSTM(

            in_feat=in_feat,

            num_labels=len(data.get_labels()),

            dropout_prob=dropout_prob,

            w2v_mapping=w2v_model

        )

    elif model_name == "SemAttention":

        model = SemAttention(

            in_feat=in_feat,

            num_labels = len(data.get_labels()),

            dropout_prob=dropout_prob,

            w2v_mapping=w2v_model

        )

    elif model_name == "Bert":

        model = BertForSequenceClassification.from_pretrained(w2v_path,num_labels=len(data.get_labels()))

    print(model)

    # model_paramters = model.parameters()

    print('Model Name: '+model_name)

    print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0))

    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)

    model.load_state_dict(checkpoint['state_dict'])

    model.to(device)

    model.eval()

    print(model)

    json_results = []

    preds = 0

    for text_example in dataloader:

        text_a = text_example.get('text_a')

        text_b = text_example.get('text_b')

        idx = text_example.get('idx')

        if model_name in model_type1_list:

            text_a_inputs_id = text_example.get("text_a_inputs_id").to(device)

            text_b_inputs_id = text_example.get("text_b_inputs_id").to(device)

            text_a_attention_mask = text_example.get("text_a_attention_mask").to(device)

            text_b_attention_mask = text_example.get("text_b_attention_mask").to(device)

        elif model_name in model_type2_list:

            input_ids = text_example.get('input_ids').to(device)

            token_type_ids = text_example.get('token_type_ids').to(device)

            attention_mask = text_example.get('attention_mask').to(device)

        with torch.no_grad():

            if model_name in model_type1_list:

                outputs = model(text_a_inputs_id,text_b_inputs_id,text_a_attention_mask,text_b_attention_mask)

            elif model_name in model_type2_list:

                outputs = model(input_ids=input_ids,token_type_ids=token_type_ids,attention_mask=attention_mask,return_dict=True).get('logits')

            probs = nn.Softmax(dim=1)(outputs)

            preds = torch.max(probs,1)[1].data.cpu()

            # print(preds)

        for i in range(outputs.size(0)):

            json_results.append({

                "id":idx[i],

                "query1":text_a[i],

                "query2":text_b[i],

                "label":id2label[preds[i].item()]

            })

            # print(json_results)

            # break

        with open(os.path.join(save_path,'results_test.json'),'w',encoding='utf-8') as f:

            json.dump(json_results,f,ensure_ascii=False,indent=2)

            f.close()

if __name__ == "__main__":

    parse = argparse.ArgumentParser()

    parse.add_argument('--model_name',type=str,default="SemAttention",help="Model name for train")

    parse.add_argument('--in_feat',type=int,default=100,help="Length of features for embbeding word")

    parse.add_argument('--dropout_prob',type=float,default=0.1,help="Dropout ratio for dropout layers")

    parse.add_argument('--batch_size',type=int,default=128,help="Batch-size for train")

    parse.add_argument('--max_length',type=int,default=32,help="Max length for setence")

    parse.add_argument('--savepath',type=str,default="./results/SemAttention",help="Save dir for trained model")

    parse.add_argument('--datadir',type=str,default='./data',help="Data path for train and test")

    parse.add_argument('--model_path',type=str,default='./results/SemAttention/best_model.pth.tar',help="Saved model path")

    parse.add_argument('--gpu',type=str,default='1',help="Gpu id for train")

    parse.add_argument('--w2v_path',type=str,default='./tencent-ailab-embedding-zh-d100-v0.2.0-s/tencent-ailab-embedding-zh-d100-v0.2.0-s.txt',help="Path for w2v_model file")

    args = parse.parse_args()

    inference(args)