import time

from aiida.orm import Group
from aiida.orm.data.cif import CifData

group_name = "test_group"

# If the group is not deleted from a previous execution, delete it
qb = QueryBuilder()
qb.append(Group, filters={'name': {'==': group_name}})
if qb.count() > 0:
	qb.all()[0][0].delete()

print "Creating a temp group"
group_name = "test_group"
group, created = Group.get_or_create(name=group_name)

qb = QueryBuilder()
qb.append(CifData)
print "Number of available Cif data", qb.count()

print "Adding nodes one by one"
start = time.time()
count = 0
prev_end = start
for cif in qb.all()[:1000]:
	group.add_nodes(cif)
	count = count + 1
	if count%100 == 0:
		curr_end = time.time()
		print "Counter: {}, Time elspsed in ms: {}, Throughput (time per 100 nodes): {}".format(count, str(curr_end-start), str(curr_end-prev_end))
		prev_end = curr_end
print "Time elapsed in ms", str(time.time() - start)

print "Deleting the created group"
group.delete()

print "Creating a temp group"
group_name = "test_group"
group, created = Group.get_or_create(name=group_name)

print "Adding nodes in batches of 100"
start = time.time()
count = 0
prev_end = start
res = [_[0] for _ in qb.all()[:1000]]
chunks = [res[x:x+100] for x in xrange(0, len(res), 100)]
for cif_chunk in chunks:
	group.add_nodes(cif_chunk)
	count = count + 1
	curr_end = time.time()
	print "Counter: {}, Time elspsed in ms: {}, Throughput (time per 100 nodes): {}, chunk size: {}".format(count, str(curr_end-start), str(curr_end-prev_end), len(cif_chunk))
	prev_end = curr_end
print "Time elapsed in ms", str(time.time() - start)