#!/usr/bin/env ruby
#
# Fri Feb 11 13:36:25 EST 2011
#
# a learning example of a Left-Leaning Red-Black Tree, in ruby
#
# you can import this as a library,
# irb -r ./rbtree.rb
# irb 0:01> run(10000)
#
# or
# run it as a script
# $> ruby rbtree.rb [sample_size]
# which will output the results of a descending trail of sample_size/2
# default sample size is 10,000
#

class RBTree
  include Comparable
  RED = true
  BLACK = false
  attr_accessor :root

  def initialize
    self.root = RBNode.new()
  end

  def insert(key,val)
    self.root = insert_local(self.root, key, val)
    self.root.color = BLACK
  end

  def insert_local(node, key, val)
    if node.nil?
      return RBNode.new(key, val)
    end
    if node.key.nil?
      return RBNode.new(key, val)
    end

    # Uncomment this, for a 2-3-4 node tree implementation
    # and comment about below
#    if (isRed(node.right) && isRed(node.left))
#      flipColors(node)
#    end

    if (key == node.key)
      #dup
    elsif (key > node.key)
      node.right = insert_local(node.right, key, val)
    else
      node.left = insert_local(node.left, key, val)
    end

    if (isRed(node.right) && !isRed(node.left))
      node = rotateLeft(node)
    end
    if (isRed(node.left) && isRed(node.left.left))
      node = rotateRight(node)
    end

    # with this test here, this acts as a 2-3 node tree
    # which i've seen a decent performance benefit of,
    # over the 2-3-4 node tree
    if (isRed(node.right) && isRed(node.left))
      flipColors(node)
    end

    return node
  end
  private :insert_local

  def search(key,val=nil)
    node = self.root

    while true
      if (key == node.key)
        return node
      elsif (key < node.key)
        node = node.left
      elsif (key > node.key)
        node = node.right
      end

      # Nothing has been found
      if node.nil?
        break
      end
    end

    return nil
  end

  def delete
  end

  def bottomLeft
    node_send = self.root
    counter = 0
    while true
      node_recv = bottomLeft_local(node_send)
      if node_send == node_recv
        return [counter, node_recv]
      else
        counter += 1
        node_send = node_recv
      end
    end
  end

  def bottomLeft_local(node)
    if node.left.nil?
      return node
    else
      return node.left
    end
  end
  private :bottomLeft_local

  def bottomRight
    node_send = self.root
    counter = 0
    while true
      node_recv = bottomRight_local(node_send)
      if node_send == node_recv
        return [counter, node_recv]
      else
        counter += 1
        node_send = node_recv
      end
    end
  end

  def bottomRight_local(node)
    if node.right.nil?
      return node
    else
      return node.right
    end
  end
  private :bottomRight_local

  def rotateLeft(node)
    x = node.right
    node.right = x.left
    x.left = node
    x.color = node.color
    node.color = RED
    return x
  end

  def rotateRight(node)
    x = node.left
    node.left = x.right
    x.right = node
    x.color = node.color
    node.color = RED
    return x
  end

  def isRed(node)
    if node.nil?
      return false
    else
      return node.color
    end
  end

  def flipColors(node)
    node.color = !node.color
    node.right.color = !node.right.color
    node.left.color = !node.left.color
  end

  def inspect
    "#<%s:0x%p root=%s>" % [self.class.name, self.object_id, self.root]
  end

  class RBNode
    attr_accessor :key, :val, :color, :left, :right
    def initialize(key = nil, val = nil)
      self.key = key
      self.val = val
      self.color = RED
    end
    def <=>(another)
      self.key <=> another.key
    end
    def inspect
      "#<%s:0x%p key=\"%s\" val=\"%s\" color=\"%s\">" % [self.class.name, self.object_id, self.key, self.val, self.color ? "RED" : "BLACK"]
    end
  end
end

def uniq_rand_set(f = 100)
  (0..f).map {|i| i }.shuffle
end

def uniq_rand(f = 100)
  if @rand_set.nil?
    @rand_set = uniq_rand_set(f)
  end
  return @rand_set.pop
end

def run(sample_size = 10000)
  tree = RBTree.new
  count = sample_size

  while count > 0
    tree.insert(uniq_rand(sample_size), nil)
    count -= 1
  end

  @rand_set = nil
  return tree
end

def run_loop(start_size)
  @results = []
  local_size = start_size

  while true
    time_start = Time.now.to_f
    t = run(local_size)
    time_stop = Time.now.to_f
    diff = time_stop - time_start


    @results << [local_size, diff, t.bottomLeft[0], t.bottomRight[0]]
    local_size = (local_size.to_f/2).to_i

    # allow a bit of garbage collection
    t = diff = time_stop = time_start = nil

    if local_size <= 1
      break
    end

  end

  return @results
end

times=10000
if ARGV[0]
  times = ARGV[0].to_i
end
p run_loop(times)