"""
    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


if !@isdefined TEST
    using Revise
end
using Mitosis
using Random, Test, LinearAlgebra, Statistics
Random.seed!(1)

# Prior (as kernel of prior mean)

μ0 = zeros(2)
σ2 = 8.0 # noise
V0 = 10*I(2)
Σ0 = σ2*V0 # prior

prior = kernel(Gaussian; μ=ConstantMap(μ0), Σ=ConstantMap(Σ0))

# Data
x = [18.25 19.75 16.5 18.25 19.50 16.25 17.25 19.00 16.25 17.50][:]
y = [36 42 33 39 43 34 37 41 27 30][:]
n = length(x)

# Model
X = [x ones(n)]
Σ = Diagonal(σ2*ones(n))
model = kernel(Gaussian; μ=LinearMap(X), Σ=ConstantMap(Σ))

# backward pass
m2, p2 = backward(BF(), model, y)
m1, p1 = backward(BF(), prior, p2)

posterior = forward(BF(), prior, m1)()

@show mean(posterior), cov(posterior)

β̂ = X\y
@test β̂ ≈ mean(p2)

@test mean(posterior) ≈ inv(X'*X + inv(V0))*X'*y
@test cov(posterior) ≈ σ2*inv(X'*X + inv(V0))


using jInv.LinearSolvers
using Test
using ConvDiffMIPDECO
using jInv.ForwardShare
using jInv.Mesh
using LinearAlgebra
using SparseArrays
using MUMPSjInv

domain = [0 1. 0 1]

N = 2*[16;32;64;128]
v = randn(2)
v = v/norm(v)
V = X -> [v[1]*ones(size(X,1)) v[2]*ones(size(X,1))]
bc = (:neu,:dir,:neu,:dir)

u  = X -> cos.(pi*X[:,1]).*cos.(pi*X[:,2])

ux = X -> -pi*sin.(pi*X[:,1]).*cos.(pi*X[:,2])
uy = X -> -pi*cos.(pi*X[:,1]).*sin.(pi*X[:,2])
sig = .001;

f  = X->  +2*sig*pi^2*u(X)+v[1]*ux(X)+v[2]*uy(X)

err = zeros(length(N))

bicg = (A,b; M=identity,tol=1e-10,maxIter=500,out=1)-> bicgstb(A,b,M1=identity,tol=tol,maxIter=maxIter,out=out)
times = zeros(length(N),3)
for j=1:2
	for k=1:length(N)
		Ainvs = (getJuliaSolver(), getMUMPSsolver(), ConvDiffMIPDECO.getBICGSTB(out=1))

	nk      = [N[k];N[k]]
	Mk  = getRegularMesh(domain,nk)
	xck = getCellCenteredGrid(Mk)

	pFork   = getConvDiffParam(Mk,V,sig=sig,gd=u,gn=ux,bc=bc,Ainv=Ainvs[j])

	dobsk,pFork = getData(f(xck), pFork)

	utruek = u(xck)
	times[k,j]= @elapsed begin
		err[k] = norm(pFork.Fields - utruek,Inf)/norm(utruek,Inf)
	end
	println("err[$k] = $(err[k])")
end
@test all(abs.(diff(log2.(err))).>0.8)
end


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


using BigCombinatorics
using Test

@test Fibonacci(10)==Fibonacci(9)+Fibonacci(8)

@test Factorial(10)==big(factorial(10))

@test FallingFactorial(10,3) == 10*9*8
@test FallingFactorial(10,10) == Factorial(10)
@test FallingFactorial(10,12) == 0

@test RisingFactorial(10,3) == 10*11*12

@test DoubleFactorial(9) == 9*7*5*3
@test DoubleFactorial(10) == 10*8*6*4*2

@test Catalan(12) == 208012

@test Derangements(1) == 0
@test Derangements(12) == 176214841

@test MultiChoose(10,1) == 10
@test MultiChoose(10,0) == 1
@test MultiChoose(10,10) == 92378

@test Multinomial(5,5,5) == div(Factorial(15), Factorial(5)^3)
@test Multinomial([6,6,6,6]) == Multinomial(6,6,6,6)

@test Bell(10)==115975
@test sum(Stirling2(10,k) for k=0:10) == Bell(10)
@test Stirling1(10,10) == 1
@test Stirling1(10,0) == 0
@test sum(Stirling1(10,k) for k=0:10) == 0

@test IntPartitions(10) == 42
@test IntPartitionsDistinct(10) == 10

@test Euler(12) == 2702765

@test PowerSum(10,3) == sum(k^3 for k=1:10)

@test BigCombinatorics.cache_clear(Fibonacci)

@test sum(Eulerian(10,k) for k=1:10) == Factorial(10)


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())


# Tests for Cost Hölder DEA Models
@testset "CostHölderDEAModel" begin

    # Test using Book data
    X = [2 2; 1 4; 4 1; 4 3; 5 5; 6 1; 2 5; 1.6 8];
    Y = [1; 1; 1; 1; 1; 1; 1; 1];
    W = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1];

    # Cost Hölder L1
    costholderl1 = deacostholder(X, Y, W, l = 1)

    @test typeof(costholderl1) == CostHolderDEAModel

    @test nobs(costholderl1) == 8
    @test ninputs(costholderl1) == 2
    @test noutputs(costholderl1) == 1
    @test ismonetary(costholderl1) == false

    @test efficiency(costholderl1) == efficiency(costholderl1, :Economic)
    @test efficiency(costholderl1, :Economic)   ≈ [0.0; 1.0; 1.0; 3.0; 6.0; 3.0; 3.0; 5.6]
    @test efficiency(costholderl1, :Technical)  ≈ [0.0; 0.0; 0.0; 2.0; 4.0; 0.0; 1.0; 0.6]
    @test efficiency(costholderl1, :Allocative) ≈ [0.0; 1.0; 1.0; 1.0; 2.0; 3.0; 2.0; 5.0]
    @test normfactor(costholderl1) ≈ [1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]

    # With monetary = true
    costholderl1mom = deacostholder(X, Y, W, l = 1, monetary = true)

    @test ismonetary(costholderl1mom) == true

    @test efficiency(costholderl1mom) ≈ efficiency(costholderl1) .* normfactor(costholderl1)
    @test efficiency(costholderl1mom, :Technical) ≈ efficiency(costholderl1, :Technical) .* normfactor(costholderl1)
    @test efficiency(costholderl1mom, :Allocative) ≈ efficiency(costholderl1, :Allocative) .* normfactor(costholderl1)

    # With CRS (Same result in this example)
    costholderl1crs = deacostholder(X, Y, W, l = 1, rts = :CRS)

    @test efficiency(costholderl1crs, :Economic)   ≈ [0.0; 1.0; 1.0; 3.0; 6.0; 3.0; 3.0; 5.6]
    @test efficiency(costholderl1crs, :Technical)  ≈ [0.0; 0.0; 0.0; 2.0; 4.0; 0.0; 1.0; 0.6]
    @test efficiency(costholderl1crs, :Allocative) ≈ [0.0; 1.0; 1.0; 1.0; 2.0; 3.0; 2.0; 5.0]

    # Cost Hölder L1 Weighted (weakly)
    costholderl1weight= deacostholder(X, Y, W, l = 1, weight = true)

    @test efficiency(costholderl1weight, :Economic)   ≈ [0.0; 0.25; 0.25; 0.75; 1.2; 0.5; 0.6; 0.7]
    @test efficiency(costholderl1weight, :Technical)  ≈ [0.0; 0.0; 0.0; 0.625; 0.8; 0.0; 0.5; 0.375]
    @test efficiency(costholderl1weight, :Allocative) ≈ [0.0; 0.25; 0.25; 0.125; 0.4; 0.5; 0.1; 0.325]

    # Cost Hölder L2
    logs, value = Test.collect_test_logs() do
        costholderl2 = deacostholder(X, Y, W, l = 2, optimizer = DEAOptimizer(:NLP))

        @test normfactor(costholderl2) ≈ sqrt(2) .* ones(8)
        show(IOBuffer(), costholderl2)
    end

    # Cost Hölder L2 Weighted (weakly)
    logs, value = Test.collect_test_logs() do
        costholderl2weight = deacostholder(X, Y, W, l = 2, weight = true, optimizer = DEAOptimizer(:NLP))

        @test normfactor(costholderl2weight) ≈ sqrt.(sum((W .* X).^2, dims = 2))
        show(IOBuffer(), costholderl2weight)
    end

    # Cost Hölder LInf
    costholderlInf = deacostholder(X, Y, W, l = Inf)

    @test efficiency(costholderlInf, :Economic)   ≈ [0.0; 0.5; 0.5; 1.5; 3.0; 1.5; 1.5; 2.8]
    @test efficiency(costholderlInf, :Technical)  ≈ [0.0; 0.0; 0.0; 4/3; 3.0; 0.0; 1.0; 0.6]
    @test efficiency(costholderlInf, :Allocative) ≈ [0.0; 0.5; 0.5; 1/6; 0.0; 1.5; 0.5; 2.2]
    @test normfactor(costholderlInf) ≈ [2.0; 2.0; 2.0; 2.0; 2.0; 2.0; 2.0; 2.0]

    # Cost Hölder LInf Weighted (weakly)
    costholderlInfweight = deacostholder(X, Y, W, l = Inf, weight = true)

    @test efficiency(costholderlInfweight, :Economic)   ≈ [0.0; 0.2; 0.2; 3/7; 0.6; 3/7; 3/7; 7/12]
    @test efficiency(costholderlInfweight, :Technical)  ≈ [0.0; 0.0; 0.0; 0.4; 0.6; 0.0; 1/3; 3/8]
    @test efficiency(costholderlInfweight, :Allocative) ≈ [0.0; 0.2; 0.2; 1/35; 0.0; 3/7; 2/21; 5/24]
    @test normfactor(costholderlInfweight) ≈ [4.0; 5.0; 5.0; 7.0; 10.0; 7.0; 7.0; 9.6]

    # Print
    show(IOBuffer(), costholderl1)
    show(IOBuffer(), costholderl1weight)
    show(IOBuffer(), costholderlInf)
    show(IOBuffer(), costholderlInfweight)

    # Test errors
    @test_throws DimensionMismatch deacostholder([1; 2 ; 3], [4 ; 5], [1; 1; 1], l = 1) #  Different number of observations
    @test_throws DimensionMismatch deacostholder([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], l = 1) # Different number of observation in prices
    @test_throws DimensionMismatch deacostholder([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], l = 1) # Different number of input prices and inputs
    @test_throws ArgumentError deacostholder([1; 2; 3], [4; 5; 6], [1; 2; 3], l = 5) # Invalid l
    @test_throws ArgumentError deacostholder([1; 2; 3], [4; 5; 6], [1; 2; 3], l = 1, rts = :Error) # Invalid returns to scale

end


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


cd(@__DIR__)
Pkg.activate("..")
using Chordal
using JuMP, Hypatia, COSMO
using SparseArrays, LinearAlgebra, Random

const CD = Chordal

rand_seed = 1234
n = 200
# optimizer = Hypatia.Optimizer
optimizer = optimizer_with_attributes(COSMO.Optimizer, "decompose" => false)

# min c^Tx st ΣF_ix_i + G ∈ PSDCONE()
# D is dual variable
c, F, G, ystar, D = CD.generate_random_sdp(n, rand_seed=rand_seed)

# ------------------------------------------------------
# ------------------ Standard Form SDP -----------------
# ------------------------------------------------------
## Without decomposition
m3 = Model(optimizer)
sparse(@variable(m3, X[1:n, 1:n] in PSDCone()))
for i in 1:n
    @constraint(m3, sum(F[i] .* X) - c[i] == 0)
end
@objective(m3, Max, -sum(G .* X))
time_non_chordal = @elapsed optimize!(m3)

Xv = value.(X)
@info "termination status: $(termination_status(m3))"
@info "solution status: $(primal_status(m3))"
@info "Optimal value: $(round(objective_value(m3), digits=3))"

## Chordal Setup
m4 = Model(optimizer)
perm, iperm, Tls, cliques = get_selectors(sparsity_pattern([F..., G]); verbose=false, ret_cliques=true)
sp = CD.get_sparsity_pattern_from_cliques(cliques)
sp_inds = zip(findnz(sp)[1:2]...)
X = @variable(m4, X[i=1:n, j=1:n; (i,j) in sp_inds])

# TODO: Package this up into a function?
for i in 1:length(c)
    @constraint(m4, CD.tr(F[i][perm, perm], X) == c[i])
end

for (p, cl) in enumerate(cliques)
    len_cl = length(cl)
    Xp = @variable(m4, [1:len_cl, 1:len_cl], base_name="X$p")
    for i in 1:len_cl, j in 1:len_cl
        Xp[i,j] = X[cl[i], cl[j]]
    end
    @constraint(m4, Xp in PSDCone())
end

@objective(m4, Max, -CD.tr(G[perm, perm], X))

time_chordal = @elapsed optimize!(m4)
@info "termination status: $(termination_status(m4))"
@info "solution status: $(primal_status(m4))"
match = objective_value(m4) ≈ objective_value(m3)
@info "Optimal value: $(round(objective_value(m4), digits=3)), equals: $(match)"
@info "Time with chordal is $(round(time_chordal, digits=3))s vs $(round(time_non_chordal, digits=3))s"

X_ = CD.reconstruct_from_sparse_varref(X, n)
@assert maximum(X_ - X_') < 1e-8
X_ = 0.5*(X_ + X_')
Xcomp = maxdet_completion(X_)[iperm, iperm]
@assert sum(-G .* Xcomp) ≈ objective_value(m4)


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