#!/usr/bin/ruby
require 'forwardable'
require 'pp'

#
# BPlusTree class
#
class BPlusTree
  extend Forwardable
  def_delegators :@tree, :root, :add, :delete, :get, :to_a, :disp
  attr_accessor :tree

  def initialize(param = {})
    @tree = BPlusTree.new(param)
  end

  def BPlusTree.new_from_assoc(assoc, param = {})
    ret = new(param)
    ret.tree = BPlusTree.new_from_assoc(assoc, param)
    return ret
  end


  class BPlusTree
    attr_accessor :root, :param

    def initialize(param = {})
      self.param = param
      self.root = Leaf.new(self.param)
    end

    def BPlusTree.new_from_assoc(assoc, param = {})
      new(param).create_bptree_from_assoc(assoc)
    end

    def add(key, value)
      ret, _ = self.root.add(key, value)
      increase_height if ret == :need_split
      return key
    end

    def delete(key)
      ret, di = self.root.delete(key)
      if ret == :need_balance and self.root.suitable_root? then
        decrease_height
      end
      return((di.nil?) ? nil : key)
    end

    def get(key)
      self.root.get(key)
    end

    def to_a
      self.root.to_a
    end

    def disp(with_data = false)
      self.root.disp(0, with_data)
    end

    def create_bptree_from_assoc(assoc)
      raise 'Multiple root node occured' if assoc.size > 1
      type, key, data = assoc[0]
      root = {:leaf => Leaf, :branch => Branch}[type].new_from_assoc(assoc[0])
      root.connect_nextp(nil)
      self.root = root
      return self
    end


    private
    def increase_height
      # create a new root
      newroot = Branch.new(self.param)

      # connect children to new root
      lbrother = self.root
      bbrother = self.root.split
      (newroot.add_child(lbrother).nil?  and
       newroot.add_child(bbrother).nil?) or
        raise("#{self.class}#increase_height: Internal error")

      self.root = newroot
    end

    def decrease_height
      self.root = self.root.get_child
    end
  end


  class Node
    attr_accessor :key, :param

    def initialize(param = {})
      self.key = []
      self.param = param

      @param[:maxentry]  ||= 3
      @param[:maxbranch] ||= 3
      @param[:minentry]  ||= (@param[:maxentry]  + 1) / 2
      @param[:minbranch] ||= (@param[:maxbranch] + 1) / 2

      return self
    end

    def maxentry
      @param[:maxentry]
    end
    def maxbranch
      @param[:maxbranch]
    end
    def minentry
      @param[:minentry]
    end
    def minbranch
      @param[:minbranch]
    end


    def min
      return self.key.first
    end

    protected
    def search_insert_pos(key)
      for i in 0..self.key.size - 1
        return i if key < self.key[i]
      end
      return self.key.size
    end
  end


  class Leaf < Node
    attr_accessor :data, :nextp

    def initialize(param = {})
      self.data = []
      self.nextp = nil
      super
    end

    def Leaf.new_from_assoc(assoc, param = {})
      new(param).create_bptree_from_assoc(assoc)
    end

    def add(key, value)
      before = search_insert_pos(key)
      self.key[before, 0]  = key
      self.data[before, 0] = value

      ret = nil
      ret = :need_split if self.key.size > self.maxentry
      return [ret, before]
    end

    def delete(key)
      pos = search_key(key)
      return [nil, nil] if pos.nil?

      self.key.delete_at(pos)
      self.data.delete_at(pos)

      ret = nil
      ret = :need_balance if self.key.size < self.minentry
      return [ret, pos]
    end

    def get(key)
      pos = search_key(key)
      return nil if pos.nil?
      return self.data[pos]
    end

    def to_a
      ret = self.key.zip(self.data)
      return ret if self.nextp.nil?
      return ret +  self.nextp.to_a
    end

    def disp(level, with_data = false)
      indent = '  ' * level
      print indent + '- Leaf' + ((level == 0) ? '(root)' : '') + ' ['
      print self.key.zip(self.data).inject('') {|before, (k, d)|
          if before.empty? then
            k.to_s + (with_data ? ':' + d.to_s : '')
          else
            before + ', ' + k.to_s + (with_data ? ':' + d.to_s : '')
          end
        }
      puts ']'
    end

    def pretty_print(pp)
      pp.group(1, '#<B+T::Leaf', '>') do
        pp.breakable
        pp.text('key:')
        pp.pp(self.key)
        pp.breakable
        pp.text('data:')
        pp.pp(self.data)
      end
    end


    def split
      range = (self.key.size / 2) .. self.key.size - 1

      sibling = Leaf.new(self.param)
      sibling.key[0, 0] = self.key[range]
      sibling.data[0, 0] = self.data[range]
      self.key[range] = nil
      self.data[range] = nil
      sibling.nextp = self.nextp
      self.nextp = sibling

      return sibling
    end

    def balance(bigbrother)
      if self.key.size + bigbrother.key.size <= self.maxentry then
        return :need_integrate
      end
      if self.key.size + bigbrother.key.size >  self.maxentry * 2 then
        raise "#{self.class}#balance: Internal error"
      end

      newkey  = self.key.dup + bigbrother.key
      newdata = self.data.dup + bigbrother.data
      keys = newkey.size
      boundary = keys / 2

      self.key        = newkey[0..boundary-1]
      bigbrother.key  = newkey[boundary..keys-1]

      self.data       = newdata[0..boundary-1]
      bigbrother.data = newdata[boundary..keys-1]

      return nil
    end

    def integrate(bigbrother)
      self.key  += bigbrother.key
      self.data += bigbrother.data
      self.nextp = bigbrother.nextp
      bigbrother.nextp = nil

      raise "#{self.class}: Internal error" if self.key.size > self.maxentry
      return nil
    end

    def get_child
      return self
    end

    def suitable_root?
      return true
    end

    def create_bptree_from_assoc(assoc)
      type, key, data = assoc
      unless type == :leaf then
        raise "#{self.class}#create_bptree_from_assoc: Internal error"
      end
      self.key = key
      self.data = data
      return self
    end

    def connect_nextp(before = nil)
      before.nextp = self unless before.nil?
      return self
    end


    private
    def search_key(key)
      for i in 0..self.key.size - 1
        return i if self.key[i] == key
      end
      return nil
    end
  end


  class Branch < Node
    attr_accessor :branch

    def initialize(param = {})
      self.branch = []
      super
    end

    def Branch.new_from_assoc(assoc, param = {})
      new(param).create_bptree_from_assoc(assoc)
    end

    def add(key, value)
      bi = search_branch(key)
      child = self.branch[bi]
      ret, before = child.add(key, value)
      self.key[bi] = child.key[before] if before == 0

      return [nil, bi] if ret.nil?
      raise "#{self.class}#add: Internal error" unless ret == :need_split

      # split child
      bbrother = child.split
      ret = self.add_child(bbrother)

      return [ret, bi] if ret.nil? or ret == :need_split
      raise "#{self.class}#add: Internal error"
    end

    def delete(key)
      bi = search_branch(key)
      child = self.branch[bi]
      ret, di = child.delete(key)
      self.key[bi] = child.key[di] if di == 0

      bj = nil
      until ret.nil?
        case ret
        when :need_balance
          # choose pair.
          bi, bj = self.choose_pair(bi)
          lbrother, bbrother = self.branch[bi..bj]

          # balance pair
          ret = lbrother.balance(bbrother)
          self.key[bj] = bbrother.min
        when :need_integrate
          # remove big-brother
          self.branch.delete_at(bj)
          self.key.delete_at(bj)

          # integrate pair
          ret = lbrother.integrate(bbrother)
        else
          raise "#{self.class}#delete: Internal error"
        end
      end

      ret = nil
      ret = :need_balance if self.key.size < self.minbranch
      return [ret, bi]
    end

    def get(key)
      bi = search_branch(key)
      return self.branch[bi].get(key)
    end

    def to_a
      return self.branch[0].to_a
    end

    def disp(level, with_data = false)
      indent = '  ' * level
      puts indent + '- Branch' + (level == 0 ? '(root)' : '') + ' [' +
        self.key.inject('') {|before, k|
          if before.empty? then
            "(#{k})"
          else
            before + ", #{k}"
          end
        } + ']'
      self.branch.map do |branch|
        branch.disp(level + 1, with_data)
      end
    end

    def pretty_print(pp)
      pp.group(1) do
        pp.text('#<B+T::Branch')
        pp.breakable
        pp.group(1, 'key:[', ']') do
          self.key.inject(false) do |before, key|
            pp.comma_breakable if before
            pp.pp(key)
            true
          end
        end

        pp.breakable
        pp.text('branch:')
      end

      pp.nest(1) do
        pp.breakable
        pp.group(1, '[', ']') do
          self.branch.inject(false) do |before, branch|
            pp.comma_breakable if before
            pp.pp(branch)
            true
          end
        end
      end

      pp.text('>')
    end


    def add_child(child)
      before = search_insert_pos(child.min)
      self.key[before, 0] = child.min
      self.branch[before, 0] = child

      return :need_split if self.key.size > self.maxbranch
      return nil
    end

    def split
      range = (self.key.size / 2) .. self.key.size - 1

      sibling = Branch.new(self.param)
      sibling.key[0, 0] = self.key[range]
      sibling.branch[0, 0] = self.branch[range]

      self.branch[range] = nil
      self.key[range] = nil

      return sibling
    end

    def balance(bigbrother)
      if self.key.size + bigbrother.key.size <= self.maxbranch then
        return :need_integrate
      end
      if self.key.size + bigbrother.key.size >  self.maxbranch * 2 then
        raise "#{self.class}#balance: Internal error"
      end

      # newkeyとnewbranchには，maxbranch * 2個の要素を置くための空間が必要
      newkey    = self.key.dup + bigbrother.key
      newbranch = self.branch.dup + bigbrother.branch
      keys = newkey.size
      boundary = keys / 2

      self.key       = newkey[0..boundary-1]
      bigbrother.key = newkey[boundary..keys-1]

      self.branch       = newbranch[0..boundary-1]
      bigbrother.branch = newbranch[boundary..keys-1]

      return nil
    end

    def integrate(bigbrother)
      for i in 0..bigbrother.key.size - 1
        self.add_child(bigbrother.branch[i])
      end

      raise "#{self.class}: Internal error" if self.key.size > self.maxbranch
      return nil
    end

    def get_child
      raise "#{self.class}#get_child: Internal error" if self.key.size != 1
      return self.branch[0]
    end

    def suitable_root?
      return true if self.key.size <= 1
      return false
    end

    def choose_pair(lbrother_idx)
      lbrother_idx -= 1 if lbrother_idx == self.key.size - 1
      bbrother_idx = lbrother_idx + 1
      raise "#{self.class}#choose_pair: Internal error" if key.size < 2
      return [lbrother_idx, bbrother_idx]
    end

    def create_bptree_from_assoc(assoc)
      type, key, branch = assoc
      unless type == :branch then
        raise "#{self.class}#create_bptree_from_assoc: Internal error"
      end
      self.key = key
      self.branch = branch.map do |child|
        {:leaf => Leaf, :branch => Branch}[child.first].new_from_assoc(child)
      end
      return self
    end

    def connect_nextp(before = nil)
      self.branch.inject(before) do |before, branch|
        branch.connect_nextp(before)
      end
    end


    private
    def search_branch(key)
      for i in 1..self.key.size - 1
        return i - 1 if key < self.key[i]
      end
      return self.key.size - 1
    end
  end
end

###########################################################################
# deleting test

puts '[Delete test 1]'
puts

assoc =
[[:branch, [1, 28, 55],
  [[:branch, [1, 10, 19],
    [[:branch, [1, 4, 7],
      [[:leaf, [1, 2, 3], ["data01", "data02", "data03"]],
       [:leaf, [4, 5, 6], ["data04", "data05", "data06"]],
       [:leaf, [7, 8, 9], ["data07", "data08", "data09"]]]],
     [:branch, [10, 13, 16],
      [[:leaf, [10, 11, 12], ["data10", "data11", "data12"]],
       [:leaf, [13, 14, 15], ["data13", "data14", "data15"]],
       [:leaf, [16, 17, 18], ["data16", "data17", "data18"]]]],
     [:branch, [19, 22, 25],
      [[:leaf, [19, 20, 21], ["data19", "data20", "data21"]],
       [:leaf, [22, 23, 24], ["data22", "data23", "data24"]],
       [:leaf, [25, 26, 27], ["data25", "data26", "data27"]]]]]],
   [:branch, [28, 37, 46],
    [[:branch, [28, 31, 34],
      [[:leaf, [28, 29, 30], ["data28", "data29", "data30"]],
       [:leaf, [31, 32, 33], ["data31", "data32", "data33"]],
       [:leaf, [34, 35, 36], ["data34", "data35", "data36"]]]],
     [:branch, [37, 40, 43],
      [[:leaf, [37, 38, 39], ["data37", "data38", "data39"]],
       [:leaf, [40, 41, 42], ["data40", "data41", "data42"]],
       [:leaf, [43, 44, 45], ["data43", "data44", "data45"]]]],
     [:branch, [46, 49, 52],
      [[:leaf, [46, 47, 48], ["data46", "data47", "data48"]],
       [:leaf, [49, 50, 51], ["data49", "data50", "data51"]],
       [:leaf, [52, 53, 54], ["data52", "data53", "data54"]]]]]],
   [:branch, [55, 64, 73],
    [[:branch, [55, 58, 61],
      [[:leaf, [55, 56, 57], ["data55", "data56", "data57"]],
       [:leaf, [58, 59, 60], ["data58", "data59", "data60"]],
       [:leaf, [61, 62, 63], ["data61", "data62", "data63"]]]],
     [:branch, [64, 67, 70],
      [[:leaf, [64, 65, 66], ["data64", "data65", "data66"]],
       [:leaf, [67, 68, 69], ["data67", "data68", "data69"]],
       [:leaf, [70, 71, 72], ["data70", "data71", "data72"]]]],
     [:branch, [73, 76, 79],
      [[:leaf, [73, 74, 75], ["data73", "data74", "data75"]],
       [:leaf, [76, 77, 78], ["data76", "data77", "data78"]],
       [:leaf, [79, 80, 81], ["data79", "data80", "data81"]]]]]]]]]

bptree = BPlusTree.new_from_assoc(assoc)
puts '[Initial state]'
bptree.disp
puts

xs = (1..81).to_a
random_sequence = []
random_sequence << xs.delete_at(rand(xs.size)) until xs.empty?

random_sequence.each do |i|
  bptree.delete(i)
  puts 'Delete: %d' % i
  bptree.disp
  puts
end
puts 'Delete test 1 passed'
puts


# deleting test 2
puts '[Delete test 2]'
puts

bptree = BPlusTree.new
count = 10000

puts 'adding...'
xs = (0..count).to_a
random_sequence = []
random_sequence << xs.delete_at(rand(xs.size)) until xs.empty?

random_sequence.each do |i|
  bptree.add(i, 'data%05d' % i)
end
puts 'done.'

xs = (0..count).to_a
random_sequence = []
random_sequence << xs.delete_at(rand(xs.size)) until xs.empty?

puts 'deleting...'
c = 0
random_sequence.each do |i|
  bptree.delete(i)
  c+=1
  puts c.to_s if c % 100 == 0
end
puts 'done.'
pp bptree

puts 'Delete test 2 passed'