import matplotlib.pyplot as plt

import pytorch_lightning as pl

from pytorch_lightning.callbacks import EarlyStopping

from pytorch_lightning.metrics.functional import accuracy

from pytorch_lightning.callbacks import LearningRateMonitor

from torch.optim.lr_scheduler import OneCycleLR

import torch

import torch.nn as nn

import os

import numpy as np

import pandas as pd

from torch.utils.data import DataLoader

from torch.utils.data.sampler import SubsetRandomSampler

import os

import torch.nn.functional as F

from torch.utils.data import Dataset

from torchvision import transforms, datasets

from torch.optim import Adam

import random

from sklearn.model_selection import train_test_split

from sklearn.metrics import confusion_matrix, classification_report

import seaborn as sns

import PIL.Image as Image

import json

import numpy as np

from matplotlib.colors import LinearSegmentedColormap

from captum.attr import IntegratedGradients

from captum.attr import GradientShap

from captum.attr import NoiseTunnel

from captum.attr import Saliency

from captum.attr import visualization as viz

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

#decay

decay = 5e-4

# training batch size

batch_size = 10

# check if cuda is available: if not available, then use cpu

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class LeukemiaDataset(Dataset):

    """

    Acute Lymphoblastic Leukemia Dataset Reader.

     Args:

           df_data: Dataframe for CSV file

           data_dir: path to Lymphoblastic Leukemia Data

           transform: transforms for performing data augmentation

    """

    def __init__(self, df_data, data_dir='./', transform=None):

        super().__init__()

        self.df = df_data.values

        self.data_dir = data_dir

        self.transform = transform

    def __len__(self):

        return len(self.df)

    def __getitem__(self, index):

        img_name, label = self.df[index]

        img_path = os.path.join(self.data_dir, img_name + '.jpg')

        image = Image.open(img_path)

        if self.transform is not None:

            image = self.transform(image)

        return image, label

def augmentation():

    """Acute Lymphoblastic Leukemia data augmentation"""

    mean, std_dev = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

    training_transforms = transforms.Compose([transforms.Resize((100, 100)),

                                              transforms.RandomRotation(30),

                                              transforms.RandomResizedCrop(100),

                                              transforms.RandomHorizontalFlip(),

                                              transforms.RandomGrayscale(p=0.1),

                                              transforms.ToTensor(),

                                              transforms.Normalize(mean=mean, std=std_dev)])

    validation_transforms = transforms.Compose([

        transforms.Resize((100, 100)),

        transforms.ToTensor(),

        transforms.Normalize(mean=mean, std=std_dev)])

    return training_transforms, validation_transforms

def predict_probability(model, transforms, image_path, use_gpu=True):

    img = Image.open(image_path)

    img = img.convert('RGB')

    img = transforms(img).unsqueeze(0)

    if use_gpu:

        image = img.cuda()

    predictions = model(image)

    predictions = torch.sigmoid(predictions)

    predictions = predictions.detach().cpu().numpy().flatten()

    return predictions

def show_prediction_confidence(prediction, class_names):

    pred_df = pd.DataFrame({

        'class_names': class_names,

        'values': prediction

    })

    sns.barplot(x='values', y='class_names', data=pred_df, orient='h')

    sns

    plt.xlim([0, 1])

def get_predictions(model, data_loader, use_gpu=True):

    model = model.eval()

    y_predictions = list()

    y_true = list()

    with torch.no_grad():

        for i, dataset in enumerate(data_loader):

            inputs, labels = dataset

            inputs, labels = inputs.to(device), labels.to(device)

        outputs = model(inputs)

        _, preds = torch.max(outputs, 1)

        y_predictions.extend(preds)

        y_true.extend(labels)

    predictions = torch.as_tensor(y_predictions).cpu()

    y_true = torch.as_tensor(y_true).cpu()

    return predictions, y_true

def confusion_matrix2(confusion_matrix, class_names, save_path):

    cm = confusion_matrix.copy()

    cell_counts = cm.flatten()

    cm_row_norm = cm / cm.sum(axis=1)[:, np.newaxis]

    row_percentages = ["{0:.2f}".format(value) for value in cm_row_norm.flatten()]

    cell_labels = [f"{cnt}\n{per}" for cnt, per in zip(cell_counts, row_percentages)]

    cell_labels = np.asarray(cell_labels).reshape(cm.shape[0], cm.shape[1])

    df_cm = pd.DataFrame(cm_row_norm, index=class_names, columns=class_names)

    hmap = sns.heatmap(df_cm, annot=cell_labels, fmt="", cmap="Blues")

    hmap.yaxis.set_ticklabels(hmap.yaxis.get_ticklabels(), rotation=0, ha='right')

    hmap.xaxis.set_ticklabels(hmap.xaxis.get_ticklabels(), rotation=30, ha='right')

    plt.ylabel('True diagnostic')

    plt.xlabel('Predicted diagnostic')

    plt.savefig(save_path)

    plt.show()

color = [(0, '#ffffff'), (0.25, '#000000'), (1, '#000000')]

name = 'custom blue'

N = 256

def linear_seg_color_map(name, color, N, gamma=1.0):

    """

    Render color map based on lookup tables

    :param name: name the color

    :param color: color code

    :param N: number of RGB quantization

    :param gamma: default is 1.0

    :return:

    """

    default_cmap = LinearSegmentedColormap.from_list(name, color, N, gamma)

    return default_cmap

def predict(model, transforms, image_path, use_cpu):

    """

    :param model:

    :param transforms: data transform

    :param image_path: inference input image path

    :param use_cpu:

    :return:

    """

    model.cpu()

    model = model.eval()

    img = Image.open(image_path)

    img = img.convert('RGB')

    transformed_img = transforms(img)

    image = transformed_img

    image = image.unsqueeze(0)

    if use_cpu:

        image = image.cpu()

    elif use_cpu == False:

        model.cuda()

        image = image.cuda

    output = model(image)

    output = torch.softmax(output, dim=1)

    prediction_score, pred_label_idx = torch.topk(output, 1)

    return image, transformed_img, prediction_score, pred_label_idx

def interpret_model(model, transforms, image_path, label_path, use_cpu=True, interpret_type=""):

    """

    :param model: our model

    :param transforms: Data transformation

    :param image_path: Image directory

    :param label_path: Json label directory

    :param use_gpu: set gpu to True

    :param interpret_type: mode for model interpretability: "integrated gradients"

    for Integrated Gradients, "gradient shap" for Gradient Shap and "occlusion" for Occlusion

    :return:

    """

    with open(label_path) as json_data:

        idx_to_labels = json.load(json_data)

    # Check if mode is Integrated Gradients

    if interpret_type == "integrated gradients":

        print('Performing Integrated Gradients Model Interpretation', interpret_type)

        image, transformed_img, prediction_score, pred_label_idx = predict(model, transforms, image_path, use_cpu)

        pred_label_idx.squeeze_()

        predicted_label = idx_to_labels[str(pred_label_idx.item())][1]

        print('Predicted:', predicted_label, '(', prediction_score.squeeze().item(), ')')

        integrated_gradients = IntegratedGradients(model)

        attributions_ig = integrated_gradients.attribute(image, target=pred_label_idx, n_steps=20)

        _ = viz.visualize_image_attr_multiple(np.transpose(attributions_ig.squeeze().cpu().detach().numpy(), (1, 2, 0)),

                                              np.transpose(transformed_img.squeeze().cpu().detach().numpy(), (1, 2, 0)),

                                              ["original_image", "heat_map"],

                                              ["all", "absolute_value"],

                                              cmap=linear_seg_color_map(name, color, N),

                                              show_colorbar=True)

    # Check if mode is Integrated Gradients with Noise Tunnel

    elif interpret_type == "integrated gradient noise":

        print('Performing Integrated Gradients Noise Tunnel Model Interpretation', interpret_type)

        image, transformed_img, prediction_score, pred_label_idx = predict(model, transforms, image_path, use_cpu)

        pred_label_idx.squeeze_()

        predicted_label = idx_to_labels[str(pred_label_idx.item())][1]

        print('Predicted:', predicted_label, '(', prediction_score.squeeze().item(), ')')

        integrated_gradients = IntegratedGradients(model)

        noise_tunnel = NoiseTunnel(integrated_gradients)

        attributions_ig_nt = noise_tunnel.attribute(image, n_samples=10, nt_type='smoothgrad_sq', target=pred_label_idx)

        _ = viz.visualize_image_attr_multiple(

            np.transpose(attributions_ig_nt.squeeze().cpu().detach().numpy(), (1, 2, 0)),

            np.transpose(transformed_img.squeeze().cpu().detach().numpy(), (1, 2, 0)),

            ["original_image", "heat_map"],

            ["all", "positive"],

            cmap=linear_seg_color_map(name, color, N),

            show_colorbar=True)

    # Check if mode is Gradient Shap

    elif interpret_type == "gradient shap":

        print('Performing Gradient Shap Model Interpretation', interpret_type)

        image, transformed_img, prediction_score, pred_label_idx = predict(model,transforms, image_path, use_cpu)

        pred_label_idx.squeeze_()

        predicted_label = idx_to_labels[str(pred_label_idx.item())][1]

        print('Predicted:', predicted_label, '(', prediction_score.squeeze().item(), ')')

        gradient_shap = GradientShap(model)

        # Defining baseline distribution of images

        rand_img_dist = torch.cat([image * 0, image * 1])

        attributions_gs = gradient_shap.attribute(image,

                                                  n_samples=50,

                                                  stdevs=0.0001,

                                                  baselines=rand_img_dist,

                                                  target=pred_label_idx)

        _ = viz.visualize_image_attr_multiple(np.transpose(attributions_gs.squeeze().cpu().detach().numpy(), (1, 2, 0)),

                                              np.transpose(transformed_img.squeeze().cpu().detach().numpy(), (1, 2, 0)),

                                              ["original_image", "heat_map"],

                                              ["all", "absolute_value"],

                                              cmap=linear_seg_color_map(name, color, N),

                                              show_colorbar=True)

        # Check if mode is Saliency

    elif interpret_type == "saliency":

        print('Performing Saliency Model Interpretation', interpret_type)

        image, transformed_img, prediction_score, pred_label_idx = predict(model, transforms, image_path, use_cpu)

        pred_label_idx.squeeze_()

        predicted_label = idx_to_labels[str(pred_label_idx.item())][1]

        print('Predicted:', predicted_label, '(', prediction_score.squeeze().item(), ')')

        saliency = Saliency(model)

        attributions_gs = saliency.attribute(image,target=pred_label_idx)

        _ = viz.visualize_image_attr_multiple(np.transpose(attributions_gs.squeeze().cpu().detach().numpy(), (1, 2, 0)),

                                              np.transpose(transformed_img.squeeze().cpu().detach().numpy(), (1, 2, 0)),

                                              ["original_image", "heat_map"],

                                              ["all", "absolute_value"],

                                              cmap=linear_seg_color_map(name, color, N),

                                              show_colorbar=True)

class LuekemiaNet(pl.LightningModule):

    def __init__(self, lr=0.01, weight_decay=5e-4):

        super(LuekemiaNet, self).__init__()

        self.save_hyperparameters()

        self.conv1 = nn.Sequential(

            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=0),

            nn.ReLU(),

            nn.BatchNorm2d(32),

            nn.Dropout(p=0.25),

            nn.AvgPool2d(2))

        self.conv2 = nn.Sequential(

            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=1),

            nn.ReLU(),

            nn.BatchNorm2d(64),

            nn.Dropout(p=0.25),

            nn.AvgPool2d(2))

        self.conv3 = nn.Sequential(

            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, stride=1, padding=1),

            nn.ReLU(),

            nn.BatchNorm2d(128),

            nn.Dropout(p=0.25),

            nn.AvgPool2d(2))

        self.fc = nn.Sequential(

            nn.Linear(128 * 10 * 10, 200),

            nn.ReLU(),

            nn.Dropout(),

            nn.Linear(200, 2))

    def forward(self, x):

        """Method for Forward Prop"""

        out = self.conv1(x)

        out = self.conv2(out)

        out = self.conv3(out)

        ######################

        # For model debugging #

        ######################

        #print(out.shape)

        out = out.view(x.shape[0], -1)

        out = self.fc(out)

        return out

    def configure_optimizers(self):

        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)

        return optimizer

    def training_step(self, batch, batch_idx):

        x, y = batch

        logits = self(x)

        loss = F.cross_entropy(logits, y)

        self.log('train_loss', loss)

        return loss

    def evaluate(self, batch, stage=None):

        x, y = batch

        logits = self(x)

        loss = F.cross_entropy(logits, y)

        preds = torch.argmax(logits, dim=1)

        acc = accuracy(preds, y)

        if stage:

            self.log(f'{stage}_loss', loss, prog_bar=True)

            self.log(f'{stage}_acc', acc, prog_bar=True)

    def validation_step(self, batch, batch_idx):

        self.evaluate(batch, 'val')

    def test_step(self, batch, batch_idx):

        self.evaluate(batch, 'test')

image_path = '/home/allen/Drive C/Peter Moss AML Leukemia Research/Dataset/all_test/Im041_0.jpg'

label_idx = '/home/allen/Drive C/Peter Moss AML Leukemia Research/ALL-PyTorch-2020/Classifier/Model/class_idx.json'

seed_everything(SEED)

# 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"]

# number of epoch

epochs = 20

# learning rate

learn_rate = 0.001

# 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.1, shuffle=True)

print(len(train),len(valid))

#data augmentation

training_transforms, validation_transforms = augmentation()

train_dataset = LeukemiaDataset(df_data=train, data_dir=train_dir, transform=training_transforms)

valid_dataset = LeukemiaDataset(df_data=valid, data_dir=train_dir, transform=validation_transforms)

train_sampler = SubsetRandomSampler(list(train.index))

valid_sampler = SubsetRandomSampler(list(valid.index))

# Prepare dataset for neural networks

train_data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

valid_data_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

model_path = "weight2.pth"

model = LuekemiaNet()

early_stop_callback = EarlyStopping(

    monitor='val_acc',

    min_delta=0.00,

    patience=3,

    verbose=False,

    mode='max'

)

trainer = pl.Trainer(max_epochs=epochs, log_every_n_steps=2, callbacks=[early_stop_callback])

trainer.fit(model, train_data_loader, valid_data_loader)

trainer.save_checkpoint(model_path)

real_model = model.load_from_checkpoint(model_path)

y_pred, y_test = get_predictions(real_model, valid_data_loader)

# Get model precision, recall and f1_score

print(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_matrix.png')

#prediction = predict_probability(real_model, validation_transforms, image_path)

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_path, label_idx, use_cpu=True, interpret_type="saliency")