/**
 * Copyright (c) 2016-present, Facebook, Inc.
 * All rights reserved.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include "qmatrix.h"

#include <assert.h>
#include <iostream>

namespace fasttext {

QMatrix::QMatrix() : qnorm_(false), m_(0), n_(0), codesize_(0) {}

QMatrix::QMatrix(const Matrix& mat, int32_t dsub, bool qnorm)
    : qnorm_(qnorm),
      m_(mat.size(0)),
      n_(mat.size(1)),
      codesize_(m_ * ((n_ + dsub - 1) / dsub)) {
  codes_.resize(codesize_);
  pq_ = std::unique_ptr<ProductQuantizer>(new ProductQuantizer(n_, dsub));
  if (qnorm_) {
    norm_codes_.resize(m_);
    npq_ = std::unique_ptr<ProductQuantizer>(new ProductQuantizer(1, 1));
  }
  quantize(mat);
}

void QMatrix::quantizeNorm(const Vector& norms) {
  assert(qnorm_);
  assert(norms.size() == m_);
  auto dataptr = norms.data();
  npq_->train(m_, dataptr);
  npq_->compute_codes(dataptr, norm_codes_.data(), m_);
}

void QMatrix::quantize(const Matrix& matrix) {
  assert(m_ == matrix.size(0));
  assert(n_ == matrix.size(1));
  Matrix temp(matrix);
  if (qnorm_) {
    Vector norms(temp.size(0));
    temp.l2NormRow(norms);
    temp.divideRow(norms);
    quantizeNorm(norms);
  }
  auto dataptr = temp.data();
  pq_->train(m_, dataptr);
  pq_->compute_codes(dataptr, codes_.data(), m_);
}

void QMatrix::addToVector(Vector& x, int32_t t) const {
  real norm = 1;
  if (qnorm_) {
    norm = npq_->get_centroids(0, norm_codes_[t])[0];
  }
  pq_->addcode(x, codes_.data(), t, norm);
}

real QMatrix::dotRow(const Vector& vec, int64_t i) const {
  assert(i >= 0);
  assert(i < m_);
  assert(vec.size() == n_);
  real norm = 1;
  if (qnorm_) {
    norm = npq_->get_centroids(0, norm_codes_[i])[0];
  }
  return pq_->mulcode(vec, codes_.data(), i, norm);
}

int64_t QMatrix::getM() const {
  return m_;
}

int64_t QMatrix::getN() const {
  return n_;
}

void QMatrix::save(std::ostream& out) {
  out.write((char*)&qnorm_, sizeof(qnorm_));
  out.write((char*)&m_, sizeof(m_));
  out.write((char*)&n_, sizeof(n_));
  out.write((char*)&codesize_, sizeof(codesize_));
  out.write((char*)codes_.data(), codesize_ * sizeof(uint8_t));
  pq_->save(out);
  if (qnorm_) {
    out.write((char*)norm_codes_.data(), m_ * sizeof(uint8_t));
    npq_->save(out);
  }
}

void QMatrix::load(std::istream& in) {
  in.read((char*)&qnorm_, sizeof(qnorm_));
  in.read((char*)&m_, sizeof(m_));
  in.read((char*)&n_, sizeof(n_));
  in.read((char*)&codesize_, sizeof(codesize_));
  codes_ = std::vector<uint8_t>(codesize_);
  in.read((char*)codes_.data(), codesize_ * sizeof(uint8_t));
  pq_ = std::unique_ptr<ProductQuantizer>(new ProductQuantizer());
  pq_->load(in);
  if (qnorm_) {
    norm_codes_ = std::vector<uint8_t>(m_);
    in.read((char*)norm_codes_.data(), m_ * sizeof(uint8_t));
    npq_ = std::unique_ptr<ProductQuantizer>(new ProductQuantizer());
    npq_->load(in);
  }
}

} // namespace fasttext
