import tensorflow as tf
from keras import metrics


class CustomMetric(metrics.Metric):

    def __init__(self, thresholds=0.5, name="custom_metric", **kwargs):
        super(CustomMetric, self).__init__(name=name, **kwargs)
        self.thresholds = thresholds
        self.true_positives = self.add_weight(name="ctp", initializer="zeros")
        self.count = self.add_weight(name="count", initializer="zeros", dtype='int32')

    @staticmethod
    def y_true_with_others(y_true):
        y_true_idx = tf.argmax(y_true, axis=1) + 1
        y_true_is_other = tf.cast(tf.math.reduce_sum(y_true, axis=1), "int64")
        y_true = tf.math.multiply(y_true_idx, y_true_is_other)
        return y_true

    def y_pred_with_others(self, y_pred):
        y_pred_idx = tf.argmax(y_pred, axis=1) + 1
        y_pred_is_other = tf.cast(tf.math.greater_equal(tf.math.reduce_max(y_pred, axis=1), self.thresholds), 'int64')
        y_pred = tf.math.multiply(y_pred_idx, y_pred_is_other)
        return y_pred

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

        # print(y_true)
        # print(y_pred)

        values = tf.cast(y_true, "int32") == tf.cast(y_pred, "int32")
        values = tf.cast(values, "float32")
        if sample_weight is not None:
            sample_weight = tf.cast(sample_weight, "float32")
            values = tf.multiply(values, sample_weight)
        self.true_positives.assign_add(tf.reduce_sum(values))
        self.count.assign_add(tf.shape(y_true)[0])

    def result(self):
        return self.true_positives / tf.cast(self.count, 'float32')

    def reset_state(self):
        # The state of the metric will be reset at the start of each epoch.
        self.true_positives.assign(0.0)
        self.count.assign(0)