################### FLIPPING UTILS ###################
from random import random
import numpy as np

def flip(q):
    #returns 1 w.p q and 0 o.w
    return int(random()<q)

def sequence(n, q):
    #returns a sequence of n coin tosses
    #HEADS(1) w.p q and TAILS(0) o.w
    return np.array([flip(q) for _ in range(n)])

################### MLE ###################
from functools import partial as curry

def mle_uniform(X):
    #returns MLE on X with uniform prior
    return X.sum()/len(X)

def mle_two_coins(e, X):
    #retuns MLE on X with prior p(1/2 - e) = p(1/2 + e) = 1/2
    return 0.5 + e if X.sum() >= len(X)/2.0 else 0.5-e

################### ERROR ###################
def zero_one_loss(q, q_hat):
    #returns 1 if q \neq q_hat and 0 o.w
    return q != q_hat

def absolute_loss(q, q_hat):
    #returns |q- q_hat|
    return np.absolute(q - q_hat)

def check_error(estimator, loss_function, n, q, trials):
    #sample a sequence with bias q
    #use estimator to estimate q
    #returns the average loss incurred in all trials
    loss = 0.0
    for _ in range(int(trials)):
        X = sequence(int(n), q)
        q_hat = estimator(X)
        loss = loss + loss_function(q, q_hat)
    return loss/trials
    

################### PLOT ###################
import matplotlib.pyplot as plt

def plot(x, y):
    plt.xlabel('Number of Coin Tosses')
    plt.ylabel('Error')
    plt.plot(x, y)
    plt.show()

n_list = np.linspace(100, 4000, 10)
loss_list = [check_error(mle_uniform, absolute_loss, n, 0.1, 100) for n in n_list]
plot(n_list, loss_list)

n_list = np.linspace(100, 300, 10)
loss_list = [check_error(curry(mle_two_coins, 0.1), zero_one_loss, n, 0.4, 10000) for n in n_list]
plt.plot(n_list, loss_list)
plt.show()