from functools import partial
import tensorflow as tf

_EPSILON = tf.keras.backend.epsilon()


def register_keras_custom_object(cls):
    tf.keras.utils.get_custom_objects()[cls.__name__] = cls
    return cls


def binary_focal_loss(y_true, y_pred, gamma, *, pos_weight=None, from_logits=False, label_smoothing=None):
    y_pred = tf.convert_to_tensor(y_pred)
    if not y_pred.dtype.is_floating:
        y_pred = tf.dtypes.cast(y_pred, dtype=tf.float32)

    if from_logits:
        return _binary_focal_loss_from_logits(labels=y_true,
                                              logits=y_pred,
                                              gamma=gamma,
                                              pos_weight=pos_weight,
                                              label_smoothing=label_smoothing)
    else:
        return _binary_focal_loss_from_probs(labels=y_true,
                                             p=y_pred,
                                             gamma=gamma,
                                             pos_weight=pos_weight,
                                             label_smoothing=label_smoothing)


@register_keras_custom_object
class BinaryFocalLoss(tf.keras.losses.Loss):
    def __init__(self, gamma, *, pos_weight=None, from_logits=False, label_smoothing=None, **kwargs):

        super().__init__(**kwargs)
        self.gamma = gamma
        self.pos_weight = pos_weight
        self.from_logits = from_logits
        self.label_smoothing = label_smoothing

    def get_config(self):
        config = super().get_config()
        config.update(gamma=self.gamma,
                      pos_weight=self.pos_weight,
                      from_logits=self.from_logits,
                      label_smoothing=self.label_smoothing)
        return config

    def call(self, y_true, y_pred):
        return binary_focal_loss(y_true=y_true,
                                 y_pred=y_pred,
                                 gamma=self.gamma,
                                 pos_weight=self.pos_weight,
                                 from_logits=self.from_logits,
                                 label_smoothing=self.label_smoothing)


# Helper functions below


def _process_labels(labels, label_smoothing, dtype):
    labels = tf.dtypes.cast(labels, dtype=dtype)
    if label_smoothing is not None:
        labels = (1 - label_smoothing) * labels + label_smoothing * 0.5
    return labels


def _binary_focal_loss_from_logits(labels, logits, gamma, pos_weight, label_smoothing):
    labels = _process_labels(labels=labels, label_smoothing=label_smoothing, dtype=logits.dtype)

    # Compute probabilities for the positive class
    p = tf.math.sigmoid(logits)

    if label_smoothing is None:
        labels_shape = labels.shape
        logits_shape = logits.shape
        if not labels_shape.is_fully_defined() or labels_shape != logits_shape:
            labels_shape = tf.shape(labels)
            logits_shape = tf.shape(logits)
            shape = tf.broadcast_dynamic_shape(labels_shape, logits_shape)
            labels = tf.broadcast_to(labels, shape)
            logits = tf.broadcast_to(logits, shape)
        if pos_weight is None:
            loss_func = tf.nn.sigmoid_cross_entropy_with_logits
        else:
            loss_func = partial(tf.nn.weighted_cross_entropy_with_logits, pos_weight=pos_weight)
        loss = loss_func(labels=labels, logits=logits)
        modulation_pos = (1 - p)**gamma
        modulation_neg = p**gamma
        mask = tf.dtypes.cast(labels, dtype=tf.bool)
        modulation = tf.where(mask, modulation_pos, modulation_neg)
        return modulation * loss

    # Terms for the positive and negative class components of the loss
    pos_term = labels * ((1 - p)**gamma)
    neg_term = (1 - labels) * (p**gamma)

    # Term involving the log and ReLU
    log_weight = pos_term
    if pos_weight is not None:
        log_weight *= pos_weight
    log_weight += neg_term
    log_term = tf.math.log1p(tf.math.exp(-tf.math.abs(logits)))
    log_term += tf.nn.relu(-logits)
    log_term *= log_weight

    # Combine all the terms into the loss
    loss = neg_term * logits + log_term
    return loss


def _binary_focal_loss_from_probs(labels, p, gamma, pos_weight, label_smoothing):
    q = 1 - p

    # For numerical stability (so we don't inadvertently take the log of 0)
    p = tf.math.maximum(p, _EPSILON)
    q = tf.math.maximum(q, _EPSILON)

    # Loss for the positive examples
    pos_loss = -(q**gamma) * tf.math.log(p)
    if pos_weight is not None:
        pos_loss *= pos_weight

    # Loss for the negative examples
    neg_loss = -(p**gamma) * tf.math.log(q)

    # Combine loss terms
    if label_smoothing is None:
        labels = tf.dtypes.cast(labels, dtype=tf.bool)
        loss = tf.where(labels, pos_loss, neg_loss)
    else:
        labels = _process_labels(labels=labels, label_smoothing=label_smoothing, dtype=p.dtype)
        loss = labels * pos_loss + (1 - labels) * neg_loss

    return loss
[((66, 92), 'tensorflow.keras.backend.epsilon', 'tf.keras.backend.epsilon', ([], {}), '()\n', (90, 92), True, 'import tensorflow as tf\n'), ((331, 359), 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['y_pred'], {}), '(y_pred)\n', (351, 359), True, 'import tensorflow as tf\n'), ((2262, 2297), 'tensorflow.dtypes.cast', 'tf.dtypes.cast', (['labels'], {'dtype': 'dtype'}), '(labels, dtype=dtype)\n', (2276, 2297), True, 'import tensorflow as tf\n'), ((2671, 2694), 'tensorflow.math.sigmoid', 'tf.math.sigmoid', (['logits'], {}), '(logits)\n', (2686, 2694), True, 'import tensorflow as tf\n'), ((4026, 4045), 'tensorflow.nn.relu', 'tf.nn.relu', (['(-logits)'], {}), '(-logits)\n', (4036, 4045), True, 'import tensorflow as tf\n'), ((4355, 4383), 'tensorflow.math.maximum', 'tf.math.maximum', (['p', '_EPSILON'], {}), '(p, _EPSILON)\n', (4370, 4383), True, 'import tensorflow as tf\n'), ((4392, 4420), 'tensorflow.math.maximum', 'tf.math.maximum', (['q', '_EPSILON'], {}), '(q, _EPSILON)\n', (4407, 4420), True, 'import tensorflow as tf\n'), ((138, 173), 'tensorflow.keras.utils.get_custom_objects', 'tf.keras.utils.get_custom_objects', ([], {}), '()\n', (171, 173), True, 'import tensorflow as tf\n'), ((414, 454), 'tensorflow.dtypes.cast', 'tf.dtypes.cast', (['y_pred'], {'dtype': 'tf.float32'}), '(y_pred, dtype=tf.float32)\n', (428, 454), True, 'import tensorflow as tf\n'), ((3497, 3534), 'tensorflow.dtypes.cast', 'tf.dtypes.cast', (['labels'], {'dtype': 'tf.bool'}), '(labels, dtype=tf.bool)\n', (3511, 3534), True, 'import tensorflow as tf\n'), ((3556, 3602), 'tensorflow.where', 'tf.where', (['mask', 'modulation_pos', 'modulation_neg'], {}), '(mask, modulation_pos, modulation_neg)\n', (3564, 3602), True, 'import tensorflow as tf\n'), ((4488, 4502), 'tensorflow.math.log', 'tf.math.log', (['p'], {}), '(p)\n', (4499, 4502), True, 'import tensorflow as tf\n'), ((4632, 4646), 'tensorflow.math.log', 'tf.math.log', (['q'], {}), '(q)\n', (4643, 4646), True, 'import tensorflow as tf\n'), ((4722, 4759), 'tensorflow.dtypes.cast', 'tf.dtypes.cast', (['labels'], {'dtype': 'tf.bool'}), '(labels, dtype=tf.bool)\n', (4736, 4759), True, 'import tensorflow as tf\n'), ((4775, 4811), 'tensorflow.where', 'tf.where', (['labels', 'pos_loss', 'neg_loss'], {}), '(labels, pos_loss, neg_loss)\n', (4783, 4811), True, 'import tensorflow as tf\n'), ((2907, 2923), 'tensorflow.shape', 'tf.shape', (['labels'], {}), '(labels)\n', (2915, 2923), True, 'import tensorflow as tf\n'), ((2951, 2967), 'tensorflow.shape', 'tf.shape', (['logits'], {}), '(logits)\n', (2959, 2967), True, 'import tensorflow as tf\n'), ((2988, 3042), 'tensorflow.broadcast_dynamic_shape', 'tf.broadcast_dynamic_shape', (['labels_shape', 'logits_shape'], {}), '(labels_shape, logits_shape)\n', (3014, 3042), True, 'import tensorflow as tf\n'), ((3064, 3094), 'tensorflow.broadcast_to', 'tf.broadcast_to', (['labels', 'shape'], {}), '(labels, shape)\n', (3079, 3094), True, 'import tensorflow as tf\n'), ((3116, 3146), 'tensorflow.broadcast_to', 'tf.broadcast_to', (['logits', 'shape'], {}), '(logits, shape)\n', (3131, 3146), True, 'import tensorflow as tf\n'), ((3280, 3352), 'functools.partial', 'partial', (['tf.nn.weighted_cross_entropy_with_logits'], {'pos_weight': 'pos_weight'}), '(tf.nn.weighted_cross_entropy_with_logits, pos_weight=pos_weight)\n', (3287, 3352), False, 'from functools import partial\n'), ((3988, 4007), 'tensorflow.math.abs', 'tf.math.abs', (['logits'], {}), '(logits)\n', (3999, 4007), True, 'import tensorflow as tf\n')]