// nnet3/nnet-diagnostics.h

// Copyright    2015  Johns Hopkins University (author: Daniel Povey)

// 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.

#ifndef KALDI_NNET3_NNET_DIAGNOSTICS_H_
#define KALDI_NNET3_NNET_DIAGNOSTICS_H_

#include "nnet3/nnet-example.h"
#include "nnet3/nnet-computation.h"
#include "nnet3/nnet-compute.h"
#include "nnet3/nnet-optimize.h"
#include "nnet3/nnet-example-utils.h"
#include "nnet3/nnet-training.h"

namespace kaldi {
namespace nnet3 {


struct SimpleObjectiveInfo {
  double tot_weight;
  double tot_objective;
  SimpleObjectiveInfo(): tot_weight(0.0),
                         tot_objective(0.0) { }
};

/* This is used to store more detailed information about the objective,
 * which will be used to compute accuracy per dimension.
 * This might be sensible only for classification tasks.
 */
struct PerDimObjectiveInfo: public SimpleObjectiveInfo {
  // Counts for each of the classes in the output.
  // In the simplest cases, this might be the number of frames for each class.
  Vector<BaseFloat> tot_weight_vec;

  // Objective contribution per-class
  Vector<BaseFloat> tot_objective_vec;
};


struct NnetComputeProbOptions {
  bool debug_computation;
  bool compute_deriv;
  bool compute_accuracy;
  // note: the component stats, if stored, will be stored in the derivative nnet
  // (c.f. GetDeriv()) if compute_deriv is true; otherwise, you should use the
  // constructor of NnetComputeProb that takes a pointer to the nnet, and the
  // stats will be stored there.
  bool store_component_stats;
  
  bool compute_per_dim_accuracy;

  NnetOptimizeOptions optimize_config;
  NnetComputeOptions compute_config;
  CachingOptimizingCompilerOptions compiler_config;
  NnetComputeProbOptions():
      debug_computation(false),
      compute_deriv(false),
      compute_accuracy(true),
      store_component_stats(false),
      compute_per_dim_accuracy(false) { }
  void Register(OptionsItf *opts) {
    // compute_deriv is not included in the command line options
    // because it's not relevant for nnet3-compute-prob.
    // store_component_stats is not included in the command line
    // options because it's not relevant for nnet3-compute-prob.
    opts->Register("debug-computation", &debug_computation, "If true, turn on "
                   "debug for the actual computation (very verbose!)");
    opts->Register("compute-accuracy", &compute_accuracy, "If true, compute "
                   "accuracy values as well as objective functions");
    opts->Register("compute-per-dim-accuracy", &compute_per_dim_accuracy,
                   "If true, compute accuracy values per-dim");

    // register the optimization options with the prefix "optimization".
    ParseOptions optimization_opts("optimization", opts);
    optimize_config.Register(&optimization_opts);
    // register the compiler options with the prefix "compiler".
    ParseOptions compiler_opts("compiler", opts);
    compiler_config.Register(&compiler_opts);
    // register the compute options with the prefix "computation".
    ParseOptions compute_opts("computation", opts);
    compute_config.Register(&compute_opts);
  }
};


/** This class is for computing cross-entropy and accuracy values in a neural
    network, for diagnostics.
    Note: because we put a "logsoftmax" component in the nnet, the actual
    objective function becomes linear at the output, but the printed messages
    reflect the fact that it's the cross-entropy objective.
 */
class NnetComputeProb {
 public:
  // does not store a reference to 'config' but does store one to 'nnet'.
  NnetComputeProb(const NnetComputeProbOptions &config,
                  const Nnet &nnet);

  // This version of the constructor may only be called if
  // config.store_component_stats == true and config.compute_deriv == false;
  // it means it will store the component stats in 'nnet'.  In this
  // case you should call ZeroComponentStats(nnet) first if you want
  // the stats to be zeroed first.
  NnetComputeProb(const NnetComputeProbOptions &config,
                  Nnet *nnet);


  // Reset the likelihood stats, and the derivative stats (if computed).
  void Reset();

  // compute objective on one minibatch.
  void Compute(const NnetExample &eg);

  // Prints out the final stats, and return true if there was a nonzero count.
  bool PrintTotalStats() const;

  // returns the objective-function info for this output name (e.g. "output"),
  // or NULL if there is no such info.
  const SimpleObjectiveInfo *GetObjective(const std::string &output_name) const;

  // This function returns the total objective over all output nodes recorded here, and
  // outputs to 'tot_weight' the total weight (typically the number of frames)
  // corresponding to it.
  double GetTotalObjective(double *tot_weight) const;

  // if config.compute_deriv == true, returns a reference to the
  // computed derivative.  Otherwise crashes.
  const Nnet &GetDeriv() const;

  ~NnetComputeProb();
 private:
  void ProcessOutputs(const NnetExample &eg,
                      NnetComputer *computer);

  NnetComputeProbOptions config_;
  const Nnet &nnet_;

  bool deriv_nnet_owned_;
  Nnet *deriv_nnet_;
  CachingOptimizingCompiler compiler_;

  // this is only for diagnostics.
  int32 num_minibatches_processed_;

  unordered_map<std::string, SimpleObjectiveInfo, StringHasher> objf_info_;

  unordered_map<std::string, PerDimObjectiveInfo, StringHasher> accuracy_info_;
};


/**
   This function computes the frame accuracy for this minibatch.  It interprets
   the supervision information in "supervision" as labels or soft labels; it
   picks the maximum element in each row and treats that as the label for
   purposes of computing the accuracy (in situations where you would care about
   the accuracy, there will normally be just one nonzero label).  The
   hypothesized labels are computed by taking the neural net output (supplied as
   a CuMatrix), and finding the maximum element in each row.
   See also the function ComputeObjectiveFunction, declared in nnet-training.h.

   @param [in] supervision  The supervision information (no elements may be
                     negative); only the maximum in each row matters (although
                     we expect that usually there will be just one nonzero
                     element in each row); and the sum of each row is
                     interpreted as a weighting factor (although we expect that
                     this sum will usually be one).

  @param [in] nnet_output   The neural net output must have the same dimensions
                     as the supervision.  Only the index of the maximum value in
                     each row matters.  Ties will be broken in an unspecified
                     way.
   @param [out] tot_weight  The sum of the values in the supervision matrix
   @param [out] tot_accuracy  The total accuracy, equal to the sum over all row
                     indexes r such that the maximum column index of row r of
                     supervision and nnet_output is the same, of the sum of 
                     the r'th row of supervision (i.e. the row's weight).
   @param [out] tot_weight_vec  If non-NULL, we write to this location
                    the counts per-class in the supervision matrix.
                    This is expected to have the same dimension as the 
                    corresponding output in the network. 
   @param [out] tot_accuracy_vec  If non-NULL, we write to this location 
                    the accuracy per-class. For index j, 
                    the value is equal to the sum 
                    over all row indexes r such that the maximum column index 
                    of row r of supervision is j and nnet_output is also j,
                    of the sum of the r'th row of supervision 
                    (i.e. the row's weight)

*/
void ComputeAccuracy(const GeneralMatrix &supervision,
                     const CuMatrixBase<BaseFloat> &nnet_output,
                     BaseFloat *tot_weight,
                     BaseFloat *tot_accuracy,
                     VectorBase<BaseFloat> *tot_weight_vec = NULL,
                     VectorBase<BaseFloat> *tot_accuracy_vec = NULL);


} // namespace nnet3
} // namespace kaldi

#endif // KALDI_NNET3_NNET_DIAGNOSTICS_H_