/**
 * @flow
 */

type Node<V> = {|
	left: ?Node<V>,
	right: ?Node<V>,
	parent: ?Node<V>,
	height: number,
	value: V,
|};

export class Tree<V, K = V> {
	root: ?Node<V>;
	cmp: (K, V) => number;
	key: V => K;

	constructor(cmp: (K, V) => number, key: V => K) {
		this.cmp = cmp;
		this.key = key;
	}

	find(k: K): ?V {
		const cmp = this.cmp;
		let node: ?Node<V> = this.root;
		while (node != null) {
			const c = cmp(k, node.value);
			if (c === 0) {
				return node.value;
			} else if (c > 0) {
				node = node.right;
			} else {
				node = node.left;
			}
		}
	}

	insert(v: V, k: ?K): Node<V> {
		if (k == null) k = this.key(v);
		if (this.root == null) {
			return (this.root = {
				value: v,
				height: 1,
				left: null,
				right: null,
				parent: null,
			});
		}
		const cmp = this.cmp;
		let node: Node<V> = this.root;
		let result;
		while (result == null) {
			const c = cmp(k, node.value);
			if (c === 0) {
				node.value = v;
				return node;
			} else if (c < 0) {
				if (node.left == null) {
					result = {
						parent: node,
						ret: (node.left = {
							value: v,
							parent: node,
							height: 1,
							left: null,
							right: null,
						}),
					};
				} else {
					node = node.left;
				}
			} else {
				if (node.right == null) {
					result = {
						parent: node,
						ret: (node.right = {
							value: v,
							parent: node,
							height: 1,
							left: null,
							right: null,
						}),
					};
				} else {
					node = node.right;
				}
			}
		}
		this.fixHeights(result.parent);
		return result.ret;
	}

	remove(k: K): ?V {
		const cmp = this.cmp;
		let node: ?Node<V> = this.root;
		while (node != null) {
			const c = cmp(k, node.value);
			if (c === 0) {
				break;
			} else if (c > 0) {
				node = node.right;
			} else {
				node = node.left;
			}
		}
		if (node == null) {
			return;
		}

		if (node.left == null && node.right == null) {
			if (node.parent == null) {
				this.root = null;
			} else if (node.parent.left === node) {
				node.parent.left = null;
			} else {
				node.parent.right = null;
			}
		} else if (node.left != null && node.right != null) {
			const oldNode = node;
			while (node.left != null) {
				node = node.left;
			}
			oldNode.value = node.value;
			node.parent.left = node.right;
			if (node.right != null) {
				node.right.parent = node.parent;
			}
		} else if (node.left == null) {
			if (node.parent == null) {
				this.root = node.right;
				node.right.parent = null;
			} else {
				if (node.parent.left === node) {
					node.parent.left = node.right;
				} else {
					node.parent.right = node.right;
				}
				node.right.parent = node.parent;
			}
		} else {
			if (node.parent == null) {
				this.root = node.left;
				node.left.parent = null;
			} else {
				if (node.parent.right === node) {
					node.parent.right = node.left;
				} else {
					node.parent.left = node.left;
				}
				node.left.parent = node.parent;
			}
		}
		if (node.parent != null) {
			const parent = node.parent;
			node.parent = null;
			node.left = null;
			node.right = null;
			this.fixHeights(parent);
		}
		return node.value;
	}

	fixHeights(node: Node<V>) {
		let parent = node;
		while (parent != null) {
			let {left, right} = parent;
			const leftHeight = left == null ? 0 : left.height;
			const rightHeight = right == null ? 0 : right.height;
			if (Math.abs(leftHeight - rightHeight) > 1) {
				const y = leftHeight > rightHeight ? left : right;
				if (y == null) throw new Error();
				let z;
				{
					const {left, right} = y;
					const leftHeight = left == null ? 0 : left.height;
					const rightHeight = right == null ? 0 : right.height;
					z = leftHeight > rightHeight ? left : right;
				}
				if (z == null) throw new Error();
				const connection: ?Node<V> = parent.parent;
				let isRightChild: ?boolean = connection && connection.right === parent;
				parent = this.restruct(parent, y, z);
				if (connection != null) {
					if (isRightChild) {
						connection.right = parent;
					} else {
						connection.left = parent;
					}
				} else {
					this.root = parent;
				}
				parent.parent = connection;
			} else {
				parent.height = Math.max(leftHeight, rightHeight) + 1;
			}
			parent = parent.parent;
		}
	}

	restruct(x: Node<V>, y: Node<V>, z: Node<V>): Node<V> {
		let a, b, c, t1, t2, t3, t4;
		if (x.left === y) {
			if (y.left === z) {
				a = z;
				b = y;
				c = x;
				t1 = a.left;
				t2 = a.right;
				t3 = b.right;
				t4 = c.right;
			} else {
				a = y;
				b = z;
				c = x;
				t1 = a.left;
				t2 = b.left;
				t3 = b.right;
				t4 = c.right;
			}
		} else {
			if (y.right === z) {
				a = x;
				b = y;
				c = z;
				t1 = a.left;
				t2 = b.left;
				t3 = c.left;
				t4 = c.right;
			} else {
				a = x;
				b = z;
				c = y;
				t1 = a.left;
				t2 = b.left;
				t3 = b.right;
				t4 = c.right;
			}
		}
		a.left = t1;
		a.right = t2;
		let h1 = 0;
		let h2 = 0;
		if (t1) {
			t1.parent = a;
			h1 = t1.height;
		}
		if (t2) {
			t2.parent = a;
			h2 = t2.height;
		}
		a.height = Math.max(h1, h2) + 1;
		h1 = 0;
		h2 = 0;
		c.left = t3;
		c.right = t4;
		if (t3) {
			t3.parent = c;
			h1 = t3.height;
		}
		if (t4) {
			t4.parent = c;
			h2 = t4.height;
		}
		c.height = Math.max(h1, h2) + 1;
		b.left = a;
		b.right = c;
		a.parent = b;
		c.parent = b;
		b.height = Math.max(a.height, c.height) + 1;
		return b;
	}

	inOrder(fn: V => void) {
		let node = this.root;
		const stack = [];
		while (stack.length !== 0 || node != null) {
			while (node != null) {
				stack.push(node);
				node = node.left;
			}
			node = stack.pop();
			fn(node.value);
			node = node.right;
		}
	}

	preOrder(fn: V => void) {
		if (this.root == null) return;
		const stack = [this.root];
		while (stack.length > 0) {
			let node = stack.pop();
			fn(node.value);
			if (node.right) stack.push(node.right);
			if (node.left) stack.push(node.left);
		}
	}

	postOrder(fn: V => void) {
		if (this.root == null) return;
		const stack = [this.root];
		const array = [];
		while (stack.length > 0) {
			let node = stack.pop();
			node = stack.pop();
			if (node.right) stack.push(node.right);
			if (node.left) stack.push(node.left);
			array.push(node.value);
		}
		for (let i = array.length - 1; i >= 0; i--) {
			fn(array[i]);
		}
	}

	range(start: K, end: K): V[] {
		let node = this.root;
		const cmp = this.cmp;
		const result: V[] = [];
		const stack = [];
		while (stack.length > 0 || node != null) {
			while (node != null) {
				stack.push(node);
				if (cmp(start, node.value) < 0) {
					node = node.left;
				} else {
					node = null;
				}
			}
			node = stack.pop();
			if (cmp(start, node.value) <= 0 && cmp(end, node.value) >= 0) {
				result.push(node.value);
			}
			if (node.right && cmp(end, node.value) > 0) {
				node = node.right;
			} else {
				node = null;
			}
		}
		return result;
	}

	first(): ?V {
		const min = this.min(this.root);
		return min && min.value;
	}

	last(): ?V {
		const max = this.max(this.root);
		return max && max.value;
	}

	min(node: ?Node<V> = this.root): ?Node<V> {
		if (node == null) return;
		while (node.left != null) {
			node = node.left;
		}
		return node;
	}

	max(node: ?Node<V> = this.root): ?Node<V> {
		if (node == null) return;
		while (node.right != null) {
			node = node.right;
		}
		return node;
	}

	next(node: Node<V>): ?Node<V> {
		if (node.right != null) {
			return this.min(node.right);
		}
		while (node.parent != null && node.parent.right === node) {
			node = node.parent;
		}
		return node.parent;
	}

	prev(node: Node<V>): ?Node<V> {
		if (node.left != null) {
			return this.min(node.left);
		}
		while (node.parent != null && node.parent.left === node) {
			node = node.parent;
		}
		return node.parent;
	}

	toString(toString: V => string, node: ?Node<V> = this.root) {
		if (node == null) return '';
		let children =
			'[' +
			this.toString(toString, node.left) +
			',' +
			this.toString(toString, node.right) +
			']';
		if (node.left == null && node.right == null) return toString(node.value);
		return toString(node.value) + (children == '[,]' ? '' : children);
	}

	checkAVL(
		node: ?Node<V> = this.root,
		map: Map<Node<V>, number> = new Map<Node<V>, number>(),
	): boolean {
		if (node == null) return true;
		const left = this.checkAVL(node.left, map);
		const right = this.checkAVL(node.right, map);
		if (!(left && right)) return false;
		const leftHeight = node.left ? map.get(node.left) : 0;
		const rightHeight = node.right ? map.get(node.right) : 0;
		if (leftHeight == null || rightHeight == null) throw new Error();
		if (Math.abs(leftHeight - rightHeight) > 1) return false;
		map.set(node, Math.max(leftHeight, rightHeight) + 1);
		return true;
	}
}