class TreeMap

# File lib/treemap/tree_map.rb, line 150
def initialize(comparator = NaturalOrder)
  @comparator = comparator
  @root = nil
  @size = 0
  @mod_count = 0
end

# File lib/treemap/tree_map.rb, line 531
def ceiling_entry(key)
  find(key, Relation::CEILING)
end

# File lib/treemap/tree_map.rb, line 536
def ceiling_key(key)
  entry = find(key, Relation::CEILING)
  entry.key if entry
end

# File lib/treemap/tree_map.rb, line 176
def clear
  @root = nil
  @size = 0
  @mod_count += 1
end

# File lib/treemap/tree_map.rb, line 168
def contains_key?(key)
  find_by_object(key)
end

# File lib/treemap/tree_map.rb, line 617
def descending_map
  BoundedMap.new(self, false, nil, Bound::NO_BOUND, nil, Bound::NO_BOUND)
end

# File lib/treemap/tree_map.rb, line 652
def each(&blk)
  if block_given?
    each_node {|node| blk.call(node.key, node.value) }
  else
    enum_for(:each)
  end
end

# File lib/treemap/tree_map.rb, line 661
def each_node
  if block_given?
    iter = NodeIterator.new(empty? ? nil : @root.first)
    while iter.has_next?
      yield iter.step_forward()
    end
  else
    enum_for(:each_node)
  end
end

# File lib/treemap/tree_map.rb, line 157
def empty?
  @size == 0
end

# File lib/treemap/tree_map.rb, line 554
def entry_set
  Set.new(each_node.to_a)
end

# File lib/treemap/tree_map.rb, line 195
def find(key, relation)
  if @root.nil?
    if relation == Relation::CREATE
      @root = Node.new(nil, key)
      @size = 1
      @mod_count += 1
      return @root
    else
      return nil
    end
  end

  nearest = @root
  while true
    comparison = @comparator.call(key, nearest.key)

    # we found the requested key
    if comparison == 0
      case relation
      when Relation::LOWER
        return nearest.prev_node
      when Relation::FLOOR, Relation::EQUAL, Relation::CREATE, Relation::CEILING
        return nearest
      when Relation::HIGHER
        return nearest.next_node
      end
    end

    child = (comparison < 0) ? nearest.left : nearest.right
    if child
      nearest = child
      next  # continue
    end

    # We found a nearest node. Every key not in the tree has up to two nearest nodes, one lower and one higher.
    if comparison < 0     # nearest.key is higher
      case relation
      when Relation::LOWER, Relation::FLOOR
        return nearest.prev_node
      when Relation::CEILING, Relation::HIGHER
        return nearest
      when Relation::EQUAL
        return nil
      when Relation::CREATE
        created = Node.new(nearest, key)
        nearest.left = created
        @size += 1
        @mod_count += 1
        rebalance(nearest, true)
        return created
      end
    else                  # comparison > 0 ; nearest.key is lower
      case relation
      when Relation::LOWER, Relation::FLOOR
        return nearest
      when Relation::CEILING, Relation::HIGHER
        return nearest.next_node
      when Relation::EQUAL
        return nil
      when Relation::CREATE
        created = Node.new(nearest, key)
        nearest.right = created
        @size += 1
        @mod_count += 1
        rebalance(nearest, true)
        return created
      end
    end
  end
end

# File lib/treemap/tree_map.rb, line 278
def find_by_entry(key, value)
  key, value = *entry
  mine = find_by_object(key)
  mine if mine && mine.value == value
end

# File lib/treemap/tree_map.rb, line 267
def find_by_object(key)
  find(key, Relation::EQUAL)
end

# File lib/treemap/tree_map.rb, line 465
def first_entry
  root.first if root
end

# File lib/treemap/tree_map.rb, line 481
def first_key
  raise "No such element." unless root
  root.first.key
end

# File lib/treemap/tree_map.rb, line 520
def floor_entry(key)
  find(key, Relation::FLOOR)
end

# File lib/treemap/tree_map.rb, line 525
def floor_key(key)
  entry = find(key, Relation::FLOOR)
  entry.key if entry
end

# File lib/treemap/tree_map.rb, line 161
def get(key)
  entry = find_by_object(key)
  entry.value if entry
end

# File lib/treemap/tree_map.rb, line 589
def head_map(*args)
  case args.count
  when 1
    to_exclusive = args.first
    BoundedMap.new(self, true, nil, Bound::NO_BOUND, to_exclusive, Bound::EXCLUSIVE)
  when 2
    to, inclusive = *args
    to_bound = inclusive ? Bound::INCLUSIVE : Bound::EXCLUSIVE
    BoundedMap.new(self, true, nil, Bound::NO_BOUND, to, to_bound)
  end
end

# File lib/treemap/tree_map.rb, line 542
def higher_entry(key)
  find(key, Relation::HIGHER)
end

# File lib/treemap/tree_map.rb, line 547
def higher_key(key)
  entry = find(key, Relation::HIGHER)
  entry.key if entry
end

# File lib/treemap/tree_map.rb, line 469
def internal_poll_first_entry
  if root
    result = root.first
    remove_internal(result)
    result
  end
end

# File lib/treemap/tree_map.rb, line 491
def internal_poll_last_entry
  if root
    result = root.last
    remove_internal(result)
    result
  end
end

# File lib/treemap/tree_map.rb, line 558
def key_set
  Set.new(each_node.map(&:key))
end

# File lib/treemap/tree_map.rb, line 487
def last_entry
  root.last if root
end

# File lib/treemap/tree_map.rb, line 503
def last_key
  raise "No such element." unless root
  root.last.key
end

# File lib/treemap/tree_map.rb, line 509
def lower_entry(key)
  find(key, Relation::LOWER)
end

# File lib/treemap/tree_map.rb, line 514
def lower_key(key)
  entry = find(key, Relation::LOWER)
  entry.key if entry
end

# File lib/treemap/tree_map.rb, line 477
def poll_first_entry
  internal_poll_first_entry
end

# File lib/treemap/tree_map.rb, line 499
def poll_last_entry
  internal_poll_last_entry
end

# File lib/treemap/tree_map.rb, line 172
def put(key, value)
  put_internal(key, value)
end

# File lib/treemap/tree_map.rb, line 187
def put_internal(key, value)
  created = find(key, Relation::CREATE)
  result = created.value
  created.value = value
  result
end

# File lib/treemap/tree_map.rb, line 362
def rebalance(unbalanced, insert)
  node = unbalanced
  while node
    left = node.left
    right = node.right
    left_height = left ? left.height : 0
    right_height = right ? right.height : 0

    delta = left_height - right_height
    if delta == -2
      right_left = right.left
      right_right = right.right
      right_right_height = right_right ? right_right.height : 0
      right_left_height = right_left ? right_left.height : 0

      right_delta = right_left_height - right_right_height
      if right_delta == -1 || (right_delta == 0 && !insert)
        rotate_left(node)
      else
        # assert (right_delta == 1)
        rotate_right(right)   # AVL right left
        rotate_left(node)
      end
      break if insert   # no further rotations will be necessary
    elsif delta == 2
      left_left = left.left
      left_right = left.right
      left_right_height = left_right ? left_right.height : 0
      left_left_height = left_left ? left_left.height : 0

      left_delta = left_left_height - left_right_height
      if left_delta == 1 || (left_delta == 0 && !insert)
        rotate_right(node)    # AVL left left
      else
        # assert (left_delta == -1)
        rotate_left(left)   # AVL left right
        rotate_right(node)
      end
      break if insert
    elsif delta == 0
      node.height = left_height + 1   # left_height == right_height
      break if insert
    else
      # assert (delta == -1 || delta == 1)
      node.height = [left_height, right_height].max + 1
      break unless insert    # the height hasn't changed, so rebalancing is done!
    end

    node = node.parent
  end
end

# File lib/treemap/tree_map.rb, line 182
def remove(key)
  node = remove_internal_by_key(key)
  node.value if node
end

# File lib/treemap/tree_map.rb, line 286
def remove_internal(node)
  left = node.left
  right = node.right
  original_parent = node.parent

  if left && right
    # To remove a node with both left and right subtrees, move an adjacent node from one of those subtrees into this node's place.
    # Removing the adjacent node may change this node's subtrees. This node may no longer have two subtrees once the adjacent node is gone!

    adjacent = left.height > right.height ? left.last : right.first
    remove_internal(adjacent)   # takes care of rebalance and size--

    left_height = 0
    left = node.left
    if left
      left_height = left.height
      adjacent.left = left
      left.parent = adjacent
      node.left = nil
    end
    right_height = 0
    right = node.right
    if right
      right_height = right.height
      adjacent.right = right
      right.parent = adjacent
      node.right = nil
    end
    adjacent.height = [left_height, right_height].max + 1
    replace_in_parent(node, adjacent)
    return
  elsif left
    replace_in_parent(node, left)
    node.left = nil
  elsif right
    replace_in_parent(node, right)
    node.right = nil
  else
    replace_in_parent(node, nil)
  end

  rebalance(original_parent, false)
  @size -= 1
  @mod_count -= 1
end

# File lib/treemap/tree_map.rb, line 332
def remove_internal_by_key(key)
  node = find_by_object(key)
  if node
    remove_internal(node)
  end
  node
end

# File lib/treemap/tree_map.rb, line 340
def replace_in_parent(node, replacement)
  parent = node.parent
  node.parent = nil
  if replacement
    replacement.parent = parent
  end

  if parent
    if parent.left == node
      parent.left = replacement
    else
      # assert (parent.right == node)
      parent.right = replacement
    end
  else
    @root = replacement
  end
end

# File lib/treemap/tree_map.rb, line 415
def rotate_left(root)
  left = root.left
  pivot = root.right
  pivot_left = pivot.left
  pivot_right = pivot.right

  # move the pivot's left child to the root's right
  root.right = pivot_left
  if pivot_left
    pivot_left.parent = root
  end

  replace_in_parent(root, pivot)

  # move the root to the pivot's left
  pivot.left = root
  root.parent = pivot

  # fix heights
  root.height = [left ? left.height : 0, pivot_left ? pivot_left.height : 0].max + 1
  pivot.height = [root.height, pivot_right ? pivot_right.height : 0].max + 1
end

# File lib/treemap/tree_map.rb, line 439
def rotate_right(root)
  pivot = root.left
  right = root.right
  pivot_left = pivot.left
  pivot_right = pivot.right

  # move the pivot's right child to the root's left
  root.left = pivot_right
  if pivot_right
    pivot_right.parent = root
  end

  replace_in_parent(root, pivot)

  # move the root to the pivot's right
  pivot.right = root
  root.parent = pivot

  # fix heights
  root.height = [right ? right.height : 0, pivot_right ? pivot_right.height : 0].max + 1
  pivot.height = [root.height, pivot_left ? pivot_left.height : 0].max + 1
end

# File lib/treemap/tree_map.rb, line 572
def sub_map(*args)
  case args.count
  when 2
    from_inclusive, to_exclusive = *args
    BoundedMap.new(self, true, from_inclusive, Bound::INCLUSIVE, to_exclusive, Bound::EXCLUSIVE)
  when 4
    from, from_inclusive, to, to_inclusive = *args
    from_bound = from_inclusive ? Bound::INCLUSIVE : Bound::EXCLUSIVE
    to_bound = to_inclusive ? Bound::INCLUSIVE : Bound::EXCLUSIVE
    BoundedMap.new(self, true, from, from_bound, to, to_bound);
  end
end

# File lib/treemap/tree_map.rb, line 605
def tail_map(*args)
  case args.count
  when 1
    from_inclusive = args.first
    BoundedMap.new(self, true, from_inclusive, Bound::INCLUSIVE, nil, Bound::NO_BOUND)
  when 2
    from, inclusive = *args
    from_bound = inclusive ? Bound::INCLUSIVE : Bound::EXCLUSIVE
    BoundedMap.new(self, true, from, from_bound, nil, Bound::NO_BOUND)
  end
end

# File lib/treemap/tree_map.rb, line 564
def values
  each_node.map(&:value)
end