# File: loopTimer.py
# Purpose: Compare the runtime of different loops
# Author: Tanya Amert

import matplotlib.pyplot as plt
import time

def runTrials(n, numTrials):
    """
    Times looping through a list of size n over numTrials trials.

    Returns: a dict mapping approach (str) to average runtime (float, in seconds)
    """
    sortedList = list(range(n))

    time_single = 0
    time_nested = 0
    time_nested_half = 0
    for i in range(numTrials):

        # A single loop
        t = time.time()
        for i in range(len(sortedList)):
            pass
        time_single += (time.time() - t)

        # A nested loop
        t = time.time()
        for i in range(len(sortedList)):
            for j in range(len(sortedList)):
                pass
        time_nested += (time.time() - t)

        # A nested loop, but just half
        t = time.time()
        for i in range(len(sortedList)):
            for j in range(i): # based on outer loop
                pass
        time_nested_half += (time.time() - t)

    return {"single": time_single,
            "nested": time_nested,
            "nested-half": time_nested_half}

def timeLoops():
    """
    Loops over lists of different sizes and plots the runtimes.
    """
    # Parameters
    nvals = [10, 20, 50, 80, 100, 200, 500, 800, 1000, 2000, 5000]
    num_trials = 1000

    # Prep lists to store results
    all_times_single = []
    all_times_nested = []
    all_times_nested_half = []

    # Try all different values of n
    for n in nvals:
        # Get data for this value of n
        print("Working on n:", n)  
        trial_results = runTrials(n, num_trials)

        # Append the average time to the appropriate lists
        all_times_single.append(trial_results["single"] / num_trials * 1000)
        all_times_nested.append(trial_results["nested"] / num_trials * 1000)
        all_times_nested_half.append(trial_results["nested-half"] / num_trials * 1000)

    # Make the plot
    plt.figure()
    plt.plot(nvals, all_times_nested, label="Nested (ms)", marker='x')
    plt.plot(nvals, all_times_nested_half, label="Nested half (ms)", marker='^')
    plt.plot(nvals, all_times_single, label="Single (ms)", marker='o')

    # Spiff up our plot
    plt.title("Time to loop as n gets really big")
    plt.xlabel("Number of elements in the list")
    plt.ylabel("Loop time (ms)")
    plt.legend()
    plt.show()

def main():
    timeLoops()

if __name__ == "__main__":
    main()