package zemberek.core.embeddings;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Random;
import zemberek.core.collections.IntVector;

class ProductQuantizer {

  //TODO: original is 8 bits but because byte is signed, for now make it 7
  int nbits_ = 7;
  int ksub_ = 1 << nbits_;
  int max_points_per_cluster_ = 128;
  int max_points_ = max_points_per_cluster_ * ksub_;
  int seed_ = 1234;
  int niter_ = 25;
  float eps_ = 1e-7f;

  int dim_;
  int nsubq_;
  int dsub_;
  int lastdsub_;

  FArray centroids_;

  Random rng;

  public ProductQuantizer() {
  }

  // used for mimicking pointer arithmetic.
  ProductQuantizer(int dim, int dsub) {
    this.dim_ = dim;
    this.nsubq_ = dim / dsub;
    dsub_ = dsub;
    centroids_ = new FArray(new float[dim * ksub_]);
    rng = new Random(seed_);
    lastdsub_ = dim_ % dsub;
    if (lastdsub_ == 0) {
      lastdsub_ = dsub_;
    } else {
      nsubq_++;
    }
  }

  float distL2(FArray x, FArray y, int d) {
    float dist = 0;
    for (int i = 0; i < d; i++) {
      float tmp = x.get(i) - y.get(i);
      dist += tmp * tmp;
    }
    return dist;
  }

  // TODO: Original code has two metohds for this, one const other not.
  FArray get_centroids(int m, byte i) {
    if (m == nsubq_ - 1) {
      return centroids_.ref(m * ksub_ * dsub_ + i * lastdsub_);
    }
    return centroids_.ref((m * ksub_ + i) * dsub_);
  }

  float assign_centroid(
      FArray x,
      FArray c0,
      BArray code,
      int d) {

    FArray c = c0.ref(0);
    float dis = distL2(x, c, d);
    code.set(0, (byte) 0);
    for (int j = 1; j < ksub_; j++) {
      c = c.ref(d);
      float disij = distL2(x, c, d);
      if (disij < dis) {
        code.set(0, (byte) j);
        dis = disij;
      }
    }
    return dis;
  }

  void Estep(
      FArray x,
      FArray centroids,
      BArray codes,
      int d,
      int n) {
    for (int i = 0; i < n; i++) {
      assign_centroid(
          x.ref(i * d), centroids, codes.ref(i), d);
    }
  }

  void MStep(FArray x0,
      FArray centroids,
      BArray codes,
      int d,
      int n) {
    int[] nelts = new int[ksub_];
    centroids.fill(0, d * ksub_, 0);

    FArray x = x0.ref(0);

    for (int i = 0; i < n; i++) {
      int k = codes.get(i);
      FArray c = centroids.ref(k * d);
      for (int j = 0; j < d; j++) {
        c.add(j, x.get(j));
      }
      nelts[k]++;
      x = x.ref(d);
    }

    FArray c = centroids.ref(0);
    for (int k = 0; k < ksub_; k++) {
      float z = (float) nelts[k];
      if (z != 0) {
        for (int j = 0; j < d; j++) {
          c.divide(j, z);
        }
      }
      c = c.ref(d);
    }

    for (int k = 0; k < ksub_; k++) {
      if (nelts[k] == 0) {
        int m = 0;
        while (rng.nextFloat() * (n - ksub_) >= nelts[m] - 1) {
          m = (m + 1) % ksub_;
        }

        centroids.arrayCopy(m * d, centroids, k * d, d);
        for (int j = 0; j < d; j++) {
          int sign = (j % 2) * 2 - 1;
          centroids.add(k * d + j, sign * eps_);
          centroids.add(m * d + j, -sign * eps_);
        }
        nelts[k] = nelts[m] / 2;
        nelts[m] -= nelts[k];
      }
    }
  }

  void kmeans(FArray x, FArray c, int n, int d) {
    int[] values = new int[n];
    IntVector perm = new IntVector(values);

    for (int i = 0; i < n; i++) {
      perm.safeSet(i, i);
    }
    perm.shuffle(rng);
    for (int i = 0; i < ksub_; i++) {
      x.arrayCopy( // x -> src
          perm.get(i) * d, // src pos
          c,     // destination
          i * d, // destination pos
          d);    //amount of data
    }
    BArray codes = new BArray(new byte[n]);
    for (int i = 0; i < niter_; i++) {
      Estep(x, c, codes, d, n);
      MStep(x, c, codes, d, n);
    }
  }

  void train(int n, float[] x) {
    if (n < ksub_) {
      throw new IllegalArgumentException(
          "Matrix too small for quantization, must have > 256 rows. But it is " + n);
    }
    IntVector perm = new IntVector(new int[n]);
    for (int i = 0; i < n; i++) {
      perm.safeSet(i, i);
    }
    int d = dsub_;
    int np = Math.min(n, max_points_);

    float[] xslice = new float[np * dsub_];
    for (int m = 0; m < nsubq_; m++) {
      if (m == nsubq_ - 1) {
        d = lastdsub_;
      }
      if (np != n) {
        perm.shuffle(rng);
      }
      for (int j = 0; j < np; j++) {
        System.arraycopy(x, perm.get(j) * dim_ + m * dsub_, xslice, j * d, d);
      }
      kmeans(new FArray(xslice),
          get_centroids(m, (byte) 0),
          np,
          d);
    }
  }

  float mulcode(Vector x, BArray codes, int t, float alpha) {
    float res = 0;
    int d = dsub_;
    BArray code = codes.ref(nsubq_ * t);
    for (int m = 0; m < nsubq_; m++) {
      FArray c = get_centroids(m, code.get(m));
      if (m == nsubq_ - 1) {
        d = lastdsub_;
      }
      for (int n = 0; n < d; n++) {
        res += x.data_[m * dsub_ + n] * c.get(n);
      }
    }
    return res * alpha;
  }

  void addcode(Vector x, BArray codes, int t, float alpha) {
    int d = dsub_;
    BArray code = codes.ref(nsubq_ * t);
    for (int m = 0; m < nsubq_; m++) {
      FArray c = get_centroids(m, code.get(m));
      if (m == nsubq_ - 1) {
        d = lastdsub_;
      }
      for (int n = 0; n < d; n++) {
        x.data_[m * dsub_ + n] += alpha * c.get(n);
      }
    }
  }

  void compute_code(FArray x, BArray code) {
    int d = dsub_;
    for (int m = 0; m < nsubq_; m++) {
      if (m == nsubq_ - 1) {
        d = lastdsub_;
      }
      assign_centroid(
          x.ref(m * dsub_),
          get_centroids(m, (byte) 0),
          code.ref(m),
          d);
    }
  }

  void compute_codes(FArray x, BArray codes, int n) {
    for (int i = 0; i < n; i++) {
      compute_code(
          x.ref(i * dim_),
          codes.ref(i * nsubq_));
    }
  }

  void save(DataOutputStream dos) throws IOException {
    dos.writeInt(dim_);
    dos.writeInt(nsubq_);
    dos.writeInt(dsub_);
    dos.writeInt(lastdsub_);
    for (int i = 0; i < centroids_.data.length; i++) {
      dos.writeFloat(centroids_.data[i]);
    }
  }

  void load(DataInputStream dis) throws IOException {
    dim_ = dis.readInt();
    nsubq_ = dis.readInt();
    dsub_ = dis.readInt();
    lastdsub_ = dis.readInt();
    float[] centroidData = new float[dim_ * ksub_];
    for (int i = 0; i < centroidData.length; i++) {
      centroidData[i] = dis.readFloat();
    }
    centroids_ = new FArray(centroidData);
  }

  static class FArray {

    int pointer;

    float[] data;

    FArray(float[] data) {
      this.data = data;
      this.pointer = 0;
    }

    FArray(int pointer, float[] data) {
      this.pointer = pointer;
      this.data = data;
    }

    void fill(int from, int to, float val) {
      Arrays.fill(data, pointer + from, pointer + to, val);
    }

    float get(int i) {
      return data[pointer + i];
    }

    void arrayCopy(
        int srcPos,
        FArray dest,
        int destPos,
        int amount) {
      System.arraycopy(
          data,
          pointer + srcPos,
          dest.data,
          dest.pointer + destPos,
          amount);
    }

    void set(int i, float value) {
      data[pointer + i] = value;
    }

    void add(int i, float value) {
      data[pointer + i] += value;
    }

    void divide(int i, float value) {
      data[pointer + i] /= value;
    }

    FArray ref(int offset) {
      return new FArray(pointer + offset, this.data);
    }

  }

  static class BArray {

    int pointer;

    byte[] data;

    BArray(byte[] data) {
      this.data = data;
      this.pointer = 0;
    }

    BArray(int pointer, byte[] data) {
      this.pointer = pointer;
      this.data = data;
    }

    byte get(int i) {
      return data[pointer + i];
    }

    void set(int i, byte value) {
      data[pointer + i] = value;
    }

    BArray ref(int offset) {
      return new BArray(pointer + offset, this.data);
    }
  }
}
