import numpy as np
import pandas as pd
import os
import torch
from torchvision import transforms
from torch.utils.data import DataLoader
import sys
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.metrics import roc_auc_score, roc_curve, log_loss, auc
from scipy import interp
from itertools import cycle
sys.path.append("/home/anjum/PycharmProjects/kaggle")
# sys.path.append("/home/anjum/rsna_code")  # GCP
from rsna_intracranial_hemorrhage_detection.datasets import ICHDataset

INPUT_DIR = "/mnt/storage_dimm2/kaggle_data/rsna-intracranial-hemorrhage-detection/"


def build_tta_loaders(img_size, dataset, phase=1, image_filter=None, batch_size=32, num_workers=1,
                      image_folder=None, png=True):
    if type(img_size) == int:
        img_size = (img_size, img_size)

    def null_transform(image):
        image = transforms.functional.to_pil_image(image)
        image = transforms.functional.resize(image, img_size)
        tensor = transforms.functional.to_tensor(image)
        # tensor = transforms.functional.normalize(tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        return tensor

    def hflip(image):
        image = transforms.functional.to_pil_image(image)
        image = transforms.functional.hflip(image)
        image = transforms.functional.resize(image, img_size)
        tensor = transforms.functional.to_tensor(image)
        # tensor = transforms.functional.normalize(tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        return tensor

    def rotate_pos(image):
        image = transforms.functional.to_pil_image(image)
        image = transforms.functional.affine(image, angle=10, translate=(0, 0), scale=1.0, shear=0)
        image = transforms.functional.resize(image, img_size)
        tensor = transforms.functional.to_tensor(image)
        # tensor = transforms.functional.normalize(tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        return tensor

    def rotate_neg(image):
        image = transforms.functional.to_pil_image(image)
        image = transforms.functional.affine(image, angle=-10, translate=(0, 0), scale=1.0, shear=0)
        image = transforms.functional.resize(image, img_size)
        tensor = transforms.functional.to_tensor(image)
        # tensor = transforms.functional.normalize(tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        return tensor

    # tta = [null_transform, hflip]
    tta = [null_transform, hflip, rotate_pos, rotate_neg]
    loaders = []

    for augmentation in tta:
        tta_dataset = ICHDataset(dataset, phase=phase, image_filter=image_filter, transforms=augmentation,
                                 image_folder=image_folder, png=png)
        loaders.append(DataLoader(tta_dataset, batch_size=batch_size, num_workers=num_workers))
    return loaders


def infer(model, loader, device, desc):
    model.eval()
    with torch.no_grad():
        predictions, targets = [], []
        for image, target in tqdm(loader, desc=desc):
            image = image.to(device)
            y_hat = model(image)
            predictions.append(y_hat.cpu())
            targets.append(target)

        predictions = torch.cat(predictions)
        targets = torch.cat(targets)
    return predictions, targets


def test_time_augmentation(model, loaders, device, desc):
    tta_predictions = []
    for i, loader in enumerate(loaders):
        predictions, targets = infer(model, loader, device, f"{desc} TTA {i}")
        tta_predictions.append(torch.unsqueeze(predictions, -1))

    tta_predictions = torch.cat(tta_predictions, -1)
    return torch.mean(tta_predictions, dim=-1), targets


class EarlyStopping:
    """
    Early stops the training if validation loss doesn't improve after a given patience.
    https://github.com/Bjarten/early-stopping-pytorch
    """
    def __init__(self, patience=7, verbose=False, delta=0, file_path='checkpoint.pt', parallel=False):
        """
        :param patience: How long to wait after last time validation loss improved. Default: 7
        :param verbose: If True, prints a message for each validation loss improvement. Default: False
        :param delta: Minimum change in the monitored quantity to qualify as an improvement. Default: 0
        :param file_path: Path to save checkpoint file
        :param parallel: If True, the multi-GPU model is saves as a single GPU model
        """

        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.delta = delta
        self.file_path = file_path
        self.parallel = parallel
        self.parallel_model = None

    def __call__(self, val_loss, model, **kwargs):

        score = -val_loss

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model, **kwargs)
        elif score < self.best_score - self.delta:
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model, **kwargs)
            self.counter = 0

    def save_checkpoint(self, val_loss, model, **save_items):
        """Saves model and addition items when validation loss decrease."""
        if self.verbose:
            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')

        if save_items == {}:
            torch.save(model.state_dict(), self.file_path)
        else:
            if self.parallel:
                self.parallel_model = model.state_dict()
                save_items['model'] = model.module.state_dict()
            else:
                save_items['model'] = model.state_dict()
            save_items["stopping_params"] = self.state_dict()
            torch.save(save_items, self.file_path)
        self.val_loss_min = val_loss

    def state_dict(self):
        state = {
            # "counter": self.counter,
            "best_score": self.best_score,
            "val_loss_min": self.val_loss_min,
        }
        return state

    def load_state_dict(self, state):
        # self.counter = state["counter"]
        self.best_score = state["best_score"]
        self.val_loss_min = state["val_loss_min"]


def plot_roc_curve(target, predictions, file_path, metric=None):

    if type(target) == torch.Tensor:
        target = target.numpy()
    if type(predictions) == torch.Tensor:
        predictions = predictions.numpy()

    plt.figure()
    fpr, tpr, _ = roc_curve(target, predictions)
    score = roc_auc_score(target, predictions)
    plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve (area = %0.4f)' % score)
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    if metric is not None:
        plt.title(f'Receiver operating characteristic. LogLoss: {metric:.4f}')
    else:
        plt.title(f'Receiver operating characteristic')
    plt.legend(loc="lower right")
    plt.savefig(file_path)


def plot_multiclass_roc_curve(target, predictions, file_path, metric=None):
    # https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html

    if type(target) == torch.Tensor:
        target = target.numpy()
    if type(predictions) == torch.Tensor:
        predictions = predictions.numpy()

    n_classes = target.shape[1]

    # Compute ROC curve and ROC area for each class
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    log_loss_vals = dict()
    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(target[:, i], predictions[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])
        log_loss_vals[i] = log_loss(target[:, i], predictions[:, i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = roc_curve(target.ravel(), predictions.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    # First aggregate all false positive rates
    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))

    # Then interpolate all ROC curves at this points
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(n_classes):
        mean_tpr += interp(all_fpr, fpr[i], tpr[i])

    # Finally average it and compute AUC
    mean_tpr /= n_classes

    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

    # Plot all ROC curves
    plt.figure()
    plt.plot(fpr["micro"], tpr["micro"],
             label='micro-average ROC curve (area = {0:0.2f})'
                   ''.format(roc_auc["micro"]),
             color='deeppink', linestyle=':', linewidth=4)

    plt.plot(fpr["macro"], tpr["macro"],
             label='macro-average ROC curve (area = {0:0.2f})'
                   ''.format(roc_auc["macro"]),
             color='navy', linestyle=':', linewidth=4)

    colors = cycle(['aqua', 'darkorange', 'cornflowerblue', 'magenta', 'lawngreen', 'gold'])
    for i, color in zip(range(n_classes), colors):
        plt.plot(fpr[i], tpr[i], color=color, lw=2, label='ROC curve of class {0} (area={1:0.2f}, log_loss={2:0.3f})'
                                                          ''.format(i, roc_auc[i], log_loss_vals[i]))

    plt.plot([0, 1], [0, 1], 'k--', lw=2)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')

    if metric is not None:
        plt.title(f'Receiver operating characteristic. LogLoss: {metric:.4f}')
    else:
        plt.title(f'Receiver operating characteristic')

    plt.legend(loc="lower right")
    plt.savefig(file_path)


def reindex_submission(df, stage="test1"):
    if stage not in ["test1", "test2"]:
        return df
    elif stage == "test1":
        sub = pd.read_csv(os.path.join(INPUT_DIR, "stage_1_sample_submission.csv"))
    else:
        sub = pd.read_csv(os.path.join(INPUT_DIR, "stage_2_sample_submission.csv"))
    return df.sort_values(by="ID").set_index(sub.sort_values(by="ID").index).sort_index()


def wide_to_long(df, id_var="ImageID"):
    categories = ["any", "epidural", "intraparenchymal", "intraventricular", "subarachnoid", "subdural"]
    df_long = df.melt(id_vars=[id_var], value_vars=categories, value_name="Label")
    df_long["ID"] = df_long[id_var] + "_" + df_long["variable"]
    df_long = df_long.sort_values(by="ID")
    return df_long[["ID", "Label"]]