// transform/fmllr-raw.cc
  
  // Copyright 2013  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.
  
  #include <utility>
  #include <vector>
  using std::vector;
  
  #include "transform/fmllr-raw.h"
  #include "transform/fmllr-diag-gmm.h"
  
  namespace kaldi {
  
  FmllrRawAccs::FmllrRawAccs(int32 raw_dim,
                             int32 model_dim,
                             const Matrix<BaseFloat> &full_transform):
      raw_dim_(raw_dim),
      model_dim_(model_dim) {
    if (full_transform.NumCols() != full_transform.NumRows() &&
        full_transform.NumCols() != full_transform.NumRows() + 1) {
      KALDI_ERR << "Expecting full LDA+MLLT transform to be square or d by d+1 "
                << "(make sure you are including rejected rows).";
    }
    if (raw_dim <= 0 || full_transform.NumRows() % raw_dim != 0)
      KALDI_ERR << "Raw feature dimension is invalid " << raw_dim
                << "(must be positive and divide feature dimension)";
    int32 full_dim = full_transform.NumRows();
    full_transform_ = full_transform.Range(0, full_dim, 0, full_dim);
    transform_offset_.Resize(full_dim);
    if (full_transform_.NumCols() == full_dim + 1)
      transform_offset_.CopyColFromMat(full_transform_, full_dim);
    
    int32 full_dim2 = ((full_dim+1)*(full_dim+2))/2;
    count_ = 0.0;
  
    temp_.Resize(full_dim + 1);
    Q_.Resize(model_dim + 1, full_dim + 1);
    S_.Resize(model_dim + 1, full_dim2);
  
    single_frame_stats_.s.Resize(full_dim + 1);
    single_frame_stats_.transformed_data.Resize(full_dim);
    single_frame_stats_.count = 0.0;
    single_frame_stats_.a.Resize(model_dim);
    single_frame_stats_.b.Resize(model_dim);
  }
  
  
  bool FmllrRawAccs::DataHasChanged(const VectorBase<BaseFloat> &data) const {
    KALDI_ASSERT(data.Dim() == FullDim());
    return !data.ApproxEqual(single_frame_stats_.s.Range(0, FullDim()), 0.0);
  }
  
  void FmllrRawAccs::CommitSingleFrameStats() {
    // Commit the stats for this from (in SingleFrameStats).
    int32 model_dim = ModelDim(), full_dim = FullDim();
    SingleFrameStats &stats = single_frame_stats_;
    if (stats.count == 0.0) return;
  
    count_ += stats.count;
  
    // a_ext and b_ext are a and b extended with the count,
    // which we'll later use to reconstruct the full stats for
    // the rejected dimensions.
    Vector<double> a_ext(model_dim + 1), b_ext(model_dim + 1);
    a_ext.Range(0, model_dim).CopyFromVec(stats.a);
    b_ext.Range(0, model_dim).CopyFromVec(stats.b);
    a_ext(model_dim) = stats.count;
    b_ext(model_dim) = stats.count;
    Q_.AddVecVec(1.0, a_ext, Vector<double>(stats.s));
  
    temp_.SetZero();
    temp_.AddVec2(1.0, stats.s);
    int32 full_dim2 = ((full_dim + 1) * (full_dim + 2)) / 2;
    SubVector<double> temp_vec(temp_.Data(), full_dim2);
    S_.AddVecVec(1.0, b_ext, temp_vec);
  }
  
  void FmllrRawAccs::InitSingleFrameStats(const VectorBase<BaseFloat> &data) {
    SingleFrameStats &stats = single_frame_stats_;
    int32 full_dim = FullDim();
    KALDI_ASSERT(data.Dim() == full_dim);
    stats.s.Range(0, full_dim).CopyFromVec(data);
    stats.s(full_dim) = 1.0;
    stats.transformed_data.AddMatVec(1.0, full_transform_, kNoTrans, data, 0.0);
    stats.transformed_data.AddVec(1.0, transform_offset_);
    stats.count = 0.0;
    stats.a.SetZero();
    stats.b.SetZero();
  }
  
  
  BaseFloat FmllrRawAccs::AccumulateForGmm(const DiagGmm &gmm,
                                           const VectorBase<BaseFloat> &data,
                                           BaseFloat weight) {
    int32 model_dim = ModelDim(), full_dim = FullDim();
    KALDI_ASSERT(data.Dim() == full_dim &&
                 "Expect raw, spliced data, which should have same dimension as "
                 "full transform.");
    if (DataHasChanged(data)) {
      // this is part of our mechanism to accumulate certain sub-parts of
      // the computation for each frame, to avoid excessive compute.
      CommitSingleFrameStats();
      InitSingleFrameStats(data);
    }
    SingleFrameStats &stats = single_frame_stats_;
  
    SubVector<BaseFloat> projected_data(stats.transformed_data, 0, model_dim);
  
    int32 num_gauss = gmm.NumGauss();
    Vector<BaseFloat> posterior(num_gauss);
    BaseFloat log_like = gmm.ComponentPosteriors(projected_data, &posterior);
    posterior.Scale(weight);
    // Note: AccumulateFromPosteriors takes the original, spliced data,
    // and returns the log-like of the rejected dimensions.
    AccumulateFromPosteriors(gmm, data, posterior);
  
    // Add the likelihood of the rejected dimensions to the objective function
    // (assume zero-mean, unit-variance Gaussian; the LDA should have any offset
    // required to ensure this).
    if (full_dim > model_dim) {
      SubVector<BaseFloat> rejected_data(stats.transformed_data,
                                         model_dim, full_dim - model_dim);
      log_like += -0.5 * (VecVec(rejected_data, rejected_data)
                          + (full_dim - model_dim) * M_LOG_2PI);
    }
    return log_like;
  }
  
  /*
    // Extended comment here.
    //
    // Let x_t(i) be the fully processed feature, dimension i (with fMLLR transform
    //  and LDA transform), but *without* any offset term from the LDA, which
    //  it's more convenient to view as an offset in the model.
    //
    //
    // For a given dimension i (either accepted or rejected), the auxf can
    // be expressed as a quadratic function of x_t(i).  We ultimately will want to
    // express x_t(i) as a linear function of the parameters of the linearized
    // fMLLR transform matrix.  Some notation:
    //    Let l be the linearized transform matrix, i.e. the concatenation of the
    //       m rows, each of length m+1, of the fMLLR transform.
    //    Let n be the number of frames we splice together each time.
    //    Let s_t be the spliced-together features on time t, with a one appended;
    //       it will have n blocks each of size m, followed by a 1.  (dim is n*m + 1).
    //     
    // x(i) [note, this is the feature without any LDA offset], is bilinear in the
    //      transform matrix and the features, so:
    //
    // x(i) = l^T M_i s_t, where s_t is the spliced features on time t,
    //          with a 1 appended
    //   [we need to compute M_i but we know the function is bilinear so it exists].
    //
    // The auxf can be written as:
    // F = sum_i sum_t  a_{ti} x(i) - 0.5  b_{ti} x(i)^2 
    //   = sum_i sum_t  a_{ti} x(i) - 0.5  b_{ti} x(i)^2
    //   = sum_i sum_t  a_{ti} (l^T M_i s_t)  -  0.5 b_{ti} (l^T M_i s_t )^2
    //   = sum_i l^T M_i q_i  +  l^T M_i S_i M_i^T l 
    //  where
    //     q_i = sum_t a_{ti} s_t, and
    //     S_i = sum_t b_{ti} s_t s_t^T
    //   [Note that we only need store S_i for the model-dim plus one, because
    //    all the rejected dimensions have the same value]
    //
    //     We define a matrix Q whose rows are q_d, with
    //       Q = \sum_t d_t s_t^T
    //    [The Q we actually store as stats will use a modified form of d that
    //     has a 1 for all dimensions past the model dim, to avoid redundancy;
    //     we'll reconstruct the true Q from this later on.]
    //     
    //
    // What is M_i?  Working it out is a little tedious.
    //  Note: each M_i (from i = 0 ... full_dim) is of
    //    dimension (raw_dim*(raw_dim+1)) by full_dim + 1
    // 
    // We want to express x(i) [we forget the subscript "t" sometimes],
    // as a bilinear function of l and s_t.
    //    We have x(i) = l^T M_i s.
    //
    // The (j,k)'th component of M_i is the term in x(i) that corresponds to the j'th
    // component of l and the k'th of s.
  
    // Before defining M_i, let us define N_i, where l^t N_i s will equal the spliced and
    // transformed pre-LDA features of dimension i.  the N's have the same dimensions as the
    // M's.
    //
    // We'll first define the j,k'th component of N_i, as this is easier; we'll then define the M_i
    // as combinations of N_i.
    //
    // For a given i, j and k, the value of n_{i,j,k} will be as follows:
    //   We first decompose index j into j1, j2 (both functions of
    //    the original index j), where
    //    j1 corresponds to the row-index of the fMLLR transform, j2 to the col-index.
    //   We next decompose i into i1, i2, where i1 corresponds to the splicing number
    //   (0...n-1), and i2 corresponds to the cepstral index.
    //
    //   If (j1 != i2) then n_{ijk} == 0.
    //
    //   Elsif k corresponds to the last element [i.e. k == m * n], then this m_{ijk} corresponds
    //   to the effect of the j'th component of l for zero input, so:
    //     If j2 == m (i.e. this the offset term in the fMLLR matrix), then
    //       n_{ijk} = 1.0,
    //     Else
    //       n_{ijk} = 0.0
    //     Fi
    //
    //   Else:
    //     Decompose k into k1, k2, where k1 = 0.. n-1 is the splicing index, and k2 = 0...m-1 is
    //      the cepstral index.
    //     If k1 != i1 then
    //       n_{ijk} = 0.0
    //     elsif k2 != j2 then
    //       n_{ijk} = 0.0
    //     else
    //       n_{ijk} = 1.0
    //     fi
    //    Endif
    //    Now,  M_i will be defined as sum_i T_{ij} N_j, where T_{ij} are the elements of the
    //     LDA+MLLT transform (but excluding any linear offset, which gets accounted for by
    //     c_i, above).
    //
    //  Now suppose we want to express the auxiliary function in a simpler form
    //  as l^T v - 0.5 l^T W l, where v and W are the "simple" linear and quadratic stats,
    //  we can do so with:
    //     v = \sum_i M_i q_i   
    //  and
    //     W = \sum_i M_i S_i M_i^T
    //
    */
  
  void FmllrRawAccs::AccumulateFromPosteriors(
      const DiagGmm &diag_gmm,
      const VectorBase<BaseFloat> &data,
      const VectorBase<BaseFloat> &posterior) {
    // The user may call this function directly, even though we also
    // call it from AccumulateForGmm(), so check again:
    if (DataHasChanged(data)) { 
      CommitSingleFrameStats();
      InitSingleFrameStats(data);
    }
    
    int32  model_dim = ModelDim();
  
    SingleFrameStats &stats = single_frame_stats_;
    
    // The quantities a and b describe the diagonal auxiliary function
    // for each of the retained dimensions in the transformed space--
    // in the format F = \sum_d alpha(d) x(d)  -0.5 beta(d) x(d)^2,
    // where x(d) is the d'th dimensional fully processed feature.
    // For d, see the comment-- it's alpha processed to take into
    // account any offset in the LDA.  Note that it's a reference.
    //
    Vector<double> a(model_dim), b(model_dim);
    
    int32 num_comp = diag_gmm.NumGauss();
    
    double count = 0.0; // data-count contribution from this frame.
  
    // Note: we could do this using matrix-matrix operations instead of
    // row by row.  In the end it won't really matter as this is not
    // the slowest part of the computation.
    for (size_t m = 0; m < num_comp; m++) {
      BaseFloat this_post = posterior(m);
      if (this_post != 0.0) {
        count += this_post;
        a.AddVec(this_post, diag_gmm.means_invvars().Row(m));
        b.AddVec(this_post, diag_gmm.inv_vars().Row(m));
      }
    }
    // Correct "a" for any offset term in the LDA transform-- we view it as
    // the opposite offset in the model [note: we'll handle the rejected dimensions
    // in update time.]  Here, multiplying the element of "b" (which is the
    // weighted inv-vars) by transform_offset_, and subtracting the result from
    // a, is like subtracting the transform-offset from the original means
    // (because a contains the means times inv-vars_.
    Vector<double> offset(transform_offset_.Range(0, model_dim));
    a.AddVecVec(-1.0, b, offset, 1.0);
    stats.a.AddVec(1.0, a);
    stats.b.AddVec(1.0, b);
    stats.count += count;
  }
  
  
  void FmllrRawAccs::Update(const FmllrRawOptions &opts,
                            MatrixBase<BaseFloat> *raw_fmllr_mat,
                            BaseFloat *objf_impr,
                            BaseFloat *count) {
    // First commit any pending stats from the last frame.
    if (single_frame_stats_.count != 0.0)
      CommitSingleFrameStats();
    
    if (this->count_ < opts.min_count) {
      KALDI_WARN << "Not updating (raw) fMLLR since count " << this->count_
                 << " is less than min count " << opts.min_count;
      *objf_impr = 0.0;
      *count = this->count_;
      return;
    }
    KALDI_ASSERT(raw_fmllr_mat->NumRows() == RawDim() &&
                 raw_fmllr_mat->NumCols() == RawDim() + 1 &&
                 !raw_fmllr_mat->IsZero());
    Matrix<double> fmllr_mat(*raw_fmllr_mat); // temporary, double-precision version
                                              // of matrix.
  
  
    Matrix<double> linear_stats; // like K in diagonal update.
    std::vector<SpMatrix<double> > diag_stats; // like G in diagonal update.
                                               // Note: we will invert these.
    std::vector<std::vector<Matrix<double> > > off_diag_stats; // these will
    // contribute to the linear term.
  
    Vector<double> simple_linear_stats;
    SpMatrix<double> simple_quadratic_stats;
    ConvertToSimpleStats(&simple_linear_stats, &simple_quadratic_stats);
    
    ConvertToPerRowStats(simple_linear_stats, simple_quadratic_stats,
                         &linear_stats, &diag_stats, &off_diag_stats);
  
    try {
      for (size_t i = 0; i < diag_stats.size(); i++) {
        diag_stats[i].Invert();
      }
    } catch (...) {
      KALDI_WARN << "Error inverting stats matrices for fMLLR "
                 << "[min-count too small?  Bad data?], not updating.";
      return;
    }
    
    int32 raw_dim = RawDim(), splice_width = SpliceWidth();
    
    double effective_beta = count_ * splice_width; // We "count" the determinant
    // splice_width times in the objective function.
  
    double auxf_orig = GetAuxf(simple_linear_stats, simple_quadratic_stats,
                               fmllr_mat);
    for (int32 iter = 0; iter < opts.num_iters; iter++) {
      for (int32 row = 0; row < raw_dim; row++) {
        SubVector<double> this_row(fmllr_mat, row);
        Vector<double> this_linear(raw_dim + 1);  // Here, k_i is the linear term
        // in the auxf expressed as a function of this row.
        this_linear.CopyFromVec(linear_stats.Row(row));
        for (int32 row2 = 0; row2 < raw_dim; row2++) {
          if (row2 != row) {
            if (row2 < row) {
              this_linear.AddMatVec(-1.0, off_diag_stats[row][row2], kNoTrans,
                                    fmllr_mat.Row(row2), 1.0);
            } else {
              // We won't have the element [row][row2] stored, but use symmetry.
              this_linear.AddMatVec(-1.0, off_diag_stats[row2][row], kTrans,
                                    fmllr_mat.Row(row2), 1.0);
            }
          }
        }
        FmllrInnerUpdate(diag_stats[row],
                         this_linear,
                         effective_beta,
                         row,
                         &fmllr_mat);
      }
      if (GetVerboseLevel() >= 2) {
        double cur_auxf = GetAuxf(simple_linear_stats, simple_quadratic_stats,
                                   fmllr_mat),
            auxf_change = cur_auxf - auxf_orig;
        KALDI_VLOG(2) << "Updating raw fMLLR: objf improvement per frame was "
                      << (auxf_change / this->count_) << " over "
                      << this->count_ << " frames, by the " << iter
                      << "'th iteration";
      }
    }
    double auxf_final = GetAuxf(simple_linear_stats, simple_quadratic_stats,
                                fmllr_mat),
        auxf_change = auxf_final - auxf_orig;
    *count = this->count_;
    KALDI_VLOG(1) << "Updating raw fMLLR: objf improvement per frame was "
                  << (auxf_change / this->count_) << " over "
                  << this->count_ << " frames.";
    if (auxf_final > auxf_orig) {
      *objf_impr = auxf_change;
      *count = this->count_;
      raw_fmllr_mat->CopyFromMat(fmllr_mat);
    } else {
      *objf_impr = 0.0;
      // don't update "raw_fmllr_mat"
    }
  }
  
  void FmllrRawAccs::SetZero() {
    count_ = 0.0;
    single_frame_stats_.count = 0.0;
    single_frame_stats_.s.SetZero();
    Q_.SetZero();
    S_.SetZero();
  }
  
  // Compute the M_i quantities, needed in the update.  This function could be
  // greatly speeded up but I don't think it's the limiting factor.
  void FmllrRawAccs::ComputeM(std::vector<Matrix<double> > *M) const {
    int32 full_dim = FullDim(), raw_dim = RawDim(),
        raw_dim2 = raw_dim * (raw_dim + 1);
    M->resize(full_dim);
    for (int32 i = 0; i < full_dim; i++)
      (*M)[i].Resize(raw_dim2, full_dim + 1);  
  
    // the N's are simpler matrices from which we'll interpolate the M's.
    // In this loop we imagine w are computing the vector of N's, but
    // when we get each element, if it's nonzero we propagate it straight
    // to the M's.
    for (int32 i = 0; i < full_dim; i++) {
      // i is index after fMLLR transform; i1 is splicing index,
      // i2 is cepstral index.
      int32 i1 = i / raw_dim, i2 = i % raw_dim;
      for (int32 j = 0; j < raw_dim2; j++) {
        // j1 is row-index of fMLLR transform, j2 is column-index
        int32 j1 = j / (raw_dim + 1), j2 = j % (raw_dim + 1);
        for (int32 k = 0; k < full_dim + 1; k++) {
          BaseFloat n_ijk;
          if (j1 != i2) {
            n_ijk = 0.0;
          } else if (k == full_dim) {
            if (j2 == raw_dim) // offset term in fMLLR matrix.
              n_ijk = 1.0;
            else
              n_ijk = 0.0;
          } else {
            // k1 is splicing index, k2 is cepstral idnex.
            int32 k1 = k / raw_dim, k2 = k % raw_dim;
            if (k1 != i1 || k2 != j2)
              n_ijk = 0.0;
            else
              n_ijk = 1.0;
          }
          if (n_ijk != 0.0)
            for (int32 l = 0; l < full_dim; l++)
              (*M)[l](j, k) += n_ijk * full_transform_(l, i);
        }
      }
    }
  }
  
  void FmllrRawAccs::ConvertToSimpleStats(
      Vector<double> *simple_linear_stats,
      SpMatrix<double> *simple_quadratic_stats) const {
    std::vector<Matrix<double> > M;
    ComputeM(&M);
  
    int32 full_dim = FullDim(), raw_dim = RawDim(), model_dim = ModelDim(),
        raw_dim2 = raw_dim * (raw_dim + 1),
        full_dim2 = ((full_dim+1)*(full_dim+2))/2;
    simple_linear_stats->Resize(raw_dim2);
    simple_quadratic_stats->Resize(raw_dim2);
    for (int32 i = 0; i < full_dim; i++) {
      Vector<double> q_i(full_dim + 1);
      SpMatrix<double> S_i(full_dim + 1);
      SubVector<double> S_i_vec(S_i.Data(), full_dim2);
      if (i < model_dim) {
        q_i.CopyFromVec(Q_.Row(i));
        S_i_vec.CopyFromVec(S_.Row(i));
      } else {
        q_i.CopyFromVec(Q_.Row(model_dim)); // The last row contains stats proportional
        // to "count", which we need to modify to be correct.
        q_i.Scale(-transform_offset_(i)); // These stats are zero (corresponding to
        // a zero-mean model) if there is no offset in the LDA transform.  Note:
        // the two statements above are the equivalent, for the rejected dims,
        // of the statement "a.AddVecVec(-1.0, b, offset);" for the kept ones.
        // 
        S_i_vec.CopyFromVec(S_.Row(model_dim)); // these are correct, and
        // all the same (corresponds to unit variance).
      }
      // The equation v = \sum_i M_i q_i:
      simple_linear_stats->AddMatVec(1.0, M[i], kNoTrans, q_i, 1.0);
      // The equation W = \sum_i M_i S_i M_i^T
      // Here, M[i] is quite sparse, so AddSmat2Sp will be faster.
      simple_quadratic_stats->AddSmat2Sp(1.0, M[i], kNoTrans, S_i, 1.0);
    }
  }
  
  // See header for comment.
  void FmllrRawAccs::ConvertToPerRowStats(
      const Vector<double> &simple_linear_stats,
      const SpMatrix<double> &simple_quadratic_stats_sp,
      Matrix<double> *linear_stats,
      std::vector<SpMatrix<double> > *diag_stats,
      std::vector<std::vector<Matrix<double> > > *off_diag_stats) const {
  
    // get it as a Matrix, which makes it easier to extract sub-parts.
    Matrix<double> simple_quadratic_stats(simple_quadratic_stats_sp);
  
    linear_stats->Resize(RawDim(), RawDim() + 1);
    linear_stats->CopyRowsFromVec(simple_linear_stats);
    diag_stats->resize(RawDim());
    off_diag_stats->resize(RawDim());
  
    // Set *diag_stats
    int32 rd1 = RawDim() + 1;
    for (int32 i = 0; i < RawDim(); i++) {
      SubMatrix<double> this_diag(simple_quadratic_stats,
                                  i * rd1, rd1,
                                  i * rd1, rd1);
      (*diag_stats)[i].Resize(RawDim() + 1);
      (*diag_stats)[i].CopyFromMat(this_diag, kTakeMean);
    }    
    
    for (int32 i = 0; i < RawDim(); i++) {
      (*off_diag_stats)[i].resize(i);
      for (int32 j = 0; j < i; j++) {
        SubMatrix<double> this_off_diag(simple_quadratic_stats,
                                        i * rd1, rd1,
                                        j * rd1, rd1);
        (*off_diag_stats)[i][j] = this_off_diag;
      }
    }
  }
  
  double FmllrRawAccs::GetAuxf(const Vector<double> &simple_linear_stats,
                               const SpMatrix<double> &simple_quadratic_stats,
                               const Matrix<double> &fmllr_mat) const {
    // linearize transform...
    int32 raw_dim = RawDim(), spice_width = SpliceWidth();
    Vector<double> fmllr_vec(raw_dim * (raw_dim + 1));
    fmllr_vec.CopyRowsFromMat(fmllr_mat);
    SubMatrix<double> square_part(fmllr_mat, 0, raw_dim,
                                  0, raw_dim);
    double logdet = square_part.LogDet();
    return VecVec(fmllr_vec, simple_linear_stats) -
        0.5 * VecSpVec(fmllr_vec, simple_quadratic_stats, fmllr_vec) +
        logdet * spice_width * count_;
  }
  
  
  
  } // namespace kaldi