from copy import deepcopy

from tqdm import tqdm

import numpy as np

import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow.keras import backend, callbacks, optimizers


def clusters_to_y(y_test, labels):

    unique = np.unique(labels)

    n_clusters = len(unique)

    y_pred = deepcopy(labels)

    for i in range(n_clusters):

        ci, counts = np.unique(y_test[labels == unique[i]], return_counts=True)

        ci = max(tuple(zip(counts, ci)))[1]

        y_pred[labels == unique[i]] = ci

    return y_pred


def tf_clusters_to_y(y_true, labels):

    y_true = tf.cast(y_true, dtype=tf.float32)
    labels = tf.cast(labels, dtype=tf.float32)

    unique = tf.unique(labels)[0]

    n_clusters = tf.size(unique)

    y_pred = tf.zeros_like(y_true) + tf.cast(n_clusters, dtype=tf.float32)

    for i in range(n_clusters):

        indices = tf.where(labels == unique[i])

        true_labels = tf.gather(y_true, indices)
        true_labels = tf.squeeze(true_labels, axis=-1)

        # --------------------------------------------

        ci, _, counts = tf.unique_with_counts(true_labels)

        max_index = backend.argmax(counts)

        ci = tf.gather(ci, max_index)

        # --------------------------------------------

        values = tf.repeat([ci], tf.size(indices))

        y_pred = tf.tensor_scatter_nd_update(y_pred, indices, values)

    return y_pred


def cut_labels(multi_class, quantile=0.5, interpolation='higher'):

    unique = np.unique(multi_class)
    unique.sort()

    label = np.quantile(unique, q=quantile, interpolation=interpolation)

    k = np.where(unique == label)[0][0]

    return unique[-k::]


def tf_cut_labels(multi_class, quantile=0.5, interpolation='higher'):

    multi_class = tf.cast(multi_class, dtype=tf.int32)
    multi_class = backend.flatten(multi_class)

    unique = tf.unique(multi_class)[0]
    unique = tf.sort(unique)

    label = tfp.stats.percentile(unique, q=quantile*100.0, interpolation=interpolation)

    k = tf.where(unique == label)[0][0]

    return unique[-k::]


def classes_to_bitmask(multi_class, **kwargs):

    labels = cut_labels(multi_class, **kwargs)

    mask = (multi_class == labels[0])

    for i in range(1, len(labels)):

        mask |= (multi_class == labels[i])

    return mask


def tf_classes_to_bitmask(multi_class, **kwargs):

    labels = tf_cut_labels(multi_class, **kwargs)

    mask = tf.equal(multi_class, labels[0])

    n_labels = tf.size(labels)

    for i in range(1, n_labels):

        mask = tf.logical_or(mask, tf.equal(multi_class, labels[i]))

    return mask


def multiclass_to_bitmask(multi_class, **kwargs):

    masks = []

    for i in tqdm(range(len(multi_class))):

        binary = classes_to_bitmask(multi_class[i], **kwargs)

        masks.append(binary)

    return np.array(masks)


def tf_multiclass_to_bitmask(multi_class, **kwargs):

    def _classes_to_bitmask(elems):

        return tf_classes_to_bitmask(elems[0], **kwargs)

    masks = tf.map_fn(_classes_to_bitmask, (multi_class, ), fn_output_signature=tf.bool,
                      name='classes_to_bitmask')

    return masks


def mask_to_segment(image, mask, background=None):

    if background is None:

        background = np.zeros((image.shape[-1], ), dtype='float32')

    segment = deepcopy(image)

    segment[~mask] = np.array(background)[None, :]

    return segment


def set_channel_first(images):

    return np.transpose(images, axes=[0, 3, 1, 2])


def set_channel_last(images):

    return np.transpose(images, axes=[0, 2, 3, 1])


def sample_indices(n, sample_size=4, random_state=42):

    random = np.random.RandomState(random_state)

    indices = np.arange(0, n)

    random.shuffle(indices)

    indices = indices[:sample_size]

    return indices


class RestOptimizerState(callbacks.Callback):

    def __init__(self, identifier, steps, verbose=1):

        super(RestOptimizerState, self).__init__()

        self.steps = steps

        self.verbose = verbose

        self.__config__ = {'class_name': identifier}

    @property
    def config(self):

        self.__config__.update({'config': self.model.optimizer.get_config()})

        return self.__config__

    def on_epoch_end(self, epoch, logs=None):

        if (epoch + 1) % self.steps == 0:

            self.model.optimizer = optimizers.get(self.config)

            if self.verbose == 1:

                print(f'Current State = {epoch // self.steps}')


class RestMetricState(callbacks.Callback):

    def __init__(self, steps, verbose=1):

        super(RestMetricState, self).__init__()

        self.steps = steps

        self.verbose = verbose

    def on_epoch_end(self, epoch, logs=None):

        if (epoch + 1) % self.steps == 0:

            for i in range(len(self.model.metrics)):

                self.model.metrics[i].reset_state()


class RestRandomState(callbacks.Callback):

    def __init__(self, random_state, steps, verbose=1):

        super(RestRandomState, self).__init__()

        self.random_state = random_state

        self.steps = steps

        self.verbose = verbose

    def on_epoch_end(self, epoch, logs=None):

        if (epoch + 1) % self.steps == 0:

            tf.random.set_seed(self.random_state)
