/**
 * Copyright (c) 2013-present, Facebook, Inc.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 * @providesModule PrefixIntervalTree
 * @flow
 * @typechecks
 */

'use strict';

const invariant = require('invariant');

const parent = node => Math.floor(node / 2);

const Int32Array = global.Int32Array ||
  function(size: number): Array<number> {
    const xs = [];
    for (let i = size - 1; i >= 0; --i) {
      xs[i] = 0;
    }
    return xs;
  };

/**
 * Computes the next power of 2 after or equal to x.
 */
function ceilLog2(x: number): number {
  let y = 1;
  while (y < x) {
    y *= 2;
  }
  return y;
}

/**
 * A prefix interval tree stores an numeric array and the partial sums of that
 * array. It is optimized for updating the values of the array without
 * recomputing all of the partial sums.
 *
 *   - O(ln n) update
 *   - O(1) lookup
 *   - O(ln n) compute a partial sum
 *   - O(n) space
 *
 * Note that the sequence of partial sums is one longer than the array, so that
 * the first partial sum is always 0, and the last partial sum is the sum of the
 * entire array.
 */
class PrefixIntervalTree {
  /**
   * Number of elements in the array
   */
  _size: number;

  /**
   * Half the size of the heap. It is also the number of non-leaf nodes, and the
   * index of the first element in the heap. Always a power of 2.
   */
  _half: number;

  /**
   * Binary heap
   */
  _heap: Int32Array;

  constructor(xs: Array<number>) {
    this._size = xs.length;
    this._half = ceilLog2(this._size);
    this._heap = new Int32Array(2 * this._half);

    let i;
    for (i = 0; i < this._size; ++i) {
      this._heap[this._half + i] = xs[i];
    }

    for (i = this._half - 1; i > 0; --i) {
      this._heap[i] =  this._heap[2 * i] + this._heap[2 * i + 1];
    }
  }

  static uniform(size: number, initialValue: number): PrefixIntervalTree {
    const xs = [];
    for (let i = size - 1; i >= 0; --i) {
      xs[i] = initialValue;
    }

    return new PrefixIntervalTree(xs);
  }

  static empty(size: number): PrefixIntervalTree {
    return PrefixIntervalTree.uniform(size, 0);
  }

  set(index: number, value: number): void {
    invariant(
      0 <= index && index < this._size,
      'Index out of range %s',
      index,
    );

    let node = this._half + index;
    this._heap[node] = value;

    node = parent(node);
    for (; node !== 0; node = parent(node)) {
      this._heap[node] =
        this._heap[2 * node] + this._heap[2 * node + 1];
    }
  }

  get(index: number): number {
    invariant(
      0 <= index && index < this._size,
      'Index out of range %s',
      index,
    );

    const node = this._half + index;
    return this._heap[node];
  }

  getSize(): number {
    return this._size;
  }

  /**
   * Returns the sum get(0) + get(1) + ... + get(end - 1).
   */
  sumUntil(end: number): number {
    invariant(
      0 <= end && end < this._size + 1,
      'Index out of range %s',
      end,
    );

    if (end === 0) {
      return 0;
    }

    let node = this._half + end - 1;
    let sum = this._heap[node];
    for (; node !== 1; node = parent(node)) {
      if (node % 2 === 1) {
        sum += this._heap[node - 1];
      }
    }

    return sum;
  }

  /**
   * Returns the sum get(0) + get(1) + ... + get(inclusiveEnd).
   */
  sumTo(inclusiveEnd: number): number {
    invariant(
      0 <= inclusiveEnd && inclusiveEnd < this._size,
      'Index out of range %s',
      inclusiveEnd,
    );
    return this.sumUntil(inclusiveEnd + 1);
  }

  /**
   * Returns the sum get(begin) + get(begin + 1) + ... + get(end - 1).
   */
  sum(begin: number, end: number): number {
    invariant(begin <= end, 'Begin must precede end');
    return this.sumUntil(end) - this.sumUntil(begin);
  }

  /**
   * Returns the smallest i such that 0 <= i <= size and sumUntil(i) <= t, or
   * -1 if no such i exists.
   */
  greatestLowerBound(t: number): number {
    if (t < 0) {
      return -1;
    }

    let node = 1;
    if (this._heap[node] <= t) {
      return this._size;
    }

    while (node < this._half) {
      const leftSum = this._heap[2 * node];
      if (t < leftSum) {
        node = 2 * node;
      } else {
        node = 2 * node + 1;
        t -= leftSum;
      }
    }

    return node - this._half;
  }

  /**
   * Returns the smallest i such that 0 <= i <= size and sumUntil(i) < t, or
   * -1 if no such i exists.
   */
  greatestStrictLowerBound(t: number): number {
    if (t <= 0) {
      return -1;
    }

    let node = 1;
    if (this._heap[node] < t) {
      return this._size;
    }

    while (node < this._half) {
      const leftSum = this._heap[2 * node];
      if (t <= leftSum) {
        node = 2 * node;
      } else {
        node = 2 * node + 1;
        t -= leftSum;
      }
    }

    return node - this._half;
  }

  /**
   * Returns the smallest i such that 0 <= i <= size and t <= sumUntil(i), or
   * size + 1 if no such i exists.
   */
  leastUpperBound(t: number): number {
    return this.greatestStrictLowerBound(t) + 1;
  }

  /**
   * Returns the smallest i such that 0 <= i <= size and t < sumUntil(i), or
   * size + 1 if no such i exists.
   */
  leastStrictUpperBound(t: number): number {
    return this.greatestLowerBound(t) + 1;
  }
}

module.exports = PrefixIntervalTree;
