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

import os

import torch.nn as nn

import torch.optim as optim

from transformers import AutoTokenizer

from transformers import BertForSequenceClassification

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

model_type2_list = ['Bert']

def train(args):

    batch_size = args.batch_size

    lr = args.lr

    save_path = args.savepath

    data_dir = args.datadir

    w2v_path = args.w2v_path

    max_length = args.max_length

    epochs = args.epochs

    model_name = args.model_name

    dropout_prob = args.dropout_prob

    in_feat = args.in_feat

    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')

    save_path = os.path.join(save_path,model_name)

    if not os.path.exists(save_path):

        os.makedirs(save_path)

    data = QQR_data(data_dir)

    if model_name in model_type1_list:

        train_dataset = QQRDataset(data.get_train_data(),data.get_labels(),w2v_map=w2v_map,max_length=max_length)

        val_dataset = QQRDataset(data.get_dev_data(),data.get_labels(),w2v_map=w2v_map,max_length=max_length)

    elif model_name in model_type2_list:

        train_dataset = BertClassificationDataset(data.get_train_data(),tokenizer=tokenizer,label_list=data.get_labels(),max_length=max_length)

        val_dataset = BertClassificationDataset(data.get_dev_data(),tokenizer=tokenizer,label_list=data.get_labels(),max_length=max_length)

    train_examples_num = len(train_dataset)

    val_examples_num = len(val_dataset)

    dataset = {'train':train_dataset,'val':val_dataset}

    len_dataset = {'train':train_examples_num,'val':val_examples_num}

    if model_name == "SemNN":

        model = SemNN(

            in_feat=in_feat,

            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()

    criterion = nn.CrossEntropyLoss()

    criterion.to(device)

    model.to(device)

    optimizer = optim.SGD(model_paramters, lr=lr, momentum=0.9, weight_decay=5e-4)

    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10,gamma=0.1)

    print('Model Name: '+model_name)

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

    best_val_acc = 0.0

    for epoch in range(epochs):

        for phase in ['train','val']:

            runing_loss = 0.0

            running_corrects = 0.0

            #Train or eval

            if(phase=='train'):

                scheduler.step()

                model.train()

            else:

                model.eval()

            dataloader = DataLoader(dataset[phase],batch_size=batch_size,shuffle=True,num_workers=4)

            for text_example in tqdm(dataloader):

                if model_name in model_type1_list:

                    text_a_inputs_id = text_example["text_a_inputs_id"].to(device)

                    text_b_inputs_id = text_example["text_b_inputs_id"].to(device)

                    text_a_attention_mask = text_example["text_a_attention_mask"].to(device)

                    text_b_attention_mask = text_example["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)

                labels = text_example['labels'].to(device)

                optimizer.zero_grad()

                if(phase=='train'):

                    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')

                else:

                    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]

                # print(preds.sum())

                loss = criterion(outputs,labels)

                if(phase=='train'):

                    loss.backward()

                    optimizer.step()

                runing_loss += loss.item() * labels.size(0)

                running_corrects += torch.sum(preds==labels.data)

            epoch_loss = runing_loss/len_dataset.get(phase)

            epoch_acc = running_corrects.double()/len_dataset.get(phase)

            if(phase=='val'):

                if(best_val_acc<epoch_acc):

                    best_val_acc = epoch_acc

                    torch.save({

                        'epoch':epoch+1,

                        'state_dict':model.state_dict(),

                        'opt_dict':optimizer.state_dict()

                    },os.path.join(save_path,'best_model.pth.tar'))

            if(epoch==epochs-1):

                torch.save({

                        'epoch':epoch+1,

                        'state_dict':model.state_dict(),

                        'opt_dict':optimizer.state_dict()

                    },os.path.join(save_path,'lastest_model.pth.tar'))

            print("[{}] Epoch: {}/{} Loss: {} Acc: {}".format(phase, epoch+1, epochs, epoch_loss, epoch_acc))

if __name__ == "__main__":

    parse = argparse.ArgumentParser()

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

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

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

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

    parse.add_argument('--epochs',type=int,default=50,help="Set epochs for train")

    parse.add_argument('--lr',type=float,default=1e-3,help="Learning Rate for train")

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

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

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

    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()

    train(args)