/**
 * Copyright (c) 2016-present, Facebook, Inc.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree. An additional grant
 * of patent rights can be found in the PATENTS file in the same directory.
 */

#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);
    }
}

}
