import Node from "./node";
import { Key, Value } from "./types";

export type Comparator<Key> = (a: Key, b: Key) => number;
export type Visitor<Key, Value> = (
  node: Node<Key, Value>
) => undefined | boolean | void;
export type NodePrinter<Key, Value> = (node: Node<Key, Value>) => string;

/* follows "An implementation of top-down splaying"
 * by D. Sleator <sleator@cs.cmu.edu> March 1992
 */

function DEFAULT_COMPARE(a: Key, b: Key): number {
  return a > b ? 1 : a < b ? -1 : 0;
}

type TreeNodeList<Key, Value> = { head: Node<Key, Value> | null };

/**
 * Simple top down splay, not requiring i to be in the tree t.
 */
function splay(
  i: Key,
  t: Node<Key, Value>,
  comparator: Comparator<Key>
): Node<Key, Value> {
  const N = new Node(null, null);
  let l = N;
  let r = N;

  while (true) {
    const cmp = comparator(i, t!.key);
    //if (i < t.key) {
    if (cmp < 0) {
      if (t!.left === null) break;
      //if (i < t.left.key) {
      if (comparator(i, t.left.key) < 0) {
        const y = t.left; /* rotate right */
        t.left = y.right;
        y.right = t;
        t = y;
        if (t.left === null) break;
      }
      r.left = t; /* link right */
      r = t;
      t = t.left;
      //} else if (i > t.key) {
    } else if (cmp > 0) {
      if (t.right === null) break;
      //if (i > t.right.key) {
      if (comparator(i, t.right.key) > 0) {
        const y = t.right; /* rotate left */
        t.right = y.left;
        y.left = t;
        t = y;
        if (t.right === null) break;
      }
      l.right = t; /* link left */
      l = t;
      t = t.right;
    } else break;
  }
  /* assemble */
  l.right = t.left;
  r.left = t.right;
  t.left = N.right;
  t.right = N.left;
  return t;
}

function insert(
  i: Key,
  data: Value,
  t: Node<Key, Value>,
  comparator: Comparator<Key>
): Node<Key, Value> {
  const node = new Node(i, data);

  if (t === null) {
    node.left = node.right = null;
    return node;
  }

  t = splay(i, t, comparator);
  const cmp = comparator(i, t.key);
  if (cmp < 0) {
    node.left = t.left;
    node.right = t;
    t.left = null;
  } else if (cmp >= 0) {
    node.right = t.right;
    node.left = t;
    t.right = null;
  }
  return node;
}

function split(
  key: Key,
  v: Node<Key, Value>,
  comparator: Comparator<Key>
): {
  left: Node<Key, Value> | null;
  right: Node<Key, Value> | null;
} {
  let left = null;
  let right = null;
  if (v) {
    v = splay(key, v, comparator);

    const cmp = comparator(v.key, key);
    if (cmp === 0) {
      left = v.left;
      right = v.right;
    } else if (cmp < 0) {
      right = v.right;
      v.right = null;
      left = v;
    } else {
      left = v.left;
      v.left = null;
      right = v;
    }
  }
  return { left, right };
}

function merge(
  left: Node<Key, Value> | null,
  right: Node<Key, Value> | null,
  comparator: Comparator<Key>
) {
  if (right === null) return left;
  if (left === null) return right;

  right = splay(left.key, right, comparator);
  right.left = left;
  return right;
}

type StringCollector = (s: string) => void;

/**
 * Prints level of the tree
 */
function printRow(
  root: Node<Key, Value>,
  prefix: string,
  isTail: boolean,
  out: StringCollector,
  printNode: NodePrinter<Key, Value>
) {
  if (root) {
    out(`${prefix}${isTail ? "└── " : "├── "}${printNode(root)}\n`);
    const indent = prefix + (isTail ? "    " : "│   ");
    if (root.left) printRow(root.left, indent, false, out, printNode);
    if (root.right) printRow(root.right, indent, true, out, printNode);
  }
}

export default class Tree<Key = number, Value = any> {
  private _comparator: Comparator<Key>;
  private _root: Node<Key, Value> | null = null;
  private _size: number = 0;

  constructor(comparator = DEFAULT_COMPARE) {
    this._comparator = comparator;
  }

  /**
   * Inserts a key, allows duplicates
   */
  public insert(key: Key, data?: Value): Node<Key, Value> {
    this._size++;
    return (this._root = insert(key, data, this._root!, this._comparator));
  }

  /**
   * Adds a key, if it is not present in the tree
   */
  public add(key: Key, data?: Value): Node<Key, Value> {
    const node = new Node(key, data);

    if (this._root === null) {
      node.left = node.right = null;
      this._size++;
      this._root = node;
    }

    const comparator = this._comparator;
    const t = splay(key, this._root, comparator);
    const cmp = comparator(key, t.key);
    if (cmp === 0) this._root = t;
    else {
      if (cmp < 0) {
        node.left = t.left;
        node.right = t;
        t.left = null;
      } else if (cmp > 0) {
        node.right = t.right;
        node.left = t;
        t.right = null;
      }
      this._size++;
      this._root = node;
    }

    return this._root;
  }

  /**
   * @param  {Key} key
   * @return {Node|null}
   */
  public remove(key: Key): void {
    this._root = this._remove(key, this._root!, this._comparator);
  }

  /**
   * Deletes i from the tree if it's there
   */
  private _remove(i: Key, t: Node<Key, Value>, comparator: Comparator<Key>) {
    let x;
    if (t === null) return null;
    t = splay(i, t, comparator);
    const cmp = comparator(i, t.key);
    if (cmp === 0) {
      /* found it */
      if (t.left === null) {
        x = t.right!;
      } else {
        x = splay(i, t.left, comparator);
        x.right = t.right;
      }
      this._size--;
      return x;
    }
    return t; /* It wasn't there */
  }

  /**
   * Removes and returns the node with smallest key
   */
  public pop(): { key: Key; data: Value } | null {
    let node = this._root;
    if (node) {
      while (node.left) node = node.left;
      this._root = splay(node.key, this._root!, this._comparator);
      this._root = this._remove(node.key, this._root, this._comparator);
      return { key: node.key, data: node.data };
    }
    return null;
  }

  /**
   * Find without splaying
   */
  public findStatic(key: Key): Node<Key, Value> | null {
    let current = this._root;
    const compare = this._comparator;
    while (current) {
      const cmp = compare(key, current.key);
      if (cmp === 0) return current;
      else if (cmp < 0) current = current.left;
      else current = current.right;
    }
    return null;
  }

  public find(key: Key): Node<Key, Value> | null {
    if (this._root) {
      this._root = splay(key, this._root, this._comparator);
      if (this._comparator(key, this._root.key) !== 0) return null;
    }
    return this._root;
  }

  public contains(key: Key): boolean {
    let current = this._root;
    const compare = this._comparator;
    while (current) {
      const cmp = compare(key, current.key);
      if (cmp === 0) return true;
      else if (cmp < 0) current = current.left;
      else current = current.right;
    }
    return false;
  }

  public forEach(visitor: Visitor<Key, Value>, ctx?: any): Tree<Key, Value> {
    let current = this._root;
    const Q = []; /* Initialize stack s */
    let done = false;

    while (!done) {
      if (current !== null) {
        Q.push(current);
        current = current.left;
      } else {
        if (Q.length !== 0) {
          current = Q.pop()!;
          visitor.call(ctx, current);

          current = current.right;
        } else done = true;
      }
    }
    return this;
  }

  /**
   * Walk key range from `low` to `high`. Stops if `fn` returns a value.
   */
  public range(
    low: Key,
    high: Key,
    fn: Visitor<Key, Value>,
    ctx?: any
  ): Tree<Key, Value> {
    const Q = [];
    const compare = this._comparator;
    let node = this._root;
    let cmp;

    while (Q.length !== 0 || node) {
      if (node) {
        Q.push(node);
        node = node.left;
      } else {
        node = Q.pop()!;
        cmp = compare(node.key, high);
        if (cmp > 0) {
          break;
        } else if (compare(node.key, low) >= 0) {
          if (fn.call(ctx, node)) return this; // stop if smth is returned
        }
        node = node.right;
      }
    }
    return this;
  }

  /**
   * Returns array of keys
   */
  public keys(): Key[] {
    const keys: Key[] = [];
    this.forEach(({ key }) => {
      keys.push(key);
    });
    return keys;
  }

  /**
   * Returns array of all the data in the nodes
   */
  public values(): Value[] {
    const values: Value[] = [];
    this.forEach(({ data }) => {
      values.push(data);
    });
    return values;
  }

  public min(): Key | null {
    if (this._root) return this.minNode(this._root)!.key;
    return null;
  }

  public max(): Key | null {
    if (this._root) return this.maxNode(this._root)!.key;
    return null;
  }

  public minNode(t = this._root): Node<Key, Value> | null {
    if (t) while (t.left) t = t.left;
    return t;
  }

  public maxNode(t = this._root): Node<Key, Value> | null {
    if (t) while (t.right) t = t.right;
    return t;
  }

  /**
   * Returns node at given index
   */
  public at(index: number): Node<Key, Value> | null {
    let current = this._root;
    let done = false;
    let i = 0;
    const Q = [];

    while (!done) {
      if (current) {
        Q.push(current);
        current = current.left;
      } else {
        if (Q.length > 0) {
          current = Q.pop()!;
          if (i === index) return current;
          i++;
          current = current.right!;
        } else done = true;
      }
    }
    return null;
  }

  public next(d: Node<Key, Value>): Node<Key, Value> | null {
    let root = this._root;
    let successor = null;

    if (d.right) {
      successor = d.right;
      while (successor.left) successor = successor.left;
      return successor;
    }

    const comparator = this._comparator;
    while (root) {
      const cmp = comparator(d.key, root.key);
      if (cmp === 0) break;
      else if (cmp < 0) {
        successor = root;
        root = root.left;
      } else root = root.right;
    }

    return successor;
  }

  public prev(d: Node<Key, Value>): Node<Key, Value> | null {
    let root = this._root;
    let predecessor = null;

    if (d.left !== null) {
      predecessor = d.left;
      while (predecessor.right) predecessor = predecessor.right;
      return predecessor;
    }

    const comparator = this._comparator;
    while (root) {
      const cmp = comparator(d.key, root.key);
      if (cmp === 0) break;
      else if (cmp < 0) root = root.left;
      else {
        predecessor = root;
        root = root.right;
      }
    }
    return predecessor;
  }

  public clear(): Tree<Key, Value> {
    this._root = null;
    this._size = 0;
    return this;
  }

  public toList() {
    return toList(this._root!);
  }

  /**
   * Bulk-load items. Both array have to be same size
   */
  public load(keys: Key[], values: Value[] = [], presort: boolean = false) {
    let size = keys.length;
    const comparator = this._comparator;

    // sort if needed
    if (presort) sort(keys, values, 0, size - 1, comparator);

    if (this._root === null) {
      // empty tree
      this._root = loadRecursive(keys, values, 0, size);
      this._size = size;
    } else {
      // that re-builds the whole tree from two in-order traversals
      const mergedList = mergeLists(
        this.toList(),
        createList(keys, values),
        comparator
      );
      size = this._size + size;
      this._root = sortedListToBST({ head: mergedList }, 0, size);
    }
    return this;
  }

  public isEmpty(): boolean {
    return this._root === null;
  }

  get size(): number {
    return this._size;
  }
  get root(): Node<Key, Value> | null {
    return this._root;
  }

  public toString(
    printNode: NodePrinter<Key, Value> = (n) => String(n.key)
  ): string {
    const out: string[] = [];
    printRow(this._root!, "", true, (v) => out.push(v), printNode);
    return out.join("");
  }

  public update(key: Key, newKey: Key, newData?: Value): void {
    const comparator = this._comparator;
    let { left, right } = split(key, this._root!, comparator);
    if (comparator(key, newKey) < 0) {
      right = insert(newKey, newData, right!, comparator);
    } else {
      left = insert(newKey, newData, left!, comparator);
    }
    this._root = merge(left, right, comparator);
  }

  public split(key: Key) {
    return split(key, this._root!, this._comparator);
  }

  *[Symbol.iterator]() {
    let current = this._root;
    const Q: Node<Key, Value>[] = []; /* Initialize stack s */
    let done = false;

    while (!done) {
      if (current !== null) {
        Q.push(current);
        current = current.left;
      } else {
        if (Q.length !== 0) {
          current = Q.pop()!;
          yield current;

          current = current.right;
        } else done = true;
      }
    }
  }
}

function loadRecursive(
  keys: Key[],
  values: Value[],
  start: number,
  end: number
): Node<Key, Value> | null {
  const size = end - start;
  if (size > 0) {
    const middle = start + Math.floor(size / 2);
    const key = keys[middle];
    const data = values[middle];
    const node = new Node(key, data);
    node.left = loadRecursive(keys, values, start, middle);
    node.right = loadRecursive(keys, values, middle + 1, end);
    return node;
  }
  return null;
}

function createList(keys: Key[], values: Value[]): Node<Key, Value> {
  const head = new Node<Key, Value>(null, null);
  let p: Node<Key, Value> = head;
  for (let i = 0; i < keys.length; i++) {
    p = p.next = new Node(keys[i], values[i]);
  }
  p.next = null;
  return head.next!;
}

function toList(root: Node<Key, Value>): Node<Key, Value> {
  let current = root;
  const Q = [];
  let done = false;

  const head = new Node<Key, Value>(null, null);
  let p = head;

  while (!done) {
    if (current) {
      Q.push(current);
      current = current.left!;
    } else {
      if (Q.length > 0) {
        current = p = p.next = Q.pop()!;
        current = current.right!;
      } else done = true;
    }
  }
  p.next = null; // that'll work even if the tree was empty
  return head.next!;
}

function sortedListToBST(
  list: TreeNodeList<Key, Value>,
  start: number,
  end: number
) {
  const size = end - start;
  if (size > 0) {
    const middle = start + Math.floor(size / 2);
    const left = sortedListToBST(list, start, middle);

    const root = list.head;
    root!.left = left;

    list.head = list.head!.next;

    root!.right = sortedListToBST(list, middle + 1, end);
    return root!;
  }
  return null;
}

function mergeLists<Key, Value>(
  l1: Node<Key, Value>,
  l2: Node<Key, Value>,
  compare: Comparator<Key>
): Node<Key, Value> {
  const head: Node<Key, Value> = new Node<Key, Value>(null as Key, null); // dummy
  let p = head;

  let p1: Node<Key, Value> = l1;
  let p2: Node<Key, Value> = l2;

  while (p1 !== null && p2 !== null) {
    if (compare(p1.key, p2.key) < 0) {
      p.next = p1;
      p1 = p1.next!;
    } else {
      p.next = p2;
      p2 = p2.next!;
    }
    p = p.next;
  }

  if (p1 !== null) {
    p.next = p1;
  } else if (p2 !== null) {
    p.next = p2;
  }

  return head.next!;
}

function sort(
  keys: Key[],
  values: Value[],
  left: number,
  right: number,
  compare: Comparator<Key>
) {
  if (left >= right) return;

  const pivot = keys[(left + right) >> 1];
  let i = left - 1;
  let j = right + 1;

  while (true) {
    do i++;
    while (compare(keys[i], pivot) < 0);
    do j--;
    while (compare(keys[j], pivot) > 0);
    if (i >= j) break;

    let tmp = keys[i];
    keys[i] = keys[j];
    keys[j] = tmp;

    tmp = values[i];
    values[i] = values[j];
    values[j] = tmp;
  }

  sort(keys, values, left, j, compare);
  sort(keys, values, j + 1, right, compare);
}