import numpy as np
import torch
import matplotlib.pyplot as plt
import pickle
import random
import mlflow

from dataset import get_dataset, get_handler
from models import get_net
from train import Train
from torchvision import transforms
from datetime import datetime

from sklearn.metrics import precision_recall_fscore_support, confusion_matrix, accuracy_score, f1_score, precision_score, recall_score, fbeta_score, roc_curve, auc, roc_auc_score

seed = 123
dataset_name = 'WILDCAM'
#dataset_path = '/datapool/wildcam/wildcam_subset_sample'
dataset_path = '/datapool/wildcam/wildcam_subset_denoised'

args_pool = {
    'WILDCAM': {
        'n_restarts': 1,
        'steps': 121,
        'n_classes': 2,
        'fc_only': True,
        'model_path': "./models/",
        'transform': {
            'train': transforms.Compose([
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])]),
            'test': transforms.Compose([
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])
        },
        'loader_tr_args': {
            'batch_size': 100, 
            'num_workers': 1
        },
        'loader_te_args': {
            'batch_size': 100, 
            'num_workers': 1
        },
        'loader_sample_args': {
            'batch_size': 100, 
            'num_workers': 1
        },
        'optimizer_args': {
            'lr': 0.001,
            'l2_regularizer_weight': 0.001,
            'penalty_anneal_iters': 0, # make this 0 for ERM
            'penalty_weight': 0.0 # make this 0 for ERM
        }
    }
}

args = args_pool[dataset_name]
model_name = "wildcam_denoised_" + str(args['steps']) + "_" + str(args['optimizer_args']['lr']) + "_" + str(args['optimizer_args']['penalty_anneal_iters']) + "_" + str(args['optimizer_args']['penalty_weight']) + "_"

np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# load dataset
envs, x_test, y_test = get_dataset(dataset_name, dataset_path)
print("\nworking with", len(envs), "training environments: ")
for env in envs:
    print("env['images']: ", len(env['images']))
    print("env['labels']: ", env['labels'].shape[0])
    unique, counts = np.unique(env['labels'], return_counts=True)
    if counts[0] >= counts[1]:
        baseline_acc = counts[0]/(counts[0]+counts[1])
        print_statement = str(round(baseline_acc, 2)) + " (always coyote)."
    else:
        baseline_acc = counts[1]/(counts[0]+counts[1])
        print_statement = str(round(baseline_acc, 2)) + " (always raccoon)."        
    print("class distribution: ", counts[0], "coyotes and", counts[1], "raccoons. baseline accuracy", print_statement)

print("x_test: ", len(x_test))
print("y_test: ", y_test.shape[0])
unique, counts = np.unique(y_test, return_counts=True)
if counts[0] >= counts[1]:
    baseline_acc = counts[0]/(counts[0]+counts[1])
    print_statement = str(round(baseline_acc, 2)) + " (always coyote)."
else:
    baseline_acc = counts[1]/(counts[0]+counts[1])
    print_statement = str(round(baseline_acc, 2)) + " (always raccoon)."    
print("test class distribution: ", counts[0], "coyotes and", counts[1], "raccoons. baseline accuracy", print_statement)

# get model and data handler
net = get_net(dataset_name)
handler = get_handler(dataset_name)

# GPU enabled?
print("Using GPU - {}".format(torch.cuda.is_available()))

final_train_accs = []
final_test_accs = []
train_process = Train(envs, x_test, y_test, net, handler, args)



start = datetime.now()
with mlflow.start_run():
    
    mlflow.log_params({
        'dataset': dataset_path,
        'steps': args['steps'],
        'lr': args['optimizer_args']['lr'],
        'penalty_anneal_iters': args['optimizer_args']['penalty_anneal_iters'],
        'penalty_weight': args['optimizer_args']['penalty_weight'],
        'l2_regularizer_weight': args['optimizer_args']['l2_regularizer_weight'],
        'train_batch_size': args['loader_tr_args']['batch_size']
    })
    
    print()
    if args['optimizer_args']['penalty_weight'] > 1.0:
        mlflow.log_param('method', 'IRM')
        model_name = model_name + "IRM"
        print("========================================IRM========================================")
    else:
        mlflow.log_param('method', 'ERM')
        model_name = model_name + "ERM"
        print("========================================ERM========================================")
    print(args)

    for restart in range(args['n_restarts']):  
        print("\nRestart", restart)
    
        with mlflow.start_run(nested=True):
            train_acc, test_acc, preds, probs = train_process.train()
            preds = torch.reshape(preds, (-1,)).long().detach().cpu().numpy()
            test_prec = precision_score(y_test, preds)
            test_rec = recall_score(y_test, preds)
            fpr, tpr, thresholds = roc_curve(y_test, probs)
            roc_auc = roc_auc_score(y_test, probs)

    final_train_accs.append(train_acc)
    final_test_accs.append(test_acc)

    mlflow.log_metrics({
        'final_train_acc_mean': np.mean(final_train_accs),
        'final_train_acc_std': np.std(final_train_accs),
        'final_test_acc_mean': np.mean(final_test_accs),
        'final_test_acc_std': np.std(final_test_accs),
    })

    print('Final train acc (mean/std across restarts so far):')
    print(round(np.mean(final_train_accs), 4), round(np.std(final_train_accs), 4))
    print('Final test acc (mean/std across restarts so far):')
    print(round(np.mean(final_test_accs), 4), round(np.std(final_test_accs), 4))
    print('Final test precision:')
    print(round(test_prec, 4))
    print('Final test recall:')
    print(round(test_rec, 4))
    
    #y = np.where(y_test.numpy() == 0, 'coyote', 'raccoon')
    #y_hat = np.where(preds == 0, 'coyote', 'raccoon')
    print('confusion matrix:')
    print(confusion_matrix(y_test.numpy(), preds))
    tn, fp, fn, tp = confusion_matrix(y_test.numpy(), preds).ravel() #, labels=["coyote", "raccoon"]).ravel()
    print(f'tn = {tn}, fp = {fp}, fn = {fn}, tp = {tp}')
        
    plt.figure(constrained_layout=True, figsize=(3, 3))

    # Plot ROC curve
    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], 'k--')  
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.0])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC')
    plt.legend(loc="lower right")
    outfile = './figures/' + model_name + '.png'
    plt.savefig(outfile, dpi=300, bbox_inches='tight', pad_inches=0) # To save figure
    plt.show() # To show figure
    
end = datetime.now()

time_elapsed = end - start
print("time for training: ", time_elapsed.days, time_elapsed.min, "minutes and", time_elapsed.seconds, "seconds.")
torch.save(train_process.clf.state_dict(), args['model_path'] + model_name + ".pth")
