/**
 * 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.m_), n_(mat.n_),
        codesize_(m_ * ((n_ + dsub - 1) / dsub)) {
            if (codesize_ > 0) {
                codes_ = new uint8_t[codesize_];
            }
            pq_ = std::unique_ptr<ProductQuantizer>( new ProductQuantizer(n_, dsub));
            if (qnorm_) {
                norm_codes_ = new uint8_t[m_];
                npq_ = std::unique_ptr<ProductQuantizer>( new ProductQuantizer(1, 1));
            }
            quantize(mat);
        }

    QMatrix::~QMatrix() {
        if (codesize_ > 0) {
            delete[] codes_;
        }
        if (qnorm_) { delete[] norm_codes_; }
    }

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

    void QMatrix::quantize(const Matrix& matrix) {
        assert(n_ == matrix.n_);
        assert(m_ == matrix.m_);
        Matrix temp(matrix);
        if (qnorm_) {
            Vector norms(temp.m_);
            temp.l2NormRow(norms);
            temp.divideRow(norms);
            quantizeNorm(norms);
        }
        auto dataptr = temp.data_;
        pq_->train(m_, dataptr);
        pq_->compute_codes(dataptr, codes_, 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_, 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_, 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_, codesize_ * sizeof(uint8_t));
        pq_->save(out);
        if (qnorm_) {
            out.write((char*) norm_codes_, 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_ = new uint8_t[codesize_];
        in.read((char*) codes_, codesize_ * sizeof(uint8_t));
        pq_ = std::unique_ptr<ProductQuantizer>( new ProductQuantizer());
        pq_->load(in);
        if (qnorm_) {
            norm_codes_ = new uint8_t[m_];
            in.read((char*) norm_codes_, m_ * sizeof(uint8_t));
            npq_ = std::unique_ptr<ProductQuantizer>( new ProductQuantizer());
            npq_->load(in);
        }
    }

}
