module Env

export State, newstate, onestep, reward

# この実験における定数
const g = 9.80665   # 重力加速度
const M = 1.0       # カートの質量
const m = 0.1       # ポールの質量
const l = 0.5       # ポールの半分の長さ
const fps = 50      # frames per second
const tau = 1 / fps # 制御周期

const State = Vector{Float64}

# 最初の状態
# [x, theta, xdot, thetadot]
const initstate = [0.0, -pi, 0.0, 0.0]

# あとの計算で使う
const ml = m * l
const mass = M + m

function newstate()
    copy(initstate)
end

function onestep(s, a)
    rungekuttasolve(s, a, tau)
end

function reward(s, a)
    x, θ, xdot, θdot = s
    if abs(x) > 2.0
        return -2.0
    end
    -abs(θ) + pi / 2
end

# 状態 s で力 u を加えたときの微分
function differential(s, u)
    x, θ, xdot, θdot = s
    sinθ = sin(θ)
    cosθ = cos(θ)

    xddot = (4u / 3 + 4ml * θdot^2 * sinθ / 3 - m * g * sin(2θ) / 2) / (4mass - m * cosθ^2)
    θddot = (mass * g * sinθ - ml * θdot^2 * sinθ * cosθ - u * cosθ) /
        (4mass * l / 3 - ml * cosθ^2)

    [xdot, θdot, xddot, θddot]
end

# オイラー法を用いて微分方程式を解く
function eulersolve(s, sdot, dt)
    s + sdot * dt
end

# ルンゲクッタ法を用いて微分方程式を解く
function rungekuttasolve(s, u, dt)
    k1 = differential(s, u)
    s1 = eulersolve(s, k1, dt / 2)
    k2 = differential(s1, u)
    s2 = eulersolve(s, k2, dt / 2)
    k3 = differential(s2, u)
    s3 = eulersolve(s, k3, dt)
    k4 = differential(s3, u)

    snext = s + (k1 + 2k2 + 2k3 + k4) * dt / 6
    snext[2] = mod2pi(snext[2] + pi) - pi
    snext
end

end


using Distributions
using Iterators
using Gurobi
using JLD

using EmpiricalBayes

nreps = 50
true_dists = [MixtureModel([ Normal(-0.3,.5), Normal(1.05,.5)])]
ys = collect(range(0, stop=3, length=7))
minimax_σs = [0.1, 0.5, 1.0]
ns = [20_000; 100_000]

combs = []

for (true_dist, y, σ, n) in product(true_dists, ys, minimax_σs, ns)
    push!(combs, Dict(:dist=>true_dist, :y=>y, :σ=>σ, :n=>n) )
end

i = parse(Int64, ARGS[1])

comb = combs[i]

n = comb[:n]
σ = comb[:σ]
true_dist = comb[:dist]
y = comb[:y]

marginal_grid = collect(range(-10, stop=10, length=1001));
prior_grid = collect(range(-3, stop=3, length=51));
d_true = NormalConvolutionProblem(true_dist, marginal_grid)
f = BinnedMarginalDensity(d_true)
m = n
ds = MixingNormalConvolutionProblem(Normal, 0.5, prior_grid, marginal_grid)

sim_array = []

for l = 1:nreps
    @show l
     Xs = rand(d_true,m)
     tst = donoho_ci(Xs, marginal_grid, prior_grid, y, σ, d_true)
     tst2 = donoho_ci(Xs, marginal_grid, prior_grid, y, σ, NPMLE)
     push!(sim_array, (tst, tst2))
end

res = (i, comb, sim_array)

save("/scratch/users/ignat/sims/Apr5/mysim_$i.jld", "res", res)


# Continuous interpolation of discrete frequency-domain data
using Interpolations
import Interpolations: interpolate

function interpolate(s::Spectrum)
    l,h = limits(s)
    fn = extrapolate(interpolate(s[:], BSpline(Cubic(Natural(OnGrid())))), Flat())
    scaled = scale(fn, range(l; stop=h, length=length(s)))
    δ -> scaled(h + (l - δ))
end

function resample(s::Spectrum, shifts::AbstractArray)
    intp = interpolate(s)
    intp.(shifts)
end


module AffineArithmetic

using IntervalArithmetic

import Base: +, -, *, /, ^, ==,
            zero, one, range,
            show

export Affine, reset_affine_index

include("AF1.jl")
include("AF.jl")
include("AFF.jl")
export AF1, AF, AFF

const affine_index = [1]  # which affine vector index to use

reset_affine_index() = affine_index[1] = 1

@doc """
An affine quantity for affine arithmetic.
The usual way to create an affine quantity is from an interval `X`, via `Affine(X)`.
"""
struct Affine{T<:AbstractFloat}
    c::T   # mid-point
    γ::Vector{T}  # error terms
end

function show(io::IO, C::Affine)
    print(io, "⟨", C.c, "; ", C.γ, "⟩")
end

==(C::Affine, D::Affine) = C.c == D.c && C.γ == D.γ

@doc """
    Affine(X::Interval)

Construct a new `Affine` quantity from an `Interval`.
"""
function Affine(X::Interval)
    c = mid(X)
    r = radius(X)

    index = affine_index[1]
    affine_index[1] += 1

    γ = zeros(index)
    γ[end] = r

    return Affine(c, γ)
end

# conversion of numerical constant to affine:
Affine(c::Real) = Affine(c, Float64[])

range(C::Affine) = C.c + sum(abs.(C.γ))*(-1..1)

range(X::Interval) = X

# morally:
# +(C::Affine, D::Afine) = Affine(C.c + D.c, C.γ + D.γ)

for op in (:+, :-)
    @eval function $op(C::Affine, D::Affine)
        k = length(C.γ)
        l = length(D.γ)

        # account for unequal lengths:
        m = min(k, l)
        common_γ = $op(C.γ[1:m], D.γ[1:m])

        if l >= k
            γ = [ common_γ; D.γ[m+1:l] ]
        else
            γ = [ common_γ; C.γ[m+1:k] ]
        end

        Affine($op(C.c, D.c), γ)
    end
end

+(C::Affine, α::Real) = Affine(C.c + α, C.γ)
+(α::Real, C::Affine) = C + α

-(C::Affine, α::Real) = Affine(C.c - α, C.γ)
-(α::Real, C::Affine) = Affine(α - C.c, [-x for x in C.γ])

function *(C::Affine, D::Affine)

    c = C.c
    d = D.c

    k = length(C.γ)
    l = length(D.γ)

    # account for unequal lengths:
    m = min(k, l)
    common_γ = d * C.γ[1:m] + c * D.γ[1:m]

    error_bound = sum(abs.(C.γ)) * sum(abs.(D.γ))

    if l >= k
        γ = [ common_γ; c*D.γ[m+1:l]; error_bound ]
    else
        γ = [ common_γ; d*C.γ[m+1:k]; error_bound ]
    end

    # |γ|₁ = sum(abs.(γ))

    Affine(C.c * D.c, γ)
end

*(α::Real, C::Affine) = Affine(α*C.c, α*C.γ)
*(C::Affine, α::Real) = α * C

^(x::Affine, n::Integer) = Base.power_by_squaring(x, n)

Base.literal_pow(::typeof(^), x::Affine{T}, ::Val{p}) where {T,p} = x^p

eltype(C::Affine{T}) where T = T
zero(C::Affine{T}) where T = Affine(zero(T))
zero(::Type{Affine{T}}) where T = Affine(zero(T))

one(C::Affine{T}) where T = Affine(one(T))
one(::Type{Affine{T}}) where T = Affine(one(T))

#one(C::Affine)

end


using LinearAlgebra, NLPModels, Stopping

include("NewtonSolver.jl")
include("Armijo.jl")
include("NewtonStop.jl")
include("prepare_LS.jl")
include("NewtonStopLS.jl")


cube_distance(a::Array{Int,1}) = Int(sum(abs.(a))/2)

function puzzle11(path::String=joinpath(@__DIR__, "..", "data/11.txt"))
    instructions = readdlm(path, ',')
    location = zeros(Int, 3)

    hexsteps = Dict(
        "n"  => [0, 1, -1],
        "ne" => [1, 0, -1],
        "se" => [1, -1, 0],
        "s"  => [0, -1, 1],
        "sw" => [-1, 0, 1],
        "nw" => [-1, 1, 0]
    )

    furthest_distance = 0
    for step in instructions
        location += hexsteps[step]
        furthest_distance = max(furthest_distance, cube_distance(location))
    end

    return [cube_distance(location), furthest_distance]

end


using Match

input = joinpath(@__DIR__, "input")
cmds = readlines(input)

@enum Opmode direct waypoint

mutable struct Ship
    waypoint::Array{Int,1}
    location::Array{Int,1}
    opmode::Opmode
end

const rturn = [0 -1 ; 1 0]
const lturn = [0 1 ; -1 0]

function move!(ship, commands)
    for cmd in commands
        instruction = cmd[1:1]
        value = parse(Int, cmd[2:end])
        move!(ship, instruction, value)
    end
end

function move!(ship::Ship, instruction::String, value::Int)
    @match instruction begin
        "N", if ship.opmode == direct end => begin ship.location += [0, value] end
        "S", if ship.opmode == direct end => begin ship.location += [0, -1 * value] end
        "W", if ship.opmode == direct end => begin ship.location += [value, 0] end
        "E", if ship.opmode == direct end => begin ship.location += [-1 * value, 0] end
        "N", if ship.opmode == waypoint end => begin ship.waypoint += [0, value] end
        "S", if ship.opmode == waypoint end => begin ship.waypoint += [0, -1 * value] end
        "W", if ship.opmode == waypoint end => begin ship.waypoint += [value, 0] end
        "E", if ship.opmode == waypoint end => begin ship.waypoint += [-1 * value, 0] end
        "L" => begin ship.waypoint = lturn^(value / 90) * ship.waypoint end
        "R" => begin ship.waypoint = rturn^(value / 90) * ship.waypoint end
        "F" => begin ship.location += ship.waypoint * value end
    end
end

ship = Ship(
    [-1, 0],
    [0, 0],
    direct
)

move!(ship, cmds)
(fx, fy) = Tuple(ship.location)
p1 = abs(fx) + abs(fy)

ship = Ship(
    [-10, 1],
    [0, 0],
    waypoint
)

move!(ship, cmds)
(fx, fy) = Tuple(ship.location)
p2 = abs(fx) + abs(fy)

@assert(p1 == 445)
@assert(p2 == 42495)

println("-----------------------------------------------------------------------")
println("rain risk -- part one :: $p1")
println("rain risk -- part two :: $p2")
println("-----------------------------------------------------------------------")


function test_convergence(bquad, sampler; D=3)
    rng = Random.MersenneTwister(42)
    prior = MvNormal(0.5 * ones(D))
    likelihood = MvNormal(7.0 * ones(D))
    logintegrand = x -> logpdf(likelihood, x)
    m = BayesModel(prior, logintegrand)
    pZ, x = bquad(m, sampler; nsamples=100)
    trueZ = exp(-0.5 * (logdet(cov(prior) + cov(likelihood)) + D * log(2π)))
    @test mean(pZ) ≈ trueZ rtol = 1e-2
end


#############################################################################
# JuMP
# An algebraic modeling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# mindistortion.jl
#
# This example arises from computational geometry, in particular the problem
# of embedding a general finite metric space into a euclidean space.
#
# It is known that the 4-point metric space defined by the star graph:
#   x
#    \
#     x — x
#    /
#   x
# where distances are computed by length of the shortest path between
# vertices, cannot be exactly embedded into a euclidean space of any
# dimension.
#
# Here we will formulate and solve an SDP to compute the best possible
# embedding, that is, the embedding f() that minimizes the distortion c such
# that (1/c)*D(a,b) ≤ ||f(a)-f(b)|| ≤ D(a,b) for all points (a,b), where
# D(a,b) is the distance in the metric space.
#
# Any embedding can be characterized by its Gram matrix Q, which is PSD,
# and ||f(a)-f(b)||^2 = Q[a,a] + Q[b,b] - 2Q[a,b]
# We can therefore constrain
# D[i,j]^2 ≤ Q[i,i] + Q[j,j] - 2Q[i,j] ≤ c^2*D[i,j]^2
# and minimize c^2, which gives us the SDP formulation below.
#
# For more detail, see
# "Lectures on discrete geometry" by J. Matoušek, Springer, 2002, pp. 378-379.

using JuMP

m = Model()

D = [0.0 1.0 1.0 1.0
     1.0 0.0 2.0 2.0
     1.0 2.0 0.0 2.0
     1.0 2.0 2.0 0.0]

@variable(m, cSq >= 1.0)

@variable(m, Q[1:4,1:4], SDP)

for i in 1:4
    for j in (i+1):4
        @constraint(m, D[i,j]^2 <= Q[i,i] + Q[j,j] - 2Q[i,j])
        @constraint(m, Q[i,i] + Q[j,j] - 2Q[i,j] <= D[i,j]^2*cSq )
    end
end

@objective(m, Min, cSq)

solve(m)

println(getvalue(cSq))


"""
    SVESMtrain(svr::LIBSVM.AbstractSVR, esn::AbstractLeakyESN [, y_target::AbstractArray{Float64}])

Train the ESN using SVM regression, as described in [1].

[1] Shi, Zhiwei, and Min Han. "Support vector echo-state machine for chaotic time-series prediction." IEEE Transactions on Neural Networks 18.2 (2007): 359-372.

"""
function SVESMtrain(svr::LIBSVM.AbstractSVR,
    esn::AbstractLeakyESN;
    y_target::AbstractArray{Float64} = esn.train_data)

    states_new = nla(esn.nla_type, esn.states)

    if size(y_target, 1) == 1
        fitted_svr = LIBSVM.fit!(svr, states_new', vec(y_target))
    else
        fitted_svr = []
        for i=1:size(y_target, 1)
            push!(fitted_svr, LIBSVM.fit!(svr, states_new', y_target[i,:]))
        end
    end
    return fitted_svr
end

"""
    SVESM_direct_predict(esn::AbstractLeakyESN, test_in::AbstractArray{Float64}, fitted_svr::LIBSVM.AbstractSVR)

Given the input data return the corresponding predicted output, as described in [1].

[1] Shi, Zhiwei, and Min Han. "Support vector echo-state machine for chaotic time-series prediction." IEEE Transactions on Neural Networks 18.2 (2007): 359-372.

"""
function SVESM_direct_predict(esn::AbstractLeakyESN,
    test_in::AbstractArray{Float64},
    fitted_svr::LIBSVM.AbstractSVR)

    x = esn.states[:, end]
    prediction_states = zeros(Float64, size(esn.states, 1), size(test_in, 2))

    if esn.extended_states == false
        for i=1:size(test_in, 2)
            x_new = nla(esn.nla_type, x)
            x = leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x_new, test_in[:,i])
            prediction_states[:, i] = x
        end
    else
        for i=1:size(test_in, 2)
            x_new = nla(esn.nla_type, x)
            x = vcat(leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x_new[1:esn.res_size], test_in[:,i]), test_in[:, i])
            prediction_states[:,i] = x
        end
    end
    output = LIBSVM.predict(fitted_svr, prediction_states')
    return output
end

#predict if one-variable timeseries is provided
function SVESMpredict(esn::AbstractLeakyESN,
    predict_len::Int,
    fitted_svr::LIBSVM.AbstractSVR)

    output = zeros(Float64, esn.in_size, predict_len)
    x = esn.states[:, end]
    if esn.extended_states == false
        for i=1:predict_len
            x_new = hcat(nla(esn.nla_type, x)...)
            out = LIBSVM.predict(fitted_svr, x_new)
            output[:, i] = out
            x = leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x, out)
        end
    else
        for i=1:predict_len
            x_new = hcat(nla(esn.nla_type, x)...)
            out = LIBSVM.predict(fitted_svr, x_new)
            output[:, i] = out
            x = vcat(leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x[1:esn.res_size], out), out)
        end
    end
    return output
end

#predict for multidimensional timeseries
"""
    SVESMpredict(esn::AbstractLeakyESN, predict_len::Int, fitted_svr::AbstractArray{Any})

Return the prediction for a given length of the constructed ESN struct using SVMs.
"""
function SVESMpredict(esn::AbstractLeakyESN,
    predict_len::Int,
    fitted_svr::AbstractArray{Any})

    output = zeros(Float64, esn.in_size, predict_len)
    x = esn.states[:, end]

    if esn.extended_states == false
        for i=1:predict_len
            x_new = hcat(nla(esn.nla_type, x)...)
            out = []
            for i=1:size(fitted_svr, 1)
                push!(out, LIBSVM.predict(fitted_svr[i], x_new)[1])
            end
            output[:, i] = out
            x = leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x, out)
        end
    else
        for i=1:predict_len
            x_new = hcat(nla(esn.nla_type, x)...)
            out = Array{Float64}(undef, esn.out_size)
            for i=1:size(fitted_svr, 1)
                out[i] = LIBSVM.predict(fitted_svr[i], x_new)[1]
            end
            output[:, i] = out
            x = vcat(leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x[1:esn.res_size], out), out)
        end
    end
    return output
end

"""
    SVESMpredict_h_steps(esn::AbstractLeakyESN, predict_len::Int, h_steps::Int, test_data::AbstractArray{Float64}, fitted_svr::LIBSVM.AbstractSVR)

Return the prediction for h steps ahead of the constructed ESN struct using SVMs.
"""
function SVESMpredict_h_steps(esn::AbstractLeakyESN,
    predict_len::Int,
    h_steps::Int,
    test_data::AbstractArray{Float64},
    fitted_svr::LIBSVM.AbstractSVR)

    output = zeros(Float64, esn.in_size, predict_len)
    x = esn.states[:, end]
    if esn.extended_states == false
        for i=1:predict_len
            x_new = hcat(nla(esn.nla_type, x)...)
            out = LIBSVM.predict(fitted_svr, x_new)
            output[:, i] = out
            if mod(i, h_steps) == 0
                x = leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x, test_data[:,i])
            else
                x = leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x, out)
            end
        end
    else
        for i=1:predict_len
            x_new = hcat(nla(esn.nla_type, x)...)
            out = LIBSVM.predict(fitted_svr, x_new)
            output[:, i] = out
            if mod(i, h_steps) == 0
                x = vcat(ReservoirComputing.leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x[1:esn.res_size], test_data[:,i]), test_data[:,i])
            else
                x = vcat(ReservoirComputing.leaky_fixed_rnn(esn.activation, esn.alpha, esn.W, esn.W_in, x[1:esn.res_size], out), out)
            end
        end
    end
    return output
end


"""
    LegendreTransform(; modes)

Constructs a legendre transform with modes `modes` that
operates on the 2nd, 3rd, ...N-modes+1-st dimension of the data.

# Example
```
julia> LegendreTransform(modes = (12, 3, 4, 12))
```
"""
Base.@kwdef struct LegendreTransform{N, T} <: AbstractTransform
    modes::NTuple{N, T}
end

function forward(::LegendreTransform, x)
    L = legendre_transform_forward_matrix.(size(x)[2:(end - 1)])
    return apply_legendre(L, x)
end

function inverse(::LegendreTransform, x)
    L⁻¹ = legendre_transform_inverse_matrix.(size(x)[2:(end - 1)])
    return apply_legendre(L⁻¹, x)
end

function truncate_modes(tr::LegendreTransform{N}, coeff) where {N}
    # Returns a low-pass filtered version of coeff assuming
    # that coeff is a tensor of spectral weights.
    # Ex.: tr.modes = (3,)
    # [0, 1, 2, 3, 4, 5] -> [0, 1, 2]
    # [a, b, c, d, e, f] -> [a, b, c]

    # indices for the spectral coefficients that we need to retain
    inds = [collect(1:m) for (s, m) in zip(size(coeff)[2:(end - 1)], tr.modes)]
    coeff_truncated = OperatorFlux.mview(coeff, inds, Val(N))

    return coeff_truncated
end

function pad_modes(
    ::LegendreTransform,
    coeff,
    size_pad::NTuple{N, T},
) where {N, T}
    # return a padded-with-zeros version of coeff assuming
    # that coeff is a tensor of spectral weights, thereby inflating coeff.
    size_space = size(coeff)[2:(end - 1)] # sizes of space-like dimensions of coeff

    # generate a zero array and indices that point to the filled-in
    # locations
    coeff_padded =
        zeros(eltype(coeff), (size(coeff)[1], size_pad..., size(coeff)[end]))
    inds = [collect(1:m) for (s, m) in zip(size_pad, size(coeff)[2:(end - 1)])]

    coeff_padded_view = OperatorFlux.mview(coeff_padded, inds, Val(N))
    coeff_padded_view .= coeff

    return coeff_padded
end

# utils
function vandermonde(x, N)
    # create view to assign values
    P = zeros(length(x), N + 1)
    P⁰ = view(P, :, 0 + 1)
    @. P⁰ = 1

    # explicitly compute second coefficient
    if N == 0
        return P
    end

    P¹ = view(P, :, 1 + 1)
    @. P¹ = x

    if N == 1
        return P
    end

    for n in 1:(N - 1)
        # get views for ith, i-1th, and i-2th columns
        Pⁿ⁺¹ = view(P, :, n + 1 + 1)
        Pⁿ = view(P, :, n + 0 + 1)
        Pⁿ⁻¹ = view(P, :, n - 1 + 1)

        # compute coefficients for ith column
        @. Pⁿ⁺¹ = ((2n + 1) * x * Pⁿ - n * Pⁿ⁻¹) / (n + 1)
    end

    return P
end

function legendre_transform_inverse_matrix(N)
    # get the legendre points to construct transformation
    # matrix. For small transforms this is performant.
    x, _ = GaussQuadrature.legendre(N, GaussQuadrature.both)
    return vandermonde(x, N - 1)
end

function legendre_transform_forward_matrix(N)
    return legendre_transform_inverse_matrix(N) \ I
end

function apply_legendre(w::Tuple{S}, x) where {S}
    F1, = w
    @tullio c[s, i, b] := F1[i, ii] * x[s, ii, b]
end

function apply_legendre(w::Tuple{S, T}, x) where {S, T}
    F1, F2 = w
    @tullio c[s, i, j, b] := F1[i, ii] * F2[j, jj] * x[s, ii, jj, b]
end

function apply_legendre(w::Tuple{S, T, V}, x) where {S, T, V}
    F1, F2, F3 = w
    @tullio c[s, i, j, k, b] :=
        F1[i, ii] * F2[j, jj] * F3[k, kk] * x[s, ii, jj, kk, b]
end

# Base extensions
Base.ndims(tr::LegendreTransform) = length(tr.modes)
Base.eltype(::LegendreTransform) = Float32
Base.size(tr::LegendreTransform) = tr.modes

function Base.show(io::IO, tr::LegendreTransform)
    print(io, "LegendreTransform(modes = $(tr.modes))")
end


# Hardy-Weinberg: A1A1: (1-p2)^2, A1A2: 2p2(1-p2), A2A2: p2^2
# ADDITIVE_MODEL=1 (A1A1: 0, A1A2: 1, A2A2: 2)
@inline function convert(::Type{T}, a2freq::AbstractFloat, ::Val{1}) where T <: AbstractFloat
    T(2a2freq)
end
# DOMINANT_MODEL=2 (A1A1: 0, A1A2: 1, A2A2: 1)
@inline function convert(::Type{T}, a2freq::AbstractFloat, ::Val{2}) where T <: AbstractFloat
    T(a2freq * (2 - a2freq))
end
# RECESSIVE_MODEL=3 (A1A1: 0, A1A2: 0, A2A2: 1)
@inline function convert(::Type{T}, a2freq::AbstractFloat, ::Val{3}) where T <: AbstractFloat
    T(abs2(a2freq))
end

"""
    Base.copyto!(v, s, P, Q; model=ADDITIVE_MODEL, center=false, scale=false)

Copy SnpArray `s` to a numeric vector or matrix `v`. Impute missing genotypes
according to ADMIXTURE estimates `P` (population A2 allele frequencies of SNPs)
and `Q` (population fractions of individuals). Columns of `P` should match columns of
`s`. Columns of `Q` should match rows of `s`. Note if the inferred minor allele
frequency from `P` and `Q` is ≤0.01, that genotype will be imputed according to
the inferred minor allele frequency.

# Positional arguments
- `v`: output vector or array.
- `s`: SnpArray.
- `P`: `K x S` array of A2 popluation allele frequencies.
- `Q`: `K x N` array of individual population fractions.

# Keyword arguments
- `model::Union{Val{1}, Val{2}, Val{3}}=ADDITIVE_MODEL`: `ADDITIVE_MODEL` (default), `DOMINANT_MODEL`, or `RECESSIVE_MODEL`.
- `center::Bool=false`: center each genotype dosage by its mean.
- `scale::Bool=false`: scale each genotype dosage by its standard deviation.
"""
function Base.copyto!(
    v      :: AbstractVecOrMat{T1},
    s      :: AbstractSnpArray,
    P      :: AbstractMatrix{T2},
    Q      :: AbstractMatrix{T2};
    model  :: Union{Val{1}, Val{2}, Val{3}} = ADDITIVE_MODEL,
    center :: Bool = false,
    scale  :: Bool = false
    ) where {T1 <: AbstractFloat, T2 <: AbstractFloat}
    m, n, K = size(s, 1), size(s, 2), size(P, 1)
    # center, scale, impute
    @inbounds for j in 1:n, i in 1:m
        vij = SnpArrays.convert(T1, s[i, j], model)
        a2freq = T2(0)
        for k in 1:K
            a2freq += Q[k, i] * P[k, j]
        end
        μij = SnpArrays.convert(T1, a2freq, model)
        σij = model == ADDITIVE_MODEL ? sqrt(μij * (1 - μij / 2)) : sqrt(μij * (1 - μij))
        v[i, j] = (isnan(vij) || a2freq ≤ 0.01 || a2freq ≥ 0.99) ? μij : vij
        center && (v[i, j] -= μij)
        scale  && (v[i, j] /= σij)
    end
    v
end

function Base.convert(
    ::Type{T},
    s::AbstractSnpArray,
    P::AbstractMatrix,
    Q::AbstractMatrix;
    kwargs...) where T <: Array
    T(s, P, Q; kwargs...)
end
Array{T,N}(s::AbstractSnpArray, P::AbstractMatrix, Q::AbstractMatrix; kwargs...) where {T,N} =
copyto!(Array{T,N}(undef, size(s)), s, P, Q; kwargs...)

function grm_admixture(
    s::AbstractSnpArray,
    P::AbstractMatrix,
    Q::AbstractMatrix,
    ::Type{T} = Float64
    ) where T <: AbstractFloat
    m, n = size(s)
    # convert genotype
    tic = time()
    G = Mmap.mmap(Matrix{T}, m, n) # slightly faster than G = Matrix{T}(undef, m, n)
    Base.copyto!(G, s, P, Q, model=ADDITIVE_MODEL, center=true, scale=true)
    @printf("convert genotype: %.2f seconds\n", time() - tic)
    # GG'
    tic = time()
    Φ = G * transpose(G) # m x m
    @printf("Φ = GG': %.2f seconds\n", time() - tic)
    # convert G to {0,1} matrix
    tic = time()
    @inbounds for (idx, g) in enumerate(G)
        G[idx] = ifelse(iszero(g), T(0), T(1))
        # iszero(g) || (G[idx] = T(1))
    end
    @printf("convert G to {0,1} matrix: %.2f seconds\n", time() - tic)
    # Sij = # SNPs observed in both individuals i and j
    tic = time()
    S = G * transpose(G)
    @printf("S = GG': %.2f seconds\n", time() - tic)
    # Φ /= 2S
    @inbounds for j in 1:m, i in 1:j
        Φ[i, j] /= 2S[i, j]
    end
    copytri!(Φ, 'U')
end # function grm_admixture


function DE(x, p, t, derivatives, scalings, constants, BCs)
    ν₀, ν₋, ΔRi, Riᶜ, Pr = p

    Nz, H, τ, f, g, α  = constants.Nz, constants.H, constants.τ, constants.f, constants.g, constants.α
    σ_uw, σ_vw, σ_wT = scalings.uw.σ, scalings.vw.σ, scalings.wT.σ
    μ_u, μ_v, σ_u, σ_v, σ_T = scalings.u.μ, scalings.v.μ, scalings.u.σ, scalings.v.σ, scalings.T.σ

    D_cell = derivatives.cell
    D_face = derivatives.face

    u = @view x[1:Nz]
    v = @view x[Nz + 1:2Nz]
    T = @view x[2Nz + 1:3Nz]

    ϵ = 1f-7
    ∂u∂z = D_face * u
    ∂v∂z = D_face * v
    ∂T∂z = D_face * T
    Ri = local_richardson.(∂u∂z .+ ϵ, ∂v∂z .+ ϵ, ∂T∂z .+ ϵ, H, g, α, σ_u, σ_v, σ_T)

    ν = ν₀ .+ ν₋ .* tanh_step.((Ri .- Riᶜ) ./ ΔRi)
    # ν = zeros(Float32, 33)

    ν∂u∂z = [-(BCs.uw.bottom - scalings.uw(0f0)); σ_u / σ_uw / H .* ν[2:end-1] .* ∂u∂z[2:end-1]; -(BCs.uw.top - scalings.uw(0f0))]
    ν∂v∂z = [-(BCs.vw.bottom - scalings.vw(0f0)); σ_v / σ_vw / H .* ν[2:end-1] .* ∂v∂z[2:end-1]; -(BCs.vw.top - scalings.vw(0f0))]
    ν∂T∂z = [-(BCs.wT.bottom - scalings.wT(0f0)); σ_T / σ_wT / H .* ν[2:end-1] ./ Pr .* ∂T∂z[2:end-1]; -(BCs.wT.top - scalings.wT(0f0))]

    ∂u∂t = τ / H * σ_uw / σ_u .* derivatives.cell * ν∂u∂z .+ f * τ / σ_u .* (σ_v .* v .+ μ_v)
    ∂v∂t = τ / H * σ_vw / σ_v .* derivatives.cell * ν∂v∂z .- f * τ / σ_v .* (σ_u .* u .+ μ_u)
    ∂T∂t = τ / H * σ_wT / σ_T .* derivatives.cell * ν∂T∂z

    return [∂u∂t; ∂v∂t; ∂T∂t]
end

function optimise_modified_pacanowski_philander(train_files, tsteps, timestepper, optimizers, maxiters, FILE_PATH;
                                                n_simulations, ν₀ = 1f-4, ν₋ = 1f-1, ΔRi=0.1f0, Riᶜ=0.25f0, Pr=1f0,
                                                train_gradient=true, gradient_scaling=5f-3,
                                                training_fractions=nothing)
    @info "Loading training data"
    𝒟 = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

    @info "Preparing constants and parameters"

    ν₀_scaling = 1 / ν₀
    ν₋_scaling = 1 / ν₋
    ΔRi_scaling = 1 / ΔRi
    Riᶜ_scaling = 1 / Riᶜ
    Pr_scaling = 1 / Pr

    function scale_parameter(parameter, scaling)
        return parameter * scaling
    end

    function unscale_parameter(parameter, scaling)
        return parameter / scaling
    end

    function prepare_parameters()
        Nz = length(𝒟.u.z)
        H = abs(𝒟.uw.z[end] - 𝒟.uw.z[1])
        τ = abs(𝒟.t[:,1][end] - 𝒟.t[:,1][1])
        u_scaling = 𝒟.scalings["u"]
        v_scaling = 𝒟.scalings["v"]
        T_scaling = 𝒟.scalings["T"]
        uw_scaling = 𝒟.scalings["uw"]
        vw_scaling = 𝒟.scalings["vw"]
        wT_scaling = 𝒟.scalings["wT"]

        constants = (H=H, τ=τ, Nz=Nz, f=1f-4, α=2f-4, g=9.80665f0)
        scalings = (u=u_scaling, v=v_scaling, T=T_scaling, uw=uw_scaling, vw=vw_scaling, wT=wT_scaling, parameters=[ν₀_scaling, ν₋_scaling, ΔRi_scaling, Riᶜ_scaling, Pr_scaling])
        derivatives = (cell=Float32.(Dᶜ(Nz, 1 / Nz)), face=Float32.(Dᶠ(Nz, 1 / Nz)))
        parameters = [ν₀, ν₋, ΔRi, Riᶜ, Pr]
        scaled_parameters = scale_parameter.(parameters, scalings.parameters)

        return constants, scalings, derivatives, scaled_parameters
    end

    constants, scalings, derivatives, scaled_parameters = prepare_parameters()

    Nz = constants.Nz

    D_face = derivatives.face

    n_steps = Int(length(@view(𝒟.t[:,1])) / n_simulations)

    uvT₀s = [𝒟.uvT_scaled[:,n_steps * i + tsteps[1]] for i in 0:n_simulations - 1]
    t_train = 𝒟.t[:,1][tsteps] ./ constants.τ
    tspan_train = (t_train[1], t_train[end])
    uvT_trains = [𝒟.uvT_scaled[:,n_steps * i + 1:n_steps * (i + 1)][:, tsteps] for i in 0:n_simulations - 1]

    u_trains, v_trains, T_trains = split_u.(uvT_trains, Nz), split_v.(uvT_trains, Nz), split_T.(uvT_trains, Nz)

    if train_gradient
        u_trains_gradients = [∂_∂z(sol, D_face) for sol in u_trains]
        v_trains_gradients = [∂_∂z(sol, D_face) for sol in v_trains]
        T_trains_gradients = [∂_∂z(sol, D_face) for sol in T_trains]
    end

    @info "Setting up BCs"

    @inline function BC(i)
        index = n_steps * i + tsteps[1]
        return (uw=(bottom=𝒟.uw.scaled[1,index], top=𝒟.uw.scaled[end, index]),
                vw=(bottom=𝒟.vw.scaled[1,index], top=𝒟.vw.scaled[end, index]),
                wT=(bottom=𝒟.wT.scaled[1,index], top=𝒟.wT.scaled[end, index]))
    end

    BCs = [BC(i) for i in 0:n_simulations - 1]

    @info "Setting up differential equations"

    prob_NDEs = [ODEProblem((x, p, t) -> DE(x, p, t, derivatives, scalings, constants, BCs[i]), uvT₀s[i], tspan_train) for i in 1:n_simulations]

    function determine_loss_scalings()
        if training_fractions === nothing
            loss_scalings = (u=1, v=1, T=1, ∂u∂z=gradient_scaling, ∂v∂z=gradient_scaling, ∂T∂z=gradient_scaling)
        else
            unscaled_parameters = unscale_parameter.(scaled_parameters, scalings.parameters)
            sols = [Array(solve(prob_NDEs[i], timestepper, p=unscaled_parameters, reltol=1f-3, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=t_train)) for i in 1:n_simulations]
            u_sols, v_sols, T_sols = split_u.(sols, Nz), split_v.(sols, Nz), split_T.(sols, Nz)

            u_loss = mean(loss.(u_trains, u_sols))
            v_loss = mean(loss.(v_trains, v_sols))
            T_loss = mean(loss.(T_trains, T_sols))

            if train_gradient
                u_sols_gradients = [∂_∂z(sol, D_face) for sol in u_sols]
                v_sols_gradients = [∂_∂z(sol, D_face) for sol in v_sols]
                T_sols_gradients = [∂_∂z(sol, D_face) for sol in T_sols]

                ∂u∂z_loss = mean(loss.(u_trains_gradients, u_sols_gradients))
                ∂v∂z_loss = mean(loss.(v_trains_gradients, v_sols_gradients))
                ∂T∂z_loss = mean(loss.(T_trains_gradients, T_sols_gradients))
            else
                ∂u∂z_loss = 0
                ∂v∂z_loss = 0
                ∂T∂z_loss = 0
            end

            losses = (u=u_loss, v=v_loss, T=T_loss, ∂u∂z=∂u∂z_loss, ∂v∂z=∂v∂z_loss, ∂T∂z=∂T∂z_loss)
            loss_scalings = calculate_loss_scalings(losses, training_fractions, train_gradient)
        end
        return loss_scalings
    end

    @info "Determining training scalings"

    loss_scalings = determine_loss_scalings()

    function loss_mpp(parameters, p)
        unscaled_parameters = unscale_parameter.(parameters, scalings.parameters)
        sols = [Array(solve(prob_NDEs[i], timestepper, p=unscaled_parameters, reltol=1f-3, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=t_train)) for i in 1:n_simulations]

        u_sols, v_sols, T_sols = split_u.(sols, Nz), split_v.(sols, Nz), split_T.(sols, Nz)
        u_loss = mean(loss.(u_trains, u_sols))
        v_loss = mean(loss.(v_trains, v_sols))
        T_loss = mean(loss.(T_trains, T_sols))

        losses = (u=u_loss, v=v_loss, T=T_loss, ∂u∂z=0, ∂v∂z=0, ∂T∂z=0)
        scaled_losses = apply_loss_scalings(losses, loss_scalings)

        return sum(scaled_losses), scaled_losses
    end

    function loss_gradient_mpp(parameters, p)
        unscaled_parameters = unscale_parameter.(parameters, scalings.parameters)
        sols = [Array(solve(prob_NDEs[i], timestepper, p=unscaled_parameters, reltol=1f-3, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=t_train)) for i in 1:n_simulations]

        u_sols, v_sols, T_sols = split_u.(sols, Nz), split_v.(sols, Nz), split_T.(sols, Nz)

        u_sols_gradients = [∂_∂z(sol, D_face) for sol in u_sols]
        v_sols_gradients = [∂_∂z(sol, D_face) for sol in v_sols]
        T_sols_gradients = [∂_∂z(sol, D_face) for sol in T_sols]

        u_loss = mean(loss.(u_trains, u_sols))
        v_loss = mean(loss.(v_trains, v_sols))
        T_loss = mean(loss.(T_trains, T_sols))

        ∂u∂z_loss = mean(loss.(u_trains_gradients, u_sols_gradients))
        ∂v∂z_loss = mean(loss.(v_trains_gradients, v_sols_gradients))
        ∂T∂z_loss = mean(loss.(T_trains_gradients, T_sols_gradients))

        losses = (u=u_loss, v=v_loss, T=T_loss, ∂u∂z=∂u∂z_loss, ∂v∂z=∂v∂z_loss, ∂T∂z=∂T∂z_loss)
        scaled_losses = apply_loss_scalings(losses, loss_scalings)

        return sum(scaled_losses), scaled_losses
    end

    @info "Setting up optimization objective"

    if train_gradient
        f_loss = OptimizationFunction(loss_gradient_mpp, GalacticOptim.AutoZygote())
    else
        f_loss = OptimizationFunction(loss_mpp, GalacticOptim.AutoZygote())
    end

    prob_loss = OptimizationProblem(f_loss, scaled_parameters, lb=[0f0, 0f0, 0f0, 0f0, 0f0], ub=[10f0, 10f0, 10f0, 10f0, 10f0])

    @info "Writing metadata"

    write_metadata_modified_pacanowski_philander_optimisation(FILE_PATH, train_files, maxiters, tsteps, unscale_parameter.(scaled_parameters, scalings.parameters), loss_scalings, optimizers)

    @info "Beginning Training"

    res = 0
    for i in 1:length(optimizers)
        iter = 1
        opt = optimizers[i]
        function cb(args...)
            if iter <= maxiters
                parameters = args[1]
                unscaled_parameters = unscale_parameter.(parameters, scalings.parameters)
                ν₀, ν₋, ΔRi, Riᶜ, Pr = unscaled_parameters
                losses = args[3]
                profile_loss = losses.u + losses.v + losses.T
                gradient_loss = losses.∂u∂z + losses.∂v∂z + losses.∂T∂z
                total_loss = profile_loss + gradient_loss
                @info "ν₀ = $ν₀, ν₋ = $ν₋, ΔRi = $ΔRi, Riᶜ = $Riᶜ, Pr = $Pr, loss = $(total_loss): profile = $profile_loss, gradient = $gradient_loss (u = $(losses.u), v = $(losses.v), T = $(losses.T)) , ∂u∂z = $(losses.∂u∂z), ∂v∂z = $(losses.∂v∂z), ∂T∂z = $(losses.∂T∂z)), optimizer $i/$(length(optimizers)), iteration = $iter/$maxiters"
                write_data_modified_pacanowski_philander_optimisation(FILE_PATH, losses, unscaled_parameters)
            end
            iter += 1
            false
        end

        res = solve(prob_loss, opt, cb=cb, maxiters=maxiters)
        scaled_parameters .= res.minimizer
    end

    ν₀, ν₋, ΔRi, Riᶜ = unscale_parameter.(scaled_parameters, scalings.parameters)
    @info "ν₀ = $ν₀, ν₋ = $ν₋, ΔRi = $ΔRi, Riᶜ = $Riᶜ, Pr = $Pr, loss = $loss"
end


# Update_mixcomps.jl

export rfullcond_fstar!;

function update_μ_σ2!(mixcomp::MixComponentNormal, suffstat::Dict{Symbol, T},
    prior::PriorMixcomponent_Normal) where T <: Real

    ## mean
    v1 = 1.0 / ( 1.0 / prior.μ.σ2 + suffstat[:sumWinv] / mixcomp.σ2 )
    sd1 = sqrt(v1)
    μ1 = v1 * ( prior.μ.μ / prior.μ.σ2 + suffstat[:sumWinv] * suffstat[:yhat] / mixcomp.σ2 )

    mixcomp.μ = rand(mixcomp.rng, Distributions.Normal(μ1, sd1) )

    ## variance
    a1 = 0.5 * ( prior.σ2.ν + length(mixcomp.ζon_indx) )
    ss = suffstat[:sumWinv]*(suffstat[:yhat] - mixcomp.μ)^2 + suffstat[:s2]
    b1 = 0.5 * ( prior.σ2.ν * prior.σ2.s0 + ss )

    mixcomp.σ2 = rand(mixcomp.rng, Distributions.InverseGamma(a1, b1) )

    return nothing
end

# function rpost_f!(y::Vector{T}, μ::T, σ2::T, κ::T, Cor::Symmetric{Float64, Matrix{Float64}},
#     rng::Union{MersenneTwister, Threefry4x}) where T <: Real
#     # conjugate update for MvNormal mean with prior mean μ, known general
#     #   correlation matrix Cor, signal-to-noise ratio κ, and indep. noisy obs with variance σ2.
#
#     Corinv = inv(Cor)
#     J = PDMat_adj((Corinv ./ κ + I) ./ σ2, 1.0e-3)
#
#     h = ( y .- μ ) ./ σ2
#
#     return rand(rng, Distributions.MvNormalCanon(h, J) )
# end
# function rpost_f!(y::Vector{T}, μ::T, σ2::T, κ::T, Cor::Symmetric{Float64, Matrix{Float64}},
#     rng::Union{MersenneTwister, Threefry4x}) where T <: Real
#     # conjugate update for MvNormal mean with prior mean μ, known general
#     #   correlation matrix Cor, signal-to-noise ratio κ, and indep. noisy obs with variance σ2.
#
#     J = inv(PDMat_adj(κ*σ2*Cor)) + I/σ2
#     h = ( y .- μ ) ./ σ2
#
#     return rand(rng, Distributions.MvNormalCanon(h, J) )
# end
function rpost_f!(y::Vector{T}, μ::T, σ2::T, κ::T, Cor::Symmetric{Float64, Matrix{Float64}},
    rng::Union{MersenneTwister, Threefry4x}) where T <: Real
    # conjugate update for MvNormal mean with prior mean μ, known general
    #   correlation matrix Cor, signal-to-noise ratio κ, and indep. noisy obs with variance σ2.

    # Rasmussen & Williams '06, p.46
    K = κ * Cor
    K1 = (K + I) \ K
    Σ0 = K * (I - K1)
    Σ1 = PDMat_adj(Symmetric( σ2*Σ0 ))
    μ1 = Σ1 * ( y .- μ ) ./ σ2

    return rand(rng, Distributions.MvNormal(μ1, Σ1) )
end

function rfullcond_fstar!(f_ell::T, Celel::T, Cstst::T, Cstel::T,
    rng::Union{MersenneTwister, Threefry4x}, tol=1.0e-10) where T <: Real

    ## here C is the **Covariance** matrix

    μ = Cstel * ( Celel \ f_ell )
    Σ = Cstst - Cstel * ( Celel \ Cstel )

    if Σ < 0.0 && abs(Σ) < tol
        Σ = abs(Σ)
    end

    return sqrt(Σ) * randn(rng) + μ
end
function rfullcond_fstar!(f_ell::T, Celel::T,
    Cstst::Matrix{T}, Cstel::Vector{T},
    rng::Union{MersenneTwister, Threefry4x}, tol=1.0e-10) where T <: Real

    ## here C is the **Covariance** matrix

    μ = Cstel .* f_ell ./ Celel
    Σ = PDMat_adj(Symmetric(Cstst - Cstel * Cstel' ./ Celel), tol)

    return rand(rng, Distributions.MvNormal( μ, Σ ) )
end
function rfullcond_fstar!(f_ell::Vector{T}, Celel::Matrix{T},
    Cstst::T, Cstel::Vector{T},
    rng::Union{MersenneTwister, Threefry4x}, tol=1.0e-10) where T <: Real

    ## here C is the **Covariance** matrix

    Celel = PDMat_adj(Celel, tol)

    μ = Cstel' * ( Celel \ f_ell )
    Σ = Cstst - Cstel' * ( Celel \ Cstel )

    if Σ < 0.0 && abs(Σ) < tol
        Σ = abs(Σ)
    end

    return sqrt(Σ) * randn(rng) + μ
end
# function rfullcond_fstar!(f_ell::Vector{T}, Celel::Matrix{T},
#     Cstst::Matrix{T}, Cstel::Matrix{T},
#     rng::Union{MersenneTwister, Threefry4x}) where T <: Real
#
#     ## here C is the **Covariance** matrix
#
#     μ = Cstel * ( Celel \ f_ell )
#     Σ = PDMat_adj(Symmetric(Cstst - Cstel * ( Celel \ Cstel' )))
#
#     return rand(rng, Distributions.MvNormal( μ, Σ ) )
# end
function rfullcond_fstar!(f_ell::Vector{T}, Celel::Matrix{T},
    Cstst::Matrix{T}, Cstel::Matrix{T},
    rng::Union{MersenneTwister, Threefry4x}, tol=1.0e-10) where T <: Real

    ## here C is the **Covariance** matrix

    Celel = PDMat_adj(Celel, tol)

    μ = Cstel * ( Celel \ f_ell )
    Σ = PDMat_adj(Symmetric(Cstst - X_invA_Xt(Celel, Cstel) ), tol)

    return rand(rng, Distributions.MvNormal( μ, Σ ) )
end

# function rfullcond_fstar!_canon(fx::Vector{T}, ## already indexed
#     Cov::PDMat, on_indx::Vector{Int}, off_indx::Vector{Int},
#     rng::Union{MersenneTwister, Threefry4x}) where T <: Real
#
#     ## assumes zero mean GP
#
#     Λ = inv(Cov)
#     J = PDMat_adj( Λ.mat[off_indx, off_indx] )
#
#     η = - vec( Λ.mat[off_indx, on_indx] * fx )
#
#     return rand(rng, Distributions.MvNormalCanon( η, J ) )
# end

### worker function that operates on its own
function update_mixcomp!(mixcomp::MixComponentNormal, prior::PriorMixcomponent_Normal,
    y::Vector{T}, update::Vector{Symbol}=[:μ, :σ2, :κ, :corParams, :fx], tol_posdef=1.0e-9) where T <: Real

    nζ = length(mixcomp.ζon_indx)

    if (:κ in update) || (:corParams in update)
        suffstat = update_Cov!(mixcomp, y, prior, SNR_Hyper_ScInvChiSq(), MaternHyper_ScInvChiSq())
    elseif nζ > 0
        W_now, sumWinv_now = getW(mixcomp.κ, mixcomp.D[mixcomp.ζon_indx, mixcomp.ζon_indx], mixcomp.corParams)
        yhat_now, s2_now = getSufficients_W(y[mixcomp.ζon_indx], W_now, sumWinv_now)
        suffstat = Dict(:sumWinv => sumWinv_now, :yhat => yhat_now, :s2 => s2_now)
    end

    if nζ > 0

        if (:μ in update) || (:σ2 in update)
            update_μ_σ2!(mixcomp, suffstat, prior)
        end

        if (:fx in update)
            mixcomp.fx[mixcomp.ζon_indx] = rpost_f!(y[mixcomp.ζon_indx], mixcomp.μ, mixcomp.σ2,
                mixcomp.κ, Symmetric(mixcomp.Cor[mixcomp.ζon_indx, mixcomp.ζon_indx]), mixcomp.rng)

            ζoff_indx = setdiff( 1:length(y) , mixcomp.ζon_indx )
            Cov = PDMat_adj(mixcomp.Cor .* (mixcomp.κ * mixcomp.σ2))

            # mixcomp.fx[ζoff_indx] = rfullcond_fstar_canon(mixcomp.fx[mixcomp.ζon_indx], Cov,
            #     mixcomp.ζon_indx, ζoff_indx, mixcomp.rng)

            mixcomp.fx[ζoff_indx] = rfullcond_fstar!(mixcomp.fx[mixcomp.ζon_indx],
                Cov.mat[mixcomp.ζon_indx, mixcomp.ζon_indx], Cov.mat[ζoff_indx, ζoff_indx],
                Cov.mat[ζoff_indx, mixcomp.ζon_indx], mixcomp.rng, tol_posdef)
        end

    else

        if (:μ in update) || (:σ2 in update)
            mixcomp.μ = rand(mixcomp.rng, Normal(prior.μ.μ, sqrt(prior.μ.σ2)))
            mixcomp.σ2 = rand(mixcomp.rng, InverseGamma(0.5*prior.σ2.ν, 0.5*prior.σ2.ν*prior.σ2.s0))
        end

        if (:fx in update)
            Cov = PDMat_adj(mixcomp.Cor .* (mixcomp.κ * mixcomp.σ2), tol_posdef)
            mixcomp.fx = rand(mixcomp.rng, Distributions.MvNormal( Cov.mat ) )
        end
    end

    return mixcomp
end

### master function that calls parallel updates
function update_mixcomps!(state::State_GPMTD, prior::Prior_GPMTD, y::Vector{T},
    update::Vector{Symbol}=[:μ, :σ2, :κ, :corParams, :fx];
    n_procs::Int=1, tol_posdef=1.0e-9) where T <: Real

    L = length(state.mixcomps)

    for ℓ = 1:L
        state.mixcomps[ℓ].adapt = state.adapt
    end

    if n_procs == 1
        for ℓ = 1:L
            state.mixcomps[ℓ] = update_mixcomp!(state.mixcomps[ℓ], prior.mixcomps[ℓ], y, update, tol_posdef)
        end
    elseif n_procs > 1
        state.mixcomps = pmap(update_mixcomp!, state.mixcomps, prior.mixcomps,
            fill(y, L), fill(update, L), fill(tol_posdef, L)) # pmap will update in place on each worker, but we also need to return the mixcomps
    end

    state.accpt = [ deepcopy(state.mixcomps[ℓ].accpt) for ℓ = 1:L ]
    state.adapt_iter = deepcopy(state.mixcomps[1].adapt_iter)

    return nothing
end


using LinearAlgebra
using SparseArrays
using NetworkLearning
using LightGraphs: Graph, add_edge!
using SimpleWeightedGraphs: SimpleWeightedGraph
using LearnBase: ObsDim, ObsDimension
using Statistics: mean
using Test
using Future: copy!
using DelimitedFiles: readdlm

# Test components
include("t_components.jl")
Test.@testset "Network Learning (various components)" begin
	t_components();
end

# Test observation-based learning
include("t_observation_networklearner.jl")
Test.@testset "Network Learning (observation-based)" begin
	t_observation_networklearner();
end

# Test entity-based learning
include("t_entity_networklearner.jl")
Test.@testset "Network Learning (entity-based)" begin
	t_entity_networklearner();
end


"""
    Dirichlet

The [Dirichlet distribution](http://en.wikipedia.org/wiki/Dirichlet_distribution) is often
used the conjugate prior for Categorical or Multinomial distributions.
The probability density function of a Dirichlet distribution with parameter
``\\alpha = (\\alpha_1, \\ldots, \\alpha_k)`` is:

```math
f(x; \\alpha) = \\frac{1}{B(\\alpha)} \\prod_{i=1}^k x_i^{\\alpha_i - 1}, \\quad \\text{ with }
B(\\alpha) = \\frac{\\prod_{i=1}^k \\Gamma(\\alpha_i)}{\\Gamma \\left( \\sum_{i=1}^k \\alpha_i \\right)},
\\quad x_1 + \\cdots + x_k = 1
```

```julia
# Let alpha be a vector
Dirichlet(alpha)         # Dirichlet distribution with parameter vector alpha

# Let a be a positive scalar
Dirichlet(k, a)          # Dirichlet distribution with parameter a * ones(k)
```
"""
struct Dirichlet{T<:Real} <: ContinuousMultivariateDistribution
    alpha::Vector{T}
    alpha0::T
    lmnB::T

    function Dirichlet{T}(alpha::Vector{T}) where T
        alpha0::T = zero(T)
        lmnB::T = zero(T)
        for i in 1:length(alpha)
            ai = alpha[i]
            ai > 0 || throw(ArgumentError("Dirichlet: alpha must be a positive vector."))
            alpha0 += ai
            lmnB += lgamma(ai)
        end
        lmnB -= lgamma(alpha0)
        new{T}(alpha, alpha0, lmnB)
    end

    function Dirichlet{T}(d::Integer, alpha::T) where T
        alpha0 = alpha * d
        new{T}(fill(alpha, d), alpha0, lgamma(alpha) * d - lgamma(alpha0))
    end
end

Dirichlet(alpha::Vector{T}) where {T<:Real} = Dirichlet{T}(alpha)
Dirichlet(d::Integer, alpha::T) where {T<:Real} = Dirichlet{T}(d, alpha)
Dirichlet(alpha::Vector{T}) where {T<:Integer} = Dirichlet{Float64}(convert(Vector{Float64},alpha))
Dirichlet(d::Integer, alpha::Integer) = Dirichlet{Float64}(d, Float64(alpha))

struct DirichletCanon
    alpha::Vector{Float64}
end

length(d::DirichletCanon) = length(d.alpha)

#### Conversions
convert(::Type{Dirichlet{Float64}}, cf::DirichletCanon) = Dirichlet(cf.alpha)
convert(::Type{Dirichlet{T}}, alpha::Vector{S}) where {T<:Real, S<:Real} = Dirichlet(convert(Vector{T}, alpha))
convert(::Type{Dirichlet{T}}, d::Dirichlet{S}) where {T<:Real, S<:Real} = Dirichlet(convert(Vector{T}, d.alpha))

Base.show(io::IO, d::Dirichlet) = show(io, d, (:alpha,))

# Properties

length(d::Dirichlet) = length(d.alpha)
mean(d::Dirichlet) = d.alpha .* inv(d.alpha0)
params(d::Dirichlet) = (d.alpha,)
@inline partype(d::Dirichlet{T}) where {T<:Real} = T

function var(d::Dirichlet)
    α = d.alpha
    α0 = d.alpha0
    c = 1.0 / (α0 * α0 * (α0 + 1.0))

    k = length(α)
    v = Vector{Float64}(undef, k)
    for i = 1:k
        @inbounds αi = α[i]
        @inbounds v[i] = αi * (α0 - αi) * c
    end
    return v
end

function cov(d::Dirichlet)
    α = d.alpha
    α0 = d.alpha0
    c = 1.0 / (α0 * α0 * (α0 + 1.0))

    k = length(α)
    C = Matrix{Float64}(undef, k, k)

    for j = 1:k
        αj = α[j]
        αjc = αj * c
        for i = 1:j-1
            @inbounds C[i,j] = - α[i] * αjc
        end
        @inbounds C[j,j] = αj * (α0 - αj) * c
    end

    for j = 1:k-1, i = j+1:k
        @inbounds C[i,j] = C[j,i]
    end
    return C
end

function entropy(d::Dirichlet)
    α = d.alpha
    α0 = d.alpha0
    k = length(α)

    en = d.lmnB + (α0 - k) * digamma(α0)
    for j in 1:k
        @inbounds αj = α[j]
        en -= (αj - 1.0) * digamma(αj)
    end
    return en
end

function dirichlet_mode!(r::Vector{T}, α::Vector{T}, α0::T) where T <: Real
    k = length(α)
    s = α0 - k
    for i = 1:k
        @inbounds αi = α[i]
        if αi <= one(T)
            error("Dirichlet has a mode only when alpha[i] > 1 for all i" )
        end
        @inbounds r[i] = (αi - one(T)) / s
    end
    return r
end

dirichlet_mode(α::Vector{T}, α0::T) where {T <: Real} = dirichlet_mode!(Vector{T}(undef, length(α)), α, α0)

mode(d::Dirichlet) = dirichlet_mode(d.alpha, d.alpha0)
mode(d::DirichletCanon) = dirichlet_mode(d.alpha, sum(d.alpha))

modes(d::Dirichlet) = [mode(d)]

# Evaluation

function insupport(d::Dirichlet, x::AbstractVector{T}) where T<:Real
    n = length(x)
    if length(d.alpha) != n
        return false
    end
    s = 0.0
    for i in 1:n
        xi = x[i]
        if xi < 0.0
            return false
        end
        s += xi
    end
    if abs(s - 1.0) > 1e-8
        return false
    end
    return true
end

function _logpdf(d::Dirichlet, x::AbstractVector{T}) where T<:Real
    a = d.alpha
    s = 0.
    for i in 1:length(a)
        @inbounds s += (a[i] - 1.0) * log(x[i])
    end
    return s - d.lmnB
end

# sampling

function _rand!(d::Union{Dirichlet,DirichletCanon}, x::AbstractVector{T}) where T<:Real
    s = 0.0
    n = length(x)
    α = d.alpha
    for i in 1:n
        @inbounds s += (x[i] = rand(Gamma(α[i])))
    end
    multiply!(x, inv(s)) # this returns x
end

#######################################
#
#  Estimation
#
#######################################

struct DirichletStats <: SufficientStats
    slogp::Vector{Float64}   # (weighted) sum of log(p)
    tw::Float64              # total sample weights

    DirichletStats(slogp::Vector{Float64}, tw::Real) = new(slogp, Float64(tw))
end

length(ss::DirichletStats) = length(s.slogp)

mean_logp(ss::DirichletStats) = ss.slogp * inv(ss.tw)

function suffstats(::Type{Dirichlet}, P::AbstractMatrix{Float64})
    K = size(P, 1)
    n = size(P, 2)
    slogp = zeros(K)
    for i = 1:n
        for k = 1:K
            @inbounds slogp[k] += log(P[k,i])
        end
    end
    DirichletStats(slogp, n)
end

function suffstats(::Type{Dirichlet}, P::AbstractMatrix{Float64}, w::AbstractArray{Float64})
    K = size(P, 1)
    n = size(P, 2)
    if length(w) != n
        throw(DimensionMismatch("Inconsistent argument dimensions."))
    end

    tw = 0.
    slogp = zeros(K)

    for i = 1:n
        @inbounds wi = w[i]
        tw += wi
        for k = 1:K
            @inbounds slogp[k] += log(P[k,i]) * wi
        end
    end
    DirichletStats(slogp, tw)
end

# fit_mle methods

## Initialization

function _dirichlet_mle_init2(μ::Vector{Float64}, γ::Vector{Float64})
    K = length(μ)

    α0 = 0.
    for k = 1:K
        @inbounds μk = μ[k]
        @inbounds γk = γ[k]
        ak = (μk - γk) / (γk - μk * μk)
        α0 += ak
    end
    α0 /= K

    multiply!(μ, α0)
end

function dirichlet_mle_init(P::AbstractMatrix{Float64})
    K = size(P, 1)
    n = size(P, 2)

    μ = Vector{Float64}(undef, K)  # E[p]
    γ = Vector{Float64}(undef, K)  # E[p^2]

    for i = 1:n
        for k = 1:K
            @inbounds pk = P[k, i]
            @inbounds μ[k] += pk
            @inbounds γ[k] += pk * pk
        end
    end

    c = 1.0 / n
    for k = 1:K
        μ[k] *= c
        γ[k] *= c
    end

    _dirichlet_mle_init2(μ, γ)
end

function dirichlet_mle_init(P::AbstractMatrix{Float64}, w::AbstractArray{Float64})
    K = size(P, 1)
    n = size(P, 2)

    μ = Vector{Float64}(undef, K)  # E[p]
    γ = Vector{Float64}(undef, K)  # E[p^2]
    tw = 0.

    for i = 1:n
        @inbounds wi = w[i]
        tw += wi
        for k = 1:K
            pk = P[k, i]
            @inbounds μ[k] += pk * wi
            @inbounds γ[k] += pk * pk * wi
        end
    end

    c = 1.0 / tw
    for k = 1:K
        μ[k] *= c
        γ[k] *= c
    end

    _dirichlet_mle_init2(μ, γ)
end

## Newton-Ralphson algorithm

function fit_dirichlet!(elogp::Vector{Float64}, α::Vector{Float64};
    maxiter::Int=25, tol::Float64=1.0e-12, debug::Bool=false)
    # This function directly overrides α

    K = length(elogp)
    length(α) == K || throw(DimensionMismatch("Inconsistent argument dimensions."))

    g = Vector{Float64}(undef, K)
    iq = Vector{Float64}(undef, K)
    α0 = sum(α)

    if debug
        objv = dot(α - 1.0, elogp) + lgamma(α0) - sum(lgamma(α))
    end

    t = 0
    converged = false
    while !converged && t < maxiter
        t += 1

        # compute gradient & Hessian
        # (b is computed as well)

        digam_α0 = digamma(α0)
        iz = 1.0 / trigamma(α0)
        gnorm = 0.
        b = 0.
        iqs = 0.

        for k = 1:K
            @inbounds ak = α[k]
            @inbounds g[k] = gk = digam_α0 - digamma(ak) + elogp[k]
            @inbounds iq[k] = - 1.0 / trigamma(ak)

            @inbounds b += gk * iq[k]
            @inbounds iqs += iq[k]

            agk = abs(gk)
            if agk > gnorm
                gnorm = agk
            end
        end
        b /= (iz + iqs)

        # update α

        for k = 1:K
            @inbounds α[k] -= (g[k] - b) * iq[k]
            @inbounds if α[k] < 1.0e-12
                α[k] = 1.0e-12
            end
        end
        α0 = sum(α)

        if debug
            prev_objv = objv
            objv = dot(α - 1.0, elogp) + lgamma(α0) - sum(lgamma(α))
            @printf("Iter %4d: objv = %.4e  ch = %.3e  gnorm = %.3e\n",
                t, objv, objv - prev_objv, gnorm)
        end

        # determine convergence

        converged = gnorm < tol
    end

    Dirichlet(α)
end

function fit_mle(::Type{Dirichlet}, P::AbstractMatrix{Float64};
    init::Vector{Float64}=Float64[], maxiter::Int=25, tol::Float64=1.0e-12)

    α = isempty(init) ? dirichlet_mle_init(P) : init
    elogp = mean_logp(suffstats(Dirichlet, P))
    fit_dirichlet!(elogp, α; maxiter=maxiter, tol=tol)
end

function fit_mle(::Type{Dirichlet}, P::AbstractMatrix{Float64}, w::AbstractArray{Float64};
    init::Vector{Float64}=Float64[], maxiter::Int=25, tol::Float64=1.0e-12)

    n = size(P, 2)
    length(w) == n || throw(DimensionMismatch("Inconsistent argument dimensions."))

    α = isempty(init) ? dirichlet_mle_init(P, w) : init
    elogp = mean_logp(suffstats(Dirichlet, P, w))
    fit_dirichlet!(elogp, α; maxiter=maxiter, tol=tol)
end


function broutine2x2!(S::AbstractMatrix{Complex{T}}, U::AbstractMatrix, locs::Int) where T
    step_1 = 1 << (locs - 1)
    step_2 = 1 << locs
    U11_re = real(U[1, 1]); U11_im = imag(U[1, 1])
    U12_re = real(U[1, 2]); U12_im = imag(U[1, 2])
    U21_re = real(U[2, 1]); U21_im = imag(U[2, 1])
    U22_re = real(U[2, 2]); U22_im = imag(U[2, 2])

    Hmax = size(S, 2)
    Sr = reinterpret(reshape, T, S)

    if step_1 == 1
        for _j in 0:(Hmax>>>locs)-1, b in axes(S, 1)
            j = _j << locs

            ST1_re = U11_re * Sr[1, b, j+1] - U11_im * Sr[2, b, j+1] +
                U12_re * Sr[1, b, j+2] - U12_im * Sr[2, b, j+2]
            ST1_im = U11_re * Sr[2, b, j+1] + U11_im * Sr[1, b, j+1] +
                U12_re * Sr[2, b, j+2] + U12_im * Sr[1, b, j+2]
            ST2_re = U21_re * Sr[1, b, j+1] - U21_im * Sr[2, b, j+1] +
                U22_re * Sr[1, b, j+2] - U22_im * Sr[2, b, j+2]
            ST2_im = U21_re * Sr[2, b, j+1] + U21_im * Sr[1, b, j+1] +
                U22_re * Sr[2, b, j+2] + U22_im * Sr[1, b, j+2]

            Sr[1, b, j+1] = ST1_re
            Sr[2, b, j+1] = ST1_im
            Sr[1, b, j+2] = ST2_re
            Sr[2, b, j+2] = ST2_im
        end
        return S
    end
end

function broutine2x2_avx!(S::AbstractMatrix{Complex{T}}, U::AbstractMatrix, locs::Int) where T
    step_1 = 1 << (locs - 1)
    step_2 = 1 << locs
    U11_re = real(U[1, 1]); U11_im = imag(U[1, 1])
    U12_re = real(U[1, 2]); U12_im = imag(U[1, 2])
    U21_re = real(U[2, 1]); U21_im = imag(U[2, 1])
    U22_re = real(U[2, 2]); U22_im = imag(U[2, 2])

    Hmax = size(S, 2)
    Sr = reinterpret(reshape, T, S)

    if step_1 == 1
        @turbo for _j in 0:(Hmax>>>locs)-1, b in axes(S, 1)
            j = _j << locs

            ST1_re = U11_re * Sr[1, b, j+1] - U11_im * Sr[2, b, j+1] +
                U12_re * Sr[1, b, j+2] - U12_im * Sr[2, b, j+2]
            ST1_im = U11_re * Sr[2, b, j+1] + U11_im * Sr[1, b, j+1] +
                U12_re * Sr[2, b, j+2] + U12_im * Sr[1, b, j+2]
            ST2_re = U21_re * Sr[1, b, j+1] - U21_im * Sr[2, b, j+1] +
                U22_re * Sr[1, b, j+2] - U22_im * Sr[2, b, j+2]
            ST2_im = U21_re * Sr[2, b, j+1] + U21_im * Sr[1, b, j+1] +
                U22_re * Sr[2, b, j+2] + U22_im * Sr[1, b, j+2]

            Sr[1, b, j+1] = ST1_re
            Sr[2, b, j+1] = ST1_im
            Sr[1, b, j+2] = ST2_re
            Sr[2, b, j+2] = ST2_im
        end
        return S
    end
end

@testset "Quantum" begin
  N = 10
  S = rand(ComplexF64, 100, 1 << N)
  U = rand(ComplexF64, 2, 2)

  S1 = broutine2x2!(copy(S), U, 1)
  S2 = broutine2x2_avx!(copy(S), U, 1)
  @test S1 ≈ S2 # 239
end


p = 2
n = 10000

function ftrefethen(p)
  x = p[1]; y = p[2]
  return exp(sin(50 * x)) + sin(60 * exp(y)) + sin(70 * sin(x)) +
             sin(sin(80 * y)) - sin(10 * (x + y)) + (x^2 + y^2)/4
end
f = ftrefethen

alpha = 0.602
gamma = 0.101

theta_lo = -100*ones(p,1)
theta_hi = 100*ones(p,1)
theta = 0.40 * (theta_hi - theta_lo) * rand(p, 1)
theta=min(theta,theta_hi)
theta=max(theta,theta_lo)

# Input to gain settings calculations
initial_step_length = 0.40 * maximum(theta_hi - theta_lo)
gavg = 2 # How many gradient estimates per iteration
n_for_gain_calc = n/20
estimated_std_dev_measurement_noise = 0.0

A = 0.10 * n / (2*gavg)
c = max(estimated_std_dev_measurement_noise/(gavg^0.5), 0.0001)

gbar=zeros(p, 1)
for i=1:int(n_for_gain_calc/(2*gavg))
   ghat=zeros(p,1)
   for j=1:gavg
      delta=2*round(rand(p,1))-1;
      thetaplus=theta+c*delta;
      thetaminus=theta-c*delta;
      yplus=f(thetaplus);
      yminus=f(thetaminus);
      ghat=(yplus-yminus)./(2*c*delta)+ghat;
   end
   gbar=gbar+abs(ghat/gavg);
end
meangbar=mean(gbar)/(n_for_gain_calc/(2*gavg))
a = initial_step_length * ((A+1)^alpha)/meangbar;

println("a = $(a), c = $(c), A = $(A)")

#a = 0.0017
#c = 0.00001 # No noise so we can select c small
#A = 5

min_f = Inf
best = 1

for k=0:n-1

    ak = a/(k+1+A)^alpha
    ck = c/(k+1)^gamma

    delta = 2*round(rand(p,1))-1

    thetaplus = theta + ck*delta
    thetaminus = theta - ck*delta

    yplus=f(thetaplus)
    yminus=f(thetaminus)

    ghat = (yplus - yminus)./(2*ck*delta)
    theta = theta - ak * ghat

    theta=min(theta,theta_hi)
    theta=max(theta,theta_lo)

    tf = f(theta)

    if tf < min_f
      min_f = tf
      best = theta
    end

    println("theta = ", theta, ", fitness = ", )

end
println("a = $(a), c = $(c), A = $(A)")
println("best = ", best, ", fitness = ", min_f)


using Distributed, SharedArrays, LinearAlgebra

using ForwardDiff
using Distributions

struct KSDTest <: KernelGoodnessOfFitTest
    kernel::Kernel
    m::Integer
    α::Float64
end

struct KSDResult <: KernelGoodnessOfFitResult
    stat::Float64

    "p-value of test"
    p_val
    "size of test"
    α # size of test
    "`:reject` or `:accept` to indicate rejection or acceptance of H₀"
    result::Symbol # :reject or :accept
end

u_p(k::Kernel, x::AbstractVector, y::AbstractVector, s_x::AbstractVector, s_y::AbstractVector) = begin
    kernel(k, x, y) * dot(s_x, s_y) +
        dot(s_x, k_dy(k, x, y)) +
        dot(k_dx(k, x, y), s_y) +
        tr(k_dxdy(k, x, y))
end

u_p(k::Kernel, p::Distribution, x::AbstractVector, y::AbstractVector) = begin
    s_x = gradlogpdf(p, x)
    s_y = gradlogpdf(p, y)

    u_p(k, x, y, s_x, s_y)
end

# compute unbiased estimate
Sᵤ_estimate(samples::AbstractArray, f) = begin
    n = size(samples, 2)
    res = 0.0

    for i = 1:n
        for j = 1:n
            if i == j
                continue
            end
            r = f(samples[:, i], samples[:, j])
            res += (r / (n * (n - 1)))
        end
    end

    res
end

"Performs a goodness-of-fit test based on the Kernelized Stein Discrepancy (KSD)."
perform(t::KSDTest, q, xs) = begin
    bootstrapped = SharedArray{Float64}(t.m);

    f = (x, y) -> u_p(t.kernel, q, x, y)

    @sync @distributed for b = 1:t.m
        bootstrapped[b] = bootstrap_degenerate(xs, f)

        if (b % 10000) == 0
            println("[$b / $m] Done!")
        end
    end

    # Unbiased estimate
    # println("Computing unbiased estimate")
    Ŝᵤ = Sᵤ_estimate(xs, f)

    # compute average times this occurs
    pval = reduce(+, bootstrapped .> Ŝᵤ) / t.m

    res = :accept
    if pval < t.α
        res = :reject;
    end

    # TODO: return some result containing information
    KSDResult(Ŝᵤ, pval, t.α, res)
end


## test_fem.jl : test GRF defined on FEM

@testset "unstructured grid        " begin

    ## Generate a set of random points ##
    pts = rand(512, 2)
    m = Matern(0.2,2.0)

    ## Cholesky ##
    grf = GaussianRandomField(CovarianceFunction(2,m),Cholesky(),pts)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,Matern)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{Cholesky})
    @test length(grf.pts) == 2
    @test size(grf.pts[1],1) == 2
    @test size(grf.pts[1],2) == size(pts,1)

    ## Spectral ##
    grf = GaussianRandomField(CovarianceFunction(2,m),Spectral(),pts)
    @test isa(grf,GaussianRandomField{Spectral})

    ## CirculantEmbedding ##
    @test_throws ArgumentError GaussianRandomField(CovarianceFunction(2,m),CirculantEmbedding(),pts)

    ## KarhunenLoeve ##
    grf = GaussianRandomField(CovarianceFunction(2,m),KarhunenLoeve(128),pts)
    @test isa(grf,GaussianRandomField{KarhunenLoeve{128}})

    ## Anisotropic random field ##
    a = AnisotropicExponential([1000 0; 0 1000])
    cov = CovarianceFunction(2,a)
    grf = GaussianRandomField(cov,KarhunenLoeve(500),pts,quad=GaussLegendre())
    @test isa(grf.cov.cov,AnisotropicExponential)

    ## GaussianRandomField with mean value ##
    grf = GaussianRandomField(1,CovarianceFunction(2,m),Cholesky(),pts)
    @test isa(grf,GaussianRandomField{Cholesky})
    @test length(grf.mean) == size(pts,1)
    @test all(grf.mean.==1)

    grf = GaussianRandomField(ones(size(pts,1)),CovarianceFunction(2,m),Spectral(),pts)
    @test isa(grf,GaussianRandomField{Spectral})
    @test length(grf.mean) == size(pts,1)
    @test all(grf.mean.==1)

    ## Test spherical domain ##
    pts = 4*rand(1024, 3) .- 2
    d = pts[:,1].*pts[:,1] + pts[:,2].*pts[:,2] + pts[:,3].*pts[:,3]
    pts = pts[d .< 4, :]
    cov = CovarianceFunction(3, Matern(1, 3/2))
    grf = GaussianRandomField(cov, KarhunenLoeve(128), pts)
    z = sample(grf)
    @test isa(grf,GaussianRandomField{KarhunenLoeve{128}})
    grf = GaussianRandomField(cov, Spectral(), pts)
    @test isa(grf,GaussianRandomField{Spectral})
    grf = GaussianRandomField(cov, Cholesky(), pts)
    @test isa(grf,GaussianRandomField{Cholesky})

end


using Random
using Test
using FFTW
using CUDA
using TransformPlayground

Random.seed!(0)

FT = Float64
ε = eps(FT)

for N in (4, 5)
    A_cpu = randn(FT, N, N)
    A_gpu = CuArray(A_cpu)

    B_cpu = FFTW.r2r(A_cpu, FFTW.REDFT10)
    B_gpu = TransformPlayground.dct_makhoul_2d(A_gpu)

    @test isapprox.(B_cpu, Array(B_gpu), atol=16ε) |> all

    C_cpu = FFTW.r2r(B_cpu, FFTW.REDFT01)
    C_cpu = C_cpu ./ (2N)^2

    C_gpu = TransformPlayground.idct_makhoul_2d(B_gpu)

    @test isapprox.(C_cpu, Array(C_gpu), atol=ε) |> any

    @test isapprox.(A_cpu, C_cpu, atol=4ε) |> all
    @test isapprox.(Array(A_gpu), Array(C_gpu), atol=4ε) |> all
end


#!/usr/bin/julia
# julia examples from LJR

f(x) = x^2 * cos(x)*tan(x)*besselj(x, 1.5)

@time println(f(2.5))
@time println(f(2.5))

A = rand(3, 3)
A * A
A .* A

A = rand(100, 10)
b = rand(100)
@time A\b
@time A\b

x = 10
println(2x)

x = y = z = 10
println( :(x=y=z=10) )
println( typeof( :(x=y=z=10) ))

println( :(x^2x) )
println( :(x^2*x) )
println( :(2x) )

println( macroexpand( :(@evalpoly x 2 3) ) )

println( @evalpoly x 1 2 3 )
println( 1 + 2x + 3x^2 )

function f(n)
   sum(0:3:n) + sum(0:5:n) - sum(0:15:n)
end

@time println( f(Int128(2)^64))
@time println( f(Int128(2)^64))

function nextfib(n)
   a, b = zero(n), one(n)
   while b < n
      a, b = b, a+b
   end
   return b
end

@time println( nextfib( Int128(2)^120 ))
@time println( nextfib( Int128(2)^120 ))

@code_llvm nextfib( Int128(2)^120 )

nheads = @parallel (+) for i=1:200000000
   Int(rand(Bool))
end
println(nheads)

x = 0
@simd for i=1:20000000
   x += Int(rand(Bool))
end
println(x)

a = SharedArray(Float64, 10)
@parallel for i=1:10
   a[i] = i
end
println(a)

#using Threads
#
#println(Threads.nthreads())
#Threads.@threads for i = 1:10
#   a[i] = Threads.threadid()
#end
#println(a)

#using ArrayFire
#a = rand(AFArray{Float32}, 100, 100)
#@time a^log(a)
#@time a^log(a)

x1 = rand(1000000);
x2 = rand(1000000);
@time dot(x1, x2);


using OctreeBH
using Test

include("test_ngbs.jl")
include("test_gravtree.jl")
include("test_balance.jl")


###########################
# An n-dimensional simplex
###########################

abstract type Simplex{T,C} <: Domain{T} end

isclosedset(::Simplex{T,:closed}) where {T} = true
isclosedset(::Simplex{T,:open}) where {T} = false

isopenset(d::Simplex) = !isclosedset(d)

abstract type UnitSimplex{T,C} <: Simplex{T,C} end
const ClosedUnitSimplex{T} = UnitSimplex{T,:closed}
const OpenUnitSimplex{T} = UnitSimplex{T,:open}

UnitSimplex(n::Int) = DynamicUnitSimplex(n)
UnitSimplex(::Val{N}) where {N} = EuclideanUnitSimplex{N}()

UnitSimplex{T}(n::Int) where {T <: StaticTypes} = StaticUnitSimplex{T}(n)
UnitSimplex{T}(::Val{N}) where {N,T} = StaticUnitSimplex{T}(Val(N))
UnitSimplex{T}() where {T <: StaticTypes} = StaticUnitSimplex{T}()
UnitSimplex{T}(n::Int) where {T} = DynamicUnitSimplex{T}(n)

UnitSimplex{T,C}(n::Int) where {T <: StaticTypes,C} = StaticUnitSimplex{T,C}(n)
UnitSimplex{T,C}(::Val{N}) where {N,T,C} = StaticUnitSimplex{T,C}(Val(N))
UnitSimplex{T,C}() where {T <: StaticTypes,C} = StaticUnitSimplex{T,C}()
UnitSimplex{T,C}(n::Int) where {T,C} = DynamicUnitSimplex{T,C}(n)

insimplex_closed(x) = mapreduce( t-> t >= 0, &, x) && norm(x,1) <= 1
insimplex_open(x) = mapreduce( t-> t > 0, &, x) && norm(x,1) < 1
insimplex_closed(x, tol) = mapreduce( t-> t >= -tol, &, x) && norm(x,1) <= 1+tol
insimplex_open(x, tol) = mapreduce( t-> t > -tol, &, x) && norm(x,1) < 1+tol

indomain(x, ::ClosedUnitSimplex) = insimplex_closed(x)
indomain(x, ::OpenUnitSimplex) = insimplex_open(x)
approx_indomain(x, ::ClosedUnitSimplex, tolerance) = insimplex_closed(x, tolerance)
approx_indomain(x, ::OpenUnitSimplex, tolerance) = insimplex_open(x, tolerance)

isempty(::UnitSimplex) = false

==(d1::UnitSimplex, d2::UnitSimplex) =
    isclosedset(d1)==isclosedset(d2) && dimension(d1)==dimension(d2)

boundingbox(d::UnitSimplex{T}) where {T} = UnitCube{T}(dimension(d))

struct StaticUnitSimplex{T,C} <: UnitSimplex{T,C}
end

StaticUnitSimplex(::Val{N}) where {N} = StaticUnitSimplex{SVector{N,Float64}}()

StaticUnitSimplex{T}() where {T} = StaticUnitSimplex{T,:closed}()

StaticUnitSimplex{T}(n::Int) where {T} =
    (@assert n == euclideandimension(T); StaticUnitSimplex{T}())
StaticUnitSimplex{T}(::Val{N}) where {N,T} =
    (@assert N == euclideandimension(T); StaticUnitSimplex{T}())

StaticUnitSimplex{T,C}(n::Int) where {T,C} =
    (@assert n == euclideandimension(T); StaticUnitSimplex{T,C}())
StaticUnitSimplex{T,C}(::Val{N}) where {N,T,C} =
    (@assert N == euclideandimension(T); StaticUnitSimplex{T,C}())

const EuclideanUnitSimplex{N,T,C} = StaticUnitSimplex{SVector{N,T},C}

EuclideanUnitSimplex{N}() where {N} = EuclideanUnitSimplex{N,Float64}()

struct DynamicUnitSimplex{T,C} <: UnitSimplex{T,C}
    dimension   ::  Int

    DynamicUnitSimplex{T,C}(n::Int) where {T,C} = new(n)
    DynamicUnitSimplex{T,C}(n::Int) where {T<:StaticTypes,C} =
        (@assert n == euclideandimension(T); new(n))
end

DynamicUnitSimplex(n::Int) = DynamicUnitSimplex{Vector{Float64}}(n)
DynamicUnitSimplex{T}(n::Int) where {T} = DynamicUnitSimplex{T,:closed}(n)

dimension(d::DynamicUnitSimplex) = d.dimension

const VectorUnitSimplex{T,C} = DynamicUnitSimplex{Vector{T},C}

VectorUnitSimplex(dimension) = VectorUnitSimplex{Float64}(dimension)

show(io::IO, d::EuclideanUnitSimplex{N,Float64,:closed}) where {N} = print(io, "UnitSimplex(Val($(N)))")
show(io::IO, d::VectorUnitSimplex{Float64,:closed}) = print(io, "UnitSimplex($(dimension(d)))")

center(d::EuclideanUnitSimplex{N,T}) where {N,T} = ones(SVector{N,T})/N
center(d::VectorUnitSimplex{T}) where {T} = ones(T, dimension(d))/dimension(d)

corners(::EuclideanUnitSimplex{N,T}) where {N,T} =
    vcat([zero(SVector{N,T})], [ SVector{N,T}(ntuple( i -> i==j, N)) for j in 1:N])
corners(d::VectorUnitSimplex{T}) where {T} =
    vcat([zeros(T,dimension(d))], [ (z = zeros(T, dimension(d)); z[j]=1; z) for j in 1:dimension(d)])

interior(d::EuclideanUnitSimplex{N,T}) where {N,T} = EuclideanUnitSimplex{N,T,:open}()
closure(d::EuclideanUnitSimplex{N,T}) where {N,T} = EuclideanUnitSimplex{N,T,:closed}()
interior(d::VectorUnitSimplex{T}) where {T} = VectorUnitSimplex{T,:open}(dimension(d))
closure(d::VectorUnitSimplex{T}) where {T} = VectorUnitSimplex{T,:closed}(dimension(d))

# We pick the center point, because it belongs to the domain regardless of
# whether it is open or closed.
point_in_domain(d::UnitSimplex) = center(d)

similardomain(d::StaticUnitSimplex{S,C}, ::Type{T}) where {S,T,C} =
    StaticUnitSimplex{T,C}()
similardomain(d::DynamicUnitSimplex{S,C}, ::Type{T}) where {S,T,C} =
    DynamicUnitSimplex{T,C}(d.dimension)


#=
datasets.jl:
- Julia version:
- Author: Jirik
- Date: 2019-03-03
=#

function random_image(shape, obj_min, obj_max, level)
    # print("Random image start")
    im = rand(shape...)

    for i=obj_min[1]:obj_max[1]
        for j=obj_min[2]:obj_max[2]
            for k=obj_min[3]:obj_max[3]
                im[i,j,k] = im[i,j,k] + level
            end
        end
    end
    return im
end

function generate_slope(data_size)
    data = zeros(Int8, data_size[1], data_size[2], data_size[3])
    for i=1:data_size[1]
        for j=1:data_size[2]
            for k=1:data_size[3]
                if ((j - k) < i) & ((j - k + 3) >= i)
                    data[i,j,k] = 1
                end
#                 if i > (j - k)
#                 if (k - j) < i
#                     data[i,j,k] = 1
#                 end
            end
        end
    end
    return data
end

function generate_truncated_sphere(r=10, data_size=nothing)
	if data_size == nothing
		data_size = [Int(r*2 + 3), Int(r*2 + 3), Int(r*2 + 3)]
	end
	center = data_size / 2
    data = zeros(Int8, data_size[1], data_size[2], data_size[3])
    for i=1:data_size[1]
        for j=1:data_size[2]
            for k=1:data_size[3]
				coords = [i, j, k]
                if ((j - k) < i) # & ((j - k + 3) >= i)
					dist = sum((center - coords).^2)^0.5
					if dist < r
                    	data[i,j,k] = 1
					end
                end
#                 if i > (j - k)
#                 if (k - j) < i
#                     data[i,j,k] = 1
#                 end
            end
        end
    end
    return data
end

function data234()
    @info "data234 init"
    data_size = [2,3,4]
    segmentation = zeros(Int8, 2, 3, 4)
    segmentation[1:2,2:3,3:4] .= 1
    return segmentation
end

function tetris_brick()
    @info "data tetris init"
    segmentation = zeros(Int8, 3, 4, 5)
    segmentation[2,2,2:4] .= 1
    segmentation[2,3,3] = 1
    return segmentation
end


using JLD
using CSV
using Printf
using Clustering


#=
Philosophy and convention:
A node is defined as the center of a grid cell. Grid cells may be those
in a primary grid or those in a dual grid.

The definitions "dual" and "primal" are made based on their conventional use
in holding fluid dynamic data. For example, pressure is held in primal nodes,
vorticity and streamfunction at dual nodes.

The definitions here are made with a "dual grid" in mind. That is, the grid is
defined by an integer number of dual cells in each direction. If a "primal grid"
is needed, then all of the defintions can be swapped (primal -> dual, dual -> primal).

Also, note that there might be dual cells that are "ghosts" (i.e. lie outside
the grid), but these are not distinguished in these basic definitions and operators.
=#

#=
Notes on data types:
All data types can be reduced to one of two different types in each direction.
(1) Lying on dual cell centers (C), numbered 1 through N
(2) Lying on dual cell edges (E), numbered 1 through N-1 (i.e. lying midway between C)

For example, Nodes{Dual} by definition lie on dual cell centers in both directions,
so they would be C x C. Nodes{Primal} are aligned with the corners of the dual
cells, so they lie along dual cell edges in both directions, E x E.

Nodes{Dual}    -- C x C
Nodes{Primal}  -- E x E
XEdges{Dual}   -- E x C
YEdges{Dual}   -- C x E
XEdges{Primal} -- C x E
YEdges{Primal} -- E x C

This is important when considering the interpolation and differentiation operations,
both of which must account for the relative indexing between E and C. The
key to remember is that E[1] is bounded by C[1] and C[2] in our indexing convention,
so that, for either interpolation or differentiation:
      E[i] <- C[i], C[i+1] for i in 1:N-1
and
      C[i] <- E[i-1], E[i] for i in 2:N-1
and
      C[i] <- C[i] for i in 1:N
and
      E[i] <- E[i] for i in 1:N-1
=#

module Fields

import Base: @propagate_inbounds, show, summary

#using Compat
using FFTW
using SpecialFunctions
using Statistics

using LinearAlgebra
using SparseArrays

import LinearAlgebra: mul!, ldiv!, cross, ×, dot, ⋅
import Base: parentindices
const GAMMA = MathConstants.γ

export Primal, Dual, ScalarGridData, VectorGridData, GridData,
       Edges, Nodes, XEdges, YEdges,
       EdgeGradient, NodePair,
       Points, ScalarData, VectorData, TensorData,
       celltype,
       diff!,interpolate!,
       curl, curl!, Curl, divergence, divergence!, Divergence,
       grad, grad!, Grad,
       laplacian, laplacian!, laplacian_symm!, plan_laplacian, plan_laplacian!,
       helmholtz, helmholtz!, plan_helmholtz, plan_helmholtz!,
       plan_intfact,plan_intfact!,Identity,
       product, product!, ∘,
       directional_derivative!, directional_derivative_conserve!, curl_cross!,
       convective_derivative!, convective_derivative_rot!,
       coordinates,
       DDF, GradDDF,
       Regularize, RegularizationMatrix, InterpolationMatrix,
       CircularConvolution

abstract type CellType end
abstract type Primal <: CellType end
abstract type Dual <: CellType end

abstract type GridData{NX,NY,T} <: AbstractMatrix{T} end

abstract type ScalarGridData{NX,NY,T} <: GridData{NX,NY,T} end

abstract type VectorGridData{NX,NY,T} <: GridData{NX,NY,T} end

macro wraparray(wrapper, field, N)
    S = eval(wrapper)
    @assert S <: AbstractArray "Wrapped type must be a subtype of AbstractArray"
    while supertype(S) <: AbstractArray
        S = supertype(S)
    end
    #T = supertype(eval(wrapper))
    #@assert T <: AbstractArray "Wrapped type must be a subtype of AbstractArray"
    #el_type, N = S.parameters

    quote
        Base.parent(A::$wrapper) = A.$field
        Base.size(A::$wrapper) = size(A.$field)
        parentindices(A::$wrapper) = parentindices(A.$field)

        if $N > 1
          function Base.show(io::IO, m::MIME"text/plain", A::$wrapper)
            println(io, "$(typeof(A)) data")
            println(io, "Printing in grid orientation (lower left is (1,1))")
            show(io,m, reverse(transpose(A.$field),dims=1))
          end
          #function Base.summary(io::IO, A::$wrapper)
          #  println(io, "$(typeof(A)) data")
          #  print(io, "Printing in grid orientation (lower left is (1,1))")
          #end
        end

        @propagate_inbounds Base.getindex(A::$wrapper, i::Int) = A.$field[i]
        @propagate_inbounds Base.setindex!(A::$wrapper, v, i::Int) = A.$field[i] = convert(eltype(A.$field), v)
        if $N > 1
          @propagate_inbounds Base.getindex(A::$wrapper, I::Vararg{Int, $N}) = A.$field[I...]
          @propagate_inbounds Base.setindex!(A::$wrapper, v, I::Vararg{Int, $N}) = A.$field[I...] = convert(eltype(A.$field), v)
        end
    end
end

celltype(w::GridData) = typeof(w).parameters[1]

function othertype end

macro othertype(celltype, k)
    esc(quote
        Fields.othertype(::$celltype) = $k
        Fields.othertype(::Type{$celltype}) = $k
    end)
end

@othertype Primal Dual
@othertype Dual Primal
@othertype CellType CellType

# This macro allows us to access scalar grid data via just the wrapper itself
@wraparray ScalarGridData data 2

include("fields/nodes.jl")
include("fields/edges.jl")
include("fields/collections.jl")

# (ctype,dnx,dny,shiftx,shifty)
scalarlist = ((:(Nodes{Primal,NX,NY,T}), 1,1,0.0,0.0),
              (:(Nodes{Dual,NX,NY,T}),   0,0,0.5,0.5),
              (:(XEdges{Primal,NX,NY,T}),0,1,0.5,0.0),
              (:(YEdges{Primal,NX,NY,T}),1,0,0.0,0.5),
              (:(XEdges{Dual,NX,NY,T}),  1,0,0.0,0.5),
              (:(YEdges{Dual,NX,NY,T}),  0,1,0.5,0.0))

# (ctype,dunx,duny,dvnx,dvny,shiftux,shiftuy,shiftvx,shiftvy)
vectorlist = ((:(Edges{Primal,NX,NY,T}),          0,1,1,0,0.5,0.0,0.0,0.5),
              (:(Edges{Dual,NX,NY,T}),            1,0,0,1,0.0,0.5,0.5,0.0),
              (:(NodePair{Dual,Dual,NX,NY,T}),    0,0,0,0,0.5,0.5,0.5,0.5),
              (:(NodePair{Primal,Dual,NX,NY,T}),  1,1,0,0,0.0,0.0,0.5,0.5),
              (:(NodePair{Dual,Primal,NX,NY,T}),  0,0,1,1,0.5,0.5,0.0,0.0),
              (:(NodePair{Primal,Primal,NX,NY,T}),1,1,1,1,0.0,0.0,0.0,0.0))

tensorlist = ((:(EdgeGradient{Dual,Primal,NX,NY,T}), 0,0,1,1,0.5,0.5,0.0,0.0),
              (:(EdgeGradient{Primal,Dual,NX,NY,T}), 1,1,0,0,0.0,0.0,0.5,0.5))

include("fields/basicoperations.jl")
include("fields/points.jl")

CollectedData = Union{EdgeGradient{R,S,NX,NY,T},NodePair{R,S,NX,NY,T}} where {R,S,NX,NY,T}

"""
    coordinates(w::GridData;[dx=1.0],[I0=(1,1)])

Return a tuple of the ranges of the physical coordinates in each direction for grid
data `w`. If `w` is of `Nodes` type, then it returns a tuple of the form
`xg,yg`. If `w` is of `Edges` or `NodePair` type, then it returns a tuple of
the form `xgu,ygu,xgv,ygv`.

The optional keyword argument `dx` sets the grid spacing; its default is `1.0`. The
optional keyword `I0` accepts a tuple of integers to set the index pair of the
primal nodes that coincide with the origin. The default is `(1,1)`.

# Example

```jldoctest
julia> w = Nodes(Dual,(12,22));

julia> xg, yg = coordinates(w,dx=0.1)
(-0.05:0.1:1.05, -0.05:0.1:2.0500000000000003)
```
"""
function coordinates end

for (ctype,dnx,dny,shiftx,shifty) in scalarlist
   @eval coordinates(w::$ctype;dx::Float64=1.0,I0::Tuple{Int,Int}=(1,1)) where {NX,NY,T} =
    dx.*((1-I0[1]-$shiftx):(NX-$dnx-I0[1]-$shiftx),
         (1-I0[2]-$shifty):(NY-$dny-I0[2]-$shifty))

end

for (ctype,dunx,duny,dvnx,dvny,shiftux,shiftuy,shiftvx,shiftvy) in vectorlist
   @eval coordinates(w::$ctype;dx::Float64=1.0,I0::Tuple{Int,Int}=(1,1)) where {NX,NY,T} =
    dx.*((1-I0[1]-$shiftux):(NX-$dunx-I0[1]-$shiftux),
         (1-I0[2]-$shiftuy):(NY-$duny-I0[2]-$shiftuy),
         (1-I0[1]-$shiftvx):(NX-$dvnx-I0[1]-$shiftvx),
         (1-I0[2]-$shiftvy):(NY-$dvny-I0[2]-$shiftvy))

end

include("fields/physicalgrid.jl")
include("fields/operators.jl")

end


using Distributions
using ForwardDiff
using Test

isnan_type(::Type{T}, v) where {T} = isnan(v) && v isa T

@testset "LogNormal" begin
    @test isa(convert(LogNormal{Float64}, Float16(0), Float16(1)),
              LogNormal{Float64})
    @test logpdf(LogNormal(0, 0), 1) === Inf
    @test logpdf(LogNormal(), Inf) === -Inf
    @test iszero(logcdf(LogNormal(0, 0), 1))
    @test iszero(logcdf(LogNormal(), Inf))
    @test logdiffcdf(LogNormal(), Float32(exp(5)), Float32(exp(3))) ≈ -6.6079385945968929 rtol=1e-12
    @test logdiffcdf(LogNormal(), Float64(exp(5)), Float64(exp(3))) ≈ -6.6079385945968929 rtol=1e-12
    let d = LogNormal(Float64(0), Float64(1)), x = Float64(exp(-60)), y = Float64(exp(-60.001))
        float_res = logdiffcdf(d, x, y)
        big_float_res = log(cdf(d, BigFloat(x, 100)) - cdf(d, BigFloat(y, 100)))
        @test float_res ≈ big_float_res
    end

    @test logccdf(LogNormal(0, 0), 1) === -Inf
    @test iszero(logccdf(LogNormal(eps(), 0), 1))

    @test iszero(quantile(LogNormal(), 0))
    @test isone(quantile(LogNormal(), 0.5))
    @test quantile(LogNormal(), 1) === Inf

    @test iszero(quantile(LogNormal(0, 0), 0))
    @test isone(quantile(LogNormal(0, 0), 0.75))
    @test quantile(LogNormal(0, 0), 1) === Inf

    @test iszero(quantile(LogNormal(0.25, 0), 0))
    @test quantile(LogNormal(0.25, 0), 0.95) == exp(0.25)
    @test quantile(LogNormal(0.25, 0), 1) === Inf

    @test cquantile(LogNormal(), 0) === Inf
    @test isone(cquantile(LogNormal(), 0.5))
    @test iszero(cquantile(LogNormal(), 1))

    @test cquantile(LogNormal(0, 0), 0) === Inf
    @test isone(cquantile(LogNormal(0, 0), 0.75))
    @test iszero(cquantile(LogNormal(0, 0), 1))

    @test cquantile(LogNormal(0.25, 0), 0) === Inf
    @test cquantile(LogNormal(0.25, 0), 0.95) == exp(0.25)
    @test iszero(cquantile(LogNormal(0.25, 0), 1))

    @test iszero(invlogcdf(LogNormal(), -Inf))
    @test isnan_type(Float64, invlogcdf(LogNormal(), NaN))

    @test invlogccdf(LogNormal(), -Inf) === Inf
    @test isnan_type(Float64, invlogccdf(LogNormal(), NaN))

    # test for #996 being fixed
    let d = LogNormal(0, 1), x = exp(1), ∂x = exp(2)
        @inferred cdf(d, ForwardDiff.Dual(x, ∂x)) ≈ ForwardDiff.Dual(cdf(d, x), ∂x * pdf(d, x))
    end
end

@testset "LogNormal type inference" begin
    # pdf
    @test @inferred(pdf(LogNormal(0.0, 0.0), 1.0)) === Inf
    @test @inferred(pdf(LogNormal(0.0, 0.0), 0.5)) === 0.0
    @test @inferred(pdf(LogNormal(0.0, 0.0), 0.0)) === 0.0
    @test @inferred(pdf(LogNormal(0.0, 0.0), -0.5)) === 0.0

    @test @inferred(pdf(LogNormal(0.0, 0.0), 1.0f0)) === Inf
    @test @inferred(pdf(LogNormal(0.0f0, 0.0f0), 1.0)) === Inf
    @test @inferred(pdf(LogNormal(0.0f0, 0.0f0), 1.0f0)) === Inf32

    @test isnan_type(Float64, @inferred(pdf(LogNormal(0.0, 0.0), NaN)))
    @test isnan_type(Float64, @inferred(pdf(LogNormal(NaN, 0.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(pdf(LogNormal(NaN, 0.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(pdf(LogNormal(NaN, 0.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(pdf(LogNormal(NaN32, 0.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(pdf(LogNormal(NaN32, 0.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(pdf(LogNormal(NaN32, 0.0f0), -1.0f0)))

    @test @inferred(pdf(LogNormal(0 // 1, 0 // 1), 1 // 1)) === Inf
    @test @inferred(pdf(LogNormal(0 // 1, 0 // 1), 0 // 1)) === 0.0
    @test @inferred(pdf(LogNormal(0 // 1, 0 // 1), -1 // 1)) === 0.0
    @test isnan_type(Float64, @inferred(pdf(LogNormal(0 // 1, 0 // 1), NaN)))

    @test @inferred(pdf(LogNormal(0.0, 0.0), BigInt(1))) == big(Inf)
    @test @inferred(pdf(LogNormal(0.0, 0.0), BigInt(0))) == big(0.0)
    @test @inferred(pdf(LogNormal(0.0, 0.0), BigInt(-1))) == big(0.0)
    @test @inferred(pdf(LogNormal(0.0, 0.0), BigFloat(1))) == big(Inf)
    @test @inferred(pdf(LogNormal(0.0, 0.0), BigFloat(0))) == big(0.0)
    @test @inferred(pdf(LogNormal(0.0, 0.0), BigFloat(-1))) == big(0.0)
    @test isnan_type(BigFloat, @inferred(pdf(LogNormal(0.0, 0.0), BigFloat(NaN))))

    # logpdf
    @test @inferred(logpdf(LogNormal(0.0, 0.0), 1.0)) === Inf
    @test @inferred(logpdf(LogNormal(0.0, 0.0), 0.5)) === -Inf
    @test @inferred(logpdf(LogNormal(0.0, 0.0), 0.0)) === -Inf
    @test @inferred(logpdf(LogNormal(0.0, 0.0), -0.5)) === -Inf

    @test @inferred(logpdf(LogNormal(0.0, 0.0), 1.0f0)) === Inf
    @test @inferred(logpdf(LogNormal(0.0f0, 0.0f0), 1.0)) === Inf
    @test @inferred(logpdf(LogNormal(0.0f0, 0.0f0), 1.0f0)) === Inf32

    @test isnan_type(Float64, @inferred(logpdf(LogNormal(0.0, 0.0), NaN)))
    @test isnan_type(Float64, @inferred(logpdf(LogNormal(NaN, 0.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(logpdf(LogNormal(NaN, 0.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(logpdf(LogNormal(NaN, 0.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(logpdf(LogNormal(NaN32, 0.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(logpdf(LogNormal(NaN32, 0.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(logpdf(LogNormal(NaN32, 0.0f0), -1.0f0)))

    @test @inferred(logpdf(LogNormal(0 // 1, 0 // 1), 1 // 1)) === Inf
    @test @inferred(logpdf(LogNormal(0 // 1, 0 // 1), 0 // 1)) === -Inf
    @test @inferred(logpdf(LogNormal(0 // 1, 0 // 1), -1 // 1)) === -Inf
    @test isnan_type(Float64, @inferred(logpdf(LogNormal(0 // 1, 0 // 1), NaN)))

    @test @inferred(logpdf(LogNormal(0.0, 0.0), BigInt(1))) == big(Inf)
    @test @inferred(logpdf(LogNormal(0.0, 0.0), BigInt(0))) == big(-Inf)
    @test @inferred(logpdf(LogNormal(0.0, 0.0), BigInt(-1))) == big(-Inf)
    @test @inferred(logpdf(LogNormal(0.0, 0.0), BigFloat(1))) == big(Inf)
    @test @inferred(logpdf(LogNormal(0.0, 0.0), BigFloat(0))) == big(-Inf)
    @test @inferred(logpdf(LogNormal(0.0, 0.0), BigFloat(-1))) == big(-Inf)
    @test isnan_type(BigFloat, @inferred(logpdf(LogNormal(0.0, 0.0), BigFloat(NaN))))

    # cdf
    @test @inferred(cdf(LogNormal(0.0, 0.0), 1.0)) === 1.0
    @test @inferred(cdf(LogNormal(0.0, 0.0), 0.5)) === 0.0
    @test @inferred(cdf(LogNormal(0.0, 0.0), 0.0)) === 0.0
    @test @inferred(cdf(LogNormal(0.0, 0.0), -0.5)) === 0.0

    @test @inferred(cdf(LogNormal(0.0, 0.0), 1.0f0)) === 1.0
    @test @inferred(cdf(LogNormal(0.0f0, 0.0f0), 1.0)) === 1.0
    @test @inferred(cdf(LogNormal(0.0f0, 0.0f0), 1.0f0)) === 1.0f0

    @test isnan_type(Float64, @inferred(cdf(LogNormal(0.0, 0.0), NaN)))
    @test isnan_type(Float64, @inferred(cdf(LogNormal(NaN, 0.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(cdf(LogNormal(NaN, 0.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(cdf(LogNormal(NaN, 0.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(cdf(LogNormal(NaN32, 0.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(cdf(LogNormal(NaN32, 0.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(cdf(LogNormal(NaN32, 0.0f0), -1.0f0)))

    @test @inferred(cdf(LogNormal(0 // 1, 0 // 1), 1 // 1)) === 1.0
    @test @inferred(cdf(LogNormal(0 // 1, 0 // 1), 0 // 1)) === 0.0
    @test @inferred(cdf(LogNormal(0 // 1, 0 // 1), -1 // 1)) === 0.0
    @test isnan_type(Float64, @inferred(cdf(LogNormal(0 // 1, 0 // 1), NaN)))

    @test @inferred(cdf(LogNormal(0.0, 0.0), BigInt(1))) == big(1.0)
    @test @inferred(cdf(LogNormal(0.0, 0.0), BigInt(0))) == big(0.0)
    @test @inferred(cdf(LogNormal(0.0, 0.0), BigInt(-1))) == big(0.0)
    @test @inferred(cdf(LogNormal(0.0, 0.0), BigFloat(1))) == big(1.0)
    @test @inferred(cdf(LogNormal(0.0, 0.0), BigFloat(0))) == big(0.0)
    @test @inferred(cdf(LogNormal(0.0, 0.0), BigFloat(-1))) == big(0.0)
    @test isnan_type(BigFloat, @inferred(cdf(LogNormal(0.0, 0.0), BigFloat(NaN))))

    # logcdf
    @test @inferred(logcdf(LogNormal(0.0, 0.0), 1.0)) === -0.0
    @test @inferred(logcdf(LogNormal(0.0, 0.0), 0.5)) === -Inf
    @test @inferred(logcdf(LogNormal(0.0, 0.0), 0.0)) === -Inf
    @test @inferred(logcdf(LogNormal(0.0, 0.0), -0.5)) === -Inf

    @test @inferred(logcdf(LogNormal(0.0, 0.0), 1.0f0)) === -0.0
    @test @inferred(logcdf(LogNormal(0.0f0, 0.0f0), 1.0)) === -0.0
    @test @inferred(logcdf(LogNormal(0.0f0, 0.0f0), 1.0f0)) === -0.0f0

    @test isnan_type(Float64, @inferred(logcdf(LogNormal(0.0, 0.0), NaN)))
    @test isnan_type(Float64, @inferred(logcdf(LogNormal(NaN, 0.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(logcdf(LogNormal(NaN, 0.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(logcdf(LogNormal(NaN, 0.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(logcdf(LogNormal(NaN32, 0.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(logcdf(LogNormal(NaN32, 0.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(logcdf(LogNormal(NaN32, 0.0f0), -1.0f0)))

    @test @inferred(logcdf(LogNormal(0 // 1, 0 // 1), 1 // 1)) === -0.0
    @test @inferred(logcdf(LogNormal(0 // 1, 0 // 1), 0 // 1)) === -Inf
    @test @inferred(logcdf(LogNormal(0 // 1, 0 // 1), -1 // 1)) === -Inf
    @test isnan_type(Float64, @inferred(logcdf(LogNormal(0 // 1, 0 // 1), NaN)))

    @test @inferred(logcdf(LogNormal(0.0, 0.0), BigInt(1))) == big(0.0)
    @test @inferred(logcdf(LogNormal(0.0, 0.0), BigInt(0))) == big(-Inf)
    @test @inferred(logcdf(LogNormal(0.0, 0.0), BigInt(-1))) == big(-Inf)
    @test @inferred(logcdf(LogNormal(0.0, 0.0), BigFloat(1))) == big(0.0)
    @test @inferred(logcdf(LogNormal(0.0, 0.0), BigFloat(0))) == big(-Inf)
    @test @inferred(logcdf(LogNormal(0.0, 0.0), BigFloat(-1))) == big(-Inf)
    @test isnan_type(BigFloat, @inferred(logcdf(LogNormal(0.0, 0.0), BigFloat(NaN))))

    # ccdf
    @test @inferred(ccdf(LogNormal(0.0, 0.0), 1.0)) === 0.0
    @test @inferred(ccdf(LogNormal(0.0, 0.0), 0.5)) === 1.0
    @test @inferred(ccdf(LogNormal(0.0, 0.0), 0.0)) === 1.0
    @test @inferred(ccdf(LogNormal(0.0, 0.0), -0.5)) === 1.0

    @test @inferred(ccdf(LogNormal(0.0, 0.0), 1.0f0)) === 0.0
    @test @inferred(ccdf(LogNormal(0.0f0, 0.0f0), 1.0)) === 0.0
    @test @inferred(ccdf(LogNormal(0.0f0, 0.0f0), 1.0f0)) === 0.0f0

    @test isnan_type(Float64, @inferred(ccdf(LogNormal(0.0, 0.0), NaN)))
    @test isnan_type(Float64, @inferred(ccdf(LogNormal(NaN, 0.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(ccdf(LogNormal(NaN, 0.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(ccdf(LogNormal(NaN, 0.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(ccdf(LogNormal(NaN32, 0.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(ccdf(LogNormal(NaN32, 0.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(ccdf(LogNormal(NaN32, 0.0f0), -1.0f0)))

    @test @inferred(ccdf(LogNormal(0 // 1, 0 // 1), 1 // 1)) === 0.0
    @test @inferred(ccdf(LogNormal(0 // 1, 0 // 1), 0 // 1)) === 1.0
    @test @inferred(ccdf(LogNormal(0 // 1, 0 // 1), -1 // 1)) === 1.0
    @test isnan_type(Float64, @inferred(ccdf(LogNormal(0 // 1, 0 // 1), NaN)))

    @test @inferred(ccdf(LogNormal(0.0, 0.0), BigInt(1))) == big(0.0)
    @test @inferred(ccdf(LogNormal(0.0, 0.0), BigInt(0))) == big(1.0)
    @test @inferred(ccdf(LogNormal(0.0, 0.0), BigInt(-1))) == big(1.0)
    @test @inferred(ccdf(LogNormal(0.0, 0.0), BigFloat(1))) == big(0.0)
    @test @inferred(ccdf(LogNormal(0.0, 0.0), BigFloat(0))) == big(1.0)
    @test @inferred(ccdf(LogNormal(0.0, 0.0), BigFloat(-1))) == big(1.0)
    @test isnan_type(BigFloat, @inferred(ccdf(LogNormal(0.0, 0.0), BigFloat(NaN))))

    # logccdf
    @test @inferred(logccdf(LogNormal(0.0, 0.0), 1.0)) === -Inf
    @test @inferred(logccdf(LogNormal(0.0, 0.0), 0.5)) === -0.0
    @test @inferred(logccdf(LogNormal(0.0, 0.0), 0.0)) === -0.0
    @test @inferred(logccdf(LogNormal(0.0, 0.0), -0.5)) === -0.0

    @test @inferred(logccdf(LogNormal(0.0, 0.0), 1.0f0)) === -Inf
    @test @inferred(logccdf(LogNormal(0.0f0, 0.0f0), 1.0)) === -Inf
    @test @inferred(logccdf(LogNormal(0.0f0, 0.0f0), 1.0f0)) === -Inf32

    @test isnan_type(Float64, @inferred(logccdf(LogNormal(0.0, 0.0), NaN)))
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(NaN, 0.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(NaN, 0.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(NaN, 0.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(logccdf(LogNormal(NaN32, 0.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(logccdf(LogNormal(NaN32, 0.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(logccdf(LogNormal(NaN32, 0.0f0), -1.0f0)))

    @test @inferred(logccdf(LogNormal(0 // 1, 0 // 1), 1 // 1)) === -Inf
    @test @inferred(logccdf(LogNormal(0 // 1, 0 // 1), 0 // 1)) === -0.0
    @test @inferred(logccdf(LogNormal(0 // 1, 0 // 1), -1 // 1)) === -0.0
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(0 // 1, 0 // 1), NaN)))

    @test @inferred(logccdf(LogNormal(0.0, 0.0), BigInt(1))) == big(-Inf)
    @test @inferred(logccdf(LogNormal(0.0, 0.0), BigInt(0))) == big(0.0)
    @test @inferred(logccdf(LogNormal(0.0, 0.0), BigInt(-1))) == big(0.0)
    @test @inferred(logccdf(LogNormal(0.0, 0.0), BigFloat(1))) == big(-Inf)
    @test @inferred(logccdf(LogNormal(0.0, 0.0), BigFloat(0))) == big(0.0)
    @test @inferred(logccdf(LogNormal(0.0, 0.0), BigFloat(-1))) == big(0.0)
    @test isnan_type(BigFloat, @inferred(logccdf(LogNormal(0.0, 0.0), BigFloat(NaN))))

    # quantile
    @test @inferred(quantile(LogNormal(1.0, 0.0), 0.0f0)) === 0.0
    @test @inferred(quantile(LogNormal(1.0, 0.0f0), 1.0)) === Inf
    @test @inferred(quantile(LogNormal(1.0f0, 0.0), 0.5)) ===  exp(1)
    @test isnan_type(Float64, @inferred(quantile(LogNormal(1.0f0, 0.0), NaN)))
    @test @inferred(quantile(LogNormal(1.0f0, 0.0f0), 0.0f0)) === 0.0f0
    @test @inferred(quantile(LogNormal(1.0f0, 0.0f0), 1.0f0)) === Inf32
    @test @inferred(quantile(LogNormal(1.0f0, 0.0f0), 0.5f0)) === exp(1.0f0)
    @test isnan_type(Float32, @inferred(quantile(LogNormal(1.0f0, 0.0f0), NaN32)))
    @test @inferred(quantile(LogNormal(1//1, 0//1), 1//2)) === exp(1)

    # cquantile
    @test @inferred(cquantile(LogNormal(1.0, 0.0), 0.0f0)) === Inf
    @test @inferred(cquantile(LogNormal(1.0, 0.0f0), 1.0)) === 0.0
    @test @inferred(cquantile(LogNormal(1.0f0, 0.0), 0.5)) === exp(1)
    @test isnan_type(Float64, @inferred(cquantile(LogNormal(1.0f0, 0.0), NaN)))
    @test @inferred(cquantile(LogNormal(1.0f0, 0.0f0), 0.0f0)) === Inf32
    @test @inferred(cquantile(LogNormal(1.0f0, 0.0f0), 1.0f0)) === 0.0f0
    @test @inferred(cquantile(LogNormal(1.0f0, 0.0f0), 0.5f0)) === exp(1.0f0)
    @test isnan_type(Float32, @inferred(cquantile(LogNormal(1.0f0, 0.0f0), NaN32)))
    @test @inferred(cquantile(LogNormal(1//1, 0//1), 1//2)) === exp(1)

    # gradlogpdf
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), 1.0)) === -1.0
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), exp(-1))) === 0.0
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), 0.0)) === 0.0
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), -0.5)) === 0.0

    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), 1.0f0)) === -1.0
    @test @inferred(gradlogpdf(LogNormal(0.0f0, 1.0f0), 1.0)) === -1.0
    @test @inferred(gradlogpdf(LogNormal(0.0f0, 1.0f0), 1.0f0)) === -1.0f0

    @test isnan_type(Float64, @inferred(logccdf(LogNormal(0.0, 1.0), NaN)))
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(NaN, 1.0), 1.0f0)))
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(NaN, 1.0), 0.0f0)))
    @test isnan_type(Float64, @inferred(logccdf(LogNormal(NaN, 1.0), -1.0f0)))

    @test isnan_type(Float32, @inferred(logccdf(LogNormal(NaN32, 1.0f0), 1.0f0)))
    @test isnan_type(Float32, @inferred(logccdf(LogNormal(NaN32, 1.0f0), 0.0f0)))
    @test isnan_type(Float32, @inferred(logccdf(LogNormal(NaN32, 1.0f0), -1.0f0)))

    @test @inferred(gradlogpdf(LogNormal(0 // 1, 1 // 1), 1 // 1)) === -1.0
    @test @inferred(gradlogpdf(LogNormal(0 // 1, 1 // 1), 0 // 1)) === 0.0
    @test @inferred(gradlogpdf(LogNormal(0 // 1, 1 // 1), -1 // 1)) === 0.0
    @test isnan_type(Float64, @inferred(gradlogpdf(LogNormal(0 // 1, 1 // 1), NaN)))

    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigInt(1))) == big(-1.0)
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigInt(0))) == big(0.0)
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigInt(-1))) == big(0.0)
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigFloat(1))) == big(-1.0)
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigFloat(0))) == big(0.0)
    @test @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigFloat(-1))) == big(0.0)
    @test isnan_type(BigFloat, @inferred(gradlogpdf(LogNormal(0.0, 1.0), BigFloat(NaN))))
end