using Distributions
using StatsPlots
using Turing
using Logging
using LaTeXStrings

default(labels=false)

Code 2.1

ways = [0, 3, 8, 9, 0]
ways = ways ./ sum(ways)
println(ways)

[0.0, 0.15, 0.4, 0.45, 0.0]

Code 2.2

b = Binomial(9, 0.5)
pdf(b, 6)

0.16406250000000056

Code 2.3

# size of the grid
size = 20

# grid and prior
p_grid = range(0, 1; length=size)
prior = repeat([1.0], size)

# compute likelihood at each value in grid
likelyhood = [pdf(Binomial(9, p), 6) for p in p_grid]

# compute product of likelihood and prior
unstd_posterior = likelyhood .* prior

# standardize the posterior, so it sums to 1
posterior = unstd_posterior / sum(unstd_posterior);

Code 2.4

plot(p_grid, posterior; 
    xlabel="probability of water", 
    ylabel="posterior probability",
    title="$size points",
    markershape=:circle
)

Code 2.5

size = 20
p_grid = range(0, 1; length=size)

# prior is different - 0 if p < 0.5, 1 if >= 0.5
prior = convert(Vector{AbstractFloat}, p_grid .>= 0.5)

# another prior to try (uncomment the line below)
# prior = exp.(-5*abs.(p_grid .- 0.5))

# the rest is the same
likelyhood = [pdf(Binomial(9, p), 6) for p in p_grid]
unstd_posterior = likelyhood .* prior
posterior = unstd_posterior / sum(unstd_posterior);

plot(p_grid, posterior; 
    xlabel="probability of water", 
    ylabel="posterior probability",
    title="$size points",
    markershape=:circle
)

Code 2.6

@model function water_land(W, L)
    p ~ Uniform(0, 1)
    W ~ Binomial(W + L, p)
end

Logging.disable_logging(Logging.Warn)
chain = sample(water_land(6, 3), NUTS(0.65), 1000)
display(chain)

Chains MCMC chain (1000×13×1 Array{Float64, 3}):

Iterations        = 501:1:1500
Number of chains  = 1
Samples per chain = 1000
Wall duration     = 6.05 seconds
Compute duration  = 6.05 seconds
parameters        = p
internals         = lp, n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size

Summary Statistics
  parameters      mean       std   naive_se      mcse        ess      rhat   e ⋯
      Symbol   Float64   Float64    Float64   Float64    Float64   Float64     ⋯

           p    0.6370    0.1436     0.0045    0.0071   407.9428    1.0023     ⋯
                                                                1 column omitted

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

           p    0.3377    0.5444    0.6397    0.7456    0.8901

Code 2.7

# analytical calculation
W = 6
L = 3
x = range(0, 1; length=101)

b = Beta(W+1, L+1)
plot(x, pdf.(b, x); label = L"\mathcal{Beta}")

# quadratic approximation
b = Normal(0.6374, 0.1414)
plot!(x, pdf.(b, x); style=:dash, label=L"\mathcal{N}(0.64, 0.14)")

Code 2.8

n_samples = 1000
p = Vector{Float64}(undef, n_samples)
p[1] = 0.5
W, L = 6, 3

for i ∈ 2:n_samples
    p_old = p[i-1]
    p_new = rand(Normal(p_old, 0.1))
    if p_new < 0
        p_new = abs(p_new)
    elseif p_new > 1
        p_new = 2-p_new
    end

    q0 = pdf(Binomial(W+L, p_old), W)
    q1 = pdf(Binomial(W+L, p_new), W)
    u = rand(Uniform())
    p[i] = (u < q1 / q0) ? p_new : p_old
end

Code 2.9

density(p; label = "MCMC")
b = Beta(W+1, L+1)
plot!(x, pdf.(b, x); label = L"\mathcal{Beta}", style=:dash)