# pylint: disable=missing-docstring
import os
import time
import numpy
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf
import math
import numpy as np
import gdown
from sklearn.metrics import roc_auc_score, roc_curve
import scikitplot as skplt
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure
from os.path import join
import seaborn as sns
sns.set_style("darkgrid")

REMOTE_CKPTS = {
    'cancer_fine_tuned': {'url': '1Zc8KdEz9JUfkT1ZsG9ELYReUPbVapbQC', 'md5': 'cd5271617e090859f73a727da81cc2e3'},
    'i3d_imagenet': {'url': '1FMWHGFYPjuvpgzkGm-_gYKdXpmv5fOq2',  'md5': 'f1408b50e5871153516fe87932121745'}
}

def load_pretrained_ckpt(ckpt, data_dir):
    if ckpt in REMOTE_CKPTS:
        download_ckpt(data_dir, ckpt, REMOTE_CKPTS[ckpt])

    # Load a pre-defined ckpt or a ckpt from path
    predefined = join(data_dir, 'checkpoints', ckpt)
    ckpt_dir = predefined if os.path.exists(predefined) else ckpt

    pre_trained_ckpt = join(ckpt_dir, 'model.ckpt')
    print('\nINFO: Loading pre-trained model:', pre_trained_ckpt)
    return pre_trained_ckpt

def download_ckpt(data_dir, name, download_info):
    if os.path.exists(join(data_dir, 'checkpoints', name)):
        print('\nINFO: {} model already downloaded.'.format(name))
    else:
        print('\nINFO: Downloading {} model...'.format(name))
        url = 'https://drive.google.com/uc?id=' + download_info['url']
        zip_output = join(data_dir, 'checkpoints', name + '.zip')
        md5 = download_info['md5']
        gdown.cached_download(url, zip_output, md5=md5, postprocess=gdown.extractall, quiet=True)
        os.remove(zip_output)

def pretty_print_floats(lst):
    return ',  '.join(['{:.3f}'.format(_) for _ in lst])

def load_npz_as_list(base_dir, npz_file):
    return np.load(join(base_dir, npz_file))['arr_0'].tolist()

def plot_loss(val_loss, tr_loss, plots_dir):
    figure(num=None, figsize=(16, 8), dpi=100)
    title = 'Training and Validation Loss'
    epochs = range(1, len(val_loss) + 1)
    plt.plot(epochs, val_loss, label='Val. Loss')
    plt.plot(epochs, tr_loss, label='Train Loss')
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()
    plt.savefig(join(plots_dir, title + '.png'), bbox_inches='tight')

def plot_acc_auc(val_acc, tr_acc, val_auc, tr_auc, plots_dir):
    figure(num=None, figsize=(16, 8), dpi=100)
    title = 'Accuracy and AUC'
    epochs = range(1, len(val_acc) + 1)
    plt.plot(epochs, val_acc, label='Val. Accuracy')
    plt.plot(epochs, tr_acc, label='Train Accuracy')
    plt.plot(epochs, val_auc, label='Val. AUC')
    plt.plot(epochs, tr_auc, label='Train AUC')
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel('Score')
    plt.legend()
    plt.savefig(join(plots_dir, title + '.png'), bbox_inches='tight')

def calc_plot_epoch_auc_roc(y, y_probs, title, plots_dir, verbose=False):
    y_prob_2_classes = [(1 - p, p) for p in y_probs]
    fpr, tpr, th = roc_curve(y, y_probs)
    if verbose:
        print('TPR:', pretty_print_floats(tpr))
        print('FPR:', pretty_print_floats(fpr))
        print('TH: ', pretty_print_floats(th), '\n')
    auc = roc_auc_score(y, y_probs)
    title = title + ',  AUC={:.3f}'.format(auc)
    skplt.metrics.plot_roc(y, y_prob_2_classes, classes_to_plot=[], 
                           title= title,
                           figsize=(7, 7), plot_micro=False, plot_macro=True, 
                           title_fontsize=15, text_fontsize=13)
    plt.show()
    plt.savefig(join(plots_dir, title) + '.png', bbox_inches='tight')

def load_and_plot_epoch_auc(metrics_dir, epoch, val_true, tr_true, plots_dir):
    val_preds_epoch = load_npz_as_list(metrics_dir, 'val_preds/epoch_' + str(epoch) + '.npz')
    calc_plot_epoch_auc_roc(val_true, val_preds_epoch, 
                            'Val. ROC for Epoch {}'.format(epoch), plots_dir)

    tr_preds_epoch = load_npz_as_list(metrics_dir, 'tr_preds/epoch_' + str(epoch) + '.npz')
    calc_plot_epoch_auc_roc(tr_true, tr_preds_epoch, 
                            'Train ROC for Epoch {}'.format(epoch), plots_dir)

def plot_metrics(epoch, metrics_dir, plots_dir):
    val_loss = load_npz_as_list(metrics_dir, 'val_loss.npz')
    val_acc = load_npz_as_list(metrics_dir, 'val_acc.npz')
    val_auc = load_npz_as_list(metrics_dir, 'val_auc.npz')
    val_true = load_npz_as_list(metrics_dir, 'val_true.npz')

    tr_loss = load_npz_as_list(metrics_dir, 'tr_loss.npz')
    tr_acc = load_npz_as_list(metrics_dir, 'tr_acc.npz')
    tr_auc = load_npz_as_list(metrics_dir, 'tr_auc.npz')
    tr_true = load_npz_as_list(metrics_dir, 'tr_true.npz')

    plot_loss(val_loss, tr_loss, plots_dir)
    plot_acc_auc(val_acc, tr_acc, val_auc, tr_auc, plots_dir)
    load_and_plot_epoch_auc(metrics_dir, epoch, val_true, tr_true, plots_dir)

def write_metrics(metrics, tr_metrics, val_metrics, metrics_dir, epoch, verbose=False):
    for (loss, acc, auc, preds, _), ds in ((tr_metrics, 'tr'), (val_metrics, 'val')):
        for metric, key in [(loss, 'loss'), (acc, 'acc'), (auc, 'auc'), (preds, 'preds')]:
            name = ds + '_' + key
            metrics[name].append(metric)
            write_number_list(metrics[name], join(metrics_dir, name))
        write_number_list(preds, join(metrics_dir, ds + '_preds', 'epoch_{}'.format(epoch)), verbose=verbose)

def apply_window(volume, axis=4):
    # Windowing
    # Our values currently range from -1024 to around 2000. 
    # Anything above 400 is not interesting to us, as these are simply bones with different radiodensity.  
    # A commonly used set of thresholds in Lungs LDCT to normalize between are -1000 and 400. 
    min_bound = -1000.0
    max_bound = 400.0
    volume = (volume - min_bound) / (max_bound - min_bound)
    volume[volume>1] = 1.
    volume[volume<0] = 0.

    # Normalize rgb values to [-1, 1]
    volume = (volume * 2) - 1
    res =  np.stack((volume, volume, volume), axis=axis)
    return res.astype(np.float32)

def write_number_list(lst, f_name, verbose=False):
    if verbose:
        print('INFO: Saving ' + f_name + '.npz ...')
        print(lst)
    np.savez(f_name + '.npz', np.array(lst))       

def batcher(iterable, batch_size=1):
    iter_len = len(iterable)
    for i in range(0, iter_len, batch_size):
        yield iterable[i: min(i + batch_size, iter_len)]

def load_data_list(path):
    coupled_data = []
    with open(path) as file_list_fp:
        for line in file_list_fp:
            volume_path, label = line.split()
            coupled_data.append((volume_path, int(label)))
    return coupled_data

def get_list_labels(coupled_data):
        return np.array([l for _, l in coupled_data]).astype(np.int64)

def placeholder_inputs(num_slices, crop_size, rgb_channels=3):
    """Generate placeholder variables to represent the input tensors.

    These placeholders are used as inputs by the rest of the model building
    code and will be fed from the downloaded data in the .run() loop, below.

    Args:
    num_slices: The num of slices per volume.
    crop_size: The crop size of per volume.
    channels: The number of RGB input channels per volume.

    Returns:
    volumes_placeholder: volumes placeholder.
    labels_placeholder: Labels placeholder.
    """
    # Note that the shapes of the placeholders match the shapes of the full
    # volume and label tensors, except the first dimension is now batch_size
    # rather than the full size of the train or test data sets.
    volumes_placeholder = tf.placeholder(tf.float32, shape=(None,
                                                           num_slices,
                                                           crop_size,
                                                           crop_size,
                                                           rgb_channels))
    labels_placeholder = tf.placeholder(tf.int64, shape=(None))
    is_training = tf.placeholder(tf.bool)
    return volumes_placeholder, labels_placeholder, is_training

def focal_loss(logits, labels, alpha=0.75, gamma=2):
    """Compute focal loss for binary classification.

    Args:
      labels: A int32 tensor of shape [batch_size].
      logits: A float32 tensor of shape [batch_size].
      alpha: A scalar for focal loss alpha hyper-parameter. If positive samples number
      > negtive samples number, alpha < 0.5 and vice versa.
      gamma: A scalar for focal loss gamma hyper-parameter.
    Returns:
      A tensor of the same shape as `labels`.
    """
    y_pred = tf.nn.sigmoid(logits)
    labels = tf.to_float(labels)
    losses = -(labels * (1 - alpha) * ((1 - y_pred) * gamma) * tf.log(y_pred)) - \
        (1 - labels) * alpha * (y_pred ** gamma) * tf.log(1 - y_pred)
    return tf.reduce_mean(losses)

def cross_entropy_loss(logits, labels):
    # pylint: disable=no-value-for-parameter
    # pylint: disable=unexpected-keyword-arg
    cross_entropy_mean = tf.reduce_mean(
                  tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
                  )
    return cross_entropy_mean

def accuracy(logit, labels):
    correct_pred = tf.equal(tf.argmax(logit, 1), labels)
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    return accuracy

def get_preds(preds):
    return preds[:, 1]

def get_logits(logits):
    return logits