package com.carrotsearch.hppc;

import com.carrotsearch.hppc.cursors.LongCursor;
import com.carrotsearch.hppc.procedures.LongProcedure;
import java.util.Arrays;
import java.util.Iterator;

/**
 * Space-efficient index that enables fast rank/range search operations on a sorted sequence of
 * <code>long</code>.
 *
 * <p>Implementation of the PGM-Index described at <a
 * href="https://pgm.di.unipi.it/">https://pgm.di.unipi.it/</a>, based on the paper
 *
 * <pre>
 *   Paolo Ferragina and Giorgio Vinciguerra.
 *   The PGM-index: a fully-dynamic compressed learned index with provable worst-case bounds.
 *   PVLDB, 13(8): 1162-1175, 2020.
 * </pre>
 *
 * It provides {@code rank} and {@code range} search operations. {@code indexOf()} is faster than
 * B+Tree, and the index is much more compact. {@code contains()} is between 4x to 7x slower than
 * {@code IntHashSet#contains()}, but between 2.5x to 3x faster than {@link Arrays#binarySearch}.
 *
 * <p>Its compactness (40KB for 200MB of keys) makes it efficient for very large collections, the
 * index fitting easily in the L2 cache. The {@code epsilon} parameter should be set according to
 * the desired space-time trade-off. A smaller value makes the estimation more precise and the range
 * smaller but at the cost of increased space usage. In practice, {@code epsilon} 64 is a good sweet
 * spot.
 *
 * <p>Internally the index uses an optimal piecewise linear mapping from keys to their position in
 * the sorted order. This mapping is represented as a sequence of linear models (segments) which are
 * themselves recursively indexed by other piecewise linear mappings.
 */
@com.carrotsearch.hppc.Generated(date = "2024-06-04T15:20:17+0200", value = "KTypePgmIndex.java")
public class LongPgmIndex implements Accountable {

  /** Empty immutable LongPgmIndex. */
  public static final LongPgmIndex EMPTY = new LongEmptyPgmIndex();

  /**
   * Epsilon approximation range when searching the list of keys. Controls the size of the returned
   * search range, strictly greater than 0. It should be set according to the desired space-time
   * trade-off. A smaller value makes the estimation more precise and the range smaller but at the
   * cost of increased space usage.
   *
   * <p>With EPSILON=64 the benchmark with 200MB of keys shows that this PGM index requires only 2%
   * additional memory on average (40KB). It depends on the distribution of the keys. This epsilon
   * value is good even for 2MB of keys. With EPSILON=32: +5% speed, but 4x space (160KB).
   */
  public static final int EPSILON = 64;

  /**
   * Epsilon approximation range for the segments layers. Controls the size of the search range in
   * the hierarchical segment lists, strictly greater than 0.
   */
  public static final int EPSILON_RECURSIVE = 32;

  /** Size of a key, measured in {@link Integer#BYTES} because the key is stored in an int[]. */
  public static final int KEY_SIZE =
      RamUsageEstimator.primitiveSizes.get(long.class) / Integer.BYTES;

  /** 2x {@link #KEY_SIZE}. */
  public static final int DOUBLE_KEY_SIZE = KEY_SIZE * 2;

  /**
   * Data size of a segment, measured in {@link Integer#BYTES}, because segments are stored in an
   * int[].
   */
  public static final int SEGMENT_DATA_SIZE = KEY_SIZE * 3;

  /** Initial value of the exponential jump when scanning out of the epsilon range. */
  public static final int BEYOND_EPSILON_JUMP = 16;

  /**
   * The list of keys for which this index is built. It is sorted and may contain duplicate
   * elements.
   */
  public final LongArrayList keys;

  /** The size of the key set. That is, the number of distinct elements in {@link #keys}. */
  public final int size;

  /** The lowest key in {@link #keys}. */
  public final long firstKey;

  /** The highest key in {@link #keys}. */
  public final long lastKey;

  /** The epsilon range used to build this index. */
  public final int epsilon;

  /** The recursive epsilon range used to build this index. */
  public final int epsilonRecursive;

  /** The offsets in {@link #segmentData} of the first segment of each segment level. */
  public final int[] levelOffsets;

  /** The index data. It contains all the segments for all the levels. */
  public final int[] segmentData;

  private LongPgmIndex(
      LongArrayList keys,
      int size,
      int epsilon,
      int epsilonRecursive,
      int[] levelOffsets,
      int[] segmentData) {
    assert keys.size() > 0;
    assert size > 0 && size <= keys.size();
    assert epsilon > 0;
    assert epsilonRecursive > 0;
    this.keys = keys;
    this.size = size;
    firstKey = keys.get(0);
    lastKey = keys.get(keys.size() - 1);
    this.epsilon = epsilon;
    this.epsilonRecursive = epsilonRecursive;
    this.levelOffsets = levelOffsets;
    this.segmentData = segmentData;
  }

  /** Empty set constructor. */
  private LongPgmIndex() {
    keys = new LongArrayList(0);
    size = 0;
    firstKey = 0L;
    lastKey = 0L;
    epsilon = 0;
    epsilonRecursive = 0;
    levelOffsets = new int[0];
    segmentData = levelOffsets;
  }

  /** Returns the size of the key set. That is, the number of distinct elements in {@link #keys}. */
  public int size() {
    return size;
  }

  /** Returns whether this key set is empty. */
  public boolean isEmpty() {
    return size() == 0;
  }

  /** Returns whether this key set contains the given key. */
  public boolean contains(long key) {
    return indexOf(key) >= 0;
  }

  /**
   * Searches the specified key, and returns its index in the element list. If multiple elements are
   * equal to the specified key, there is no guarantee which one will be found.
   *
   * @return The index of the searched key if it is present; otherwise, {@code (-(<i>insertion
   *     point</i>) - 1)}. The <i>insertion point</i> is defined as the point at which the key would
   *     be inserted into the list: the index of the first element greater than the key, or {@link
   *     #keys}#{@code size()} if all the elements are less than the specified key. Note that this
   *     guarantees that the return value will be &gt;= 0 if and only if the key is found.
   */
  public int indexOf(long key) {
    if (key < firstKey) {
      return -1;
    }
    if (key > lastKey) {
      return -keys.size() - 1;
    }
    final int[] segmentData = this.segmentData;
    int segmentDataIndex = findSegment(key);
    int nextIntercept = (int) getIntercept(segmentDataIndex + SEGMENT_DATA_SIZE, segmentData);
    int index =
        Math.min(
            approximateIndex(key, segmentDataIndex, segmentData),
            Math.min(nextIntercept, keys.size() - 1));
    assert index >= 0 && index < keys.size();
    long k = keys.get(index);
    if (key < k) {
      // Scan sequentially before the approximated index, within epsilon range.
      final int fromIndex = Math.max(index - epsilon - 1, 0);
      while (--index >= fromIndex) {
        k = keys.get(index);
        if (key > k) {
          return -index - 2;
        }
        if (((key) == (k))) {
          return index;
        }
      }
      // Continue scanning out of the epsilon range.
      // This might happen in rare cases of precision error during the approximation
      // computation for longs (we don't have long double 128 bits in Java).
      // This might also happen in rare corner cases of large duplicate elements
      // sequence at the epsilon range boundary.
      index++;
      int jump = BEYOND_EPSILON_JUMP;
      do {
        int loIndex = Math.max(index - jump, 0);
        if (key >= keys.get(loIndex)) {
          return Arrays.binarySearch(keys.buffer, loIndex, index, key);
        }
        index = loIndex;
        jump <<= 1;
      } while (index > 0);
      return -1;
    } else if (((key) == (k))) {
      return index;
    } else {
      // Scan sequentially after the approximated index, within epsilon range.
      final int toIndex = Math.min(index + epsilon + 3, keys.size());
      while (++index < toIndex) {
        k = keys.get(index);
        if (key < k) {
          return -index - 1;
        }
        if (((key) == (k))) {
          return index;
        }
      }
      // Continue scanning out of the epsilon range.
      int jump = BEYOND_EPSILON_JUMP;
      do {
        int hiIndex = Math.min(index + jump, keys.size());
        if (key <= keys.get(hiIndex)) {
          return Arrays.binarySearch(keys.buffer, index, hiIndex, key);
        }
        index = hiIndex;
        jump <<= 1;
      } while (index < keys.size());
      return -keys.size() - 1;
    }
  }

  /**
   * Returns, for any value {@code x}, the number of keys in the sorted list which are smaller than
   * {@code x}. It is equal to {@link #indexOf} if {@code x} belongs to the list, or -{@link
   * #indexOf}-1 otherwise.
   *
   * <p>If multiple elements are equal to the specified key, there is no guarantee which one will be
   * found.
   *
   * @return The index of the searched key if it is present; otherwise, the {@code insertion point}.
   *     The <i>insertion point</i> is defined as the point at which the key would be inserted into
   *     the list: the index of the first element greater than the key, or {@link #keys}#{@code
   *     size()} if all the elements are less than the specified key. Note that this method always
   *     returns a value &gt;= 0.
   */
  public int rank(long x) {
    int index = indexOf(x);
    return index >= 0 ? index : -index - 1;
  }

  /**
   * Returns the number of keys in the list that are greater than or equal to {@code minKey}
   * (inclusive), and less than or equal to {@code maxKey} (inclusive).
   */
  public int rangeCardinality(long minKey, long maxKey) {
    int fromIndex = rank(minKey);
    int maxIndex = indexOf(maxKey);
    int toIndex = maxIndex >= 0 ? maxIndex + 1 : -maxIndex - 1;
    return Math.max(toIndex - fromIndex, 0);
  }

  /**
   * Returns an iterator over the keys in the list that are greater than or equal to {@code minKey}
   * (inclusive), and less than or equal to {@code maxKey} (inclusive).
   */
  public Iterator<LongCursor> rangeIterator(long minKey, long maxKey) {
    int fromIndex = rank(minKey);
    return new RangeIterator(keys, fromIndex, maxKey);
  }

  /**
   * Applies {@code procedure} to the keys in the list that are greater than or equal to {@code
   * minKey} (inclusive), and less than or equal to {@code maxKey} (inclusive).
   */
  public <T extends LongProcedure> T forEachInRange(T procedure, long minKey, long maxKey) {
    final long[] buffer = keys.buffer;
    long k;
    for (int i = rank(minKey), size = keys.size(); i < size && (k = buffer[i]) <= maxKey; i++) {
      procedure.apply(k);
    }
    return procedure;
  }

  /**
   * Estimates the allocated memory. It does not count the memory for the list of keys, only for the
   * index itself.
   */
  @Override
  public long ramBytesAllocated() {
    // int: size, epsilon, epsilonRecursive
    // long: firstKey, lastKey
    return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
        + 3 * Integer.BYTES
        + 2L * KEY_SIZE * Integer.BYTES
        // + keys.ramBytesAllocated()
        + RamUsageEstimator.shallowSizeOfArray(levelOffsets)
        + RamUsageEstimator.shallowSizeOfArray(segmentData);
  }

  /**
   * Estimates the bytes that are actually used. It does not count the memory for the list of keys,
   * only for the index itself.
   */
  @Override
  public long ramBytesUsed() {
    // int: size, epsilon, epsilonRecursive
    // long: firstKey, lastKey
    return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
        + 3 * Integer.BYTES
        + 2L * KEY_SIZE * Integer.BYTES
        // + keys.ramBytesUsed()
        + RamUsageEstimator.shallowSizeOfArray(levelOffsets)
        + RamUsageEstimator.shallowSizeOfArray(segmentData);
  }

  /**
   * Finds the segment responsible for a given key, that is, the rightmost segment having its first
   * key <= the searched key.
   *
   * @return the segment data index; or -1 if none.
   */
  private int findSegment(long key) {
    assert key >= firstKey && key <= lastKey;
    final int epsilonRecursive = this.epsilonRecursive;
    final int[] levelOffsets = this.levelOffsets;
    final int[] segmentData = this.segmentData;
    int level = levelOffsets.length - 1;
    int segmentDataIndex = levelOffsets[level] * SEGMENT_DATA_SIZE;
    while (--level >= 0) {
      int nextIntercept = (int) getIntercept(segmentDataIndex + SEGMENT_DATA_SIZE, segmentData);
      int index = Math.min(approximateIndex(key, segmentDataIndex, segmentData), nextIntercept);
      assert index >= 0 && index <= levelOffsets[level + 1] - levelOffsets[level] - 1;
      int sdIndex = (levelOffsets[level] + index) * SEGMENT_DATA_SIZE;
      if (getKey(sdIndex, segmentData) <= key) {
        // Scan sequentially segments after the approximated index, within the epsilon range.
        final int levelNumSegments = levelOffsets[level + 1] - levelOffsets[level] - 1;
        final int toIndex = Math.min(index + epsilonRecursive + 3, levelNumSegments);
        while (index++ < toIndex && getKey(sdIndex + SEGMENT_DATA_SIZE, segmentData) <= key) {
          sdIndex += SEGMENT_DATA_SIZE;
        }
      } else {
        // Scan sequentially segments before the approximated index, within the epsilon range.
        final int fromIndex = Math.max(index - epsilonRecursive - 1, 0);
        while (index-- > fromIndex) {
          sdIndex -= SEGMENT_DATA_SIZE;
          if (getKey(sdIndex, segmentData) <= key) {
            break;
          }
        }
      }
      segmentDataIndex = sdIndex;
    }
    assert segmentDataIndex >= 0;
    return segmentDataIndex;
  }

  private int approximateIndex(long key, int segmentDataIndex, int[] segmentData) {
    long intercept = getIntercept(segmentDataIndex, segmentData);
    long sKey = getKey(segmentDataIndex, segmentData);
    double slope = getSlope(segmentDataIndex, segmentData);
    int index = (int) (slope * ((double) key - sKey) + intercept);
    return Math.max(index, 0);
  }

  private static long getIntercept(int segmentDataIndex, int[] segmentData) {
    return PgmIndexUtil.getIntercept(segmentDataIndex, segmentData, KEY_SIZE);
  }

  private long getKey(int segmentDataIndex, int[] segmentData) {
    return PgmIndexUtil.getKey(segmentDataIndex + KEY_SIZE, segmentData, 0L);
  }

  private static double getSlope(int segmentDataIndex, int[] segmentData) {
    return PgmIndexUtil.getSlope(segmentDataIndex + DOUBLE_KEY_SIZE, segmentData, KEY_SIZE);
  }

  /** Empty immutable PGM Index. */
  private static class LongEmptyPgmIndex extends LongPgmIndex {

    private final Iterator<LongCursor> emptyIterator = new LongEmptyIterator();

    @Override
    public int indexOf(long key) {
      return -1;
    }

    @Override
    public Iterator<LongCursor> rangeIterator(long minKey, long maxKey) {
      return emptyIterator;
    }

    @Override
    public <T extends LongProcedure> T forEachInRange(T procedure, long minKey, long maxKey) {
      return procedure;
    }

    private static class LongEmptyIterator extends AbstractIterator<LongCursor> {
      @Override
      protected LongCursor fetch() {
        return done();
      }
    }
  }

  /** Iterator over a range of elements in a sorted array. */
  protected static class RangeIterator extends AbstractIterator<LongCursor> {
    private final long[] buffer;
    private final int size;
    private final LongCursor cursor;
    private final long maxKey;

    /** Range iterator from {@code fromIndex} (inclusive) to {@code maxKey} (inclusive). */
    protected RangeIterator(LongArrayList keys, int fromIndex, long maxKey) {
      this.buffer = keys.buffer;
      this.size = keys.size();
      this.cursor = new LongCursor();
      this.cursor.index = fromIndex;
      this.maxKey = maxKey;
    }

    @Override
    protected LongCursor fetch() {
      if (cursor.index >= size) {
        return done();
      }
      cursor.value = buffer[cursor.index++];
      if (cursor.value > maxKey) {
        cursor.index = size;
        return done();
      }
      return cursor;
    }
  }

  /** Builds a {@link LongPgmIndex} on a provided sorted list of keys. */
  public static class LongBuilder implements PlaModel.SegmentConsumer, Accountable {

    protected LongArrayList keys;
    protected int epsilon = EPSILON;
    protected int epsilonRecursive = EPSILON_RECURSIVE;
    protected PlaModel plam;
    protected int size;
    protected IntArrayList segmentData;
    protected int numSegments;

    /** Sets the sorted list of keys to build the index for; duplicate elements are allowed. */
    public LongBuilder setSortedKeys(LongArrayList keys) {
      this.keys = keys;
      return this;
    }

    /** Sets the sorted array of keys to build the index for; duplicate elements are allowed. */
    public LongBuilder setSortedKeys(long[] keys, int length) {
      LongArrayList keyList = new LongArrayList(0);
      keyList.buffer = keys;
      keyList.elementsCount = length;
      return setSortedKeys(keyList);
    }

    /** Sets the epsilon range to use when learning the segments for the list of keys. */
    public LongBuilder setEpsilon(int epsilon) {
      if (epsilon <= 0) {
        throw new IllegalArgumentException("epsilon must be > 0");
      }
      this.epsilon = epsilon;
      return this;
    }

    /**
     * Sets the recursive epsilon range to use when learning the segments for the segment levels.
     */
    public LongBuilder setEpsilonRecursive(int epsilonRecursive) {
      if (epsilonRecursive <= 0) {
        throw new IllegalArgumentException("epsilonRecursive must be > 0");
      }
      this.epsilonRecursive = epsilonRecursive;
      return this;
    }

    /** Builds the {@link LongPgmIndex}; or {@link #EMPTY} if there are no keys in the list. */
    public LongPgmIndex build() {
      if (keys == null || keys.size() == 0) {
        return (LongPgmIndex) EMPTY;
      }
      plam = new PlaModel(epsilon);

      int segmentsInitialCapacity =
          Math.min(
              Math.max(keys.size() / (2 * epsilon * epsilon) * SEGMENT_DATA_SIZE, 16), 1 << 19);
      segmentData = new IntArrayList(segmentsInitialCapacity);
      IntArrayList levelOffsets = new IntArrayList(16);

      int levelOffset = 0;
      levelOffsets.add(levelOffset);
      int levelNumSegments = buildFirstLevel();
      while (levelNumSegments > 1) {
        int nextLevelOffset = numSegments;
        levelOffsets.add(nextLevelOffset);
        levelNumSegments = buildUpperLevel(levelOffset, levelNumSegments);
        levelOffset = nextLevelOffset;
      }

      int[] segmentDataFinal = segmentData.toArray();
      int[] levelOffsetsFinal = levelOffsets.toArray();
      return new LongPgmIndex(
          keys, size, epsilon, epsilonRecursive, levelOffsetsFinal, segmentDataFinal);
    }

    private int buildFirstLevel() {
      assert numSegments == 0;
      int numKeys = keys.size();
      int size = 0;
      long key = keys.get(0);
      size++;
      plam.addKey(key, 0, this);
      for (int i = 1; i < numKeys; i++) {
        long nextKey = keys.get(i);
        if (!((nextKey) == (key))) {
          key = nextKey;
          plam.addKey(key, i, this);
          size++;
        }
      }
      plam.finish(this);
      addSentinelSegment(numKeys);
      this.size = size;
      return numSegments - 1;
    }

    private int buildUpperLevel(int levelOffset, int levelNumSegments) {
      plam.setEpsilon(epsilonRecursive);
      assert numSegments > 0;
      int initialNumSegments = numSegments;
      int segmentDataIndex = levelOffset * SEGMENT_DATA_SIZE;
      long key = getKey(segmentDataIndex, segmentData.buffer);
      plam.addKey(key, 0, this);
      for (int i = 1; i < levelNumSegments; i++) {
        segmentDataIndex += SEGMENT_DATA_SIZE;
        long nextKey = getKey(segmentDataIndex, segmentData.buffer);
        if (!((nextKey) == (key))) {
          key = nextKey;
          plam.addKey(key, i, this);
        }
      }
      plam.finish(this);
      addSentinelSegment(levelNumSegments);
      return numSegments - initialNumSegments - 1;
    }

    private long getKey(int segmentDataIndex, int[] segmentData) {
      return PgmIndexUtil.getKey(segmentDataIndex + KEY_SIZE, segmentData, 0L);
    }

    /**
     * Adds a sentinel segment that is used to give a limit for the position approximation, but does
     * not count in the number of segments per level.
     */
    private void addSentinelSegment(int endIndex) {
      // This sentinel segment is used in findSegment().
      accept(Double.MAX_VALUE, 0d, endIndex);
    }

    @Override
    public void accept(double firstKey, double slope, long intercept) {
      PgmIndexUtil.addIntercept(intercept, segmentData, KEY_SIZE);
      PgmIndexUtil.addKey((long) firstKey, segmentData);
      PgmIndexUtil.addSlope(slope, segmentData, KEY_SIZE);
      numSegments++;
      assert segmentData.size() == numSegments * SEGMENT_DATA_SIZE;
    }

    /**
     * Estimates the allocated memory. It does not count the memory for the list of keys, only for
     * the builder itself.
     */
    @Override
    public long ramBytesAllocated() {
      // int: epsilon, epsilonRecursive, size, numSegments
      return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
          + 4 * Integer.BYTES
          // + keys.ramBytesAllocated()
          + plam.ramBytesAllocated()
          + segmentData.ramBytesAllocated();
    }

    /**
     * Estimates the bytes that are actually used. It does not count the memory for the list of
     * keys, only for the builder itself.
     */
    @Override
    public long ramBytesUsed() {
      // int: epsilon, epsilonRecursive, size, numSegments
      return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
          + 4 * Integer.BYTES
          // + keys.ramBytesUsed()
          + plam.ramBytesUsed()
          + segmentData.ramBytesUsed();
    }
  }
}