import os
import random

import numpy as np
import pandas as pd
import torch
import torch.optim as optim
import tqdm
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import train_test_split
from torch import nn
from torch.optim.adam import Adam
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler

from Classes.Data import LeukemiaDataset, augmentation
from Classes.Helpers import Helpers
from Classes.interpretability import interpret_model
from Classes.Model_2020 import LuekemiaNet, train_model
from Classes.model_api import (confusion_matrix2, get_predictions,
                               plot_training_history)

SEED = 323


def seed_everything(seed=SEED):
    random.seed(seed)
    os.environ['PYHTONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True


image_path = '/home/allen/Drive C/Peter Moss AML Leukemia Research/ALL-PyTorch-2020/Classifier/Model/Data/Test/Im047_0.jpg'
image_path2 = '/home/allen/Drive C/Peter Moss AML Leukemia Research/ALL-PyTorch-2020/Classifier/Model/Data/Test/Im006_1.jpg'
label_idx = '/home/allen/Drive C/Peter Moss AML Leukemia Research/Dataset/pred_label.json'

seed_everything(SEED)
# helper class
helper = Helpers("Test Model", False)
# train data directory
train_dir = '/home/allen/Drive C/Peter Moss AML Leukemia Research/Dataset/all_train/'
# train label directoy
train_csv = '/home/allen/Drive C/Peter Moss AML Leukemia Research/Dataset/train.csv'
# labels
class_name = ["zero", "one"]

# training batch size
batch_size = helper.config["classifier"]["train"]["batch"]
# accuracy and loss save directory
acc_loss_png = helper.config["classifier"]["model_params"]["plot_loss_and_acc"]
# confusion matrix save directory
confusion_png = helper.config["classifier"]["model_params"]["confusion_matrix"]
# number of epoch
epochs = helper.config["classifier"]["train"]["epochs"]
# learning rate
learn_rate = helper.config["classifier"]["train"]["learning_rate_adam"]
# decay
decay = helper.config["classifier"]["train"]["decay_adam"]
# read train CSV file
labels = pd.read_csv(train_csv)
# print label count
labels_count = labels.label.value_counts()
print(labels_count)
# print 5 label header
print(labels.head())
# splitting data into training and validation set
train, valid = train_test_split(labels, stratify=labels.label, test_size=0.2, shuffle=True)
print(len(train), len(valid))

# data augmentation
training_transforms, validation_transforms = augmentation()

# Read Acute Lymphoblastic Leukemia dataset from disk
trainset = LeukemiaDataset(df_data=train, data_dir=train_dir, transform=training_transforms)
validset = LeukemiaDataset(df_data=valid, data_dir=train_dir, transform=validation_transforms)

train_size, valid_size = len(trainset), len(validset)
print(train_size, valid_size)

train_sampler = SubsetRandomSampler(list(train.index))
valid_sampler = SubsetRandomSampler(list(valid.index))

# Prepare dataset for neural networks
train_data_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
valid_data_loader = DataLoader(validset, batch_size=batch_size, shuffle=False)

# Checking the dataset
print('Training Set:\n')
for images, labels in train_data_loader:
    print('Image batch dimensions:', images.size())
    print('Image label dimensions:', labels.size())
    break
print("\n")

#
for images, labels in valid_data_loader:
    print("The labels: ", labels)

# Define model
model = LuekemiaNet()
# check if CUDA is available, else use CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Move our model to available hardware
model.to(device)

# Cross entropy loss function
criterion = nn.CrossEntropyLoss()
# specify optimizer (stochastic gradient descent) and learning rate = 0.001
optimizer = Adam(params=model.parameters(), lr=learn_rate, weight_decay=decay)
# scheduler = CyclicLR(optimizer, base_lr=lr, max_lr=0.01, step_size=5, mode='triangular2')
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7,
                                      gamma=helper.config["classifier"]["model_params"]["gamma"])

# begin training
real_model, history = train_model(model, train_data_loader,
                                  valid_data_loader,
                                  optimizer, scheduler,
                                  criterion, train_size,
                                  valid_size,
                                  device=device, n_epochs=epochs)
# plot model loss and accuracy
plot_training_history(history, save_path=acc_loss_png)
# Get model prediction
y_pred, y_test = get_predictions(real_model, valid_data_loader, device)
# Get model precision, recall and f1_score
helper.logger.info(classification_report(y_test, y_pred, target_names=class_name))
# Get model confusion matrix
cm = confusion_matrix(y_test, y_pred)
confusion_matrix2(cm, class_name, save_path=confusion_png)

interpret_model(real_model, validation_transforms, image_path, label_idx, use_cpu=True,
                interpret_type="integrated gradients")
interpret_model(real_model, validation_transforms, image_path, label_idx, use_cpu=True, interpret_type="gradient shap")

interpret_model(real_model, validation_transforms, image_path2, label_idx, use_cpu=True,
                interpret_type="integrated gradients")
interpret_model(real_model, validation_transforms, image_path2, label_idx, use_cpu=True, interpret_type="gradient shap")