from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import scipy.optimize
import scipy.special
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp


def find_scaling_temperature(labels, logits, temp_range=(1e-5, 1e5)):
    """
  Code is adapted from https://github.com/google-research/google-research/tree/master/uq_benchmark_2019/
  Find max likelihood scaling temperature using binary search.

  Args:
    labels: Integer labels (shape=[num_samples]).
    logits: Floating point softmax inputs (shape=[num_samples, num_classes]).
    temp_range: 2-tuple range of temperatures to consider.
  Returns:
    Floating point temperature value.
  """
    if not tf.executing_eagerly():
        raise NotImplementedError(
            'find_scaling_temperature() not implemented for graph-mode TF')
    if len(labels.shape) != 1:
        raise ValueError('Invalid labels shape=%s' % str(labels.shape))
    if len(logits.shape) not in (1, 2):
        raise ValueError('Invalid logits shape=%s' % str(logits.shape))
    if len(labels.shape) != 1 or len(labels) != len(logits):
        raise ValueError('Incompatible shapes for logits (%s) vs labels (%s).' %
                         (logits.shape, labels.shape))

    @tf.function(autograph=False, experimental_relax_shapes=True)
    def grad_fn(temperature):
        """Returns gradient of log-likelihood WRT a logits-scaling temperature."""
        dist = tfp.distributions.Bernoulli(logits=logits / temperature)
        nll = -dist.log_prob(labels)
        nll = tf.reduce_sum(nll, axis=0)
        grad, = tf.gradients(nll, [temperature])
        return grad

    tmin, tmax = temp_range
    return scipy.optimize.bisect(lambda t: grad_fn(t * tf.ones([])).numpy(), tmin, tmax)


def inverse_sigmoid(probs, eps=1e-7):
    """ compute the logit"""
    uniform = np.ones_like(probs) / probs.shape[-1]
    probs = eps * uniform + (1. - eps) * probs
    return -tf.math.log(1. / probs - 1.)


def apply_temperature_scaling(temperature, probs):
    """Apply temperature scaling to an array of probabilities."""
    logits_t = inverse_sigmoid(probs) / temperature
    return tf.sigmoid(logits_t).numpy()
