content: package edu.berkeley.cs.succinct.util.stream;
  
  import edu.berkeley.cs.succinct.util.CommonUtils;
  import edu.berkeley.cs.succinct.util.dictionary.Tables;
  import edu.berkeley.cs.succinct.util.stream.serops.BitMapOps;
  import org.apache.hadoop.fs.FSDataInputStream;
  
  import java.io.IOException;
  
  import static edu.berkeley.cs.succinct.util.DictionaryUtils.*;
  
  public class WaveletTreeStream {
  
    private FSDataInputStream stream;
    private long startPos;
  
    public WaveletTreeStream(FSDataInputStream stream, long startPos) throws IOException {
      this.stream = stream;
      this.startPos = startPos;
    }
  
  
    public long lookup(int contextPos, int cellPos, int startIdx, int endIdx) throws IOException {
      stream.seek(startPos);
      return waveletTreeLookup(contextPos, cellPos, startIdx, endIdx);
    }
  
    private long waveletTreeLookup(int contextPos, int cellPos, int startIdx, int endIdx)
      throws IOException {
      byte m = stream.readByte();
      int left = (int) stream.readLong();
      int right = (int) stream.readLong();
      int dictPos = (int) stream.getPos();
      long p, v;
  
      if (contextPos > m && contextPos <= endIdx) {
        if (right == 0) {
          return select1(dictPos, cellPos);
        }
        stream.seek(startPos + right);
        p = waveletTreeLookup(contextPos, cellPos, m + 1, endIdx);
  
        v = select1(dictPos, (int) p);
      } else {
        if (left == 0) {
          return select0(dictPos, cellPos);
        }
        stream.seek(startPos + left);
        p = waveletTreeLookup(contextPos, cellPos, startIdx, m);
  
        v = select0(dictPos, (int) p);
      }
  
      return v;
    }
  
    private long select0(long dictPos, int i) throws IOException {
      assert (i >= 0);
  
      RandomAccessLongStream dictBuf = new RandomAccessLongStream(stream, dictPos, Integer.MAX_VALUE);
  
      long size = dictBuf.get();
  
      long val = i + 1;
      int sp = 0;
      int ep = (int) (size / CommonUtils.two32);
      int m;
      long r;
      int pos = 0;
      int blockClass, blockOffset;
      long sel;
      int lastBlock;
      long rankL12, posL12;
  
      int l3Size = (int) ((size / CommonUtils.two32) + 1);
      int l12Size = (int) ((size / 2048) + 1);
  
      int basePos = (int) dictBuf.position();
  
      while (sp <= ep) {
        m = (sp + ep) / 2;
        r = (m * CommonUtils.two32 - dictBuf.get(basePos + m));
        if (val > r) {
          sp = m + 1;
        } else {
          ep = m - 1;
        }
      }
  
      ep = Math.max(ep, 0);
      val -= (ep * CommonUtils.two32 - dictBuf.get(basePos + ep));
      pos += dictBuf.get(basePos + l3Size + ep);
      sp = (int) (ep * CommonUtils.two32 / 2048);
      ep = (int) (Math.min(((ep + 1) * CommonUtils.two32 / 2048), Math.ceil((double) size / 2048.0))
        - 1);
  
      assert (val <= CommonUtils.two32);
      assert (pos >= 0);
  
      dictBuf.position(basePos + 2 * l3Size);
      basePos = (int) dictBuf.position();
  
      while (sp <= ep) {
        m = (sp + ep) / 2;
        r = m * 2048 - GETRANKL2(dictBuf.get(basePos + m));
        if (val > r) {
          sp = m + 1;
        } else {
          ep = m - 1;
        }
      }
  
      ep = Math.max(ep, 0);
      sel = (long) (ep) * 2048L;
      rankL12 = dictBuf.get(basePos + ep);
      posL12 = dictBuf.get(basePos + l12Size + ep);
      val -= (ep * 2048 - GETRANKL2(rankL12));
      pos += GETPOSL2(posL12);
  
      assert (val <= 2048);
      assert (pos >= 0);
  
      r = (512 - GETRANKL1(rankL12, 1));
      if (sel + 512 < size && val > r) {
        pos += GETPOSL1(posL12, 1);
        val -= r;
        sel += 512;
        r = (512 - GETRANKL1(rankL12, 2));
        if (sel + 512 < size && val > r) {
          pos += GETPOSL1(posL12, 2);
          val -= r;
          sel += 512;
          r = (512 - GETRANKL1(rankL12, 3));
          if (sel + 512 < size && val > r) {
            pos += GETPOSL1(posL12, 3);
            val -= r;
            sel += 512;
          }
        }
      }
  
      dictBuf.position(basePos + 2 * l12Size);
  
      assert (val <= 512);
      assert (pos >= 0);
  
      dictBuf.get(); // TODO: Could remove this field altogether
  
      while (true) {
        blockClass = (int) BitMapOps.getValPos(dictBuf, pos, 4);
        short offsetSize = (short) Tables.offsetBits[blockClass];
        pos += 4;
        blockOffset = (int) ((blockClass == 0) ? BitMapOps.getBit(dictBuf, pos) * 16 : 0);
        pos += offsetSize;
  
        if (val <= (16 - (blockClass + blockOffset))) {
          pos -= (4 + offsetSize);
          break;
        }
  
        val -= (16 - (blockClass + blockOffset));
        sel += 16;
      }
  
      blockClass = (int) BitMapOps.getValPos(dictBuf, pos, 4);
      pos += 4;
      blockOffset = (int) BitMapOps.getValPos(dictBuf, pos, Tables.offsetBits[blockClass]);
      lastBlock = Tables.decodeTable[blockClass][blockOffset];
  
      long count = 0;
      for (i = 0; i < 16; i++) {
        if (((lastBlock >> (15 - i)) & 1) == 0) {
          count++;
        }
        if (count == val) {
          return sel + i;
        }
      }
  
      return sel;
    }
  
    private long select1(long dictPos, int i) throws IOException {
      assert (i >= 0);
  
      RandomAccessLongStream dictBuf = new RandomAccessLongStream(stream, dictPos, Integer.MAX_VALUE);
  
      long size = dictBuf.get();
  
      long val = i + 1;
      int sp = 0;
      int ep = (int) (size / CommonUtils.two32);
      int m;
      long r;
      int pos = 0;
      int blockClass, blockOffset;
      long sel;
      int lastBlock;
      long rankL12, posL12;
  
      int l3Size = (int) ((size / CommonUtils.two32) + 1);
      int l12Size = (int) ((size / 2048) + 1);
      int basePos = (int) dictBuf.position();
  
      while (sp <= ep) {
        m = (sp + ep) / 2;
        r = dictBuf.get(basePos + m);
        if (val > r) {
          sp = m + 1;
        } else {
          ep = m - 1;
        }
      }
  
      ep = Math.max(ep, 0);
      val -= dictBuf.get(basePos + ep);
      pos += dictBuf.get(basePos + l3Size + ep);
      sp = (int) (ep * CommonUtils.two32 / 2048);
      ep = (int) (Math.min(((ep + 1) * CommonUtils.two32 / 2048), Math.ceil((double) size / 2048.0))
        - 1);
  
      assert (val <= CommonUtils.two32);
      assert (pos >= 0);
  
      dictBuf.position(basePos + 2 * l3Size);
      basePos = (int) dictBuf.position();
  
      while (sp <= ep) {
        m = (sp + ep) / 2;
        r = GETRANKL2(dictBuf.get(basePos + m));
        if (val > r) {
          sp = m + 1;
        } else {
          ep = m - 1;
        }
      }
  
      ep = Math.max(ep, 0);
      sel = (long) (ep) * 2048L;
      rankL12 = dictBuf.get(basePos + ep);
      posL12 = dictBuf.get(basePos + l12Size + ep);
      val -= GETRANKL2(rankL12);
      pos += GETPOSL2(posL12);
  
      assert (val <= 2048);
      assert (pos >= 0);
  
      r = GETRANKL1(rankL12, 1);
      if (sel + 512 < size && val > r) {
        pos += GETPOSL1(posL12, 1);
        val -= r;
        sel += 512;
        r = GETRANKL1(rankL12, 2);
        if (sel + 512 < size && val > r) {
          pos += GETPOSL1(posL12, 2);
          val -= r;
          sel += 512;
          r = GETRANKL1(rankL12, 3);
          if (sel + 512 < size && val > r) {
            pos += GETPOSL1(posL12, 3);
            val -= r;
            sel += 512;
          }
        }
      }
  
      dictBuf.position(basePos + 2 * l12Size);
  
      assert (val <= 512);
      assert (pos >= 0);
  
      dictBuf.get(); // TODO: Could remove this field altogether
  
      while (true) {
        blockClass = (int) BitMapOps.getValPos(dictBuf, pos, 4);
        short offsetSize = (short) Tables.offsetBits[blockClass];
        pos += 4;
        blockOffset = (int) ((blockClass == 0) ? BitMapOps.getBit(dictBuf, pos) * 16 : 0);
        pos += offsetSize;
  
        if (val <= (blockClass + blockOffset)) {
          pos -= (4 + offsetSize);
          break;
        }
  
        val -= (blockClass + blockOffset);
        sel += 16;
      }
  
      blockClass = (int) BitMapOps.getValPos(dictBuf, pos, 4);
      pos += 4;
      blockOffset = (int) BitMapOps.getValPos(dictBuf, pos, Tables.offsetBits[blockClass]);
      lastBlock = Tables.decodeTable[blockClass][blockOffset];
  
      long count = 0;
      for (i = 0; i < 16; i++) {
        if (((lastBlock >>> (15 - i)) & 1) == 1) {
          count++;
        }
        if (count == val) {
          return sel + i;
        }
      }
  
      return sel;
    }
  
    public void close() throws IOException {
      stream.close();
    }
  }
  