text: __author__ = 'sherk'
  __version__ = '0.2.3'
  """
  
  """
  
  import datetime
  
  from os.path import join
  from itertools import chain
  
  import numpy as np
  
  import tensorflow as tf
  import tensorflow_hub as hub
  from tensorflow.contrib import slim
  
  from audhelper.read import nread
  
  _EPSILON = tf.keras.backend.epsilon()
  
  
  def sparse_softmax_categorical_focal_loss(
    logits, labels, gamma=2.0, alpha=1.0
  ) -> tf.Tensor:
    r"""Focal loss function for multiclass classification with integer labels.
    This loss function generalizes multiclass softmax cross-entropy by
    introducing a hyperparameter called the *focusing parameter* that allows
    hard-to-classify examples to be penalized more heavily relative to
    easy-to-classify examples.
    See :meth:`~focal_loss.binary_focal_loss` for a description of the focal
    loss in the binary setting, as presented in the original work [1]_.
    In the multiclass setting, with integer labels :math:`y`, focal loss is
    defined as
    .. math::
      L(y, \hat{\mathbf{p}})
      = -\left(1 - \hat{p}_y\right)^\gamma \log(\hat{p}_y)
    where
    * :math:`y \in \{0, \ldots, K - 1\}` is an integer class label (:math:`K`
      denotes the number of classes),
    * :math:`\hat{\mathbf{p}} = (\hat{p}_0, \ldots, \hat{p}_{K-1})
      \in [0, 1]^K` is a vector representing an estimated probability
      distribution over the :math:`K` classes,
    * :math:`\gamma` (gamma, not :math:`y`) is the *focusing parameter* that
      specifies how much higher-confidence correct predictions contribute to
      the overall loss (the higher the :math:`\gamma`, the higher the rate at
      which easy-to-classify examples are down-weighted).
    The usual multiclass softmax cross-entropy loss is recovered by setting
    :math:`\gamma = 0`.
    Parameters
    ----------
    labels : (batch_size) tensor-like
      Integer class labels.
    logits : (bach_size, num_cls) tensor-like
      output, need to do softmax inside
    gamma : float or tensor-like of shape (K,)
      The focusing parameter :math:`\gamma`. Higher values of `gamma` make
      easy-to-classify examples contribute less to the loss relative to
      hard-to-classify examples. Must be non-negative. This can be a
      one-dimensional tensor, in which case it specifies a focusing parameter
      for each class.
    alpha :
      compensate ratio for loss reduction
    Returns
    -------
    :class:`tf.Tensor`
      The focal loss for each example.
    Examples
    --------
    This function computes the per-example focal loss between a one-dimensional
    integer label vector and a two-dimensional prediction matrix:
    >>> import numpy as np
    >>> from focal_loss import sparse_categorical_focal_loss
    >>> labels = [0, 0, 0, 1, 2]
    >>> preds = [[0.1, 0.8, 0.1], [0.5, 0.4, 0.1], [0.8, 0.1, 0.1], [0.2, 0.7, 0.1], [0.2, 0.2, 0.6]]
    >>> pass
    Warnings
    --------
    This function does not reduce its output to a scalar, so it cannot be passed
    to :meth:`tf.keras.Model.compile` as a `loss` argument. Instead, use the
    wrapper class :class:`~focal_loss.SparseCategoricalFocalLoss`.
    References
    ----------
    [1] T. Lin, P. Goyal, R. Girshick, K. He and P. Dollár. Focal loss for
      dense object detection. IEEE Transactions on Pattern Analysis and
      Machine Intelligence, 2018.
      (`DOI <https://doi.org/10.1109/TPAMI.2018.2858826>`__)
      (`arXiv preprint <https://arxiv.org/abs/1708.02002>`__)
    See Also
    --------
    :meth:`~focal_loss.SparseCategoricalFocalLoss`
      A wrapper around this function that makes it a
      :class:`tf.keras.losses.Loss`.
    """
    _probs = tf.nn.softmax(logits)
    _xloss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)
    
    _mask = tf.one_hot(labels, depth=_probs.shape[1]) # TODO: add num_class support
    _p = tf.add(tf.reduce_max(tf.multiply(_probs, _mask), axis=1), _EPSILON)
  
  
    _focal_modulation = - alpha * tf.log(_p) * ( (1 - _p) ** gamma )
    _loss = _focal_modulation * _xloss
  
    return _loss
  
  class BaseKWS(object):
    def __init__(self, cfg, training=False):
      self.initailized          = False
      self.summary_dir          = cfg['summary_dir'] if 'summary_dir' in cfg else None
      self.training_dir         = cfg['training_dir'] if 'training_dir' in cfg else None
      self.wanted_words         = ['_unknown_'] + cfg['wanted_words']
      self.num_classes          = len(self.wanted_words)
      self.samples              = cfg['samples']
      self.sample_rate          = cfg['sample_rate']
      self.loss_type            = cfg['loss_type']
      self.batch_size           = cfg['batch_size']
      self.learning_rates       = cfg['learning_rates']
      self.training_steps       = cfg['training_steps']
      self.eval_step_interval   = cfg['eval_step_interval']
      self.training             = training
      self.name                 = None
  
      # initialize graph and sess
      self.graph = tf.Graph()
      self.sess = tf.Session(graph=self.graph)
  
    def initialize(self):
      with self.graph.as_default():
        # initialize self.loss
        if self.loss_type == 'Focal':
          self.loss = sparse_softmax_categorical_focal_loss
        elif self.loss_type == 'CE':
          self.loss = tf.nn.sparse_softmax_cross_entropy_with_logits
        else:
          raise Exception('Unsupported loss type: %s' % self.loss_type)
  
        self.__audios = tf.placeholder(tf.float32, [None, self.samples])
        self.__labels = tf.placeholder(tf.int64, [None])
  
        self.__logits = self.module(self.__audios)
        self.__preds = tf.nn.softmax(self.__logits)
  
        self.__accuracy = tf.contrib.metrics.accuracy(tf.arg_max(self.__preds, 1), self.__labels)
        tf.summary.scalar('accuracy', self.__accuracy)
  
        if self.training:
          # we have only loss in training
          self.__loss = tf.reduce_mean(self.loss(logits=self.__logits, labels=self.__labels))
          tf.summary.scalar('loss', self.__loss)
  
          update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
          with tf.name_scope('train'), tf.control_dependencies(update_ops):
            self.__lr = tf.placeholder(tf.float32, [], name='learning_rate_input')
            self.__optimizer = tf.train.AdamOptimizer(self.__lr)
            self.__train_step = slim.learning.create_train_op(self.__loss, self.__optimizer)
          self.__global_step = tf.train.get_or_create_global_step()
          self.__saver = tf.train.Saver(tf.global_variables())
  
          # merge all summaries
          self.__summaries = tf.summary.merge_all()
          self.__train_writer = tf.summary.FileWriter(join(self.summary_dir, 'train'))
          self.__val_writer = tf.summary.FileWriter(join(self.summary_dir, 'val'))
  
        with self.sess.as_default():
          self.sess.run(tf.local_variables_initializer())
          self.sess.run(tf.global_variables_initializer())
  
          # Parameter counts
          params = tf.trainable_variables()
          num_params = sum(map(lambda t: np.prod(tf.shape(t.value()).eval()), params))
          print('Training parameter numbers: %d' % num_params)
      self.initailized = True
    
    def __infer(self, auds):
      return self.sess.run(self.__preds, feed_dict={self.__audios: auds})
  
    def infer(self, audio):
      if not self.initailized:
        print('Model not initailized!')
        self.initialize()
  
      aud = nread(
        data=audio,
        samples=self.samples, sample_rate=self.sample_rate,
        shuffle=False, aug=None
      ).reshape((-1, self.samples))
  
      return self.__infer(aud)
    
    def batch_infer(self, audios):
      if not self.initailized:
        print('Model not initailized!')
        self.initialize()
      
      # assume the audios are already normalized!
      return self.__infer(audios)
  
    def pred(self, audios):
      if not self.initailized:
        print('Model not initailized!')
        self.initialize()
  
      assert len(audios) <= self.batch_size, f'Input length {len(audios)} is larger than batch size{self.batch_size}'    
  
      auds = np.empty((len(audios), self.samples), dtype=np.float32)
  
      for i, audio in enumerate(audios):
        auds[i] = nread(
          data=audio,
          samples=self.samples, sample_rate=self.sample_rate,
          shuffle=False, aug=None
        )
  
      return self.sess.run(self.__preds, feed_dict={self.__audios: auds})
  
    def train(self, train_dataset, val_dataset=None):
      assert self.initailized, 'Model not initailized!'
  
      train_dataset.eval(self)
  
      best_accuracy = 0.
      lrs = list(chain.from_iterable([[lr] * ep for lr, ep in zip(self.learning_rates, self.training_steps)]))
  
      for epoch, learning_rate_value in enumerate(lrs):
        for _b_auds, _b_labels in train_dataset:
          train_summary, train_accuracy, train_loss, curr_step, _ = self.sess.run(
            [self.__summaries, self.__accuracy, self.__loss, self.__global_step, self.__train_step],
            feed_dict={
              self.__audios: _b_auds,
              self.__labels: _b_labels,
              self.__lr: learning_rate_value
            }
          )
          self.__train_writer.add_summary(train_summary, curr_step)
          print('%d-%d: %.2f-%.4f' % (epoch, curr_step, train_accuracy*100, train_loss))
  
          # validation
          if val_dataset and (curr_step + 1) % self.eval_step_interval == 0:
            total_accuracy, total_loss = self.test(val_dataset, curr_step=curr_step)
            
            if total_accuracy > best_accuracy:
              best_accuracy = total_accuracy
              save_dir = join(self.training_dir, '%d-%d' % (int(best_accuracy*1e4), curr_step))
              self.save_module(save_dir)
              print('Validation %d-%d: %.2f-%.4f' % (epoch, curr_step, total_accuracy*100, total_loss))
      save_dir = join(self.training_dir, 'last')
      self.save_module(save_dir)
  
    def batch_train(self, audios, labels, lr):
      train_summary, train_accuracy, train_loss, curr_step, _ = self.sess.run(
        [self.__summaries, self.__accuracy, self.__loss, self.__global_step, self.__train_step],
        feed_dict={
          self.__audios: audios,
          self.__labels: labels,
          self.__lr: lr
        }
      )
      self.__train_writer.add_summary(train_summary, curr_step)
      print('%d, learning_rate: %.5f: %.2f-%.4f' % (curr_step, lr, train_accuracy*100, train_loss))
  
    def test(self, test_dataset, curr_step=0):
      assert self.initailized, 'Model not initailized!'
  
      test_dataset.eval(self)
      total_accuracy = 0.
      total_loss = 0.
      for test_auds, test_labels in test_dataset:
        if not self.training:
          test_accuracy = self.sess.run(
            self.__accuracy, feed_dict={self.__audios: test_auds, self.__labels: test_labels}
          )
        else:
          test_summary, test_accuracy, test_loss = self.sess.run(
            [self.__summaries, self.__accuracy, self.__loss],
            feed_dict={
              self.__audios: test_auds,
              self.__labels: test_labels
            }
          )
          self.__val_writer.add_summary(test_summary, curr_step)
          total_loss += (test_loss * len(test_labels))
        total_accuracy += (test_accuracy * len(test_labels))
      total_accuracy /= (test_dataset.size + _EPSILON)
      total_loss /= (test_dataset.size + _EPSILON)
  
      return total_accuracy, total_loss
  
    def save_module(self, save_dir):
      raise NotImplementedError
  
    def get_name(self):
      return self.name
  
  if __name__ == "__main__":
    pass
  
  
  
  
  
  
  