import csv

import os

import numpy as np

import pandas as pd

from sklearn.metrics import roc_auc_score, accuracy_score

import torch

from sklearn.metrics import confusion_matrix

from sklearn import metrics

import seaborn as sns

import matplotlib.pyplot as plt

def image_to_csv(file_path, save_path, mode=""):

    """

    Create CSV from Acute Lymphoblastic Leukemia dataset

    Args:

        file_path: Acute Lymphoblastic Leukemia data path

        save_path: path to save the created CSV

        mode:

    """

    columns = ['data', 'label']

    csv_data = ""

    imagefiles = list()  # create a list to store image names

    if mode == "train":

        csv_data = "train.csv"

    if mode == "test":

        csv_data = "test.csv"

    with open(os.path.join(save_path, csv_data), 'w', newline='') as csvfile:

        for root, dirs, files in os.walk(file_path): #scan through the file path

            for file in files: # loop through all files

                if '.tif' or '.jpg' or '.png' in file: #chech if there are files with *.tif, *.jpg or *.png

                    imagefiles.append(os.path.splitext(file)[0])#retrieve file names and add to imagefiles list

        writer = csv.writer(csvfile, dialect='excel')  # Create a writer from csv module

        writer.writerow(columns)#write down the columns

        for image in imagefiles:# loop through all image names in the imagefiles list

            label = os.path.basename(image)

            if "_0" in label:

                label = 0

            elif "_1" in label:

                label = 1

            writer.writerow([image, label])

    print("done")

def print_accuracy_and_classification_report(labels, prediction):

    """Print model accuracy and classification report.

    Args:

        labels (numpy.array): Truth label

        prediction (numpy.array): Model predictions

    """

    print('Cross validation accuracy:')

    print('\t', metrics.accuracy_score(labels, prediction))

    print('\nCross validation classification report\n')

    print(metrics.classification_report(labels, prediction))

def f1_loss(y_true: torch.Tensor, y_pred: torch.Tensor, is_training=False) -> torch.Tensor:

    '''Calculate F1 score. Can work with gpu tensors

    The original implmentation is written by Michal Haltuf on Kaggle.

    Returns

    -------

    torch.Tensor

        `ndim` == 1. 0 <= val <= 1

    Reference

    ---------

    - https://www.kaggle.com/rejpalcz/best-loss-function-for-f1-score-metric

    - https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html#sklearn.metrics.f1_score

    - https://discuss.pytorch.org/t/calculating-precision-recall-and-f1-score-in-case-of-multi-label-classification/28265/6

    '''

    assert y_true.ndim == 1

    assert y_pred.ndim == 1 or y_pred.ndim == 2

    if y_pred.ndim == 2:

        y_pred = y_pred.argmax(dim=1)

    tp = (y_true * y_pred).sum().to(torch.float32)

    tn = ((1 - y_true) * (1 - y_pred)).sum().to(torch.float32)

    fp = ((1 - y_true) * y_pred).sum().to(torch.float32)

    fn = (y_true * (1 - y_pred)).sum().to(torch.float32)

    epsilon = 1e-7

    precision = tp / (tp + fp + epsilon)

    recall = tp / (tp + fn + epsilon)

    f1 = 2 * (precision * recall) / (precision + recall + epsilon)

    f1.requires_grad = is_training

    return f1

def sigmoid(x):

    return 1.0 / (1.0 + np.exp(-x))

def accuracy_mini_batch(predicted, true, i, acc, tpr, tnr):

    predicted = predicted.cpu()

    true = true.cpu()

    predicted = (sigmoid(predicted.data.numpy()) > 0.5)

    true = true.data.numpy()

    accuracy = np.sum(predicted == true) / true.shape[0]

    true_positive_rate = np.sum((predicted == 1) * (true == 1)) / np.sum(true == 1)

    true_negative_rate = np.sum((predicted == 0) * (true == 0)) / np.sum(true == 0)

    acc = acc * (i) / (i + 1) + accuracy / (i + 1)

    tpr = tpr * (i) / (i + 1) + true_positive_rate / (i + 1)

    tnr = tnr * (i) / (i + 1) + true_negative_rate / (i + 1)

    return acc, tpr, tnr

def accuracy(predicted, true):

    predicted = predicted.cpu()

    true = true.cpu()

    predicted = (sigmoid(predicted.data.numpy()) > 0.5)

    true = true.data.numpy()

    accuracy = np.sum(predicted == true) / true.shape[0]

    true_positive_rate = np.sum((predicted == 1) * (true == 1)) / np.sum(true == 1)

    true_negative_rate = np.sum((predicted == 0) * (true == 0)) / np.sum(true == 0)

    return accuracy, true_positive_rate, true_negative_rate

def model_confusion_matrix(y_true, y_pred, classes=[]):

    cm = confusion_matrix(y_true=y_true, y_pred=y_pred)

    df_cm = pd.DataFrame(cm, index=classes, columns=classes)

    hmap = sns.heatmap(df_cm, annot=True, fmt='d')

    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 Label')

    plt.xlabel('Predicted label')

    plt.show()