// ivector/logistic-regression.cc
  
  // Copyright 2014  David Snyder
  
  // 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 "ivector/logistic-regression.h"
  #include "gmm/model-common.h" // For GetSplitTargets()
  #include <numeric> // For std::accumulate
  
  namespace kaldi {
  
  void LogisticRegression::Train(const Matrix<BaseFloat> &xs,
                                 const std::vector<int32> &ys,
                                 const LogisticRegressionConfig &conf) {
  
    int32 xs_num_rows = xs.NumRows(), xs_num_cols = xs.NumCols(),
                       num_ys = ys.size();
    KALDI_ASSERT(xs_num_rows == num_ys);
  
    // Adding on extra column for each x to handle the prior.
    Matrix<BaseFloat> xs_with_prior(xs_num_rows, xs_num_cols + 1);
    SubMatrix<BaseFloat> sub_xs(xs_with_prior, 0, xs_num_rows, 0, xs_num_cols);
    sub_xs.CopyFromMat(xs);
  
    int32 num_classes = *std::max_element(ys.begin(), ys.end()) + 1;
    weights_.Resize(num_classes, xs_num_cols + 1);
    Matrix<BaseFloat> xw(xs_num_rows, num_classes);
  
    // Adding on extra column for each x to handle the prior.
    for (int32 i = 0; i < xs_num_rows; i++) {
      xs_with_prior(i, xs_num_cols) = 1.0;
    }
  
    // At the beginning of training we have no mixture components,
    // therefore class_ is the "identity" mapping, that is
    // class_[i] = i.
    for (int32 i = 0; i < num_classes; i++) {
      class_.push_back(i);
    }
  
    weights_.SetZero();
    TrainParameters(xs_with_prior, ys, conf, &xw);
    KALDI_LOG << "Finished training parameters without mixture components.";
  
    // If we are using mixture components, we add those components
    // in MixUp and retrain with the extra weights.
    if (conf.mix_up > num_classes) {
      MixUp(ys, num_classes, conf);
      Matrix<BaseFloat> xw(xs_num_rows, weights_.NumRows());
      TrainParameters(xs_with_prior, ys, conf, &xw);
      KALDI_LOG << "Finished training mixture components.";
    }
  }
  
  
  void LogisticRegression::MixUp(const std::vector<int32> &ys,
                                 const int32 &num_classes,
                                 const LogisticRegressionConfig &conf) {
  
    Vector<BaseFloat> counts(num_classes);
    for (int32 i = 0; i < ys.size(); i++) {
      counts(ys[i]) += 1.0;
    }
  
    // TODO: Figure out what min_count should be
    int32 min_count = 1;
    std::vector<int32> targets;
    GetSplitTargets(counts, conf.mix_up, conf.power, min_count, &targets);
    int32 new_dim = std::accumulate(targets.begin(), targets.end(),
                                    static_cast<int32>(0));
  
    KALDI_LOG << "Target number mixture components was " << conf.mix_up
              << ". Training " << new_dim << " mixture components.";
  
    int32 old_dim = weights_.NumRows(),
          num_components = old_dim,
          num_feats = weights_.NumCols();
  
    Matrix<BaseFloat> old_weights(weights_);
    weights_.Resize(new_dim, num_feats);
    SubMatrix<BaseFloat> sub_weights(weights_, 0, num_classes, 0, num_feats);
    // We need to retain the original weights
    sub_weights.CopyFromMat(old_weights);
    class_.resize(new_dim);
    // For each class i
    for (int32 i = 0; i < targets.size(); i++) {
      int32 mixes = targets[i];
      // We start at j = 1 since one copy of the components already
      // exists in weights_.
      for (int32 j = 1; j < mixes; j++) {
        int32 offset = num_components;
        weights_.Row(offset).CopyRowFromMat(weights_, i);
        Vector<BaseFloat> noise(num_feats);
        noise.SetRandn();
        weights_.Row(offset).AddVec(1.0e-05, noise);
        class_[offset] = i; // The class i maps to the row at offset
        num_components += 1;
      }
    }
  }
  
  void LogisticRegression::TrainParameters(const Matrix<BaseFloat> &xs,
      const std::vector<int32> &ys, const LogisticRegressionConfig &conf,
      Matrix<BaseFloat> *xw) {
    int32 max_steps = conf.max_steps;
    BaseFloat normalizer = conf.normalizer;
    LbfgsOptions lbfgs_opts;
    lbfgs_opts.minimize = false;
    // Get initial w vector
    Vector<BaseFloat> init_w(weights_.NumRows() * weights_.NumCols());
    init_w.CopyRowsFromMat(weights_);
    OptimizeLbfgs<BaseFloat> lbfgs(init_w, lbfgs_opts);
  
    for (int32 step = 0; step < max_steps; step++) {
      DoStep(xs, xw, ys, &lbfgs, normalizer);
    }
  
    Vector<BaseFloat> best_w(lbfgs.GetValue());
    weights_.CopyRowsFromVec(best_w);
  }
  
  void LogisticRegression::GetLogPosteriors(const Matrix<BaseFloat> &xs,
                                            Matrix<BaseFloat> *log_posteriors) {
    int32 xs_num_rows = xs.NumRows(),
        xs_num_cols = xs.NumCols(),
        num_mixes = weights_.NumRows();
  
    int32 num_classes = *std::max_element(class_.begin(), class_.end()) + 1;
  
    log_posteriors->Resize(xs_num_rows, num_classes);
    Matrix<BaseFloat> xw(xs_num_rows, num_mixes);
  
    Matrix<BaseFloat> xs_with_prior(xs_num_rows, xs_num_cols + 1);
    SubMatrix<BaseFloat> sub_xs(xs_with_prior, 0, xs_num_rows, 0, xs_num_cols);
    sub_xs.CopyFromMat(xs);
    // Adding on extra column for each x to handle the prior.
    for (int32 i = 0; i < xs_num_rows; i++) {
      xs_with_prior(i, xs_num_cols) = 1.0;
    }
    xw.AddMatMat(1.0, xs_with_prior, kNoTrans, weights_,
                 kTrans, 0.0);
  
    log_posteriors->Set(-std::numeric_limits<BaseFloat>::infinity());
  
    // i is the training example
    for (int32 i = 0; i < xs_num_rows; i++) {
      for (int32 j = 0; j < num_mixes; j++) {
        int32 k = class_[j];
        (*log_posteriors)(i,k) = LogAdd((*log_posteriors)(i,k), xw(i, j));
      }
      // Normalize the row.
      log_posteriors->Row(i).Add(-xw.Row(i).LogSumExp());
    }
  }
  
  void LogisticRegression::GetLogPosteriors(const Vector<BaseFloat> &x,
                                            Vector<BaseFloat> *log_posteriors) {
    int32 x_dim = x.Dim();
    int32 num_classes = *std::max_element(class_.begin(), class_.end()) + 1,
        num_mixes = weights_.NumRows();
    log_posteriors->Resize(num_classes);
    Vector<BaseFloat> xw(weights_.NumRows());
  
    Vector<BaseFloat> x_with_prior(x_dim + 1);
    SubVector<BaseFloat> sub_x(x_with_prior, 0, x_dim);
    sub_x.CopyFromVec(x);
    // Adding on extra element to handle the prior
    x_with_prior(x_dim) = 1.0;
  
    xw.AddMatVec(1.0, weights_, kNoTrans, x_with_prior, kNoTrans);
  
    log_posteriors->Set(-std::numeric_limits<BaseFloat>::infinity());
  
    for (int32 i = 0; i < num_mixes; i++) {
      int32 j = class_[i];
      (*log_posteriors)(j) = LogAdd((*log_posteriors)(j), xw(i));
    }
    log_posteriors->Add(-log_posteriors->LogSumExp());
  }
  
  BaseFloat LogisticRegression::DoStep(const Matrix<BaseFloat> &xs,
      Matrix<BaseFloat> *xw,
      const std::vector<int32> &ys, OptimizeLbfgs<BaseFloat> *lbfgs,
      BaseFloat normalizer) {
    Matrix<BaseFloat> gradient(weights_.NumRows(), weights_.NumCols());
    // Vector form of the above matrix
    Vector<BaseFloat> grad_vec(weights_.NumRows() * weights_.NumCols());
  
    // Calculate XW.T. The rows correspond to the x
    // training examples and the columns to the class labels.
    xw->AddMatMat(1.0, xs, kNoTrans, weights_, kTrans, 0.0);
  
    // Calculate both the gradient and the objective function.
    BaseFloat objf = GetObjfAndGrad(xs, ys, *xw, &gradient, normalizer);
  
    // Convert gradient (a matrix) into a vector of size
    // gradient.NumCols * gradient.NumRows.
    grad_vec.CopyRowsFromMat(gradient);
  
    // Compute next step in L-BFGS.
    lbfgs->DoStep(objf, grad_vec);
  
    // Update weights
    Vector<BaseFloat> new_w(lbfgs->GetProposedValue());
    weights_.CopyRowsFromVec(new_w);
    KALDI_LOG << "Objective function is " << objf;
    return objf;
  }
  
  BaseFloat LogisticRegression::GetObjfAndGrad(
      const Matrix<BaseFloat> &xs,
      const std::vector<int32> &ys, const Matrix<BaseFloat> &xw,
      Matrix<BaseFloat> *grad, BaseFloat normalizer) {
    BaseFloat raw_objf = 0.0;
    int32 num_classes = *std::max_element(ys.begin(), ys.end()) + 1;
    std::vector< std::vector<int32> > class_to_cols(num_classes, std::vector<int32>());
    for (int32 i = 0; i < class_.size(); i++) {
      class_to_cols[class_[i]].push_back(i);
    }
    // For each training example class
    for (int32 i = 0; i < ys.size(); i++) {
      Vector<BaseFloat> row(xw.NumCols());
      row.CopyFromVec(xw.Row(i));
      row.ApplySoftMax();
      // Identify the rows of weights_ (which are a set of columns in wx)
      // which correspond to class ys[i]
      const std::vector<int32> &cols = class_to_cols[ys[i]];
      SubVector<BaseFloat> x = xs.Row(i);
      BaseFloat class_sum = 0.0;
      for (int32 j = 0; j < cols.size(); j++) {
        class_sum += row(cols[j]);
      }
      if (class_sum < 1.0e-20) class_sum = 1.0e-20;
      raw_objf += Log(class_sum);
      // Iterate over weights for each component. If there are no
      // mixtures each row corresponds to a class.
      for (int32 k = 0; k < weights_.NumRows(); k++) {
        // p(y = k | x_i) where k is a component.
        BaseFloat p = row(k);
        if (class_[k] == ys[i]) {
          // If the classes aren't split into mixture components
          // then p/class_sum = 1.0.
          grad->Row(k).AddVec(p/class_sum - p, x);
        } else {
          grad->Row(k).AddVec(-1.0 * p, x);
        }
      }
    }
    // Scale and add regularization term.
    grad->Scale(1.0/ys.size());
    grad->AddMat(-1.0 * normalizer, weights_);
    raw_objf /= ys.size();
    BaseFloat regularizer = - 0.5 * normalizer
                            * TraceMatMat(weights_, weights_, kTrans);
    KALDI_VLOG(2) << "Objf is " << raw_objf << " + " << regularizer
                  << " = " << (raw_objf + regularizer);
    return raw_objf + regularizer;
  }
  
  void LogisticRegression::SetWeights(const Matrix<BaseFloat> &weights,
                                      const std::vector<int32> classes) {
    weights_.Resize(weights.NumRows(), weights.NumCols());
    weights_.CopyFromMat(weights);
    class_.resize(classes.size());
    for (int32 i = 0; i < class_.size(); i++)
      class_[i] = classes[i];
  }
  
  void LogisticRegression::ScalePriors(const Vector<BaseFloat> &scales) {
    Vector<BaseFloat> log_scales(scales);
    log_scales.ApplyLog();
  
    for (int32 i = 0; i < weights_.NumRows(); i++)
      weights_(i, weights_.NumCols() - 1) += log_scales(class_[i]);
  }
  
  void LogisticRegression::Write(std::ostream &os, bool binary) const {
    WriteToken(os, binary, "<LogisticRegression>");
    WriteToken(os, binary, "<weights>");
    weights_.Write(os, binary);
    WriteToken(os, binary, "<class>");
    WriteIntegerVector(os, binary, class_);
    WriteToken(os, binary, "</LogisticRegression>");
  }
  
  void LogisticRegression::Read(std::istream &is, bool binary) {
    ExpectToken(is, binary, "<LogisticRegression>");
    ExpectToken(is, binary, "<weights>");
    weights_.Read(is, binary);
    std::string token;
    ReadToken(is, binary, &token);
    if (token == "<class>") {
      ReadIntegerVector(is, binary, &class_);
    } else {
      int32 num_classes = weights_.NumRows();
      for (int32 i = 0; i < num_classes; i++) {
        class_.push_back(i);
      }
    }
    ExpectToken(is, binary, "</LogisticRegression>");
  }
  
  }