import { CompressedBitvector } from './CompressedBitvector.js';

class WaveletTree {
  constructor(sequence, alphabetSize, alphabetMin = 0) {
    this.sequence = sequence;
    this.length = sequence.length;
    this.alphabetMin = alphabetMin;
    this.alphabetMax = alphabetMin + alphabetSize - 1;
    this.alphabetSize = alphabetSize;
    this.isLeaf = alphabetSize <= 1;

    this.bitvector = null;
    this.leftChild = null;
    this.rightChild = null;

    if (this.isLeaf) {
      return;
    }

    const mid = (this.alphabetMin + this.alphabetMax) >> 1;
    const bitvectorBits = new Uint8Array(this.length);
    const leftSeq = [];
    const rightSeq = [];

    for (let i = 0; i < this.length; i++) {
      const c = sequence[i];
      if (c <= mid) {
        bitvectorBits[i] = 0;
        leftSeq.push(c);
      } else {
        bitvectorBits[i] = 1;
        rightSeq.push(c);
      }
    }

    this.bitvector = new CompressedBitvector(bitvectorBits);

    const leftSize = mid - this.alphabetMin + 1;
    const rightSize = this.alphabetMax - mid;

    if (leftSeq.length) {
      this.leftChild = new WaveletTree(new Uint32Array(leftSeq), leftSize, this.alphabetMin);
    }
    if (rightSeq.length) {
      this.rightChild = new WaveletTree(new Uint32Array(rightSeq), rightSize, mid + 1);
    }
  }

  access(i) {
    if (i < 0 || i >= this.length) {
      throw new Error(`Index ${i} out of bounds [0, ${this.length})`);
    }

    if (this.isLeaf) {
      return this.alphabetMin;
    }

    const bit = this.bitvector.access(i);
    if (bit === 0) {
      const pos = this.bitvector.rank0(i) - 1;
      if (pos < 0 || !this.leftChild) return this.alphabetMin;
      return this.leftChild.access(pos);
    }

    const pos = this.bitvector.rank1(i) - 1;
    if (pos < 0 || !this.rightChild) return this.alphabetMax;
    return this.rightChild.access(pos);
  }

  rank(c, i) {
    if (i < 0) return 0;
    if (i >= this.length) i = this.length - 1;

    if (this.isLeaf) {
      return c === this.alphabetMin ? i + 1 : 0;
    }

    const mid = (this.alphabetMin + this.alphabetMax) >> 1;
    if (c <= mid) {
      const pos = this.bitvector.rank0(i) - 1;
      if (pos < 0 || !this.leftChild) return 0;
      return this.leftChild.rank(c, pos);
    }

    const pos = this.bitvector.rank1(i) - 1;
    if (pos < 0 || !this.rightChild) return 0;
    return this.rightChild.rank(c, pos);
  }

  select(c, k) {
    if (k < 1) return -1;
    if (this.isLeaf) {
      return c === this.alphabetMin && k <= this.length ? k - 1 : -1;
    }

    const mid = (this.alphabetMin + this.alphabetMax) >> 1;
    if (c <= mid) {
      if (!this.leftChild) return -1;
      const pos = this.leftChild.select(c, k);
      if (pos === -1) return -1;
      return this.bitvector.select0(pos + 1);
    }

    if (!this.rightChild) return -1;
    const pos = this.rightChild.select(c, k);
    if (pos === -1) return -1;
    return this.bitvector.select1(pos + 1);
  }

  getMemoryUsage() {
    let total = 0;
    const stack = [this];

    while (stack.length) {
      const node = stack.pop();
      if (!node) continue;
      if (node.bitvector) {
        total += node.bitvector.getMemoryUsage();
      }
      if (node.leftChild) stack.push(node.leftChild);
      if (node.rightChild) stack.push(node.rightChild);
    }

    return total;
  }
}

export { WaveletTree };
