using InstantiateFromURL
# optionally add arguments to force installation: instantiate = true, precompile = true
github_project("QuantEcon/quantecon-notebooks-julia", version = "0.8.0")

using LinearAlgebra, Statistics
using Plots, Distributions, Random, Statistics
gr(fmt = :png, size = (900, 500))

function ksl(distribution, n = 100)
    title = nameof(typeof(distribution))
    observations = rand(distribution, n)
    sample_means = cumsum(observations) ./ (1:n)
    μ = mean(distribution)
    plot(repeat((1:n)', 2),
         [zeros(1, n); observations'], label = "", color = :grey, alpha = 0.5)
    plot!(1:n, observations, color = :grey, markershape = :circle,
          alpha = 0.5, label = "", linewidth = 0)
    if !isnan(μ)
        hline!([μ], color = :black, linewidth = 1.5, linestyle = :dash, grid = false,
               label = ["Mean"])
    end
    plot!(1:n, sample_means, linewidth = 3, alpha = 0.6, color = :green,
          label = "Sample mean")
    return plot!(title = title)
end

distributions = [TDist(10), Beta(2, 2), Gamma(5, 2), Poisson(4), LogNormal(0.5),
                 Exponential(1)]

ksl(Normal())

Random.seed!(0); # reproducible results

plot(ksl.(sample(distributions, 3, replace = false))..., layout = (3, 1), legend = false)

Random.seed!(0); # reproducible results

ksl(Cauchy())

Random.seed!(0); # reproducible results

function plot_means(n = 1000)
    sample_mean = cumsum(rand(Cauchy(), n)) ./ (1:n)
    plot(1:n, sample_mean, color = :red, alpha = 0.6, label = "Sample Mean", linewidth = 3)
    return hline!([0], color = :black, linestyle = :dash, label = "", grid = false)
end

plot_means()

binomial_pdf(n) =
    bar(0:n, pdf.(Binomial(n), 0:n),
        xticks = 0:10, ylim = (0, 1), yticks = 0:0.1:1,
        label = "Binomial($n, 0.5)", legend = :topleft)

plot(binomial_pdf.((1,2,4,8))...)

using StatsPlots

function simulation1(distribution, n = 250, k = 10_000)
    σ = std(distribution)
    y = rand(distribution, n, k)
    y .-= mean(distribution)
    y = mean(y, dims = 1)
    y = √n * vec(y)
    density(y, label = "Empirical Distribution")
    return plot!(Normal(0, σ), linestyle = :dash, color = :black,
                 label = "Normal(0.00, $(σ^2))")
end

simulation1(Exponential(0.5))

function simulation2(distribution = Beta(2, 2), n = 5, k = 10_000)
    y = rand(distribution, k, n)
    for col in 1:n
        y[:,col] += rand([-0.5, 0.6, -1.1], k)
    end
    y = (y .- mean(distribution)) ./ std(distribution)
    y = cumsum(y, dims = 2) ./ sqrt.(1:5)' # return grid of data
end

ys = simulation2()
plots = [] # would preallocate in optimized code
for i in 1:size(ys, 2)
    p = density(ys[:, i], linealpha = i, title = "n = $i")
    push!(plots, p)
end

plot(plots..., legend = false)

function exercise1(distribution = Uniform(0, π/2); n = 250, k = 10_000, g = sin, g′ = cos)
    μ, σ = mean(distribution), std(distribution)
    y = rand(distribution, n, k)
    y = mean(y, dims = 1)
    y = vec(y)
    error_obs = sqrt(n) .* (g.(y) .- g.(μ))
    density(error_obs, label = "Empirical Density")
    return plot!(Normal(0, g′(μ) .* σ), linestyle = :dash, label = "Asymptotic",
                 color = :black)
end
exercise1()

function exercise2(;n = 250, k = 50_000, dw = Uniform(-1, 1), du = Uniform(-2, 2))
    vw = var(dw)
    vu = var(du)
    Σ = [vw    vw
         vw vw + vu]
    Q = inv(sqrt(Σ))
    function generate_data(dw, du, n)
        dw = rand(dw, n)
        X = [dw dw + rand(du, n)]
        return sqrt(n) * mean(X, dims = 1)
    end
    X = mapreduce(x -> generate_data(dw, du, n), vcat, 1:k)
    X = Q * X'
    X = sum(abs2, X, dims = 1)
    X = vec(X)
    density(X, label = "", xlim = (0, 10))
    return plot!(Chisq(2), color = :black, linestyle = :dash,
                 label = "Chi-squared with 2 degrees of freedom", grid = false)
end
exercise2()