"""
Name   : train_rnn_agnews.py
Author : Zhijie Wang
Time   : 2021/7/23
"""


from data.data_utils import CommentDataset
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from transformers import BertTokenizerFast
from torchtext.datasets import AG_NEWS
import argparse
import pandas as pd
import numpy as np
from model.simple_rnn import SimpleGRUMultiClassification


parser = argparse.ArgumentParser()
parser.add_argument('--dataset_name', dest='dataset_name', default='agnews', help='name of dataset')
parser.add_argument('--train_file', dest='train_file', default='./file/data/train_subset.csv', help='path to data file')
parser.add_argument('--out_path', dest='out_path', default='./file/checkpoints/', help='output path')
parser.add_argument('--usegpu', dest='usegpu', default=True, type=bool, help='use gpu or not')
parser.add_argument('--batch_size', dest='batch_size', default=10, type=int, help='batch size')
parser.add_argument('--num_epochs', dest='num_epochs', default=40, type=int, help='num of epochs')
parser.add_argument('--lr', dest='lr', default=1e-4, type=float, help='learning rate')
parser.add_argument('--num_layers', dest='num_layers', default=2, type=int, help='num of rnn layers')
parser.add_argument('--word_vec_size', dest='word_vec_size', default=200, type=int, help='word embedding size')
parser.add_argument('--rnn_size', dest='rnn_size', default=256, type=int, help='hidden state dim')
parser.add_argument('--dense_hidden_dim', dest='dense_hidden_dim', default=[64], type=list, help='hidden state dim')


if __name__ == '__main__':
    args = parser.parse_args()
    out_path = args.out_path

    train_iter = AG_NEWS(root='./file/data/', split='train')
    X_train = {'comment_text': [], 'target': []}
    for label, line in train_iter:
        X_train['comment_text'] += [line]
        X_train['target'] += [label - 1]
    X_train = pd.DataFrame(data=X_train)

    test_iter = AG_NEWS(root='./file/data/', split='test')
    X_val = {'comment_text': [], 'target': []}
    for label, line in test_iter:
        X_val['comment_text'] += [line]
        X_val['target'] += [label - 1]
    X_val = pd.DataFrame(data=X_val)

    tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
    max_sentence_length = 200

    train_dataset = CommentDataset(df=X_train, tokenizer=tokenizer, max_length=max_sentence_length,
                                   data_col="comment_text", target="target", is_testing=False)

    train_dataloader = DataLoader(train_dataset, num_workers=0, batch_size=args.batch_size, shuffle=True)

    val_dataset = CommentDataset(df=X_val, tokenizer=tokenizer, max_length=max_sentence_length,
                                 data_col="comment_text", target="target", is_testing=False)

    val_dataloader = DataLoader(val_dataset, num_workers=0, batch_size=args.batch_size, shuffle=False)

    device = torch.device('cuda' if args.usegpu else 'cpu')

    model = SimpleGRUMultiClassification(word_vec_size=args.word_vec_size, rnn_size=args.rnn_size,
                                         embedding_size=tokenizer.vocab_size, num_layers=args.num_layers,
                                         dense_hidden_dim=args.dense_hidden_dim, target_size=4)
    model = model.to(device)
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    for epoch in range(args.num_epochs):
        model.train()
        avg_loss = 0
        batch_num = 0
        for batch in train_dataloader:
            input_tensor, target_tensor = batch['input_ids'].to(device).long(), batch['y'].to(device)
            y_pred = model(input_tensor)
            loss = loss_function(y_pred, target_tensor)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            avg_loss += loss.item()  # / len(train_loader)
            batch_num += 1
            if batch_num > -1 and batch_num % 100 == 0:
                print('Train Epoch: %d, Batch: %d/%d, Loss: %6f' % (
                epoch + 1, batch_num + 1, len(train_dataloader), avg_loss / (batch_num + 1)))
        print('Train Epoch: %d, Loss: %6f' % (epoch + 1, avg_loss / (len(train_dataloader))))
        best_loss = float('inf')
        best_acc = 0
        model.eval()
        with torch.no_grad():
            avg_loss = 0
            correct = 0
            all_pred = 0
            for batch in val_dataloader:
                input_tensor, target_tensor = batch['input_ids'].to(device).long(), batch['y'].to(device)
                y_pred = model(input_tensor)
                loss = loss_function(y_pred, target_tensor)
                avg_loss += loss.item()  # / len(train_loader)
                y_pred = torch.argmax(y_pred, dim=-1)
                correct += torch.sum(y_pred == target_tensor).item()
                all_pred += y_pred.size(0)
            acc = correct / all_pred
            print('Val Epoch: %d, Loss: %6f, Acc: %4f' % (epoch + 1, avg_loss / (len(val_dataloader)), acc))
            if (avg_loss / (len(val_dataloader))) < best_loss or acc > best_acc:
                best_loss = avg_loss / (len(val_dataloader))
                best_acc = acc
                result = {'model': model.state_dict(),
                          'loss': avg_loss / (len(val_dataloader)),
                          'acc': acc,
                          'model_args': (model.rnn_size, model.word_vec_size, model.embedding_size,
                                         model.dense_hidden_dim, model.dropoutProb, model.num_layers,
                                         model.target_size)}
                torch.save(result, '%s%s_ckpt_best.pth' % (out_path, args.dataset_name))