from tqdm import tqdm

import numpy as np

import tensorflow as tf
from tensorflow.keras import metrics, backend

from . import op


def accuracy_score(y_true, y_pred):

    score = (y_true.squeeze() == y_pred.squeeze()).mean()

    return score


def compute_binary_iou(y, y_hat):

    intersection = np.logical_and(y, y_hat).sum()
    union = np.logical_or(y, y_hat).sum()

    return intersection / union


def compute_binary_dice(y, y_hat):

    intersection = np.logical_and(y, y_hat).sum()
    total_cardinality = y.sum() + y_hat.sum()

    return 2.0 * intersection / total_cardinality


def tf_compute_binary_iou(y, y_hat):

    intersection = tf.logical_and(y, y_hat)
    intersection = backend.sum(tf.cast(intersection, dtype=tf.float32))

    union = tf.logical_or(y, y_hat)
    union = backend.sum(tf.cast(union, dtype=tf.float32))

    return intersection / union


def tf_compute_binary_dice(y, y_hat):

    intersection = tf.logical_and(y, y_hat)
    intersection = backend.sum(tf.cast(intersection, dtype=tf.float32))

    total_cardinality = backend.sum(tf.cast(y, dtype=tf.float32))
    total_cardinality += backend.sum(tf.cast(y_hat, dtype=tf.float32))

    return 2.0 * intersection / total_cardinality


def binary_iou(y_true, y_pred):

    size = len(y_true)

    score = 0.0

    for i in tqdm(range(size)):

        score += compute_binary_iou(y_true[i], y_pred[i])

    return score / size


def binary_dice(y_true, y_pred):

    size = len(y_true)

    score = 0.0

    for i in tqdm(range(size)):

        score += compute_binary_dice(y_true[i], y_pred[i])

    return score / size


def tf_binary_iou(y_true, y_pred):

    def _iou(elems):

        return tf_compute_binary_iou(elems[0], elems[1])

    score = backend.mean(tf.map_fn(_iou, (y_true, y_pred), fn_output_signature=tf.float32))

    return score


def tf_binary_dice(y_true, y_pred):

    def _dice(elems):

        return tf_compute_binary_dice(elems[0], elems[1])

    score = backend.mean(tf.map_fn(_dice, (y_true, y_pred), fn_output_signature=tf.float32))

    return score


def cluster_level(counts, max_index=None):

    proba = counts / counts.sum()

    if max_index is None:

        max_index = np.argmax(proba)

    pos_proba = proba[max_index]

    pos_entropy = -pos_proba * np.log(pos_proba)

    total_entropy = - proba * np.log(proba)

    error = np.sum(total_entropy) - pos_entropy

    return pos_proba, error


def clustering_accuracy(y_true, labels):

    size = len(y_true)

    y_pred = op.clusters_to_y(y_true, labels)

    score = np.sum(y_true == y_pred) / size

    return score


def tf_partitions_reduction(output):

    labels = tf.cast(output, dtype=tf.int32)

    partitions = backend.shape(tf.unique(backend.flatten(labels))[0])[0]

    return partitions


class ClusteringAccuracy(metrics.Accuracy):

    def __init__(self, name='clustering_accuracy', **kwargs):

        super(ClusteringAccuracy, self).__init__(name=name, **kwargs)

    def update_state(self, y_true, y_pred, sample_weight=None):

        y_pred = op.tf_clusters_to_y(y_true, y_pred)

        super(ClusteringAccuracy, self).update_state(y_true, y_pred, sample_weight=sample_weight)
