# Import the plotting package for Julia
using PyPlot

# Declaring some global variables for multiple uses
N = 1000
sequence = [2, 4, 10, 50, 100, 500, 10000, 100000]
Qs = [0.4, 0.1, 0.25, 0.5]
ks = [2, 10, 50, 100, 500, 10000, 100000];

# Returns the number of heads in n flips
# rand() generates a random number between 0 and 1, that number is then rounded to either 0 or 1. 
# By summing them up, we essentially are counting the number of heads
flip_coin(n) = sum(round(rand(n)))

heads = flip_coin(N)
figure(figsize=(5, 3))
bar([0, 1], [heads, N-heads], width=0.2, align="center")
xticks([0, 1], ["Heads", "Tails"])
title("1000 coin tosses");

# Simulates a flip of n coins m times
flip_runs(n, m) = [flip_coin(n) for i = 1:m]

figure(figsize=(5, 3))
plt.hist(flip_runs(N, 1000))
title("Heads frequency in 1000 runs of 1000 coin tosses");

for i = sequence
    figure(figsize=(5, 3))
    plt.hist(flip_runs(i, 1000))
    title("Heads frequency in 1000 runs of $i coin tosses");
end

scaled_run(n, m) = [flip_coin(n) - n/2 for i = 1:m]

for i = sequence
    figure(figsize=(5, 3))
    plt.hist(scaled_run(i, 1000));
end

for i = sequence
    figure(figsize=(5, 3))
    plt.hist(scaled_run(i, 1000));
    xlim(-500, 500)
end

for i = sequence
    figure(figsize=(5, 3))
    plt.hist(flip_runs(i, 1000));
    xlim(0, i)
end

rand_one() = rand() < 0.5 ? -1 : 1

# An n-step random walk starts at 0 by default
function path(n, walk=[0])
    for i = 1:n
        push!(walk, rand_one() + walk[end])
    end 
    return walk
end

figure(figsize=(5, 3))
title("20 random walks of 1000 coin flips")
for _ = 1:20
    plot(path(1000))
end

figure(figsize=(5, 3))
title("100 normalized random walks of 1000 coin flips")
for _ = 1:100
    plot([k/n for (n, k) = enumerate(path(1000))])
end

identity(x) = x

# runs and tosses are the number of runs and tosses respectively
# q is the threshold which we're interested in finding the fraction of heads smaller than
# func is the function we want to apply to the vertical axis values
# For this part, that function should do nothing (identity). For next part, it's the log function
function multi_toss_runs(runs, tosses, q, func)
    runs_table = Dict()
    for k = 1:tosses
        runs_table[k] = Float64[]
    end
    for _ = 1:runs
        heads = 0
        for k = 1:tosses
            heads += round(rand())[1]
            push!(runs_table[k], heads)
        end
    end
    
    # Computes the frequency of the fraction of heads is less than q
    # Code explanation: for every run whose fraction of heads is less than q, give that run a 1, otherwise 0
    # Then sum up all the values to get the number of runs that satisfies this requirement, and divides
    # it by the total number of runs. Repeat for each of the k tosses
    return [func(sum([heads/k <= q ? 1 : 0 for heads = runs_table[k]]) / runs) for k = 1:tosses]
end

figure(figsize=(5, 3))
for q = Qs
    plot(multi_toss_runs(10000, 1000, q, identity));
end

figure(figsize=(5, 3))
for q = Qs
    plot(multi_toss_runs(10000, 1000, q, log));
end

function q_curve(m, k)
    res = [flip_coin(k) / k for _ = 1:m]
    sort!(res)
    first_quartile = res[m * 0.25]
    second_quartile = res[m * 0.5]
    third_quartile = res[m * 0.75]
    return res, first_quartile, second_quartile, third_quartile
end

figure(figsize=(5, 3))
title("Ratio of heads as a function of q")
curve, _, _, _ = q_curve(10000, 1000) # don't care about the quartile values here yet
plot(curve, 1:10000);

firsts, seconds, thirds = Float64[], Float64[], Float64[]

figure(figsize=(5, 3))
title("Ratio of heads in k tosses as a function of q")
for k = ks
    curve, first_quartile, second_quartile, third_quartile = q_curve(10000, k)
    plot(curve, 1:10000)
    
    # Save the quartile values
    push!(firsts, first_quartile)
    push!(seconds, second_quartile)
    push!(thirds, third_quartile)
end

# Ignore k = 2, so we start indexing from 2. Recall that Julia is 1-indexed
function plot_quartiles()
    plot(ks[2:], firsts[2:]);
    plot(ks[2:], seconds[2:]);
    plot(ks[2:], thirds[2:]);
end

figure(figsize=(5, 3))
title("q quartile markers as a function of k")
plot_quartiles();

# Computes the gap between the 0.75 marker and the 0.25 marker
gap = thirds - firsts

function plot_quartiles(plot_function)
    plot_function(ks[2:], gap[2:]);
end

figure(figsize=(5, 3))
title("Linear-linear plot")
plot_quartiles(plot);

figure(figsize=(5, 3))
title("Log-linear plot")
plot_quartiles(semilogx);

figure(figsize=(5, 3))
title("Linear-log plot")
plot_quartiles(semilogy);

figure(figsize=(5, 3))
title("Log-log plot")
plot_quartiles(loglog);

function q_curve_norm(m, k)
    res = [(flip_coin(k)/k - 0.5) * sqrt(k) + 0.5 for _ = 1:m]
    sort!(res)
    return res
end

figure(figsize=(5, 3))
title("Scaled ratio of heads in k tosses as a function of q")
for k = ks
    plot(q_curve_norm(10000, k), 1:10000)
end