using Test
using Statistics

include("../src/UncValue.jl")
using ..UncValue

a = Value(3.1415, 0.0012)
b = Value(2.7182818, 3.4e-6)
c = Value(36458.246, 25.64)

@testset "constructor" begin
    @test a.x == 3.1415
    @test a.σ == 0.0012

    @test_throws ErrorException Value(4.1, -2.4)

    d = Value(5.248)
    @test d.x == 5.248
    @test d.σ == 0

    d = Value(8.0, 1)
    @test typeof(d.x) == typeof(d.σ)

    d = Value(8, 0.2)
    @test typeof(d.x) == typeof(d.σ)
end

@testset "precision" begin
    @test precision(a) == -4
    @test precision(b) == -7
    @test precision(c) == 0
end

@testset "addition" begin
    @test typeof(a + 2) <: Value
    @test typeof(2 + a) <: Value
    s = a + b
    @test typeof(s) <: Value
    @test s.x == (a.x + b.x)
    @test s.σ ≈ hypot(a.σ, b.σ)

    n = -a
    @test typeof(n) <: Value
    @test n.x == (-a.x)
    @test n.σ == a.σ
end

@testset "substraction" begin
    @test typeof(a - 2) <: Value
    @test typeof(2 - a) <: Value
    s = a - b
    @test typeof(s) <: Value
    @test s.x ≈ (a.x - b.x)
    @test s.σ ≈ hypot(a.σ, b.σ)
end

@testset "product" begin
    @test typeof(a * 2) <: Value
    @test typeof(2 * a) <: Value
    s = a * b
    @test typeof(s) <: Value
    @test s.x ≈ (a.x * b.x)
    @test s.σ ≈ hypot(a.σ * b.x, a.x * b.σ)
end

@testset "division" begin
    @test typeof(a / 2) <: Value
    @test typeof(2 / a) <: Value
    s = a / b
    @test typeof(s) <: Value
    @test s.x ≈ (a.x / b.x)
    @test s.σ ≈ hypot(a.σ / b.x, a.x * b.σ / b.x^2)
end

@testset "left division" begin
    s = a \ 2
    @test typeof(s) <: Value
    @test s.x ≈ (a.x \ 2)
    @test s.σ ≈ a.x^2 \ 2 * a.σ

    s = 2 \ a
    @test typeof(s) <: Value
    @test s.x ≈ (2 \ a.x)
    @test s.σ ≈ 2 \ a.σ

    s = a \ b
    @test typeof(s) <: Value
    @test s.x ≈ (a.x \ b.x)
    @test s.σ ≈ hypot(a.x \ b.σ, a.x^2 \ b.x * a.σ)
end

@testset "integer division" begin
    j = Value{Int32}(2, 1)
    k = Value{Int32}(7, 2)

    s = j ÷ k.x
    @test typeof(s) <: Value
    @test s.x ≈ j.x ÷ k.x
    @test s.σ ≈ j.σ ÷ k.x

    s = k.x ÷ j
    @test typeof(s) <: Value
    @test s.x ≈ k.x ÷ j.x
    @test s.σ ≈ k.x * j.σ ÷ j.x^2

    s = k ÷ j
    @test typeof(s) <: Value
    @test s.x ≈ k.x ÷ j.x
    @test s.σ ≈ hypot(k.x * j.σ ÷ j.x^2, j.σ ÷ k.x)

    s = fld(j, k.x)
    @test typeof(s) <: Value
    @test s.x ≈ fld(j.x, k.x)
    @test s.σ ≈ fld(j.σ, k.x)

    s = fld(k.x, j)
    @test typeof(s) <: Value
    @test s.x ≈ fld(k.x, j.x)
    @test s.σ ≈ fld(k.x * j.σ, j.x^2)

    s = fld(k, j)
    @test typeof(s) <: Value
    @test s.x ≈ fld(k.x, j.x)
    @test s.σ ≈ floor(hypot(fld(k.x * j.σ, j.x^2), fld(j.σ, k.x)))

    s = cld(j, k.x)
    @test typeof(s) <: Value
    @test s.x ≈ cld(j.x, k.x)
    @test s.σ ≈ cld(j.σ, k.x)

    s = cld(k.x, j)
    @test typeof(s) <: Value
    @test s.x ≈ cld(k.x, j.x)
    @test s.σ ≈ cld(k.x * j.σ, j.x^2)

    s = cld(k, j)
    @test typeof(s) <: Value
    @test s.x ≈ cld(k.x, j.x)
    @test s.σ ≈ ceil(hypot(cld(k.x * j.σ, j.x^2), cld(j.σ, k.x)))
end

@testset "power" begin
    @test typeof(a^3) <: Value
    @test typeof(3^a) <: Value
    s = a^b
    @test typeof(s) <: Value
    @test s.x ≈ a.x^b.x
    @test s.σ ≈ hypot(b.x * a.x^(b.x-1) * a.σ, a.x^b.x * log(abs(a.x)) * b.σ)
end

@testset "equality" begin
    @test a == 3.1415
    @test (a == 3.141) == false
    @test a != 3.141

    @test 3.1415 == a
    @test (3.141 == a) == false
    @test 3.141 != a

    @test (a == b) == false
    @test a != b
end

@testset "inequality" begin
    @test a < 3.1416
    @test (a < 3.141) == false

    @test 3.1416 > a
    @test (3.141 > a) == false

    @test a <= 3.1415
    @test a <= 3.1416
    @test (a <= 3.141) == false

    @test 3.1416 >= a
    @test 3.1415 >= a
    @test (3.141 >= a) == false

    @test a > 3.141
    @test (a > 3.1416) == false

    @test 3.141 < a
    @test (3.1416 < a) == false

    @test a >= 3.1415
    @test a >= 3.1414
    @test (a >= 3.1416) == false

    @test 3.1414 <= a
    @test 3.1415 <= a
    @test (3.146 <= a) == false

    @test a > b
    @test b < a

    @test a >= a
    @test a >= b
    @test (b >= a) == false
    @test a <= a
    @test b <= a
    @test (a <= b) == false
end

@testset "abs" begin
    s = abs(a)
    @test typeof(s) <: Value
    @test s.x == a.x
    @test s.σ == a.σ

    s = abs(-a)
    @test s.x == a.x
    @test s.σ == a.σ
end

@testset "abs2" begin
    s = abs2(a)
    @test typeof(s) <: Value
    @test s.x == a.x^2
    @test s.σ == 2 * a.x * a.σ

    s = abs2(-a)
    @test s.x == a.x^2
    @test s.σ == 2 * a.x * a.σ
end

@testset "sin/cos/tan" begin
    s = sin(b)
    @test typeof(s) <: Value
    @test s.x ≈ sin(b.x)
    @test s.σ ≈ abs(cos(b.x)) * b.σ

    c = cos(b)
    @test typeof(c) <: Value
    @test c.x ≈ cos(b.x)
    @test c.σ ≈ abs(sin(b.x)) * b.σ

    t = tan(b)
    @test typeof(t) <: Value
    @test t.x ≈ tan(b.x)
    @test t.σ ≈ sec(b.x)^2 * b.σ
end

@testset "sind/cosd/tand" begin
    s = sind(b)
    @test typeof(s) <: Value
    @test s.x ≈ sind(b.x)
    @test s.σ ≈ abs(cosd(b.x)) * b.σ

    c = cosd(b)
    @test typeof(c) <: Value
    @test c.x ≈ cosd(b.x)
    @test c.σ ≈ abs(sind(b.x)) * b.σ

    t = tand(b)
    @test typeof(t) <: Value
    @test t.x ≈ tand(b.x)
    @test t.σ ≈ secd(b.x)^2 * b.σ
end

@testset "sinh/cosh/tanh" begin
    s = sinh(b)
    @test typeof(s) <: Value
    @test s.x ≈ sinh(b.x)
    @test s.σ ≈ cosh(b.x) * b.σ

    c = cosh(b)
    @test typeof(c) <: Value
    @test c.x ≈ cosh(b.x)
    @test c.σ ≈ abs(sinh(b.x)) * b.σ

    t = tanh(b)
    @test typeof(t) <: Value
    @test t.x ≈ tanh(b.x)
    @test t.σ ≈ sech(b.x)^2 * b.σ
end

@testset "sinpi/cospi" begin
    s = sinpi(b)
    @test typeof(s) <: Value
    @test s.x ≈ sinpi(b.x)
    @test s.σ ≈ π * abs(cospi(b.x)) * b.σ

    c = cospi(b)
    @test typeof(c) <: Value
    @test c.x ≈ cospi(b.x)
    @test c.σ ≈ π * abs(sinpi(b.x)) * b.σ
end

@testset "asin/acos/atan" begin
    d = Value(0.2435, 0.0658)

    s = asin(d)
    @test typeof(s) <: Value
    @test s.x ≈ asin(d.x)
    @test s.σ ≈ d.σ / sqrt(1 - d.x^2)

    c = acos(d)
    @test typeof(c) <: Value
    @test c.x ≈ acos(d.x)
    @test c.σ ≈ d.σ / sqrt(1 - d.x^2)

    t = atan(d)
    @test typeof(t) <: Value
    @test t.x ≈ atan(d.x)
    @test t.σ ≈ d.σ / (1 + d.x^2)

    t = atan(a, b)
    @test t.x ≈ atan(a.x, b.x)
    @test t.σ ≈ hypot(b.x * a.σ, a.x * b.σ) / hypot(a.x, b.x)^2

    t = atan(a, 0.1)
    @test t.x ≈ atan(a.x, 0.1)
    @test t.σ ≈ 0.1 * a.σ / hypot(a.x, 0.1)^2

    t = atan(0.1, a)
    @test t.x ≈ atan(0.1, a.x)
    @test t.σ ≈ 0.1 * a.σ / hypot(a.x, 0.1)^2
end

@testset "asind/acosd/atand" begin
    d = Value(0.2435, 0.0658)

    s = asind(d)
    @test typeof(s) <: Value
    @test s.x ≈ asind(d.x)
    @test s.σ ≈ d.σ / sqrt(1 - d.x^2)

    c = acosd(d)
    @test typeof(c) <: Value
    @test c.x ≈ acosd(d.x)
    @test c.σ ≈ d.σ / sqrt(1 - d.x^2)

    t = atand(d)
    @test typeof(t) <: Value
    @test t.x ≈ atand(d.x)
    @test t.σ ≈ d.σ / (1 + d.x^2)

    t = atand(a, b)
    @test t.x ≈ atand(a.x, b.x)
    @test t.σ ≈ hypot(b.x * a.σ, a.x * b.σ) / hypot(a.x, b.x)^2

    t = atand(a, 0.1)
    @test t.x ≈ atand(a.x, 0.1)
    @test t.σ ≈ 0.1 * a.σ / hypot(a.x, 0.1)^2

    t = atand(0.1, a)
    @test t.x ≈ atand(0.1, a.x)
    @test t.σ ≈ 0.1 * a.σ / hypot(a.x, 0.1)^2
end

@testset "asinh/acosh/atanh" begin
    d = Value(1.2435, 0.0658)

    s = asinh(d)
    @test typeof(s) <: Value
    @test s.x ≈ asinh(d.x)
    @test s.σ ≈ d.σ / sqrt(1 + d.x^2)

    c = acosh(d)
    @test typeof(c) <: Value
    @test c.x ≈ acosh(d.x)
    @test c.σ ≈ d.σ / sqrt(d.x^2 - 1)

    d -= 1
    t = atanh(d)
    @test typeof(t) <: Value
    @test t.x ≈ atanh(d.x)
    @test t.σ ≈ d.σ / (1 - d.x^2)
end

@testset "csc/sec/cot" begin
    s = csc(b)
    @test typeof(s) <: Value
    @test s.x ≈ csc(b.x)
    @test s.σ ≈ abs(cot(b.x) * csc(b.x)) * b.σ

    c = sec(b)
    @test typeof(c) <: Value
    @test c.x ≈ sec(b.x)
    @test c.σ ≈ abs(tan(b.x) * sec(b.x)) * b.σ

    t = cot(b)
    @test typeof(t) <: Value
    @test t.x ≈ cot(b.x)
    @test t.σ ≈ csc(b.x)^2 * b.σ
end

@testset "cscd/secd/cotd" begin
    s = cscd(b)
    @test typeof(s) <: Value
    @test s.x ≈ cscd(b.x)
    @test s.σ ≈ abs(cotd(b.x) * cscd(b.x)) * b.σ

    c = secd(b)
    @test typeof(c) <: Value
    @test c.x ≈ secd(b.x)
    @test c.σ ≈ abs(tand(b.x) * secd(b.x)) * b.σ

    t = cotd(b)
    @test typeof(t) <: Value
    @test t.x ≈ cotd(b.x)
    @test t.σ ≈ cscd(b.x)^2 * b.σ
end

@testset "csch/sech/coth" begin
    s = csch(b)
    @test typeof(s) <: Value
    @test s.x ≈ csch(b.x)
    @test s.σ ≈ abs(coth(b.x) * csch(b.x)) * b.σ

    c = sech(b)
    @test typeof(c) <: Value
    @test c.x ≈ sech(b.x)
    @test c.σ ≈ abs(tanh(b.x) * sech(b.x)) * b.σ

    t = coth(b)
    @test typeof(t) <: Value
    @test t.x ≈ coth(b.x)
    @test t.σ ≈ csch(b.x)^2 * b.σ
end

@testset "acsc/asec/acot" begin
    s = acsc(b)
    @test typeof(s) <: Value
    @test s.x ≈ acsc(b.x)
    @test s.σ ≈ b.σ / sqrt(b.x^2 - 1)

    c = asec(b)
    @test typeof(c) <: Value
    @test c.x ≈ asec(b.x)
    @test c.σ ≈ b.σ / sqrt(b.x^2 - 1)

    t = acot(b)
    @test typeof(t) <: Value
    @test t.x ≈ acot(b.x)
    @test t.σ ≈ b.σ / (1 + b.x^2)
end

@testset "acscd/asecd/acotd" begin
    s = acscd(b)
    @test typeof(s) <: Value
    @test s.x ≈ acscd(b.x)
    @test s.σ ≈ b.σ / sqrt(b.x^2 - 1)

    c = asecd(b)
    @test typeof(c) <: Value
    @test c.x ≈ asecd(b.x)
    @test c.σ ≈ b.σ / sqrt(b.x^2 - 1)

    t = acotd(b)
    @test typeof(t) <: Value
    @test t.x ≈ acotd(b.x)
    @test t.σ ≈ b.σ / (1 + b.x^2)
end

@testset "acsch/asech/acoth" begin
    d = Value(0.12435, 0.0658)
    s = acsch(b)
    @test typeof(s) <: Value
    @test s.x ≈ acsch(b.x)
    @test s.σ ≈ b.σ / sqrt(b.x^2 + 1)

    c = asech(d)
    @test typeof(c) <: Value
    @test c.x ≈ asech(d.x)
    @test c.σ ≈ d.σ / sqrt(1 - d.x^2)

    t = acoth(b)
    @test typeof(t) <: Value
    @test t.x ≈ acoth(b.x)
    @test t.σ ≈ b.σ / abs(1 - b.x^2)
end

@testset "deg2rad" begin
    d = deg2rad(a)
    @test typeof(d) <: Value
    @test d.x == deg2rad(a.x)
    @test d.σ == deg2rad(a.σ)

    d = rad2deg(a)
    @test typeof(d) <: Value
    @test d.x == rad2deg(a.x)
    @test d.σ == rad2deg(a.σ)
end

@testset "exp" begin
    r = exp(a)
    @test typeof(r) <: Value
    @test r.x == exp(a.x)
    @test r.σ ≈ r.x * a.σ

    r = exp2(a)
    @test typeof(r) <: Value
    @test r.x == exp2(a.x)
    @test r.σ ≈ r.x * a.σ * log(2)

    r = exp10(a)
    @test typeof(r) <: Value
    @test r.x == exp10(a.x)
    @test r.σ ≈ r.x * a.σ * log(10)
end

@testset "log" begin
    r = log(a)
    @test typeof(r) <: Value
    @test r.x == log(a.x)
    @test r.σ ≈ a.σ / a.x

    r = log2(a)
    @test typeof(r) <: Value
    @test r.x == log2(a.x)
    @test r.σ ≈ a.σ / (a.x * log(2))

    r = log10(a)
    @test typeof(r) <: Value
    @test r.x == log10(a.x)
    @test r.σ ≈ a.σ / (a.x * log(10))

    r = log(3, a)
    @test typeof(r) <: Value
    @test r.x == log(3, a.x)
    @test r.σ ≈ a.σ / (a.x * log(3))

    r = log1p(a)
    @test typeof(r) <: Value
    @test r.x == log1p(a.x)
    @test r.σ ≈ a.σ / (a.x + 1)
end

@testset "roots" begin
    r = sqrt(a)
    @test typeof(r) <: Value
    @test r.x == sqrt(a.x)
    @test r.σ == a.σ / (2 * sqrt(a.x))

    r = cbrt(a)
    @test typeof(r) <: Value
    @test r.x == cbrt(a.x)
    @test r.σ == a.σ / (3 * cbrt(a.x)^2)

    h = hypot(a, 3)
    @test typeof(h) <: Value
    @test h.x == hypot(a.x, 3)
    @test h.σ == a.σ * a.x / h

    h = hypot(3, a)
    @test typeof(h) <: Value
    @test h.x == hypot(3, a.x)
    @test h.σ == a.σ * a.x / h

    h = hypot(a, b)
    @test typeof(h) <: Value
    @test h.x == hypot(a.x, b.x)
    @test h.σ == hypot(a.σ * a.x, b.σ * b.x) / h
end

@testset "sign" begin
    @test sign(a) == sign(a.x)
    @test sign(Value(-2.4, 2.4)) == sign(-2.4)

    @test signbit(a) == signbit(a.x)
    @test signbit(Value(-2.4, 2.4)) == signbit(-2.4)
end

@testset "inv" begin
    r = inv(a)
    @test typeof(r) <: Value
    @test r.x == inv(a.x)
    @test r.σ == a.σ / a.x^2
end

@testset "approx" begin
    @test a ≈ 3.141
    @test (a ≈ 3.145) == false
    @test isapprox(a, 3.145; significance=3)

    @test 3.141 ≈ a
    @test (3.145 ≈ a) == false
    @test isapprox(3.145, a; significance=3)

    @test (a ≈ b) == false
end

@testset "cmp" begin
    @test cmp(a, 3) == 1
    @test cmp(3, a) == -1
    @test cmp(a, a) == 0
    @test cmp(a, b) == 1
    @test cmp(b, a) == -1
end

@testset "isless" begin
    @test isless(a, b) == false
    @test isless(b, a)
    @test isless(a, 3) == false
    @test isless(3, a)
end

@testset "clamp" begin
    @test clamp(3, a) == a.x - a.σ
    @test clamp(4, a) == a.x + a.σ
    @test clamp(3.1414, a) == 3.1414
end

@testset "min/max" begin
    @test min(a, b) == min(b, a) == b
    @test max(a, b) == max(b, a) == a
    @test min(a, 3) == min(3, a) == 3
    @test max(a, 4) == max(4, a) == 4

    @test typeof(min(a, 4)) <: Value
    @test typeof(max(a, 3)) <: Value

    @test max(a) == a.x + a.σ
    @test min(a) == a.x - a.σ

    @test min(a, b, c, -2, 10) == -2
    @test min(a, b, c, 34, 13, 9) == b
    @test max(a, b, c, 2, 10) == c
    @test max(a, 4, b, 2, 1) == 4
end

@testset "val/unc" begin
    @test val(a.x) == a.x
    @test unc(a.x) == 0

    @test val(a) == a.x
    @test unc(a) == a.σ

    A = fill(Value(a.x, a.σ), (5, 2, 6))

    @test eltype(A) <: Value
    @test eltype(val(A)) <: Real
    @test mean(val(A)) ≈ a.x
    @test mean(unc(A)) ≈ a.σ
end

@testset "set_unc" begin
    A = zeros(3, 7, 2)

    uncA = set_unc(A, 0.2)
    @test eltype(uncA) <: Value
    @test mean(unc(uncA)) ≈ 0.2

    uncA = set_unc(A, fill(0.3, (3, 7, 2)))
    @test eltype(uncA) <: Value
    @test mean(unc(uncA)) ≈ 0.3

    A = fill(Value(a.x, a.σ), (5, 2, 6))
    uncA = set_unc(A, 0.2)
    @test eltype(uncA) <: Value
    @test mean(unc(uncA)) ≈ 0.2

    uncA = set_unc(A, fill(0.3, (5, 2, 6)))
    @test eltype(uncA) <: Value
    @test mean(unc(uncA)) ≈ 0.3
end


using UnPack, OrderedCollections, LinearAlgebra, JLD2, SparseArrays

# Load guesses
sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "..", "test", "reference", "crw_sss.jld2"), "r")

zguess = vec(sssout["z_rss"])
yguess = vec(sssout["y_rss"])
Psiguess = sssout["Psi_rss"]

mutable struct CoeurdacierReyWinant{T <: Real}
    σr::T # SD of interest rate shocks
    σy::T # SD of endowment shocks
    β::T  # intertemporal discount rate
    γ::T  # risk aversion coefficient
    θ::T
    ρr::T # persistence of interest rate
    ρy::T # persistence of endowment
    rr::T # long-run interest rate
    yy::T # long-run endowment
end

function CoeurdacierReyWinant(; σr::T = .025, σy::T = .025, β::T = .96, γ::T = 2.,
                              θ::T = 1., ρr::T = .9, ρy::T = .9, rr::T = .01996, yy::T = log(θ)) where {T <: Real}
    return CoeurdacierReyWinant{T}(σr, σy, β, γ, θ, ρr, ρy, rr, yy)
end

function crw(m::CoeurdacierReyWinant{T}; Ψ = nothing, sparse_jacobian::Vector{Symbol} = Symbol[],
             sparse_arrays::Bool = false) where {T <: Real}
    @unpack σr, σy, β, γ, θ, ρr, ρy, rr, yy = m

    # Nₜ = exp(rₜ) * Aₜ₋₁ + Yₜ, where Aₜ is foreign assets and Yₜ is the endowment
    # The jump variables are consumption, expected return on assets Xₜ = 𝔼ₜ[Rₜ₊₁], and
    # Wₜ = 𝔼ₜ[Yₜ₊₁]
    S  = OrderedDict{Symbol, Int}(:N => 1, :r => 2, :y => 3) # State variables
    J  = OrderedDict{Symbol, Int}(:c => 1, :x => 2, :w => 3) # Jump variables
    SH = OrderedDict{Symbol, Int}(:εr => 1, :εy => 2)        # Exogenous shocks
    Nz = length(S)
    Ny = length(J)
    Nε = length(SH)

    function μ(F, z, y) # note that y here refers to jump variables
        F[S[:N]] = exp(y[J[:w]]) + exp(y[J[:x]]) * (z[S[:N]] - exp(y[J[:c]]))
        F[S[:r]] = (1 - ρr) * rr + ρr * z[S[:r]]
        F[S[:y]] = (1 - ρy) * yy + ρy * z[S[:y]]
    end

    function ξ(F, z, y)
        F[J[:c]] = log(β) + γ * y[J[:c]]    # Euler equation
        F[J[:x]] = -y[J[:x]]                # rₜ₊₁ - xₜ, rational expectations
        F[J[:w]] = -y[J[:w]]                # yₜ₊₁ - wₜ
    end

    # The cache is initialized as zeros so we only need to fill non-zero elements
    function Λ(F, z, y)
        F[S[:N], J[:x]] = z[S[:N]] - exp(y[J[:c]])
        F[S[:N], J[:w]] = 1.
    end

    # The cache is initialized as zeros so we only need to fill non-zero elements
    function Σ(F, z, y)
        F[S[:r], SH[:εr]] = σr
        F[S[:y], SH[:εy]] = σy
    end

    Γ₅ = zeros(T, Ny, Nz)
    Γ₅[J[:c], S[:r]] = 1.
    Γ₅[J[:x], S[:r]] = 1.
    Γ₅[J[:w], S[:y]] = 1.

    Γ₆ = zeros(T, Ny, Ny)
    Γ₆[J[:c], J[:c]] = -γ

    if sparse_arrays
        Γ₅ = sparse(Γ₅)
        Γ₆ = sparse(Γ₆)
    end

    z = zguess
    y = yguess
    if isnothing(Ψ)
        Ψ = Psiguess
    end

    if sparse_arrays
        return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, crw_ccgf, z, y, Ψ, Nε; sparse_jacobian = sparse_jacobian,
                                         Λ_cache_init = dims -> spzeros(dims...),
                                         Σ_cache_init = dims -> spzeros(dims...),
                                         jump_dependent_shock_matrices = true)
    else
        return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, crw_ccgf, z, y, Ψ, Nε; sparse_jacobian = sparse_jacobian,
                                         jump_dependent_shock_matrices = true)
    end
end

crw_cₜ(m, zₜ) = exp(m.y[1] + (m.Ψ * (zₜ - m.z))[1])

# Evaluates m_{t + 1} + r_{t + 1}
function crw_logSDFxR(m, zₜ, εₜ₊₁, Cₜ)
    zₜ₊₁, yₜ₊₁ = simulate(m, εₜ₊₁, zₜ)

    return log(m_crw.β) - m_crw.γ * (yₜ₊₁[1] - log(Cₜ)) + zₜ₊₁[2]
end

# Calculate 𝔼ₜ[exp(mₜ₊₁ + rₜ₊₁)] via quadrature
std_norm_mean = zeros(2)
std_norm_sig  = ones(2)
crw_𝔼_quadrature(f::Function) = gausshermite_expectation(f, std_norm_mean, std_norm_sig, 10)

# Calculate implied state variable(s)
function crw_endo_states(m, zₜ, zₜ₋₁, c_impl)
    # rₜ, yₜ are exogenous while Nₜ = exp(rₜ) * Aₜ₋₁ + Yₜ is entirely pre-determined.
    # Thus, our implied state variable will be foreign asset Aₜ = Nₜ - Cₜ.

    # zₜ₋₁ may be the previous period's implied state, so we start from there
    # to calculate Aₜ₋₁.
    yₜ₋₁ = m.y + m.Ψ * (zₜ₋₁ - m.z) # Calculate implied jump variables last period
    Cₜ₋₁ = exp(yₜ₋₁[1])             # to get the implied consumption last period.
    Aₜ₋₁ = zₜ₋₁[1] - Cₜ₋₁           # Given that consumption, we compute implied foreign assets yesterday.
    Nₜ   = exp(zₜ[2]) * Aₜ₋₁ + exp(zₜ[3]) # Now we can get implied resources available today.

    return vcat(zₜ, Nₜ - exp(c_impl)) # This gives us implied foreign assets today, along with other state variables
end

function crw_ccgf(F, α, z)
    # F .= .5 * diag(α * α') # slower but this is the underlying math
    sum!(F, α.^2) # faster implementation
    F .*= .5
end


module WaterVaporMod
using PhysCon
#-----------------------------------------------------------------------
# DESCRIPTION:
# Calculate saturation vapor pressure and latent heat of vaporization
#
export SatVap

function SatVap(t)
    #
    # !DESCRIPTION:
    # Compute saturation vapor pressure and change in saturation vapor pressure
    # with respect to temperature. Polynomial approximations are from:
    # Flatau et al (1992) Polynomial fits to saturation vapor pressure.
    # Journal of Applied Meteorology 31:1507-1513
    #
    # !USES:
    # tfrz = 273.15;                # Freezing point of water (K)
    #
    # !ARGUMENTS:

    # t        ! Temperature (K)
    # es       ! Vapor pressure (Pa)
    # d(es)/d(t) (Pa/K)

    # For water vapor (temperature range is 0C to 100C)
    a0 =  6.11213476;
    a1 =  0.444007856;
    a2 =  0.143064234e-01;
    a3 =  0.264461437e-03;
    a4 =  0.305903558e-05;
    a5 =  0.196237241e-07;
    a6 =  0.892344772e-10;
    a7 = -0.373208410e-12;
    a8 =  0.209339997e-15;

    # and for derivative
    b0 =  0.444017302;
    b1 =  0.286064092e-01;
    b2 =  0.794683137e-03;
    b3 =  0.121211669e-04;
    b4 =  0.103354611e-06;
    b5 =  0.404125005e-09;
    b6 = -0.788037859e-12;
    b7 = -0.114596802e-13;
    b8 =  0.381294516e-16;

    # For ice (temperature range is -75C to 0C)
    c0 =  6.11123516;
    c1 =  0.503109514;
    c2 =  0.188369801e-01;
    c3 =  0.420547422e-03;
    c4 =  0.614396778e-05;
    c5 =  0.602780717e-07;
    c6 =  0.387940929e-09;
    c7 =  0.149436277e-11;
    c8 =  0.262655803e-14;

    # and for derivative
    d0 =  0.503277922;
    d1 =  0.377289173e-01;
    d2 =  0.126801703e-02;
    d3 =  0.249468427e-04;
    d4 =  0.313703411e-06;
    d5 =  0.257180651e-08;
    d6 =  0.133268878e-10;
    d7 =  0.394116744e-13;
    d8 =  0.498070196e-16;

    tc = t - physcon.tfrz
    # Not sure these are needed but why not (shouldn't really exceeed them)
    if (tc > 100.0) tc = 100.0 end
    if (tc < -75.0) tc = -75.0 end

    if (tc >= 0.0)
       es    = a0 + tc*(a1 + tc*(a2 + tc*(a3 + tc*(a4 + tc*(a5 + tc*(a6 + tc*(a7 + tc*a8)))))))
       desdt = b0 + tc*(b1 + tc*(b2 + tc*(b3 + tc*(b4 + tc*(b5 + tc*(b6 + tc*(b7 + tc*b8)))))))
    else
       es    = c0 + tc*(c1 + tc*(c2 + tc*(c3 + tc*(c4 + tc*(c5 + tc*(c6 + tc*(c7 + tc*c8)))))))
       desdt = d0 + tc*(d1 + tc*(d2 + tc*(d3 + tc*(d4 + tc*(d5 + tc*(d6 + tc*(d7 + tc*d8)))))))
    end

    es    = es    * 100.;            # Convert from mb to Pa
    desdt = desdt * 100.;            # Convert from mb to Pa
    return es, desdt

end # SatVap
end # Module


struct TrimConditions
    # TODO: improve typing
    tas::Number  # m/s
    ψ::Number  # rad
    position::AbstractArray  # [x, y, z] earth (m)
    ψ_dot::Number  # rad/s
    γ::Number  # rad
    aircraft::Aircraft
    atmosphere::Atmosphere
    gravity::Gravity
end

struct TrimSolution
    # TODO: improve typing
    α::Number
    β::Number
    controls::AbstractArray
    θ::Number
    ϕ::Number
    p::Number
    q::Number
    r::Number
end

function trim(
    dss_guess, controls_guess, aircraft, atmosphere, gravity, γ=0.0, ψ_dot=0.0;
    show_trace=false,
    ftol=1e-16,
    iterations=5000
)

    # TRIMMING SOLUTION
    sol_gues = [
        get_α(dss_guess),  # alpha (rad)
        get_β(dss_guess),  # beta (rad)
        controls_guess..., # thtl  (0-1), el (deg), ail (deg), rdr (deg)
    ]

    # CONSTS
    trim_conditions = TrimConditions(
        get_tas(dss_guess),  # TAS (m/s)
        get_euler_angles(dss_guess)[1],  # psi (rad)
        get_earth_position(dss_guess),  # north, east, down (m)
        ψ_dot,  # ψ_dot (rad/s)
        γ,  # γ (rad)
        aircraft,
        atmosphere,
        gravity,
    )

    ds_type = typeof(dss_guess)

    f_opt(sol) = trim_cost_function(ds_type, sol, trim_conditions; full_output=false)

    result = nlsolve(
        f_opt, sol_gues;
        ftol=ftol, show_trace=show_trace, iterations=iterations
    )

     if show_trace
            println(result)
     end

    # TODO: return a DSStateDot instead of x
    sol = result.zero
    dssd, controls, outputs, cost = trim_cost_function(
        ds_type, sol, trim_conditions; full_output=true
    )
    return dssd, controls, outputs, cost
end

function trim_cost_function(ds::Type{T}, sol, trim_conditions; full_output=false) where {T<:DSState}

    tas = trim_conditions.tas
    ψ = trim_conditions.ψ
    x, y, z = trim_conditions.position
    ψ_dot = trim_conditions.ψ_dot
    γ = trim_conditions.γ
    aircraft = trim_conditions.aircraft
    atmosphere = trim_conditions.atmosphere
    gravity = trim_conditions.gravity

    gd = get_gravity_accel(gravity)

    ϕ, θ, p, q, r = apply_trimmer_constrains(
        tas,  # m/s
        sol[1],  # α (rad)
        sol[2],  # β (rad)
        γ,  # rad
        ψ_dot,  # rad/s
        gd,  # m/s²
    )

    trim_solution = TrimSolution(
        sol[1],  # α (rad)
        sol[2],  # β (rad)
        sol[3:6],  # controls  -> thtl  (0-1), el (deg), ail (deg), rdr (deg)
        θ,  # rad
        ϕ,  # rad
        p,  # rad/s
        q,  # rad/s
        r,  # rad/s
    )
    # Construct state vector
    dss = ds(trim_conditions, trim_solution, aircraft)
    dssd, outputs = f(time, dss, trim_solution.controls, aircraft, atmosphere, gravity)

    x_dot = get_xdot(dssd)
    cost = [x_dot[1:3]..., x_dot[7:9]...]

    if full_output
        return dssd, trim_solution.controls, outputs, cost
    else
        return cost
    end
end

function apply_trimmer_constrains(tas, α, β, γ, ψ_dot, gd)
    # Coordinated turn bank --> phi
    ϕ = coordinated_turn_bank(ψ_dot, α, β, tas, γ, gd)
    # Climb -> theta
    θ = rate_of_climb_constrain_no_wind(γ, α, β, ϕ)
    # Angular kinemtic -> p, q, r
    p, q, r = ψθϕ_dot_2_pqr(ψ_dot, 0, 0, θ, ϕ)

    return ϕ, θ, p, q, r
end

function SixDOFAeroEuler{T}(tc::TrimConditions, ts::TrimSolution, ac::Aircraft) where {T}
    x = [
        tc.tas,
        ts.α,
        ts.β,
        ts.ϕ,
        ts.θ,
        tc.ψ,
        ts.p,
        ts.q,
        ts.r,
        tc.position...,
        tgear(ac, ts.controls[1])
    ]
    return SixDOFAeroEuler(x)
end

function SixDOFBodyEuler{T}(tc::TrimConditions, ts::TrimSolution, ac::Aircraft) where {T}
    x = [
        wind2body(tc.tas, 0, 0, ts.α, ts.β)...,
        ts.ϕ,
        ts.θ,
        tc.ψ,
        ts.p,
        ts.q,
        ts.r,
        tc.position...,
        tgear(ac, ts.controls[1])
    ]
    return SixDOFBodyEuler(x)
end


using MLDatasets
using CUDA
using Flux
using Flux: Data.DataLoader
using Flux: onehotbatch, onecold, @epochs
using Statistics: mean

x_train, y_train = MNIST.traindata()
x_test, y_test = MNIST.testdata()
x_train = reshape(x_train, 28, 28, 1, :) .|> Float32
x_test = reshape(x_test, 28, 28, 1, :) .|> Float32
# one hot batch
y_train = onehotbatch(y_train, 0:9)
y_test = onehotbatch(y_test, 0:9)
x_train, x_test, y_train, y_test = if CUDA.functional()
    x_train |> gpu, x_test |> gpu, y_train |> gpu, y_test |> gpu
else
    x_train, x_test, y_train, y_test
end
data = DataLoader((x_train, y_train), batchsize=40, shuffle=true)

model = Chain(
  Conv((5, 5), 1=>6, relu),
  MaxPool((2, 2)),
  Conv((5, 5), 6=>16, relu),
  MaxPool((2, 2)),
  flatten,
  Dense(256, 120, relu),
  Dense(120, 84, relu),
  Dense(84, 10),
  softmax,
)
model = if CUDA.functional() model |> gpu else model end

loss(x, y) = Flux.crossentropy(model(x), y)
accurate(x, y) = mean(onecold(model(x)) .== onecold(y))
opt = ADAM()

@epochs 10 begin
  Flux.train!(loss, params(model), data, opt)
  @show(loss(x_test, y_test))
  @show(accurate(x_test, y_test))
end


# Functions for preforming gradient descent on the likelihood of an embedding for a given graph

""" Likelihood of an edge given the embedding coordinate """
function edge_probability(f1, i::Int, j::Int, coordinates)
    yi, yj = coordinates[[i,j], :]
    return f1((norm(yi - yj))^2)
end

""" A likelihood function """
f1(x) = 1/(1+exp(x))

""" Derivative of the likelihood function """
f1prime(x) = -exp(x)/((exp(x)+1)^2)

""" Gradient of the log-likelihood """
function diff_log_sig(x,y)
    z = (norm(x - y))^2
    return (f1prime(z)/f1(z))*2*(x-y)
end

""" Gradient of the log-likelihood with negative edge samples """
function diff_f(x,y,neg_sample)
    x1 = diff_log_sig(x,y)
    x2 = diff_log_sig(y,x)
    for sample in neg_sample
        x2 = x2 + diff_log_sig(y,sample)
    end
    return [x1,x2]
end

"""
    diff_likelihood(coordinates,g,nodes, γ, weight2,num_negative_edges,negsource,negdestin) -> coordinates_revised
Gradient descent step on the embedding `coordinates` given the graph `g`.
"""
function diff_likelihood(coordinates,g,nodes, γ, weight2,num_negative_edges,negsource,negdestin)
    coordinates_revised = [[0.0,0.0] for x in nodes]
    for ed in edges(g)
        ei, ej = src(ed), dst(ed)

        indi = findall(x->x==ei, nodes)[1]
        indj = findall(x->x==ej, nodes)[1]

        yi = coordinates[indi, :]
        yj = coordinates[indj, :]

        x1 = diff_log_sig(yi,yj)
        x2 = diff_log_sig(yj,yi)

        neg_ind = findall(x->x==ej,negdestin)

        num_of_neg_ej = length(neg_ind)
        neg_ind_weight = weight2[neg_ind]
        neg_ind_weight_norm = neg_ind_weight./(sum(neg_ind_weight)+0.000001)

        for i in 1:num_negative_edges
            edge_ind = sample(neg_ind, Weights(neg_ind_weight_norm))

            indii = findall(x->x==negsource[edge_ind][1], nodes)[1]
            indjj = findall(x->x==negdestin[edge_ind][1], nodes)[1]

            yii = coordinates[indii, :]
            yjj = coordinates[indjj, :]

            x2 = x2 .+ diff_log_sig(yj,yii)
        end

        coordinates_revised[indi] = coordinates_revised[indi] .+ x1
        coordinates_revised[indj] = coordinates_revised[indj] .+ x2

    end

    coordinates_revised = [[x[1]/ne(g),x[2]/ne(g)] for x in coordinates_revised]

    return coordinates_revised
end


function diffeq(du,u,p,t)
    CycE = p[C.CycElevel] - u[V.CycEp27]

    Vdp27 = p[C.kd27] + (p[C.kd27e]*CycE)*u[V.Skp2]
    Vdcyce = p[C.kdcyce] + p[C.kdcycee]*CycE
    Vdskp2 = p[C.kdskp2] + p[C.kdskp2c1]*u[V.Cdh1]

    Vicdh1 = p[C.kicdh1e]*CycE

    du[V.p27T] = p[C.ks27] - Vdp27*u[V.p27T]

    du[V.Skp2] = p[C.ksskp2] - Vdskp2*u[V.Skp2]

    du[V.CycEp27] = p[C.kasse]*(p[C.CycElevel]-u[V.CycEp27])*(u[V.p27T]-u[V.CycEp27])-(p[C.kdise]+Vdp27+Vdcyce)*u[V.CycEp27]

    du[V.EmiC] = p[C.kasec]*(p[C.Cdh1T]-u[V.EmiC])*(p[C.Emi1T]-u[V.EmiC]) - (p[C.kdiec]+p[C.kdemi1])*u[V.EmiC]

    du[V.Cdh1dp] = p[C.kacdh1]*(p[C.Cdh1T]-u[V.Cdh1dp]) - Vicdh1*u[V.Cdh1dp]

    du[V.Cdh1] = (p[C.kdiec]+p[C.kdemi1])*(u[V.Cdh1dp]-u[V.Cdh1]) - p[C.kasec]*u[V.Cdh1]*(p[C.Emi1T]-u[V.EmiC])+p[C.kacdh1]*(p[C.Cdh1T]-u[V.EmiC]-u[V.Cdh1])-Vicdh1*u[V.Cdh1]

end

function get_steady_state(p::Vector{Float64})
    tspan::Tuple{Float64,Float64} = (0.0,Inf)
    u0::Vector{Float64} = zeros(V.NUM)

    prob = ODEProblem(diffeq,u0,tspan,p)
    prob = SteadyStateProblem(prob)
    sol = solve(prob,DynamicSS(CVODE_BDF()),dt=1.0)

    return sol.u
end


using NURBS.Bases
import LinearAlgebra: dot, cross, norm

struct Surface{BU<:Basis1D, BV<:Basis1D, T<:Real} <: Object2D
    basisU::BU
    basisV::BV
    d1bU::BU
    d2bU::BU
    d3bU::BU
    d1bV::BV
    d2bV::BV
    d3bV::BV
    points::Array{T, 3}

    function Surface(basisU::BU, basisV::BV, points::Array{T, 3}) where {BU<:Basis1D, BV<:Basis1D, T<:Real}
        d1bU = deriv(basisU)
        d2bU = deriv(d1bU)
        d3bU = deriv(d2bU)
        d1bV = deriv(basisV)
        d2bV = deriv(d1bV)
        d3bV = deriv(d2bV)
        new{BU, BV, T}(basisU, basisV, d1bU, d2bU, d3bU, d1bV, d2bV, d3bV, points)
    end

    function Surface(srf::Surface{BU, BV, T}, moves::Array{T, 3}) where {BU<:Basis1D, BV<:Basis1D, T<:Real}
        points = srf.points + moves
        new{BU, BV, T}(
            srf.basisU, srf.basisV, srf.d1bU, srf.d2bU, srf.d3bU, srf.d1bV, srf.d2bV, srf.d3bV, points
        )
    end
end

const BSurface = Surface{BSplineBasis, BSplineBasis}
const NSurface = Surface{NURBSBasis, NURBSBasis}

function evalpoint(srf::Surface, u, v)
    return srf.basisV(v, srf.basisU(u, srf.points))'
end

function evalnormal(srf::Surface, u, v)
    (du, dv, duu, duv, dvv) = evalderiv2(srf, u, v)
    E = dot(du, du)
    F = dot(du, dv)
    G = dot(dv, dv)
    n = cross(du, dv) ./ norm(cross(du, dv))
    # n ./= norm(n)
    L = dot(duu, n)
    M = dot(duv, n)
    N = dot(dvv, n)
    dS2 = E*G - F^2
    K = (L*N - M^2) / dS2
    H = (E*N - 2*F*M + G*L) / (2*dS2)
    Δ = sqrt(H^2 - K)
    if H > 0  # Rhino orders principals by absolute value
        κ1 = H + Δ
        κ2 = H - Δ
    else
        κ1 = H - Δ
        κ2 = H + Δ
    end
    λ1 = (κ1*F - M) / (N - κ1*G)
    λ2 = (κ2*F - M) / (N - κ2*G)
    d1 = (du + λ1*dv) ./ norm(du + λ1*dv)
    # d1 ./= norm(d1)
    d2 = (du + λ2*dv) ./ norm(du + λ2*dv)
    # d2 ./= norm(d2)
    (n, κ1, κ2, du, dv, d1, d2, dS2)
end

function evalderiv1(srf::Surface, u, v)
    du = srf.basisV(v, srf.d1bU(u, srf.points))
    dv = srf.d1bV(v, srf.basisU(u, srf.points))
    (du', dv')
end

function evalderiv2(srf::Surface, u, v)
    duu = srf.basisV(v, srf.d2bU(u, srf.points))
    du = srf.basisV(v, srf.d1bU(u, srf.points))
    duv = srf.d1bV(v, srf.d1bU(u, srf.points))
    dv = srf.d1bV(v, srf.basisU(u, srf.points))
    dvv = srf.d2bV(v, srf.basisU(u, srf.points))
    (du', dv', duu', duv', dvv')
end

function evalderiv3(srf::Surface, u, v)
    duuu = srf.basisV(v, srf.d3bU(u, srf.points))
    duu = srf.basisV(v, srf.d2bU(u, srf.points))
    duuv = srf.d1bV(v, srf.d2bU(u, srf.points))
    duv = srf.d1bV(v, srf.d1bU(u, srf.points))
    du = srf.basisV(v, srf.d1bU(u, srf.points))
    duvv = srf.d2bV(v, srf.d1bU(u, srf.points))
    dv = srf.d1bV(v, srf.basisU(u, srf.points))
    dvv = srf.d2bV(v, srf.basisU(u, srf.points))
    dvvv = srf.d3bV(v, srf.basisU(u, srf.points))
    (du', dv', duu', duv', dvv', duuu', duuv', duvv', dvvv')
end

function BSurface(orderU::Int, orderV::Int, points::Array{T, 3}) where {T<:Real}
    bu = BSplineBasis(range(-1, stop=1, length=size(points, 1)-orderU+2), orderU)
    bv = BSplineBasis(range(-1, stop=1, length=size(points, 2)-orderV+2), orderV)
    Surface(bu, bv, points)
end

function BSurface(KnotsU::Vector{T}, KnotsV::Vector{T}, orderU::Int, orderV::Int, points) where {T<:Real}
    bu = BSplineBasis(2 .* (KnotsU .- minimum(KnotsU)) ./ (maximum(KnotsU) - minimum(KnotsU)) .- 1, orderU)
    bv = BSplineBasis(2 .* (KnotsV .- minimum(KnotsV)) ./ (maximum(KnotsV) - minimum(KnotsV)) .- 1, orderV)
    Surface(bu, bv, points)
end

function isocurve(srf::Surface, dir::Bool, u)
    b = dir ? srf.basisV : srf.basisU
    p = dir ? srf.basisU(u, srf.points) : srf.basisV(u, permutedims(srf.points, [2, 1, 3]))
    return Curve(b, p)
end

NSurface(srf::BSurface) = Surface(NURBSBasis(srf.basisU), NURBSBasis(srf.basisV))
NSurface(points::Array{T, 3}, degU::Int, degV::Int) where {T<:Real} = NSurface(BSurface(points, degU, degV))


function estimatelearningcore(y,x,Choice,bstart,N,T,S,J,BigN,Delta,sig,Resid,resid,tresid,Psi1,abilsub,Csum)
    abil=zeros(N*S,J)

    idelta   = inv(Delta)
    vtemp2   = zeros(J,J)
    vabil    = zeros(N*S,J)
    vabilw   = zeros(N*S,J)
    for j=1:J
        tresid[:,j] = (sum(reshape(Resid[:,j],(T,N*S));dims=1))'
    end

    for i=1:S*N
        psit=Psi1[i,:]
        Psi = zeros(J,J)
        for j=1:J
            Psi[j,j] = psit[j]./sig[j]
        end

        vtemp=inv(idelta.+Psi)

        vectres = zeros(J)
        for j=1:J
            vectres[j] = tresid[i,j]./sig[j]
        end
        temp=(vtemp*vectres)'
        abil[i,:]=temp

        for j=1:J
            vabil[i,j] = vtemp[j,j] # later need PTypesl[i]*vtemp[j,j]
        end

        vtemp2.+=(vtemp.+temp'*temp) # later need PTypesl[i]*(vtemp+temp'*temp)
    end

    Delta=vtemp2./N

    vabilw    = deepcopy(vabil)
    Abil      = kron(abil,ones(T,1))
    Vabil     = kron(vabil,ones(T,1))

    sigdem = zeros(J)
    for j=1:J
        abilsub[j] = Abil[vec(Choice.==j),j]
        sigdem[j]  = sum((resid[j].-abilsub[j]).^2) # later need sum(PTypesub[j].*(resid[j].-abilsub[j]).^2)
    end

    sig=((sum(Csum.*vabilw;dims=1))'.+sigdem)./BigN

    for j=1:J
        flag                     = (y[:,j].!=999)
        bstart[:,j]              = x[flag,:,j]\(y[flag,j].-abilsub[j]) # later weight by PTypesub[j]
        resid[j]                 = y[flag,j].-x[flag,:,j]*bstart[:,j]
        Resid[vec(Choice.==j),j] = resid[j]
    end

    return bstart,Delta,sig,Resid,resid,tresid,abilsub
end


using MixedModels, Test

@testset "linear" begin
    m1 = fit(MixedModel, @formula(yield ~ 1 + (1|batch)), MixedModels.dataset(:dyestuff); progress=false)
    @test first(m1.θ) ≈ 0.7525806757718846 rtol=1.0e-5
end

@testset "generalized" begin
    gm1 = fit(MixedModel, @formula(use ~ 1 + urban + livch + age + abs2(age) + (1|dist)),
              MixedModels.dataset(:contra), Bernoulli(); progress=false)
    @test deviance(gm1) ≈ 2372.7286 atol=1.0e-3
end

@testset "Normal-IdentityLink" begin
    @test isa(fit(MixedModel, @formula(yield ~ 1 + (1|batch)), MixedModels.dataset(:dyestuff), Normal(); progress=false),
              LinearMixedModel)
    @test_throws(ArgumentError("use LinearMixedModel for Normal distribution with IdentityLink"),
                 fit(GeneralizedLinearMixedModel,
                     @formula(yield ~ 1 + (1|batch)),
                     MixedModels.dataset(:dyestuff); progress=false))
end

@testset "Normal Distribution GLMM" begin
    @test isa(fit(MixedModel, @formula(yield ~ 1 + (1|batch)), MixedModels.dataset(:dyestuff),
                         Normal(), SqrtLink(); progress=false),
                     GeneralizedLinearMixedModel)
end


using Distributions

𝕀(b) = b ? 1 : 0 # indicator function

function sierra_mixture(; 𝛍=[0,0], σ=3, 𝐬=[+σ,-σ], δ=2, η=6, decay=true)
    𝚺 = [σ 0.0; 0.0 σ]
    decay = 𝕀(decay)
    origin = [0.0, 0.0]
    𝐒 = MvNormal[MvNormal(𝛍, 𝚺/(σ*η))]

    for g in [[+δ, +δ], [+δ, -δ], [-δ, +δ], [-δ, -δ]]
        for (i,𝐩) in enumerate([origin, [+1, +1], [+2, 0], [+3, +1], [0, +2], [+1, +3]])
            for s in 𝐬
                push!(𝐒, MvNormal(g + s*𝐩 + 𝛍, i^decay/η * 𝚺))
            end
        end
    end

    return MixtureModel(𝐒)
end

global MMCACHE = Dict() # hyperparameters => MixtureModel

function sierra(x; minimize=true, kwargs...)
    global MMCACHE

    hpk = kwargs
    if haskey(MMCACHE, hpk)
        𝐒 = MMCACHE[hpk]
    else
        𝐒 = sierra_mixture(; kwargs...)
    end

    # mixture model with equal weights
    sgn = minimize ? -1 : +1
    return sgn*pdf(𝐒, x)
end


module Dynamics
    using OrbitalTrajectories.SpiceUtils

    using SPICE
    using Plots
    using Unitful
    using DifferentialEquations
    using DiffEqBase
    using DiffEqSensitivity
    using SciMLBase
    using LinearAlgebra
    using ModelingToolkit
    using NLsolve
    using LineSearches
    using ForwardDiff
    using FiniteDiff
    using Memoize
    using SimpleTraits
    using StaticArrays
    using ProgressMeter
    using RecipesBase
    using PhysicalConstants.CODATA2014: NewtonianConstantOfGravitation

    # Required to ensure that we can precompile ODEFunctions (which needs to be
    # done in our own cache).
    using RuntimeGeneratedFunctions
    RuntimeGeneratedFunctions.init(@__MODULE__)

    #----------------#
    # ABSTRACT TYPES #
    #----------------#

    # Dynamical models
    abstract type Abstract_DynamicalModel end
    abstract type Abstract_ModelODEFunctions <: ModelingToolkit.AbstractODESystem end
    abstract type Abstract_R3BPModel <: Abstract_DynamicalModel end
    abstract type Abstract_R4BPModel <: Abstract_DynamicalModel end

    # Reference frames
    abstract type Abstract_ReferenceFrame end

    # Differential correctors
    abstract type Abstract_DifferentialCorrector end

    #----------#
    # INCLUDES #
    #----------#
    include("orbital_trajectories.jl")
    include("reference_frames.jl")
    include("sensitivities.jl")
    include("dynamical_models.jl")
    include("differential_correction.jl")
    include("plotting.jl")
end


######################### Boundary Conditions #########################
# The job of the BC functions is to calculate/enforce face fluxes at boundary faces.
# The JST method only calculates fluxes at internal faces, then these conditions are applied to calculate them at the boundaries

# InletConditions: [ Static Pressure, Static Temperture, Ux, Uy, Uz ]
function supersonicInletBoundary(mesh::Mesh, sln::SolutionState, boundaryNumber, inletConditions, fluid::Fluid)
    P, T, Ux, Uy, Uz = inletConditions

    rho = idealGasRho(T, P, fluid.R)
    xMom, yMom, zMom = [Ux, Uy, Uz] .* rho
    e = calPerfectEnergy(T, fluid)
    eV2 = rho*(e + (mag([Ux, Uy, Uz])^2)/2)

    boundaryFluxes = Vector{Float64}(undef, 15)
    calculateFluxes3D!(boundaryFluxes, [P, T, Ux, Uy, Uz],  [rho, xMom, yMom, zMom, eV2])
    currentBoundary = mesh.boundaryFaces[boundaryNumber]
    @inbounds for face in currentBoundary
        sln.faceFluxes[face,:] = boundaryFluxes
    end
end

# InletConditions: [ totalPressure, totalTemp, nx, ny, nz ]
# Where n is the unit vector representing the direction of inlet velocity
# Using method from FUN3D solver
function subsonicInletBoundary(mesh, sln::SolutionState, boundaryNumber, inletConditions, fluid::Fluid)
    Pt, Tt, nx, ny, nz = inletConditions
    Cp, R, gamma = (fluid.Cp, fluid.R, fluid.gamma)

    boundaryFluxes = Vector{Float64}(undef, 15)
    primitives = Vector{Float64}(undef, 5)
    state = Vector{Float64}(undef, 5)
    velUnitVector = [ nx, ny, nz ]
    currentBoundary = mesh.boundaryFaces[boundaryNumber]

    @inbounds @fastmath for face in currentBoundary
        ownerCell = mesh.faces[face][1]
        @views adjustedVelocity = dot(velUnitVector, sln.cellPrimitives[ownerCell, 3:5])
        primitives[2] = Tt - (gamma-1)/2 * (adjustedVelocity^2)/(gamma*R)
        machNum = abs(adjustedVelocity) / sqrt(gamma * R * primitives[2])
        primitives[1] = Pt*(1 + (gamma-1)/2 * machNum^2)^(-gamma/(gamma-1))
        primitives[3:5] = adjustedVelocity .* velUnitVector

        # Calculate state variables
        state[1] = idealGasRho(primitives[2], primitives[1], R)
        state[2:4] .= primitives[3:5] .* state[1]
        e = calPerfectEnergy(primitives[2], fluid)
        state[5] = state[1]*(e + (mag(primitives[3:5])^2)/2 )

        calculateFluxes3D!(boundaryFluxes, primitives, state)
        sln.faceFluxes[face, :] = boundaryFluxes
    end
end

function pressureOutletBoundary(mesh, sln::SolutionState, boundaryNumber, outletPressure, fluid::Fluid)
    nFluxes = size(sln.cellFluxes, 2)

    # Directly extrapolate cell center flux to boundary (zero gradient between the cell center and the boundary)
    currentBoundary = mesh.boundaryFaces[boundaryNumber]
    @inbounds @fastmath for face in currentBoundary
        ownerCell = mesh.faces[face][1]
        # Mass fluxes are unaffected by pressure boundary
        for flux in 1:nFluxes
            sln.faceFluxes[face, flux] = sln.cellFluxes[ownerCell, flux]
        end

        origP = sln.cellPrimitives[ownerCell, 1]
        # Adjust the momentum fluxes containing pressure
        for flux in [4, 8, 12]
            sln.faceFluxes[face, flux] += outletPressure - origP
        end
        #Adjust the energy fluxes
        for d in 1:3
            sln.faceFluxes[face, 12+d] += sln.cellPrimitives[ownerCell, 2+d]*(outletPressure - origP)
        end
    end
end

function zeroGradientBoundary(mesh::Mesh, sln::SolutionState, boundaryNumber, _, __)
    nFluxes = size(sln.cellFluxes, 2)

    # Directly extrapolate cell center flux to boundary (zero gradient between the cell center and the boundary)
    currentBoundary = mesh.boundaryFaces[boundaryNumber]
    @inbounds for face in currentBoundary
        ownerCell = mesh.faces[face][1] #One of these will be -1 (no cell), the other is the boundary cell we want
        for flux in 1:nFluxes
            sln.faceFluxes[face, flux] = sln.cellFluxes[ownerCell, flux]
        end
    end
end

function wallBoundary(mesh::Mesh, sln::SolutionState, boundaryNumber, _, __)
    currentBoundary = mesh.boundaryFaces[boundaryNumber]
    @inbounds for f in currentBoundary
        ownerCell = max(mesh.faces[f][1], mesh.faces[f][2]) #One of these will be -1 (no cell), the other is the boundary cell we want

        faceP = sln.cellPrimitives[ownerCell, 1]
        # Momentum flux is Pressure in each of the normal directions (dot product)
        sln.faceFluxes[f, 4] = faceP
        sln.faceFluxes[f, 8] = faceP
        sln.faceFluxes[f, 12] = faceP

        # Mass Flux is zero
        sln.faceFluxes[f, 1:3] .= 0.0
        # Energy Flux is zero
        sln.faceFluxes[f, 13:15] .= 0.0
    end
end

# Symmetry is identical to a wall in Euler simulations
symmetryBoundary = wallBoundary

function emptyBoundary(mesh::Mesh, sln::SolutionState, boundaryNumber, _, __)
    return
end

function applyBoundaryConditions(mesh::Mesh, sln::SolutionState, boundaryConditions, nBoundaries, fluid::Fluid)
    for boundaryNumber in 1:nBoundaries
        bFunctionIndex = 2*boundaryNumber-1
        boundaryConditionFunction = boundaryConditions[bFunctionIndex]
        boundaryParameters = boundaryConditions[bFunctionIndex+1]

        boundaryConditionFunction(mesh, sln, boundaryNumber, boundaryParameters, fluid)
    end
end


################################
######     Fixations      ######
################################

# E[Dn]  = LT(E[Dn+] + E[Dn-] + E[Dns])

# Neutral fixation corrected by background selection
"""

 - fixNeut()

Expected neutral fixations rate reduce by B value.

```math
\\mathbb{E}[D_{s}] = (1 - p_{-} - p_{+}) B \\frac{1}{2N}
```
# Returns
 - `Float64`: expected rate of neutral fixations.

"""
function fixNeut(param::parameters)
	# Synonymous probabilty * (fixation probabilty corrected by BGS value)
	out::Float64 = 0.25*(1.0/(param.B*param.NN))
	return out
end

# Negative fixations corrected by background selection
"""

	fixNegB(ppos)

Expected fixation rate from negative DFE.

```math
\\mathbb{E}[D_{n-}] =  p_{-}\\left(2^-\\alpha\\beta^\\alpha\\left(-\\zeta[\\alpha,\\frac{2+\\beta}{2}] + \\zeta[\\alpha,1/2(2-\\frac{1}{N+\\beta})]\\right)\\right)
```

# Arguments
 - `ppos::Float64`: positive selected alleles probabilty.

# Returns
 - `Float64`: expected rate of fixations from negative DFE.

"""
function fixNegB(param::parameters,ppos::Float64)
	# Non-synonymous proportion * negative alleles probability * fixation probability from gamma distribution
	out::Float64 = 0.75*(1-ppos)*(2^(-param.al))*(param.B^(-param.al))*(param.be^param.al)*(-zeta(param.al,1.0+param.be/(2.0*param.B))+zeta(param.al,0.5*(2-1.0/(param.N*param.B)+param.be/param.B)))
	return out
end

# Positive fixations
"""
	pFix()

Expected positive fixation rate.

```math
\\mathbb{E}[D_{n+}] =  p_{+} \\cdot B \\cdot (1 - e^{(-2s)})
```

# Arguments
 - `ppos::Float64`: positive selected alleles probabilty.

# Returns
 - `Float64`: expected rate of positive fixation.

"""
function pFix(param::parameters,gam::Int64)

	# Fixation probability
	s::Float64    = gam/(param.NN)
	pfix::Float64 = (1.0-ℯ^(-2.0*s))/(1.0-ℯ^(-2.0*gam))

	# Correcting pFix for large s following Uricchio et al. 2014
	if s >= 0.1
		pfix = ℯ^(-(1.0+s))
		lim = 0
		while(lim < 200)
			pfix = ℯ^((1.0+s)*(pfix-1.0))
			lim  = lim + 1
		end
		pfix = 1 -pfix
	end

	return pfix
end

# Positive fixations after apply Φ. reduction of positive fixations due deleterious linkage given a value B of background selection
"""

	fixPosSim(gamma,ppos)

Expected positive fixations rate reduced due to the impact of background selection and linkage. The probabilty of fixation of positively selected alleles is reduced by a factor Φ across all deleterious linked sites [`Analytical.phiReduction`](@ref).

```math
\\mathbb{E}[D_{n+}] =  \\Phi \\cdot \\mathbb{E}[D_{n+}]
```

# Arguments
 - `ppos::Float64`: positive selected alleles probabilty.

# Returns
 - `Float64`: expected rate of positive fixations under background selection

"""
function fixPosSim(param::parameters,gamma::Int64,ppos::Float64)

	redPlus = phiReduction(param,gamma)

	# Non-synonymous * positive alleles probability * B reduction * fixation probility
	out::Float64 = 0.75*ppos*redPlus*pFix(param,gamma)
	return out
end


function coreMVI(logl::Function, gradlogl::Function, LAposteriors; seed = 1, S = 100, optimiser = Optim.LBFGS(), iterations = 1, numerical_verification = false, Stest=0, show_every=-1, inititerations=0)

    D = length(mean(LAposteriors[1]))

    @assert(D == size(cov(LAposteriors[1]), 1) == size(cov(LAposteriors[1]), 2))

    @printf("Running MVI with S=%d, D=%d for %d iterations\n", S, D, iterations)

    #----------------------------------------------------
    # generate latent variables
    #----------------------------------------------------

    Ztrain = generatelatentZ(S = S, D = D, seed = seed)

    Ztest  = generatelatentZ(S = Stest, D = D, seed = seed+1)

    #----------------------------------------------------
    function unpack(param)
    #----------------------------------------------------

        @assert(length(param) == 2*D)

        local μ     = @view param[1:D]

        local Esqrt = @view param[D+1:2*D]

        return μ, Esqrt

    end

    #----------------------------------------------------
    function minauxiliary(V, param)
    #----------------------------------------------------

        local μ, Esqrt = unpack(param)

        return -1.0 * elbo(μ, Esqrt, V, Ztrain)

    end

    #----------------------------------------------------
    function minauxiliary_grad(V, param)
    #----------------------------------------------------

        local μ, Esqrt = unpack(param)

        return -1.0 * elbo_grad(μ, Esqrt, V, Ztrain)

    end

    #----------------------------------------------------
    function getcov(V, Esqrt)
    #----------------------------------------------------

        local aux = V * Diagonal(Esqrt)

        return aux*aux'

    end

    #----------------------------------------------------
    function getcovroot(V, Esqrt)
    #----------------------------------------------------

        return V * Diagonal(Esqrt)

    end

    #----------------------------------------------------
    function elbo(μ, Esqrt, V, Z)
    #----------------------------------------------------

        local C = getcovroot(V, Esqrt)

        mean(map(z -> logl(μ .+ C*z), Z)) + ApproximateVI.entropy(Esqrt)

    end

    #----------------------------------------------------
    function elbo_grad(μ, Esqrt, V, Z)
    #----------------------------------------------------

        local grad_μ     = zeros(D)

        local grad_Esqrt = zeros(D)

        local C          = getcovroot(V, Esqrt)

        local S          = length(Z)

        for i=1:S

            local aux = gradlogl(μ .+ C * Z[i])

            grad_μ += aux / S

            grad_Esqrt .+= V' * aux .* Z[i] / S

        end

        # logdet contribution in entropy

        grad_Esqrt .+= 1.0./Esqrt

        return [grad_μ; grad_Esqrt]

    end

    gradhelper(V, storage, param) = copyto!(storage, minauxiliary_grad(V, param))

    #----------------------------------------------------
    # Numerically verify gradient
    #----------------------------------------------------

    if numerical_verification

        local V, Esqrt = eigendecomposition(cov(LAposteriors[1]))
        adgrad = ForwardDiff.gradient(x->minauxiliary(V, x), [mean(LAposteriors[1]); Esqrt])
        angrad = minauxiliary_grad(V, [mean(LAposteriors[1]); Esqrt])
        @printf("gradient from AD vs analytical gradient\n")
        display([vec(adgrad) vec(angrad)])
        @printf("maximum absolute difference is %f\n", maximum(abs.(vec(adgrad) - vec(angrad))))

    end

    #----------------------------------------------------
    # Evaluate initial solutions for few iterations
    #----------------------------------------------------

    function initoptimise(q)

        local V, Esqrt = eigendecomposition(cov(q))

        Optim.optimize(x->minauxiliary(V,x), (x1,x2)->gradhelper(V,x1,x2), [mean(q); Esqrt], optimiser, Optim.Options(iterations = inititerations))

    end

    results = if inititerations > 0
        @showprogress "Initial search with random start " map(initoptimise, LAposteriors)
    else
        map(initoptimise, LAposteriors)
    end

    bestindex = argmin(map(r -> r.minimum, results))

    bestinitialsolution = results[bestindex].minimizer

    V, Esqrt = eigendecomposition(cov(LAposteriors[bestindex]))

    #----------------------------------------------------
    # Call optimiser
    #----------------------------------------------------

    options = Optim.Options(extended_trace = false, store_trace = false, show_trace = true, show_every=show_every, iterations = iterations, g_tol = 1e-6)

    result  = Optim.optimize(x->minauxiliary(V, x), (x1,x2)->gradhelper(V,x1,x2), bestinitialsolution, optimiser, options)

    μopt, Esqrtopt = unpack(result.minimizer)

    #----------------------------------------------------
    # Return results
    #----------------------------------------------------

    Σopt = getcov(V, Esqrtopt)

    return MvNormal(μopt, Σopt), elbo(μopt, Esqrtopt, V, generatelatentZ(S = 10*S, D = D, seed = seed+2))

end


𝑃, 𝑂 = spacecraft_rendezvous(T=200.0, orderT=10, orderQ=2, abs_tol=1e-28, max_steps=5000);

𝑂jets = Options(:orderT=>orderT, :orderQ=>orderQ, :abs_tol=>abs_tol, :max_steps=>max_steps)

sol = solve(𝑃, 𝑂, TMJets(𝑂jets), LazyDiscretePost(:check_invariant_intersection=>true))

# first mode
#p = IVP(𝑃.s.modes[1], 𝑃.x0[1][2])
#solve(p, 𝑂, op=TMJets(𝑂jets))
# compute projection onto the plot variables (project_reachset doesn't currently
# work for a limitation in the hybrid solve, which does not preserve the type of opC)


using SimpleNeRF: sinusodial_emb, NeRFConfig, NeRFModel, total
using Optimisers
using LinearAlgebra
using Test

bs = 16
freq = 4
coords = randn(Float32, 3, bs)
comb = sinusodial_emb(coords, freq)
@test size(comb) == (2*3*freq, bs)

cfg = NeRFConfig()
model = NeRFModel(;config=cfg)
x = randn(Float32, 3, bs)
d = randn(Float32, 3, bs)
density, rgb = model(x, d)
@test size(density) == (1, bs)
@test size(rgb) == (3, bs)

params = Optimisers.setup(Optimisers.ADAM(), model)
@test sqrt(total(x->sum(x.^2), model)) > 1.0f0


using Logging, Random
using JuMP, Gurobi
using DecisionProgramming

include("min_expected_value.jl")

Random.seed!(13)

N = 4
L = [1]
R_k = [k + 1 for k in 1:N]
A_k = [(N + 1) + k for k in 1:N]
F = [2*N + 2]
T = [2*N + 3]
L_states = ["high", "low"]
R_k_states = ["high", "low"]
A_k_states = ["yes", "no"]
F_states = ["failure", "success"]
c_k = rand(N)
b = 0.03
fortification(k, a) = [c_k[k], 0][a]

@info("Creating the influence diagram.")
S = States([
    (length(L_states), L),
    (length(R_k_states), R_k),
    (length(A_k_states), A_k),
    (length(F_states), F)
])
C = Vector{ChanceNode}()
D = Vector{DecisionNode}()
V = Vector{ValueNode}()
X = Vector{Probabilities}()
Y = Vector{Consequences}()

for j in L
    I_j = Vector{Node}()
    X_j = zeros(S[I_j]..., S[j])
    X_j[1] = rand()
    X_j[2] = 1.0 - X_j[1]
    push!(C, ChanceNode(j, I_j))
    push!(X, Probabilities(j, X_j))
end

for j in R_k
    I_j = L
    x, y = rand(2)
    X_j = zeros(S[I_j]..., S[j])
    X_j[1, 1] = max(x, 1-x)
    X_j[1, 2] = 1.0 - X_j[1, 1]
    X_j[2, 2] = max(y, 1-y)
    X_j[2, 1] = 1.0 - X_j[2, 2]
    push!(C, ChanceNode(j, I_j))
    push!(X, Probabilities(j, X_j))
end

for (i, j) in zip(R_k, A_k)
    I_j = [i]
    push!(D, DecisionNode(j, I_j))
end

for j in F
    I_j = L ∪ A_k
    x, y = rand(2)
    X_j = zeros(S[I_j]..., S[j])
    for s in paths(S[A_k])
        d = exp(b * sum(fortification(k, a) for (k, a) in enumerate(s)))
        X_j[1, s..., 1] = max(x, 1-x) / d
        X_j[1, s..., 2] = 1.0 - X_j[1, s..., 1]
        X_j[2, s..., 1] = min(y, 1-y) / d
        X_j[2, s..., 2] = 1.0 - X_j[2, s..., 1]
    end
    push!(C, ChanceNode(j, I_j))
    push!(X, Probabilities(j, X_j))
end

for j in T
    I_j = A_k ∪ F
    Y_j = zeros(S[I_j]...)
    for s in paths(S[A_k])
        cost = sum(-fortification(k, a) for (k, a) in enumerate(s))
        Y_j[s..., 1] = cost + 0
        Y_j[s..., 2] = cost + 100
    end
    push!(V, ValueNode(j, I_j))
    push!(Y, Consequences(j, Y_j))
end

validate_influence_diagram(S, C, D, V)
sort!.((C, D, V, X, Y), by = x -> x.j)

P = DefaultPathProbability(C, X)
U = DefaultPathUtility(V, Y)

@info("Creating the decision model.")
U⁺ = PositivePathUtility(S, U)

k = length(X)  # choose failure node and its probability
# Strategy changes between ϵ 0.8 and 0.9
ϵ = 0.01  # we can try different parameters values here
deviations = [
    s_I => Deviation(X[k].data[s_I..., :], ϵ) for s_I in paths(S[C[k].I_j])
]

model = Model()
z = DecisionVariables(model, S, D)
(min_ev, mpevs) = min_expected_value(deviations, model, S, C, X, U⁺, z, k)
@objective(model, Max, min_ev)

@info("Starting the optimization process.")
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol"      => 1e-9,
    # "LazyConstraints" => 1,
)
set_optimizer(model, optimizer)
optimize!(model)

@info("Construct minimizing path probability")
X′ = min_probabilities(mpevs, X, k)
P = DefaultPathProbability(C, X′)

@info("Reoptimize model with minimum path probability")
model2 = Model()
z2 = DecisionVariables(model2, S, D)
π_s = PathProbabilityVariables(model2, z2, S, P; hard_lower_bound=false)
probability_cut(model2, π_s, P)
# active_paths_cut(model2, π_s, S, P)
ev = expected_value(model2, π_s, U⁺)
@objective(model2, Max, ev)

optimizer2 = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol"      => 1e-9,
    "LazyConstraints" => 1,
)
set_optimizer(model2, optimizer2)
optimize!(model2)

@info("Extracting results.")
Z = DecisionStrategy(z)

@info("Printing decision strategy:")
print_decision_strategy(S, Z)

@info("Printing state probabilities:")
sprobs = StateProbabilities(S, P, Z)
print_state_probabilities(sprobs, L)
print_state_probabilities(sprobs, R_k)
print_state_probabilities(sprobs, A_k)
print_state_probabilities(sprobs, F)

@info("Computing utility distribution.")
udist = UtilityDistribution(S, P, U, Z)

@info("Printing utility distribution.")
print_utility_distribution(udist)

@info("Printing statistics")
print_statistics(udist)


export  PlaneStress
mutable struct PlaneStress
    H::Array{Float64}
    E::Float64
    ν::Float64
    ρ::Float64
    σ0::Array{Float64} # stress at last time step
    σ0_::Array{Float64} # σ0 to be updated in `commitHistory`
    ε0::Array{Float64}
    ε0_::Array{Float64}
end

function PlaneStress(prop::Dict{String, Any})
    E = prop["E"]; ν = prop["nu"]; ρ = prop["rho"]
    H = zeros(3,3)

    H[1,1] = E/(1. -ν*ν)
    H[1,2] = H[1,1]*ν
    H[2,1] = H[1,2]
    H[2,2] = H[1,1]
    H[3,3] = E/(2.0*(1.0+ν))
    # @show H
    # # ! remove it!!!
    # H = [ 4.09627e12  2.88053e11  -1.77231
    #     2.88053e11  8.29483e11   0.616748
    # -1.77231     0.616748     3.43662e11]
    σ0 = zeros(3); σ0_ = zeros(3); ε0 = zeros(3); ε0_ = zeros(3)
    PlaneStress(H, E, ν, ρ, σ0, σ0_,ε0,ε0_)
end

function getStress(self::PlaneStress, strain::Array{Float64}, Dstrain::Array{Float64}, Δt::Float64 = 0.0)
    sigma = self.H * strain
    self.σ0_ = copy(sigma)
    self.ε0_ = copy(strain)
    return sigma, self.H
end

function getTangent(self::PlaneStress)
    self.H
end

function commitHistory(self::PlaneStress)
    self.σ0 = self.σ0_
    self.ε0 = self.ε0_
end


using BandedMatrices, Base.Test

for T in (Float64,Complex128,Float32,Complex64)
    A=brand(T,10,10,3,2)
    Q,R=qr(A)
    @test_approx_eq full(Q)*full(R) A
    b=rand(T,10)
    @test_approx_eq A_mul_B!(similar(b),Q,Ac_mul_B!(similar(b),Q,b)) b
    for j=1:size(A,2)
        @test_approx_eq Ac_mul_B(Q,A[:,j]) R[:,j]
    end
    A=brand(T,14,10,3,2)

    Q,R=qr(A)

    for k=1:size(Q,1),j=1:size(Q,2)
        @test_approx_eq Q[k,j] full(Q)[k,j]
    end

    @test_approx_eq full(Q)*full(R) A

    A=brand(T,10,14,3,2)

    Q,R=qr(A)

    for k=1:size(Q,1),j=1:size(Q,2)
        @test_approx_eq Q[k,j] full(Q)[k,j]
    end

    @test_approx_eq full(Q)*full(R) A
    A=brand(T,100,100,3,4)
    Q,R=qr(A)
    b=rand(T,100)
    @test_approx_eq R\(Q'*b) full(A)\b
end

A=brand(10,10,3,2)
b=rand(Complex128,10)
Q,R=qr(A)
@test_approx_eq R\(Q'*b) full(A)\b

A=brand(Complex128,10,10,3,2)
b=rand(10)
Q,R=qr(A)
@test_approx_eq R\(Q'*b) full(A)\b

println("Large dimension QR solve")
A=brand(100000,100000,3,4)
b=rand(size(A,1))
@time Q,R=qr(A)
@time R\(Q'*b)

println("Time should be   0.030149 seconds (23 allocations: 17.548 MB)")
println("Time should be   0.008210 seconds (29 allocations: 8.012 MB)")

A=brand(2000,2000,400,400)
b=rand(size(A,1))
@time Q,R=qr(A)
@time R\(Q'*b)

println("Time should be   0.401671 seconds (24 allocations: 36.665 MB)")
println("Time should be   0.009276 seconds (30 allocations: 12.262 MB)")


push!(LOAD_PATH, pwd()*"/../../src/")
using Distributions
using PyPlot
using BVB

x = linspace(-15, 15, 500)
init_approx_d = Normal(0, 20)

lower = -10.
upper = 10.
max_iter = 200

# example 1
@show true_d = TriangularDist(-10, 10, 0)
ll(x) = logpdf(true_d, x)

# include("run_laplacian_gradient_boosting.jl")
include("run_laplacian_gradient_boosting_exclusive.jl")
show()

println("Done.")
readline(STDIN)


function extend_to_dc(inFreqs, inArray)
   deltaf = inFreqs[2]-inFreqs[1]
   numextrapts = round(Int, inFreqs[1]/deltaf)
   phasechange = angle(inArray[1])/numextrapts # amount by which to reverse the phase (in radians)

   # extend the inArray down to DC, rolling down the phase
   extrapts = zeros(Complex64,numextrapts)
   for j = 1:numextrapts
      extrapts[numextrapts-j+1] = inArray[1]*e^(-im*j*phasechange)
   end
   extrapts[1,1] = abs(extrapts[1,1]) # (imag(inArray(1,1)) < 10^-12) -> 0
   outArray = vcat(extrapts, inArray)

   # extend the Freqsuencies
   newFreqs = (inFreqs[1]-numextrapts*deltaf:deltaf:inFreqs[1]-deltaf)
   outFreqs = vcat(newFreqs, inFreqs)
   return (outFreqs, outArray)
end


unpackconstruct(Construct::Ordered_Construct) = map((x->getfield(x, :m)|>String |> (x -> parse(Int64, x))),Construct)

quad(x,Q,λ) = transpose(x)*Q*x - λ*norm(x,1)
# ref https://github.com/baptistar/BOCS/blob/master/test_problems/Quadratic/quad_mat.m
function compute_decay(rng,n_vars,Lc)
	K = (s,t) -> exp(-1*(s-t)^2/Lc)
	decay = Array{Float64,2}(undef,n_vars,n_vars)

	for i=1:n_vars
		for j=1:n_vars
			decay[i,j] = K(i,j);
		end
	end

	#Generate random quadratic model
	#and apply exponential decay to Q
	Q = randn(rng,n_vars, n_vars);
	Qa = Q.*decay;
end


#
function q_Sigma_cd_grid(q::Int, c_grid::Array{Float64,1}, d_grid::Array{Float64,1})
    mkΣ_T = Int(1000)#numerical integration grid
    Tmax = mkΣ_T + 2
# Compute cosine transforms and weights
  fvec = collect(π*(1:1:q))'
  tvec = collect((0.5:1:mkΣ_T)/mkΣ_T)
  v = zeros(Tmax, q)
  v[1:mkΣ_T, 1:q] = sqrt(2)*cos(tvec*fvec)
  v_1 = v
#
  x = (1/mkΣ_T)*collect(1:Tmax-1)
  n_c = length(c_grid)
  n_d = length(d_grid)
  qΣcd_grid = NaN*zeros(n_c, n_d, q, q)
    #
  for ic in 1:n_c
    println("c iter $ic of $n_c")
    c = c_grid[ic]
    (c == 0) && (c = 0.01) # corner case
    cx = c*x
    # Summation by parts formula used in d0 < 0.5;
    vqd = NaN*zeros(Tmax, q)
    vqd[1:Tmax-1, :] = (c/mkΣ_T)*v_1[1:Tmax-1, :] - (v_1[2:Tmax, :] - v_1[1:Tmax-1, :])
    vqd[1, :] = vqd[1, :] - v_1[1, :]
    vqd[Tmax, :] = 0
    v_0 = mkΣ_T*vqd
    for id in 1:n_d
      d = d_grid[id]
      vqd = NaN*zeros(Tmax, q)
      d0 = d;
      (d < 0.5) && (d = d+1)
      (d == 0.5) && (d = 0.51)
      if (d0 < 0.5)
        v = v_0
      else
        v = v_1
      end
      #
      ga_0 = c^(1-2*d)*sqrt(π)*gamma(-0.5+d)/gamma(d)
      bselval = besselk(d-0.5, cx)
      ga = ((2^(1.5-d))*sqrt(π)/gamma(d))*((x/c).^(d-0.5)).*bselval
      vqd[1, :] = [ga_0 ga']*v
      for t in 1:Tmax-2
        vqd[t+1, :] = [flipdim(ga[1:t], 1)' ga_0 ga[1:Tmax-t-1]']*v
      end
      vqd[Tmax, :] = [flipdim(ga, 1)' ga_0]*v
      qΣcd = v'*vqd/(mkΣ_T^2)
      qΣcd = qΣcd/(2*π)
      #Normalize to have average value equal to unity
      tmp = mean(diag(qΣcd))
      qΣcd = qΣcd/tmp
      qΣcd_grid[ic, id, :, :] = qΣcd
    end
  end
  return qΣcd_grid
end
#
function interp_q_Sigma_cd_grid(qΣcd_grid, c, d, c_grid, d_grid)
  q = size(qΣcd_grid, 4)
  qΣcd = NaN*zeros(q,q)
  for i in 1:q
    for j in 1:q
      tmp = qΣcd_grid[:, :, i, j] # these are the values of the covariance between transforms i and j for each value of c and d
      spl = Spline2D(c_grid, d_grid, tmp)
      qΣcd[i, j] = evalgrid(spl, [c], [d])[1]
    end
  end
  tmp = mean(diag(qΣcd))
  qΣcd = qΣcd/tmp
  return qΣcd
end
#
function estimate_d(lr::bcdModel)
  q = lr.q
  cvec = lr.c.grid
  dvec = lr.d.grid
  Sigma_d_mat=NaN*zeros(lr.q, lr.q, length(dvec))
  Sigma_inv_d_mat = Sigma_d_mat
  #"Estimate" d via grid through likelihood calibration
  lvec = NaN*zeros(length(dvec))
  X_norm = lr.cost
  for i in 1:length(dvec)
    tic()
    d = dvec[i]
    Sigma = Sigma_Compute2(0, 0, d, lr.q)
    sig_inv = inv(Sigma)
    lvec[i, :] = den_invariant(X_norm, sig_inv)
    Sigma_d_mat[:, :, i] = Sigma
    Sigma_inv_d_mat[:, :, i] = sig_inv
    toc()
  end
  lr_stat = mean(lvec)./lvec
  pct_vec = 1-[2/3 0.9 0.95]'
  lr_d_cv = NaN*zeros(length(dvec), length(pct_vec))
  nsim = 50000
  Ysim = randn(lr.q, nsim)
  lvec_sim = NaN*zeros(size(dvec,1),nsim);
  for i in 1:length(dvec)
    tic()
    sig_ho = Sigma_d_mat[:, :, i]
    chol_ho = chol(Hermitian(sig_ho))
    Xsim = chol_ho*Ysim
    tmp = sqrt(sum(Xsim.^2, 1))
    Xsim_norm = Xsim./repmat(tmp,q,1);
    lvec_sim = NaN*zeros(size(dvec,1),nsim);
    for j = 1:size(dvec,1);
         sig_inv = Sigma_inv_d_mat[:,:,j];
         lvec_sim[j,:] = den_invariant(Xsim_norm,sig_inv);
     end;
     lr_stat_ho = mean(lvec_sim)./lvec_sim[i,:];
     tmp = pctile(lr_stat_ho',pct_vec);
     lr_d_cv[i, :] = tmp';
     toc()
   end
   tmp, ii = findmax(lvec)
   println("MLE (from grid) for d: $(dvec[ii])")
   condmat = repmat(log(lr_stat), 1, size(lr_d_cv, 2)) .> log(lr_d_cv)
   println("Confidence sets for d")
   cset = repmat(pct_vec, 1, 2)
   for i in 1:length(pct_vec)
     tmp = dvec[condmat[:, i].==true]
     #println(" Level = $(pct_vec[i]): $(tmp[1]) $(tmp[end])")
     #cset[i, :] = [tmp[1] tmp[end]]
   end
   #tmp, ii = findmin(abs(dvec-0))
   #lvec = lvec
   lr.d.val = dvec[ii]
   lr.stuf[:lvec] = lvec
   lr.d.lik = lvec
  #return lvec, lr_stat, Sigma_d_mat, Sigma_inv_d_mat, lr_d_cv, lvec_sim, cset
  return lr
end
#
function Sigma_Compute2(b, c, d, q; iconst = false)
  T_general = 150; T_special = 600;
  cov_Id = 0
  if c < 0.0001
    T = T_special;
    cov_Id = frac_cov(T,d)
  else
    if d == 1
      T = T_special;
      rho = 1 - c/T;
      rho = max(0, rho);
      cov_Id, tmp = ar1_cov(T,rho);
    else
      error("not implemented")
      T = T_general;
      rho = 1 - c/T;
      rho = max(0, rho);
      #cov_Id = frac_dif_cov(T, rho, d);
      cov_Id, tmp = ar1_cov(T,rho);
    end
  end
  #Add I(0) components
  scale_I0 = (b*(T^d))^2;
  #println(cov_Id); println(scale_I0); println(eye(T))
  covmat = cov_Id + scale_I0*eye(T);
  psi = psi_compute(T, q);
  if iconst == true
    psi = [ones(T,1)/T psi];
  end
  # Step 3: compute Sigma
  Sigma = psi'*covmat*psi;
  expoente = 1. - 2*d;
  Sigma = (T^expoente)*Sigma;
  #= Verify Sigma is positive definite
  ier = 0;
  tmp, tmp2 = eig(Sigma);
  if minimum(tmp) < 0
    ier = 1;
    Sigma_2 = Sigma*Sigma';
    V, D2 = eig(Sigma_2);
    D_vec = sqrt(diag(D2));
    D_mat = diagm(D_vec);
    Sigma = V*D_mat*V';
  end =#
  return Sigma
end
#


struct Symplectic2ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  b1::T
  b2::T
end

function PseudoVerletLeapfrogConstantCache(T,T2)
  a1 = T(1)
  a2 = T(0)
  b1 = T(1//2)
  b2 = T(1//2)
  Symplectic2ConstantCache{T,T2}(a1,a2,b1,b2)
end

function McAte2ConstantCache(T,T2)
  a2 = T(1 - (1/2)*sqrt(T(2)))
  a1 = T(1 - a2)
  b2 = T(1/(2*(1-a2)))
  b1 = T(1 - b2)
  Symplectic2ConstantCache{T,T2}(a1,a2,b1,b2)
end

function VerletLeapfrogConstantCache(T,T2)
  a1 = T(1//2)
  a2 = T(1//2)
  b1 = T(0)
  b2 = T(1)
  Symplectic2ConstantCache{T,T2}(a1,a2,b1,b2)
end

struct Symplectic3ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  b1::T
  b2::T
  b3::T
end

function Ruth3ConstantCache(T,T2)
  a1 = T(2//3)
  a2 = T(-2//3)
  a3 = T(1)
  b1 = T(7//24)
  b2 = T(3//4)
  b3 = T(-1//24)
  Symplectic3ConstantCache{T,T2}(a1,a2,a3,b1,b2,b3)
end

function McAte3ConstantCache(T,T2)
  a1 = T(0.9196615230173999)
  a2 = T(0.25/a1 - a1/2)
  a3 = T(1 - a1 - a2)
  b1 = T(a3)
  b2 = T(a2)
  b3 = T(a1)
  Symplectic3ConstantCache{T,T2}(a1,a2,a3,b1,b2,b3)
end

struct Symplectic4ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  b1::T
  b2::T
  b3::T
  b4::T
end

function CandyRoz4ConstantCache(T,T2)
  a1 = T((2 + T(2)^(1//3) + T(2)^(-1//3))/6)
  a2 = T((1 - T(2)^(1//3) - T(2)^(-1//3))/6)
  a3 = T(a2)
  a4 = T(a1)
  b1 = T(0)
  b2 = T((2 - T(2)^(1//3))^-1)
  b3 = T((1-T(2)^(2//3))^-1)
  b4 = T(b2)
  Symplectic4ConstantCache{T,T2}(a1,a2,a3,a4,b1,b2,b3,b4)
end

function McAte4ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1 =T( 0.515352837431122936)
  a2 =T(-0.085782019412973646)
  a3 =T( 0.441583023616466524)
  a4 =T( 0.128846158365384185)
  b1 =T( 0.134496199277431089)
  b2 =T(-0.224819803079420806)
  b3 =T( 0.756320000515668291)
  b4 =T( 0.334003603286321425)
  Symplectic4ConstantCache{T,T2}(a1,a2,a3,a4,b1,b2,b3,b4)
end

function McAte4ConstantCache(T::Type,T2::Type)
  a1 =T(parse(BigFloat,"0.515352837431122936"))
  a2 =T(parse(BigFloat,"-0.085782019412973646"))
  a3 =T(parse(BigFloat," 0.441583023616466524"))
  a4 =T(parse(BigFloat," 0.128846158365384185"))
  b1 =T(parse(BigFloat," 0.134496199277431089"))
  b2 =T(parse(BigFloat,"-0.224819803079420806"))
  b3 =T(parse(BigFloat," 0.756320000515668291"))
  b4 =T(parse(BigFloat," 0.334003603286321425"))
  Symplectic4ConstantCache{T,T2}(a1,a2,a3,a4,b1,b2,b3,b4)
end

struct Symplectic45ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
end

function CalvoSanz4ConstantCache(T,T2)
  a1= T(0.205177661542290)
  a2= T(0.403021281604210)
  a3=-T(0.12092087633891)
  a4= T(0.512721933192410)
  a5= T(0.0)
  b1= T(0.061758858135626)
  b2= T(0.33897802655364)
  b3= T(0.61479130717558)
  b4=-T(0.14054801465937)
  b5= T(0.12501982279453)
  Symplectic45ConstantCache{T,T2}(a1,a2,a3,a4,a5,b1,b2,b3,b4,b5)
end

# Broken
# http://epubs.siam.org/doi/pdf/10.1137/0916010
# On the numerical integration of ordinary differential equations by symmetric composition methods
function McAte42ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1 =T(0.40518861839525227722)
  a2 =T(-0.28714404081652408900)
  a3 = 1 - 2a1 - 2a2
  a4 = a2
  a5 = a1
  b1 = T(-3//73)
  b2 = T(17//59)
  b3 = 1 - 2b1 - 2b2
  b4 = b2
  b5 = b1
  Symplectic45ConstantCache{T,T2}(a1,a2,a3,a4,a5,b1,b2,b3,b4,b5)
end

function McAte42ConstantCache(T::Type,T2::Type)
  a1 =T(parse(BigFloat,"0.40518861839525227722"))
  a2 =T(parse(BigFloat,"-0.28714404081652408900"))
  a3 = 1 - 2a1 - 2a2
  a4 = a2
  a5 = a1
  b1 = T(-3//73)
  b2 = T(17//59)
  b3 = 1 - 2b1 - 2b2
  b4 = b2
  b5 = b1
  Symplectic45ConstantCache{T,T2}(a1,a2,a3,a4,a5,b1,b2,b3,b4,b5)
end

struct Symplectic5ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  a6::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
  b6::T
end

function McAte5ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1=T(0.339839625839110000)
  a2=T(-0.088601336903027329)
  a3=T(0.5858564768259621188)
  a4=T(-0.603039356536491888)
  a5=T(0.3235807965546976394)
  a6=T(0.4423637942197494587)
  b1=T(0.1193900292875672758)
  b2=T(0.6989273703824752308)
  b3=T(-0.1713123582716007754)
  b4=T(0.4012695022513534480)
  b5=T(0.0107050818482359840)
  b6=T(-0.0589796254980311632)
  Symplectic5ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,b1,b2,b3,b4,b5,b6)
end

function McAte5ConstantCache(T::Type,T2::Type)
  a1=T(parse(BigFloat,"0.339839625839110000"))
  a2=T(parse(BigFloat,"-0.088601336903027329"))
  a3=T(parse(BigFloat,"0.5858564768259621188"))
  a4=T(parse(BigFloat,"-0.603039356536491888"))
  a5=T(parse(BigFloat,"0.3235807965546976394"))
  a6=T(parse(BigFloat,"0.4423637942197494587"))
  b1=T(parse(BigFloat,"0.1193900292875672758"))
  b2=T(parse(BigFloat,"0.6989273703824752308"))
  b3=T(parse(BigFloat,"-0.1713123582716007754"))
  b4=T(parse(BigFloat,"0.4012695022513534480"))
  b5=T(parse(BigFloat,"0.0107050818482359840"))
  b6=T(parse(BigFloat,"-0.0589796254980311632"))
  Symplectic5ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,b1,b2,b3,b4,b5,b6)
end

struct Symplectic6ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  a6::T
  a7::T
  a8::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
  b6::T
  b7::T
  b8::T
end

function Yoshida6ConstantCache(T,T2)
  a1=T( 0.78451361047756)
  a2=T( 0.23557321335936)
  a3=T(-1.1776799841789)
  a4=T( 1.3151863206839)
  a5=T( a3)
  a6=T( a2)
  a7=T( a1)
  a8=T( 0.0)
  b1= a1/2
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T( b4)
  b6=T( b3)
  b7=T( b2)
  b8=T( b1)
  Symplectic6ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,b1,b2,b3,b4,b5,b6,b7,b8)
end

struct Symplectic62ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  a6::T
  a7::T
  a8::T
  a9::T
  a10::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
  b6::T
  b7::T
  b8::T
  b9::T
  b10::T
end

function KahanLi6ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1=T(0.39216144400731413927925056)
  a2=T(0.33259913678935943859974864)
  a3=T(-0.70624617255763935980996482)
  a4=T(0.08221359629355080023149045)
  a5=T(0.79854399093482996339895035)
  a6= a4
  a7= a3
  a8= a2
  a9= a1
  a10 = T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=b5
  b7=b4
  b8=b3
  b9=b2
  b10=b1
  Symplectic62ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,b1,b2,b3,b4,b5,b6,b7,b8,b9,b10)
end

function KahanLi6ConstantCache(T::Type,T2::Type)
  a1=T(parse(BigFloat,"0.39216144400731413927925056"))
  a2=T(parse(BigFloat,"0.33259913678935943859974864"))
  a3=T(parse(BigFloat,"-0.70624617255763935980996482"))
  a4=T(parse(BigFloat,"0.08221359629355080023149045"))
  a5=T(parse(BigFloat,"0.79854399093482996339895035"))
  a6= a4
  a7= a3
  a8= a2
  a9= a1
  a10 = T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=b5
  b7=b4
  b8=b3
  b9=b2
  b10=b1
  Symplectic62ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,b1,b2,b3,b4,b5,b6,b7,b8,b9,b10)
end

struct McAte8ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  a6::T
  a7::T
  a8::T
  a9::T
  a10::T
  a11::T
  a12::T
  a13::T
  a14::T
  a15::T
  a16::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
  b6::T
  b7::T
  b8::T
  b9::T
  b10::T
  b11::T
  b12::T
  b13::T
  b14::T
  b15::T
  b16::T
end

function McAte8ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1=T(0.74167036435061295344822780)
  a2=T(-0.40910082580003159399730010)
  a3=T(0.19075471029623837995387626)
  a4=T(-0.57386247111608226665638773)
  a5=T(0.29906418130365592384446354)
  a6=T(0.33462491824529818378495798)
  a7=T(0.31529309239676659663205666)
  a8=T(-0.79688793935291635401978884)
  a9=a7
  a10=a6
  a11=a5
  a12=a4
  a13=a3
  a14=a2
  a15=a1
  a16=T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=T((a5+a6)/2)
  b7=T((a6+a7)/2)
  b8=T((a7+a8)/2)
  b9=b8
  b10=b7
  b11=b6
  b12=b5
  b13=b4
  b14=b3
  b15=b2
  b16=b1
  McAte8ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15,a16,
                            b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15,b16)
end

function McAte8ConstantCache(T::Type,T2::Type)
  a1=T(parse(BigFloat,"0.74167036435061295344822780"))
  a2=T(parse(BigFloat,"-0.40910082580003159399730010"))
  a3=T(parse(BigFloat,"0.19075471029623837995387626"))
  a4=T(parse(BigFloat,"-0.57386247111608226665638773"))
  a5=T(parse(BigFloat,"0.29906418130365592384446354"))
  a6=T(parse(BigFloat,"0.33462491824529818378495798"))
  a7=T(parse(BigFloat,"0.31529309239676659663205666"))
  a8=T(parse(BigFloat,"-0.79688793935291635401978884"))
  a9=a7
  a10=a6
  a11=a5
  a12=a4
  a13=a3
  a14=a2
  a15=a1
  a16=T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=T((a5+a6)/2)
  b7=T((a6+a7)/2)
  b8=T((a7+a8)/2)
  b9=b8
  b10=b7
  b11=b6
  b12=b5
  b13=b4
  b14=b3
  b15=b2
  b16=b1
  McAte8ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15,a16,
                            b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15,b16)
end

struct KahanLi8ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  a6::T
  a7::T
  a8::T
  a9::T
  a10::T
  a11::T
  a12::T
  a13::T
  a14::T
  a15::T
  a16::T
  a17::T
  a18::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
  b6::T
  b7::T
  b8::T
  b9::T
  b10::T
  b11::T
  b12::T
  b13::T
  b14::T
  b15::T
  b16::T
  b17::T
  b18::T
end

function KahanLi8ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1=T(0.13020248308889008087881763)
  a2=T(0.56116298177510838456196441)
  a3=T(-0.38947496264484728640807860)
  a4=T(0.15884190655515560089621075)
  a5=T(-0.39590389413323757733623154)
  a6=T(0.18453964097831570709183254)
  a7=T(0.25837438768632204729397911)
  a8=T(0.29501172360931029887096624)
  a9=T(-0.60550853383003451169892108)
  a10=a8
  a11=a7
  a12=a6
  a13=a5
  a14=a4
  a15=a3
  a16=a2
  a17=a1
  a18=T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=T((a5+a6)/2)
  b7=T((a6+a7)/2)
  b8=T((a7+a8)/2)
  b9=T((a8+a9)/2)
  b10=b9
  b11=b8
  b12=b7
  b13=b6
  b14=b5
  b15=b4
  b16=b3
  b17=b2
  b18=b1
  KahanLi8ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15,a16,a17,a18,
                              b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15,b16,b17,b18)
end

function KahanLi8ConstantCache(T::Type,T2::Type)
  a1=T(parse(BigFloat,"0.13020248308889008087881763"))
  a2=T(parse(BigFloat,"0.56116298177510838456196441"))
  a3=T(parse(BigFloat,"-0.38947496264484728640807860"))
  a4=T(parse(BigFloat,"0.15884190655515560089621075"))
  a5=T(parse(BigFloat,"-0.39590389413323757733623154"))
  a6=T(parse(BigFloat,"0.18453964097831570709183254"))
  a7=T(parse(BigFloat,"0.25837438768632204729397911"))
  a8=T(parse(BigFloat,"0.29501172360931029887096624"))
  a9=T(parse(BigFloat,"-0.60550853383003451169892108"))
  a10=a8
  a11=a7
  a12=a6
  a13=a5
  a14=a4
  a15=a3
  a16=a2
  a17=a1
  a18=T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=T((a5+a6)/2)
  b7=T((a6+a7)/2)
  b8=T((a7+a8)/2)
  b9=T((a8+a9)/2)
  b10=b9
  b11=b8
  b12=b7
  b13=b6
  b14=b5
  b15=b4
  b16=b3
  b17=b2
  b18=b1
  KahanLi8ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15,a16,a17,a18,
                              b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15,b16,b17,b18)
end

struct SofSpa10ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  a1::T
  a2::T
  a3::T
  a4::T
  a5::T
  a6::T
  a7::T
  a8::T
  a9::T
  a10::T
  a11::T
  a12::T
  a13::T
  a14::T
  a15::T
  a16::T
  a17::T
  a18::T
  a19::T
  a20::T
  a21::T
  a22::T
  a23::T
  a24::T
  a25::T
  a26::T
  a27::T
  a28::T
  a29::T
  a30::T
  a31::T
  a32::T
  a33::T
  a34::T
  a35::T
  a36::T
  b1::T
  b2::T
  b3::T
  b4::T
  b5::T
  b6::T
  b7::T
  b8::T
  b9::T
  b10::T
  b11::T
  b12::T
  b13::T
  b14::T
  b15::T
  b16::T
  b17::T
  b18::T
  b19::T
  b20::T
  b21::T
  b22::T
  b23::T
  b24::T
  b25::T
  b26::T
  b27::T
  b28::T
  b29::T
  b30::T
  b31::T
  b32::T
  b33::T
  b34::T
  b35::T
  b36::T
end

function SofSpa10ConstantCache(::Type{T},::Type{T2}) where {T<:CompiledFloats,T2<:CompiledFloats}
  a1=T(0.07879572252168641926390768)
  a2=T(0.31309610341510852776481247)
  a3=T(0.02791838323507806610952027)
  a4=T(-0.22959284159390709415121340)
  a5=T(0.13096206107716486317465686)
  a6=T(-0.26973340565451071434460973)
  a7=T(0.07497334315589143566613711)
  a8=T(0.11199342399981020488957508)
  a9=T(0.36613344954622675119314812)
  a10=T(-0.39910563013603589787862981)
  a11=T(0.10308739852747107731580277)
  a12=T(0.41143087395589023782070412)
  a13=T(-0.00486636058313526176219566)
  a14=T(-0.39203335370863990644808194)
  a15=T(0.05194250296244964703718290)
  a16=T(0.05066509075992449633587434)
  a17=T(0.04967437063972987905456880)
  a18=T(0.04931773575959453791768001)
  a19=a17
  a20=a16
  a21=a15
  a22=a14
  a23=a13
  a24=a12
  a25=a11
  a26=a10
  a27=a9
  a28=a8
  a29=a7
  a30=a6
  a31=a5
  a32=a4
  a33=a3
  a34=a2
  a35=a1
  a36=T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=T((a5+a6)/2)
  b7=T((a6+a7)/2)
  b8=T((a7+a8)/2)
  b9=T((a8+a9)/2)
  b10=T((a9+a10)/2)
  b11=T((a10+a11)/2)
  b12=T((a11+a12)/2)
  b13=T((a12+a13)/2)
  b14=T((a13+a14)/2)
  b15=T((a14+a15)/2)
  b16=T((a15+a16)/2)
  b17=T((a16+a17)/2)
  b18=T((a17+a18)/2)
  b19=b18
  b20=b17
  b21=b16
  b22=b15
  b23=b14
  b24=b13
  b25=b12
  b26=b11
  b27=b10
  b28=b9
  b29=b8
  b30=b7
  b31=b6
  b32=b5
  b33=b4
  b34=b3
  b35=b2
  b36=b1
  SofSpa10ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15,a16,a17,a18,
                              a19,a20,a21,a22,a23,a24,a25,a26,a27,a28,a29,a30,a31,a32,a33,a34,
                              a35,a36,
                              b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15,b16,b17,b18,
                              b19,b20,b21,b22,b23,b24,b25,b26,b27,b28,b29,b30,b31,b32,b33,b34,
                              b35,b36)
end

function SofSpa10ConstantCache(T::Type,T2::Type)
  a1=T(parse(BigFloat,"0.07879572252168641926390768"))
  a2=T(parse(BigFloat,"0.31309610341510852776481247"))
  a3=T(parse(BigFloat,"0.02791838323507806610952027"))
  a4=T(parse(BigFloat,"-0.22959284159390709415121340"))
  a5=T(parse(BigFloat,"0.13096206107716486317465686"))
  a6=T(parse(BigFloat,"-0.26973340565451071434460973"))
  a7=T(parse(BigFloat,"0.07497334315589143566613711"))
  a8=T(parse(BigFloat,"0.11199342399981020488957508"))
  a9=T(parse(BigFloat,"0.36613344954622675119314812"))
  a10=T(parse(BigFloat,"-0.39910563013603589787862981"))
  a11=T(parse(BigFloat,"0.10308739852747107731580277"))
  a12=T(parse(BigFloat,"0.41143087395589023782070412"))
  a13=T(parse(BigFloat,"-0.00486636058313526176219566"))
  a14=T(parse(BigFloat,"-0.39203335370863990644808194"))
  a15=T(parse(BigFloat,"0.05194250296244964703718290"))
  a16=T(parse(BigFloat,"0.05066509075992449633587434"))
  a17=T(parse(BigFloat,"0.04967437063972987905456880"))
  a18=T(parse(BigFloat,"0.04931773575959453791768001"))
  a19=a17
  a20=a16
  a21=a15
  a22=a14
  a23=a13
  a24=a12
  a25=a11
  a26=a10
  a27=a9
  a28=a8
  a29=a7
  a30=a6
  a31=a5
  a32=a4
  a33=a3
  a34=a2
  a35=a1
  a36=T(0)
  b1= T(a1/2)
  b2=T((a1+a2)/2)
  b3=T((a2+a3)/2)
  b4=T((a3+a4)/2)
  b5=T((a4+a5)/2)
  b6=T((a5+a6)/2)
  b7=T((a6+a7)/2)
  b8=T((a7+a8)/2)
  b9=T((a8+a9)/2)
  b10=T((a9+a10)/2)
  b11=T((a10+a11)/2)
  b12=T((a11+a12)/2)
  b13=T((a12+a13)/2)
  b14=T((a13+a14)/2)
  b15=T((a14+a15)/2)
  b16=T((a15+a16)/2)
  b17=T((a16+a17)/2)
  b18=T((a17+a18)/2)
  b19=b18
  b20=b17
  b21=b16
  b22=b15
  b23=b14
  b24=b13
  b25=b12
  b26=b11
  b27=b10
  b28=b9
  b29=b8
  b30=b7
  b31=b6
  b32=b5
  b33=b4
  b34=b3
  b35=b2
  b36=b1
  SofSpa10ConstantCache{T,T2}(a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15,a16,a17,a18,
                              a19,a20,a21,a22,a23,a24,a25,a26,a27,a28,a29,a30,a31,a32,a33,a34,
                              a35,a36,
                              b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15,b16,b17,b18,
                              b19,b20,b21,b22,b23,b24,b25,b26,b27,b28,b29,b30,b31,b32,b33,b34,
                              b35,b36)
end


using MatrixCompletion

@testset "$(format("Exponential Family: forward_map[poisson]"))" begin
    let
        tc1 = rand(100)
        @test check(:l2diff, forward_map(Poisson(),tc1),exp.(tc1),0)
        @test check(:l2diff, forward_map(:Poisson,tc1),exp.(tc1),0)
        tc2 = rand(100,100)
        @test check(:l2diff, forward_map(Poisson(),tc2),exp.(tc2),0)
        @test check(:l2diff, forward_map(:Poisson,tc2),exp.(tc2),0)
    end
end

@testset "$(format("Exponential Family: forward_map[gamma]"))" begin
    let
        tc1 = rand(100)
        @test check(:l2diff, forward_map(Gamma(),tc1),1 ./ tc1,0)
        @test check(:l2diff, forward_map(:Gamma,tc1),1 ./ tc1,0)
        tc2 = rand(100,100)
        @test check(:l2diff, forward_map(Gamma(),tc2),1 ./ tc2,0)
        @test check(:l2diff, forward_map(:Gamma,tc2),1 ./ tc2,0)
    end
end

@testset "$(format("Exponential Family: forward_map[gaussian]"))" begin
    let
        tc1 = rand(100)
        @test check(:l2diff, forward_map(Gaussian(),tc1),tc1,0)
        @test check(:l2diff, forward_map(:Gaussian,tc1), tc1,0)
        tc2 = rand(100,100)
        @test check(:l2diff, forward_map(Gaussian(),tc2),tc2,0)
        @test check(:l2diff, forward_map(:Gaussian,tc2),tc2,0)
    end
end

@testset "$(format("Exponential Family: forward_map[bernoulli]"))" begin
    let
        logit = (x) -> log.(x./(1 .- x))
        tc1 = rand(100)
        tc1_logit = logit(tc1)
        @test check(:l2diff, forward_map(Bernoulli(),tc1_logit),tc1,0)
        @test check(:l2diff, forward_map(:Bernoulli,tc1_logit), tc1,0)
        tc2 = rand(100,100)
        tc2_logit = logit(tc2)
        @test check(:l2diff, forward_map(Bernoulli(),tc2_logit),tc2,0)
        @test check(:l2diff, forward_map(:Bernoulli,tc2_logit), tc2,0)
    end
end

@testset "$(format("Exponential Family: predict[poisson]"))" begin
    let
        tc1 = rand(2:20,100)
        log_tc1 = log.(tc1)
        @test check(:l2diff, predict(Poisson(),forward_map(Poisson(),log_tc1)),tc1,0.0)
        @test check(:l2diff, predict(:Poisson,forward_map(:Poisson,log_tc1)),tc1,0.0)
        @test predict(:Poisson,[1,1];custom_prediction_function=1) == -1 #
        tc2 = rand(2:20,100,100)
        log_tc2 = log.(tc2)
        @test check(:l2diff, predict(Poisson(),forward_map(Poisson(),log_tc2)),tc2,0.0)
        @test check(:l2diff, predict(:Poisson,forward_map(:Poisson,log_tc2)),tc2,0.0)
        @test predict(:Poisson,[1 1;1 1];custom_prediction_function=1) == -1 #
    end
end

@testset "$(format("Exponential Family: predict[bernoulli]"))" begin
    let
        logit = (x) -> log.(x./(1 .- x))
        tc1 = rand(100)
        tc1_int = Int.(tc1 .> 0.5)
        logit_tc1 = logit(tc1)
        @test check(:l2diff, predict(Bernoulli(), forward_map(Bernoulli(),logit_tc1)), tc1_int, 0.0)
        @test check(:l2diff, predict(:Bernoulli, forward_map(:Bernoulli, logit_tc1)), tc1_int, 0.0)
        @test predict(:Bernoulli,[1,1];custom_prediction_function=1) == -1 #
        tc2 = rand(100,100)
        tc2_int = Int.(tc2 .> 0.5)
        logit_tc2 = logit(tc2)
        @test check(:l2diff, predict(Bernoulli(),forward_map(Bernoulli(),logit_tc2)),tc2_int,0.0)
        @test check(:l2diff, predict(:Bernoulli,forward_map(:Bernoulli,logit_tc2)),tc2_int,0.0)
        @test predict(:Bernoulli,[1 1;1 1];custom_prediction_function=1) == -1 #
    end
end

@testset "$(format("Exponential Family: predict[gaussian]"))" begin
    let
        tc1 = rand(100)
        @test check(:l2diff, predict(Gaussian(),forward_map(Gaussian(),tc1)),tc1,0.0)
        @test check(:l2diff, predict(:Gaussian,forward_map(:Gaussian,tc1)),tc1,0.0)
        @test predict(:Gaussian,[1,1];custom_prediction_function=1) == -1 #
        tc2 = rand(100,100)
        @test check(:l2diff, predict(Gaussian(),forward_map(Gaussian(),tc2)),tc2,0.0)
        @test check(:l2diff, predict(:Gaussian,forward_map(:Gaussian,tc2)),tc2,0.0)
        @test predict(:Poisson,[1 1;1 1];custom_prediction_function=1) == -1 #
    end
end

@testset "$(format("Exponential Family: predict[gaussian]"))" begin
    let
        tc1 = rand(100)
        inv_tc1 = 1 ./ tc1
        @test check(:l2diff, predict(Gamma(),forward_map(Gamma(),1 ./ tc1)),tc1,0.0)
        @test check(:l2diff, predict(:Gamma,forward_map(:Gamma,1 ./ tc1)),tc1,0.0)
        @test predict(:Gamma,[1,1];custom_prediction_function=1) == -1 #
        tc2 = rand(100,100)
        inv_tc2 = 1 ./ tc2
        @test check(:l2diff, predict(Gamma(),forward_map(Gamma(),inv_tc2)),tc2,0.0)
        @test check(:l2diff, predict(:Gamma,forward_map(:Gamma,inv_tc2)),tc2,0.0)
        @test predict(:Poisson,[1 1;1 1];custom_prediction_function=1) == -1 #
    end
end


function simple_effect_probabilities(chain, data)
    γpost_x = collect(hcat(data.Xunique...)')
    γpost_mins = [minimum(γpost_x[:, i]) for i in 1:size(γpost_x, 2)]
    γpost_x = γpost_x - ones(size(γpost_x, 1)) * (γpost_mins .- 1)'

    ch_gamma_mat = collect(hcat(chain.gamma...)')
    nvar = size(γpost_x, 2)
    p_effects = zeros(nvar)
    for i in 1:nvar
        subset = deepcopy(γpost_x)
        subset[:, i] = 2 .- (subset[:, i] .!= 1)
        subset = maximum(subset, dims = 2)[:, 1] .== 1
        if sum(subset) > 0
            p_effects[i] = mean(maximum(ch_gamma_mat[:, subset], dims = 2))
        end
    end
    return DataFrame(var1 = ["x$i" for i in 1:nvar], prob = p_effects)
end


################################################################################
# This set of functions provides various genetic utilities.
################################################################################
#
# Required external packages.
#
# using Statistics, Distributions

export map_function, inverse_map_function
export hardy_weinberg_test, xlinked_hardy_weinberg_test

"""
Calculate recombination fractions based on Haldane's or Kosambi's formula.
"""
function map_function(d::T where T <: Real, choice::String)
#
  one_d = one(typeof(d))
  if choice == "Haldane"
    theta = (one_d - exp(-2d)) / 2
  elseif choice == "Kosambi"
    theta = tanh(2d) / 2
  else
  throw(ArgumentError(
    "Only Haldane's or Kosambi's map function is allowed.\n \n"))
  end
  return theta
end # function map_function

"""
Calculates genetic map distances based on Haldane's or Kosambi's formula.
"""
function inverse_map_function(theta::T where T <: Real, choice::String)
#
  one_th = one(typeof(theta))
  if choice == "Haldane"
    if theta >= one_th / 2
      d = Inf
    else
      d = -(log(one_th - 2 * theta)) / 2
    end
  elseif choice == "Kosambi"
    if theta >= one_th / 2
      d = Inf
    else
      d = log((one_th + 2 * theta) / (one_th - 2 * theta)) / 4
    end
  else
    throw(ArgumentError(
      "Only Haldane's or Kosambi's map function is allowed.\n \n"))
  end
  return d
end # function inverse_map_function

"""
Test for Hardy-Weinberg equilibrium at a SNP.
The genotype vector conveys for each person his or her number of
reference alleles. Thus, all genotypes belong to {0, 1, 2, Missing}.
The pvalue of the chi-square statistic is returned.
"""
function hardy_weinberg_test(genotype::Vector)
#
  (expected, observed) = (zeros(3), zeros(3))
  #
  # Tally observations per cell.
  #
  n = 0
  for i = 1:length(genotype)
    if !ismissing(genotype[i])
      n = n + 1
      g = round(Int, genotype[i])
      j = g + 1
      observed[j] = observed[j] + 1.0
    end
  end
  if n == 0
    return 1.0
  end
  #
  # Compute allele frequencies.
  #
  p = (2.0 * observed[1] + observed[2]) / (2 * n)
  if p <= 0.0 || p >= 1.0
    return 1.0
  else
    q = 1.0 - p
    #
    # Sum over cells of the contingency table.
    #
    expected[1] = n * p^2
    expected[2] = n * 2.0 * p * q
    expected[3] = n * q^2
    test = 0.0
    for i = 1:3
    test = test + (observed[i] - expected[i])^2 / expected[i]
  end
  return ccdf(Chisq(1), test)
  end
end # function hardy_weinberg_test

"""
Test for Hardy-Weinberg equilibrium at an xlinked SNP.
The genotype vector conveys for each person his or her number of
reference alleles. Thus, all genotypes belong to {0, 1, 2, Missing}.
The p-value of the chi-square statistic is returned.
"""
function xlinked_hardy_weinberg_test(genotype::Vector,
  male::BitArray{1})

  (expected_female, observed_female) = (zeros(3), zeros(3))
  (expected_male, observed_male) = (zeros(2), zeros(2))
  #
  # Tally observations per cell and count members of each sex.
  #
  (females, males) = (0, 0)
  for i = 1:length(genotype)
    if !ismissing(genotype[i])
      g = round(Int, genotype[i])
      if male[i]
        males = males + 1
        j = min(g + 1, 2)
        observed_male[j] = observed_male[j] + 1.0
      else
        females = females + 1
        j = g + 1
        observed_female[j] = observed_female[j] + 1.0
      end
    end
  end
  if females + males == 0
    return 1.0
  end
  #
  # Compute allele frequencies.
  #
  p = 2.0 * observed_female[1] + observed_female[2] + observed_male[1]
  p = p / (2 * females + males)
  if p <= 0.0 || p >= 1.0
    return 1.0
  end
  q = 1.0 - p
  #
  # Sum over female cells of the contingency table.
  #
  expected_female[1] = females * p^2
  expected_female[2] = females * 2.0 * p * q
  expected_female[3] = females * q^2
  test = 0.0
  for i = 1:3
    term = (observed_female[i] - expected_female[i])^2 / expected_female[i]
    test = test + term
  end
  #
  # Sum over male cells of the contingency table.
  #
  expected_male[1] = males * p
  expected_male[2] = males * q
  for i = 1:2
    term = (observed_male[i] - expected_male[i])^2 / expected_male[i]
    test = test + term
  end
  return ccdf(Chisq(2), test)
end # function xlinked_hardy_weinberg_test

# Testing Routines
#
# d = rand()
# choice = "Haldane";
# theta = map_function(d, choice)
# di = inverse_map_function(theta, choice);
# println(di);
# d = rand()
# choice = "Kosambi";
# theta = map_function(d, choice)
# di= inverse_map_function(theta, choice);
# println(di);
# p = 1 / 3;
# n = 100;
# x = zeros(n);
# for i = 1:n
#   u = rand(2)
#   if u[1] <= p; x[i] = x[i] + 1.0; end
#   if u[2] <= p; x[i] = x[i] + 1.0; end
# end
# pvalue = hardy_weinberg_test(x);
# println("pvalue = ",pvalue);
# male = falses(n);
# male[1:div(n, 2)] .= true;
# for i = 1:div(n, 2)
#   x[i] = 0.0
#   u = rand(1)
#   if u[1] <= p; x[i] = x[i] + 1.0; end
# end
# pvalue = xlinked_hardy_weinberg_test(x, male);
# println("pvalue = ",pvalue);


#computes the energy of a sequence
#energy(epistasis, parameters, sequence)
function Energy(epistasis::EpistasisMu, x::Vector{Float64}, variant)
	return energy(epistasis, x, variant)
end

#computes the binding probability of a sequence
#pb(epistasis, parameters, sequence)
function pb(epistasis::EpistasisMu, x::Vector{Float64}, variant)
	E = Energy(epistasis, x, variant)
	return 1/(1+exp(E))
end

export Energy, pb


using ForwardDiff
using LinearAlgebra

let
	f(x) = x+x
	g(x) = 5*x
	h(x) = dot(x, x)

	# Once we get ForwardDiff working, this should work
	# h(x::VecE) = dot(x, x)

	@test ForwardDiff.jacobian(f, VecE2(1.0, 2.0)) == [2.0 0.0; 0.0 2.0]
	@test ForwardDiff.jacobian(g, VecE2(1.0, 2.0)) == [5.0 0.0; 0.0 5.0]
	@test ForwardDiff.gradient(h, VecE2(1.0, 2.0)) == [2.0, 4.0]
	@test ForwardDiff.hessian(h, VecE2(1.0, 2.0)) == [2.0 0.0; 0.0 2.0]

	@test ForwardDiff.jacobian(f, VecE3(1.0, 2.0, 3.0)) == 2.0*Matrix{Float64}(I, 3, 3)
	@test ForwardDiff.jacobian(g, VecE3(1.0, 2.0, 3.0)) == 5.0*Matrix{Float64}(I, 3, 3)
	@test ForwardDiff.gradient(h, VecE3(1.0, 2.0, 3.0)) == [2.0, 4.0, 6.0]
	@test ForwardDiff.hessian(h, VecE3(1.0, 2.0, 3.0)) == 2.0*Matrix{Float64}(I, 3, 3)

	@test ForwardDiff.jacobian(f, VecSE2(1.0, 2.0, 3.0)) == 2.0*Matrix{Float64}(I, 3, 3)

	# Once we get ForwardDiff working, this should work...
	# h(x::VecSE2) = normsquared(VecE2(x))
	# @test ForwardDiff.jacobian(g, VecSE2(1.0, 2.0, 3.0)) == diagm([5.0, 5.0, 1.0])
	# @test ForwardDiff.gradient(h, VecSE2(1.0, 2.0, 3.0)) == [2.0, 4.0, 0.0]
	# @test ForwardDiff.hessian(h, VecSE2(1.0, 2.0, 3.0)) == [2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 0.0]
end


include("normcol_lessequal.jl")
include("diagadd!.jl")

function add_sub!(aMat::Matrix{Float64}, bMat::Matrix{Float64}, cMat::Matrix{Float64})
   # aMat <- aMat + bMat - cMat
    Dim = size(aMat)
    for j=1:Dim[2], i=1:Dim[1]
        @inbounds aMat[i,j] += bMat[i,j] - cMat[i,j]
    end
end

function updateD!(D::Array{Any,1}, A::Array{Any,1}, DataMat::Array{Any,1})
   # Update D by Eq. (12)
    for (i, Aᵢ) in enumerate(A)
        @inbounds TempData::Matrix{Float64} = DataMat[i]
        ρ::Float64 = 1.
        rate_ρ::Float64 = 1.2
        @inbounds TempS::Matrix{Float64} = D[i]
        TempT::Matrix{Float64} = zeros(TempS)
        @inbounds preD::Matrix{Float64} = D[i]
        Iter::Int = 1
        lossᴰ::Float64 = 1.
        while lossᴰ > 1e-8 && Iter < 100

            tempMat::Matrix{Float64} = TempData*Aᵢ'
            tempMat += (TempS-TempT) .* ρ # tempMat <- tempMat + ρ(TempS - TempT)
            tempMatCoef::Matrix{Float64} = Aᵢ*Aᵢ'
            diagadd!(tempMatCoef, ρ)
            TempD::Matrix{Float64} = tempMat/tempMatCoef

            TempS = normcol_lessequal(TempD+TempT)
            add_sub!(TempT, TempD, TempS) # TemP <- TemP + (TempD-TempS)
            ρ *= rate_ρ
            lossᴰ = mean(abs2(preD-TempD))
            preD = TempD
            Iter += 1
        end
        @inbounds D[i] = preD
    end
end


for n in 4:10
    # setup square matrices
    S = [BlockMatrix(rand(2n, 2n), [n, n], [n, n]) for _ in 1:3]
    id = BlockMatrix(Matrix(1.0I, 2n, 2n), [n, n], [n, n])
    @testset "verify identity" for M in S
        @test M == rred_star(M, id) == rred_star(id, M)
    end
    @testset "verify associativity" for A in S, B in S, C in S
        @test (rred_star(A, rred_star(B, C)) ≈ rred_star(rred_star(A, B), C))
    end
    @testset "verify inverse" for M in S
        M⁻¹ = inv(M)
        @test rred_star(M, M⁻¹) ≈ rred_star(M⁻¹, M) ≈ id
    end
end


using JLD2; using PyPlot; using Statistics; using DelimitedFiles

# Load in the posteriors:
@load "../../../data/T1_hmc_total_02212020.jld2" state_total

clf()
# Make a confidence plot of eccentricity vectors, and
# show the longitudes of periastron.

nplanet = 7

e2 = 0.015; e1 = -0.015
ncos = 30
nsin = 30
evecgrid = zeros(nplanet,ncos,nsin)
ecosgrid = zeros(ncos)
esingrid = zeros(nsin)
for i=1:ncos
  ecosgrid[i] = (i-0.5)*(e2-e1)/ncos+e1
  esingrid[i] = (i-0.5)*(e2-e1)/nsin+e1
end

nsamp = 224000

conf = [0.95, 0.683]
nconf = length(conf)
level_ecc = zeros(nplanet,nconf)
planet_name=["b","c","d","e","f","g","h"]

cname = ["C0","C1","C2","C3","C4","C5","C6","C7"]

for j=1:nplanet
  ncosp = 10000
  # Set up vectors for samples:
  ecossamp= zeros(nsamp)
  esinsamp = zeros(nsamp)
  # Loop over number of samples:
  for isamp=1:nsamp
    # Draw from posterior distribution:
    # Compute the mass of planet in Earth mass units:
    ecossamp[isamp] = state_total[(j-1)*5+4,isamp]
    esinsamp[isamp] = state_total[(j-1)*5+5,isamp]
    # Now compute density distribution:
    if ecossamp[isamp] > e1 && ecossamp[isamp] < e2 && esinsamp[isamp] > e1 && esinsamp[isamp] < e2
      ir = ceil(Int64,(esinsamp[isamp]-e1)/(e2-e1)*ncos)
      im = ceil(Int64,(ecossamp[isamp]-e1)/(e2-e1)*nsin)
      evecgrid[j,ir,im] +=1.0
    end
  end
  # Set levels:
  evecgrid[j,:,:] /= float(sum(evecgrid[j,:,:]))
  evecgrid_sort = sort(vec(evecgrid[j,:,:]))
  i0 = nsin*ncos
  evecgrid_cum = evecgrid_sort[i0]
  while evecgrid_cum < conf[1] && i0 > 1
    i0 -= 1
    for k=1:nconf
      if evecgrid_cum < conf[k] && (evecgrid_cum+evecgrid_sort[i0]) > conf[k]
        level_ecc[j,k] = 0.5*(evecgrid_sort[i0]+evecgrid_sort[i0+1])
      end
    end
    evecgrid_cum += evecgrid_sort[i0]
  end
  println(j," ",mean(ecossamp)," ",std(ecossamp)," fractional error: ",std(ecossamp)/mean(ecossamp)*100.0,"%")
end
mopt = readdlm("../../../data/optimum_values.txt")
figure(figsize=(6,6))
# Plot maximum likelihood values:
for j=1:nplanet
  cs = contour(ecosgrid,esingrid,evecgrid[j,:,:],levels = level_ecc[j,:],colors=cname[j],linewidth=2.0)
  plot(mopt[(j-1)*5+4],mopt[(j-1)*5+5],"o",label=planet_name[j],color=cname[j])
  text(mopt[(j-1)*5+4]+0.0005,mopt[(j-1)*5+5]-0.0005,planet_name[j],color=cname[j],fontsize=15)
end
#legend(fontsize=15)
xlabel(L"$e \cos{\omega}$",fontsize=15)
ylabel(L"$e \sin{\omega}$",fontsize=15)
axis([-0.0125,0.0125,-0.0125,0.0125])
plot([-0.015,0.015],[0,0],linestyle=":",color="k")
plot([0,0],[-0.015,0.015],linestyle=":",color="k")
tight_layout()
savefig("../esin_vs_ecos.pdf",bbox_inches="tight")