// chain/chain-generic-numerator.cc

// Copyright      2017   Hossein Hadian

// 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 "chain/chain-generic-numerator.h"
#include "chain/chain-kernels-ansi.h"

#include <iterator>
#include <limits>
#include <algorithm>

namespace kaldi {
namespace chain {

// GenericNumeratorComputation is responsible for the forward-backward of the
// end-to-end 'supervision' (numerator) FST. It is used in chain-training.cc
// (similar to NumeratorComputation) to compute the numerator derivatives
// for end-to-end training 'supervision's.

GenericNumeratorComputation::GenericNumeratorComputation(
    const Supervision &supervision,
    const CuMatrixBase<BaseFloat> &nnet_output):
    supervision_(supervision),
    nnet_output_(nnet_output) {
  KALDI_ASSERT(supervision.num_sequences *
               supervision.frames_per_sequence == nnet_output.NumRows() &&
               supervision.label_dim == nnet_output.NumCols());

  using std::vector;
  int num_sequences = supervision_.num_sequences;
  KALDI_ASSERT(supervision_.e2e_fsts.size() == num_sequences);

  // Find the maximum number of HMM states and then
  // initialize final probs, alpha, and beta.
  int max_num_hmm_states = 0;
  for (int i = 0; i < num_sequences; i++) {
    KALDI_ASSERT(supervision_.e2e_fsts[i].Properties(fst::kIEpsilons, true)
                 == 0);
    if (supervision_.e2e_fsts[i].NumStates() > max_num_hmm_states)
      max_num_hmm_states = supervision_.e2e_fsts[i].NumStates();
  }
  final_probs_.Resize(num_sequences, max_num_hmm_states);

  // Initialize incoming transitions for easy access
  in_transitions_.resize(num_sequences);  // indexed by seq, state
  out_transitions_.resize(num_sequences);  // indexed by seq, state
  for (int seq = 0; seq < num_sequences; seq++) {
    in_transitions_[seq] = vector<vector<DenominatorGraphTransition> >(
        supervision_.e2e_fsts[seq].NumStates());
    out_transitions_[seq] = vector<vector<DenominatorGraphTransition> >(
        supervision_.e2e_fsts[seq].NumStates());
  }

  offsets_.Resize(num_sequences);
  std::unordered_map<int32, MatrixIndexT> pdf_to_index;
  int32 pdf_stride = nnet_output_.Stride();
  int32 view_stride = nnet_output_.Stride() * num_sequences;
  pdf_to_index.reserve(view_stride);
  nnet_output_stride_ = pdf_stride;
  for (int seq = 0; seq < num_sequences; seq++) {
    for (int32 s = 0; s < supervision_.e2e_fsts[seq].NumStates(); s++) {
      final_probs_(seq, s)= -supervision_.e2e_fsts[seq].Final(s).Value();
      BaseFloat offset = 0.0;
      if (s == 0) {
        for (fst::ArcIterator<fst::StdVectorFst> aiter(
                 supervision_.e2e_fsts[seq], s);
             !aiter.Done();
             aiter.Next())
          if (aiter.Value().weight.Value() > offset)
            offset = aiter.Value().weight.Value();
        offsets_(seq) = offset;
      }

      for (fst::ArcIterator<fst::StdVectorFst> aiter(
             supervision_.e2e_fsts[seq], s);
           !aiter.Done();
           aiter.Next()) {
        const fst::StdArc &arc = aiter.Value();
        DenominatorGraphTransition transition;
        transition.transition_prob = -(arc.weight.Value() - offset);

        int32 pdf_id = arc.ilabel - 1;  // note: the FST labels were pdf-id plus one.

        // remap  to a unique index in the remapped space
        pdf_id = pdf_id + seq * pdf_stride;
        KALDI_ASSERT(pdf_id < view_stride);

        if (pdf_to_index.find(pdf_id) == pdf_to_index.end()) {
          index_to_pdf_.push_back(pdf_id);
          pdf_to_index[pdf_id] = index_to_pdf_.size() - 1;
        }

        transition.pdf_id = pdf_to_index[pdf_id];
        transition.hmm_state = s;
        in_transitions_[seq][arc.nextstate].push_back(transition);
        transition.hmm_state = arc.nextstate;
        out_transitions_[seq][s].push_back(transition);
      }
    }
  }
}


void GenericNumeratorComputation::AlphaFirstFrame(int seq,
                                                  Matrix<BaseFloat> *alpha) {
  const int32 num_frames = supervision_.frames_per_sequence,
              num_states = supervision_.e2e_fsts[seq].NumStates();
  alpha->Resize(num_frames + 1,  num_states + 1, kSetZero);
  alpha->Set(-std::numeric_limits<BaseFloat>::infinity());
  (*alpha)(0, 0) = 0.0;
  (*alpha)(0, num_states) = 0.0;
}


void GenericNumeratorComputation::CopySpecificPdfsIndirect(
                                    const CuMatrixBase<BaseFloat> &nnet_output,
                                    const std::vector<MatrixIndexT> &indices,
                                    Matrix<BaseFloat> *out) {
  KALDI_ASSERT(nnet_output_stride_ == nnet_output_.Stride());
  const int32 num_sequences = supervision_.num_sequences,
              frames_per_sequence = supervision_.frames_per_sequence;

  const BaseFloat *starting_ptr = nnet_output.RowData(0);
  const int view_stride = num_sequences * nnet_output.Stride();

  const CuSubMatrix<BaseFloat> sequence_view(starting_ptr,
                                             frames_per_sequence,
                                             view_stride,
                                             view_stride);

  CuArray<MatrixIndexT> indices_gpu(indices);
  CuMatrix<BaseFloat> required_pdfs(frames_per_sequence,
                                        indices.size());

  required_pdfs.CopyCols(sequence_view, indices_gpu);
  out->Swap(&required_pdfs);
}

// The alpha computation for some 0 < t <= num_time_steps_.
BaseFloat GenericNumeratorComputation::AlphaRemainingFrames(int seq,
                                              const Matrix<BaseFloat> &probs,
                                              Matrix<BaseFloat> *alpha) {
  // Define some variables to make things nicer
  const int32 num_sequences = supervision_.num_sequences,
              num_frames = supervision_.frames_per_sequence;

  KALDI_ASSERT(seq >= 0 && seq < num_sequences);

  // variables for log_likelihood computation
  double log_scale_product = 0,
         log_prob_product = 0;

  for (int t = 1; t <= num_frames; ++t) {
    const BaseFloat *probs_tm1 = probs.RowData(t - 1);
    BaseFloat *alpha_t = alpha->RowData(t);
    const BaseFloat *alpha_tm1 = alpha->RowData(t - 1);

    for (int32 h = 0; h < supervision_.e2e_fsts[seq].NumStates(); h++) {
      for (auto tr = in_transitions_[seq][h].begin();
          tr != in_transitions_[seq][h].end(); ++tr) {
        BaseFloat transition_prob = tr->transition_prob;
        int32 pdf_id = tr->pdf_id,
              prev_hmm_state = tr->hmm_state;
        BaseFloat prob = probs_tm1[pdf_id];
        alpha_t[h] = LogAdd(alpha_t[h],
            alpha_tm1[prev_hmm_state] + transition_prob + prob);
      }
    }
    double sum = alpha_tm1[alpha->NumCols() - 1];
    SubMatrix<BaseFloat> alpha_t_mat(*alpha, t, 1, 0,
                                      alpha->NumCols() - 1);
    alpha_t_mat.Add(-sum);
    sum = alpha_t_mat.LogSumExp();

    alpha_t[alpha->NumCols() - 1] = sum;
    log_scale_product += sum;
  }
  SubMatrix<BaseFloat> last_alpha(*alpha, alpha->NumRows() - 1, 1,
                                  0, alpha->NumCols() - 1);
  SubVector<BaseFloat> final_probs(final_probs_.RowData(seq),
                                   alpha->NumCols() - 1);

  // adjust last_alpha
  double sum = (*alpha)(alpha->NumRows() - 1, alpha->NumCols() - 1);
  log_scale_product -= sum;
  last_alpha.AddVecToRows(1.0, final_probs);
  sum = last_alpha.LogSumExp();
  (*alpha)(alpha->NumRows() - 1, alpha->NumCols() - 1) = sum;

  // second part of criterion
  log_prob_product = sum - offsets_(seq);

  return log_prob_product + log_scale_product;
}

bool GenericNumeratorComputation::ForwardBackward(
                                 BaseFloat *total_loglike,
                                 CuMatrixBase<BaseFloat> *nnet_output_deriv) {
  KALDI_ASSERT(total_loglike != NULL);
  KALDI_ASSERT(nnet_output_deriv != NULL);
  KALDI_ASSERT(nnet_output_deriv->NumCols() == nnet_output_.NumCols());
  KALDI_ASSERT(nnet_output_deriv->NumRows() == nnet_output_.NumRows());

  BaseFloat partial_loglike = 0;
  const int32 num_sequences = supervision_.num_sequences;

  bool ok = true;
  Matrix<BaseFloat> alpha;
  Matrix<BaseFloat> beta;
  Matrix<BaseFloat> probs;
  Matrix<BaseFloat> derivs;

  // We selectively copy only those pdfs we need
  CopySpecificPdfsIndirect(nnet_output_, index_to_pdf_, &probs);

  derivs.Resize(probs.NumRows(), probs.NumCols());
  derivs.Set(-std::numeric_limits<BaseFloat>::infinity());

  for (int seq = 0; seq < num_sequences; ++seq) {
    // Forward part
    AlphaFirstFrame(seq, &alpha);
    partial_loglike += AlphaRemainingFrames(seq, probs, &alpha);

    // Backward part
    BetaLastFrame(seq, alpha, &beta);
    BetaRemainingFrames(seq, probs, alpha, &beta, &derivs);
    if (GetVerboseLevel() >= 1)
      ok = ok && CheckValues(seq, probs, alpha, beta, derivs);
  }
  // Transfer and add the derivatives to the values in the matrix
  AddSpecificPdfsIndirect(&derivs, index_to_pdf_, nnet_output_deriv);
  *total_loglike = partial_loglike;
  return ok;
}

BaseFloat GenericNumeratorComputation::ComputeObjf() {
  BaseFloat partial_loglike = 0;
  const int32 num_sequences = supervision_.num_sequences;

  Matrix<BaseFloat> alpha;
  Matrix<BaseFloat> probs;

  // We selectively copy only those pdfs we need
  CopySpecificPdfsIndirect(nnet_output_, index_to_pdf_, &probs);

  for (int seq = 0; seq < num_sequences; ++seq) {
    // Forward part
    AlphaFirstFrame(seq, &alpha);
    partial_loglike += AlphaRemainingFrames(seq, probs, &alpha);
  }
  return partial_loglike;
}

BaseFloat GenericNumeratorComputation::GetTotalProb(
                                          const Matrix<BaseFloat> &alpha) {
  return alpha(alpha.NumRows() - 1, alpha.NumCols() - 1);
}

void GenericNumeratorComputation::BetaLastFrame(int seq,
                                                const Matrix<BaseFloat> &alpha,
                                                Matrix<BaseFloat> *beta) {
  // Sets up the beta quantity on the last frame (frame ==
  // frames_per_sequence_).  Note that the betas we use here contain a
  // 1/(tot-prob) factor in order to simplify the backprop.
  const int32 num_frames = supervision_.frames_per_sequence,
              num_states = supervision_.e2e_fsts[seq].NumStates();
  float tot_prob = GetTotalProb(alpha);

  beta->Resize(2, num_states);
  beta->Set(-std::numeric_limits<BaseFloat>::infinity());

  SubVector<BaseFloat> beta_mat(beta->RowData(num_frames % 2), num_states);
  SubVector<BaseFloat> final_probs(final_probs_.RowData(seq), num_states);

  BaseFloat inv_tot_prob = -tot_prob;
  beta_mat.Set(inv_tot_prob);
  beta_mat.AddVec(1.0, final_probs);
}

void GenericNumeratorComputation::BetaRemainingFrames(int seq,
                                                const Matrix<BaseFloat> &probs,
                                                const Matrix<BaseFloat> &alpha,
                                                Matrix<BaseFloat> *beta,
                                                Matrix<BaseFloat> *derivs) {
  const int32
      num_sequences = supervision_.num_sequences,
      num_frames = supervision_.frames_per_sequence,
      num_states = supervision_.e2e_fsts[seq].NumStates();
  KALDI_ASSERT(seq >= 0 && seq < num_sequences);

  for (int t = num_frames - 1; t >= 0; --t) {
    const BaseFloat *alpha_t = alpha.RowData(t),
        *beta_tp1 = beta->RowData((t + 1) % 2),
        *probs_t = probs.RowData(t);
    BaseFloat *log_prob_deriv_t = derivs->RowData(t),
        *beta_t = beta->RowData(t % 2);

    BaseFloat inv_arbitrary_scale = alpha_t[num_states];
    for (int32 h = 0; h < supervision_.e2e_fsts[seq].NumStates(); h++) {
      BaseFloat tot_variable_factor;
      tot_variable_factor = -std::numeric_limits<BaseFloat>::infinity();
      for (auto tr = out_transitions_[seq][h].begin();
               tr != out_transitions_[seq][h].end(); ++tr) {
        BaseFloat transition_prob = tr->transition_prob;
        int32 pdf_id = tr->pdf_id,
            next_hmm_state = tr->hmm_state;
        BaseFloat variable_factor = transition_prob +
            beta_tp1[next_hmm_state] +
            probs_t[pdf_id] - inv_arbitrary_scale;
        tot_variable_factor = LogAdd(tot_variable_factor,
                                     variable_factor);

        BaseFloat occupation_prob = variable_factor + alpha_t[h];
        log_prob_deriv_t[pdf_id] = LogAdd(log_prob_deriv_t[pdf_id],
                                           occupation_prob);
      }
      beta_t[h] = tot_variable_factor;
    }
  }
}


void GenericNumeratorComputation::AddSpecificPdfsIndirect(
                                 Matrix<BaseFloat> *logprobs,
                                 const std::vector<MatrixIndexT> &indices,
                                 CuMatrixBase<BaseFloat> *output) {
  const int32 num_sequences = supervision_.num_sequences,
              frames_per_sequence = supervision_.frames_per_sequence;

  const int view_stride = output->Stride() * num_sequences;

  KALDI_ASSERT(frames_per_sequence * num_sequences == output->NumRows());

  CuMatrix<BaseFloat> specific_pdfs;
  specific_pdfs.Swap(logprobs);
  specific_pdfs.ApplyExp();
  specific_pdfs.Scale(supervision_.weight);

  std::vector<MatrixIndexT> indices_expanded(view_stride, -1);
  for (int i = 0; i < indices.size(); ++i) {
    int pdf_index = indices[i];
    int sequence_local_pdf_index = pdf_index % nnet_output_stride_;
    int sequence_index = pdf_index / nnet_output_stride_;
    pdf_index = sequence_local_pdf_index
                + sequence_index * output->Stride();
    KALDI_ASSERT(pdf_index < view_stride);
    KALDI_ASSERT(i < specific_pdfs.NumCols());
    indices_expanded[pdf_index] = i;
  }

  CuArray<MatrixIndexT> cu_indices(indices_expanded);
  CuSubMatrix<BaseFloat> out(output->Data(), frames_per_sequence,
                             view_stride, view_stride);

  out.AddCols(specific_pdfs, cu_indices);
}

bool GenericNumeratorComputation::CheckValues(int seq,
                                            const Matrix<BaseFloat> &probs,
                                            const Matrix<BaseFloat> &alpha,
                                            const Matrix<BaseFloat> &beta,
                                            const Matrix<BaseFloat> &derivs) const {
  const int32 num_frames = supervision_.frames_per_sequence;
  // only check the derivs for the first and last frames
  const std::vector<int32> times = {0, num_frames - 1};
  for (const int32 t: times) {
    BaseFloat deriv_sum = 0.0;
    for (int32 n = 0; n < probs.NumCols(); n++) {
      int32 pdf_stride = nnet_output_.Stride();
      int32 pdf2seq = index_to_pdf_[n] / pdf_stride;
      if (pdf2seq != seq)  // this pdf is not in the space of this sequence
        continue;
      deriv_sum += Exp(derivs(t, n));
    }

    if (!ApproxEqual(deriv_sum, 1.0)) {
      KALDI_WARN << "On time " << t
                 << " for seq " << seq << ", deriv sum "
                 << deriv_sum << " != 1.0";
      if (fabs(deriv_sum - 1.0) > 0.05 || deriv_sum - deriv_sum != 0) {
        KALDI_WARN << "Excessive error detected, will abandon this minibatch";
        return false;
      }
    }
  }
  return true;
}

}  // namespace chain
}  // namespace kaldi