// transform/compressed-transform-stats.cc

// Copyright 2012  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/compressed-transform-stats.h"

namespace kaldi {

void CompressedAffineXformStats::CopyFromAffineXformStats(
    const AffineXformStats &input) {
  int32 dim = input.Dim();
  beta_ = input.beta_;
  if (beta_ == 0.0) { // empty; no stats.
    K_.Resize(dim, dim+1); // Will set to zero.
    // This stores the dimension.  Inefficient but this shouldn't happen often.
    Matrix<float> empty;
    G_.CopyFromMat(empty); // Sets G empty.
    return;
  }
  KALDI_ASSERT(input.G_.size() == dim && input.K_.NumCols() == dim+1
               && input.K_.NumRows() == dim && input.G_[0].NumRows() == dim+1);
  // OK, we have valid, nonempty stats.
  // We first slightly change the format of G.
  Matrix<double> Gtmp(dim, 1 + (((dim+1)*(dim+2))/2));
  // Gtmp will be compressed into G_.  The first element of each
  // row of Gtmp is the trace of the corresponding G[i], divided
  // by (beta * dim).  [this division is so we expect it to be
  // approximately 1, to keep things in a good range so they
  // can be more easily compressed.]  The next (((dim+1)*(dim+2))/2))
  // elements are the linearized form of the symmetric (d+1) by (d+1) matrix
  // input.G_[i], normalized appropriately using that trace.

  Matrix<double> K_corrected(input.K_); // This K_corrected matrix is a version of the
  // K_ matrix that we will correct to ensure that the derivative of the
  // objective function around the default matrix stays the same after
  // compression.

  SpMatrix<double> Gi_tmp(dim+1);
  for (int32 i = 0; i < dim; i++) {
    SubVector<double> this_row(Gtmp, i);
    PrepareOneG(input.G_[i], beta_, &this_row);
    ExtractOneG(this_row, beta_, &Gi_tmp);

    // At this stage we use the difference betwen Gi and Gi_tmp to
    // make a correction to K_.
    Vector<double> old_g_row(dim+1), new_g_row(dim+1);
    old_g_row.CopyRowFromSp(input.G_[i], i); // i'th row of old G_i.
    new_g_row.CopyRowFromSp(Gi_tmp, i); // i'th row of compressed+reconstructed G_i.
    // The auxiliary function for the i'th row of the transform, v_i, is as follows
    // [ignoring the determinant], where/ k_i is the i'th row of K:
    //  v_i . k_i - 0.5 v_i^T G_i u_i.
    // Let u_i be the unit vector in the i'th dimension.  This is the "default" value
    // of v_i.  The derivative of the auxf w.r.t. v_i, taken around this point, is:
    // k_i - G_i u_i
    // which is the same as k_i minus the i'th row (or column) of G_i
    // we want the derivative to be unchanged after compression:
    // new_ki - new_G_i u_i = old_ki - old_G_i u_i
    // new_ki = old_ki - old_G_i u_i + new_G_i u_i.
    // new_ki = old_ki - (i'th row of old G_i) + (i'th row of new G_i).
    
    SubVector<double> Ki(K_corrected, i);
    Ki.AddVec(-1.0, old_g_row);
    Ki.AddVec(+1.0, new_g_row);
  }
  K_.Resize(dim, dim+1);
  K_.CopyFromMat(K_corrected);
  G_.CopyFromMat(Gtmp);
}

void CompressedAffineXformStats::CopyToAffineXformStats(
    AffineXformStats *output) const {
  int32 dim = K_.NumRows();
  if (dim == 0) {
    output->Init(0, 0);
    return;
  }
  if (output->Dim() != dim || output->G_.size() != dim || beta_ == 0.0)
    output->Init(dim, dim);
  if (beta_ == 0.0) return; // Init() will have cleared it.
  output->beta_ = beta_;
  output->K_.CopyFromMat(K_);
  Matrix<double> Gtmp(G_.NumRows(), G_.NumCols());  // CopyToMat no longer
  // resizes, we have to provide correctly-sized matrix
  G_.CopyToMat(&Gtmp);
  for (int32 i = 0; i < dim; i++) {
    SubVector<double> this_row(Gtmp, i);
    ExtractOneG(this_row, beta_, &(output->G_[i]));
  }
}

void CompressedAffineXformStats::Write(std::ostream &os, bool binary) const {
  WriteToken(os, binary, "<CompressedAffineXformStats>");
  WriteBasicType(os, binary, beta_);
  K_.Write(os, binary);
  G_.Write(os, binary);
  WriteToken(os, binary, "</CompressedAffineXformStats>");
}

void CompressedAffineXformStats::Read(std::istream &is, bool binary) {
  ExpectToken(is, binary, "<CompressedAffineXformStats>");
  ReadBasicType(is, binary, &beta_);
  K_.Read(is, binary);
  G_.Read(is, binary);
  ExpectToken(is, binary, "</CompressedAffineXformStats>");
}

// Convert one G matrix into linearized, normalized form ready
// for compression.  A static function.
void CompressedAffineXformStats::PrepareOneG(const SpMatrix<double> &Gi,
                                             double beta,
                                             SubVector<double> *linearized) {
  KALDI_ASSERT(beta != 0.0);
  int32 dim = Gi.NumRows() - 1;
  double raw_trace = Gi.Trace();
  double norm_trace = (raw_trace / (beta * dim));
  (*linearized)(0) = norm_trace; // should be around 1.
  SubVector<double> linearized_matrix((*linearized), 1, ((dim+1)*(dim+2))/2);
  TpMatrix<double> C(dim+1);
  C.Cholesky(Gi); // Get the Cholesky factor: after we compress and uncompress
  // this and re-create Gi, it's bound to be +ve semidefinite, which is a Good Thing.
  C.Scale(sqrt(dim / raw_trace)); // This is the scaling that is equivalent
  // to scaling Gi by dim / raw_trace, which would make the diagonals
  // of Gi average to 1.  We can reverse this when we decompress.
  linearized_matrix.CopyFromPacked(C);  
}

// Reverse the process of PrepareOneG.  A static function.
void CompressedAffineXformStats::ExtractOneG(const SubVector<double> &linearized,
                                             double beta,
                                             SpMatrix<double> *Gi) {
  int32 dim = Gi->NumRows() - 1;
  KALDI_ASSERT(dim > 0);
  double norm_trace = linearized(0);
  double raw_trace = norm_trace * beta * dim;
  TpMatrix<double> C(dim+1);
  C.CopyFromVec(linearized.Range(1, ((dim+1)*(dim+2))/2));
  Gi->AddTp2(raw_trace / dim, C, kNoTrans, 0.0);
}



} // namespace kaldi