code: package zemberek.core.quantization;
  
  import com.google.common.primitives.Doubles;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.Map;
  import zemberek.core.math.DoubleArrays;
  
  /**
   * A Quantizer that applies binning algorithm to input values to put them in a smaller range. There
   * are two binning implementation
   */
  public class BinningQuantizer implements Quantizer {
  
    Map<Double, Integer> lookup;
    double[] means;
  
    public BinningQuantizer(Map<Double, Integer> lookup, double[] means) {
      this.lookup = lookup;
      this.means = means;
    }
  
    public static BinningQuantizer linearBinning(float[] dataToQuantize, int bits) {
      return linearBinning(DoubleArrays.convert(dataToQuantize), bits);
    }
  
    public static BinningQuantizer linearBinning(double[] dataToQuantize, int bits) {
      checkRange(bits);
      final int dataLength = dataToQuantize.length;
      int range = 1 << bits;
      Map<Double, Integer> lookup = new HashMap<>(dataToQuantize.length);
      double means[] = new double[range];
      if (range >= dataLength) {
        means = new double[dataLength];
        int i = 0;
        for (double v : dataToQuantize) {
          lookup.put(v, i);
          means[i] = v;
          i++;
        }
        return new BinningQuantizer(lookup, means);
      }
      RawData[] data = new RawData[dataLength];
      for (int i = 0; i < data.length; i++) {
        data[i] = new RawData(dataToQuantize[i], i);
      }
      Arrays.sort(data);
      double binStep = (double) dataLength / range;
      int i = 0;
      int start = 0;
      double cursor = 0;
      int end = 0;
  
      while (cursor < dataLength) {
        start = (int) cursor;
        cursor += binStep;
        end = (int) cursor;
        if (end >= dataLength) {
          end = dataLength;
        }
        double total = 0;
        for (int k = start; k < end; k++) {
          total += data[k].value;
          lookup.put(data[k].value, i);
        }
        double mean = total / (end - start);
        means[i] = mean;
        i++;
      }
      return new BinningQuantizer(lookup, means);
    }
  
    private static void checkRange(int bits) {
      if (bits < 2 || bits > 24) {
        throw new IllegalArgumentException(
            "Bit count cannot be less than 4 or larger than 24" + bits);
      }
    }
  
    public static BinningQuantizer logCountBinning(float[] dataToQuantize, int counts[], int bits) {
      return logCountBinning(DoubleArrays.convert(dataToQuantize), counts, bits);
    }
  
    public static BinningQuantizer logCountBinning(double[] dataToQuantize, int counts[], int bits) {
      checkRange(bits);
      int range = 1 << bits;
      Map<Double, Integer> lookup = new HashMap<>(dataToQuantize.length);
      final int dataLength = dataToQuantize.length;
      double means[] = new double[range];
  
      if (range >= dataLength) {
        means = new double[dataLength];
        int i = 0;
        for (double v : dataToQuantize) {
          lookup.put(v, i);
          means[i] = v;
          i++;
        }
        return new BinningQuantizer(lookup, means);
      }
  
      RawDataWithCount[] data = new RawDataWithCount[dataLength];
  
      int totalLogCount = 0;
  
      for (int i = 0; i < data.length; i++) {
        data[i] = new RawDataWithCount(dataToQuantize[i], i, counts[i]);
        totalLogCount += data[i].logCount;
      }
  
      Arrays.sort(data);
  
      int i = 0;
      int k = 0;
  
      while (k < range) {
        double binStep = (double) totalLogCount / (range - k);
        int binLogCountTotal = 0;
  
        int start = i;
        while (binLogCountTotal < binStep) {
          binLogCountTotal += data[i].logCount;
          i++;
        }
        int end = i;
  
        double binWeightedValueTotal = 0;
        for (int j = start; j < end; j++) {
          binWeightedValueTotal += (data[j].value * data[j].logCount);
          lookup.put(data[j].value, k);
        }
        double weightedAverage = binWeightedValueTotal / binLogCountTotal;
        means[k] = weightedAverage;
  
        totalLogCount -= binLogCountTotal;
        k++;
      }
      return new BinningQuantizer(lookup, means);
    }
  
    double calculateError(double[] data, int[] counts) {
      double totalError = 0;
      int i = 0;
      for (double v : data) {
        totalError += Math.abs(v - means[lookup.get(v)]) * counts[i];
        i++;
      }
      return totalError;
    }
  
    double calculateError(double[] data) {
      double totalError = 0;
      int i = 0;
      for (double v : data) {
        totalError += Math.abs(v - means[lookup.get(v)]);
        i++;
      }
      return totalError;
    }
  
    public int getQuantizationIndex(double value) {
      return lookup.get(value);
    }
  
    public double getQuantizedValue(double value) {
      return means[lookup.get(value)];
    }
  
    public DoubleLookup getDequantizer() {
      return new DoubleLookup(means);
    }
  
    private static class RawData implements Comparable<RawData> {
  
      double value;
      int index;
  
      private RawData(double value, int index) {
        this.value = value;
        this.index = index;
      }
  
      public int compareTo(RawData o) {
        return Doubles.compare(value, o.value);
      }
    }
  
    private static class RawDataWithCount implements Comparable<RawDataWithCount> {
  
      double value;
      int index;
      int logCount;
  
      private RawDataWithCount(double value, int index, int count) {
        this.value = value;
        this.index = index;
        this.logCount = (int) (Math.log(count) + 1);
      }
  
      public int compareTo(RawDataWithCount o) {
        return Doubles.compare(value, o.value);
      }
    }
  }
  