using Distributions
using Plots
using Random
using SymPy

@vars x μ # variables
n = 4
X = symbols("X1:$(n+1)") # sample of size n = 4

negtwologp(x, μ) = log(2PI) + (x - μ)^2

X̄ = mean(X).factor() # sample mean

VX = mean((X .- X̄).^2) # sample variance

A = sum(negtwologp(X[i], μ) for i in 1:n).expand().simplify()

B = (n*((μ - X̄)^2 + VX)).expand()

A - B

ecdf(A, x) = count(a ≤ x for a in A)/length(A)

function hack!(tmp; dist = Normal(), c = 2.0, nstep = 1, maxiters = 1000)
    μ = mean(dist)
    σ = std(dist)
    S = zero(eltype(tmp))
    n = 0
    iter = 0
    rng = Random.default_rng()
    while abs(S - n*μ) ≤ c*σ*√n
        iter > maxiters && break
        rand!(rng, dist, tmp)
        S += sum(tmp)
        n += nstep
        iter += 1
    end
    iter
end
function hack!(; dist = Normal(), c = 2.0, nstep = 1, maxiters = 1000)
    tmp = rand(dist, nstep)
    hack!(tmp; dist = Normal(), c = 2.0, nstep = 1, maxiters = 1000)
end

function plot_hack(; dist_str = "Normal()", c = 2.0, 
        nstep = 1, maxiters = 1000, L = 10^6, 
        xticklength = 20, ytickstep=0.05)
    dist = eval(Meta.parse(dist_str))
    tmp = rand(dist, nstep)
    @time H = [hack!(tmp; dist, c, maxiters, nstep) for _ in 1:L]
    
    xtick = range(0, maxiters; length=xticklength+1)
    ytick = 0:ytickstep:1
    title = "Stopping condition: |X̅ - μ| > $(round(c, digits=2))σ/√n  (X_i ∼ $dist_str)"
   
    n = 0:maxiters
    plot(n, ecdf.(Ref(H), n); label="ecdf", legend=:bottomright)
    plot!(; xlabel="iterations (nstep = $nstep)", xtick, ytick)
    plot!(; title, titlefontsize=11)
endえn

plot_hack(c = quantile(Normal(), 1 - 0.05/2))

plot_hack(c = quantile(Normal(), 1 - 0.05/2), maxiters = 2000, xticklength = 10)

plot_hack(c = quantile(Normal(), 1 - 0.05/2), nstep = 10, maxiters = 200)

plot_hack(c = quantile(Normal(), 1 - 0.01/2), ytickstep = 0.01)

plot_hack(c = quantile(Normal(), 1 - 0.001/2), ytickstep=0.002)

plot_hack(dist_str = "Exponential()", c = quantile(Normal(), 1 - 0.05/2))

function pval(a, b, X)
    n = length(X)
    X̄ = mean(X)
    if a ≤ X̄ ≤ b
        1.0
    else
        2ccdf(Normal(0.0, 1/√n), max(a - X̄, X̄ - b))
    end
end

function plot_pval(a, b, X;
        μ₀ = 0.0, 
        x = range(-0.7, 0.7; length=701),
        xtick = range(extrema(x)...; step=0.1),
        n = length(X),
        title = "p-values of $a < μ + x < $b for μ₀ = $μ₀,  n = $n"
    )
    ylim = (-0.02, 1.02)
    ytick = 0:0.05:1
    p = pval.(a .- x, b .- x, Ref(X))
    plot(x, p; label="", xtick, ytick, ylim, xlabel="x")
    plot!(; title, titlefontsize=12)
end

n = 2^16
X = randn(n)
a, b = -0.1, 0.2

plot_pval(a, b, X[1:2^4])

plot_pval(a, b, X[1:2^6])

plot_pval(a, b, X[1:2^10])

plot_pval(a, b, X)

ts = range(log2(2^4), log2(n); length=200)
ts = [fill(log2(2^4), 20); ts; fill(log2(n), 20); reverse(ts)]
@gif for t in ts
    k = round(Int, 2^t)
    plot_pval(a, b, X[1:k])
end

ts = range(log2(2^4), log2(n); length=200)
ts = [fill(log2(2^4), 20); ts; fill(log2(n), 20); reverse(ts)]
@gif for t in ts
    k = round(Int, 2^t)
    plot_pval(0, 0, X[1:k];
        x = range(-0.7, 0.7; length=2801),
        title = "p-values of μ + x = 0 for μ₀ = 0.0,  n = $k"
    )
end