import datetime
import logging
import random
import sys

import riak.benchmark as benchmark

from multiprocessing import cpu_count
from riak import RiakClient

# logger = logging.getLogger()
# logger.level = logging.DEBUG
# logger.addHandler(logging.StreamHandler(sys.stdout))

epoch = datetime.datetime.utcfromtimestamp(0)
onesec = datetime.timedelta(0, 1)

rowcount = 32768
batchsz = 32
if rowcount % batchsz != 0:
    raise AssertionError('rowcount must be divisible by batchsz')

weather = ['typhoon', 'hurricane', 'rain', 'wind', 'snow']
rows = []
keys = []
for i in range(rowcount):
    ts = datetime.datetime(2016, 1, 1, 12, 0, 0) + \
        datetime.timedelta(seconds=i)
    family_idx = i % 4
    series_idx = i % 4
    family = 'hash{:d}'.format(family_idx)
    series = 'user{:d}'.format(series_idx)
    w = weather[i % len(weather)]
    temp = (i % 100) + random.random()
    row = [family, series, ts, w, temp]
    key = [family, series, ts]
    rows.append(row)
    keys.append(key)

print("Benchmarking timeseries:")
print("      CPUs: {0}".format(cpu_count()))
print("      Rows: {0}".format(len(rows)))
print()

tbl = 'GeoCheckin'
h = 'riak-test'
n = [
    {'host': h, 'pb_port': 10017},
    {'host': h, 'pb_port': 10027},
    {'host': h, 'pb_port': 10037},
    {'host': h, 'pb_port': 10047}
]
client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': False})
table = client.table(tbl)

with benchmark.measure() as b:
    with b.report('populate'):
        for i in range(0, rowcount, batchsz):
            x = i
            y = i + batchsz
            r = rows[x:y]
            ts_obj = table.new(r)
            result = ts_obj.store()
            if result is not True:
                raise AssertionError("expected success")
    with b.report('get'):
        for k in keys:
            ts_obj = client.ts_get(tbl, k)
            if ts_obj is None:
                raise AssertionError("expected obj")
            if len(ts_obj.rows) != 1:
                raise AssertionError("expected one row, got: %d" % len(tsobj.rows))
            row = ts_obj.rows[0]
            if len(row) != 5:
                raise AssertionError("expected row to have five items")

@@ -275,86 +275,159 @@
 ))%0A%0A
-epoch = datetime.datetime.utcfromtimestamp(0)%0Aonesec = datetime.timedelta(0, 1
+# batch sizes 8, 16, 32, 64, 128, 256%0Aif len(sys.argv) != 3:%0A    raise AssertionError('first arg is batch size, second arg is true / false for use_ttb'
 )%0A%0Ar
@@ -452,18 +452,32 @@
 tchsz = 
-32
+int(sys.argv%5B1%5D)
 %0Aif rowc
@@ -562,16 +562,138 @@
 atchsz')
+%0Ause_ttb =  sys.argv%5B2%5D.lower() == 'true'%0A%0Aepoch = datetime.datetime.utcfromtimestamp(0)%0Aonesec = datetime.timedelta(0, 1)
 %0A%0Aweathe
@@ -759,18 +759,8 @@
  %5B%5D%0A
-keys = %5B%5D%0A
 for 
@@ -889,25 +889,31 @@
 y_idx = i %25 
-4
+batchsz
 %0A    series_
@@ -922,17 +922,23 @@
 x = i %25 
-4
+batchsz
 %0A    fam
@@ -1184,29 +1184,8 @@
 row)
-%0A    keys.append(key)
 %0A%0Apr
@@ -1229,20 +1229,99 @@
 (%22  
+ Use TTB: %7B%7D%22.format(use_ttb))%0Aprint(%22Batch Size: %7B%7D%22.format(batchsz))%0Aprint(%22
     
+  
 CPUs: %7B
-0
 %7D%22.f
@@ -1360,17 +1360,16 @@
  Rows: %7B
-0
 %7D%22.forma
@@ -1570,16 +1570,51 @@
 : 10047%7D
+,%0A    %7B'host': h, 'pb_port': 10057%7D
 %0A%5D%0Aclien
@@ -1687,13 +1687,15 @@
 b': 
-False
+use_ttb
 %7D)%0At
@@ -1751,16 +1751,44 @@
 ) as b:%0A
+    for i in (1, 2, 3):%0A    
     with
@@ -1806,20 +1806,31 @@
 populate
-'):%0A
+-%25d' %25 i):%0A    
         
@@ -1883,14 +1883,22 @@
     
+    
 x = i%0A
+    
     
@@ -1917,24 +1917,28 @@
 i + batchsz%0A
+    
             
@@ -1955,32 +1955,36 @@
 :y%5D%0A            
+    
 ts_obj = table.n
@@ -1989,16 +1989,20 @@
 .new(r)%0A
+    
         
@@ -2041,16 +2041,20 @@
         
+  
+  
 if resul
@@ -2076,32 +2076,36 @@
                 
+    
 raise AssertionE
@@ -2129,441 +2129,8 @@
 cces
-s%22)%0A    with b.report('get'):%0A        for k in keys:%0A            ts_obj = client.ts_get(tbl, k)%0A            if ts_obj is None:%0A                raise AssertionError(%22expected obj%22)%0A            if len(ts_obj.rows) != 1:%0A                raise AssertionError(%22expected one row, got: %25d%22 %25 len(tsobj.rows))%0A            row = ts_obj.rows%5B0%5D%0A            if len(row) != 5:%0A                raise AssertionError(%22expected row to have five item
 s%22)%0A

Add args to benchmark program.