#!/usr/bin/env ruby
#
# Note: Ruby 1.9 is faster than 1.8, as expected.

$LOAD_PATH[0,0] = File.join(File.dirname(__FILE__), '..', 'lib')
require 'mongo'

include Mongo

TRIALS = 2
PER_TRIAL = 5000
BATCH_SIZE = 100

SMALL = {}
MEDIUM = {
  'integer' => 5,
  'number' => 5.05,
  'boolean' => false,
  'array' => ['test', 'benchmark']
}
LARGE = {
  'base_url' => 'http://www.example.com/test-me',
  'total_word_count' => 6743,
  'access_time' => Time.now,
  'meta_tags' => {
    'description' => 'i am a long description string',
    'author' => 'Holly Man',
    'dynamically_created_meta_tag' => 'who know\n what'
  },
  'page_structure' => {
    'counted_tags' => 3450,
    'no_of_js_attached' => 10,
    'no_of_images' => 6
  },
  'harvested_words' => ['10gen','web','open','source','application','paas',
                        'platform-as-a-service','technology','helps',
                        'developers','focus','building','mongodb','mongo'] * 20
}

require 'benchmark'

def report(str, t)
  printf("%s%d\n", str.ljust(60, '.'), PER_TRIAL / t)
end

def profile(str)
  if ENV['MONGO_PROFILE']
    require 'rubygems'
    require 'ruby-prof'
    prof_results = RubyProf.profile do
      yield
    end
    File.open("benchmark/#{str}.html", "w") do |f|
      RubyProf::GraphHtmlPrinter.new(prof_results).print(f, :min_percent => 5)
    end
  else
    yield
  end
end

def benchmark(str, n, db, coll_name, object)
  coll = db.collection(coll_name)
  profile(str) do
    GC.start
    tms = Benchmark.measure do
      n.times { |i| yield(coll, object.is_a?(Hash) ? object.dup : object, i) }
    end
    report(str, tms.real)
  end
end

host = ENV['MONGO_RUBY_DRIVER_HOST'] || 'localhost'
port = ENV['MONGO_RUBY_DRIVER_PORT'] || Connection::DEFAULT_PORT

connection = Connection.new(host, port)
connection.drop_database("benchmark")
@db = connection.db('benchmark')

def benchmark_insert(desc, data_desc, data)
  benchmark(desc, PER_TRIAL, @db, data_desc, data) do |coll, object, i| 
    object['x'] = i
    coll.insert(object)
  end
end

def benchmark_insert_index(desc, data_desc, data)
  benchmark(desc, PER_TRIAL, @db, data_desc, data) do |coll, object, i| 
    coll.create_index('x')
    object['x'] = i
    coll.insert(object)
  end
end

benchmark_insert('insert (small, no index)', 'small_none', SMALL)  
benchmark_insert('insert (medium, no index)', 'medium_none', MEDIUM)
benchmark_insert('insert (large, no index)', 'large_none', LARGE)

benchmark_insert_index('insert (small, index)', 'small_indexed', SMALL)  
benchmark_insert_index('insert (medium, index)', 'medium_indexed', MEDIUM)
benchmark_insert_index('insert (large, index)', 'large_indexed', LARGE)
 
def benchmark_insert_batch(desc, data_desc, data)
  benchmark(desc, PER_TRIAL / BATCH_SIZE, @db, data_desc, data) do |coll, object, i| 
    object['x'] = i
    coll.insert([object] * BATCH_SIZE)
  end
end

benchmark_insert_batch('insert batch (small, index)', 'small_bulk', SMALL)  
benchmark_insert_batch('insert batch (medium, index)', 'medium_bulk', MEDIUM)
benchmark_insert_batch('insert batch (large, index)', 'large_bulk', LARGE)

def benchmark_find_one(desc, data_desc, data)
  benchmark(desc, PER_TRIAL, @db, data_desc, data) do |coll, object, i| 
    coll.find_one('x' => data_desc)
  end
end

benchmark_find_one('find_one (small, no index)', 'small_none', PER_TRIAL / 2)
benchmark_find_one('find_one (medium, no index)', 'medium_none', PER_TRIAL / 2)
benchmark_find_one('find_one (large, no index)', 'large_none', PER_TRIAL / 2)
benchmark_find_one('find_one (small, indexed)', 'small_indexed', PER_TRIAL / 2)
benchmark_find_one('find_one (medium, indexed)', 'medium_indexed', PER_TRIAL / 2)
benchmark_find_one('find_one (large, indexed)', 'large_indexed', PER_TRIAL / 2)

def benchmark_find_all(desc, data_desc, data)
  benchmark(desc, PER_TRIAL, @db, data_desc, data) do |coll, object, i| 
    coll.find('x' => data_desc).each {}
  end
end

benchmark_find_all('find_all (small, no index)', 'small_none', PER_TRIAL / 2)
benchmark_find_all('find_all (medium, no index)', 'medium_none', PER_TRIAL / 2)
benchmark_find_all('find_all (large, no index)', 'large_none', PER_TRIAL / 2)
benchmark_find_all('find_all (small, indexed)', 'small_indexed', PER_TRIAL / 2)
benchmark_find_all('find_all (medium, indexed)', 'medium_indexed', PER_TRIAL / 2)
benchmark_find_all('find_all (large, indexed)', 'large_indexed', PER_TRIAL / 2)

def benchmark_find_range(desc, data_desc, data)
  benchmark(desc, PER_TRIAL, @db, data_desc, data) do |coll, object, i| 
    coll.find('x' => data_desc).each {}
  end
end

benchmark_find_range('find_range (small, indexed)', 'small_indexed', 
                     {"$gt" => PER_TRIAL / 2, "$lt" => PER_TRIAL / 2 + BATCH_SIZE})
benchmark_find_range('find_range (medium, indexed)', 'medium_indexed', 
                     {"$gt" => PER_TRIAL / 2, "$lt" => PER_TRIAL / 2 + BATCH_SIZE})
benchmark_find_range('find_range (large, indexed)', 'large_indexed', 
                     {"$gt" => PER_TRIAL / 2, "$lt" => PER_TRIAL / 2 + BATCH_SIZE})