// nnet3/nnet-discriminative-diagnostics.cc
  
  // Copyright  2012-2015    Johns Hopkins University (author: Daniel Povey)
  // Copyright  2014-2015    Vimal Manohar
  
  // 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 "nnet3/nnet-discriminative-diagnostics.h"
  #include "nnet3/nnet-utils.h"
  #include "nnet3/discriminative-training.h"
  
  namespace kaldi {
  namespace nnet3 {
  
  NnetDiscriminativeComputeObjf::NnetDiscriminativeComputeObjf(
      const NnetComputeProbOptions &nnet_config,
      const discriminative::DiscriminativeOptions &discriminative_config,
      const TransitionModel &tmodel,
      const VectorBase<BaseFloat> &priors,
      const Nnet &nnet):
      nnet_config_(nnet_config),
      discriminative_config_(discriminative_config),
      tmodel_(tmodel),
      log_priors_(priors),
      nnet_(nnet),
      compiler_(nnet, nnet_config_.optimize_config),
      deriv_nnet_(NULL),
      num_minibatches_processed_(0) {
    log_priors_.ApplyLog();
    if (nnet_config_.compute_deriv) {
      deriv_nnet_ = new Nnet(nnet_);
      ScaleNnet(0.0, deriv_nnet_);
      SetNnetAsGradient(deriv_nnet_); // force simple update
    }
  }
  
  const Nnet& NnetDiscriminativeComputeObjf::GetDeriv() const {
    if (deriv_nnet_ == NULL)
      KALDI_ERR << "GetDeriv() called when no derivatives were requested.";
    return *deriv_nnet_;
  }
  
  NnetDiscriminativeComputeObjf::~NnetDiscriminativeComputeObjf() {
    delete deriv_nnet_;  // delete does nothing if pointer is NULL.
  }
  
  void NnetDiscriminativeComputeObjf::Reset() {
    num_minibatches_processed_ = 0;
    objf_info_.clear();
    if (deriv_nnet_) {
      ScaleNnet(0.0, deriv_nnet_);
      SetNnetAsGradient(deriv_nnet_);
    }
  }
  
  void NnetDiscriminativeComputeObjf::Compute(const NnetDiscriminativeExample &eg) {
    bool need_model_derivative = nnet_config_.compute_deriv,
        store_component_stats = false;
    bool use_xent_regularization = (discriminative_config_.xent_regularize != 0.0),
        use_xent_derivative = false;
  
    ComputationRequest request;
    GetDiscriminativeComputationRequest(nnet_, eg,
                                        need_model_derivative,
                                        store_component_stats,
                                        use_xent_regularization, use_xent_derivative,
                                        &request);
    std::shared_ptr<const NnetComputation> computation = compiler_.Compile(request);
    NnetComputer computer(nnet_config_.compute_config, *computation,
                          nnet_, deriv_nnet_);
    // give the inputs to the computer object.
    computer.AcceptInputs(nnet_, eg.inputs);
    computer.Run();
    this->ProcessOutputs(eg, &computer);
    if (nnet_config_.compute_deriv)
      computer.Run();
  }
  
  void NnetDiscriminativeComputeObjf::ProcessOutputs(
                                      const NnetDiscriminativeExample &eg,
                                      NnetComputer *computer) {
    // There will normally be just one output here, named 'output',
    // but the code is more general than this.
    std::vector<NnetDiscriminativeSupervision>::const_iterator iter = eg.outputs.begin(),
        end = eg.outputs.end();
    for (; iter != end; ++iter) {
      const NnetDiscriminativeSupervision &sup = *iter;
      int32 node_index = nnet_.GetNodeIndex(sup.name);
      if (node_index < 0 ||
          !nnet_.IsOutputNode(node_index))
        KALDI_ERR << "Network has no output named " << sup.name;
  
      const CuMatrixBase<BaseFloat> &nnet_output = computer->GetOutput(sup.name);
  
      bool use_xent = (discriminative_config_.xent_regularize != 0.0);
      std::string xent_name = sup.name + "-xent";  // typically "output-xent".
      CuMatrix<BaseFloat> nnet_output_deriv, xent_deriv;
  
      if (nnet_config_.compute_deriv)
        nnet_output_deriv.Resize(nnet_output.NumRows(), nnet_output.NumCols(),
                                 kUndefined);
  
      if (use_xent)
        xent_deriv.Resize(nnet_output.NumRows(), nnet_output.NumCols(),
                          kUndefined);
  
      if (objf_info_.count(sup.name) == 0)
        objf_info_.insert(std::make_pair(sup.name,
            discriminative::DiscriminativeObjectiveInfo(discriminative_config_)));
  
      discriminative::DiscriminativeObjectiveInfo *stats = &(objf_info_[sup.name]);
  
      discriminative::ComputeDiscriminativeObjfAndDeriv(discriminative_config_,
                                                        tmodel_, log_priors_,
                                                        sup.supervision, nnet_output,
                                                        stats,
                                                        (nnet_config_.compute_deriv ?
                                                         &nnet_output_deriv : NULL),
                                                        (use_xent ? &xent_deriv : NULL));
  
      if (nnet_config_.compute_deriv)
        computer->AcceptInput(sup.name, &nnet_output_deriv);
  
      if (use_xent) {
        if (objf_info_.count(xent_name) == 0)
          objf_info_.insert(std::make_pair(xent_name,
            discriminative::DiscriminativeObjectiveInfo(discriminative_config_)));
        discriminative::DiscriminativeObjectiveInfo &xent_stats = objf_info_[xent_name];
  
        // this block computes the cross-entropy objective.
        const CuMatrixBase<BaseFloat> &xent_output = computer->GetOutput(xent_name);
        // at this point, xent_deriv is posteriors derived from the numerator
        // computation.  note, xent_deriv has a factor of 'supervision.weight',
        // but so does tot_weight.
        BaseFloat xent_objf = TraceMatMat(xent_output, xent_deriv, kTrans);
        xent_stats.tot_t_weighted += stats->tot_t_weighted;
        xent_stats.tot_objf += xent_objf;
      }
  
      num_minibatches_processed_++;
    }
  }
  
  bool NnetDiscriminativeComputeObjf::PrintTotalStats() const {
    bool ans = false;
    unordered_map<std::string, discriminative::DiscriminativeObjectiveInfo, StringHasher>::const_iterator
        iter, end;
    iter = objf_info_.begin();
    end = objf_info_.end();
    for (; iter != end; ++iter) {
      const std::string &name = iter->first;
      int32 node_index = nnet_.GetNodeIndex(name);
      KALDI_ASSERT(node_index >= 0);
      const discriminative::DiscriminativeObjectiveInfo &info = iter->second;
      BaseFloat tot_weight = info.tot_t_weighted;
      BaseFloat tot_objective = info.TotalObjf(
          discriminative_config_.criterion);
  
      info.PrintAll(discriminative_config_.criterion);
  
      if (info.tot_l2_term == 0.0) {
        KALDI_LOG << "Overall " << discriminative_config_.criterion
                  << " objective for '"
                  << name << "' is "
                  << (tot_objective / tot_weight)
                  << " per frame, "
                  << "over " << tot_weight << " frames.";
      } else {
        KALDI_LOG << "Overall " << discriminative_config_.criterion
                  << " objective for '"
                  << name << "' is "
                  << (tot_objective / tot_weight)
                  << " + " << (info.tot_l2_term / tot_weight)
                  << " per frame, "
                  << "over " << tot_weight << " frames.";
      }
  
      if (tot_weight > 0)
        ans = true;
    }
    return ans;
  }
  
  const discriminative::DiscriminativeObjectiveInfo* NnetDiscriminativeComputeObjf::GetObjective(
      const std::string &output_name) const {
    unordered_map<std::string, discriminative::DiscriminativeObjectiveInfo, StringHasher>::const_iterator
        iter = objf_info_.find(output_name);
    if (iter != objf_info_.end())
      return &(iter->second);
    else
      return NULL;
  }
  
  } // namespace nnet3
  } // namespace kaldi