from tqdm import tqdm

import numpy as np

import tensorflow as tf
from tensorflow.keras import backend

from utils import op, evaluate


def g_loss(rho):

    @tf.function
    def reconstruction_loss(y_true, y_pred):

        loss = 0.5 * backend.mean((y_pred - y_true) ** 2)

        loss = tf.cond(loss < rho, true_fn=lambda: 0.0, false_fn=lambda: loss, name='lower_bound')

        return loss

    return reconstruction_loss


@tf.function
def adversarial_partition_loss(latent):

    state = backend.random_uniform(shape=backend.shape(latent), minval=0.0,
                                   maxval=2.0, dtype=tf.float32)

    loss = backend.binary_crossentropy(state, latent)
    loss = backend.sum(loss)

    return loss


@tf.function
def zero_loss(*_):

    return 0.0


@tf.function
def iou(y_true, y_pred):

    y_true = tf.cast(y_true, dtype=tf.bool)

    y_pred = tf.cast(y_pred, dtype=tf.int32)
    y_pred = op.tf_multiclass_to_bitmask(y_pred)

    score = evaluate.tf_binary_iou(y_true, y_pred)

    return score


@tf.function
def partitions_reduction(_, y_pred):

    return evaluate.tf_partitions_reduction(y_pred)


def benchmark_evaluate(model, tf_data):

    k = 0

    accuracy_score = 0.0
    iou_score = 0.0

    for x, y in tqdm(tf_data.as_numpy_iterator()):

        output = model.predict_on_batch(x)

        y_pred = output.astype('int32')
        y_pred = op.multiclass_to_bitmask(y_pred)

        accuracy_score += evaluate.accuracy_score(y, y_pred)
        iou_score += evaluate.binary_iou(y, y_pred)

        k += 1

    accuracy_score /= k
    iou_score /= k

    return accuracy_score, iou_score


def benchmark_clusters_count(model, tf_data):

    counts = []

    for x, _ in tqdm(tf_data.as_numpy_iterator()):

        output = model.predict_on_batch(x)

        y_pred = output.astype('int32')

        for i in range(len(y_pred)):

            counts.append(len(np.unique(y_pred[i])))

    return counts
