// nnet/nnet-loss.cc
  
  // Copyright 2011-2015  Brno University of Technology (author: Karel Vesely)
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //  http://www.apache.org/licenses/LICENSE-2.0
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  #include <sstream>
  #include <iterator>
  #include <algorithm>
  #include <iomanip>
  
  #include "nnet/nnet-loss.h"
  #include "nnet/nnet-utils.h"
  #include "cudamatrix/cu-math.h"
  #include "hmm/posterior.h"
  
  namespace kaldi {
  namespace nnet1 {
  
  
  /* Xent */
  
  /**
   * Helper function of Xent::Eval,
   * calculates number of matching elemente in 'hyp', 'ref' weighted by 'weights'.
   */
  template <typename T>
  inline void CountCorrectFramesWeighted(const CuArray<T> &hyp,
                                         const CuArray<T> &ref,
                                         const CuVectorBase<BaseFloat> &weights,
                                         Vector<double> *correct) {
    KALDI_ASSERT(hyp.Dim() == ref.Dim());
    KALDI_ASSERT(hyp.Dim() == weights.Dim());
    int32 dim = hyp.Dim();
    // Get GPU data to host,
    std::vector<T> hyp_h(dim), ref_h(dim);
    hyp.CopyToVec(&hyp_h);
    ref.CopyToVec(&ref_h);
    Vector<BaseFloat> w(dim);
    weights.CopyToVec(&w);
    // Accumulate weighted counts of correct frames,
    for (int32 i = 0; i < dim; i++) {
      KALDI_ASSERT(ref_h[i] < correct->Dim());
      (*correct)(ref_h[i]) += w(i) * (hyp_h[i] == ref_h[i] ? 1.0 : 0.0);
    }
  }
  
  
  void Xent::Eval(const VectorBase<BaseFloat> &frame_weights,
                  const CuMatrixBase<BaseFloat> &net_out,
                  const CuMatrixBase<BaseFloat> &targets,
                  CuMatrix<BaseFloat> *diff) {
    // check inputs,
    KALDI_ASSERT(net_out.NumCols() == targets.NumCols());
    KALDI_ASSERT(net_out.NumRows() == targets.NumRows());
    KALDI_ASSERT(net_out.NumRows() == frame_weights.Dim());
  
    KALDI_ASSERT(KALDI_ISFINITE(frame_weights.Sum()));
    KALDI_ASSERT(KALDI_ISFINITE(net_out.Sum()));
    KALDI_ASSERT(KALDI_ISFINITE(targets.Sum()));
  
    // buffer initialization,
    int32 num_classes = targets.NumCols();
    if (frames_.Dim() == 0) {
      frames_.Resize(num_classes);
      xentropy_.Resize(num_classes);
      entropy_.Resize(num_classes);
      correct_.Resize(num_classes);
    }
  
    // get frame_weights to GPU,
    frame_weights_ = frame_weights;
  
    // There may be frames for which the sum of targets is zero.
    // This happens in multi-lingual training when the frame
    // has target class in the softmax of another language.
    // We 'switch-off' such frames by masking the 'frame_weights_',
    target_sum_.Resize(targets.NumRows());
    target_sum_.AddColSumMat(1.0, targets, 0.0);
    frame_weights_.MulElements(target_sum_);
  
    // compute derivative wrt. activations of last layer of neurons,
    *diff = net_out;
    diff->AddMat(-1.0, targets);
    diff->MulRowsVec(frame_weights_);  // weighting,
  
    // count frames per class,
    frames_aux_ = targets;
    frames_aux_.MulRowsVec(frame_weights_);
    frames_.AddRowSumMat(1.0, CuMatrix<double>(frames_aux_));
  
    // evaluate the frame-level classification,
    net_out.FindRowMaxId(&max_id_out_);  // find max in nn-output
    targets.FindRowMaxId(&max_id_tgt_);  // find max in targets
    CountCorrectFramesWeighted(max_id_out_, max_id_tgt_,
                               frame_weights_, &correct_);
  
    // calculate cross_entropy (in GPU),
    xentropy_aux_ = net_out;  // y
    xentropy_aux_.Add(1e-20);  // avoid log(0)
    xentropy_aux_.ApplyLog();  // log(y)
    xentropy_aux_.MulElements(targets);  // t*log(y)
    xentropy_aux_.MulRowsVec(frame_weights_);  // w*t*log(y)
    xentropy_.AddRowSumMat(-1.0, CuMatrix<double>(xentropy_aux_));
  
    // caluculate entropy (in GPU),
    entropy_aux_ = targets;  // t
    entropy_aux_.Add(1e-20);  // avoid log(0)
    entropy_aux_.ApplyLog();  // log(t)
    entropy_aux_.MulElements(targets);  // t*log(t)
    entropy_aux_.MulRowsVec(frame_weights_);  // w*t*log(t)
    entropy_.AddRowSumMat(-1.0, CuMatrix<double>(entropy_aux_));
  
    // progressive loss reporting
    if (opts_.loss_report_frames > 0) {
      frames_progress_ += frame_weights_.Sum();
      xentropy_progress_ += -xentropy_aux_.Sum();
      entropy_progress_ += -entropy_aux_.Sum();
  
      KALDI_ASSERT(KALDI_ISFINITE(xentropy_progress_));
      KALDI_ASSERT(KALDI_ISFINITE(entropy_progress_));
  
      if (frames_progress_ > opts_.loss_report_frames) {
        // loss value,
        double progress_value =
          (xentropy_progress_ - entropy_progress_) / frames_progress_;
  
        // time-related info (fps is weighted),
        double time_now = timer_.Elapsed();
        double fps = frames_progress_ / (time_now - elapsed_seconds_);
        double elapsed_hours = time_now / 3600;
        elapsed_seconds_ = time_now; // store,
  
        // print,
        KALDI_LOG << "ProgressLoss[last "
                  << static_cast<int>(frames_progress_/100/3600) << "h of "
                  << static_cast<int>(frames_.Sum()/100/3600) << "h]: "
                  << progress_value << " (Xent)"
                  << ", fps=" << fps
                  << std::setprecision(3)
                  << ", elapsed " << elapsed_hours << "h";
        // store,
        loss_vec_.push_back(progress_value);
        // reset,
        frames_progress_ = 0;
        xentropy_progress_ = 0.0;
        entropy_progress_ = 0.0;
      }
    }
  }
  
  
  void Xent::Eval(const VectorBase<BaseFloat> &frame_weights,
                  const CuMatrixBase<BaseFloat> &net_out,
                  const Posterior &post,
                  CuMatrix<BaseFloat> *diff) {
    int32 num_frames = net_out.NumRows(),
      num_pdf = net_out.NumCols();
    KALDI_ASSERT(num_frames == post.size());
  
    // convert posterior to matrix,
    PosteriorToMatrix(post, num_pdf, &tgt_mat_);
  
    // call the other eval function,
    Eval(frame_weights, net_out, tgt_mat_, diff);
  }
  
  
  std::string Xent::Report() {
    double loss_value =
      (xentropy_.Sum() - entropy_.Sum()) / frames_.Sum();
    std::ostringstream oss;
    oss << "AvgLoss: " << loss_value << " (Xent), "
        << "[AvgXent: " << xentropy_.Sum() / frames_.Sum()
        << ", AvgTargetEnt: " << entropy_.Sum() / frames_.Sum()
        << "]" << std::endl;
  
    oss << "progress: [";
    std::copy(loss_vec_.begin(), loss_vec_.end(),
              std::ostream_iterator<float>(oss, " "));
    oss << "]" << std::endl;
  
    double frame_accuracy = 100.0 * correct_.Sum() / frames_.Sum();
    oss << "FRAME_ACCURACY >> " << frame_accuracy << "% <<" << std::endl;
  
    return oss.str();
  }
  
  
  std::string Xent::ReportPerClass() {
    std::ostringstream oss;
    oss << "PER-CLASS PERFORMANCE:" << std::endl;
    oss << "@@@ Frames per-class:" << frames_;
    // get inverted counts,
    CuVector<double> inv_frames(frames_);
    inv_frames.Add(0.5);  // avoid 0-frames,
    inv_frames.ApplyPow(-1.0);
    // loss, kl = xentropy-entropy,
    CuVector<double> loss(xentropy_);
    loss.AddVec(-1.0, entropy_);
    loss.MulElements(inv_frames);
    oss << "@@@ Loss per-class:" << loss;
    // frame accuracy (assuming targets are binary),
    CuVector<double> frm_accu(correct_);
    frm_accu.MulElements(inv_frames);
    frm_accu.Scale(100.0);
    oss << "@@@ Frame-accuracy per-class:" << frm_accu;
    //
    return oss.str();
  }
  
  
  /* Mse */
  
  void Mse::Eval(const VectorBase<BaseFloat> &frame_weights,
                 const CuMatrixBase<BaseFloat>& net_out,
                 const CuMatrixBase<BaseFloat>& target,
                 CuMatrix<BaseFloat>* diff) {
    // check inputs,
    KALDI_ASSERT(net_out.NumCols() == target.NumCols());
    KALDI_ASSERT(net_out.NumRows() == target.NumRows());
    KALDI_ASSERT(net_out.NumRows() == frame_weights.Dim());
  
    KALDI_ASSERT(KALDI_ISFINITE(frame_weights.Sum()));
    KALDI_ASSERT(KALDI_ISFINITE(net_out.Sum()));
    KALDI_ASSERT(KALDI_ISFINITE(target.Sum()));
  
    int32 num_frames = frame_weights.Sum();
    KALDI_ASSERT(num_frames >= 0.0);
  
    // get frame_weights to GPU,
    frame_weights_ = frame_weights;
  
    // compute derivative w.r.t. neural nerwork outputs
    *diff = net_out;  // y
    diff->AddMat(-1.0, target);  // (y - t)
    diff->MulRowsVec(frame_weights_);  // weighting,
  
    // Compute MeanSquareError loss of mini-batch
    diff_pow_2_ = *diff;
    diff_pow_2_.MulElements(diff_pow_2_);  // (y - t)^2
    diff_pow_2_.MulRowsVec(frame_weights_);  // w*(y - t)^2
    double mean_square_error = 0.5 * diff_pow_2_.Sum();  // sum the matrix,
  
    KALDI_ASSERT(KALDI_ISFINITE(mean_square_error));
  
    // accumulate
    loss_ += mean_square_error;
    frames_ += num_frames;
  
    // progressive loss reporting
    if (opts_.loss_report_frames > 0) {
      frames_progress_ += num_frames;
      loss_progress_ += mean_square_error;
      if (frames_progress_ > opts_.loss_report_frames) {
        KALDI_LOG << "ProgressLoss[last "
                  << static_cast<int>(frames_progress_/100/3600) << "h of "
                  << static_cast<int>(frames_/100/3600) << "h]: "
                  << loss_progress_/frames_progress_ << " (Mse)";
        // store
        loss_vec_.push_back(loss_progress_/frames_progress_);
        // reset
        frames_progress_ = 0;
        loss_progress_ = 0.0;
      }
    }
  }
  
  
  void Mse::Eval(const VectorBase<BaseFloat> &frame_weights,
                 const CuMatrixBase<BaseFloat>& net_out,
                 const Posterior& post,
                 CuMatrix<BaseFloat>* diff) {
    int32 num_frames = net_out.NumRows(),
      num_nn_outputs = net_out.NumCols();
    KALDI_ASSERT(num_frames == post.size());
  
    // convert posterior to matrix,
    PosteriorToMatrix(post, num_nn_outputs, &tgt_mat_);
  
    // call the other eval function,
    Eval(frame_weights, net_out, tgt_mat_, diff);
  }
  
  
  std::string Mse::Report() {
    // compute root mean square,
    int32 num_tgt = diff_pow_2_.NumCols();
    BaseFloat root_mean_square = sqrt(loss_/frames_/num_tgt);
    // build the message,
    std::ostringstream oss;
    oss << "AvgLoss: " << loss_/frames_ << " (Mse), "
        << "[RMS " << root_mean_square << ", frames "
        << frames_ << "]" << std::endl;
    oss << "progress: [";
    std::copy(loss_vec_.begin(), loss_vec_.end(),
              std::ostream_iterator<float>(oss, " "));
    oss << "]" << std::endl;
    return oss.str();
  }
  
  
  /* MultiTaskLoss */
  
  void MultiTaskLoss::InitFromString(const std::string& s) {
    std::vector<std::string> v;
    SplitStringToVector(s, ",:" /* delimiter */, false, &v);
  
    KALDI_ASSERT((v.size()-1) % 3 == 0);  // triplets,
    KALDI_ASSERT(v[0] == "multitask");  // header,
  
    // parse the definition of multitask loss,
    std::vector<std::string>::iterator it(v.begin()+1);  // skip header,
    for ( ; it != v.end(); ++it) {
      // type,
      if (*it == "xent") {
        loss_vec_.push_back(new Xent(opts_));
      } else if (*it == "mse") {
        loss_vec_.push_back(new Mse(opts_));
      } else {
        KALDI_ERR << "Unknown objective function code : " << *it;
      }
      ++it;
      // dim,
      int32 dim;
      if (!ConvertStringToInteger(*it, &dim)) {
        KALDI_ERR << "Cannot convert 'dim' " << *it << " to integer!";
      }
      loss_dim_.push_back(dim);
      ++it;
      // weight,
      BaseFloat weight;
      if (!ConvertStringToReal(*it, &weight)) {
        KALDI_ERR << "Cannot convert 'weight' " << *it << " to integer!";
      }
      KALDI_ASSERT(weight >= 0.0);
      loss_weights_.push_back(weight);
    }
  
    // build vector with starting-point offsets,
    loss_dim_offset_.resize(loss_dim_.size()+1, 0);  // 1st zero stays,
    for (int32 i = 1; i <= loss_dim_.size(); i++) {
      loss_dim_offset_[i] = loss_dim_offset_[i-1] + loss_dim_[i-1];
    }
  
    // sanity check,
    KALDI_ASSERT(loss_vec_.size() > 0);
    KALDI_ASSERT(loss_vec_.size() == loss_dim_.size());
    KALDI_ASSERT(loss_vec_.size() == loss_weights_.size());
  }
  
  void MultiTaskLoss::Eval(const VectorBase<BaseFloat> &frame_weights,
              const CuMatrixBase<BaseFloat>& net_out,
              const Posterior& post,
              CuMatrix<BaseFloat>* diff) {
    int32 num_frames = net_out.NumRows(),
      num_output = net_out.NumCols();
    KALDI_ASSERT(num_frames == post.size());
    KALDI_ASSERT(num_output == loss_dim_offset_.back());  // sum of loss-dims,
  
    // convert posterior to matrix,
    PosteriorToMatrix(post, num_output, &tgt_mat_);
  
    // allocate diff matrix,
    diff->Resize(num_frames, num_output);
  
    /// One vector of frame_weights per loss-function,
    /// The original frame weights are multiplied with
    /// a mask of `defined targets' according to the 'Posterior'.
    std::vector<Vector<BaseFloat> > frmwei_have_tgt;
    for (int32 l = 0; l < loss_vec_.size(); l++) {
      // copy original weights,
      frmwei_have_tgt.push_back(Vector<BaseFloat>(frame_weights));
      // We need to mask-out the frames for which the 'posterior' is not defined (= is empty):
      int32 loss_beg = loss_dim_offset_[l];   // first column of loss target,
      int32 loss_end = loss_dim_offset_[l+1]; // (last+1) column of loss target,
      for (int32 f = 0; f < num_frames; f++) {
        bool tgt_defined = false;
        for (int32 p = 0; p < post[f].size(); p++) {
          if (post[f][p].first >= loss_beg && post[f][p].first < loss_end) {
            tgt_defined = true;
            break;
          }
        }
        if (!tgt_defined) {
          frmwei_have_tgt[l](f) = 0.0; // set zero_weight for the frame with no targets!
        }
      }
    }
  
    // call the vector of loss functions,
    CuMatrix<BaseFloat> diff_aux;
    for (int32 l = 0; l < loss_vec_.size(); l++) {
      loss_vec_[l]->Eval(frmwei_have_tgt[l],
        net_out.ColRange(loss_dim_offset_[l], loss_dim_[l]),
        tgt_mat_.ColRange(loss_dim_offset_[l], loss_dim_[l]),
        &diff_aux);
      // Scale the gradients,
      diff_aux.Scale(loss_weights_[l]);
      // Copy to diff,
      diff->ColRange(loss_dim_offset_[l], loss_dim_[l]).CopyFromMat(diff_aux);
    }
  }
  
  std::string MultiTaskLoss::Report() {
    // calculate overall loss (weighted),
    BaseFloat overall_loss = AvgLoss();
    // copy the loss-values into a vector,
    std::vector<BaseFloat> loss_values;
    for (int32 i = 0; i < loss_vec_.size(); i++) {
      loss_values.push_back(loss_vec_[i]->AvgLoss());
    }
  
    // build the message,
    std::ostringstream oss;
    oss << "MultiTaskLoss, with " << loss_vec_.size()
        << " parallel loss functions." << std::endl;
    // individual loss reports first,
    for (int32 i = 0; i < loss_vec_.size(); i++) {
      oss << "Loss " << i+1 << ", " << loss_vec_[i]->Report() << std::endl;
    }
  
    // overall loss is last,
    oss << "Loss (OVERALL), "
        << "AvgLoss: " << overall_loss << " (MultiTaskLoss), "
        << "weights " << loss_weights_ << ", "
        << "values " << loss_values << std::endl;
  
    return oss.str();
  }
  
  BaseFloat MultiTaskLoss::AvgLoss() {
    BaseFloat ans(0.0);
    for (int32 i = 0; i < loss_vec_.size(); i++) {
      BaseFloat val = loss_weights_[i] * loss_vec_[i]->AvgLoss();
      if (!KALDI_ISFINITE(val)) {
        KALDI_WARN << "Loss " << i+1 << ", has bad objective function value '"
                   << val << "', using 0.0 instead.";
        val = 0.0;
      }
      ans += val;
    }
    return ans;
  }
  
  }  // namespace nnet1
  }  // namespace kaldi