ENV["LINES"] = 200

using Distributions

using Plots
pyplot(fmt=:svg)
pyplotclf() = if backend() == Plots.PyPlotBackend(); PyPlot.clf(); end

using Base64
displayfile(mime, file; tag="img") = open(file) do f
    display("text/html", """<$tag src="data:$mime;base64,$(base64encode(f))"/>""")
end

using Printf
using Roots
using SpecialFunctions

rd(x, digits=3) = round(x; digits=digits)
safemul(x, y) = iszero(x) ? x : x*y
safexlogy(x, y) = iszero(x) ? x : x*log(y)

displayfile("image/png", "Bernoulli model.png")

# mean(Beta(a, b)) → a/(a+b)
mean(Beta(3, 7))

# lim_{b→0} mean(a, b) = a/a = 1
mean(Beta(1, 1e-8)), mean(Beta(2, 1e-8))

# the improper "Beta(0, 0)" distribution
plot(x -> 1/(x*(1-x)), 0, 1; ylim=(-0.1, 100), label="1/(x(1-x))", size=(400, 250))

# the improper "Beta(1, 0)" distribution
plot(x -> 1/(1-x), 0, 1; ylim=(-0.1, 100), label="1/(1-x)", size=(400, 250))

# Jeffreys' prior
plot(x -> pdf(Beta(1/2, 1/2), x), 0, 1; ylim=(-0.1, 10), label="Beta(1/2, 1/2)", size=(400, 250))

pred_MLE(n, k) = Bernoulli(k/n) # the predictive distribution of maximum likelihood estimation 
pred_Bayes(n, k; β=1.0, a=0.5, b=0.5) = Bernoulli((β*k+a)/(β*n+a+b)) # a=b=0.5 for the Jeffreys prior

n = 10
@show [pred_MLE(n, k).p   |> rd for k in 0:n]
@show [pred_Bayes(n, k).p |> rd for k in 0:n]
@show [pred_Bayes(n, k; a=1.0, b=1.0).p |> rd for k in 0:n]
@show [pred_Bayes(n, k; a=0.01, b=0.01).p |> rd for k in 0:n]
@show [pred_Bayes(n, k; β=0.0).p |> rd for k in 0:n]
@show [pred_Bayes(n, k; β=1e8).p |> rd for k in 0:n]
;

G(q, p) = -sum(safemul(pdf(q, x), logpdf(p, x)) for x in support(q)) # generalization error (or loss)
S(q) = G(q, q) # Shannon's information
KL(q, p) = G(q, p) - S(q) # Kullback-Leibler information, or prediction error of simulation of q by p

[(p, KL(Bernoulli(0.3), Bernoulli(p))) for p in 0:0.1:1]

make_samples(q::Bernoulli, n, L) = rand(Binomial(n, q.p), L)
ecdf(sample, x) = mean(sample .≤ x)

function plot_esitmation_errors(true_prob, n; L=10^6, bin=0:0.02:0.3, β=1.0, a=0.5, b=0.5)
    q = Bernoulli(true_prob)
    sample = make_samples(q, n, L)
    p_MLE   = pred_MLE.(n, sample)
    p_Bayes = pred_Bayes.(n, sample; β=β, a=a, b=b)
    KL_MLE   = KL.(q, p_MLE)
    KL_Bayes = KL.(q, p_Bayes)
    @show true_prob
    @show n
    
    P1 = plot(title="normalized histogram", titlefontsize=8)
    histogram!(KL_MLE;   bin=bin, label="KL_MLE",   alpha=0.5, norm=true)
    histogram!(KL_Bayes; bin=bin, label="KL_Bayes(β=$(β), a=$a, b=$b)", alpha=0.2, norm=true)
    x = range(0.01maximum(bin), maximum(bin), length=400)
    plot!(x, 2n*pdf.(Chisq(1), 2n*x);    label="2n pdf(2nx|χ²(1))", ls=:dot, color=:black)

    P2 = plot(title="empirical cumulative distribution function", titlefontsize=8)
    plot!(x, ecdf.(Ref(KL_MLE),   x); label="KL_MLE")
    plot!(x, ecdf.(Ref(KL_Bayes), x); label="KL_Bayes(β=$(β), a=$a, b=$b)", ls=:dash)
    x = range(extrema(bin)..., length=400)
    plot!(x, cdf.(Chisq(1), 2n*x);    label="cdf(2nx|χ²(1))", ls=:dot, color=:black)

    sleep(0.1)
    plot(P1, P2; size=(800, 250), layout=(1,2))
end

plot_esitmation_errors(0.3, 10)

plot_esitmation_errors(0.3, 100; bin=0:0.0025:0.05)

MSE(q, p) = sum((p.p - q.p)^2*pdf(q, k) for k in support(q)) # mean squared error

# expextation value
E(dist, f) = sum(f(x)*pdf(dist, x) for x in support(dist))

# expextation value (extremum cases eliminated)
function EE(dist, f)
    s = support(dist)
    sum(f(x)*pdf(dist, x) for x in s[2:end-1])/(1 - pdf(dist, s[1]) - pdf(dist, s[end]))
end

function print_estimation_errors(true_prob, n; β=1.0, a=0.5, b=0.5)
    q = Bernoulli(true_prob)
    q_sample = Binomial(n, q.p)
    f_KL_Bayes(x) = KL(q, pred_Bayes(n, x; β=β, a=a, b=b))
    f_KL_MLE(x) = KL(q, pred_MLE(n, x))
    f_MSE_Bayes(x) = MSE(q, pred_Bayes(n, x; β=β, a=a, b=b))
    f_MSE_MLE(x) = MSE(q, pred_MLE(n, x))

    @show true_prob
    @show n
    @show β
    @show a
    @show b
    println()
    @show E(q_sample, f_KL_Bayes)
    @show E(q_sample, f_KL_MLE)
    @show EE(q_sample, f_KL_Bayes)
    @show EE(q_sample, f_KL_MLE)
    println()
    @show E(q_sample, f_MSE_Bayes)
    @show E(q_sample, f_MSE_MLE)
    @show EE(q_sample, f_MSE_Bayes)
    @show EE(q_sample, f_MSE_MLE)
    nothing
end

print_estimation_errors(0.3, 10)

print_estimation_errors(0.3, 100)

function plot_mean_estimation_errors(n; 
        legend_KL=true, legend_MSE=true, 
        β₁=1.0, a₁=0.5, b₁=0.5,
        β₂=1.0, a₂=1.0, b₂=1.0
    )
    f_KL_MLE(true_prob, x)    = KL(Bernoulli(true_prob), pred_MLE(n, x))
    f_KL_BayesJ(true_prob, x) = KL(Bernoulli(true_prob), pred_Bayes(n, x; β=β₁, a=a₁, b=b₁))
    f_KL_Bayes(true_prob, x)  = KL(Bernoulli(true_prob), pred_Bayes(n, x; β=β₂, a=a₂, b=b₂))

    f_MSE_MLE(true_prob, x)    = MSE(Bernoulli(true_prob), pred_MLE(n, x))
    f_MSE_BayesJ(true_prob, x) = MSE(Bernoulli(true_prob), pred_Bayes(n, x; β=β₁, a=a₁, b=b₁))
    f_MSE_Bayes(true_prob, x)  = MSE(Bernoulli(true_prob), pred_Bayes(n, x; β=β₂, a=a₂, b=b₂))
    
    EE_KL_MLE(true_prob)    = EE(Binomial(n, true_prob), x->f_KL_MLE(true_prob, x))
    EE_KL_BayesJ(true_prob) = EE(Binomial(n, true_prob), x->f_KL_BayesJ(true_prob, x))
    EE_KL_Bayes(true_prob)  = EE(Binomial(n, true_prob), x->f_KL_Bayes(true_prob, x))

    E_MSE_MLE(true_prob)    = EE(Binomial(n, true_prob), x->f_MSE_MLE(true_prob, x))
    E_MSE_BayesJ(true_prob) = EE(Binomial(n, true_prob), x->f_MSE_BayesJ(true_prob, x))
    E_MSE_Bayes(true_prob)  = EE(Binomial(n, true_prob), x->f_MSE_Bayes(true_prob, x))

    true_prob = 0.01:0.01:0.99
    
    P1 = plot(title="mean prediction error (extremum cases eliminated)", titlefontsize=10)
    plot!(xlabel="true success probability", guidefontsize=9)
    plot!(xtick=0:0.1:1)
    plot!(legend=legend_KL, legendfontsize=7)
    plot!(true_prob, EE_KL_MLE.(true_prob);    label="MLE")
    plot!(true_prob, EE_KL_BayesJ.(true_prob); label="Bayes($β₁; $a₁, $b₁)", ls=:dash)
    plot!(true_prob, EE_KL_Bayes.(true_prob);  label="Bayes($β₂; $a₂, $b₂)", ls=:dashdot)
    
    P2 = plot(title="mean squared error (n = $n)", titlefontsize=10)
    plot!(xlabel="true success probability", guidefontsize=9)
    plot!(xtick=0:0.1:1)
    plot!(legend=legend_MSE, legendfontsize=7)
    plot!(true_prob, E_MSE_MLE.(true_prob);    label="MLE")
    plot!(true_prob, E_MSE_BayesJ.(true_prob); label="Bayes($β₁; $a₁, $b₁)", ls=:dash)
    plot!(true_prob, E_MSE_Bayes.(true_prob);  label="Bayes($β₂; $a₂, $b₂)", ls=:dashdot)
    
    plot(P1, P2; size=(800, 250), layout=(1,2))
end

plot_mean_estimation_errors(10)

plot_mean_estimation_errors(20)

plot_mean_estimation_errors(50; legend_KL=:top)

plot_mean_estimation_errors(100)

# 予測分布の対数尤度の -1/n 倍
function T(n, k; β=1.0, a=0.5, b=0.5, w̄=(β*k+a)/(β*n+a+b))
    p = k/n
    -(safexlogy(p, w̄) + safexlogy(1-p, 1-w̄))
end

# 事後分布に関する log p(X_i|w) の分散のiに関する和の 1/n 倍
function V(n, k; β=1.0, a=0.5, b=0.5)
    p = k/n
    p*trigamma(β*k+a) + (1-p)*trigamma(β*(n-k)+b) - trigamma(β*n+a+b)
end

aic(n, k) = T(n, k; a=0.0, b=0.0) + 1/n
waic(n, k; β=1.0, a=0.5, b=0.5) = T(n, k; β=β, a=a, b=b) + β*V(n, k; β=β, a=a, b=b)
function loocv_old(n, k; β=1.0, a=0.5, b=0.5)
    p = k/n
    (
        -   p  *(logbeta((1-β)+β*k+a, β*(n-k)+b) - logbeta(-β+β*k+a, β*(n-k)+b))
        - (1-p)*(logbeta(β*k+a, (1-β)+β*(n-k)+b) - logbeta(β*k+a, -β+β*(n-k)+b))
    )
end
function loocv(n, k; β=1.0, a=0.5, b=0.5)
    p = k/n
    -safexlogy(p, -β+β*k+a) - safexlogy(1-p, -β+β*(n-k)+b) + log(-β+β*n+a+b)
end

AIC(n, k)  = 2n*aic(n, k)
WAIC(n, k; β=1.0, a=0.5, b=0.5) = 2n*waic(n, k; β=β, a=a, b=b)
LOOCV(n, k; β=1.0, a=0.5, b=0.5) = 2n*loocv(n, k; β=β, a=a, b=b)

# https://ja.wikipedia.org/wiki/%E8%B5%A4%E6%B1%A0%E6%83%85%E5%A0%B1%E9%87%8F%E8%A6%8F%E6%BA%96
# 
#    AICc = AIC + 2nparam(nparam + 1)/(samplesize - nparam - 1)
#
# nparam = 1,  samplesize = n
#
AICc(n, k) = AIC(n, k) + (2*1^2 + 2*1)/(n-1-1)

ge(pred_dist, q, n, k) = G(q, pred_dist(n, k))
pe(pred_dist, q, n, k) = KL(q, pred_dist(n, k))

GE(pred_dist, q, n, k) = 2n*ge(pred_dist, q, n, k)
PE(pred_dist, q, n, k) = 2n*pe(pred_dist, q, n, k)

aic0(q, n, k) = T(n, k; a=0.0, b=0.0, w̄=q.p)
waic0(q, n, k; β=1.0, a=0.5, b=0.5) = T(n, k; β=β, a=a, b=b, w̄=q.p)
loocv0(q, n, k; β=1.0, a=0.5, b=0.5) = waic0(q, n, k; β=β, a=a, b=b)

AIC0(q, n, k) = 2n*aic0(q, n, k)
WAIC0(q, n, k; β=1.0, a=0.5, b=0.5) = 2n*waic0(q, n, k; β=β, a=a, b=b)
LOOCV0(q, n, k; β=1.0, a=0.5, b=0.5) = 2n*loocv0(q, n, k; β=β, a=a, b=b)

function supp(dist, m=2)
    μ, σ = mean(dist), std(dist)
    a = max(minimum(dist), floor(Int, μ - m*σ))
    b = min(maximum(dist),  ceil(Int, μ + m*σ))
    a:b
end

sample_supp(q, n, m=2) = supp(Binomial(n, q.p), m)

E(q, n, f) = E(Binomial(n, q.p), f)
EE(q, n, f) = EE(Binomial(n, q.p), f)

loocv_old(10, 0) # error

loocv(10, 0)

function plot_EEs_and_Es(n; f=[fill(true, 5); false])
    w = range(0.005, 0.995, step=0.005)
    q = Bernoulli.(w)
    EE_GE_MLE  = (q -> EE(q, n, k->GE(pred_MLE, q, n, k))).(q)
    EE_AIC     = (q -> EE(q, n, k->AIC(n, k))).(q)
    EE_AICc    = (q -> EE(q, n, k->AICc(n, k))).(q)
    E_GE_Bayes = (q -> E(q, n, k->GE(pred_Bayes, q, n, k))).(q)
    E_WAIC     = (q -> E(q, n, k->WAIC(n, k))).(q)
    E_LOOCV    = (q -> E(q, n, k->LOOCV(n, k))).(q)

    pal = palette(:default)
    P = plot(title="n = $n", titlefontsize=10)
    plot!(xlabel="w", ylabel="expectation value", xtick=0:0.1:1)
    f[1] && plot!(w, EE_GE_MLE;  label="EE[GE_MLE]",   ls=:dot,     color=pal[1])
    f[2] && plot!(w, EE_AIC;     label="EE[AIC]",      ls=:dash,    color=pal[2])
    f[3] && plot!(w, E_GE_Bayes; label="E[GE_Bayes]",  ls=:solid,   color=pal[3])
    f[4] && plot!(w, E_WAIC;     label="E[WAIC]",      ls=:dash,    color=pal[4])
    f[5] && plot!(w, E_LOOCV;    label="E[LOOCV]",     ls=:dashdot, color=pal[5])
    f[6] && plot!(w, EE_AICc;     label="EE[AICc]",    ls=:dash,    color=pal[6])
    
    plot(P; size=(400, 250))
end

plot_EEs_and_Es(10)

plot_EEs_and_Es(20)

plot_EEs_and_Es(50)

plot_EEs_and_Es(100)

plot_EEs_and_Es(10; f=[true; true; fill(false, 3); true])

plot_EEs_and_Es(20; f=[true; true; fill(false, 3); true])

function make_xtick(n, w)
    q = Bernoulli(w)
    q_sample = Binomial(n, q.p)
    
    t3 = round(Int, mean(q_sample))
    t2, t4 = extrema(sample_supp(q, n, 1))
    t1, t5 = extrema(sample_supp(q, n, 2))
    xticklabel = fill("", t5 - t1 + 1)
    xticklabel[1] = "$t1"
    xticklabel[t2-t1+1] = "$t2"
    xticklabel[t3-t1+1] = "$t3"
    xticklabel[t4-t1+1] = "$t4"
    xticklabel[end] = "$t5"
    tmp = cdf.(q_sample, t1-1:t5)
    xtick = @views @. (tmp[1:end-1] + tmp[2:end])/2
    
    xtick, xticklabel
end

function plot_AIC(n, w;
        legend = :top,
        ylim = (-4, 6),
        ytick = ylim[1]:ylim[2]
    )

    q = Bernoulli(w)
    q_sample = Binomial(n, q.p)

    x = range(0, 1, length=801)
    k = quantile.(q_sample, x)
    xtick, xticklabel = make_xtick(n, w)
    
    pred_error = @. PE(pred_MLE, q, n, k)
    AIC_minus_AIC0 = @. AIC(n,k) - AIC0(q,n,k)
    
    P = plot(title="n = $n,  w = $w"; titlefontsize=10)
    plot!(xlim=(0, 1), ylim=ylim)
    plot!(xtick=(xtick, xticklabel), ytick=ytick)
    plot!(xlabel="k", guidefontsize=10)
    plot!(legend=legend)
    plot!(x, pred_error; label="Prediction Error of MLE")
    plot!(x, AIC_minus_AIC0; label="AIC - AIC0", ls=:dash, lw=1.3)
    hline!([0]; color=:black, ls=:dot, alpha=0.8, label="")
    
    plot(P; size=(500, 350))
end

function plot_WAIC_LOOCV(n, w;
        legend = :top,
        ylim = (-4, 6),
        ytick = ylim[1]:ylim[2]
    )

    q = Bernoulli(w)
    q_sample = Binomial(n, w)

    x = range(0, 1, length=801)
    k = quantile.(q_sample, x)
    xtick, xticklabel = make_xtick(n, w)
    
    pred_error = @. PE(pred_Bayes, q, n, k)
    WAIC_minus_WAIC0 = @. WAIC(n,k) - WAIC0(q,n,k)
    LOOCV_minus_LOOCV0 = @. LOOCV(n,k) - LOOCV0(q,n,k)
    
    P = plot(title="n = $n,  w = $w"; titlefontsize=10)
    plot!(xlim=(0, 1), ylim=ylim)
    plot!(xtick=(xtick, xticklabel), ytick=ytick)
    plot!(xlabel="k", guidefontsize=10)
    plot!(legend=legend)
    plot!(x, pred_error; label="Prediction Error of Bayes")
    plot!(x, WAIC_minus_WAIC0; label="WAIC - WAIC0", ls=:dash, lw=1.3)
    plot!(x, LOOCV_minus_LOOCV0; label="LOOCV - LOOCV0", ls=:dot, lw=1.3)
    hline!([0]; color=:black, ls=:dot, alpha=0.8, label="")
    
    plot(P; size=(500, 350))
end

n = 10
w = 0.4
@show n
@show w
q = Bernoulli(w)
prob_fail_AIC = E(q, n, k->AIC(n,k) < AIC0(q,n,k)) |> rd
@show prob_fail_AIC
println()
for k in sample_supp(q, n, 3)
    pred_error = PE(pred_MLE, q, n, k)
    AIC_minus_AIC0 = AIC(n,k) - AIC0(q,n,k)
    (
        k = k,
        prob = pdf(Binomial(n, q.p), k)|>rd, 
        PE_MLE = pred_error |>rd, 
        AIC_minus_AIC0 = AIC_minus_AIC0 |>rd, 
        PE_plus_AIC_minus_AIC0 = pred_error + AIC_minus_AIC0 |>rd
    ) |> println
end

n = 10
w = 0.4
@show n
@show w
q = Bernoulli(w)
prob_fail_WAIC = E(q, n, k -> WAIC(n,k) < WAIC0(q,n,k)) |> rd
@show prob_fail_WAIC
println()
for k in sample_supp(q, n, 3)
    pred_error = PE(pred_Bayes, q, n, k)
    WAIC_minus_WAIC0 = WAIC(n,k) - WAIC0(q,n,k)
    (
        k = k, 
        prob = pdf(Binomial(n, q.p), k)|>rd, 
        PE_Bayes = pred_error |>rd, 
        WAIC_minus_WAIC0 = WAIC_minus_WAIC0 |>rd, 
        PE_plus_WAIC_minus_WAIC0 = pred_error + WAIC_minus_WAIC0 |>rd
    ) |> println
end

n = 10
w = 0.4
@show n
@show w
q = Bernoulli(w)
prob_fail_LOOCV = EE(q, n, k->LOOCV(n,k) < LOOCV0(q,n,k)) |> rd
@show prob_fail_LOOCV
println()
for k in sample_supp(q, n, 3)
    pred_error = PE(pred_Bayes, q, n, k)
    LOOCV_minus_LOOCV0 = LOOCV(n,k) - LOOCV0(q,n,k)
    (
        k = k, 
        prob = pdf(Binomial(n, q.p), k)|>rd, 
        PE_Bayes = pred_error |>rd, 
        LOOCV_minus_LOOCVC0 = LOOCV_minus_LOOCV0 |>rd, 
        PE_plus_LOOCV_minus_LOOCV0 = pred_error + LOOCV_minus_LOOCV0 |>rd
    ) |> println
end

n = 10
w = 0.4
plot_AIC(n, w) |> display
plot_WAIC_LOOCV(n, w) |> display
pyplotclf()

n = 20
w = 0.4
plot_AIC(n, w) |> display
plot_WAIC_LOOCV(n, w) |> display
pyplotclf()

n = 100
w = 0.4
@show n
@show w
q = Bernoulli(w)
prob_fail_AIC = E(q, n, k->AIC(n,k) < AIC0(q,n,k)) |> rd
@show prob_fail_AIC
println()
for k in sample_supp(q, n, 3)
    pred_error = PE(pred_MLE, q, n, k)
    AIC_minus_AIC0 = AIC(n,k) - AIC0(q,n,k)
    (
        k = k,
        prob = pdf(Binomial(n, q.p), k)|>rd, 
        PE_MLE = pred_error |>rd, 
        AIC_minus_AIC0 = AIC_minus_AIC0 |>rd, 
        PE_plus_AIC_minus_AIC0 = pred_error + AIC_minus_AIC0 |>rd
    ) |> println
end

n = 100
w = 0.4
@show n
@show w
q = Bernoulli(w)
prob_fail_WAIC = E(q, n, k -> WAIC(n,k) < WAIC0(q,n,k)) |> rd
@show prob_fail_WAIC
println()
for k in sample_supp(q, n, 3)
    pred_error = PE(pred_Bayes, q, n, k)
    WAIC_minus_WAIC0 = WAIC(n,k) - WAIC0(q,n,k)
    (
        k = k, 
        prob = pdf(Binomial(n, q.p), k)|>rd, 
        PE_Bayes = pred_error |>rd, 
        WAIC_minus_WAIC0 = WAIC_minus_WAIC0 |>rd, 
        PE_plus_WAIC_minus_WAIC0 = pred_error + WAIC_minus_WAIC0 |>rd
    ) |> println
end

n = 100
w = 0.4
@show n
@show w
q = Bernoulli(w)
prob_fail_LOOCV = EE(q, n, k->LOOCV(n,k) < LOOCV0(q,n,k)) |> rd
@show prob_fail_LOOCV
println()
for k in sample_supp(q, n, 3)
    pred_error = PE(pred_Bayes, q, n, k)
    LOOCV_minus_LOOCV0 = LOOCV(n,k) - LOOCV0(q,n,k)
    (
        k = k, 
        prob = pdf(Binomial(n, q.p), k)|>rd, 
        PE_Bayes = pred_error |>rd, 
        LOOCV_minus_LOOCVC0 = LOOCV_minus_LOOCV0 |>rd, 
        PE_plus_LOOCV_minus_LOOCV0 = pred_error + LOOCV_minus_LOOCV0 |>rd
    ) |> println
end

n = 100
w = 0.4
plot_AIC(n, w) |> display
plot_WAIC_LOOCV(n, w) |> display
pyplotclf()

n = 1000
w = 0.4
plot_AIC(n, w) |> display
plot_WAIC_LOOCV(n, w) |> display
pyplotclf()

w = 0.4
@show w
q = Bernoulli(w)
for n in round.(Int, 10 .^ (1:0.5:6))
    prob_fail_AIC = E(q, n, k -> AIC(n,k) < AIC0(q,n,k))
    (n=n, prob_fail_AIC=prob_fail_AIC|>x->rd(x,6)) |> println
end

w = 0.4
@show w
q = Bernoulli(w)
for n in round.(Int, 10 .^ (1:0.5:6))
    prob_fail_WAIC = E(q, n, k -> WAIC(n,k) < WAIC0(q,n,k))
    (n=n, prob_fail_WAIC=prob_fail_WAIC|>x->rd(x,6)) |> println
end

w = 0.4
@show w
q = Bernoulli(w)
for n in round.(Int, 10 .^ (1:0.5:6))
    prob_fail_LOOCV = EE(q, n, k -> LOOCV(n,k) < LOOCV0(q,n,k))
    (n=n, prob_fail_LOOCV=prob_fail_LOOCV|>x->rd(x,6)) |> println
end

ccdf(Chisq(1), 2)|>x->rd(x,6)

bic(n, k) = T(n, k; a=0.0, b=0.0) + (1/2)*log(n)/n
bic0(q, n, k) = aic0(q, n, k)

BIC(n, k) = 2n*bic(n, k)
BIC0(q, n, k) = 2n*bic0(q, n, k)

function wbic(n, k; a=0.5, b=0.5)
    β = 1/log(n)
    p = k/n
    -p*digamma(β*k+a) - (1-p)*digamma(β*(n-k)+b) + digamma(β*n+a+b)
end
wbic0(q, n, k; a=0.5, b=0.5) = waic0(q, n, k; β=1.0, a=a, b=b)

WBIC(n, k; a=0.5, b=0.5) = 2n*wbic(n, k; a=a, b=b)
WBIC0(q, n, k; a=0.5, b=0.5) = 2n*wbic0(q, n, k; a=a, b=b)

FE(n, k; β=1.0, a=0.5, b=0.5) = (2/β)*(
    - (loggamma(a+k) - loggamma(a)) 
    - (loggamma(b+n-k) - loggamma(b))
    + (loggamma(a+b+n) - loggamma(a+b))
)
FE0(q, n, k; β=1.0) = (2/β)*(- safemul(k, log(q.p)) - safemul(n-k, log(1-q.p)))

function plot_BIC_WBIC_FE(n, w;
        legend = :bottom,
        ylim = (-4, 6),
        ytick = ylim[1]:ylim[2]
    )

    q = Bernoulli(w)
    q_sample = Binomial(n, q.p)

    x = range(0, 1, length=801)
    k = quantile.(q_sample, x)
    xtick, xticklabel = make_xtick(n, w)
    
    FE_minus_FE0 = @. FE(n,k) - FE0(q,n,k)
    BIC_minus_BIC0 = @. BIC(n,k) - BIC0(q,n,k)
    WBIC_minus_WBIC0 = @. WBIC(n,k) - WBIC0(q,n,k)
    
    P = plot(title="n = $n,  w = $w"; titlefontsize=10)
    plot!(xlim=(0, 1), ylim=ylim)
    plot!(xtick=(xtick, xticklabel), ytick=ytick)
    plot!(xlabel="k", guidefontsize=10)
    plot!(legend=legend)
    plot!(x, FE_minus_FE0; label="FE - FE0")
    plot!(x, WBIC_minus_WBIC0; label="WBIC - WBIC0", ls=:dash, lw=1.2)
    plot!(x, BIC_minus_BIC0; label="BIC - BIC0", ls=:dot, lw=1.5)
    hline!([0]; color=:black, ls=:dot, alpha=0.8, label="")
    
    plot(P; size=(500, 350))
end

n = 10
w = 0.4
q = Bernoulli(w)
@show n
@show w
for k in sample_supp(q, n, 3)
    (
        k=k, 
        prob=pdf(Binomial(n, q.p), k)|>rd, 
        FE0=FE0(q,n,k)|>rd, 
        BIC0=BIC0(q,n,k)|>rd, 
        WBIC0=WBIC0(q,n,k)|>rd, 
    ) |> println
end

n = 10
w = 0.4
q = Bernoulli(w)
@show n
@show w
for k in sample_supp(q, n, 3)
    F0 = FE0(q,n,k)
    (
        k=k, 
        prob=pdf(Binomial(n, q.p), k)|>x->rd(x,4), 
        FE_minus_FE0=FE(n,k)-F0|>rd, 
        BIC_minus_FE0=BIC(n,k)-F0|>rd,
        WBIC_minus_FE0=WBIC(n,k)-F0|>rd
    ) |> println
end

plot_BIC_WBIC_FE(n, w)

n = 20
w = 0.4
q = Bernoulli(w)
@show n
@show w
for k in sample_supp(q, n, 3)
    F0 = FE0(q,n,k)
    (
        k=k, 
        prob=pdf(Binomial(n, q.p), k)|>x->rd(x,4), 
        FE_minus_FE0=FE(n,k)-F0|>rd, 
        BIC_minus_FE0=BIC(n,k)-F0|>rd,
        WBIC_minus_FE0=WBIC(n,k)-F0|>rd
    ) |> println
end

plot_BIC_WBIC_FE(n, w)

n = 50
w = 0.4
q = Bernoulli(w)
@show n
@show w
for k in sample_supp(q, n, 3)
    F0 = FE0(q,n,k)
    (
        k=k, 
        prob=pdf(Binomial(n, q.p), k)|>x->rd(x,4), 
        FE_minus_FE0=FE(n,k)-F0|>rd, 
        BIC_minus_FE0=BIC(n,k)-F0|>rd,
        WBIC_minus_FE0=WBIC(n,k)-F0|>rd
    ) |> println
end

plot_BIC_WBIC_FE(n, w)

n = 100
w = 0.4
q = Bernoulli(w)
@show n
@show w
for k in sample_supp(q, n, 3)
    F0 = FE0(q,n,k)
    (
        k=k, 
        prob=pdf(Binomial(n, q.p), k)|>x->rd(x,4), 
        FE_minus_FE0=FE(n,k)-F0|>rd, 
        BIC_minus_FE0=BIC(n,k)-F0|>rd,
        WBIC_minus_FE0=WBIC(n,k)-F0|>rd
    ) |> println
end

plot_BIC_WBIC_FE(n, w)

n = 1000
w = 0.4
q = Bernoulli(w)
@show n
@show w
for k in sample_supp(q, n, 3)
    F0 = FE0(q,n,k)
    (
        k=k, 
        prob=pdf(Binomial(n, q.p), k)|>x->rd(x,4), 
        FE_minus_FE0=FE(n,k)-F0|>rd, 
        BIC_minus_FE0=BIC(n,k)-F0|>rd,
        WBIC_minus_FE0=WBIC(n,k)-F0|>rd
    ) |> println
end

plot_BIC_WBIC_FE(n, w; ylim=(-4, 8))

w = 0.4
@show w
q = Bernoulli(w)
for n in round.(Int, 10 .^ (1:0.5:6))
    prob_fail_BIC = E(q, n, k -> BIC(n,k) < BIC0(q,n,k))
    (n=n, prob_fail_BIC=prob_fail_BIC|>x->rd(x,6)) |> println
end

w = 0.4
@show w
q = Bernoulli(w)
for n in round.(Int, 10 .^ (1:0.5:6))
    prob_fail_WBIC = E(q, n, k -> WBIC(n,k) < WBIC0(q,n,k))
    (n=n, prob_fail_WBIC=prob_fail_WBIC|>x->rd(x,6)) |> println
end

w = 0.4
@show w
q = Bernoulli(w)
for n in round.(Int, 10 .^ (1:0.5:6))
    prob_fail_FE = E(q, n, k -> FE(n,k) < FE0(q,n,k))
    (n=n, prob_fail_FE=prob_fail_FE|>x->rd(x,6)) |> println
end

# ジェネリックな信頼区間函数

is_in(x, CI) = CI[1] ≤ x ≤ CI[2]

function conf_int(pval_func, n, k; α=0.05)
    f(w) = pval_func(Bernoulli(w), n, k) - α
    CI = find_zeros(f, 0, 1)
    if k == 0
        [0.0, CI[1]]
    elseif k == n
        [CI[1], 1.0]
    else
        CI
    end
end

# 二項分布の正規分布近似による仮説検定のP値

function pval_normal(q, n, k)
    q_sample = Binomial(n, q.p)
    μ = mean(q_sample)
    σ = std(q_sample)
    z = (k - μ)/σ
    2ccdf(Normal(), abs(z))
end

n = 10
[(k=k, conf_int=conf_int(pval_normal, n, k).|>rd) for k in 0:n]

# ベイズ信用区間

post_dist(n, k; a=0.5, b=0.5) = Beta(a+k, b+n-k)

cred_int(n, k; a=0.5, b=0.5, α=0.05) = quantile.(post_dist(n, k; a=a, b=b), [α/2, 1-α/2])

n = 10
[(k=k, cred_int=cred_int(n, k).|>rd) for k in 0:n]

# ベイズ信用区間の別の実装

# 以下の函数は事後分布を使ったP値の類似物
function pval_post(q, n, k; a=0.5, b=0.5)
    d = post_dist(n, k; a=a, b=b)
    min(2cdf(d, q.p), 2ccdf(d, q.p), 1)
end

# 上の計算結果と一致する
n = 10
[(k=k, conf_int=conf_int(pval_post, n, k).|>rd) for k in 0:n]

# AICを使ったP値の類似物 (対数尤度比のχ²検定のP値)

function pval_AIC(q, n, k)
    x = AIC0(q, n, k) - (AIC(n, k) - 2)
    ccdf(Chisq(1), x)
end

n = 10
[(k=k, conf_int=conf_int(pval_AIC, n, k).|>rd) for k in 0:n]

# WAICを使ったP値の類似物

function pval_WAIC(q, n, k; a=0.5, b=0.5)
    #x = WAIC0(q, n, k; a=a, b=b) - (WAIC(n, k; a=a, b=b) - 2)
    x = WAIC0(q, n, k; a=a, b=b) - 2n*T(n, k; a=a, b=b)
    ccdf(Chisq(1), x)
end

n = 10
[(k=k, conf_int=conf_int(pval_WAIC, n, k).|>rd) for k in 0:n]

n = 40
w = 0.29
q = Bernoulli(w)
sample = rand(Binomial(n, q.p), 10^5)
y1 = WAIC0.(q, n, sample) - WAIC.(n, sample) .+ 2
y2 = WAIC0.(q, n, sample) - 2n*T.(n, sample)
P1 = histogram(y1, xlim=(-1, 10), bin=-1:1:10, alpha=0.5, norm=true, label="WAIC0 - (WAIC - 2)")
plot!(x->pdf(Chisq(1),x), 0.2, 10; label="Chisq(1)")
P2 = histogram(y2, xlim=(-1, 10), bin=-1:1:10, alpha=0.5, norm=true, label="WAIC0 - 2n T")
plot!(x->pdf(Chisq(1),x), 0.2, 10; label="Chisq(1)")
plot(P1, P2, size=(800, 250))

n = 100
w = 0.3
q = Bernoulli(w)
@show n
@show w
[(k=k, pval_normal=pval_normal(q, n, k)|>rd, pval_post=pval_post(q, n, k)|>rd, pval_AIC=pval_AIC(q, n, k)|>rd, pval_WAIC=pval_WAIC(q, n, k)|>rd) for k in sample_supp(q, n, 3)]

n = 30
[
    (
        k=k, 
        Normal=conf_int(pval_normal, n, k).|>rd, 
        Bayes=cred_int(n, k).|>rd, 
        AIC=conf_int(pval_AIC, n, k).|>rd, 
        WAIC=conf_int(pval_WAIC, n, k).|>rd
    ) for k in 0:n
]

function animate_pvalue_fuctions(n; fps=n/10, f=fill(true,4), pal=palette(:default))
    anim = @animate for k in 0:n
        w = 0.01:0.01:0.99
        y_normal = pval_normal.(Bernoulli.(w), n, k)
        y_Bayes  = pval_post.(Bernoulli.(w), n, k)
        y_AIC    = pval_AIC.(Bernoulli.(w), n, k)
        y_WAIC   = pval_WAIC.(Bernoulli.(w), n, k)
        P = plot(title="p-value functions for n=$n, k=" * @sprintf("%03d",k), titlefontsize=10)
        plot!(ylim=(-0.02, 1.02))
        plot!(xtick=0:0.1:1, ytick=0:0.1:1)
        f[1] && plot!(w, y_normal; label="Normal", ls=:solid, color=pal[1])
        f[2] && plot!(w, y_Bayes; label="Posterior", ls=:dash, color=pal[2])
        f[3] && plot!(w, y_AIC; label="AIC", ls=:dashdot, color=pal[3])
        f[4] && plot!(w, y_WAIC; label="WAIC", ls=:dot, lw=1.3, color=pal[4])
        plot(P; size=(500, 350))
    end
    pyplotclf()

    g = Int.(f)
    gifname = "pvalue_functions_$(n)_$(g[1])$(g[2])$(g[3])$(g[4]).gif"
    gif(anim, gifname; fps=fps)
    displayfile("image/gif", gifname)
end

@time animate_pvalue_fuctions(10)

@time animate_pvalue_fuctions(20)

@time animate_pvalue_fuctions(50)

@time animate_pvalue_fuctions(100)

@time animate_pvalue_fuctions(10; f=[true, true, false, false])

@time animate_pvalue_fuctions(20; f=[true, true, false, false])

@time animate_pvalue_fuctions(50; f=[true, true, false, false])

@time animate_pvalue_fuctions(100; f=[true, true, false, false])