### Python code to produce plots and experiments used in statistics lecture.
### The entry point is the lecture() function, which is called at the end
### of this file.

import matplotlib.pyplot as plt
import numpy as np
import math
import random

mike_rolls = [5, 6, 6, 5, 6, 6]
student_rolls = [2, 2, 4, 5, 1, 2]
student_name = "Brett"

def normalize(list):
    """Return a new list in which the elements are scaled to sum to 1.
    The input list is a list of numbers.  It is not modified by this routine.
    """
    list_sum = float(sum(list))
    return [elt/list_sum for elt in list]

assert normalize([3, 1, 0, 3, 2, 1]) == ([.3, .1, 0, .3, .2, .1])

def rolls_to_hist(rolls):
    """Convert a list of single-die rolls to a list of counts.
    The input is a list, all of whose numbers are between 1 and 6 inclusive.
    In the result list, result[i] is the number of times i appeared in the input list.
    In the result list, result[0] == 0.
    """
    result = [0] * 7
    for roll in rolls:
        result[roll] = result[roll] + 1
    return result

assert rolls_to_hist([1, 5, 4, 6, 2, 1, 1, 3, 2, 5]) == [0, 3, 2, 1, 1, 2, 1]

def plot_die_rolls(rolls, label):
    """Plot a die-roll normalized histogram.
    The input is a list of rolls, where each element is the sum of visible pips."""
    hist = rolls_to_hist(rolls)
    plot_die_hist(hist, label)

def plot_die_hist(hist, label):
    """Plot a die-roll normalized histogram.
    The input is a histogram indicating, for each number of pips,
    its normalized frequency.
    """
    plt.plot(normalize(hist), label=label)
    plt.plot(normalize(hist), 'o')
    plt.axis([1,len(hist)-1,0,0.7])
    plt.xlabel("Pips")
    plt.ylabel("Frequency")
    plt.legend()

def show_two_hist(rolls1, name1, rolls2, name2):
    plot_die_rolls(rolls1, name1)
    plot_die_rolls(rolls2, name2)
    plt.show()

def plot_one_hist(rolls):
    numdice = 1
    plt.hist(rolls, bins=np.linspace(numdice-.5, 6*numdice+.5, num=5*numdice+2))
    plt.axis([0.5,6.5,0,max(rolls_to_hist(rolls))+1])
    plt.xlabel("Pips")
    plt.ylabel("Frequency")

def show_one_hist(rolls):
    plot_one_hist(rolls)
    plt.show()

def show_one_hist_and_line(rolls):
    plot_one_hist(rolls)
    hist = rolls_to_hist(rolls)
    plt.plot(hist)
    plt.plot(hist, 'o')
    plt.axis([0.5,6.5,0,max(hist)+1])
    plt.show()

def die_roll():
    return random.randrange(6) + 1

def dice_roll(numdice):
    result = 0
    for i in range(numdice):
        result = result + die_roll()
    return result

def dice_rolls(numdice, numrolls):
    """Returns a list of numrolls numbers, each the sum of rolling numdice dice."""
    result = []
    for i in range(numrolls):
        result.append(dice_roll(numdice))
    return result

random10_rolls = dice_rolls(1, 10)

def show_notation():
    show_one_hist(dice_rolls(1, 10))
    show_one_hist(dice_rolls(1, 10))
    random10_rolls = dice_rolls(1, 10)
    show_one_hist(random10_rolls)
    show_one_hist_and_line(random10_rolls)

def plot_random(numdice, numrolls):
    rolls = dice_rolls(numdice, numrolls)
    plot_die_rolls(rolls, "Random (" + str(numrolls) + ")")

def show_random_1000():
    plot_random(1, 1000)
    plt.show()

def show_random_many():
    plot_random(1, 10)
    plot_random(1, 100)
    plot_random(1, 1000)
    plot_random(1, 10000)
    plt.show()

def show_random():
    plot_die_rolls([1, 2, 3, 4, 5, 6], "Ideal")
    plt.show()
    plot_die_rolls([1, 2, 3, 4, 5, 6], "Ideal")
    show_random_1000()
    show_random_many()

zurich_first_roll = [25, 28, 45, 49, 106, 136]
zurich_second_roll = [5, 4, 6, 10, 28, 57]

def show_zurich_students():
    plot_die_hist(zurich_first_roll, "First roll (389 students)")
    plt.xlabel("Payoff (CHF)")
    plt.axis([0,5,0,.6])
    plt.title("Zurich students")
    plt.show()
    plot_die_hist(zurich_first_roll, "First roll (389 students)")
    plot_die_hist(zurich_second_roll, "Second roll (110 students)")
    plt.xlabel("Payoff (CHF)")
    plt.axis([0,5,0,.6])
    plt.title("Zurich students")
    plt.show()
    
def dice_rolls(numdice, numrolls):
    dice_hist = [0] * (numdice * 6 + 1)
    result = []
    for i in range(numrolls):
        result.append(dice_roll(numdice))
    return result

def make_dice_histogram(therolls=None, therollname=None, numdice=None, numrolls=100000):
    if numdice == None:
        numdice = len(therolls)
    plt.hist(dice_rolls(numdice, numrolls), bins=np.linspace(numdice-.5, 6*numdice+.5, num=5*numdice+2))
    plt.xlabel("Pips")
    plt.ylabel("Frequency")
    plt.title(str(numdice) + " dice")

def annotate_dice_histogram(therolls, therollname, numrolls=100000):
    thepips = sum(therolls)
    plt.annotate(therollname + "'s rolls", xy=(thepips, numrolls/100), xytext=(thepips, numrolls/30),
                 arrowprops=dict(facecolor='black', shrink=0.2),
                )        

def show_annotated_dice_histogram():
    # show_dice_histogram(numdice=3)
    # show_dice_histogram(numdice=6)
    make_dice_histogram(therolls=mike_rolls)
    annotate_dice_histogram(therolls=mike_rolls, therollname="Mike")
    annotate_dice_histogram(therolls=student_rolls, therollname=student_name)
    plt.show()

def dice_likelihood(therolls, numdice, numrolls):
    thesum = sum(therolls)
    rolls = dice_rolls(numdice, numrolls)
    as_unlikely = 0
    for roll in rolls:
        if thesum <= roll:
            as_unlikely = as_unlikely + 1
    return float(as_unlikely) / numrolls

def print_dice_likelihood(therolls, numrolls, name):
    print "Likelihood of " + name + "'s result or better:", dice_likelihood(therolls, len(therolls), numrolls)

def print_two_dice_likelihood():
    print_dice_likelihood(mike_rolls, 100000, "Mike")
    print_dice_likelihood(student_rolls, 100000, student_name)


def lecture():
    """The main entry point, which shows all the lecture's plots in turn."""
    show_notation()

    show_two_hist(mike_rolls, "Mike", student_rolls, student_name)

    show_random()

    show_zurich_students()

    show_annotated_dice_histogram()

    print_two_dice_likelihood()


lecture()