"""
training Recurrent Hierarchical model.
Sheng Liu
All rights reserved
Report bugs to shengliu@nyu.edu
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch.optim as optim
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import random
from training_tools import *
from model import *
import logging
#%matplotlib inline

def data_split(patient_list, downsampling_rate = 0.3):
	# Argments:
	#       patient_list : in format [[patient_id, label]] where death label = 1 otherwise label = 0
	#       downsampling_rate : the percentage of negative sample in the final training data
	# Output:
	#        train_data, validation data and test_data
	
	while True:
		
		np.random.rand(patient_list.shape[0]).argsort()
		np.take(patient_list,np.random.rand(patient_list.shape[0]).argsort(),axis=0,out=patient_list)
		patient_list = np.array(patient_list)
		num_patients = len(patient_list[:,0])
		train_data = patient_list[:int(0.7*num_patients)]
		val_data   = patient_list[int(0.7*num_patients):int(0.8*num_patients)]
		test_data  = patient_list[int(0.8*num_patients):num_patients]
		
		#downsampling
		if sum(train_data[:,1]) > 10:
			break
		else:
			np.random.rand(patient_list.shape[0]).argsort()
			np.take(patient_list,np.random.rand(patient_list.shape[0]).argsort(),axis=0,out=patient_list)
	
	if sum(train_data[:,1])/len(train_data[:,1]) <= 0.3:

		downsampling_size = int(sum(train_data[:,1])*(1 - downsampling_rate)/downsampling_rate)

		train_data_survive = train_data[train_data[:,1] != 1][:downsampling_size]
		train_data_dead = train_data[train_data[:,1] == 1]
		train_data = np.vstack((train_data_survive,train_data_dead))
		#random.shuffle(train_data)
		print('The percentage of negative sample after downsampling is {:.1%}'.format(sum(train_data[:,1])/len(train_data[:,1])))
		return train_data[:,0], val_data[:,0], test_data[:,0]

	else:
		print('The percentage of negative sample is {:.1%}'.format(sum(train_data[:,1])/len(train_data[:,1])))
		return train_data[:,0], val_data[:,0], test_data[:,0]

batch_size = 16
num_tokens = 181444
embed_size = 300
word_gru_hidden = 100
sent_gru_hidden = 100
note_gru_hidden = 100
number_of_classes = 2
num_epoch = 40000
print_val_loss_every = 100
print_loss_every = 50
learning_rate = 0.1

word_attn = AttentionWordRNN(batch_size=batch_size, num_tokens=num_tokens, embed_size=embed_size, 
							 word_gru_hidden=word_gru_hidden, bidirectional= True)

sent_attn = AttentionSentRNN(batch_size=batch_size, sent_gru_hidden= sent_gru_hidden, word_gru_hidden=word_gru_hidden, 
							 n_classes=number_of_classes, bidirectional= True)

word_attn.cuda()
sent_attn.cuda()

logger_name = "mortality_prediction"
logger = logging.getLogger(logger_name)
logger.setLevel(logging.INFO)

# file handler
fh = logging.FileHandler('./' + 'Mortality_prediction')
fh.setLevel(logging.INFO)
logger.addHandler(fh)

# stream handler
console = logging.StreamHandler()
console.setLevel(logging.INFO)
logger.addHandler(console)

#    torch.cuda.set_device(args.gpu_id)

#for arg in vars(args):
#    logger.info(str(arg) + ' ' + str(getattr(args, arg)))

# load train/dev/test data
# train data
logger.info('loading data...')
train = np.load('./train_15m.npy')
val = np.load('./val_15m.npy')
test = np.load('./test_15m.npy')
X_train = train.item()['DATA']
y_train = train.item()['MORTALITY_LABEL']
X_val = val.item()['DATA']
y_val = val.item()['MORTALITY_LABEL']
X_test = test.item()['DATA']
y_test = test.item()['MORTALITY_LABEL']
#y = np.load('../mordality_pred/label.npy')
#y = y.item()['MORTALITY_LABEL']
#patient_index = np.vstack((range(len(y)), y)).T.astype(int)
#train, val, test = data_split(patient_index, downsampling_rate = 0.3)
#X_train = X[train]
#y_train = y[train]
#X_val   = X[val]
#y_val   = y[val]
#X_test  = X[test]
#y_test  = y[test]
logger.info('train size # sent ' + str(len(X_train)))
logger.info('dev size # sent ' + str(len(X_val)))
logger.info('test size # sent ' + str(len(X_test)))

logger.info('train size # sent ' + str(len(X_train)))
logger.info('dev size # sent ' + str(len(X_val)))
logger.info('test size # sent ' + str(len(X_test)))
logger.info('batch_size: ' +str(batch_size))
logger.info('num_tokens: ' + str(num_tokens))
logger.info('embed_size: ' + str(embed_size))
logger.info('word_gru_hidden: ' + str(word_gru_hidden))
logger.info('sent_gru_hidden: ' + str(sent_gru_hidden))
logger.info('note_gru_hidden: ' + str(note_gru_hidden))
logger.info('number_of_classes: ' + str(number_of_classes))
logger.info('num_epoch: ' +str(num_epoch))
logger.info('print_val_loss_every' + str(print_val_loss_every))
logger.info('print_loss_every' + str(print_loss_every))
logger.info('Optimizer learning rate is ' + str(learning_rate))

momentum = 0.9
word_optmizer = torch.optim.SGD(word_attn.parameters(), lr=learning_rate, momentum= momentum)
sent_optimizer = torch.optim.SGD(sent_attn.parameters(), lr=learning_rate, momentum= momentum)
criterion = nn.NLLLoss()

logger.info('Start to train...')

loss_full= train_early_stopping(batch_size, X_train, y_train, X_test, y_test, word_attn, sent_attn, word_optmizer, sent_optimizer, 
							criterion, num_epoch, print_val_loss_every, print_loss_every)
logger.info('Training end!')
logger.info('Test!')
test_acc,test_auc = test_accuracy_full_batch(X_test, y_test, 64, model)
logger.info('Test accuracy is %.2f' % test_acc)
logger.info('Test auc is %.2f' % test_auc)
np.save('loss_full','loss_full')
