// chain/chain-numerator.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_CHAIN_CHAIN_NUMERATOR_H_
  #define KALDI_CHAIN_CHAIN_NUMERATOR_H_
  
  #include <vector>
  #include <map>
  
  #include "base/kaldi-common.h"
  #include "util/common-utils.h"
  #include "fstext/fstext-lib.h"
  #include "tree/context-dep.h"
  #include "lat/kaldi-lattice.h"
  #include "matrix/kaldi-matrix.h"
  #include "hmm/transition-model.h"
  #include "chain/chain-supervision.h"
  #include "cudamatrix/cu-matrix.h"
  #include "cudamatrix/cu-array.h"
  
  namespace kaldi {
  namespace chain {
  
  
  // This class is responsible for the forward-backward of the 'supervision'
  // (numerator) FST.
  //
  // note: the supervision.weight is ignored by this class, you have to apply
  // it externally.
  // Because the supervision FSTs are quite skinny, i.e. have very few paths for
  // each frame, it's feasible to do this computation on the CPU, and that's what
  // we do.  We transfer from/to the GPU only the things that we need.
  
  class NumeratorComputation {
  
   public:
  
    /// Initialize the objcect.  Note: we expect the 'nnet_output' to have the
    /// same number of rows as supervision.num_frames * supervision.num_sequences,
    /// and the same number of columns as the 'label-dim' of the supervision
    /// object (which will be the NumPdfs() of the transition model); but the
    /// ordering of the rows of 'nnet_output' is not the same as the ordering of
    /// frames in paths in the 'supervision' object (which has all frames of the
    /// 1st sequence first, then the 2nd sequence, and so on).  Instead, the
    /// frames in 'nnet_output' are ordered as: first the first frame of each
    /// sequence, then the second frame of each sequence, and so on.  This is more
    /// convenient both because the nnet3 code internally orders them that way,
    /// and because this makes it easier to order things in the way that class
    /// SingleHmmForwardBackward needs (we can just transpose, instead of doing a
    /// 3d tensor rearrangement).
    NumeratorComputation(const Supervision &supervision,
                         const CuMatrixBase<BaseFloat> &nnet_output);
  
    // TODO: we could enable a Viterbi mode.
  
    // Does the forward computation.  Returns the total log-prob multiplied
    // by supervision_.weight.
    BaseFloat Forward();
  
    // Does the backward computation and (efficiently) adds the derivative of the
    // nnet output w.r.t. the (log-prob times supervision_.weight times
    // deriv_weight) to 'nnet_output_deriv'.
    void Backward(CuMatrixBase<BaseFloat> *nnet_output_deriv);
  
   private:
  
    const Supervision &supervision_;
  
    // state times of supervision_.fst.
    std::vector<int32> fst_state_times_;
  
  
    // the exp of the neural net output.
    const CuMatrixBase<BaseFloat> &nnet_output_;
  
  
    // 'fst_output_indexes' contains an entry for each arc in the supervision FST, in
    // the order you'd get them if you visit each arc of each state in order.
    // the contents of fst_output_indexes_ are indexes into nnet_output_indexes_
    // and nnet_logprobs_.
    std::vector<int32> fst_output_indexes_;
  
    // nnet_output_indexes is a list of (row, column) indexes that we need to look
    // up in nnet_output_ for the forward-backward computation.  The order is
    // arbitrary, but indexes into this vector appear in fst_output_indexes;
    // and it's important that each pair only appear once (in order for the
    // derivatives to be summed properly).
    CuArray<Int32Pair> nnet_output_indexes_;
  
    // the log-probs obtained from lookup in the nnet output, on the CPU.  This
    // vector has the same size as nnet_output_indexes_.  In the backward
    // computation, the storage is re-used for derivatives.
    Vector<BaseFloat> nnet_logprobs_;
  
    // derivatives w.r.t. the nnet logprobs.  These can be interpreted as
    // occupation probabilities.
    Vector<BaseFloat> nnet_logprob_derivs_;
  
    // The log-alpha value (forward probability) for each state in the lattices.
    Vector<double> log_alpha_;
  
    // The total pseudo-log-likelihood from the forward-backward.
    double tot_log_prob_;
  
    // The log-beta value (backward probability) for each state in the lattice
    Vector<double> log_beta_;
  
    // This function creates fst_output_indexes_ and nnet_output_indexes_.
    void ComputeLookupIndexes();
  
    // convert time-index in the FST to a row-index in the nnet-output (to account
    // for the fact that the sequences are interleaved in the nnet-output).
    inline int32 ComputeRowIndex(int32 t, int32 frames_per_sequence,
                                 int32 num_sequences) {
      return t / frames_per_sequence +
          num_sequences * (t % frames_per_sequence);
    }
  
  };
  
  
  
  
  }  // namespace chain
  }  // namespace kaldi
  
  #endif  // KALDI_CHAIN_CHAIN_NUMERATOR_H_