"""
WLASL evaluation metric
"""
import tensorflow as tf


class ClassSpecificMetrics(tf.keras.metrics.Metric):
    """Evaluate classification for each class separately
    """
    def __init__(self, classes, **kwargs):
        super().__init__(**kwargs)
        # keep track of true positives, false positives, true negatives and false negatives
        self.true_positives = self.add_weight(shape=(classes, ), name='tp', initializer='zeros')
        self.false_positives = self.add_weight(shape=(classes, ), name='fp', initializer='zeros')
        self.true_negatives = self.add_weight(shape=(classes, ), name='tn', initializer='zeros')
        self.false_negatives = self.add_weight(shape=(classes, ), name='fn', initializer='zeros')
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        # true labels
        y_true = tf.cast(y_true, tf.bool)
        # convert predicted distribution to one hot encoding
        y_pred = tf.cast(tf.one_hot(tf.argmax(y_pred, axis=1), tf.shape(y_true)[1], axis=1), dtype=tf.bool)
        y_pred = tf.reshape(y_pred, (tf.shape(y_pred)[0], -1))
        # calculate TP, FP, TN and FN
        tp = tf.cast(tf.logical_and(y_true, y_pred), self.dtype)
        fp = tf.cast(tf.logical_and(tf.logical_not(y_true), y_pred), self.dtype)
        tn = tf.cast(tf.logical_not(tf.logical_or(y_true, y_pred)), self.dtype)
        fn = tf.cast(tf.logical_and(y_true, tf.logical_not(y_pred)), self.dtype)
        # update stats
        self.true_positives.assign_add(tf.reduce_sum(tp, axis=0))
        self.false_positives.assign_add(tf.reduce_sum(fp, axis=0))
        self.true_negatives.assign_add(tf.reduce_sum(tn, axis=0))
        self.false_negatives.assign_add(tf.reduce_sum(fn, axis=0))

    def result(self):
        # output results
        return {w.name: w for w in self.weights}
