from tqdm import tqdm
import pandas as pd
import numpy as np, argparse, time, pickle, random, os, datetime

import torch

import torch.optim as optim
from model import MaskedNLLLoss, BC_LSTM
from dataloader import MELDDataLoader


from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report

def setup_seed(seed):
    """ Manually Fix the random seed to get deterministic results.
    """
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.benchmark = False
    torch.backends.cudnn.deterministic = True

def train_or_eval_model(model, loss_function, dataloader, epoch, optimizer=None, mode='train'):
    losses, preds, labels, masks, losses_sense  = [], [], [], [], []
    
    max_sequence_len = []
    assert mode != 'train' or optimizer != None
    if mode == 'train':
        model.train()
    else:
        model.eval()
    
    with tqdm(dataloader) as td:
        for data in td:

            if mode == 'train':
                optimizer.zero_grad()
                
            textf, acouf, mask, label = [d.cuda() for d in data[:-1]] if args.cuda else data[:-1]
            
            
            log_prob, _ = model(textf, None, acouf, None, mask)

            lp_ = log_prob.transpose(0,1).contiguous().view(-1, log_prob.size()[2]) # batch*seq_len, n_classes
            labels_ = label.view(-1) # batch*seq_len
            loss = loss_function(lp_, labels_, mask)

            pred_ = torch.argmax(lp_,1) # batch*seq_len
            preds.append(pred_.data.cpu().numpy())
            labels.append(labels_.data.cpu().numpy())
            masks.append(mask.view(-1).cpu().numpy())
            losses.append(loss.item()*masks[-1].sum())

            if mode == 'train':
                total_loss = loss
                total_loss.backward()
                
                optimizer.step()

    if preds!=[]:
        preds  = np.concatenate(preds)
        labels = np.concatenate(labels)
        masks  = np.concatenate(masks)
    else:
        return float('nan'), float('nan'), float('nan'), [], [], [], float('nan'),[]

    avg_loss = round(np.sum(losses)/np.sum(masks), 4)
    avg_sense_loss = round(np.sum(losses_sense)/np.sum(masks), 4)

    avg_accuracy = round(accuracy_score(labels,preds, sample_weight=masks)*100, 2)
    avg_fscore = round(f1_score(labels,preds, sample_weight=masks, average='weighted')*100, 2)
    if mode == 'test':
        class_report = classification_report(labels, preds, sample_weight=masks, target_names=['neutral', 'surprise', 'fear', 'sadness', 'joy', 'disgust', 'anger'], digits=6)
        print(class_report)
    return avg_loss, avg_accuracy, labels, preds, masks, [avg_fscore]


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_workers', type=int, default=0,
                        help='num workers of loading data')
    
    # dataloader settings 
    parser.add_argument('--batch-size', type=int, default=32, metavar='BS', help='batch size')
    parser.add_argument('--data_path', type=str, default='../TextCnn/dataset/MELD_features_raw.pkl')

    # model settings.
    parser.add_argument('--attention_type', type=str, default='general2')
    parser.add_argument('--utterance_dim', type=int, default=600,
                        help='embedding dims to use')
    parser.add_argument('--emotion_state_dim', type=int, default=100)
    parser.add_argument('--hidden_layer_dim', type=int, default=100)
    parser.add_argument('--dropout', type=float, default=0.25)
    parser.add_argument('--n_classes', type=int, default=7)
    # late fusion module.
    parser.add_argument('--lateFusionModule', type=str, default='concat')
    parser.add_argument('--input_features', type=tuple, default=(100, 300))
    parser.add_argument('--pre_fusion_hidden_dims', type=tuple, default=(24, 7))
    parser.add_argument('--pre_fusion_dropout', type=float, default=0.4)
    parser.add_argument('--post_fusion_dropout', type=float, default=0.3)

    # train settings.
    parser.add_argument('--lr', type=float, default=1e-4, metavar='LR', help='learning rate')
    parser.add_argument('--l2', type=float, default=1e-5, metavar='L2', help='L2 regularization weight')
    parser.add_argument('--epochs', type=int, default=100, metavar='E', help='number of epochs')

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()

    args.cuda = torch.cuda.is_available()
    if args.cuda:
        print('Running on GPU')
    else:
        print('Running on CPU')

    for seed in [1, 11, 111, 1111, 11111]:
        setup_seed(seed)
        args.seed = seed
        
        print(args)

        model = BC_LSTM(args)
        print('MELD BC_LSTM MODULE ...')

        if args.cuda:
            model.cuda()
        
        loss_weights = torch.FloatTensor([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
        
        loss_function  = MaskedNLLLoss(loss_weights.cuda() if args.cuda else loss_weights)
        
        optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2)
        lf = open('logs/cnn_meld_logs.txt', 'a')
        
        dataloader = MELDDataLoader(args) 

        valid_losses, valid_fscores = [], []
        test_fscores, test_accuracys, test_losses = [], [], []
        best_loss, best_label, best_pred, best_mask = None, None, None, None

        for e in range(args.epochs):
            start_time = time.time()
            train_loss, train_acc, _, _, _, train_fscore = train_or_eval_model(model, loss_function, dataloader['train'], e, optimizer, mode='train')
            valid_loss, valid_acc, _, _, _, valid_fscore = train_or_eval_model(model, loss_function, dataloader['valid'], e, mode='valid')
            test_loss, test_acc, test_label, test_pred, test_mask, test_fscore = train_or_eval_model(model, loss_function, dataloader['test'], e, mode='test')
                
            valid_losses.append(valid_loss)
            valid_fscores.append(valid_fscore)
            test_losses.append(test_loss)
            test_accuracys.append(test_acc)
            test_fscores.append(test_fscore)
                
            x = 'epoch: {}, train_loss: {}, acc: {}, fscore: {}, valid_loss: {}, acc: {}, fscore: {}, test_loss: {}, acc: {}, fscore: {}, time: {} sec'.format(e+1, train_loss, train_acc, train_fscore, valid_loss, valid_acc, valid_fscore, test_loss, test_acc, test_fscore, round(time.time()-start_time, 2))
            
            print (x)
            lf.write(x + '\n')

        valid_fscores = np.array(valid_fscores).transpose()
        test_fscores = np.array(test_fscores).transpose() # [1, epoches]
        test_accuracys = np.array(test_accuracys).transpose() # [epoches]

        f1_score1 = test_fscores[0][np.argmin(valid_losses)]
        acc_score1 = test_accuracys[np.argmin(valid_losses)]
        f1_score2 = test_fscores[0][np.argmax(valid_fscores[0])]
        acc_score2 = test_accuracys[np.argmax(valid_fscores[0])]
        scores = [acc_score1, f1_score1, acc_score2, f1_score2]
        scores = [str(item) for item in scores]

        print ('Test Scores: Weighted F1')
        print('@Best Valid Loss: Test Acc: {}, Test F1: {}'.format(acc_score1, f1_score1))
        print('@Best Valid F1: Test Acc: {}, Test F1: {}'.format(acc_score2, f1_score2))

        rf = open('results/cnn_meld_results.txt', 'a')
        rf.write('\t'.join(scores) + '\t' + str(args) + '\n')
        rf.close()

[((17, 4, 17, 27), 'torch.manual_seed', 'torch.manual_seed', ({(17, 22, 17, 26): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((18, 4, 18, 32), 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', ({(18, 27, 18, 31): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((19, 4, 19, 36), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(19, 31, 19, 35): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((20, 4, 20, 24), 'numpy.random.seed', 'np.random.seed', ({(20, 19, 20, 23): 'seed'}, {}), '(seed)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((21, 4, 21, 21), 'random.seed', 'random.seed', ({(21, 16, 21, 20): 'seed'}, {}), '(seed)', False, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((81, 13, 81, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((114, 16, 114, 41), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((35, 9, 35, 25), 'tqdm.tqdm', 'tqdm', ({(35, 14, 35, 24): 'dataloader'}, {}), '(dataloader)', False, 'from tqdm import tqdm\n'), ((63, 17, 63, 38), 'numpy.concatenate', 'np.concatenate', ({(63, 32, 63, 37): 'preds'}, {}), '(preds)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((64, 17, 64, 39), 'numpy.concatenate', 'np.concatenate', ({(64, 32, 64, 38): 'labels'}, {}), '(labels)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((65, 17, 65, 38), 'numpy.concatenate', 'np.concatenate', ({(65, 32, 65, 37): 'masks'}, {}), '(masks)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((75, 23, 75, 174), 'sklearn.metrics.classification_report', 'classification_report', (), '', False, 'from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report\n'), ((126, 16, 126, 29), 'model.BC_LSTM', 'BC_LSTM', ({(126, 24, 126, 28): 'args'}, {}), '(args)', False, 'from model import MaskedNLLLoss, BC_LSTM\n'), ((132, 23, 132, 77), 'torch.FloatTensor', 'torch.FloatTensor', ({(132, 41, 132, 76): '[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]'}, {}), '([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])', False, 'import torch\n'), ((139, 21, 139, 41), 'dataloader.MELDDataLoader', 'MELDDataLoader', ({(139, 36, 139, 40): 'args'}, {}), '(args)', False, 'from dataloader import MELDDataLoader\n'), ((50, 20, 50, 39), 'torch.argmax', 'torch.argmax', ({(50, 33, 50, 36): 'lp_', (50, 37, 50, 38): '1'}, {}), '(lp_, 1)', False, 'import torch\n'), ((69, 21, 69, 35), 'numpy.sum', 'np.sum', ({(69, 28, 69, 34): 'losses'}, {}), '(losses)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((69, 36, 69, 49), 'numpy.sum', 'np.sum', ({(69, 43, 69, 48): 'masks'}, {}), '(masks)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((70, 27, 70, 47), 'numpy.sum', 'np.sum', ({(70, 34, 70, 46): 'losses_sense'}, {}), '(losses_sense)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((70, 48, 70, 61), 'numpy.sum', 'np.sum', ({(70, 55, 70, 60): 'masks'}, {}), '(masks)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((72, 25, 72, 74), 'sklearn.metrics.accuracy_score', 'accuracy_score', (), '', False, 'from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report\n'), ((73, 23, 73, 86), 'sklearn.metrics.f1_score', 'f1_score', (), '', False, 'from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report\n'), ((146, 25, 146, 36), 'time.time', 'time.time', ({}, {}), '()', False, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((162, 24, 162, 47), 'numpy.array', 'np.array', ({(162, 33, 162, 46): 'valid_fscores'}, {}), '(valid_fscores)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((163, 23, 163, 45), 'numpy.array', 'np.array', ({(163, 32, 163, 44): 'test_fscores'}, {}), '(test_fscores)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((164, 25, 164, 49), 'numpy.array', 'np.array', ({(164, 34, 164, 48): 'test_accuracys'}, {}), '(test_accuracys)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((166, 36, 166, 59), 'numpy.argmin', 'np.argmin', ({(166, 46, 166, 58): 'valid_losses'}, {}), '(valid_losses)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((167, 36, 167, 59), 'numpy.argmin', 'np.argmin', ({(167, 46, 167, 58): 'valid_losses'}, {}), '(valid_losses)', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((168, 36, 168, 63), 'numpy.argmax', 'np.argmax', ({(168, 46, 168, 62): 'valid_fscores[0]'}, {}), '(valid_fscores[0])', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((169, 36, 169, 63), 'numpy.argmax', 'np.argmax', ({(169, 46, 169, 62): 'valid_fscores[0]'}, {}), '(valid_fscores[0])', True, 'import numpy as np, argparse, time, pickle, random, os, datetime\n'), ((157, 278, 157, 289), 'time.time', 'time.time', ({}, {}), '()', False, 'import numpy as np, argparse, time, pickle, random, os, datetime\n')]