import random
import time

# import plotting functionality
from matplotlib import pyplot

def generate_set(set_size):
    data = set()
        
    while len(data) < set_size:
        data.add(random.randint(0, set_size*10))
    
    return data

def generate_list(set_size):
    data = []
    
    for i in range(set_size):
        data.append(random.randint(0, set_size*10));
    
    return data

def query_data(data, num_queries, largest):
    # we'll count them for fun, but won't do anything with the count
    count = 0 
    
    begin = time.time();
    
    for i in range(num_queries):
        # just check to see if a random number is in there
        random_num = random.randint(0, largest)
        
        if random_num in data:
            count += 1
            
    return time.time() - begin
    
def speed_data_old(num_queries, set_size_min, set_size_max, step_size):
    print "size\tlist\tset"
    
    for set_size in range(set_size_min, set_size_max, step_size):
        # generate random list
        list_data = generate_list(set_size)
    
        # generate random set
        set_data = generate_set(set_size)
    
        # time the list querying
        list_elapsed = query_data(list_data, num_queries, set_size*10)

        # time the list querying
        set_elapsed = query_data(set_data, num_queries, set_size*10)
        
        print "%d\t%f\t%f" % (set_size, list_elapsed, set_elapsed)
        
def speed_data(num_queries, set_size_min, set_size_max, step_size):
    set_sizes = []
    list_times = []
    set_times = []
    
    for set_size in range(set_size_min, set_size_max, step_size):
        # generate random list
        list_data = generate_list(set_size)
    
        # generate random set
        set_data = generate_set(set_size)
        
        # time the list querying
        list_elapsed = query_data(list_data, num_queries, set_size*10)

        # time the list querying
        set_elapsed = query_data(set_data, num_queries, set_size*10)
        
        set_sizes.append(set_size)
        list_times.append(list_elapsed)
        set_times.append(set_elapsed)
        
    return (set_sizes, list_times, set_times)

def print_speed_data(set_sizes, list_times, set_times):
    print "set_size\tlist\t\tset"
    
    for i in range(len(query_times)):
        print "%d\t%f\t%f" % (set_sizes[i], list_times[i], set_times[i])

def plot_speed_data(set_sizes, list_times, set_times):
    # we have the data, now plot it
    pyplot.plot(set_sizes, list_times, "r", set_sizes, set_times, "b")
    pyplot.xlabel("Collection size")
    pyplot.ylabel("Time elapsed (seconds)")
    pyplot.title("Timing 'in' for lists vs. sets")
    pyplot.legend(('lists', 'sets'), 'upper left')

    pyplot.show()
    
def run_experiment():
    query_times, list_times, set_times = speed_data(1000, 1000, 10000, 500)
    plot_speed_data(query_times, list_times, set_times)