// nnet3/natural-gradient-online-test.cc

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

#include "nnet3/natural-gradient-online.h"
#include "util/common-utils.h"

namespace kaldi {
namespace nnet3 {

// Simple version of OnlineNaturalGradient that we use to make
// sure it is behaving as advertised.
class OnlineNaturalGradientSimple {
 public:
  OnlineNaturalGradientSimple(): rank_(40), num_samples_history_(2000.0), alpha_(4.0),
                                epsilon_(1.0e-10), delta_(5.0e-04) { }

  void SetRank(int32 rank) { rank_ = rank; }

  void PreconditionDirections(
      CuMatrixBase<BaseFloat> *R,
      CuVectorBase<BaseFloat> *row_prod,
      BaseFloat *scale);


 private:
  BaseFloat Eta(int32 N) const;

  void PreconditionDirectionsCpu(
      MatrixBase<double> *R,
      VectorBase<double> *row_prod,
      BaseFloat *scale);


  void Init(const MatrixBase<double> &R0);

  void InitDefault(int32 D);

  int32 rank_;
  double num_samples_history_;
  double alpha_;
  double epsilon_;
  double delta_;

  // Fisher matrix defined as F_t = R_t^T diag(d_t) R_t + rho_t I.
  Vector<double> d_t_;
  Matrix<double> R_t_;
  double rho_t_;
};


void OnlineNaturalGradientSimple::PreconditionDirections(
      CuMatrixBase<BaseFloat> *R,
      CuVectorBase<BaseFloat> *row_prod,
      BaseFloat *scale) {
  Matrix<BaseFloat> R_cpu(*R);
  Vector<BaseFloat> row_prod_cpu(*row_prod);
  Matrix<double> R_cpu_dbl(R_cpu);
  Vector<double> row_prod_cpu_dbl(row_prod_cpu);
  PreconditionDirectionsCpu(&R_cpu_dbl,
                            &row_prod_cpu_dbl,
                            scale);
  row_prod_cpu.CopyFromVec(row_prod_cpu_dbl);
  R_cpu.CopyFromMat(R_cpu_dbl);
  R->CopyFromMat(R_cpu);
  row_prod->CopyFromVec(row_prod_cpu);
}

void OnlineNaturalGradientSimple::InitDefault(int32 D) {
  if (rank_ >= D) {
    KALDI_WARN << "Rank " << rank_ << " of online preconditioner is >= dim " << D
               << ", setting it to "
               << (D - 1) << " (but this is probably still too high)";
    rank_ = D - 1;
  }
  int32 R = rank_;
  R_t_.Resize(R, D);
  for (int32 r = 0; r < R; r++) {
    std::vector<int32> cols;
    for (int32 c = r; c < D; c += R)
      cols.push_back(c);
    for (int32 i = 0; i < cols.size(); i++) {
      int32 c = cols[i];
      R_t_(r, c) = (i == 0 ? 1.1 : 1.0) /
          sqrt(1.1 * 1.1 + cols.size() - 1);
    }
  }
  d_t_.Resize(R);
  d_t_.Set(epsilon_);
  rho_t_ = epsilon_;
}

void OnlineNaturalGradientSimple::Init(const MatrixBase<double> &R0) {
  int32 D = R0.NumCols(), N = R0.NumRows();
  InitDefault(D);
  int32 num_init_iters = 3;
  for (int32 i = 0; i < num_init_iters; i++) {
    CuMatrix<BaseFloat> R0_copy(R0);
    CuVector<BaseFloat> row_products(N);
    BaseFloat scale;
    PreconditionDirections(&R0_copy, &row_products, &scale);
  }
}

BaseFloat OnlineNaturalGradientSimple::Eta(int32 N) const {
  KALDI_ASSERT(num_samples_history_ > 0.0);
  BaseFloat ans = 1.0 - exp(-N / num_samples_history_);
  if (ans > 0.9) ans = 0.9;
  return ans;
}


void OnlineNaturalGradientSimple::PreconditionDirectionsCpu(
    MatrixBase<double> *X_t,
    VectorBase<double> *row_prod,
    BaseFloat *scale) {
  if (R_t_.NumRows() == 0)
    Init(*X_t);
  int32 R = R_t_.NumRows(), D = R_t_.NumCols(), N = X_t->NumRows();
  BaseFloat eta = Eta(N);

  SpMatrix<double> F_t(D);
  // F_t =(def) R_t^T D_t R_t + \rho_t I
  F_t.AddToDiag(rho_t_);
  F_t.AddMat2Vec(1.0, R_t_, kTrans, d_t_, 1.0);

  // Make sure F_t is +ve definite.
  {
    KALDI_ASSERT(d_t_.Min() > 0);
    Vector<double> eigs(D);
    F_t.Eig(&eigs, NULL);
    KALDI_ASSERT(eigs.Min() > 0);
  }

  // S_t =(def) 1/N X_t^T X_t.
  SpMatrix<double> S_t(D);
  S_t.AddMat2(1.0 / N, *X_t, kTrans, 0.0);

  // T_t =(def) \eta S_t + (1-\eta) F_t
  SpMatrix<double> T_t(D);
  T_t.AddSp(eta, S_t);
  T_t.AddSp(1.0 - eta, F_t);

  // Y_t =(def) R_t T_t
  Matrix<double> Y_t(R, D);
  Y_t.AddMatSp(1.0, R_t_, kNoTrans, T_t, 0.0);

  // Z_t =(def) Y_t Y_t^T
  SpMatrix<double> Z_t(R);
  Z_t.AddMat2(1.0, Y_t, kNoTrans, 0.0);

  Matrix<double> U_t(R, R);
  Vector<double> c_t(R);
  // decompose Z_t = U_t C_t U_t^T
  Z_t.Eig(&c_t, &U_t);
  SortSvd(&c_t, &U_t);
  double c_t_floor = pow(rho_t_ * (1.0 - eta), 2);
  int32 nf;
  c_t.ApplyFloor(c_t_floor, &nf);
  if (nf > 0) {
    KALDI_WARN << "Floored " << nf << " elements of c_t.";
  }
  // KALDI_LOG << "c_t is " << c_t;
  // KALDI_LOG << "U_t is " << U_t;
  // KALDI_LOG << "Z_t is " << Z_t;

  Vector<double> sqrt_c_t(c_t);
  sqrt_c_t.ApplyPow(0.5);
  Vector<double> inv_sqrt_c_t(sqrt_c_t);
  inv_sqrt_c_t.InvertElements();
  Matrix<double> R_t1(R, D);
  // R_{t+1} = C_t^{-0.5} U_t^T Y_t
  R_t1.AddMatMat(1.0, U_t, kTrans, Y_t, kNoTrans, 0.0);
  R_t1.MulRowsVec(inv_sqrt_c_t);

  double rho_t1 = (1.0 / (D - R)) *
      (eta * S_t.Trace() + (1.0 - eta) * (D * rho_t_ + d_t_.Sum()) - sqrt_c_t.Sum());

  Vector<double> d_t1(sqrt_c_t);
  d_t1.Add(-rho_t1);

  double floor_val = std::max(epsilon_, delta_ * sqrt_c_t.Max());
  if (rho_t1 < floor_val) {
    KALDI_WARN << "flooring rho_{t+1} to " << floor_val << ", was " << rho_t1;
    rho_t1 = floor_val;
  }
  d_t1.ApplyFloor(floor_val, &nf);
  if (nf > 0) {
    KALDI_VLOG(3) << "d_t1 was " << d_t1;
    KALDI_WARN << "Floored " << nf << " elements of d_{t+1}.";
  }
  // a check.
  if (nf == 0 && rho_t1 > floor_val) {
    double tr_F_t1 = D * rho_t1 + d_t1.Sum(), tr_T_t = T_t.Trace();
    AssertEqual(tr_F_t1, tr_T_t);
  }

  // G_t = F_t + alpha/D tr(F_t)
  SpMatrix<double> G_t(F_t);
  G_t.AddToDiag(alpha_ / D * F_t.Trace());
  SpMatrix<double> G_t_inv(G_t);
  G_t_inv.Invert();

  double beta_t = rho_t_ + alpha_/D * F_t.Trace();
  // X_hat_t = beta_t X_t G_t^{-1}.
  Matrix<double> X_hat_t(N, D);
  X_hat_t.AddMatSp(beta_t, *X_t, kNoTrans, G_t_inv, 0.0);

  double tr_x_x = TraceMatMat(*X_t, *X_t, kTrans),
      tr_Xhat_Xhat = TraceMatMat(X_hat_t, X_hat_t, kTrans);
  double gamma = (tr_Xhat_Xhat == 0 ? 1.0 : sqrt(tr_x_x / tr_Xhat_Xhat));

  X_t->CopyFromMat(X_hat_t);
  row_prod->AddDiagMat2(1.0, *X_t, kNoTrans, 0.0);
  *scale = gamma;

  // Update the parameters
  rho_t_ = rho_t1;
  d_t_.CopyFromVec(d_t1);
  R_t_.CopyFromMat(R_t1);

  KALDI_VLOG(3) << "rho_t_ = " << rho_t_;
  KALDI_VLOG(3) << "d_t_ = " << d_t_;
  KALDI_VLOG(3) << "R_t_ = " << R_t_;


  { // check that R_t_ R_t_^T = I.
    SpMatrix<double> unit(R);
    unit.AddMat2(1.0, R_t_, kNoTrans, 0.0);
    if (!unit.IsUnit(1.0e-03)) {
      KALDI_WARN  << "R is not orthogonal, reorthogonalizing.";
      for (int32 i = 0; i < R; i++) {
        SubVector<double> row(R_t_, i);
        for (int32 j = 0; j < i; j++) {
          SubVector<double> row_j(R_t_, j);
          row.AddVec(-VecVec(row_j, row), row_j);
        }
        row.Scale(1.0 / row.Norm(2.0));
      }
    }
    unit.AddMat2(1.0, R_t_, kNoTrans, 0.0);
    KALDI_ASSERT(unit.IsUnit(1.0e-03));
  }
}


void UnitTestPreconditionDirectionsOnline() {
  MatrixIndexT R = 1 + Rand() % 30,  // rank of correction
      N = (2 * R) + Rand() % 30,  // batch size
      D = R + 1 + Rand() % 20; // problem dimension.  Must be > R.

  // Test sometimes with features that are all-zero or all-one; this will
  // help to make sure low-rank or zero input doesn't crash the code.
  bool zero = false;
  bool one = false;
  if (Rand() % 3 == 0) zero = true;
  //else if (Rand() % 2 == 0) one = true;

  CuVector<BaseFloat> row_prod1(N);
  BaseFloat gamma1, gamma2;
  BaseFloat big_eig_factor = RandInt(1, 20);
  big_eig_factor = big_eig_factor * big_eig_factor;
  Vector<BaseFloat> big_eig_vector(D);
  big_eig_vector.SetRandn();
  big_eig_vector.Scale(big_eig_factor);

  OnlineNaturalGradientSimple preconditioner1;
  OnlineNaturalGradient preconditioner2;
  preconditioner1.SetRank(R);
  preconditioner2.SetRank(R);
  preconditioner2.TurnOnDebug();

  int32 num_iters = 100;
  for (int32 iter = 0; iter < num_iters; iter++) {
    Matrix<BaseFloat> M_cpu(N, D);
    if (one) M_cpu.Set(1.0);
    else if (!zero) {
      M_cpu.SetRandn();
      Vector<BaseFloat> rand_vec(N);
      rand_vec.SetRandn();
      M_cpu.AddVecVec(1.0, rand_vec, big_eig_vector);
    }
    CuMatrix<BaseFloat> M(M_cpu);

    CuMatrix<BaseFloat> Mcopy1(M), Mcopy2(M);

    preconditioner1.PreconditionDirections(&Mcopy1, &row_prod1, &gamma1);

    preconditioner2.PreconditionDirections(&Mcopy2, &gamma2);

    BaseFloat trace1 = TraceMatMat(M, M, kTrans),
        trace2 = TraceMatMat(Mcopy1, Mcopy1, kTrans);
    AssertEqual(trace1, trace2 * gamma2 * gamma2, 1.0e-02);

    AssertEqual(Mcopy1, Mcopy2);
    AssertEqual(gamma1, gamma2, 1.0e-02);

    // make sure positive definite
    CuVector<BaseFloat> inner_prods(M.NumRows());
    inner_prods.AddDiagMatMat(1.0, M, kNoTrans, Mcopy1, kTrans, 0.0);
    KALDI_ASSERT(inner_prods.Min() >= 0.0);
  }
  return;
}


} // namespace nnet3
} // namespace kaldi


int main() {
  using namespace kaldi;
  using namespace kaldi::nnet3;
  for (int32 loop = 0; loop < 2; loop++) {
#if HAVE_CUDA == 1
    CuDevice::Instantiate().SetDebugStrideMode(true);
    if (loop == 0)
      CuDevice::Instantiate().SelectGpuId("no"); // -1 means no GPU
    else
      CuDevice::Instantiate().SelectGpuId("optional"); // -2 .. automatic selection
#endif
    for (int32 i = 0; i < 5; i++) {
      UnitTestPreconditionDirectionsOnline();
    }
  }
}