using Pkg; Pkg.activate("../."); Pkg.instantiate();
using IJulia; try IJulia.clear_output(); catch _ end


using Distributions, StatsPlots, SpecialFunctions

# computes log10 of Gamma function
function log10gamma(num)
    num = convert(BigInt, num)
    return log10(gamma(num))
end


μ  = 0.4; # specify model parameter
n_tosses = 200                   # specify number of coin tosses
samples = rand(n_tosses) .<= μ   # Flip 200 coins

function handle_coin_toss(prior :: Beta, observation :: Bool)
    posterior = Beta(prior.α + observation, prior.β + (1 - observation))
    return posterior
end

function log_evidence_prior(prior :: Beta, N :: Int64, n :: Int64)
    log_evidence = log10gamma(prior.α + prior.β) - log10gamma(prior.α) - log10gamma(prior.β) + log10gamma(n+prior.α) + log10gamma((N-n)+prior.β) - log10gamma(N+prior.α+prior.β)
    return log_evidence
end

priors = [Beta(100., 500.), Beta(5., 1.), Beta(8, 13)]  #specify prior distributions (you can add additional beta distributions here) model 3 is the "best" prior. Can you see why?                              #save prior distributions to compute log-evidence

n_models = length(priors)

posterior_distributions = [[d] for d in priors]     #save a sequence of posterior distributions for every prior, starting with the prior itself
log_evidences = [[] for _ in priors]                    #maintain a vector of log evidences to plot later

for (N, sample) in enumerate(samples)                                               #for every sample we want to update our posterior
    for (i, prior) in enumerate(priors)                                             #at every sample we want to update all distributions
        posterior = handle_coin_toss(prior, sample)                                 #do bayesian updating
        push!(posterior_distributions[i], posterior)                                #add posterior to vector of posterior distributions
        log_evidence = log_evidence_prior(posterior_distributions[i][N], N, sum(samples[1:N]))    #compute log evidence and add to vector
        push!(log_evidences[i], log_evidence)
        priors[i] = posterior                                                       #the prior for the next sample is the posterior from the current sample
    end
end



#animate posterior distributions over time in a gif

anim = @animate for i in 1:n_tosses
    p = plot(title=string("n = ", i))
    for j in 1:n_models
        plot!(posterior_distributions[j][i+1], xlims = (0, 1), fill=(0, .2,), label=string("Posterior ", j), linewidth=2, ylims=(0,28), xlabel="μ")
    end
end

gif(anim, "anim_bay_ct.gif")

using LaTeXStrings

evidences = map(model -> exp.(model), log_evidences)

anim = @animate for i in 1:n_tosses
    p = plot(title=string(L"\frac{p_i(\mathbf{x}_{1:n})}{\sum_i p_i(\mathbf{x}_{1:n})}","   n = ", i), ylims=(0, 1))
    total = sum([evidences[j][i] for j in 1:n_models])
    bar!([(evidences[j][i] / total) for j in 1:n_models], group=["Model $i" for i in 1:n_models])
end

gif(anim, "anim_bay_ct_log_evidence.gif")

using Distributions, StatsPlots, Plots.PlotMeasures, LaTeXStrings

function kullback_leibler(q :: Normal, p :: Normal)                                 #Calculates the KL Divergence between two gaussians (see https://stats.stackexchange.com/questions/7440/kl-divergence-between-two-univariate-gaussians for calculations)
    return log((p.σ / q.σ)) + ((q.σ)^2 + (p.μ - q.μ)^2) / (2*p.σ^2) - (1. / 2.)
end

μ_p = 0             #statistics of distributions we'll keep constant (we'll vary the mean of q). Feel free to change these and see what happens
σ_p = 1
σ_q = 1

p = Normal(μ_p, σ_p)

anim = @animate for i in 1:100
    μ_seq = [(j / 10.) - 5. + μ_p for j in 1:i]                                                                                                         #compute the sequence of means tested so far (to compute sequence of KL divergences)
    kl = [kullback_leibler(Normal(μ, σ_q), p) for μ in μ_seq]                                                                                           #compute KL divergence data
    viz = plot(right_margin=8mm, title=string(L"D_{KL}(Q || P) = ", round(100 * kl[i]) / 100.), legend=:topleft)                                        #build plot and format title and margins
    μ_q = μ_seq[i]                                                                                                                                      #extract mean of current frame from mean sequence
    q = Normal(μ_q, σ_q)
    plot!(p, xlims = (μ_p - 8, μ_p + 8), fill=(0, .2,), label=string("P"), linewidth=2, ylims=(0,0.5))                                                  #plot p
    plot!(q, fill=(0, .2,), label=string("Q"), linewidth=2, ylims=(0,0.5))                                                                              #plot q
    plot!(twinx(), μ_seq, kl, xticks=:none, ylims=(0, maximum(kl) + 3), linewidth = 3,                                                                  #plot KL divergence data with different y-axis scale and different legend
        legend=:topright,xlims = (μ_p - 8, μ_p + 8), color="green", label=L"D_{KL}(Q || P)")
end
gif(anim, "anim_lat_kl.gif")

open("../../styles/aipstyle.html") do f display("text/html", read(f, String)) end