/**
 * 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 "productquantizer.h"

#include <algorithm>
#include <iostream>
#include <numeric>
#include <stdexcept>
#include <string>

namespace fasttext
{

  real distL2(const real *x, const real *y, int32_t d)
  {
    real dist = 0;
    for (auto i = 0; i < d; i++)
    {
      auto tmp = x[i] - y[i];
      dist += tmp * tmp;
    }
    return dist;
  }

  ProductQuantizer::ProductQuantizer(int32_t dim, int32_t dsub)
      : dim_(dim),
        nsubq_(dim / dsub),
        dsub_(dsub),
        centroids_(dim * ksub_),
        rng(seed_)
  {
    lastdsub_ = dim_ % dsub;
    if (lastdsub_ == 0)
    {
      lastdsub_ = dsub_;
    }
    else
    {
      nsubq_++;
    }
  }

  const real *ProductQuantizer::get_centroids(int32_t m, uint8_t i) const
  {
    if (m == nsubq_ - 1)
    {
      return &centroids_[m * ksub_ * dsub_ + i * lastdsub_];
    }
    return &centroids_[(m * ksub_ + i) * dsub_];
  }

  real *ProductQuantizer::get_centroids(int32_t m, uint8_t i)
  {
    if (m == nsubq_ - 1)
    {
      return &centroids_[m * ksub_ * dsub_ + i * lastdsub_];
    }
    return &centroids_[(m * ksub_ + i) * dsub_];
  }

  real ProductQuantizer::assign_centroid(
      const real *x,
      const real *c0,
      uint8_t *code,
      int32_t d) const
  {
    const real *c = c0;
    real dis = distL2(x, c, d);
    code[0] = 0;
    for (auto j = 1; j < ksub_; j++)
    {
      c += d;
      real disij = distL2(x, c, d);
      if (disij < dis)
      {
        code[0] = (uint8_t)j;
        dis = disij;
      }
    }
    return dis;
  }

  void ProductQuantizer::Estep(
      const real *x,
      const real *centroids,
      uint8_t *codes,
      int32_t d,
      int32_t n) const
  {
    for (auto i = 0; i < n; i++)
    {
      assign_centroid(x + i * d, centroids, codes + i, d);
    }
  }

  void ProductQuantizer::MStep(
      const real *x0,
      real *centroids,
      const uint8_t *codes,
      int32_t d,
      int32_t n)
  {
    std::vector<int32_t> nelts(ksub_, 0);
    memset(centroids, 0, sizeof(real) * d * ksub_);
    const real *x = x0;
    for (auto i = 0; i < n; i++)
    {
      auto k = codes[i];
      real *c = centroids + k * d;
      for (auto j = 0; j < d; j++)
      {
        c[j] += x[j];
      }
      nelts[k]++;
      x += d;
    }

    real *c = centroids;
    for (auto k = 0; k < ksub_; k++)
    {
      real z = (real)nelts[k];
      if (z != 0)
      {
        for (auto j = 0; j < d; j++)
        {
          c[j] /= z;
        }
      }
      c += d;
    }

    std::uniform_real_distribution<> runiform(0, 1);
    for (auto k = 0; k < ksub_; k++)
    {
      if (nelts[k] == 0)
      {
        int32_t m = 0;
        while (runiform(rng) * (n - ksub_) >= nelts[m] - 1)
        {
          m = (m + 1) % ksub_;
        }
        memcpy(centroids + k * d, centroids + m * d, sizeof(real) * d);
        for (auto j = 0; j < d; j++)
        {
          int32_t sign = (j % 2) * 2 - 1;
          centroids[k * d + j] += sign * eps_;
          centroids[m * d + j] -= sign * eps_;
        }
        nelts[k] = nelts[m] / 2;
        nelts[m] -= nelts[k];
      }
    }
  }

  void ProductQuantizer::kmeans(const real *x, real *c, int32_t n, int32_t d)
  {
    std::vector<int32_t> perm(n, 0);
    std::iota(perm.begin(), perm.end(), 0);
    std::shuffle(perm.begin(), perm.end(), rng);
    for (auto i = 0; i < ksub_; i++)
    {
      memcpy(&c[i * d], x + perm[i] * d, d * sizeof(real));
    }
    auto codes = std::vector<uint8_t>(n);
    for (auto i = 0; i < niter_; i++)
    {
      Estep(x, c, codes.data(), d, n);
      MStep(x, c, codes.data(), d, n);
    }
  }

  void ProductQuantizer::train(int32_t n, const real *x)
  {
    if (n < ksub_)
    {
      throw std::invalid_argument(
          "Matrix too small for quantization, must have at least " +
          std::to_string(ksub_) + " rows");
    }
    std::vector<int32_t> perm(n, 0);
    std::iota(perm.begin(), perm.end(), 0);
    auto d = dsub_;
    auto np = std::min(n, max_points_);
    auto xslice = std::vector<real>(np * dsub_);
    for (auto m = 0; m < nsubq_; m++)
    {
      if (m == nsubq_ - 1)
      {
        d = lastdsub_;
      }
      if (np != n)
      {
        std::shuffle(perm.begin(), perm.end(), rng);
      }
      for (auto j = 0; j < np; j++)
      {
        memcpy(
            xslice.data() + j * d,
            x + perm[j] * dim_ + m * dsub_,
            d * sizeof(real));
      }
      kmeans(xslice.data(), get_centroids(m, 0), np, d);
    }
  }

  real ProductQuantizer::mulcode(
      const Vector &x,
      const uint8_t *codes,
      int32_t t,
      real alpha) const
  {
    real res = 0.0;
    auto d = dsub_;
    const uint8_t *code = codes + nsubq_ * t;
    for (auto m = 0; m < nsubq_; m++)
    {
      const real *c = get_centroids(m, code[m]);
      if (m == nsubq_ - 1)
      {
        d = lastdsub_;
      }
      for (auto n = 0; n < d; n++)
      {
        res += x[m * dsub_ + n] * c[n];
      }
    }
    return res * alpha;
  }

  void ProductQuantizer::addcode(
      Vector &x,
      const uint8_t *codes,
      int32_t t,
      real alpha) const
  {
    auto d = dsub_;
    const uint8_t *code = codes + nsubq_ * t;
    for (auto m = 0; m < nsubq_; m++)
    {
      const real *c = get_centroids(m, code[m]);
      if (m == nsubq_ - 1)
      {
        d = lastdsub_;
      }
      for (auto n = 0; n < d; n++)
      {
        x[m * dsub_ + n] += alpha * c[n];
      }
    }
  }

  void ProductQuantizer::compute_code(const real *x, uint8_t *code) const
  {
    auto d = dsub_;
    for (auto m = 0; m < nsubq_; m++)
    {
      if (m == nsubq_ - 1)
      {
        d = lastdsub_;
      }
      assign_centroid(x + m * dsub_, get_centroids(m, 0), code + m, d);
    }
  }

  void ProductQuantizer::compute_codes(const real *x, uint8_t *codes, int32_t n)
      const
  {
    for (auto i = 0; i < n; i++)
    {
      compute_code(x + i * dim_, codes + i * nsubq_);
    }
  }

  void ProductQuantizer::save(std::ostream &out) const
  {
    out.write((char *)&dim_, sizeof(dim_));
    out.write((char *)&nsubq_, sizeof(nsubq_));
    out.write((char *)&dsub_, sizeof(dsub_));
    out.write((char *)&lastdsub_, sizeof(lastdsub_));
    out.write((char *)centroids_.data(), centroids_.size() * sizeof(real));
  }

  void ProductQuantizer::load(std::istream &in)
  {
    in.read((char *)&dim_, sizeof(dim_));
    in.read((char *)&nsubq_, sizeof(nsubq_));
    in.read((char *)&dsub_, sizeof(dsub_));
    in.read((char *)&lastdsub_, sizeof(lastdsub_));
    centroids_.resize(dim_ * ksub_);
    for (size_t i = 0; i < centroids_.size(); i++)
    {
      in.read((char *)&centroids_[i], sizeof(real));
    }
  }

} // namespace fasttext
