using LinearAlgebra

using GLMakie; GLMakie.activate!()
using AbstractPlotting

using Delaunay
using DEC2D

# convert vector of vectors into a matrix
vv2m(vv) = reduce(vcat, transpose.(vv))

# Simples case
begin
    vertices = vcat([[0, 0]], [[cos(2π * i / 6), sin(2π * i / 6)] for i = 1:6])
    edges = [
        [1, 2],
        [1, 3],
        [1, 4],
        [1, 5],
        [1, 6],
        [1, 7],
        [2, 3],
        [3, 4],
        [4, 5],
        [5, 6],
        [6, 7],
        [7, 2],
    ]
    triangles = [[1, 2, 3], [1, 3, 4], [1, 4, 5], [1, 5, 6], [1, 6, 7], [1, 7, 2]]

    # return the mesh and the indices of the interior vertices
    vint, mesh = DECMesh(vertices, edges, triangles)
    vext = setdiff(1:length(vertices), vint)

    # d₀ = boundary1(mesh)
    # d₁ = boundary2d(mesh)
    # *₁ = hodgestar1(mesh)
    # *₂ = hodgestar2(mesh)
    # Δ = *₂ * d₁ * *₁ * d₀
    Δ = laplacian(mesh) # calculate the laplacian operator

    b = [-1 - 10 * sum(Δ[i, vext]) for i = 1:length(vint)]

    # numerical solution
    zf = Δ[vint, vint] \ b

    x = vv2m(vertices)[:, 1]
    y = vv2m(vertices)[:, 2]
    zn = similar(vertices, Float64)
    zn[vint] = zf
    zn[vext] .= 10

    # exact solution
    f(x, y) = (1 - x^2 - y^2) / 4 + 10
    z = [f(v...) for v in vertices]

    println(
        "La norma de la diferencia entre la solución numérica y la exacta es ",
        norm(zn - z),
    )

    #   ____  _       _   _   _
    #  |  _ \| | ___ | |_| |_(_)_ __   __ _
    #  | |_) | |/ _ \| __| __| | '_ \ / _` |
    #  |  __/| | (_) | |_| |_| | | | | (_| |
    #  |_|   |_|\___/ \__|\__|_|_| |_|\__, |
    #                                 |___/

    # plot numerical solution
    scene = Scene()
    mesh!(
        scene,
        hcat(x, y, zn),
        vv2m(mesh.tri.triangles),
        color = zn,
        scale_plot = false,
    )
    xlims!(scene, -1, 1)
    ylims!(scene, -1, 1)
    zlims!(scene, 10, 10.25)
    wireframe!(scene[end][1], color = (:black, 0.6), linewidth = 2)
    display(scene)
    sleep(5.0)
end

# Caso medio
begin
    # we get the vertices of a more fine mesh
    vertices = [
        [cos(2π * i / (50 * r + 1)), sin(2π * i / (50 * r + 1))] .* r for r = 0:0.1:1
        for i = 1:round(Int, 50 * r + 1, RoundDown)
    ]

    # get thetriangulation using Delaunay algorithm
    meshd = delaunay(vv2m(vertices))
    triangles = [meshd.simplices[i, :] for i = 1:size(meshd.simplices, 1)]
    edges = reduce(
        union,
        [[sort([t[1], t[2]]), sort([t[2], t[3]]), sort([t[3], t[1]])] for t in triangles],
    )

    vint, mesh = DECMesh(vertices, edges, triangles)
    vext = setdiff(1:length(vertices), vint)

    Δ = laplacian(mesh)

    b = [-1 - 10 * sum(Δ[i, vext]) for i = 1:length(vint)]
    zf = Δ[vint, vint] \ b

    x = vv2m(vertices)[:, 1]
    y = vv2m(vertices)[:, 2]
    zn = similar(vertices, Float64)
    zn[vint] = zf
    zn[vext] .= 10

    f(x, y) = (1 - x^2 - y^2) / 4 + 10
    z = [f(v...) for v in vertices]

    println(
        "La norma de la diferencia entre la solución numérica y la exacta es ",
        norm(zn - z),
    )

    scene = Scene()
    mesh!(
        scene,
        hcat(x, y, zn),
        vv2m(mesh.tri.triangles),
        color = zn,
        scale_plot = false,
    )
    xlims!(scene, -1, 1)
    ylims!(scene, -1, 1)
    zlims!(scene, 10, 10.25)
    wireframe!(scene[end][1], color = (:black, 0.6), linewidth = 2)
    display(scene)
    sleep(5.0)
end

# Caso denso
begin
    # we get the vertices of a more fine mesh
    vertices = [
        [cos(2π * i / (200 * r + 1)), sin(2π * i / (200 * r + 1))] .* r for r = 0:0.025:1 for i = 1:round(Int, 200 * r + 1, RoundDown)
    ]

    # get thetriangulation using Delaunay algorithm
    meshd = delaunay(vv2m(vertices))
    triangles = [meshd.simplices[i, :] for i = 1:size(meshd.simplices, 1)]
    edges = reduce(
        union,
        [[sort([t[1], t[2]]), sort([t[2], t[3]]), sort([t[3], t[1]])] for t in triangles],
    )

    vint, mesh = DECMesh(vertices, edges, triangles)
    vext = setdiff(1:length(vertices), vint)

    Δ = laplacian(mesh)

    b = [-1 - 10 * sum(Δ[i, vext]) for i = 1:length(vint)]
    zf = Δ[vint, vint] \ b

    x = vv2m(vertices)[:, 1]
    y = vv2m(vertices)[:, 2]
    zn = similar(vertices, Float64)
    zn[vint] = zf
    zn[vext] .= 10

    f(x, y) = (1 - x^2 - y^2) / 4 + 10
    z = [f(v...) for v in vertices]

    println(
        "La norma de la diferencia entre la solución numérica y la exacta es ",
        norm(zn - z),
    )

    scene = Scene()
    mesh!(
        scene,
        hcat(x, y, zn),
        vv2m(mesh.tri.triangles),
        color = zn,
        scale_plot = false,
    )
    xlims!(scene, -1, 1)
    ylims!(scene, -1, 1)
    zlims!(scene, 10, 10.25)
    wireframe!(scene[end][1], color = (:black, 0.6), linewidth = 2)
    display(scene)
end


using PyCall
@pyimport matplotlib.pyplot as plt

using Gen

#import Gen: Distribution, logpdf

# Example from Section 4 of Reversible jump Markov chain Monte Carlo
# computation and Bayesian model determination

########################
# custom distributions #
########################

# minimum of k draws from uniform_continuous(lower, upper)

# we can sequentially sample the order statistics of a collection of K uniform
# continuous samples on the interval [a, b], by:
# x1 ~ min_uniform_continuous(a, b, K)
# x2 | x1 ~ min_uniform_continuous(x1, b, K-1)
# ..
# xK | x1 .. x_{K-1} ~ min_uniform_continuous(x_{K-1}, b, 1)

struct MinUniformContinuous <: Distribution{Float64} end
const min_uniform_continuous = MinUniformContinuous()

function Gen.logpdf(::MinUniformContinuous, x::Float64, lower::T, upper::U, k::Int) where {T<:Real,U<:Real}
    if x > lower && x < upper
        (k-1) * log(upper - x) + log(k) - k * log(upper - lower)
    else
        -Inf
    end
end

function Gen.random(::MinUniformContinuous, lower::T, upper::U, k::Int) where {T<:Real,U<:Real}
    # inverse CDF method
    p = rand()
    upper - (upper - lower) * (1. - p)^(1. / k)
end

# piecewise homogenous Poisson process

# n intervals - n + 1 bounds
# (b_1, b_2]
# (b_2, b_3]
# ..
# (b_n, b_{n+1}]

function compute_total(bounds, rates)
    num_intervals = length(rates)
    if length(bounds) != num_intervals + 1
        error("Number of bounds does not match number of rates")
    end
    total = 0.
    bounds_ascending = true
    for i=1:num_intervals
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        len = upper - lower
        if len <= 0
            bounds_ascending = false
        end
        total += len * rate
    end
    (total, bounds_ascending)
end

struct PiecewiseHomogenousPoissonProcess <: Distribution{Vector{Float64}} end
const piecewise_poisson_process = PiecewiseHomogenousPoissonProcess()

function Gen.logpdf(::PiecewiseHomogenousPoissonProcess, x::Vector{Float64}, bounds::Vector{Float64}, rates::Vector{Float64})
    cur = 1
    upper = bounds[cur+1]
    lpdf = 0.
    for xi in sort(x)
        if xi < bounds[1] || xi > bounds[end]
            error("x ($xi) lies outside of interval")
        end
        while xi > upper
            cur += 1
            upper = bounds[cur+1]
        end
        lpdf += log(rates[cur])
    end
    (total, bounds_ascending) = compute_total(bounds, rates)
    if bounds_ascending
        lpdf - total
    else
        -Inf
    end
end

function Gen.random(::PiecewiseHomogenousPoissonProcess, bounds::Vector{Float64}, rates::Vector{Float64})
    x = Vector{Float64}()
    num_intervals = length(rates)
    for i=1:num_intervals
        lower = bounds[i]
        upper = bounds[i+1]
        rate = (upper - lower) * rates[i]
        n = random(poisson, rate)
        for j=1:n
            push!(x, random(uniform_continuous, lower, upper))
        end
    end
    x
end

#########
# model #
#########

@gen function model(T::Float64)

    # prior on number of change points
    k = @addr(poisson(3.), "k")

    # prior on the location of (sorted) change points
    change_pts = Vector{Float64}(undef, k)
    lower = 0.
    for i=1:k
        cp = @addr(min_uniform_continuous(lower, T, k-i+1), "cp$i")
        change_pts[i] = cp
        lower = cp
    end

    # k + 1 rate values
    # h$i is the rate for cp$(i-1) to cp$i where cp0 := 0 and where cp$(k+1) := T
    alpha = 1.
    beta = 200.
    rates = Float64[@addr(Gen.gamma(alpha, 1. / beta), "h$i") for i=1:k+1]

    # poisson process
    bounds = vcat([0.], change_pts, [T])
    @addr(piecewise_poisson_process(bounds, rates), "points")
end

function render(trace; ymax=0.02)
    T = get_call_record(trace).args[1]
    assignment = get_assignment(trace)
    k = assignment["k"]
    bounds = vcat([0.], sort([assignment["cp$i"] for i=1:k]), [T])
    rates = [assignment["h$i"] for i=1:k+1]
    for i=1:length(rates)
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        plt.plot([lower, upper], [rate, rate], color="black", linewidth=2)
    end
    points = assignment["points"]
    plt.scatter(points, -rand(length(points)) * (ymax/5.), color="black", s=5)
    ax = plt.gca()
    xlim = [0., T]
    plt.plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
end

function show_prior_samples()
    plt.figure(figsize=(16,16))
    T = 40000.
    for i=1:16
        plt.subplot(4, 4, i)
        trace = simulate(model, (T,))
        render(trace; ymax=0.015)
    end
    plt.tight_layout(pad=0)
    plt.savefig("prior_samples.pdf")
end

#############################
# height and position moves #
#############################

@gen function height_proposal(prev, i::Int)
    prev_assignment = get_assignment(prev)
    height = prev_assignment["h$i"]
    @addr(uniform_continuous(height/2., height*2.), "h$i")
end

@gen function position_proposal(prev, i::Int)
    prev_assignment = get_assignment(prev)
    k = prev_assignment["k"]
    lower = (i == 1) ? 0. : prev_assignment["cp$(i-1)"]
    upper = (i == k) ? T : prev_assignment["cp$(i+1)"]
    @addr(uniform_continuous(lower, upper), "cp$i")
end

function height_move(trace)
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k+1)
    mh(model, height_proposal, (i,), trace)
end

function position_move(trace)
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k)
    mh(model, position_proposal, (i,), trace)
end

######################
# birth / death move #
######################

# insert a new change point at i, where 1 <= i <= k+1
# the current change point at i, and all after, will be shifted right
# the new change point will be placed between cp$(i-1) and the current cp$i
@gen function birth_proposal(prev, T, i::Int)
    prev_assignment = get_assignment(prev)
    k = prev_assignment["k"]
    lower = (i == 1) ? 0. : prev_assignment["cp$(i-1)"]
    upper = (i == k+1) ? T : prev_assignment["cp$i"]
    @addr(uniform_continuous(lower, upper), "new-cp")
    @addr(uniform_continuous(0., 1.), "u")
end

@gen function death_proposal(prev) end

function birth_move_new_heights(cur_height, new_cp, prev_cp, next_cp, u)
    d_prev = new_cp - prev_cp
    d_next = next_cp - new_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_cur_height = log(cur_height)
    log_ratio = log(1 - u) - log(u)
    new_h_prev = exp(log_cur_height - (d_prev / d_total) * log_ratio)
    new_h_next = exp(log_cur_height + (d_next / d_total) * log_ratio)
    @assert new_h_prev > 0.
    @assert new_h_next > 0.
    (new_h_prev, new_h_next)
end

const MODEL = :model
const PROPOSAL = :proposal

@inj function birth_injection(T, i::Int)

    # increment k
    k = @read(MODEL => "k")
    @write(k+1, MODEL => "k")

    # changepoints
    for j=1:i-1
        @copy(MODEL => "cp$j", MODEL => "cp$j")
    end
    @copy(PROPOSAL => "new-cp", MODEL => "cp$i")
    for j=i:k
        @copy(MODEL => "cp$j", MODEL => "cp$(j+1)")
    end

    # compute new heights
    cur_height = @read(MODEL => "h$i")
    prev_cp = (i == 1) ? 0. : @read(MODEL => "cp$(i-1)")
    next_cp = (i == k+1) ? T : @read(MODEL => "cp$i")
    new_cp = @read(PROPOSAL => "new-cp")
    u = @read(PROPOSAL => "u")
    (new_h_prev, new_h_next) = birth_move_new_heights(cur_height, new_cp, prev_cp, next_cp, u)

    # heights
    for j=1:i-1
        @copy(MODEL => "h$j", MODEL => "h$j")
    end
    @write(new_h_prev, MODEL => "h$i")
    @write(new_h_next, MODEL => "h$(i+1)")
    for j=i+1:k+1
        @copy(MODEL => "h$j", MODEL => "h$(j+1)")
    end

    @copy(MODEL => "points", MODEL => "points")
end

function death_move_u_new_height(prev_height, next_height, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_prev_height = log(prev_height)
    log_next_height = log(next_height)
    new_height = exp((d_prev / d_total) * log_prev_height + (d_next / d_total) * log_next_height)
    u = prev_height / (prev_height + next_height)
    @assert new_height > 0.
    (new_height, u)
end

@inj function death_injection(T, i::Int)

    # decrement k
    k = @read(MODEL => "k")
    @assert k > 0
    @write(k-1, MODEL => "k")

    # change points
    for j=1:i-1
        @copy(MODEL => "cp$j", MODEL => "cp$j")
    end
    for j=i+1:k
        @copy(MODEL => "cp$j", MODEL => "cp$(j-1)")
    end

    # compute new height
    cur_cp = @read(MODEL => "cp$i")
    prev_cp = (i == 1) ? 0. : @read(MODEL => "cp$(i-1)")
    next_cp = (i == k) ? T : @read(MODEL => "cp$(i+1)")
    prev_height = @read(MODEL => "h$i")
    next_height = @read(MODEL => "h$(i+1)")
    (new_height, u) = death_move_u_new_height(prev_height, next_height, cur_cp, prev_cp, next_cp)

    # heights
    for j=1:i-1
        @copy(MODEL => "h$j", MODEL => "h$j")
    end
    @write(new_height, MODEL => "h$i")
    for j=i+2:k+1
        @copy(MODEL => "h$j", MODEL => "h$(j-1)")
    end

    @copy(MODEL => "points", MODEL => "points")

    @copy(MODEL => "cp$i", PROPOSAL => "new-cp")
    @write(u, PROPOSAL => "u")
end

function birth_move(trace)
    # if k > 0, then prob_b = 0.25 and prob_d = 0.25
    # the probability that the one we introduce will get deleted, given that
    # we choose a death move, is: 1/(k+1); which is also the probability that
    # we choose the i that we choose here
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k+1)
    # prob_b = 1, but prob_d = 0.25, so we correct by -log(4)
    correction = (new_trace) -> (k == 0 ? -log(4) : 0)
    T = get_call_record(trace).args[1]
    rjmcmc(model,
        birth_proposal, (T, i),
        death_proposal, (),
        birth_injection, (T, i),
        trace, correction)
end

function death_move(trace)
    k = get_assignment(trace)["k"]
    @assert k > 0
    i = random(uniform_discrete, 1, k)
    correction = (new_trace) -> (k == 0 ? log(4) : 0)
    T = get_call_record(trace).args[1]
    rjmcmc(model,
        death_proposal, (),
        birth_proposal, (T, i),
        death_injection, (T, i),
        trace, correction)
end

##########################
# Generic MCMC inference #
##########################

function resimulation_mh(selection, trace)
    model_args = get_call_record(trace).args
    (new_trace, weight) = regenerate(model, model_args, NoChange(), trace, selection)
    if log(rand()) < weight
        # accept
        return new_trace
    else
        # reject
        return trace
    end
end

k_selection = DynamicAddressSet()
push_leaf_node!(k_selection, "k")

function generic_mcmc_step(trace)
    k = get_assignment(trace)["k"]
    if k > 0
        prob_h = 1./3
        prob_p = 1./3
        prob_change_k = 1./3
    else
        prob_h = 0.
        prob_p = 0.
        prob_change_k = 1
    end
    move_type = random(categorical, [prob_h, prob_p, prob_change_k])
    if move_type == 1
        height_move(trace)
    elseif move_type == 2
        position_move(trace)
    else
        resimulation_mh(k_selection, trace)
    end
end

#########################
# RJMCMC MCMC inference #
#########################0

function mcmc_step(trace)
    k = get_assignment(trace)["k"]
    if k > 0
        prob_h = 0.25
        prob_p = 0.25
        prob_b = 0.25
        prob_d = 0.25
    else
        prob_h = 0.
        prob_p = 0.
        prob_b = 1.
        prob_d = 0.
    end
    move_type = random(categorical, [prob_h, prob_p, prob_b, prob_d])
    if move_type == 1
        height_move(trace)
    elseif move_type == 2
        position_move(trace)
    elseif move_type == 3
        birth_move(trace)
    elseif move_type == 4
        death_move(trace)
    else
        error("Unknown move type $move_type")
    end
end

function do_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $(get_assignment(trace)["k"])")
        end
        #trace = mcmc_step(trace)
        trace = generic_mcmc_step(trace)
    end
    trace
end

########################
# inference experiment #
########################

Gen.load_generated_functions()

import Random
Random.seed!(1)

# load data set
import CSV
function load_data_set()
    df = CSV.read("coal.csv")
    dates = df[1]
    dates = dates .- minimum(dates)
    dates * 365.25 # convert years to days
end

const points = load_data_set()
const T = maximum(points)
const observations = DynamicAssignment()
observations["points"] = points

function show_posterior_samples()
    plt.figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        tic()
        plt.subplot(4, 4, i)
        trace = do_mcmc(T, 5000)#10000)
        toc()
        render(trace; ymax=0.015)
    end
    plt.tight_layout(pad=0)
    plt.savefig("posterior_samples.pdf")
end

function get_rate_vector(trace, test_points)
    assignment = get_assignment(trace)
    k = assignment["k"]
    cps = [assignment["cp$i"] for i=1:k]
    hs = [assignment["h$i"] for i=1:k+1]
    rate = Vector{Float64}()
    cur_h_idx = 1
    cur_h = hs[cur_h_idx]
    next_cp_idx = 1
    upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
    for x in test_points
        while x > upper
            next_cp_idx += 1
            upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
            cur_h_idx += 1
            cur_h = hs[cur_h_idx]
        end
        push!(rate, cur_h)
    end
    rate
end

# compute posterior mean rate curve

function plot_posterior_mean_rate()
    test_points = collect(1.0:10.0:T)
    rates = Vector{Vector{Float64}}()
    num_samples = 0
    num_steps = 5000 # 20000
    for reps=1:10
        (trace, _) = generate(model, (T,), observations)
        for iter=1:num_steps
            if iter % 1000 == 0
                println("iter $iter of $num_steps, k: $(get_assignment(trace)["k"])")
            end
            trace = mcmc_step(trace)
            if iter > 4000
                num_samples += 1
                rate_vector = get_rate_vector(trace, test_points)
                @assert length(rate_vector) == length(test_points)
                push!(rates, rate_vector)
            end
        end
    end
    posterior_mean_rate = zeros(length(test_points))
    for rate in rates
        posterior_mean_rate += rate / Float64(num_samples)
    end
    ymax = 0.010
    plt.figure()
    plt.plot(test_points, posterior_mean_rate, color="black")
    plt.scatter(points, -rand(length(points)) * (ymax/6.), color="black", s=5)
    ax = plt.gca()
    xlim = [0., T]
    plt.plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
    plt.savefig("posterior_mean_rate.pdf")
end

function plot_trace_plot()
    # show the number of clusters
    (trace, _) = generate(model, (T,), observations)
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        (trace, accept) = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.figure()
    plt.plot(num_clusters_vec)
    ax = plt.gca()
    plt.savefig("trace_plot_rjmcmc.pdf")
end

function plot_trace_plot()
    plt.figure(figsize=(8, 4))

    # generic
    (trace, _) = generate(model, (T,), observations)
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        trace = generic_mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.subplot(2, 1, 1)
    plt.plot(num_clusters_vec, "r")

    # reversible jump
    (trace, _) = generate(model, (T,), observations)
    height1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        trace = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.subplot(2, 1, 2)
    plt.plot(num_clusters_vec, "b")

    ax = plt.gca()
    plt.savefig("trace_plot.pdf")
end

println("showing prior samples...")
show_prior_samples()

println("showing posterior samples...")
show_posterior_samples()

println("estimating posterior mean rate...")
plot_posterior_mean_rate()

println("making trace plot...")
plot_trace_plot()


"""
	MMCK(μ, λ, c, k)

Tạo mô hình M/M/c/k
"""
struct MMCK{T} <:AbstractMMCK
	λ::T
	μ::T
	c::Union{T, Integer}
	k::Union{T, Integer}
	r::Union{T, Real}
	ρ::Union{T, Real}
	function MMCK(λ, μ, c, k)
		T = Union{typeof(μ), typeof(λ)}
		new{T}(λ, μ, c,	k, λ/μ,	λ/μ/c)
	end
end

function pn(m::MMCK, n::Integer)
	r = m.r
	k = m.k
	c = m.c
	if n == 0
		1 / (
		     sum(r^i / factorial(i) for i = 0:c-1) +
		     sum(r^i / (c^(i-c) * factorial(c)) for i = c:k)
		     )
	elseif n < c
		pn(m, 0) * r^n / factorial(n)
	elseif n <= k
		pn(m, 0) * r^n / (c^(n-c) * factorial(c))
	else 0
	end
end

function L(m::MMCK)
	sum(i * pn(m, i) for i=0:m.k)
end

function Lq(m::MMCK)
	sum((i - m.c) * pn(m, i) for i=m.c:m.k)
end

function W(m::MMCK)
	L(m) / λe(m)
end

function Wq(m::MMCK)
	Lq(m) / λe(m)
end


# This file is a part of Julia. License is MIT: http://julialang.org/license

# https://github.com/JuliaLang/julia/issues/4707

#----------------- Get the A matrix
function getDivGrad(n1,n2,n3)
        D1 = kron(speye(n3),kron(speye(n2),ddx(n1)))
        D2 = kron(speye(n3),kron(ddx(n2),speye(n1)))
        D3 = kron(ddx(n3),kron(speye(n2),speye(n1)))

        # DIV from faces to cell-centers
        Div = [D1 D2 D3]
        return Div*Div';
end

#----------------- 1D finite difference on staggered grid
function ddx(n)
    return spdiags(ones(n)*[-1 1],[0,1],n,n+1)
end

#------------- Build a diagonal matrix
function spdiags(B,d,m,n)
    d = d[:]
    p = length(d)

    len = zeros(Int, p+1, 1)
    for k = 1:p
        len[k+1] = len[k] + length(max(1,1-d[k]):min(m,n-d[k]))
    end

    a = zeros(Int, len[p+1], 3)
    for k = 1:p
        # Append new d[k]-th diagonal to compact form
        i = max(1,1-d[k]):min(m,n-d[k])
        a[(len[k]+1):len[k+1],:] = [i i+d[k] B[i+(m>=n)*d[k],k]]
    end

    return sparse(a[:,1],a[:,2],a[:,3],m,n)
end


export Zernike,
       zernike_norm,
       zernike_nm_to_fringe,
       zernike_nm_to_ansi_j,
       zernike_ansi_j_to_nm,
       zernike_noll_to_nm,
       zernike_fringe_to_nm,
       zernike_zero_separation

"""
    Zernike([T, ]n, m, norm=true)

Zernike polynomial of orders `n` and `m`. If `norm` is true the output will be normalized. The output type can be specified with `T`, which will default to `Float64`.

# Examples

```jldoctest
julia> z = Zernike(10, 3)
Zernike{Float64}(n=10, m=3)

julia> z(0.8, 0.2)
-1.6860955965619917

julia> Zernike(BigFloat, 10, 3)(0.8, 0.2)
-1.68609559656199170518675600760616362094879150390625
```
"""
struct Zernike{T<:AbstractFloat,F1,F2}
    n::Int
    m::Int
    basis::F1
    trig_func::F2
    Zernike{T}(n, m, basis, trig_func) where {T} = new{T,typeof(basis),typeof(trig_func)}(n, m, basis, trig_func)
end

function Zernike(T, n, m, norm::Bool=true)
    norm_val = norm ? zernike_norm(n, m) : 1.0
    trig_func = m < 0 ? sin : cos
    basis = x -> norm_val * jacobi((n - m) ÷ 2, 0, abs(m), x)
    Zernike{T}(n, m, basis, trig_func)
end

Zernike(n, m, norm::Bool=true) = Zernike(Float64, n, m, norm)

function Base.show(io::IO, z::Zernike{T}) where T
    print(io, "Zernike{$T}(n=$(z.n), m=$(z.m))")
end

function (z::Zernike{T})(ρ, θ) where T
    x = ρ^2 - 1
    return T(ρ^abs(z.m) * z.basis(x) * z.trig_func(z.m * θ))
end

"""
    kronecker(i,j)

1 if i==j, else 0; mathematical kronecker function
"""
kronecker(i, j) = Int(i == j)

"""
    zernike_norm(n, m)

Norm of Zernike polynomial of radial order n, azimuthal order m.

The norm is the average squared distance to zero.  By multiplying a zernike
value by the norm, the term is given unit stdev or RMS.
"""
function zernike_norm(n, m)
    num = √(2 * (n+1)) / (1 + kronecker(m, 0))
end

"""
    zernike_nm_to_fringe(n, m)

Map (n,m) ANSI indices to a single fringe index.
"""
function zernike_nm_to_fringe(n, m)
    term1 = (1 + (n + abs(m))/2)^2
    term2 = 2*abs(m)
    term3 = (1 + sign(m)) / 2
    return int(term2 - term2 - term3) + 1
end

"""
    zernike_nm_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_ansi_j_to_nm`](@ref) (reciprocal of this function)
"""
function zernike_nm_to_ansi_j(n, m)
    return int((n * (n + 2) + m) / 2)
end

"""
    zernike_ansi_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_nm_to_ansi_j`](@ref) (reciprocal of this function)
"""
function zernike_ansi_j_to_nm(j)
    n = int(ceil((-3 + √(9 + 8j))/2))
    m = 2j - n * (n + 2)
    return n, m
end

"""
    zernike_noll_to_nm(j)

Map j Noll index to ANSI (n,m) indices.
"""
function zernike_noll_to_nm(j)
    n = int(ceil((-1 + √(1 + 8j))/2) - 1)
    if n == 0
        m = 0
    else
        nseries = int((n+1) * (n+2) / 2)
        residual = j - nseries - 1

        if isodd(j)
            sign = -1
        else
            sign = 1
        end

        if isodd(n)
            ms = [1,1]
        else
            ms = [0]
        end

        for i=0:n÷2
            push!(ms, ms[end]+2)
            push!(ms, ms[end])
        end

        m = ms[residual] * sign
    end
    return n, m
end

"""
    zernike_fringe_to_nm(j)

Map j Fringe index to ANSI (n,m) indices.
"""
function zernike_fringe_to_nm(j)
    m_n = 2 * ceil(√j - 1)
    g_s = (m_n / 2)^2 + 1
    n = m_n / 2 + floor((j-g_s)/2)
    m = m_n - n * (1 - mod(j-g_s, 2) * 2)
    return int(n), int(m)
end

"""
    zernike_zero_separation(n)

Minimum zero separation of Zernike polynomial of radial order n.  Useful for
computing sample count requirements.
"""
function zernike_zero_separation(n)
    return 1 / n^2
end

"""
    zernike(n, m, ρ, θ[; norm])

Zernike polynomial of radial order n and azimuthal order m, evaluated at the
point (ρ, θ).  No normalization is required of (ρ, θ), though the polynomials
are orthogonal only over the unit disk.

norm is a boolean flag indicating whether the result should be orthonormalized
(scaled to unit RMS) or not.
"""
function zernike(n, m, ρ, θ; norm::Bool=true)
    x = ρ^2 - 1
    n_j = (n - m) / 2
    am = abs(m)
    # α=0, β=|m|
    # there is a second syntax where you have x reversed, 1 - ρ^2,
    # in which ase you swap α and β.  It makes absolutely no difference
    out = jacobi(n_j, 0, am, x)
    if m != 0
        if sign(m) == -1
            f = sin
        else
            f = cos
        end
        out *= (ρ^am * f(m*θ))
    end
	if norm
		out *= zernike_norm(n,m)
	end
    return out
end


using DataFrames, GLM, Distributions, LinearAlgebra, Random
using Plots, LaTeXStrings;pyplot()
Random.seed!(0)

beta0, beta1 = 2.0, 1.5
sigma = 2.5
n, N = 10, 10^4
alpha = 0.05

function coefEst()
    xVals = collect(1:n)
    yVals = beta0 .+ beta1*xVals + rand(Normal(0,sigma),n)
    data = DataFrame([xVals,yVals],[:X,:Y])
    model = lm(@formula(Y ~ X), data)
    coef(model)
end

ests = [coefEst() for _ in 1:N]

xBar = mean(1:n)
sXX = sum([(x - xBar)^2 for x in 1:n])
sx2 = sum([x^2 for x in 1:n])
var0 = sigma^2 * sx2/(n*sXX)
var1 = sigma^2/sXX
cv = -sigma^2*xBar/sXX

mu = [beta0, beta1]
Sigma = [var0 cv; cv var1]

A = cholesky(Sigma).L
Ai = inv(A)

r = quantile(Rayleigh(),1-alpha)
isInEllipse(x) = norm(Ai*(x-mu)) <= r
estIn = isInEllipse.(ests)

println("Proportion of points inside ellipse: ", sum(estIn)/N)

scatter(first.(ests[estIn]),last.(ests[estIn]),c=:green, ms=2, msw=0)
scatter!(first.(ests[.!estIn]),last.(ests[.!estIn]),c=:blue, ms=2, msw=0)

ellipsePts = [r*A*[cos(t),sin(t)] + mu  for t in 0:0.01:2pi]
scatter!([beta0],[beta1],c=:red, ms=5, msw = 0)
plot!(first.(ellipsePts),last.(ellipsePts),
	c=:red, lw=2, legend=:none,
	xlabel=L"\hat{\beta}_0", ylabel=L"\hat{\beta}_1")


#Runs a simple Metropolis-Hastings (ie MCMC) algoritm to simulate an
#exponential distribution, which has the probability density
#p(x)=exp(-x/m^2), where m>0.
#
#Author: H. Paul Keeler, 2019.
# Website: hpaulkeeler.com
# Repository: github.com/hpaulkeeler/posts

#Note: Need the .+ for adding a scalar to an array
#Also need . for sqrt, exp, cos, sin etc and assinging scalars to arrays

#clearconsole(); #for clearing Julia REPL console

using Distributions #for random simulations
using PyPlot #for plotting
using StatsBase #for histograms etc
using Random
using LinearAlgebra
PyPlot.close("all");  # close all PyPlot figures

#set random seed for reproducibility
#Random.seed!(1234)

numbSim=10^4; #number of random variables simulated
numbSteps=25; #number of steps for the Markov process
numbBins=50; #number of bins for histogram

sigma=1; #standard deviation for normal random steps
m=2; #parameter (ie mean) for distribution to be simulated

function fun_p(x)
    return ((exp.(-x./m)./m).*(x.>0));
end

xRand=rand(numbSim); #random intial values
probCurrent=fun_p(xRand); #current transition probabilities

for jj=1:numbSteps
    zRand= xRand.+sigma.*rand(Normal(),numbSim);#take a (normally distributed) random step
    #zRand= xRand +2*sigma*(rand(size(xRand))-0.5);#take a (uniformly distributed) random step
    probProposal=fun_p(zRand); #proposed probability

    #acceptance rejection step
    booleAccept=rand(numbSim) .< probProposal./probCurrent;
    #update state of random walk/Markov chain
    xRand[booleAccept]=zRand[booleAccept];
    #update transition probabilities
    probCurrent[booleAccept]=probProposal[booleAccept];

end

#histogram section: empirical probability density
histX=fit(Histogram, xRand,nbins=numbBins); #find histogram data
histX=normalize(histX,mode=:pdf); #normalize histogram
binEdges=histX.edges; #retrieve bin edges
xValues=(binEdges[1][2:end]+binEdges[1][1:end-1])./2; #mid-points of bins
pdfEmp=(histX.weights)
#analytic solution of probability density
pdfExact=fun_p(xValues);

# Plotting
PyPlot.plot(xValues, pdfExact)
PyPlot.scatter(xValues, pdfEmp, marker="x", c="r");
PyPlot.grid(true);
PyPlot.xlabel("x");
PyPlot.ylabel("p(x)");
PyPlot.show();


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ f6d978ea-b349-11eb-309e-f5a517efc907
begin
	using LinearAlgebra
	using Plots
	using QuadGK
end

# ╔═╡ 99268461-9e3f-4685-b530-e24febfb539e
#=
IDEA OF MULTIGRID :

	G is a mesh of a 1D domain.
	G1 is another, but coarser.
	G1's operator has higher frequency eigenvalues as opposed to  G.
	Use G1 to weed out these values.
	Use smoothers alternatively.
=#

# ╔═╡ b422b21d-b9d3-4aad-b29b-d0669331abc8

begin
	e = 5; #  Tune e for different basis functions
	N = 10;
	Ni = 100;
	xmin = 0;
	xmax = 1;
	x = LinRange(xmin,xmax,N);
	xi = LinRange(x[e-1],x[e+1],Ni);
	h = x[2] - x[1];
end

# ╔═╡ b11596d4-ee5d-41ec-af8f-96734ba2f579
# This shows a heirarchy of grids
grid1 = plot(x,zeros(length(x)),marker="*")

# ╔═╡ c1463b68-12e6-4e8c-9902-c2cc7e6caa91
begin
	# Plot the basis functions
	y = zeros(length(xi));
	for i = 1:length(xi);
		if(xi[i]>=xi[1] && xi[i]<=x[e])
			y[i] = 1/h * (xi[i]-x[e-1]);
		elseif (xi[i]>x[e] && xi[i]<x[e+1])
			y[i] = 1/h * (x[e+1]-xi[i]);
		end
	end
	plot!(xi,y,title="Linear Basis Functions");
end

# ╔═╡ d5da54e6-0a4f-4f43-a7ab-8791574b32c6
# Now we use the linear basis functions and obtain an approximation of the Laplacian operator.

begin
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function θ(x,i,N)
		n = length(x);
		y = zeros(n);
		if i==1
			for k=1:N
				if x[k]<= x[i+1]
					y[k] = (N-1)*(x[i+1]-x[k]);
				end
			end
		elseif i==N
			for l = 1:N
				if x[l] >= x[i-1]
					y[l] = (N-1) * (x[l]-x[i]);
				end
			end
		else
			for j=2:n-1
				if(x[j]>=x[i-1] && x[j]<=x[i])
					y[j] = (N-1) * (x[j]-x[i-1]);
				elseif(x[j]>=x[i] && x[j]<=x[i+1])
					y[j] = (N-1) * (x[i+1]-x[j]);
				end
			end
		end
		return y;
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function θ_scalar(x,i,xi,N)
			if(x>=xi[i-1] && x<=xi[i])
				y = (N-1) * (x-xi[i-1]);
			elseif(x>=xi[i] && x<=xi[i+1])
				y = (N-1) * (xi[i+1]-x);
			else
				y=0;
			end
		return y;
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function stiffness_matrix(n)
		md = 2.0*ones(n-2);
		ld = -1.0*ones(n-3);
		ud = -1.0*ones(n-3);
		return (n-1)*Tridiagonal(ld,md,ud);
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function intergrate(x,i,N)
		domain = LinRange(0,1,N+2);
		return sin(pi*x) * θ_scalar(x,i,domain,N);
	end

	function compute_rhs(n)
		n = n-2;
		integral = zeros(n);
		for i=1:length(integral)
		integral[i] = 1.5*quadgk(q->intergrate(q,i+1,n),0,1,rtol=1e-8)[1];
		end
		return integral;
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
end

# ╔═╡ a5f1a911-3f71-4630-a872-99514d52e5f3
begin
	nk = 12;
	a = LinRange(xmin,xmax,nk);
	z = θ(a,3,nk);
	plot(a,z)
end

# ╔═╡ feba23c9-1b66-400e-9ffb-0dc5fc65896f
begin
	# Get the stiffness matrix and the RHS : The RHS is evaluated and plotted.
	n = 7;
	test_A = stiffness_matrix(n)
	test_b = compute_rhs(n);
	plot(test_b)
	#test_A
	#test_b
end

# ╔═╡ b8b7b670-6039-4bd9-9f95-8652f6e03528
md"""

	#NESTED MULTIGRID

	Start on a coarse grid h_0

	Iterate using a smoother and remove one range of high frequencies.

	Interpolate solution to a finer grid - h_1

	Start step 1 with a smaller time step h_1.
"""

# ╔═╡ a217bef6-3c9f-4d12-8b3f-eb937987c709
# Description

md"""

To implement a nested multigrid method.

We would need access to a smoother. The kind of smoother to choose depends on the nature of the matrix A that we solve.

For the finite element approximation of the poisson equation we have a matrix A that is not diagonally dominant.

So solve the problem we need a method for which convergence could be proven.

In comes the Gauss Seidel method.

"""

# ╔═╡ f2a7eb59-3617-47fe-be19-f6873d0f3c56
function gauss_seidel_relaxation(A,b,x0,max_iterations)
	x = x0;
	res = b - A*x;
	for iter = 1:max_iterations
		for i = 1:length(x)
			x[i] = (1/A[i,i])*(b[i] - dot(A[i,:],x) + A[i,i]*x[i]);
		end
		res = b - A*x;
	end
	return x
end

# ╔═╡ abf72f94-406b-41a6-bbdd-fa3ac2919531
# To perform a nested iteration
begin
	n_nodes = 7;
	A = stiffness_matrix(n_nodes);
	b = compute_rhs(n_nodes);
	x0 = zeros(n_nodes-2);

	y_appx = gauss_seidel_relaxation(A,b,x0,10);
	y_appx = vcat(0,y_appx,0);
	println(y_appx)

	x_nodes = LinRange(0,1,100);
	y_actual = zeros(length(x_nodes));
	for i=1:length(x_nodes)
		y_actual[i] = 1/pi^2 * sin(pi*x_nodes[i]);
	end
	plot(x_nodes,y_actual,label="Actual Solution");
	plot!(LinRange(0,1,length(y_appx)),y_appx,label="Approximate solution");
end

# ╔═╡ Cell order:
# ╠═f6d978ea-b349-11eb-309e-f5a517efc907
# ╠═99268461-9e3f-4685-b530-e24febfb539e
# ╠═b422b21d-b9d3-4aad-b29b-d0669331abc8
# ╠═b11596d4-ee5d-41ec-af8f-96734ba2f579
# ╠═c1463b68-12e6-4e8c-9902-c2cc7e6caa91
# ╠═d5da54e6-0a4f-4f43-a7ab-8791574b32c6
# ╠═a5f1a911-3f71-4630-a872-99514d52e5f3
# ╠═feba23c9-1b66-400e-9ffb-0dc5fc65896f
# ╟─b8b7b670-6039-4bd9-9f95-8652f6e03528
# ╟─a217bef6-3c9f-4d12-8b3f-eb937987c709
# ╠═f2a7eb59-3617-47fe-be19-f6873d0f3c56
# ╠═abf72f94-406b-41a6-bbdd-fa3ac2919531


# Some built-in functions
println("println is a function")
println("parse(Int64, \"123\"): $(parse(Int64, "123"))")
println("parse(Float64, \"123.456\"): $(parse(Float64, "123.456"))")
println("trunc(Int64, 3.67123): $(trunc(Int64, 3.67123))")
println("trunc(Int64, -3.67123): $(trunc(Int64, -3.67123))")
println("Other examples: ", float(32), "\t", string(32), "\t", string(123.456))
println("Natural log and base, log(exp(1)): $(log(exp(1)))")
println("Other log bases, log2(10), log10(10): $(log2(10)), $(log10(10))")
println("Trigonometry, sin(pi), cos(pi), tan(pi): $(sin(pi)), $(cos(pi)), $(tan(pi))")
println("sqrt(2) = $(sqrt(2))")

# Defning functions
function printlines()
    println("This is the first line.")
    println("This is the second line.")
end

# printlines()

function printmorelines()
    printlines()
    printlines()
end

# printmorelines()

function printargument(arg)
    println(arg)
end

printargument(123)
printargument("123456")
printargument(cos(pi/2))

# Locality of variables and parameters
function sumthree(a, b, c)
    a = a + b
    total = a
    total += c
    return total
end
d, e, f = 1, 2, 3
println("sumthree(d, e, f): $(sumthree(d, e, f))")
println("d, e, f: $d, $e, $f")
# These lines will cause an error:
# println(a, b, c)
# println(total)

# Example of void function that returns nothing:
println(printlines())

# Example of a docstring
"""
polyline(t, n, len, angle)

Draws n line segments with the given length and
angle (in degrees) between them.  t is a turtle.
"""
function polyline(t, n, len, angle)
end

# Example of including errors
function fact(n)
    if !(n isa Int64)
        error("Factorial is only defined for integers.")
    elseif n < 0
        error("Factorial is not defined for negative integers.")
    elseif n == 0
        return 1
    else
        return n * fact(n-1)
    end
end

# println(fact(1.5))
# println(fact(-1))

# Shorthand notation for function definition
function func_ex(a=1, b=2)
    return a + 2b
end

# The above function definition is equivalent to the following:
func(a, b) = a + 2b
func(a) = func(a, 2)
func() = func(1, 2)
println("func(7, 5): $(func(7, 5))")
println("func(2): $(func(2))")
println("func(): $(func())")


# p13.jl - solve linear BVP u_xx = exp(4x), u(-1)=u(1)=0

N = 16;
(D,x) = cheb(N);
D2 = D^2;
D2 = D2[2:N,2:N];                   # boundary conditions
f = @. exp(4*x[2:N]);
u = D2\f;                           # Poisson eq. solved here
u = [0;u;0];
clf(); axes([.1,.4,.8,.5]);
plot(x,u,".",markersize=6);
xx = -1:.01:1;
uu = polyval(polyfit(x,u),xx);      # interpolate grid data
plot(xx,uu);
grid(true);
exact = @. ( exp(4*xx) - sinh(4)*xx - cosh(4) )/16;
title("max err = $(signif(norm(uu-exact,Inf),4))",fontsize=12);


# This file calculates scattering from the gravity Helmholtz equation via the adaptive spectral method.
# Δu + (E+x_2)u = 0,
# u|Γ = 0,
# u^i = Φ(⋅,⋅,E)|Γ
# u = u^i + u^s.
# The normal derivative ∂u/∂n of the entire wave is calculated on the sound-soft boundaries.
# The scattered wave is calculated by convolving ∂u/∂n with the fundamental solution.
# Then, the total wave is obtained by summing the incident and the scattered waves.

using ApproxFun, SingularIntegralEquations
include("Scatteraux.jl")

E = 20.
ω = 2π
ui = (x,y) ->  lhelmfs(complex(x,y),-5.0im,E)

# The gravity Helmholtz Green's function.
g3 = (x,y) ->  lhelmfs(x,y,E)
r = (x,y) ->  lhelm_riemann(x,y,E)

dom = ∪(Segment.([-10.0-3.0im,5.0,-2+5im],[-5.0+0.0im,10.0-3im,2+5im]))
sp = Space(dom)
cwsp = CauchyWeight(sp⊗sp,0)
uiΓ,⨍ = Fun(t->ui(real(t),imag(t)),sp),DefiniteLineIntegral(dom)

@time G = GreensFun(g3,cwsp;method=:unsplit)

@time ∂u∂n = ⨍[G]\uiΓ
println("The length of ∂u∂n is: ",ncoefficients(∂u∂n))
us = (x,y) ->  -linesum(g3,∂u∂n,complex(x,y))


# Local Tangent Space Alignment (LTSA)
# ---------------------------
# Principal Manifolds and Nonlinear Dimension Reduction via Local Tangent Space Alignment,
# Zhang, Zhenyue; Hongyuan Zha (2004),  SIAM Journal on Scientific Computing 26 (1): 313–338.
# doi:10.1137/s1064827502419154.

"""
    LTSA{NN <: AbstractNearestNeighbors, T <: Real} <: AbstractDimensionalityReduction

The `LTSA` type represents a local tangent space alignment model constructed for `T` type data with a help of the `NN` nearest neighbor algorithm.
"""
struct LTSA{NN <: AbstractNearestNeighbors, T <: Real} <: AbstractDimensionalityReduction
    λ::AbstractVector{T}
    proj::Projection{T}
    nearestneighbors::NN
    component::AbstractVector{Int}
end

## properties
outdim(R::LTSA) = size(R.proj, 1)
eigvals(R::LTSA) = R.λ
neighbors(R::LTSA) = R.nearestneighbors.k
vertices(R::LTSA) = R.component

## show
summary(io::IO, R::LTSA) = print(io, "LTSA(outdim = $(outdim(R)), neighbors = $(neighbors(R)))")

## interface functions
"""
    fit(LTSA, data; k=12, maxoutdim=2, nntype=BruteForce)

Fit a local tangent space alignment model to `data`.

# Arguments
* `data`: a matrix of observations. Each column of `data` is an observation.

# Keyword arguments
* `k`: a number of nearest neighbors for construction of local subspace representation
* `maxoutdim`: a dimension of the reduced space.
* `nntype`: a nearest neighbor construction class (derived from `AbstractNearestNeighbors`)

# Examples
```julia
M = fit(LTSA, rand(3,100)) # construct LTSA model
R = transform(M)           # perform dimensionality reduction
```
"""
function fit(::Type{LTSA}, X::AbstractMatrix{T};
        k::Int=12, maxoutdim::Int=2, ɛ::Real=1.0, nntype=BruteForce) where {T<:Real}
    # Construct NN graph
    NN = fit(nntype, X, k)
    D, E = knn(NN, X)
    G, C = largest_component(SimpleWeightedGraph(adjmat(D,E)))
    XX = @view X[:, C]
    n = length(C)

    S = ones(k)./sqrt(k)
    B = spzeros(T, n,n)
    for i=1:n
        II = @view E[:,i]
        VX = view(XX, :, II)

        # re-center points in neighborhood
        μ = mean(VX, dims=2)
        δ_x = VX .- μ

        # Compute orthogonal basis H of θ'
        θ_t = svd(δ_x).V[:,1:maxoutdim]

        # Construct alignment matrix
        G = hcat(S, θ_t)
        B[II, II] .+= Diagonal(fill(one(T), k)) .- G*transpose(G)
    end

    # Align global coordinates
    λ, V = decompose(B, maxoutdim)
    return LTSA{nntype, T}(λ, transpose(V), NN, C)
end

"""
    transform(R::LTSA)

Transforms the data fitted to the local tangent space alignment model `R` into a reduced space representation.
"""
transform(R::LTSA) = R.proj


"""
    swapping_even_odd_bits(n::Int64))

For a given integer value, all odd bits will be swapped with even bits.

# Arguments
- `n::Int64`: Integer value of the number, which has to be swapped.

# Examples
```julia-repl
julia> import ClassicAlgorithmsCollections
julia> ClassicAlgorithmsCollections.function swapping_even_odd_bits(23)
43
```

# Notes
---
The definition of the even and odd bits:
```julia-repl
even_bits = n & 0xAAAAAAAA
odd_bits = n & 0x55555555
```
"""
function swapping_even_odd_bits(n::Int64)
    # Defining even and odd bits
    even_bits = n & 0xAAAAAAAA
    odd_bits = n & 0x55555555

    # Shifting the even and odd bits
    even_bits >>= 1
    odd_bits <<= 1
    return (even_bits | odd_bits)
end


module Percolation

export clustgyration, hklabel, hkmaxclust!, genclust, clustfractal

using DataStructures

"""
Finds the radius of gyration for the largest closed cluster in the percolation grid

Percolation is assumed to occur from top to bottom. Uses the Hoshen-Kopelman algorithm.

# Argument

`grid::Matrix{Bool}`: Represents the grid cells; True if porous (part of the path;
    can pass through) and False if blocked (can't pass through; no path passes through)

# Returns

The gyration radius of the largest closed cluster (the largest cluster that doesn't connect
the top and bottom of the grid)
"""
function clustgyration(grid::BitMatrix)
    labels, list, clusters = hklabel(grid)
    if length(clusters) == 0
        return 0
    end

    maxlabel = hkmaxclust!(labels, list, clusters)
    if maxlabel == 0
        return 0
    end

    maxclust = Tuple.(findall(x -> x != 0 && hkfind(list, x) == maxlabel, labels))
    mass = length(maxclust)
    if mass == 0
        return 0
    end

    centerofmass = reduce(.+, maxclust) ./ mass
    return √(sum(pos -> reduce(.+, (pos .- centerofmass).^2),
        maxclust) / mass)
end

"""
Uses the Hoshen-Kopelman algorithm to label the grid

returns the labels, label mapping list, and cluster lengths
"""
@inline function hklabel(grid::BitMatrix)
    height, width = size(grid)
    labels = zeros(Integer, height, width)
    list = Vector{Integer}()
    clusters = Vector{Integer}()

    # first row, special case because of top boundary
    hkfirstrow!(grid, labels, list, clusters)

    for row in 2:height
        # first column, special case because of left boundary
        if grid[row, 1]
            hkfirstcol!(labels, list, clusters, row)
        end

        # "middle" points, with appropriate neighbors
        for col in 2:width
            if grid[row, col]
                hkiter!(labels, list, clusters, row, col)
            end
        end
    end

    return labels, list, clusters
end

"""
Helper function for `hkpercolate`; Adds new label.
"""
@inline function hknew!(labels::Matrix{<:Integer}, list::Vector{<:Integer},
        clusters::Vector{<:Integer}, row::Integer, col::Integer)
    labels[row, col] = length(list) + 1
    push!(list, labels[row, col])
    push!(clusters, 0)
end

"""
Helper function for `hkpercolate`; Applies the Hoshen-Kopelman algorithm to the first row.
"""
@inline function hkfirstrow!(grid::BitMatrix, labels::Matrix{<:Integer},
        list::Vector{<:Integer}, clusters::Vector{<:Integer})
    width = size(labels)[2]
    if grid[1, 1]
        hknew!(labels, list, clusters, 1, 1)
        clusters[labels[1, 1]] += 1
    end
    for col in 2:width
        if grid[1, col]
            if labels[1, col - 1] != 0
                labels[1, col] = hkfind!(list, labels[1, col - 1])
            else
                hknew!(labels, list, clusters, 1, col)
            end
            clusters[labels[1, col]] += 1
        end
    end
end

"""
Helper function for `hkpercolate`; Applies the Hoshen-Kopelman algorithm to the first column.
"""
@inline function hkfirstcol!(labels::Matrix{<:Integer}, list::Vector{<:Integer},
        clusters::Vector{<:Integer}, row::Integer)
    if labels[row - 1, 1] != 0
        labels[row, 1] = hkfind!(list, labels[row - 1, 1])
    else
        hknew!(labels, list, clusters, row, 1)
    end
    clusters[labels[row, 1]] += 1
end

"""
Helper function for `hkpercolate`; Applies one iteration of the Hoshen-Kopelman algorithm.
"""
@inline function hkiter!(labels::Matrix{<:Integer}, list::Vector{<:Integer},
        clusters::Vector{<:Integer}, row::Integer, col::Integer)
    above, left = labels[row - 1, col], labels[row, col - 1]
    if above == 0
        if left == 0
            hknew!(labels, list, clusters, row, col)
        else
            labels[row, col] = hkfind!(list, left)
        end
    elseif left == 0
        labels[row, col] = hkfind!(list, above)

    else
        labels[row, col] = hkunion!(list, clusters, above, left)
    end
    clusters[labels[row, col]] += 1
end

"""
Union function of the union-find algorithm used by the Hoshen-Kopelman algorithm
"""
@inline function hkunion!(list::Vector{<:Integer}, clusters::Vector{<:Integer},
        x::Integer, y::Integer)
    root = hkfind!(list, x)
    dest = hkfind!(list, y)
    if dest != root
        list[dest] = root
        clusters[root] += clusters[dest]
        clusters[dest] = 0
    end
    return root
end

"""
Find function of the union-find algorithm used by the Hoshen-Kopelman algorithm
"""
@inline function hkfind!(list::Vector{<:Integer}, x::Integer)
    y = x;
    while list[y] != y
        y = list[y];
    end
    while list[x] != x
        z = list[x];
        list[x] = y;
        x = z;
    end
    return y;
end

"""
Same as hkfind!, but doesn't update labels
"""
@inline function hkfind(list::Vector{<:Integer}, x::Integer)
    while x != list[x]
        x = list[x];
    end
    return x
end

"""
Finds the label of the largest closed cluster (a.k.a. finite cluster).

Also, sets the size of open clusters (a.k.a. infinite cluster) to zero, hence the !
"""
@inline function hkmaxclust!(labels::Matrix{Integer}, list::Vector{Integer},
        clusters::Vector{Integer})
    for root in labels[1, :]
        if root > 0
            for dest in labels[end, :]
                if dest > 0
                    toplabel = hkfind(list, root)
                    bottomlabel = hkfind(list, dest)
                    if toplabel == bottomlabel
                        clusters[toplabel] = 0
                    end
                end
            end
        end
    end

    if maximum(clusters) > 0
        return argmax(clusters)
    else
        return 0
    end
end

"""
Uses a depth-first search algorithm to generate a cluster with the given site probability

# Arguments

- `probability`: the probability a site is generated when growing the cluster
- `maxsize::Integer` (optional, defaults to 10000): the maximum size of the cluster
    (breaks the loop after the cluster is grown to this size)

# Returns

an integer matrix (8-bit) with 0 where no site exists, 1 where a cell is generated, and 2
where the growth is stopped
"""
function genclust(probability::Float64; maxsize::Integer=10000, crop::Bool=true)
    cluster = zeros(Int8, maxsize, maxsize)
    # create "fences" to stop the cluster form reaching the bounds of the matrix
    cluster[1, :] .= 2
    cluster[:, 1] .= 2
    cluster[end, :] .= 2
    cluster[:, end] .= 2

    stack = Stack{Tuple{Integer, Integer}}()
    seed = (maxsize÷2, maxsize÷2)
    push!(stack, seed)
    cluster[seed...] = 1

    clustdfs!(stack, cluster, probability, maxsize)

    if crop
        return cropclust!(cluster)
    else
        return cluster
    end
end

"""
The depth-first search algorithm used for `genclust`
"""
@inline function clustdfs!(stack::Stack, cluster::Matrix{<:Integer},
        probability::Float64, maxsize::Integer)
    size = 1
    while !isempty(stack) && size < maxsize
        site = pop!(stack)
        for step in ((-1, 0), (0, -1), (0, 1), (1, 0))
            neighbor = site .+ step
            if cluster[neighbor...] == 0
                if rand() <= probability
                    cluster[neighbor...] = 1
                    push!(stack, neighbor)
                    size += 1
                else
                    cluster[neighbor...] = 2
                end
            end
        end
    end
end

"""
finds the mass and radius of gyration of the cluster from the output of `genclust`
"""
function clustfractal(cluster::Matrix{Int8})
    clust = Tuple.(findall(x -> x == 1, cluster))
    mass = length(clust)
    if mass == 0
        return 0
    end

    centerofmass = reduce(.+, clust) ./ mass
    return mass, √(sum(pos -> reduce(.+, (pos .- centerofmass).^2), clust) / mass)
end

"""
crop the `genclust` function output to remove excess zeros
"""
@inline function cropclust!(cluster::Matrix{<:Integer})
    height, width = size(cluster)
    top = bottom = height ÷ 2
    left = right = width ÷ 2

    while 1 in cluster[top + 1, 2:end-1] || 2 in cluster[top + 1, 2:end-1]
        top += 1
    end
    while 1 in cluster[bottom - 1, 2:end-1] || 2 in cluster[bottom - 1, 2:end-1]
        bottom -= 1
    end
    while 1 in cluster[2:end-1, left - 1] || 2 in cluster[2:end-1, left - 1]
        left -= 1
    end
    while 1 in cluster[2:end-1, right + 1] || 2 in cluster[2:end-1, right + 1]
        right += 1
    end

    return cluster[bottom:top, left:right]
end

end


"""
    ExactReach

ExactReach performs exact reachability analysis to compute the output reachable set for a network.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: HPolytope
3. Output: HPolytope

# Return
`ReachabilityResult`

# Method
Exact reachability analysis.

# Property
Sound and complete.

# Reference
[W. Xiang, H.-D. Tran, and T. T. Johnson,
"Reachable Set Computation and Safety Verification for Neural Networks with ReLU Activations,"
*ArXiv Preprint ArXiv:1712.08163*, 2017.](https://arxiv.org/abs/1712.08163)
"""
struct ExactReach end

function solve(solver::ExactReach, problem::Problem)
    reach = forward_network(solver, problem.network, problem.input)
    return check_inclusion(reach, problem.output)
end

function forward_layer(solver::ExactReach, layer::Layer, input::Vector{HPolytope})
    output = Vector{HPolytope}(undef, 0)
    for i in 1:length(input)
        input[i] = linear_transformation(layer, input[i])
        append!(output, forward_partition(layer.activation, input[i]))
    end
    return output
end

function forward_layer(solver::ExactReach, layer::Layer, input::HPolytope)
    input = linear_transformation(layer, input)
    return forward_partition(layer.activation, input)
end

function forward_partition(act::ReLU, input::HPolytope)
    n = dim(input)
    output = Vector{HPolytope}(undef, 0)
    C, d = tosimplehrep(input)
    dh = [d; zeros(n)]
    for h in 0:2^n-1
        P = getP(h, n)
        Ch = [C; I - 2P]
        input_h = HPolytope(Ch, dh)
        if !isempty(input_h)
            push!(output, linear_transformation(Matrix(P), input_h))
        end
    end
    return output
end

function getP(h::Int64, n::Int64)
    str = string(h-1, pad = n, base = 2)
    vec = Vector{Int64}(undef, n)
    for i in 1:n
        vec[i] = ifelse(str[i] == '1', 1, 0)
    end
    return Diagonal(vec)
end


# Axis-aligned multivariate normal (i.e., independent entries, i.e., diagonal covariance matrix) using conjugate prior.
module MVNaaC

module MVNaaCmodel # submodule for component family definitions
export Theta, Data, log_marginal, new_theta, Theta_clear!, Theta_adjoin!, Theta_remove!,
       Hyperparameters, construct_hyperparameters, update_hyperparameters!

const Data = Array{Float64,1}

type Theta
    n::Int64                 # number of data points assigned to this cluster
    sum_x::Array{Float64,1}  # sum of the data points x assigned to this cluster
    sum_xx::Array{Float64,1} # sum of x.*x for the data points assigned to this cluster
    Theta(d) = (p=new(); p.n=0; p.sum_x=zeros(d); p.sum_xx=zeros(d); p)
end

new_theta(H) = Theta(H.d)

Theta_clear!(p) = (p.sum_x[:] = 0.; p.sum_xx[:] = 0.; p.n = 0)

Theta_adjoin!(p,x) = (for i=1:length(x); p.sum_x[i] += x[i]; p.sum_xx[i] += x[i]*x[i]; end; p.n += 1)
Theta_remove!(p,x) = (for i=1:length(x); p.sum_x[i] -= x[i]; p.sum_xx[i] -= x[i]*x[i]; end; p.n -= 1)

# In each dimension independently,
# X_1,...,X_n ~ Normal(mu,1/lambda) with Normal(mu|m,1/(c*lambda))Gamma(lambda|a,b) prior on mean=mu, precision=lambda.
function log_marginal(p,H)
    n = p.n
    LB = 0.0
    # For each dimension
    for i = 1:H.d
        ## update to b
        t1 = H.b + 0.5 * p.sum_xx[i] * H.alpha_wt ###
        t2 = 0.5 * p.sum_x[i]*p.sum_x[i] / n * H.alpha_wt
        t3 = 0.5 * H.c * n * H.alpha_wt * (p.sum_x[i]/n - H.m)^2 / (H.c + n * H.alpha_wt)
        #LB += log(H.b + 0.5 * p.sum_xx[i] - 0.5 * p.sum_x[i]*p.sum_x[i] / n + \
        #          0.5 * H.c * n * (p.sum_x[i]/n - H.m)^2/(H.c + n))
        LB += log(t1 - t2 + t3)
    end

    # update to a
    aa = H.a + 0.5 * n * H.alpha_wt ##
    # log term w/ update to c
    cc = 0.5 * log(H.c + n * H.alpha_wt) ##

    term = H.d * (H.constant - 0.5*n*log(2*pi)*H.alpha_wt - cc + H.log_Ga[p.n]) - aa * LB
    #return H.d * (H.constant - 0.5*n*log(2*pi) - 0.5*log(H.c+n) + H.log_Ga[n]) - (H.a+0.5*n) * LB
    return term
end

function log_marginal(x,p,H)
    Theta_adjoin!(p,x)
    result = log_marginal(p,H)
    Theta_remove!(p,x)
    return result
end

type Hyperparameters
    d::Int64    # dimension
    m::Float64  # prior mean of mu's
    c::Float64  # prior precision multiplier for mu's
    a::Float64  # prior shape of lambda's
    b::Float64  # prior rate of lambda's
    constant::Float64
    log_Ga::Array{Float64,1}
    alpha_wt::Float64 # alpha weight (not technically a hyperparameter)
end

function construct_hyperparameters(options)
    x = options.x
    n = length(x)
    d = length(x[1])
    mu = mean(x)
    v = mean([xi.*xi for xi in x]) - mu.*mu  # sample variance
    @assert(all(abs.(mu) .< 1e-10) && all(abs.(v - 1.0) .< 1e-10), "Data must be normalized to zero mean, unit variance.")
    m = 0.0
    c = 1.0
    a = 1.0
    b = 1.0

    # specify the weight
    alpha_eval = eval(parse(options.alpha_wt_str))
    alpha_wt_fn(n) = Base.invokelatest(alpha_eval, n)
    alpha_wt = alpha_wt_fn(options.n)

    log_Ga = lgamma.(a+0.5*(1:n+1) * alpha_wt)

    # normalizing constant of the log Normal-gamma pdf involving hyperparams
    constant = 0.5*log(c) + a*log(b) - lgamma(a)

    return Hyperparameters(d,m,c,a,b,constant,log_Ga, alpha_wt)
end

function update_hyperparameters!(H,theta,list,t,x,z)
    #error("update_hyperparameters is not yet implemented.")
end

end # module MVNaaCmodel
using .MVNaaCmodel

# Include generic code
include("generic.jl")

# Include core sampler code
include("coreConjugate.jl")

end # module MVNaaC


amaxmin(x::T, y::T) where {T} = ifelse( abs(x) < abs(y), (y,x), (x,y) )
function amaxmin(x::T, y::T, z::T) where {T}
     y, z = amaxmin(y, z)
     x, z = amaxmin(x, z)
     x, y = amaxmin(x, y)
     return x, y, z
end
function amaxmin(a::T, b::T, c::T, d::T) where {T}
    a, b = amaxmin(a, b)
    c, d = amaxmin(c, d)
    a, c = amaxmin(a, c)
    b, d = amaxmin(b, d)
    b, c = amaxmin(b, c)

    return a, b, c, d
end

"""
    two_sum(a, b)

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_sum(a::T, b::T) where {T<:Real}
    hi = a + b
    v  = hi - a
    lo = (a - (hi - v)) + (b - v)
    return hi, lo
end

"""
    one_sum(a, b)

Computes `fl(a+b)`.
"""
@inline function one_sum(a::T, b::T) where {T<:Real}
    return a + b
end

"""
   three_sum(a, b, c)

Computes `hi = fl(a+b+c)` and `md = err(a+b+c), lo = err(md)`.
"""
function three_sum(a::T, b::T, c::T) where {T}
    hi, md, lo = amaxmin(a, b, c)
    md, lo = two_hilo_sum(md, lo)
    hi, md = two_hilo_sum(hi, md)
    md, lo = two_hilo_sum(md, lo)
    return hi, md, lo
end

"""
    two_sum(a, b, c)

Computes `hi = fl(a+b+c)` and `lo = err(a+b+c)`.
"""
@inline function two_sum(a::T, b::T, c::T) where {T}
    hi, md, lo = amaxmin(a, b, c)
    md, lo = two_hilo_sum(md, lo)
    hi, md = two_hilo_sum(hi, md)
    md += lo
    return hi, md
end

"""
    one_sum(a, b, c)

Computes `fl(a+b+c)`
"""
@inline function one_sum(a::T,b::T,c::T) where {T}
    md, lo = two_sum(b, c)
    hi, md = two_sum(a, md)
    md, lo = two_sum(md, lo)
    hi, md = two_hilo_sum(hi, md)
    md = md + lo
    hi = hi + md
    return hi
end

"""
    four_sum(a, b, c, d)

Computes `hi = fl(a+b+c+d)` and `hm = err(a+b+c+d), ml = err(hm), lo = err(ml)`.
"""
function four_sum(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_sum(a,  b)
    t2, t3 = two_sum(c,  d)
    hi, t4 = two_sum(t0, t2)
    t5, lo = two_sum(t1, t3)
    hm, ml = two_sum(t4, t5)
    ml, lo = two_hilo_sum(ml, lo)
    hm, ml = two_hilo_sum(hm, ml)
    hi, hm = two_hilo_sum(hi,hm)
    return hi, hm, ml, lo
end

"""
    three_sum(a, b, c, d)

Computes `hi = fl(a+b+c+d)` and `md = err(a+b+c+d), lo = err(md)`.
"""
function three_sum(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_sum(a,  b)
    t2, t3 = two_sum(c,  d)
    hi, t4 = two_sum(t0, t2)
    t5, t6 = two_sum(t1, t3)
    md, lo = two_sum(t4, t5)
    lo = lo + t6
    md, lo = two_hilo_sum(md, lo)
    hi, md = two_hilo_sum(hi, md)
    return hi, md, lo
end

"""
    two_sum(a, b, c, d)

Computes `hi = fl(a+b+c+d)` and `lo = err(a+b+c+d)`.
"""
function two_sum(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_sum(a,  b)
    t2, t3 = two_sum(c,  d)
    hi, t4 = two_sum(t0, t2)
    t5, t6 = two_sum(t1, t3)
    lo, t7 = two_sum(t4, t5)
    t7 = t7 + t6
    lo, t7 = two_hilo_sum(lo, t7)
    hi, lo = two_hilo_sum(hi, lo)
    return hi, lo
end

"""
    one_sum(a, b, c, d)

Computes `fl(a+b+c+d)`.
"""
function one_sum(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_sum(a,  b)
    t2, t3 = two_sum(c,  d)
    hi, t4 = two_sum(t0, t2)
    t5, t6 = two_sum(t1, t3)
    lo, t7 = two_sum(t4, t5)
    t7 = t7 + t6
    lo, t7 = two_hilo_sum(lo, t7)
    hi = hi + lo
    return hi
end

function vec_sum(x0::T, x1::T, x2::T, x3::T) where {T}
    s3 = x3
    s2, e3 = two_sum(x2, s3)
    s1, e2 = two_sum(x1, s2)
    s0, e1 = two_sum(x0, s1)
    return s0,e1,e2,e3
end

function vsum_errbranch(x::NTuple{4,T}) where {T}
    y = zeros(T, 4)
    r = zeros(T, 4)
    e = zeros(T, 4)
    j = 1
    e[1] = x[1]
    for i = 1:2
        r[i], t = two_sum(e[i], x[i+1])
        if t !== zero(T)
            y[j] = r[i]
            e[i+1] = t
            j += 1
        else
            e[i+1] = r[i]
        end
    end
    y[j], y[j+1] = two_sum(e[3], x[4])
    return y
end

function foursum(x1::T, x2::T, x3::T, x4::T) where {T}
    a1, a2 = two_sum(x1, x2)
    b1, b2 = two_sum(x3, x4)
    c1, c2 = two_sum(a1, b1)
    d1, d2 = two_sum(a2, b2)
    e1to4 = vec_sum(c1,c2,d1,d2)
    y = vsum_errbranch(e1to4)
    return (y...,)
end

"""
    five_sum(a, b, c, d, e)

Computes `s = fl(a+b+c+d+e)` and
    `e1 = err(a+b+c+d), e2 = err(e1), e3 = err(e2), e4 = err(e3)`.
"""
function five_sum(v::T, w::T, x::T, y::T, z::T) where {T}
    t0, t4 = two_sum(y, z)
    t0, t3 = two_sum(x, t0)
    t0, t2 = two_sum(w, t0)
    a, t1  = two_sum(v, t0)
    t0, t3 = two_sum(t3, t4)
    t0, t2 = two_sum(t2, t0)
    b, t1  = two_sum(t1, t0)
    t0, t2 = two_sum(t2, t3)
    c, t1  = two_sum(t1, t0)
    d, e   = two_sum(t1, t2)
    return a, b, c, d, e
end

"""
    two_diff(a, b)

Computes `s = fl(a-b)` and `e = err(a-b)`.
"""
@inline function two_diff(a::T, b::T) where {T}
    hi = a - b
    v  = hi - a
    lo = (a - (hi - v)) - (b + v)
    return hi, lo
end

"""
    three_diff(a, b, c)

Computes `s = fl(a-b-c)` and `e1 = err(a-b-c), e2 = err(e1)`.
"""
function three_diff(a::T,b::T,c::T) where {T}
    s, t = two_diff(-b, c)
    x, u = two_sum(a, s)
    y, z = two_sum(u, t)
    x, y = two_hilo_sum(x, y)
    return x, y, z
end

"""
    four_diff(a, b, c, d)

Computes `hi = fl(a-b-c-d)` and `hm = err(a-b-c-d), ml = err(hm), lo = err(ml)`.
"""
function four_diff(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_diff(a ,  b)
    t0, t2 = two_diff(t0,  c)
    hi, t3 = two_diff(t0,  d)
    t0, t1 = two_sum(t1, t2)
    hm, t2 = two_sum(t0, t3) # here, t0 >= t3
    ml, lo = two_sum(t1, t2)
    return hi, hm, ml, lo
end

"""
    two_square(a)

Computes `hi = fl(a*a)` and `lo = fl(err(a*a))`.
"""
@inline function two_square(a::T) where {T}
    hi = a * a
    lo = fma(a, a, -hi)
    hi, lo
end

"""
    two_prod(a, b)

Computes `hi = fl(a*b)` and `lo = fl(err(a*b))`.
"""
@inline function two_prod(a::T, b::T) where {T}
    hi = a * b
    lo = fma(a, b, -hi)
    hi, lo
end

@inline max_min(a,b) = abs(a) < abs(b) ? (b,a) : (a,b)

"""
    three_prod(a, b, c)

Computes `hi = fl(a*b*c)` and `md = err(a*b*c), lo = err(md)`.
"""
function three_prod(a::T, b::T, c::T) where {T}
    a, b = max_min(a,b)
    a, c = max_min(a,c)
    b, c = max_min(b,c)
    abhi, ablo = two_prod(a, b)
    hi, abhiclo = two_prod(abhi, c)
    ablochi, abloclo = two_prod(ablo, c)
    md, lo, tmp  = three_sum(ablochi, abhiclo, abloclo)
    hi, md = two_sum(hi, md)
    return hi, md, lo
end

#=
   three_fma algorithm from
   Sylvie Boldo and Jean-Michel Muller
   Some Functions Computable with a Fused-mac
=#

"""
   three_fma(a, b, c)

Computes `hi = fl(fma(a,b,c))` and `md = fl(err(fma(a,b,c))), lo = fl(err(md))`.
"""
function three_fma(a::T, b::T, c::T) where {T}
    hi = fma(a, b, c)
    hi0, lo0 = two_prod(a, b)
    hi1, lo1 = two_sum(c, lo0)
    hi2, lo2 = two_sum(hi0, hi1)
    y = ((hi2 - hi) + lo2)
    md, lo = two_hilo_sum(y, lo1)
    return hi, md, lo
end

"""
   two_fma(a, b, c)

Computes `hi = fl(fma(a,b,c))` and `lo = fl(err(fma(a,b,c)))`.
"""
function two_fma(a::T, b::T, c::T) where {T}
    hi = fma(a, b, c)
    hi0, lo0 = two_prod(a, b)
    hi1, lo1 = two_sum(c, lo0)
    hi2, lo2 = two_sum(hi0, hi1)
    lo = ((hi2 - hi) + lo2) + lo1
    return hi, lo
end

"""
   two_muladd(a, b, c)

Computes `hi = fl(muladd(a,b,c))` and `lo = fl(err(muladd(a,b,c)))`.
"""
function two_muladd(a::T, b::T, c::T) where {T}
    hi = fma(a, b, c)
    lo = fma(a, b, c-hi)
    return hi, lo
end

# with arguments sorted by magnitude

"""
    two_hilo_sum(a, b)

*unchecked* requirement `|a| ≥ |b|`

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_hilo_sum(a::T, b::T) where {T}
    hi = a + b
    lo = b - (hi - a)
    return hi, lo
end

"""
    two_lohi_sum(a, b)

*unchecked* requirement `|b| ≥ |a|`

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_lohi_sum(a::T, b::T) where {T}
    hi = b + a
    lo = a - (hi - b)
    return hi, lo
end

"""
    two_hilo_diff(a, b)

*unchecked* requirement `|a| ≥ |b|`

Computes `hi = fl(a-b)` and `lo = err(a-b)`.
"""
@inline function two_hilo_diff(a::T, b::T) where {T}
    hi = a - b
    lo = (a - hi) - b
    hi, lo
end

"""
    two_lohi_diff(a, b)

*unchecked* requirement `|b| ≥ |a|`

Computes `hi = fl(a-b)` and `lo = err(a-b)`.
"""
@inline function two_lohi_diff(a::T, b::T) where {T}
    hi = b - a
    lo = (b - hi) - a
    hi, lo
end

"""
    two_hilo_sum(a, b, c)

*unchecked* requirement `|a| ≥ |b| ≥ |c|`

Computes `hi = fl(a+b+c)` and `lo = err(a+b+c)`.
"""
function two_hilo_sum(a::T, b::T, c::T) where {T}
    lo, t = two_hilo_sum(b, c)
    hi, lo = two_sum(a, lo)
    lo += t
    hi, lo = two_hilo_sum(hi, lo)
    return hi, lo
end

"""
    three_hilo_sum(a, b, c)

*unchecked* requirement `|a| ≥ |b| ≥ |c|`

Computes `x = fl(a+b+c)` and `y = err(a+b+c), z = err(y)`.
"""
function three_hilo_sum(a::T, b::T, c::T) where {T}
    md, lo = two_hilo_sum(b, c)
    hi, md = two_sum(a, md)
    md, lo = two_hilo_sum(md, lo)
    hi, md = two_hilo_sum(hi, md)
    return hi,md,lo
end

"""
    three_lohi_sum(a, b, c)

*unchecked* requirement `|c| ≥ |b| ≥ |a|`

Computes `x = fl(a+b+c)` and `y = err(a+b+c), z = err(y)`.
"""
function three_lohi_sum(a::T,b::T,c::T) where {T}
    md, lo = two_lohi_sum(a, b)
    hi, md = two_sum(c, md)
    md, lo = two_lohi_sum(lo, md)
    hi, md = two_lohi_sum(md, hi)
    return hi,md,lo
end

"""
    three_hilo_diff(a, b, c)

*unchecked* requirement `|a| ≥ |b| ≥ |c|`

Computes `x = fl(a-b-c)` and `y = err(a-b-c), z = err(y)`.
"""
function three_hilo_diff(a::T,b::T,c::T) where {T}
    md, lo = two_hilo_sum(-b, -c)
    hi, md = two_sum(a, md)
    md, lo = two_hilo_sum(md, lo)
    hi, md = two_hilo_sum(hi, md)
    return hi,md,lo
end

"""
    three_lohi_diff(a, b, c)

*unchecked* requirement `|c| ≥ |b| ≥ |a|`

Computes `x = fl(a-b-c)` and `y = err(a-b-c), z = err(y)`.
"""
function three_lohi_diff(c::T,b::T,a::T) where {T}
    three_lohi_sum(c, -b, -a)
end

"""
    four_hilo_sum(a, b, c, d)

*unchecked* requirement `|a| ≥ |b| ≥ |c| ≥ |d|`

Computes `hi = fl(a+b+c+d)` and `hm = err(a+b+c+d), ml = err(hm), lo = err(ml)`.
"""
function four_hilo_sum(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_hilo_sum(a ,  b)
    t0, t2 = two_hilo_sum(t0,  c)
    hi, t3 = two_hilo_sum(t0,  d)
    t0, t1 = two_hilo_sum(t1, t2)
    hm, t2 = two_hilo_sum(t0, t3) # here, t0 >= t3
    ml, lo = two_hilo_sum(t1, t2)
    return hi, hm, ml, lo
end

"""
    four_lohi_sum(a, b, c, d)

*unchecked* requirement `|d| ≥ |c| ≥ |b| ≥ |a|`

Computes `hi = fl(a+b+c+d)` and `hm = err(a+b+c+d), ml = err(hm), lo = err(ml)`.
"""
function four_lohi_sum(d::T,c::T,b::T,a::T) where {T}
    t0, t1 = two_hilo_sum(a ,  b)
    t0, t2 = two_hilo_sum(t0,  c)
    hi, t3 = two_hilo_sum(t0,  d)
    t0, t1 = two_hilo_sum(t1, t2)
    hm, t2 = two_hilo_sum(t0, t3)
    ml, lo = two_hilo_sum(t1, t2)
    return hi, hm, ml, lo
end

"""
    four_hilo_diff(a, b, c, d)

*unchecked* requirement `|a| ≥ |b| ≥ |c| ≥ |d|`

Computes `hi = fl(a-b-c-d)` and `hm = err(a-b-c-d), ml = err(hm), lo = err(ml)`.
"""
function four_hilo_diff(a::T,b::T,c::T,d::T) where {T}
    t0, t1 = two_hilo_diff(a,  b)
    t0, t2 = two_hilo_diff(t0,  c)
    hi, t3 = two_hilo_diff(t0,  d)
    t0, t1 = two_hilo_sum(t1, t2)
    hm, t2 = two_hilo_sum(t0, t3) # here, t0 >= t3
    ml, lo = two_hilo_sum(t1, t2)
    return hi, hm, ml, lo
end

"""
    four_hilo_diff(a, b, c, d)

*unchecked* requirement `|d| ≥ |c| ≥ |b| ≥ |a|`

Computes `hi = fl(a-b-c-d)` and `hm = err(a-b-c-d), ml = err(hm), lo = err(ml)`.
"""
function four_lohi_diff(d::T,c::T,b::T,a::T) where {T}
    t0, t1 = two_hilo_diff(a,  b)
    t0, t2 = two_hilo_diff(t0,  c)
    hi, t3 = two_hilo_diff(t0,  d)
    t0, t1 = two_hilo_sum(t1, t2)
    hm, t2 = two_hilo_sum(t0, t3) # here, t0 >= t3
    ml, lo = two_hilo_sum(t1, t2)
    return hi, hm, ml, lo
end


function checkSVDDone(S::Vector,
                      thresh::Float64)
  N = length(S)
  (N <= 1 || thresh < 0.0) && return (true,1)
  S1t = S[1]*thresh
  start = 2
  while start <= N
    (S[start] < S1t) && break
    start += 1
  end
  if start >= N
    return (true,N)
  end
  return (false,start)
end

function svd_recursive(M::AbstractMatrix;
                      thresh::Float64=1E-3,
                      north_pass::Int=2)
  Mr,Mc = size(M)

  if Mr > Mc
    V,S,U = svd_recursive(transpose(M))
    conj!(U)
    conj!(V)
    return U,S,V
  end

  #rho = BLAS.gemm('N','T',-1.0,M,M) #negative to sort eigenvalues greatest to smallest
  rho = -M*M' #negative to sort eigenvalues in decreasing order
  D,U = eigen(Hermitian(rho),1:size(rho,1))

  Nd = length(D)

  V = M'*U

  V,R = qr_positive(V)
  for n=1:Nd
    D[n] = R[n,n]
  end

  (done,start) = checkSVDDone(D,thresh)

  done && return U,D,V

  u = view(U,:,start:Nd)
  v = view(V,:,start:Nd)

  b = u'*(M*v)
  bu,bd,bv = svd_recursive(b,
                          thresh=thresh,
                          north_pass=north_pass)

  u .= u*bu
  v .= v*bv
  view(D,start:Nd) .= bd

  return U,D,V
end

# TODO: maybe move to another location?
# Include options for other svd algorithms
function polar(M::AbstractMatrix)
  U,S,V = svd(M) # calls LinearAlgebra.svd(_)
  return U*V', V*Diagonal(S)*V'
end


# 'cstrt_bwinja_ngon' provides the vertices coordinates (a,b) of a small
# n-gon for n = 2m and m >= 3
function cstrt_bwinja_ngon(n)
if mod(n,2) == 0 && n>=6
    # initialization
    v(u) = (pi/2-u)/(n/2-1);
    d(u) = asin(sin(u)+sin(u+3*v(u)/2)/(2*cos(v(u)/2)))-u-v(u);
    A(u) = sin(u)+sin(2*v(u))-sin(v(u)+d(u))+
    (n/2-3)*(sin(v(u))-tan(v(u)/2))+(cos(v(u)-d(u))-cos(2*v(u))-1/2)*tan(v(u)/2);
    F(u) = -A(u);
    res = optimize(F,pi/(2*n-2),pi/n);
    u0 = Optim.minimizer(res);
    v0 = (pi/2-u0)/(n/2-1);
    d0 = asin(sin(u0)+sin(u0+3*v0/2)/(2*cos(v0/2)))-u0-v0;
    x = zeros(Int(n/2));
    x[1] = u0;
    x[2] = v0+d0;
    x[3] = v0-d0;
    if n >= 8
        x[4:end] .= v0;
    end
    # construction
    a = zeros(n-1); b = zeros(n-1);
    a[Int(n/2-1)] = sin(x[1]); a[Int(n/2+1)] = -a[Int(n/2-1)];
    b[Int(n/2-1)] = cos(x[1]); b[Int(n/2+1)] = b[Int(n/2-1)];
    for i = 2:Int(n/2-1)
        a[Int(mod(i*(n-1)/2-mod(i,2)/2,n))] = a[Int(mod((i-1)*(n-1)/2-mod(i-1,2)/2,n))]-(-1)^i*sin(sum(x[1:i]));
        a[Int(n-mod(i*(n-1)/2-mod(i,2)/2,n))] = -a[Int(mod(i*(n-1)/2-mod(i,2)/2,n))];
        b[Int(mod(i*(n-1)/2-mod(i,2)/2,n))] = b[Int(mod((i-1)*(n-1)/2-mod(i-1,2)/2,n))]-(-1)^i*cos(sum(x[1:i]));
        b[Int(n-mod(i*(n-1)/2-mod(i,2)/2,n))] = b[Int(mod(i*(n-1)/2-mod(i,2)/2,n))];
    end
    a[Int(n/2)] = 0; b[Int(n/2)] = 1;
    return(a,b)
end
end


"""
    QR()
Finds a vector vanishing 'gamma' such that all of its entries sum to 1.

Input arguments
---------------
gamma::Array{Float64, 2}    This matrix verifies properties necessary for QR() to work. We know that the rank of 'Gamma' n-1 (otherwise we wouldn't reach this function).
tolerance::Float64          Same as always.

Returns
-------
lambda::Vec{Float64}        A vector vanishing 'gamma'. lambda is a column vector containing the components of a vector generating Nul(M) with all the components summing up to 1.
"""
function QR(Gamma, tolerance)
   # We know that rank(Gamma)=n-1
   #and rank(Gamma0)=n
   # In addition all columns of M are non vanishing
   ## Outcome
   # lambda is a column vector containing the components of a vector generating
   # Nul(M) with all the components summing up to 1
   ##

   n = size(Gamma, 2)
   m = size(Gamma, 1) # m>=n-1
   if m == 1
      #then n<=2 but n=1 is excluded since it corresponds to a 1-vertex boundary, thus n=2 and Gamma=[a b]
      #also notice that since Gamma0=[[a b];[1 1]] has full rank it is ruled
      #out that a=b
      lambda = [Gamma[2]; -Gamma[1]] / (Gamma[2] - Gamma[1])
   else

      # M=Q*R
      #Q is an orthogonal matrix of dimension m x m
      qr_decomposition = qr(Gamma).R # access the second element of the tuple
      # R is an upper triangular matrix of dimension m x n
      R = triu(qr_decomposition)

      # diagonal is a column vector with dimension either n-1 (if m=n-1) or n (if m>=n)
      # Set entries that are too small relative to `tolerance` to zero.
      diagonal = heaviside(tolerance .- abs.(diag(R)))
      index = round.(Int, transpose(collect(1:min(m, n))) * diagonal)[1]

      if index == 0
         # the first n-1 columns are linearly independent
         index = n - 1
      else
         #the index-th entry in diag(R) is zero
         index = index - 1
      end

      lambda = zeros(n, 1)
      lambda[1:index] = - R[1:index, 1:index] \ R[1:index, index + 1] # A\b means inv(A)*b
      lambda[index + 1] = 1
      lambda = lambda / (ones(1, n) * lambda)
   end
   lambda = lambda .* heaviside(abs.(lambda) .- tolerance)

   return(lambda)
end


# https://mrandri19.github.io/2020/09/28/bayesian-linear-regression-with-conjugate-priors.html
# https://qiita.com/haru1843/items/3956dab2fd0d448cd02b
# https://discourse.julialang.org/t/solving-ax-b-for-large-matrix-dimesnions-efficiently-in-julia/51504
# https://stackoverflow.com/questions/57270276/identity-matrix-in-julia

# bayesian inferense (using conjugate prior analyticaly) of coefficients and variance in error term for linear regression model

using Random, Plots, StatsPlots, Distributions, LinearAlgebra, DataFrames, Printf, LaTeXStrings

"""
Conjugate Priors for linear regression model (β|σ^2 ~ MvNormal, σ^2 ~ InverseGamma)

    Inputs
    ------
    y:    data, independent variable (n vector)
    X:    data, dependent variables (n * k matrix)
    β0:   mean of conditional prior (multivariate normal β0|σ^2 ~ MvNormal) for coefficients
    A0:   precision matrix (inverse of Σ) of conditional prior for coefficients
    ν0:   shape parameter of prior for variance of error term
    λ0:   scale parameter of prior for variance of error term
    prob: credible interval probability (0 < prob < 1)

    Outputs
    -------
    results: posterior summary stats table (DataFrames)
    marginal_posterior_β: vector of marginal posterior for each coefficients (single β|D ~ TDist)
    posterior_σ2: posterior for variance of error term (σ^2|D ~ InverseGamma)
"""
function regression_stats(y, X, β0, A0, ν0, λ0, prob)
    XX = X' * X
    Xy = X' * y
    # @time XX \ Xy # fast and efficient
    # @time inv(XX) * Xy
    β_ols = XX \ Xy # OLS estimator
    A_star = XX + A0
    β_star = A_star \ (Xy + A0 * β0)
    C_star = inv(inv(XX) + inv(A0))
    ν_star = size(y)[1] + ν0
    ssr = sum((y - X * β_ols).^2) # SSR sum of squared residual 残差二乗和
    # @time (β0 - β_ols)' * C_star * (β0 - β_ols)
    # @time dot(β0 - β_ols, C_star, β0 - β_ols) # faster
    λ_star = ssr + dot(β0 - β_ols, C_star, β0 - β_ols) + λ0
    H_star = (λ_star / ν_star) * inv(A_star)
    h_star = diag(sqrt(H_star))

    marginal_posterior_β = LocationScale.(β_star, h_star, TDist(ν_star)) # broadcasting through coefficients
    sd_β = std.(marginal_posterior_β)
    ci_β_upper = quantile.(marginal_posterior_β, prob - (1-prob)/2)
    ci_β_lower = quantile.(marginal_posterior_β, (1-prob)/2)
    # ci_β = [ci_β_lower, ci_β_upper]
    stats_β = [β_star β_star β_star sd_β ci_β_lower ci_β_upper]

    posterior_σ2 = InverseGamma(ν_star/2, λ_star/2)
    mean_σ2 = mean(posterior_σ2)
    median_σ2 = median(posterior_σ2)
    mode_σ2 = mode(posterior_σ2)
    sd_σ2 = std(posterior_σ2)
    ci_σ2_upper = quantile(posterior_σ2, prob - (1-prob)/2)
    ci_σ2_lower = quantile(posterior_σ2, (1-prob)/2)
    # ci_σ2 = [ci_σ2_lower, ci_σ2_upper]
    stats_σ2 = [mean_σ2 median_σ2 mode_σ2 sd_σ2 ci_σ2_lower ci_σ2_upper]
    stats = [stats_β; stats_σ2]
    stats_string = ["Mean", "Median", "Mode", "SD", @sprintf("CI(%d)_Lower", prob*100), @sprintf("CI(%d)_Upper", prob*100)]
    results = DataFrame(stats, stats_string)
    return results, marginal_posterior_β, posterior_σ2, β_ols
end

# Generate moc data
n = 50
dim = 1
Random.seed!(99)
u = rand(Normal(0, 0.7), n) # true value σ^2 = 0.49
x = rand(Uniform(-sqrt(3.0), sqrt(3.0)), n, dim) # x generated from Uniform[-√3, √3]
X = [ones(n) x] # dependent variable matrix with ones
β = Vector(1:1:dim+1) # true coefficients
y = X*β + u # true data generating process

# set prior
k = size(X, 2)
β0 = zeros(k)
τ0 = 0.2 # no ridge penalty ... τ0 = Inf ?

# τ0 = 0.001
# τ0 = 0.01
# τ0 = 0.1
# τ0 = 1
# τ0 = 10
# τ0 = 100
# τ0 = 1000

A0 = τ0 * I # Identity matrix
ν0 = 5.0
λ0 = 7.0
H0 = (λ0 / ν0) * inv(A0) # scale matrix for marginal prior of multivariate coefficient (mv Tdist)
h0 = diag(sqrt(Matrix(H0, k, k))) # scale params for marginal prior of single coefficient (T dist) H0の対角成分
marginal_prior_β = LocationScale.(β0, h0, TDist(ν0)) # broadcasting through coefficients
prior_σ2 = InverseGamma(ν0/2, λ0/2)

# calculate posterior
prob = 0.95
results, marginal_posterior_β, posterior_σ2, β_ols = regression_stats(y, X, β0, A0, ν0, λ0, prob)
print("$results\n")

# Visualize Posterior
plots = []

for i = 1:k
    label_prior = @sprintf("marginal prior: T(ν=%.1f, μ=%.1f, σ^2=%.1f)", params(params(marginal_prior_β[i])[3])[1], params(marginal_prior_β[i])[1], params(marginal_prior_β[i])[2])
    p = plot(marginal_prior_β[i], linestyle=:dash, label=label_prior, xlims=(0.0, 3.0), ylims=(0.0, 4.0), xlabel=@sprintf("β%d", i), ylabel="Probability Density")
    label_posterior = @sprintf("marginal posterior: T(ν=%.1f, μ=%.1f, σ^2=%.1f)", params(params(marginal_posterior_β[i])[3])[1], params(marginal_posterior_β[i])[1], params(marginal_posterior_β[i])[2])
    plot!(marginal_posterior_β[i], linestyle=:solid, label=label_posterior)
    push!(plots, p)
end

label_prior = @sprintf("prior: InverseGamma(α=%.1f, θ=%.1f)", params(prior_σ2)[1], params(prior_σ2)[2])
p = plot(prior_σ2, linestyle=:dash, label=label_prior, xlims=(0.0, 3.0), ylims=(0.0, 4.0), xlabel=L"\sigma^2", ylabel="Probability Density")
label_posterior = @sprintf("prior: InverseGamma(α=%.1f, θ=%.1f)", params(posterior_σ2)[1], params(posterior_σ2)[2])
plot!(posterior_σ2, linestyle=:solid, label=label_posterior)
push!(plots, p)

plot(plots..., layout=(1, k+1))

# plot ols
scatter(x, y, label="data")
# scatter(x, y, xlims=(-0.1, 0.1), ylims=(0.9, 1.5))
plot!(x, X * β, label=(@sprintf("True: y = %.d + %.dx", β[1], β[2])))
plot!(x, X * β_ols, label=(@sprintf("OLS: y = %.5f + %.5fx", β_ols[1], β_ols[2])))
β_star = results[:, "Mean"][1:k]
plot!(x, X * β_star, label=(@sprintf("Bayes: y = %.5f + %.5fx (τ0=%.2f, Ridge)", β_star[1], β_star[2], τ0)))

# significance test on coefficients H0: βj = 0, H1: βj ≠ 0
# evaluate with SDDR
"""
calculate SDDR (compare density of posterior and prior on particular value)
# https://docs.julialang.org/en/v1/manual/methods/
# https://stackoverflow.com/questions/61088294/declaring-the-name-of-argument-when-invoking-a-function

    Inputs
    ------
    β_h0: test value, (default is 0)
    posterior
    prior

    Output
    ------
    sddr
"""
function sddr(posterior, prior, ;null_hypothesis::Float64=0.0)
    return pdf(posterior, null_hypothesis) / pdf(prior, null_hypothesis)
end

print(log10.(sddr.(marginal_posterior_β, marginal_prior_β, null_hypothesis=0.0)))


# probit
# https://stats.stackexchange.com/questions/155239/marginal-effects-from-bayesian-probit

using Turing
using Distributions
using Random
using LinearAlgebra
using StatsPlots
using MCMCChains
using DataFrames
using Printf
# using Plots # reesport
# using Statistics # reesport

using StatsFuns: normcdf, norminvcdf, normpdf

# cdf of normal distribution # Φ(x'β)
plot(normcdf.(Vector(-10:0.1:10)))
# probit function, inverse-cdf of normal distribution # Φ^{-1}(y)
plot(norminvcdf.(Vector(0:0.01:1)))

# geberate moc data from logit model
n = 1_000
num_var = 2
Random.seed!(99)
X = rand(Uniform(-sqrt(3.0), 2.0*sqrt(3.0)), n, num_var) # case-specific variables Z_i Matrix
X = [ones(n) X]
β = Vector([0.0, 0.5, -0.5]) # true coefficients
# V = X * β # representative utility, # utility (latent variable) y_star = U = Xβ + ϵ = V + ϵ
p = normcdf.(X * β) # probability (する確率)
y = rand.(Bernoulli.(p)) # observable choice 0 or 1

# set prior, β ~ MvNormal(β0, A0), no σ²
n, k = size(X)
β0 = zeros(k)
A0 = 0.01 * I

@model function probit_model(X, y)
    # Set the priors only for β.
    β ~ MvNormal(β0, inv(A0))

    # Write the likelihood, can't vectorize? no multivariate distribution
    n = size(X, 1)
    for i = 1:n
        y[i] ~ Bernoulli(normcdf(X[i, :]' * β))
    end
end

# Draw random number from posterior.
model = probit_model(X, y)

n_draws = 5_000
# n_chains = 4
n_chains = 1
n_tune = 1_000
chn = sample(model, NUTS(), MCMCThreads(), n_draws, n_chains)
# chn = sample(model, NUTS(), MCMCThreads(), n_draws, n_chains)[n_tune:n_draws, :, :]
β_chn = chn[n_tune+1:n_draws, :, :] # delete initial and exclude odd (hamiltonian bluh bluh) chain
β_chn = group(β_chn, :β)

summarize(β_chn)
quantile(β_chn)
plot(β_chn)

# Marginal Effect
"""
Marginal Effect of probit regression for continuous variable

∂Pr(Y=1|X)/∂X_j = ∂Pr(Φ(Xβ))/∂X_j
                = ϕ(Xβ)β_j, where ϕ() is standard normal pdf

The marginal effect on X_j depends on β_j, X_j and other β and X. Therefore,

    β_chn:          random drawed β
    β_interest:     variables of interest, :Symbol
                    one of a symbol in `chain` argument.
    x:              independent variable vector to calculate marginal
                    effect (individual, at mean, etc...).
    Output
    ------
    dydxs:   marginal effects regarding to each random drawed β, ((n_draws - n_tune) * n_chains) matrix
"""
function marginal_effect_probit(β_chn, β_interest, x)
    β_rand = get_params(β_chn).β # get random value
    β_rand_interest = β_chn[β_interest].data
    dydxs = []
    for which_chain = chains(β_chn)
        β_rand_single_chain = reduce(hcat, [β_rand[j].data[:, which_chain] for j=1:length(β_rand)])
        β_rand_interest_single_chain = β_rand_interest[:, which_chain]
        dydx = normpdf.(β_rand_single_chain * x) .* β_rand_interest_single_chain
        push!(dydxs, dydx)
    end
    dydxs = reduce(hcat, dydxs)
    return dydxs
end

rand_margins = marginal_effect_probit(β_chn, Symbol("β[2]"), X[100, :])
density(rand_margins)

"""
Average Marginal Effect of all sample.
"""
function marginal_effect_probit_mae(β_chn, β_interest, X)
    # https://discourse.julialang.org/t/converting-a-matrix-into-an-array-of-arrays/17038/7
    # https://discourse.julialang.org/t/how-to-broadcast-over-only-certain-function-arguments/19274
    return mean(marginal_effect_probit.((β_chn,), (β_interest,), [r for r in eachrow(X)]))
end

rand_mae = marginal_effect_probit_mae(β_chn, Symbol("β[2]"), X)
density(rand_mae)

# Plot Marginal Effect regarding to different value of X_j in interest
"""
Margins plot

Plot Marginal Effect regarding to different value of X_j in interest
# https://docs.julialang.org/en/v1/manual/functions/#Optional-Arguments
# https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments
"""
function plot_marginal_effect(β_chn::Chains, β_interest::Symbol, x_interest_theoritical::Vector; at_mean::Bool=false, ci_p::Float64=0.95)
    interest_dim = findall(x -> x == β_interest, names(β_chn))
    other_dim = findall(x -> x != β_interest, names(β_chn))
    x_interest = X[:, interest_dim]
    # https://stackoverflow.com/questions/37661221/julia-select-all-but-one-element-in-array-matrix
    x_other = X[:, 1:end .!= interest_dim]
    # x_interest_theoritical = Vector(minimum(x_interest):step:maximum(x_interest))

    if at_mean
        # at mean other x
        X_theoritical_atmean = Matrix{Float64}(undef, size(x_interest_theoritical, 1), length(names(β_chn)))
        X_theoritical_atmean[:, interest_dim] = x_interest_theoritical
        for other in other_dim
            X_theoritical_atmean[:, other] .= mean(X, dims=1)[other]
        end
        rand_margins = marginal_effect_probit.((β_chn,), (β_interest,), [r for r in eachrow(X_theoritical_atmean)])
        ylabel = @sprintf("Marginal Effect at Means, CI(%d)", ci_p*100)
    else
        # Average other x
        X_theoritical_ames = []
        for x in x_interest_theoritical
            X_theoritical_ame = copy(X)
            X_theoritical_ame[:, interest_dim] .= x
            push!(X_theoritical_ames, X_theoritical_ame)
        end
        rand_margins = marginal_effect_probit_mae.((β_chn,), (β_interest,), X_theoritical_ames)
        ylabel = @sprintf("Average Marginal Effect, CI(%d)", ci_p*100)
    end

    # margins plot
    # https://discourse.julialang.org/t/asymmetric-error-bars-and-box-plots-in-julia/73647/2
    rand_margins_median = median.(rand_margins)
    ci_lower = quantile.(rand_margins, (1-ci_p)/2)
    ci_upper = quantile.(rand_margins, ci_p-(1-ci_p)/2)
    plot(x_interest_theoritical, rand_margins_median, color="black", xlabel=@sprintf("%s's x", β_interest), ylabel=ylabel,  legend=false)
    display(scatter!(x_interest_theoritical, rand_margins_median, yerror=(ci_lower, ci_upper), color="cornflowerblue"))

    return rand_margins
end

x_hypo = Vector(-2.0:0.5:3.5)

plot_marginal_effect(β_chn, Symbol("β[2]"), x_hypo, at_mean=false)
plot_marginal_effect(β_chn, Symbol("β[2]"), x_hypo, at_mean=true)
plot_marginal_effect(β_chn, Symbol("β[2]"), x_hypo, at_mean=true, ci_p=0.5)

plot_marginal_effect(β_chn, Symbol("β[3]"), x_hypo, at_mean=false)
plot_marginal_effect(β_chn, Symbol("β[3]"), x_hypo, at_mean=true)


#=
If p is the perimeter of a right angle triangle with integral length sides, {a,b,c}, there are exactly three solutions for p = 120.

{20,48,52}, {24,45,51}, {30,40,50}

For which value of p ≤ 1000, is the number of solutions maximised?
=#
function calc()
  dict = Dict{Float64, Set}()

  for i in 1:1000, j in 1:1000
    k = i^2+j^2
    key = i + j + sqrt(k)
    if sqrt(k) != floor(sqrt(k)) || key > 1000
      continue
    end

    if haskey(dict, key) == false
      dict[key] = Set()
    else
      push!(dict[key], sort!([i, j, sqrt(k)]))
    end
  end

  (maxkey, maxcount) = 0,0
  for key in keys(dict)
    count = length(dict[key])
    if count > maxcount
      maxcount, maxkey = count, key
    end
  end
  Int64(maxkey)
end
@time println(calc())


using Distributions: Normal

"""
    MomentumHarmonicWigner(ω, β)

Wigner distribution in a 1D harmonic potential for the momentum
"""
function MomentumHarmonicWigner(ω, β)
    σ = sqrt(Q(ω, β) / β)
    Normal(0, σ)
end

"""
    PositionHarmonicWigner(ω, β)

Wigner distribution in a 1D harmonic potential for the position
"""
function PositionHarmonicWigner(ω, β, m; centre=0)
    σ = sqrt(Q(ω, β) / β / m) / ω
    Normal(centre, σ)
end

"""
    VelocityHarmonicWigner(ω, β, m)

Wigner distribution in a 1D harmonic potential for the velocity
"""
function VelocityHarmonicWigner(ω, β, m)
    σ = sqrt(Q(ω, β) / β) / m
    Normal(0, σ)
end

"Quantum corrector for the Wigner distribution"
function Q(ω, β)
    halfβħω = β*ω/2
    return halfβħω / tanh(halfβħω)
end


using LinearAlgebra, ForwardDiff
using Convex, SCS, ECOS

"""
	projection onto second-order cone
"""
function Πsoc(v, s)
	if norm(v) <= -s
		# @warn "below cone"
		return zero(v), 0.0
	elseif norm(v) <= s
		# @warn "in cone"
		return v, s
	elseif norm(v) > abs(s)
		# @warn "outside cone"
		a = 0.5 * (1.0 + s / norm(v))
		return a * v, a * norm(v)
	else
		@warn "soc projection error"
		return zero(v), 0.0
	end
end

function projection_difference(v, s)
	vcat(Πsoc(v, s)...) - [v; s]
end

# projection_difference(ones(2), 1.0)

"""
 min v'b
 st norm(b) <= y
"""

# parameters
v = [1.0; 1.0e-1]
ψ = norm(v)
y = 1.0

"Convex.jl"
b = Variable(2)
prob = minimize(v' * b)
prob.constraints += norm(b) <= y
@time solve!(prob, ECOS.Optimizer)

@show prob.status
@show b.value
@show prob.constraints[1].dual
prob.optval

function lagrangian(v, y, b, ψ)
	v' * b + ψ' * projection_difference(b, y) + 0.5 * projection_difference(b, y)' * projection_difference(b, y)
end

function r(z, θ)
	b = z[1:2]
	ψ = z[3:5]
	y = θ[1]
	v = θ[2:3]

	lb(w) = lagrangian(v, y, w, ψ)
	# lψ(w) = lagrangian(v, y, b, w)

	# return ForwardDiff.gradient(lb, b)
	return [ForwardDiff.gradient(lb, b);
			projection_difference(b, y)]

end

θ = [y; v]
rz(x) = r(x, θ)

sol = levenberg_marquardt(rz, ones(5))
rθ(x) = r(sol, x)

drdz = ForwardDiff.jacobian(rz, sol)
drdθ = ForwardDiff.jacobian(rθ, θ)

# norm(drdz - drz(z, θ))
# norm(drdθ - drθ(z, θ))
rank(drdz)
# eigen(drdz).values
#
# rank(drz(z, θ))
# eigen(drz(z, θ)).values

x1 = (-drdz \ drdθ)[1:2,:]
x3 = (-(drdz' * drdz) \ (drdz' * drdθ))[1:2,:]

ρ = 1.0e-8
x2 = (-(drdz' * drdz + ρ * I) \ (drdz' * drdθ))[1:2,:]
x4 = (-drdz' * ((drdz * drdz' + ρ * I) \ drdθ))[1:2,:]

# x2 = (-drz(z,θ) \ drθ(z,θ))[1:2,:]
# x4 = (-(drz(z,θ)' * drz(z,θ)) \ (drz(z, θ)' * drθ(z, θ)))[1:2, :]
#


doc"""
    Gamma(α,θ)

The *Gamma distribution* with shape parameter `α` and scale `θ` has probability density
function

$f(x; \alpha, \beta) = \frac{x^{\alpha-1} e^{-x/\beta}}{\Gamma(\alpha) \beta^\alpha},
\quad x > 0$

```julia
Gamma()          # Gamma distribution with unit shape and unit scale, i.e. Gamma(1.0, 1.0)
Gamma(a)         # Gamma distribution with shape a and unit scale, i.e. Gamma(a, 1.0)
Gamma(a, b)      # Gamma distribution with shape a and scale b

params(d)        # Get the parameters, i.e. (a, b)
shape(d)         # Get the shape parameter, i.e. a
scale(d)         # Get the scale parameter, i.e. b
```

External links

* [Gamma distribution on Wikipedia](http://en.wikipedia.org/wiki/Gamma_distribution)

"""
immutable Gamma <: ContinuousUnivariateDistribution
    α::Float64
    θ::Float64

    function Gamma(α::Real, θ::Real)
        @check_args(Gamma, α > zero(α) && θ > zero(θ))
        new(α, θ)
    end
    function Gamma(α::Real)
        @check_args(Gamma, α > zero(α))
        new(α, 1.0)
    end
    Gamma() = new(1.0, 1.0)
end

@distr_support Gamma 0.0 Inf

#### Parameters

shape(d::Gamma) = d.α
scale(d::Gamma) = d.θ
rate(d::Gamma) = 1.0 / d.θ

params(d::Gamma) = (d.α, d.θ)

#### Statistics

mean(d::Gamma) = d.α * d.θ

var(d::Gamma) = d.α * d.θ^2

skewness(d::Gamma) = 2.0 / sqrt(d.α)

kurtosis(d::Gamma) = 6.0 / d.α

function mode(d::Gamma)
    (α, θ) = params(d)
    α >= 1.0 ? θ * (α - 1.0) : error("Gamma has no mode when shape < 1.0")
end

function entropy(d::Gamma)
    (α, θ) = params(d)
    α + lgamma(α) + (1.0 - α) * digamma(α) + log(θ)
end

mgf(d::Gamma, t::Real) = (1.0 - t * d.θ)^(-d.α)

cf(d::Gamma, t::Real) = (1.0 - im * t * d.θ)^(-d.α)

#### Evaluation & Sampling

@_delegate_statsfuns Gamma gamma α θ

gradlogpdf(d::Gamma, x::Float64) =
    insupport(Gamma, x) ? (d.α - 1.0) / x - 1.0 / d.θ : 0.0

rand(d::Gamma) = StatsFuns.Rmath.gammarand(d.α, d.θ)

#### Fit model

immutable GammaStats <: SufficientStats
    sx::Float64      # (weighted) sum of x
    slogx::Float64   # (weighted) sum of log(x)
    tw::Float64      # total sample weight

    GammaStats(sx::Real, slogx::Real, tw::Real) = new(sx, slogx, tw)
end

function suffstats{T<:Real}(::Type{Gamma}, x::AbstractArray{T})
    sx = 0.
    slogx = 0.
    for xi = x
        sx += xi
        slogx += log(xi)
    end
    GammaStats(sx, slogx, length(x))
end

function suffstats{T<:Real}(::Type{Gamma}, x::AbstractArray{T}, w::AbstractArray{Float64})
    n = length(x)
    if length(w) != n
        throw(ArgumentError("Inconsistent argument dimensions."))
    end

    sx = 0.
    slogx = 0.
    tw = 0.
    for i = 1:n
        @inbounds xi = x[i]
        @inbounds wi = w[i]
        sx += wi * xi
        slogx += wi * log(xi)
        tw += wi
    end
    GammaStats(sx, slogx, tw)
end

function gamma_mle_update(logmx::Float64, mlogx::Float64, a::Float64)
    ia = 1.0 / a
    z = ia + (mlogx - logmx + log(a) - digamma(a)) / (abs2(a) * (ia - trigamma(a)))
    1.0 / z
end

function fit_mle(::Type{Gamma}, ss::GammaStats;
    alpha0::Float64=NaN, maxiter::Int=1000, tol::Float64=1.0e-16)

    mx = ss.sx / ss.tw
    logmx = log(mx)
    mlogx = ss.slogx / ss.tw

    a::Float64 = isnan(alpha0) ? 0.5 / (logmx - mlogx) : alpha0
    converged = false

    t = 0
    while !converged && t < maxiter
        t += 1
        a_old = a
        a = gamma_mle_update(logmx, mlogx, a)
        converged = abs(a - a_old) <= tol
    end

    Gamma(a, mx / a)
end