# # Tutorial: missing value handling
#
# This tutorial illustrates the usage of Transducers.jl by stepping
# through various handling of missing values.

using Transducers

# ## Dot product
#
# Here is a simple way to compute a dot product using
# [`mapfoldl`](@ref) and [`MapSplat`](@ref):

mapfoldl(MapSplat(*), +, zip(1:3, 10:2:14))

# Let's see what it does step by step.  First we create a "printer"
# transducer using the following function (see [`Map`](@ref)):

xf_printer(label) = Map() do x
    println(label, ": ", x)
    return x  # just return it as-is
end
nothing  # hide

# This transducer just pass-through the input while printing its value
# (prefixed by a `label`).  Let's sandwich the previous `MapSplat(*)`
# with it:

mapfoldl(
    xf_printer(" input") |> MapSplat(*) |> xf_printer("output"),
    +, zip(1:3, 10:2:14))
flush(stdout)  # hide

# You can see that the input tuple `(1, 10)` is splatted into function
# `*` by `MapSplat` which then outputs `10`.  This is repeated for all
# inputs.

# Perhaps unfortunately, this way of computing a dot product
# propagates any missing values contained in the input arrays to the
# result (which may actually be desired in certain cases).

xs = [1, missing, 3, 2]
ys = [10, 14, missing, 12]
mapfoldl(MapSplat(*), +, zip(xs, ys))

# However, it is very simple to ignore any missing values using
# [`OfType`](@ref):

xf_mdot = OfType(Tuple{Vararg{Number}}) |> MapSplat(*)
mapfoldl(xf_mdot, +, zip(xs, ys))

# Here, `Tuple{Vararg{Number}}` is a type that matches with a tuple of
# any length with numbers.  It does not match with a tuple if it has a
# `missing`.

@assert (1, 0.5) isa Tuple{Vararg{Number}}
@assert (1, 0.5, 2im) isa Tuple{Vararg{Number}}
@assert !((1, missing) isa Tuple{Vararg{Number}})

# ## Covariance
#
# Transducer `xf_mdot` above can also be used to compute the
# covariance.  First, we need the number of pairs of elements in `xs`
# and `ys` that _both_ of them are not `missing`:

nonmissings = mapfoldl(OfType(Tuple{Vararg{Number}}) |> Count(),
                       right,
                       zip(xs, ys);
                       init = 0)
nonmissings  # hide
#-
@assert nonmissings == 2  # hide

# We do this by using [`Count`](@ref) and [`right`](@ref).  `Count`
# ignores input and count the number of times the input is provided.
# Since `OfType(Tuple{Vararg{Number}})` provides the inputs to the
# downstream transducer only if there is no `missing` values, this
# correctly counts the number of non-missing pairs.  Function `right`
# is simply defined as `right(l, r) = r` (and `right(r) = r`).  Thus,
# the whole `mapfoldl` returns the last output of `Count`.  In case
# `Count` never gets called (i.e., there are no non-missing pairs), we
# pass `init=0`.

ans =  # hide
mapfoldl(OfType(Tuple{Vararg{Number}}) |> Count(),
         right,
         zip(Int[], Int[]);
         init = 0)
ans  # hide
#-
@assert ans == 0  # hide

# Finally, we have to pre-process the input to `xf_mdot` by
# subtracting the average.  It's easy to do with `Map`:

using Statistics: mean

function xf_demean(xs, ys)
    xmean = mean(skipmissing(xs))
    ymean = mean(skipmissing(ys))
    return Map(((x, y),) -> (x - xmean, y - ymean))
end

mapfoldl(xf_demean(xs, ys) |> xf_mdot, +, zip(xs, ys)) / nonmissings

# ## Addition
#
# How do we use transducers for vector-to-vector transformation?  Here
# is a function to calculate ``y = x + y`` while ignoring missing
# values in ``x``.  First, mandatory input shape check:

function add_skipmissing!(ys, xs)
    length(ys) == length(xs) || error("length(ys) != length(xs)")
    firstindex(ys) == 1 || error("firstindex(ys) != 1")
#+
# For filtering out missing values from `xs` while tracking indices,
# we use [`Enumerate`](@ref) and [`Filter`](@ref).  To iterate over
# the output of the transducer, [`foreach`](@ref) is used instead of
# [`mapfoldl`](@ref) since mutating an array is better expressed as a
# side-effect than a fold.

    foreach(Enumerate() |> Filter(!(ismissing ∘ last)), xs) do (i, xi)
        @inbounds ys[i] += xi
    end
#+
# We then return the mutated value to behave like the rest of Julia
# functions (`push!`, `mul!`, etc.):

    return ys
end
nothing  # hide

# Example:

ans =  # hide
add_skipmissing!([100, 110, 120], [1, missing, 2])
#-
@assert ans == [101, 110, 122]  # hide

# ## Vectorized reduction
#
# `foldl`, `mapfoldl`, etc. in `Base` support `dims` argument.
# Transducers.jl does not support this argument as of writing.
# However, this can easily be emulated using `eachcol`, `eachrow`, or
# `eachslice` iterators in `Base`.

xs = [
    0       missing 1       2
    3       4       5       missing
    missing 6       7       missing
]

function xf_sum_columns(prototype)
    T = Base.nonmissingtype(eltype(prototype)) # subtract Missing from type
    dims = size(prototype)
    return Scan(add_skipmissing!, CopyInit(zeros(T, dims)))
end
nothing  # hide

# We use [`CopyInit`](@ref) here to allocate the "output array"
# into which the columns are added by `add_skipmissing!`.

if VERSION >= v"1.1-"  # eachcol not in Julia 1.0  #src
ans =  # hide
mapfoldl(xf_sum_columns(xs[:, 1]), right, eachcol(xs))
#-
@assert ans == [3, 12, 13]  # hide

# Above computation returns the sum over each row without taking into
# account the relationship within a column.  Another possibly useful
# reduction is the sum of the columns with no missing values.  This
# can easily be done by prepending a filter:

ans =  # hide
mapfoldl(Filter(x -> !any(ismissing, x)) |> xf_sum_columns(xs[:, 1]),
         right, eachcol(xs))
#md ans  # hide
#src `#md ans` is a hack to avoid Literate.jl to put `continued = true`.
#-
@assert ans == [1, 5, 7]  # hide

# Note that above combination of `Scan` and `right` is redundant.  For
# example, we can simply pass `add_skipmissing!` to "normal" `foldl`:

ans =  # hide
foldl(add_skipmissing!, eachcol(xs), init=zeros(Int, size(xs, 1)))
#-
@assert ans == [3, 12, 13]  # hide

# However, packaging it as a transducer is sometimes useful as it can
# be composed with other transducers and "bottom" reducing function.
# For example, vectorized version of `cumsum` can easily obtained by
# composing it with `append!` (and then `reshape` after `mapfoldl`):

result = mapfoldl(
    xf_sum_columns(xs[:, 1]),
    Completing(append!),
    eachcol(xs);
    init = Int[])
ans =  # hide
reshape(result, (size(xs, 1), :))
#-
@assert ans == [      # hide
    0  0   1   3      # hide
    3  7  12  12      # hide
    0  6  13  13      # hide
]                     # hide
end  # if VERSION >= v"1.1-"  #src

# Note that we need [`Completing`](@ref) here since `append!` does not
# have the unary method used for [`complete`](@ref
# Transducers.complete) protocol.

# ## Argmax
#
# Another useful operation to do ignoring missing values is
# `argmax`/`argmin`.  It can be implemented using `Enumerate() |>
# Filter(!(ismissing ∘ last))` (see also `add_skipmissing!` above)
# composed with [`ScanEmit`](@ref).  We first need to define a
# function to be called by `ScanEmit`:

##                     ,--- current state
##                     |
##                     |              ,-- input
##                     |              |
function argmax_step((argmax, max), (index, value))
    argmax, max = value > max ? (index, value) : (argmax, max)
    return argmax, (argmax, max)
    ##       \        \
    ##        \        \__ next state
    ##         \
    ##          \__ output
end
nothing  # hide

# This function is passed to `ScanEmit` with the initial state:

xf_argmax = Enumerate() |> Filter(!(ismissing ∘ last)) |>
    ScanEmit(argmax_step, (0, typemin(Int)))
##                          |
##                   initial state
nothing  # hide

# As [`ScanEmit`](@ref) is one of the most complex (and powerful)
# transducer, it may require some comments on how above code works:
#
# * The state `(argmax, max)` is initialized to `(0, typemin(Int))` in
#   `xf_argmax`.  This is the first value passed to the first argument
#   `(argmax, max)` of `argmax_step`.
#
# * The upstream transducer `Enumerate()` provides `(index,
#   value)`-pair which becomes the input (the second argument) of
#   `argmax_step`.
#
# * Function `argmax_step` must return a pair.  The first item becomes
#   the output of `ScanEmit`.  In this case that's the index of the
#   largest item seen so far.
#
# * The second item in the returned pair is fed back to `argmax_step`
#   in the next call.
#
# We have the argmax function by extracting the last output of
# `xf_argmax`:

mapfoldl(xf_argmax, right, [1, 3, missing, 2])

# Side note: We use `typemin(Int)` as the initial value of `max` for
# simplicity.  In practice, it should be
# `typemin(eltype(input_array))`.  A more generic solution is to
# special-case the first invocation by using a singleton like
# `nothing`.  Julia can handle small `Union` type such as this (see
# the next section).  Another solution is to use `Init(>)` from
# InitialValues.jl.

# ## Extrema
#
# Transducer `xf_argmax` in the previous section only outputs the
# index of the maximum element so far.  To output the maximum element
# as well, we can simply use [`Scan`](@ref).  Also, while we are at
# it, let's support both argmax and argmin.  To this end, we
# parametrize the function passed to `Scan` by the comparison function
# `>` and `<`.  Another problem with `xf_argmax` is that it does not
# handle non-`Int` input types.  To properly handle different input
# types, we initialize `Scan`'s state with `nothing` and special-case
# the first invocation to return the input as-is.  Following function
# `argext_step` takes the function `>` or `<` and return a function
# appropriate for `Scan.`

argext_step(should_update) =
    (old, (index, value)) ->
        if old === nothing || should_update(old[2], value)
            (index, value)
        else
            old
        end
nothing  # hide

xf_scanext(should_update) = Scan(argext_step(should_update), nothing)
nothing  # hide

# Passing `<` gives us the argmax transducer:

@time begin  #src
ans = # hide
mapfoldl(
    Enumerate() |> OfType(Tuple{Integer, Number}) |> xf_scanext(<),
    right, [1.0, 3.0, missing, 2.0])
end  #src
#-
@assert ans === (2, 3.0) # hide
@show ans  #src

# We now have transducers `xf_scanext(<)` and `xf_scanext(>)` for
# argmax and argmin, respectively.  We can compute them concurrently
# by [`Zip`](@ref)'ing them together:

xf_fullextrema = Enumerate() |> OfType(Tuple{Integer, Number}) |>
    Zip(xf_scanext(>), xf_scanext(<))

@time begin  #src
ans = # hide
mapfoldl(xf_fullextrema, right, [1.0, 3.0, -1.0, missing, 2.0])
end  #src
#-
@assert ans === ((3, -1.0), (2, 3.0))  # hide
@show ans  #src

# This transducer produces a tuple `((argmin, min), (argmax, max))`.
# To output only indices, append an appropriate `Map`:

xf_argextrema =
    xf_fullextrema |> Map() do ((argmin, min), (argmax, max))
        (argmin, argmax)
    end

@time begin  #src
ans = # hide
mapfoldl(xf_argextrema, right, [1.0, 3.0, -1.0, missing, 2.0])
end  #src
#-
@assert ans === (3, 2)  # hide
@show ans  #src


"""
    eom_anishchenko1(u, p, t)

Equations of motion for the system defined by eq. 13 in [1], which can
give strange, nonchaotic attractors.

## Equations of motion

The equations of motion are

```math
\\begin{aligned}
dx &= \\alpha (1-s \\cos (2 \\pi \\phi )) \\cdot x(1-x) \\\\
dϕ &= (\\phi + \\omega ) \\mod{1}
\\end{aligned}
```

## References

1. Anishchenko, Vadim S., and Galina I. Strelkova. "Irregular attractors."
    Discrete dynamics in Nature and Society 2.1 (1998): 53-72.
"""
function eom_anishchenko1(u, p, t)
    x, ϕ = (u...,)
    α, s, ω = (p...,)
    dx = α*(1 - s*cos(2*pi*ϕ))*x*(1 - x)
    dϕ = (ϕ + ω) % 1

    return SVector{2}(dx, dϕ)
end

function anishchenko1(u₀, α, s, ω)
    p = [α, s, ω]
    DiscreteDynamicalSystem(eom_anishchenko1, u₀, p)
end

"""
    anishchenko1(;u₀ = rand(2), α =3.277, s=0.1, ω=0.5*(sqrt(5)-1))

Initialise the system defined by eq. 13 in [1], which can give strange,
nonchaotic attractors.

## Equations of motion

The equations of motion are

```math
\\begin{aligned}
dx &= \\alpha (1-s \\cos (2 \\pi \\phi )) \\cdot x(1-x) \\\\
dϕ &= (\\phi + \\omega ) \\mod{1}
\\end{aligned}
```

## References

1. Anishchenko, Vadim S., and Galina I. Strelkova. "Irregular attractors."
    Discrete dynamics in Nature and Society 2.1 (1998): 53-72.
"""
anishchenko1(;u₀ = rand(2), α =3.277, s=0.1, ω=0.5*(sqrt(5)-1)) =
    anishchenko1(u₀, α, s, ω)


module mlmetrics
using  Statistics

"""
    struct mse{T<:AbstractFloat}
        y::Vector
        ŷ::Vector
    end
    The mean squared error is the sum of the squared differences between the
    predicted and observed values.

    struct mae{T<:AbstractFloat}
        y::Vector
        ŷ::Vector
    end
    The mean absolute error is the sum of the absolute differences between the
    predicted and observed values.

    struct rmse{T<:AbstractFloat}
        y::Vector
        ŷ::Vector
    end
    The root mean squared error is the square root of the mean squared error.

    struct rmae{T<:AbstractFloat}
        y::Vector
        ŷ::Vector
    end
    The root mean absolute error is the square root of the mean absolute error.

    struct rmsle{T<:AbstractFloat}
        y::Vector
        ŷ::Vector
    end
    The root mean squared log error is the square root of the mean squared log error.

    struct r2score{T<:AbstractFloat}
        y::Vector
        ŷ::Vector
    end
    The R2 is a statistic that will give some information about the goodness of fit of a model.
"""

struct mse{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function mse(y, ŷ)
        return mean((y - ŷ).^2)
    end
end

struct mae{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function mae(y, ŷ)
        return mean(abs.(y - ŷ))
    end
end

struct rmse{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function rmse(y, ŷ)
        return sqrt(mse(y, ŷ))
    end
end

struct rmae{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function rmae(y, ŷ)
        return sqrt(mae(y, ŷ))
    end
end

struct rmsle{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function rmsle(y, ŷ)
        return sqrt(mean((log.(y .+ 1) - log.(ŷ .+ 1)).^2))
    end
end

struct r2score{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function r2score(y, ŷ)
        squares = sum(y.^2)
        residuals = sum(y - ŷ).^2
        return 1 - residuals/squares
    end
end

struct binarycrossentropy{T<:AbstractFloat}
    y::Vector
    ŷ::Vector

    function binarycrossentropy(y::Vector, ŷ::Vector)
        ε = 1e-10
        return -mean(y.*log2.(ŷ .+ ε) + (1 .- y).*log2.(1 .- ŷ .+ ε))
    end
end

struct accuracy{T<:AbstractFloat}
    y_true::AbstractArray
    y_predict::AbstractArray

    function accuracy(y_true::AbstractArray, y_predict::AbstractArray)
        return sum(y_true.==y_predict)/length(y_true)
    end

end

struct precision{T<:AbstractFloat}
    y_true::AbstractArray
    y_predict::AbstractArray

    function precision(y_true::AbstractArray, y_predict::AbstractArray)
        true_positives = sum(y_true.== 1 .== y_predict)
        true_negatives = sum(y_true.== 0 .== y_predict)
        false_positives = sum(y_true.!= 1 .== y_predict)
        false_negatives = sum(y_true.!= 0 .== y_predict)

        return true_positives/(true_positives + false_positives)
    end

end

struct recall{T<:AbstractFloat}
    y_true::AbstractArray
    y_predict::AbstractArray

    function recall(y_true::AbstractArray, y_predict::AbstractArray)
        true_positives = sum(y_true.== 1 .== y_predict)
        true_negatives = sum(y_true.== 0 .== y_predict)
        false_positives = sum(y_true.!= 1 .== y_predict)
        false_negatives = sum(y_true.!= 0 .== y_predict)

        return true_positives/(true_positives + false_negatives)
    end

end

struct f1_score{T<:AbstractFloat}
    y_true::AbstractArray
    y_predict::AbstractArray

    function f1_score(y_true::AbstractArray, y_predict::AbstractArray)
        num =  2 * precision(y_true, y_predict) * recall(y_true, y_predict)
        den = precision(y_true, y_predict) + recall(y_true, y_predict)
        return num/den
    end

end

end


using EngEconomics, Roots

# Given
# a = parasailing
# b = kayaking
aInit = -100000
aAnnual = 15000
aN = 15

bInit = -10000
bAnnual = 2000
bN = 15
MARR = 0.10

N = aN

# Determine the better alternative
ΔInit = aInit - bInit
ΔAnnual = aAnnual - bAnnual
ΔIRR(x) = ΔInit + ΔAnnual * seriesPresentAmountFactor(x, N)
ΔIRR_ans = find_zero(ΔIRR, 0.05)

if ΔIRR_ans > MARR
	println("Choose Parasailing")
else
	println("Choose Kayaking")
end


# stoichiometric matrix: fluxes (in order) are v1-v6 (v6 is reverse of v5), b1-b20
S = [-1 0 0 1 2 -2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #citrulline
     -1 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;  #aspartate
     1 -1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #arginosuccinate
     0 1 0 0 0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #fumarate
     0 1 -1 0 -2 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #arginine
     0 0 1 -1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;  #ornithine
     0 0 1 0 0 0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #urea
     0 0 0 -1 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #carbamoyl phosphate
     -1 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #ATP
     1 0 0 0 0 0 0 0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0 0 0; #AMP
     1 0 0 0 0 0 0 0 0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0 0; #diphosphate
     0 0 0 1 0 0 0 0 0 0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0; #phosphate
     0 0 0 0 -3 3 0 0 0 0 0 0 0 0 1 -1 0 0 0 0 0 0 0 0 0 0; #NADPH
     0 0 0 0 -3 3 0 0 0 0 0 0 0 0 0 0 1 -1 0 0 0 0 0 0 0 0; #H+
     0 0 0 0 -4 4 0 0 0 0 0 0 0 0 0 0 0 0 1 -1 0 0 0 0 0 0; #O2
     0 0 0 0 2 -2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 -1 0 0 0 0; #NO2
     0 0 0 0 3 -3 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 -1 0 0; #NADP+
     0 0 -1 0 4 -4 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 -1]; #H2O

#atom matrix: C H N O P S
A = [6 13 3 3 0 0; #citrulline
    4 7 1 4 0 0;  #aspartate
    10 18 4 6 0 0; #arginosuccinate
    4 4 0 4 0 0; #fumarate
    6 14 4 2 0 0; #arginine
    5 12 2 2 0 0; #ornithine
    1 4 2 1 0 0; #urea
    1 4 1 5 1 0; #carbamoyl phosphate
    10 16 5 13 3 0; #ATP
    10 14 5 7 1 0;  #AMP
    0 4 0 7 2 0;  #diphosphate
    0 3 0 4 1 0;  #phosphate
    21 30 7 17 3 0; #NADPH
    0 1 0 0 0 0; #H+
    0 0 0 2 0 0; #O2
    0 0 1 1 0 0; #NO
    21 29 7 17 3 0; #NADP+
    0 2 0 1 0 0] #H2O

e = transpose(A)*S

#look at only the balances for v1-v6, species are not balanced over boundaries
print(e[:,1:6])


using Knet, ParticleSwarmOptimization
include("Args.jl")

hosaki(x) = (1 - 8*x[1] + 7*(x[1]^2) - 7/3*(x[1]^3) + (1/4)*(x[1]^4))*(x[2]^2)*exp(-1 * x[2])
#x* = (4,2)
#f(x*) = -2.3458
#julia hosaki.jl --num 10 --epoch 100

function main()
    gpu(false)
    args = parse_commandline()

    problem = RegularProblem(hosaki, 2, [0.0,0.0], [5.0,6.0])
    gbest = pso(problem, args["num"], args["epoch"], args["phi1"], args["phi2"], args["w"], args["vmin"], args["vmax"])

    println("\nBest: $(gbest.pbest_fitness), $(gbest.x)")
end

main()


# ---
# title: 241. Different Ways to Add Parentheses
# id: problem241
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Divide and Conquer
# link: <https://leetcode.com/problems/different-ways-to-add-parentheses/description/>
# hidden: true
# ---
#
# Given a string of numbers and operators, return all possible results from
# computing all the different possible ways to group numbers and operators. The
# valid operators are `+`, `-` and `*`.
#
# **Example 1:**
#
#
#
#     Input: "2-1-1"
#     Output: [0, 2]
#     Explanation:
#     ((2-1)-1) = 0
#     (2-(1-1)) = 2
#
# **Example 2:**
#
#
#
#     Input:"2*3-4*5"
#     Output: [-34, -14, -10, -10, 10]
#     Explanation: (2*(3-(4*5))) = -34
#     ((2*3)-(4*5)) = -14
#     ((2*(3-4))*5) = -10
#     (2*((3-4)*5)) = -10
#     (((2*3)-4)*5) = 10 ****
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using KrylovMethods
using LinearOperators
using MatrixDepot
using Base.Test

println("=== Testing LSQR ===")

# small test
A   = [1 1; 0 -1; 1 0.]
r1  = A[:,1]   # compatible
r2  = randn(3) # incompatible
r3  = [1;1;-1.] # in the nullpace of A'

x1 = lsqr(A,r1,out=2,atol=1e-8,doBidiag=true,storeInterm=true)
@test norm(x1[1][:,end]-[1;0])<1e-10
@test x1[2]==1
x2 = lsqr(A,r2,out=1,atol=1e-8)
@test x2[2]==2
x3 = lsqr(A,r3,out=2)
@test all(x3[1].==0)

# test with sparse random matrix
A   = sprandn(100,10,.2)
rhs = randn(100)
xgt = full(A)\rhs
Af = (flag,x,a=0.0,v=0.0) -> ((flag=='F') ? A*x +a*v : A'*x+a*v)

x1  = lsqr(LinearOperator(A),rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10,out=1)
x2  = lsqr(A,rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)
x3  = lsqr(Af,rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)

@test norm(xgt-x1[1])/norm(xgt) < 1e-6
@test norm(xgt-x2[1])/norm(xgt) < 1e-6
@test norm(xgt-x3[1])/norm(xgt) < 1e-6
@test norm(x1[1]-x2[1])/norm(x1[1]) < 1e-12
@test norm(x2[1]-x3[1])/norm(x3[1]) < 1e-12

# test with dense hilbert matrix
A   = matrixdepot("hilb",3,3)
rhs = [1;0;0.]
xgt = A\rhs

x1  = lsqr(LinearOperator(A),rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)
x2  = lsqr(A,rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)

@test norm(xgt-x1[1])/norm(xgt) < 1e-6
@test norm(xgt-x2[1])/norm(xgt) < 1e-6
@test norm(x1[1]-x2[1])/norm(x1[1]) < 1e-12
conA = norm(vec(A))*norm(vec(pinv(A)))
@test norm(conA-x1[3][end,end])/conA < 1e-12
@test norm(conA-x2[3][end,end])/conA < 1e-12

println("=== LSQR: All tests passed. ===")


function getwraparoundDFTfreqs(N::Int, fs::T, ν_begin::T) where T

    U0 = getDFTfreqrange(N, fs)
    out, inds = wrapfreqrange(U0, ν_begin, fs)

    return U0, out, inds
end

# Assumes U0 is sorted in ascending order.
function wrapfreqrange(U0, ν_begin::T, fs::T) where T <: Real

    N = length(U0)
    ind = findfirst(xx->(xx>ν_begin), U0)

    # TODO handle these exceptions with more grace.
    @assert typeof(ind) == Int
    @assert ind <= N

    out = Vector{T}(undef, N)
    #out[1:ind] = U0[1:ind] .+ fs
    #out[ind+1:end] = U0[ind+1:end]

    M = N-ind
    out[1:M] = U0[ind+1:end]
    out[M+1:end] = U0[1:ind] .+ fs

    inds = collect(1:N)
    inds[1:M] = collect(ind+1:N)
    inds[M+1:end] = collect(1:ind)

    return out, inds
end

function gettimerange(N::Int, fs::T) where T
    Ts::T = 1/fs

    return zero(T):Ts:(N-1)*Ts
 end

 function getDFTfreqrange(N::Int, fs::T)::LinRange{T} where T
     a = zero(T)
     b = fs-fs/N

     return LinRange(a, b, N)
 end

 # case 1D.
function computeDTFTch3eq29(h, u::T, Λ)::Complex{T} where T <: Real

    # debug_array = zeros(Complex{T}, length(Λ))
    # debug_array2 = zeros(Complex{T}, length(Λ))
    # fs = 1/(Λ[2]-Λ[1])

    running_sum = zero(T)
    for i = 1:length(Λ)
        x = Λ[i]

        running_sum += h[i]*exp(-im*2*π*u*x)

        # # debug.
        # debug_array[i] = h[i]*exp(-im*2*π*u*x)
        #
        # #debug_array2[i] = h[i]*exp(-im*2*π*u*x +im*35*2*pi)
        # debug_array2[i] = h[i]*exp(-im*2*π*(u+3*fs)*x)
    end

    #println("debug: norm(debug_array -debug_array2 ) = ", norm(debug_array -debug_array2 ) )
    return running_sum
end


"plotVF!(fx,fy,x,y) plot vector field"
function plotVF!(p,fx::Function,fy::Function,xbase::Union{Array,UnitRange,StepRangeLen},ybase::Union{Array,UnitRange,StepRangeLen};map_speed = true,log_speed = false)

    dxbase = xbase[2] - xbase[1]
    dybase = ybase[2] - ybase[1]
    dxybase = sqrt(dxbase^2+dybase^2)
    cprod = collect.(collect(Base.product(xbase,ybase))) # not writing for-loop, use Cartesian product of x base and y base, cprod is a permutation set
    cprod = vec(cprod)
    x = map(x->x[1],cprod)
    y = map(x->x[2],cprod)
    dx=map(x->fx(x[1],x[2]),cprod)
    dy=map(x->fy(x[1],x[2]),cprod)
    speed = sqrt.(dx.^2 .+ dy.^2)#[sqrt(dx[i]^2+dy[i]^2) for i in 1:length(dx)]

    if log_speed
        speed = log.(speed)
        dx=log.(dx)
        dy = log.(dy)
    end

    #normalization
    dx ./=maximum(abs.(dx))
    dx .*=dxbase
    dy ./=maximum(abs.(dy))
    dy .*=dybase
    #show(dx);show(dy)
    #how(speed)

    if map_speed
        return quiver!(p,x,y,quiver = (dx,dy))#;markerstrokecolor = qcMap_def(speed))
    end
    return quiver!(p,x,y,quiver = (dx,dy))
end

#fx(x,y)=x;fy(x,y)=y
#plotVF!(plot(),fx,fy,collect(-3:1.0:3),collect(-3:1.0:3))

#quiver(1:10,1:10,quiver=(10:-1:1,10:-1:1),markerstrokecolor = qcMap_def(1:10))

"plotGradient!(p,x) - autonomous"
function plotGradient!(p,fx::Function,xbase;t = 0:1.0:5,map_speed = true,log_speed = false,normalize = true)
    if t==Missing
        t = 1:1.0:length(x)
    end
    dt = t[2]-t[1]
    dx = xbase[2] - xbase[1]

    if isinteger(t[1])
        t = Float64.(t)
    end
    if isinteger(xbase[1])
        xbase = Float64.(xbase)
    end

    dtbase = t[2] - t[1]
    dxbase = xbase[2] - xbase[1]
    dbase = sqrt(dt^2+dt^2)
    cprod = collect.(collect(Base.product(t,xbase))) # not writing for-loop, use Cartesian product of x base and y base, cprod is a permutation set
    cprod = vec(cprod)
    x = map(x->x[2],cprod)
    t = map(x->x[1],cprod)
    dt= repeat([dtbase],length(t))
    dx=map(x->fx(x[2]),cprod)
    speed = sqrt.(dt.^2 .+ dx.^2)#[sqrt(dt[i]^2+dx[i]^2) for i in 1:length(dt)]

    if log_speed
        speed = log.(speed)
        dt=log.(dt)
        dx = log.(dx)
    end

    #normalization
    if normalize
        dx ./=maximum(abs.(dx))
        dx .*=dxbase
        #shrinks t vec also so that gradient is preserved

        dt ./=maximum(abs.(dx))
        dt .*-dtbase
    end
    #show(dt);show(dx)
    #how(speed)

    if map_speed
        return quiver!(p,t,x,quiver = (dt,dx))#;markerstrokecolor = qcMap_def(speed))
    end
    return quiver!(p,t,x,quiver = (dt,dx))
end

"plotNullCline(fx,fy)"
function plotNullCline!(p,fx::Function,fy::Function,xbase,ybase;tol = 1e-3)
    p=contour!(p,xbase,ybase,[fx(xi,yi) for yi in ybase, xi in xbase],levels = [-tol,tol],color = :red,alpha = .3)#,key = "x NullCline est")
    p=contour!(p,xbase,ybase,[fy(xi,yi) for yi in ybase, xi in xbase],levels = [-tol,tol],color =:blue ,alpha = .3)#,key = "y NullCline est")
    return p
end

#contour(1:2,1:2,[.11 .02;.1 0],levels = [-.1,.1])

"plotValongCycle!(p,fx::Function,fy::Function,cycle with r) works for 2d system only"
function plotValongCycle!(p,fx::Function,fy::Function,Cycle::Function;theta = 0.0:.3:2*pi)
    xl = zeros(length(theta))
    yl = similar(xl)
    dxl =similar(yl)
    dyl =similar(dxl)
    for i in 1:length(theta)
        err = ""
        try
            x,y = Cycle(theta[i])
            xl[i] = x
            yl[i] = y
            dxl[i] = fx(x,y)
            dyl[i] = fy(x,y)
        catch err
            println("Singularity?")
        end
    end
    p = quiver!(p,xl,yl,quiver = (dxl,dyl))
    p = plot!(p,xl,yl;alpha = .3,key = false)
    return p
end

function circularCycle(theta;center = (0.0,0.0),r = 1.0)
    [center[1]+ r * cos(theta), center[2] + r * sin(theta) ]
end

function plotValongCycle!(p,fx::Function,fy::Function;r = 1.0,center = (0.0,0.0),theta = .3:.3:2*pi)
    cycle(theta) = circularCycle(theta;center = center,r=r)
    return plotValongCycle!(p,fx,fy,cycle;theta = .3:.3:2*pi)
end


# # Differentiating a QP wrt a single variable

#md # [![](https://img.shields.io/badge/show-github-579ACA.svg)](@__REPO_ROOT_URL__/docs/src/examples/matrix-inversion-manual.jl)

# Consider the quadratic program

# ```math
# \begin{split}
# \begin{array} {ll}
# \mbox{minimize} & \frac{1}{2} x^T Q x + q^T x \\
# \mbox{subject to} & G x \leq h, x \in \mathcal{R}^2 \\
# \end{array}
# \end{split}
# ```

# where `Q`, `q`, `G` are fixed and `h` is the single parameter.

# In this example, we'll try to differentiate the QP wrt `h`, by finding its
# jacobian by hand (using Eqn (6) of [QPTH article](https://arxiv.org/pdf/1703.00443.pdf))
# and compare the results:
# - Manual compuation
# - Using JuMP and DiffOpt

# Assuming
# ```
# Q = [[4, 1], [1, 2]]
# q = [1, 1]
# G = [1, 1]
# ```
# and begining with a starting value of `h=-1`

# few values just for reference

# | variable | optimal value | note |
# |----|------|-----|
# | x* | [-0.25; -0.75] | Primal optimal |
# | 𝜆∗ | -0.75 | Dual optimal |

# ## Finding Jacobian using matrix inversion
# Lets formulate Eqn (6) of [QPTH article](https://arxiv.org/pdf/1703.00443.pdf) for our QP. If we assume `h` as the only parameter and `Q`,`q`,`G` as fixed problem data - also note that our QP doesn't involves `Ax=b` constraint - then Eqn (6) reduces to
# ```math
# \begin{gather}
#  \begin{bmatrix}
#      Q & g^T \\
#      \lambda^* G & G x^* - h
#  \end{bmatrix}
#  \begin{bmatrix}
#      dx \\
#      d \lambda
#  \end{bmatrix}
#  =
#   \begin{bmatrix}
#    0 \\
#    \lambda^* dh
#    \end{bmatrix}
# \end{gather}
# ```

# Now to find the jacobians $$ \frac{\partial x}{\partial h}, \frac{\partial \lambda}{\partial h}$$
# we substitute `dh = I = [1]` and plug in values of `Q`,`q`,`G` to get
# ```math
# \begin{gather}
#  \begin{bmatrix}
#      4 & 1 & 1 \\
#      1 & 2 & 1 \\
#      -0.75 & -0.75 & 0
#  \end{bmatrix}
#  \begin{bmatrix}
#      \frac{\partial x_1}{\partial h} \\
#      \frac{\partial x_2}{\partial h} \\
#      \frac{\partial \lambda}{\partial h}
#  \end{bmatrix}
#  =
#   \begin{bmatrix}
#    0 \\
#    0 \\
#    -0.75
#    \end{bmatrix}
# \end{gather}
# ```

# Upon solving using matrix inversion, the jacobian is
# ```math
# \frac{\partial x_1}{\partial h} = 0.25, \frac{\partial x_2}{\partial h} = 0.75, \frac{\partial \lambda}{\partial h} = -1.75
# ```

# ## Finding Jacobian using JuMP and DiffOpt

using JuMP
import DiffOpt
import Ipopt

n = 2 # variable dimension
m = 1; # no of inequality constraints

Q = [4. 1.;1. 2.]
q = [1.; 1.]
G = [1. 1.;]
h = [-1.;]   # initial values set

# Initialize empty model

model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
set_silent(model)

# Add the variables

@variable(model, x[1:2])

# Add the constraints.

@constraint(
    model,
    cons[j in 1:1],
    sum(G[j, i] * x[i] for i in 1:2)  <= h[j]
);

@objective(
    model,
    Min,
    1/2 * sum(Q[j, i] * x[i] *x[j] for i in 1:2, j in 1:2) +
    sum(q[i] * x[i] for i in 1:2)
)

# Solve problem

optimize!(model)

# primal solution

value.(x)

# dual solution

dual.(cons)

# set sentivitity

MOI.set(
    model,
    DiffOpt.ForwardInConstraint(),
    cons[1],
    0.0 * index(x[1]) - 1.0,  # to indicate the direction vector to get directional derivatives
)

# Note that `0.0 * index(x[1])` is used to make its type `typeof(0.0 * index(x[1]) - 1.0) <: MOI.AbstractScalarFunction`.
# To indicate different direction to get directional derivative, users should replace `0.0 * index(x[1]) - 1.0` as the form of `dG*x - dh`, where `dG` and `dh` correspond to the elements of direction vectors along `G` and `h` axes, respectively.

# Compute derivatives

DiffOpt.forward(model)

# Query derivative

dx = MOI.get.(
    model,
    DiffOpt.ForwardOutVariablePrimal(),
    x,
)

using Test                                  #src
@test dx ≈ [0.25 ,0.75] atol=1e-4 rtol=1e-4 #src


using QuasiCopula, Random, Statistics, Test, LinearAlgebra, StatsFuns
using LinearAlgebra: BlasReal, copytri!
### MVN only
@testset "Generate a 3 element random vector each with Y1, Y2, Y3 ~ Normal(5, 0.2). First we check the conditional mean and variance of Y2 | Y1 and then for Y3 | Y1, Y2. " begin
Random.seed!(12345)
mean_normal = 5
var_normal = 0.2
d1 = Normal(mean_normal, var_normal)
d2 = Normal(mean_normal, var_normal)
d3 = Normal(mean_normal, var_normal)

vecd = [d1, d2, d3]
variance_component_1 = 0.1
Γ = variance_component_1 * ones(3, 3)

nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)

#### Testing if the second element in the bivariate poisson follows correct mean and variance
# preallocate Y and res to simulate and store
n = 3
Y = Vector{Float64}(undef, n)
res = Vector{Float64}(undef, n)

# simulate Y1 first
i = 1
# form constants for Y1 pdf
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])
# simulate Y1 and store it to Y[1]
Random.seed!(12345)
Y[i] = rand(nonmixed_multivariate_dist.gc_obs[i])
# update residuals res[1] using Y[1] and the mean and variance (lambda for Poisson) using GLM package
res[i] = update_res!(Y[i], res[i], nonmixed_multivariate_dist.gc_obs[i])

# next, for conditional density of Y2 | Y1, form the constants for pdf
i = 2
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])

# pre allocate for sampling nsample times from conditional density
Random.seed!(12345)
nsample = 10_000
@info "sample $nsample points for the conditional Normal distribution"
s = Vector{Float64}(undef, nsample)
Random.seed!(12345)
rand!(nonmixed_multivariate_dist.gc_obs[i], s) # compile
Random.seed!(12345)
@time rand!(nonmixed_multivariate_dist.gc_obs[i], s) # get time
# 0.032922 seconds (1 allocation: 32 bytes)
println("sample mean = $(Statistics.mean(s)); theoretical mean = $(mean(nonmixed_multivariate_dist.gc_obs[i]))")
println("sample var = $(Statistics.var(s)); theoretical var = $(var(nonmixed_multivariate_dist.gc_obs[i]))")

# conditional mean and variance of Y2 given Y1
# sample mean = 5.005900458369502; theoretical mean = 5.002996732265867
# sample var = 0.04316391761735573; theoretical var = 0.04362288747145726

## now testing conditional Y3 given Y1, Y2
# simulate Y2 first
i = 2
# form constants for Y1 pdf
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])
# simulate Y1 and store it to Y[1]
Random.seed!(12345)
Y[i] = rand(nonmixed_multivariate_dist.gc_obs[i])
# update residuals res[2] using Y[2] and the mean and variance (lambda for Poisson) using GLM package
res[i] = update_res!(Y[i], res[i], nonmixed_multivariate_dist.gc_obs[i])

# next, for conditional density of Y2 | Y1, form the constants for pdf
i = 3
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])
# pre allocate for sampling nsample times from conditional density
Random.seed!(12345)
nsample = 10_000
@info "sample $nsample points for the conditional Normal distribution"
s = Vector{Float64}(undef, nsample)
Random.seed!(12345)
rand!(nonmixed_multivariate_dist.gc_obs[i], s) # compile
Random.seed!(12345)
@time rand!(nonmixed_multivariate_dist.gc_obs[i], s) # get time
# 0.040553 seconds (1 allocation: 32 bytes)
println("sample mean = $(Statistics.mean(s)); theoretical mean = $(mean(nonmixed_multivariate_dist.gc_obs[i]))")
println("sample var = $(Statistics.var(s)); theoretical var = $(var(nonmixed_multivariate_dist.gc_obs[i]))")
# conditional mean and variance of Y3 given Y1, Y2
# sample mean = 5.006034080584214; theoretical mean = 5.003125498930692
# sample var = 0.04331760554438279; theoretical var = 0.043778156954903835
# rand(mvn_dist, Y, res)
end

### Poisson only
@testset "Generate a 3 element random vector each with Y1, Y2, Y3 ~ Poisson(5). First we check the conditional mean and variance of Y2 | Y1 and then for Y3 | Y1, Y2. " begin
# variance components
Random.seed!(12345)
# simulate desired variance components (coefficients of covariance matrices V[k] for k in 1:m)
mean1 = 5
n = 3
d1 = Poisson(mean1)
d2 = Poisson(mean1)
d3 = Poisson(mean1)

vecd = [d1, d2, d3]
random_intercept_1 = 0.1
Γ = random_intercept_1 * ones(n, n)

nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)

#### Testing if the second element in the bivariate poisson follows correct mean and variance
# preallocate Y and res to simulate and store
Y = Vector{Float64}(undef, n)
res = Vector{Float64}(undef, n)

# simulate Y1 first
i = 1
# form constants for Y1 pdf
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])
# simulate Y1 and store it to Y[1]
Random.seed!(12345)
Y[i] = rand(nonmixed_multivariate_dist.gc_obs[i])
# update residuals res[1] using Y[1] and the mean and variance (lambda for Poisson) using GLM package
res[i] = update_res!(Y[i], res[i], nonmixed_multivariate_dist.gc_obs[i])

# next, for conditional density of Y2 | Y1, form the constants for pdf
i = 2
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])

# pre allocate for sampling nsample times from conditional density
Random.seed!(12345)
nsample = 10_000
@info "sample $nsample points for the conditional Poisson distribution"
s = Vector{Int64}(undef, nsample)
Random.seed!(12345)
rand!(nonmixed_multivariate_dist.gc_obs[i], s) # compile
Random.seed!(12345)
@time rand!(nonmixed_multivariate_dist.gc_obs[i], s) # get time
# 0.019684 seconds (50.00 k allocations: 12.207 MiB)
println("sample mean = $(Statistics.mean(s)); theoretical mean = $(mean(nonmixed_multivariate_dist.gc_obs[i]))")
println("sample var = $(Statistics.var(s)); theoretical var = $(var(nonmixed_multivariate_dist.gc_obs[i]))")
# conditional mean and variance of Y2 given Y1
# sample mean = 5.2485; theoretical mean = 5.219298245614036
# sample var = 5.689716721672168; theoretical var = 5.6098030163127035

## now testing conditional Y3 given Y1, Y2
# simulate Y2 first
i = 2
# form constants for Y1 pdf
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])
# simulate Y1 and store it to Y[1]
Random.seed!(12345)
Y[i] = rand(nonmixed_multivariate_dist.gc_obs[i])
# update residuals res[2] using Y[2] and the mean and variance (lambda for Poisson) using GLM package
res[i] = update_res!(Y[i], res[i], nonmixed_multivariate_dist.gc_obs[i])

# next, for conditional density of Y2 | Y1, form the constants for pdf
i = 3
nonmixed_multivariate_dist.gc_obs[i] = pdf_constants(nonmixed_multivariate_dist.Γ, res, i, nonmixed_multivariate_dist.vecd[i])
# pre allocate for sampling nsample times from conditional density
Random.seed!(12345)
nsample = 10_000
@info "sample $nsample points for the conditional Poisson distribution"
s = Vector{Int64}(undef, nsample)
Random.seed!(12345)
rand!(nonmixed_multivariate_dist.gc_obs[i], s) # compile
Random.seed!(12345)
@time rand!(nonmixed_multivariate_dist.gc_obs[i], s) # get time
# 0.024654 seconds (50.00 k allocations: 12.207 MiB)
println("sample mean = $(Statistics.mean(s)); theoretical mean = $(mean(nonmixed_multivariate_dist.gc_obs[i]))")
println("sample var = $(Statistics.var(s)); theoretical var = $(var(nonmixed_multivariate_dist.gc_obs[i]))")
# conditional mean and variance of Y3 given Y1, Y2
# sample mean = 5.2381; theoretical mean = 5.206611570247935
# sample var = 5.658774267426744; theoretical var = 5.577146369783474
end

@testset "Form mixtures of different discrete and cts distributions. Generate 3 element vector: Y1 ~ Normal(5, 0.2), Y2 ~ Exponential(1/5), Y3 ~ Poisson(5) " begin
Random.seed!(12345)
mean_normal = -5
var_normal = 0.2

d1 = Normal(mean_normal, var_normal)
d2 = Exponential(-1/mean_normal)
d3 = Poisson(exp(mean_normal))

vc = 0.1
Γ = vc * ones(3, 3)

vecd = [d1, d2, d3]
mixed_multivariate_dist = MultivariateMix(vecd, Γ)

# Y = Vector{Float64}(undef, n)
# res = Vector{Float64}(undef, n)
# rand(mixed_multivariate_dist, Y, res)
#### Testing if the second element in the bivariate poisson follows correct mean and variance
# preallocate Y and res to simulate and store
n = 3
Y = Vector{Float64}(undef, n)
res = Vector{Float64}(undef, n)

# simulate Y1 first
i = 1
# form constants for Y1 pdf
mixed_multivariate_dist.gc_obs[i] = pdf_constants(mixed_multivariate_dist.Γ, res, i, mixed_multivariate_dist.vecd[i])
# simulate Y1 and store it to Y[1]
Random.seed!(12345)
Y[i] = rand(mixed_multivariate_dist.gc_obs[i])
# update residuals res[1] using Y[1] and the mean and variance (lambda for Poisson) using GLM package
res[i] = update_res!(Y[i], res[i], mixed_multivariate_dist.gc_obs[i])

# next, for conditional density of Y2 | Y1, form the constants for pdf
i = 2
mixed_multivariate_dist.gc_obs[i] = pdf_constants(mixed_multivariate_dist.Γ, res, i, mixed_multivariate_dist.vecd[i])

# pre allocate for sampling nsample times from conditional density
Random.seed!(12345)
nsample = 10_000
@info "sample $nsample points for the conditional Exponential distribution"
s = Vector{Float64}(undef, nsample)
Random.seed!(12345)
rand!(mixed_multivariate_dist.gc_obs[i], s) # compile
Random.seed!(12345)
@time rand!(mixed_multivariate_dist.gc_obs[i], s) # get time
# 0.030101 seconds (50.00 k allocations: 12.207 MiB)
println("sample mean = $(Statistics.mean(s)); theoretical mean = $(mean(mixed_multivariate_dist.gc_obs[i]))")
println("sample var = $(Statistics.var(s)); theoretical var = $(var(mixed_multivariate_dist.gc_obs[i]))")
# conditional mean and variance of Y2 given Y1
# sample mean = 0.2235739008006706; theoretical mean = 0.22115607164570328
# sample var = 0.056128853458491966; theoretical var = 0.05527858504273324

## now testing conditional Y3 given Y1, Y2
# simulate Y2 first
i = 2
# form constants for Y1 pdf
mixed_multivariate_dist.gc_obs[i] = pdf_constants(mixed_multivariate_dist.Γ, res, i, mixed_multivariate_dist.vecd[i])
# simulate Y2 and store it to Y[2]
Random.seed!(12345)
Y[i] = rand(mixed_multivariate_dist.gc_obs[i])
# update residuals res[2] using Y[2] and the mean and variance (1/5 for Exponential) using GLM package
res[i] = update_res!(Y[i], res[i], mixed_multivariate_dist.gc_obs[i])

# next, for conditional density of Y2 | Y1, form the constants for pdf
i = 3
mixed_multivariate_dist.gc_obs[i] = pdf_constants(mixed_multivariate_dist.Γ, res, i, mixed_multivariate_dist.vecd[i])
# pre allocate for sampling nsample times from conditional density
Random.seed!(12345)
nsample = 10_000
@info "sample $nsample points for the conditional Poisson distribution"
s = Vector{Int64}(undef, nsample)
Random.seed!(12345)
rand!(mixed_multivariate_dist.gc_obs[i], s) # compile
Random.seed!(12345)
@time rand!(mixed_multivariate_dist.gc_obs[i], s) # get time
# 0.024654 seconds (50.00 k allocations: 12.207 MiB)
println("sample mean = $(Statistics.mean(s)); theoretical mean = $(mean(mixed_multivariate_dist.gc_obs[i]))")
println("sample var = $(Statistics.var(s)); theoretical var = $(var(mixed_multivariate_dist.gc_obs[i]))")
# conditional mean and variance of Y3 given Y1, Y2
# sample mean = 0.0562; theoretical mean = 0.05564357782888464
# sample var = 0.053646924692469226; theoretical var = 0.0538934802364384
#####

### bivariate Poisson correlation and covariance
mean1 = 5
mean2 = 5
d1 = Poisson(mean1)
d2 = Poisson(mean2)

vecd = [d1, d2]
Γ1 = 0.1 * ones(2, 2)
Γ2 = [0.1 -0.1; -0.1 0.1]

gc_vec1 = NonMixedMultivariateDistribution(vecd, Γ1)
gc_vec2 = NonMixedMultivariateDistribution(vecd, Γ2)

nsample = 10_000
@info "sample $nsample independent vectors for the bivariate Poisson distribution"
# compile
Y_nsample1 = simulate_nobs_independent_vectors(gc_vec1, nsample)
Y_nsample2 = simulate_nobs_independent_vectors(gc_vec2, nsample)

Y_11 = [Y_nsample1[i, 1][1] for i in 1:nsample]
Y_12 = [Y_nsample1[i, 1][2] for i in 1:nsample]
Statistics.cov(Y_11, Y_12)

@show Statistics.cor(Y_11, Y_12)

Y_21 = [Y_nsample2[i, 1][1] for i in 1:nsample]
Y_22 = [Y_nsample2[i, 1][2] for i in 1:nsample]
Statistics.cov(Y_21, Y_22)

@show Statistics.cor(Y_21, Y_22)

@test QuasiCopula.cov(gc_vec1, 1, 2) == -QuasiCopula.cov(gc_vec2, 1, 2)
@test QuasiCopula.cor(gc_vec1, 1, 2) == -QuasiCopula.cor(gc_vec2, 1, 2)

###
# bivariate mixed normal and poisson with theoretical correlation ~
##
mean_normal = 5
var_normal = 0.2

mean_poisson = 7

d1 = Normal(mean_normal, var_normal)
d2 = Poisson(mean_poisson)

vc = 0.1
Γ1 = [0.1 0.1; 0.1 0.1]
Γ2 = [0.1 -0.1; -0.1 0.1]

vecd = [d1, d2]
gc_vec1 = MultivariateMix(vecd, Γ1)
gc_vec2 = MultivariateMix(vecd, Γ2)

nsample = 10_000
@info "sample $nsample independent vectors for the normal and poisson"
# compile
Y_nsample1 = simulate_nobs_independent_vectors(gc_vec1, nsample)
Y_nsample2 = simulate_nobs_independent_vectors(gc_vec2, nsample)

Y_11 = [Y_nsample1[i, 1][1] for i in 1:nsample]
Y_12 = [Y_nsample1[i, 1][2] for i in 1:nsample]
Statistics.cov(Y_11, Y_12)

@show Statistics.cor(Y_11, Y_12)
QuasiCopula.cov(gc_vec1, 1, 2)
QuasiCopula.cor(gc_vec1, 1, 2)
println("sample cor = $(Statistics.cor(Y_11, Y_12)); theoretical cor = $(QuasiCopula.cor(gc_vec1, 1, 2))")
println("sample cov = $(Statistics.cov(Y_11, Y_12)); theoretical var = $(QuasiCopula.cov(gc_vec1, 1, 2))")

Y_21 = [Y_nsample2[i, 1][1] for i in 1:nsample]
Y_22 = [Y_nsample2[i, 1][2] for i in 1:nsample]
Statistics.cov(Y_21, Y_22)

@show Statistics.cor(Y_21, Y_22)

QuasiCopula.cov(gc_vec2, 1, 2)
QuasiCopula.cor(gc_vec2, 1, 2)
println("sample cor = $(Statistics.cor(Y_21, Y_22)); theoretical cor = $(QuasiCopula.cor(gc_vec2, 1, 2))")
println("sample cov = $(Statistics.cov(Y_21, Y_22)); theoretical var = $(QuasiCopula.cov(gc_vec2, 1, 2))")

@test QuasiCopula.cov(gc_vec1, 1, 2) == -QuasiCopula.cov(gc_vec2, 1, 2)
@test QuasiCopula.cor(gc_vec1, 1, 2) == -QuasiCopula.cor(gc_vec2, 1, 2)
####
# mixed with normal and 0, 1
mean_normal = 5
var_normal = 0.2

mean_bernoulli = 0.5

d1 = Normal(mean_normal, var_normal)
d2 = Bernoulli(mean_bernoulli)

vc = 0.1
Γ1 = [0.1 0.1; 0.1 0.1]
Γ2 = [0.1 -0.1; -0.1 0.1]

vecd = [d1, d2]
gc_vec1 = MultivariateMix(vecd, Γ1)
gc_vec2 = MultivariateMix(vecd, Γ2)

nsample = 10_000
@info "sample $nsample independent vectors for the mixed normal and binary bivariate"
# compile
Y_nsample1 = simulate_nobs_independent_vectors(gc_vec1, nsample)
Y_nsample2 = simulate_nobs_independent_vectors(gc_vec2, nsample)

Y_11 = [Y_nsample1[i, 1][1] for i in 1:nsample]
Y_12 = [Y_nsample1[i, 1][2] for i in 1:nsample]
# @show Statistics.cov(Y_11, Y_12)

# @show Statistics.cor(Y_11, Y_12)
# @show QuasiCopula.cov(gc_vec1, 1, 2)
# @show QuasiCopula.cor(gc_vec1, 1, 2)
println("sample cor = $(Statistics.cor(Y_11, Y_12)); theoretical cor = $(QuasiCopula.cor(gc_vec1, 1, 2))")
println("sample cov = $(Statistics.cov(Y_11, Y_12)); theoretical var = $(QuasiCopula.cov(gc_vec1, 1, 2))")

Y_21 = [Y_nsample2[i, 1][1] for i in 1:nsample]
Y_22 = [Y_nsample2[i, 1][2] for i in 1:nsample]
# @show Statistics.cov(Y_21, Y_22)

# @show Statistics.cor(Y_21, Y_22)

# @show QuasiCopula.cov(gc_vec2, 1, 2)
# @show QuasiCopula.cor(gc_vec2, 1, 2)
println("sample cor = $(Statistics.cor(Y_21, Y_22)); theoretical cor = $(QuasiCopula.cor(gc_vec2, 1, 2))")
println("sample cov = $(Statistics.cov(Y_21, Y_22)); theoretical var = $(QuasiCopula.cov(gc_vec2, 1, 2))")
#

@test QuasiCopula.cov(gc_vec1, 1, 2) == -QuasiCopula.cov(gc_vec2, 1, 2)
@test QuasiCopula.cor(gc_vec1, 1, 2) == -QuasiCopula.cor(gc_vec2, 1, 2)

# mixed with normal and 0, 1 with more noise in normal
mean_normal = 5
sd_normal = 0.2

mean_bernoulli = 0.5

d1 = Bernoulli(mean_bernoulli)
d2 = Normal(mean_normal, sd_normal)

vc = 0.1
Γ1 = [0.1 0.1; 0.1 0.1]
Γ2 = [0.1 -0.1; -0.1 0.1]

vecd = [d1, d2]
gc_vec1 = MultivariateMix(vecd, Γ1)
gc_vec2 = MultivariateMix(vecd, Γ2)

nsample = 10_000
@info "sample $nsample independent vectors for the mixed normal and binary bivariate"
# compile
Y_nsample1 = simulate_nobs_independent_vectors(gc_vec1, nsample)
Y_nsample2 = simulate_nobs_independent_vectors(gc_vec2, nsample)

Y_11 = [Y_nsample1[i, 1][1] for i in 1:nsample]
Y_12 = [Y_nsample1[i, 1][2] for i in 1:nsample]
# @show Statistics.cov(Y_11, Y_12)

# @show Statistics.cor(Y_11, Y_12)
# @show QuasiCopula.cov(gc_vec1, 1, 2)
# @show QuasiCopula.cor(gc_vec1, 1, 2)
println("sample cor = $(Statistics.cor(Y_11, Y_12)); theoretical cor = $(QuasiCopula.cor(gc_vec1, 1, 2))")
println("sample cov = $(Statistics.cov(Y_11, Y_12)); theoretical cov = $(QuasiCopula.cov(gc_vec1, 1, 2))")

Y_21 = [Y_nsample2[i, 1][1] for i in 1:nsample]
Y_22 = [Y_nsample2[i, 1][2] for i in 1:nsample]
# @show Statistics.cov(Y_21, Y_22)

# @show Statistics.cor(Y_21, Y_22)

# @show QuasiCopula.cov(gc_vec2, 1, 2)
# @show QuasiCopula.cor(gc_vec2, 1, 2)
println("sample cor = $(Statistics.cor(Y_21, Y_22)); theoretical cor = $(QuasiCopula.cor(gc_vec2, 1, 2))")
println("sample cov = $(Statistics.cov(Y_21, Y_22)); theoretical cov = $(QuasiCopula.cov(gc_vec2, 1, 2))")
#

@test QuasiCopula.cov(gc_vec1, 1, 2) == -QuasiCopula.cov(gc_vec2, 1, 2)
@test QuasiCopula.cor(gc_vec1, 1, 2) == -QuasiCopula.cor(gc_vec2, 1, 2)

end


@inline function two_sum(a::T, b::T) where T<:AbstractFloat
    ab =  a + b
    B  = ab - a
    ε  = (a - (ab - B)) + (b - B)
    return ab, ε
end

@inline function two_diff(a::T, b::T) where T<:AbstractFloat
    ab =  a - b
    B  = ab - a
    ε  = (a - (ab - B)) + (b - B)
    return ab, ε
end

"""
    fast_two_sum(x, y)::(a, b)

{ ab  =  a + b,  ε  = (a - ab) + b }

__unchecked precondition__: abs(a) >= abs(b).
"""
@inline function fast_two_sum(a::T, b::T) where T<:AbstractFloat
     ab =  a + b
     B  = ab - a
     ε  =  b - B
     return ab, ε
end

"""
    fast_two_diff(x, y)::(a, b)

{ ab  =  a - b,  ε  = (a - ab) + b }

__unchecked precondition__: abs(a) >= abs(b).
"""
@inline function fast_two_diff(a::T, b::T) where T<:AbstractFloat
     ab =  a - b
     B  = ab - a
     ε  =  B - b
     return ab, ε
end

@inline function two_prod(a::T, b::T) where T<:AbstractFloat
    ab =  a * b
    ε  = fma(a, b, -ab)
    return ab, ε
end

#=
   Separate an IEEEFloat (fp) into two parts, (hi, lo), assuring
      `fp === hi + lo`, `hi === fp - lo`, `lo === fp - hi`, and
      `ulp(abs(hi)) > ufp(abs(lo))`.

      ulp, ufp are acronyms for u[nit in the] {l[ast], f[irst]} p[lace]

      ulp(fp) = eps(fp)/2
      ufp(fp) = ldexp(0.5, exponent(fp)

   ref
       A floating-point technique for extending the available precision.
       T. J. Dekker.
       Numer. Math., 18:224–242, 1971
=#

splitter(::Type{Float64}) =
    Float64(one(Int64) + one(Int64) << (cld(precision(Float64),2)-1))
splitter(::Type{Float32}) =
    Float32(one(Int32) + one(Int32) << (cld(precision(Float32),2)-1))
splitter(::Type{Float16}) =
    Float16(one(Int16) + one(Int16) << (cld(precision(Float16),2)-1))

splitmax(::Type{Float64}) = realmax(Float64) / splitter(Float64)
splitmax(::Type{Float32}) = realmax(Float32) / splitter(Float32)
splitmax(::Type{Float16}) = realmax(Float16) / splitter(Float16)

# Veldkamp splitting of a floating point value
@inline function splitting(x::T) where T<:AbstractFloat
    (!isfinite(x) || abs(x) > splitmax(T)) && throw(OverflowError("$x overflows"))
    z   = x * splitter(T)
    zₕᵢ = z - (z - x)
    zₗₒ = x - zₕᵢ
    return zₕᵢ, zₗₒ
end

# Rump splitting
#    In extractscalar, a floating-point number fp is split relative to p2,
#    a fixed power of 2.
function extractscalar(fp, p2=2^27)
    hi = (p2 + fp) - p2
    lo = fp - hi
    return hi, lo
end

#=
    splitfp(fp) =?= extractscalar(2^(?+exponent(fp)), fp)

=#

for (U,F) in ((:UInt64, :Float64), (:UInt32, :Float32), (:UInt16, :Float16))
  @eval begin
    @inline function ufp(x::$F)
        u = reinterpret($U, x)
        u = (u >> (precision($F)-1)) << (precision($F)-1)
        return reinterpret($F, u)
    end
  end
end

const Float64ulp = inv(ldexp(1.0, precision(Float64)))
const Float32ulp = inv(ldexp(1.0, precision(Float32)))
const Float16ulp = inv(ldexp(1.0, precision(Float16)))

@inline ulp(x::Float64) = ufp(x) * Float64ulp
@inline ulp(x::Float32) = ufp(x) * Float32ulp
@inline ulp(x::Float16) = ufp(x) * Float16ulp

const Float64eps = inv(ldexp(1.0, precision(Float64)-1))
const Float32eps = inv(ldexp(1.0, precision(Float32)-1))
const Float16eps = inv(ldexp(1.0, precision(Float16)-1))

@inline epsi(x::Float64) = ufp(x) * Float64eps
@inline epsi(x::Float32) = ufp(x) * Float32eps
@inline epsi(x::Float16) = ufp(x) * Float16eps


@testset "partial" begin

@testset "create_∂ and apply_∂!, 1D" begin
    N = (10,)
    M = prod(N)

    Fu = rand(N...)
    Gv = similar(Fu)
    gv = zeros(M)

    Mask = similar(Fu)  # masking array
    for nw = 1:1
        sub = Vector{Int}(undef, 1)
        sub′ = Vector{Int}(undef, 1)

        Nw = N[nw]
        ∆w⁻¹ = rand(Nw)

        for ns = (-1,1), isbloch = (true,false)  # (backward,forward difference), (Bloch,symmetric)
            ∂ws = spzeros(M,M)

            for ind = 1:M
                sub .= CartesianIndices(N)[ind].I  # subscripts of diagonal entry
                indw = sub[nw]
                ∆wᵢ⁻¹ = ∆w⁻¹[indw]
                ∂ws[ind,ind] = -ns * ∆wᵢ⁻¹  # diagonal entries

                # Calculate the column index of the off-diagonal entry in the row `ind`.
                sub′ .= sub  # subscripts of off-diagonal entry
                if ns == 1  # forward difference
                    if sub′[nw] == Nw
                        sub′[nw] = 1
                    else
                        sub′[nw] += 1
                    end
                else  # backward difference
                    if sub′[nw] == 1
                        sub′[nw] = Nw
                    else
                        sub′[nw] -= 1
                    end
                end

                ind′ = LinearIndices(N)[sub′...]
                ∂ws[ind, ind′] += ns * ∆wᵢ⁻¹  # off-diagonal entry
            end

            if !isbloch  # symmetry boundary
                # Initialize the masking array.
                Mask .= 1
                Mask[Base.setindex(axes(Mask), 1, nw)...] = 0  # not .= 0 unlike code for 2D and 3D
                if ns == 1  # forward difference
                    # The input fields at the symmetry boundary should be zero, so apply the
                    # mask to the input field.
                    ∂ws = ∂ws * spdiagm(0=>Mask[:])
                else  # backward difference
                    # The output fields at the symmetry boundary should be zero, so apply
                    # mask to the output field.
                    ∂ws = spdiagm(0=>Mask[:]) * ∂ws
                end
            end

            # Test create_∂.
            @test create_∂(nw, ns==1, [N...], ∆w⁻¹, isbloch) == ∂ws

            # Test apply_∂!.
            fu = Fu[:]
            mul!(gv, ∂ws, fu)
            apply_∂!(Gv, Fu, Val(:(=)), nw, ns==1, ∆w⁻¹, isbloch)
            @test Gv[:] ≈ gv

            # print("matrix: "); @btime mul!($gv, $∂ws, $fu)
            # print("matrix-free: "); @btime apply_∂!($Gv, $Fu, $nw, $ns==1, $∆w⁻¹)
            # println()
        end
    end
end  # @testset "create_∂ and apply_∂!, 1D"

@testset "create_∂ and apply_∂!, 2D" begin
    N = (8,9)
    M = prod(N)

    Fu = rand(N...)
    Gv = similar(Fu)
    gv = zeros(M)

    Mask = similar(Fu)  # masking array
    for nw = 1:2
        sub = Vector{Int}(undef, 2)
        sub′ = Vector{Int}(undef, 2)

        Nw = N[nw]
        ∆w⁻¹ = rand(Nw)

        for ns = (-1,1), isbloch = (true,false)  # (backward,forward difference), (Bloch,symmetric)
            ∂ws = spzeros(M,M)

            for ind = 1:M
                sub .= CartesianIndices(N)[ind].I  # subscripts of diagonal entry
                indw = sub[nw]
                ∆wᵢ⁻¹ = ∆w⁻¹[indw]
                ∂ws[ind,ind] = -ns * ∆wᵢ⁻¹  # diagonal entries

                # Calculate the column index of the off-diagonal entry in the row `ind`.
                sub′ .= sub  # subscripts of off-diagonal entry
                if ns == 1  # forward difference
                    if sub′[nw] == Nw
                        sub′[nw] = 1
                    else
                        sub′[nw] += 1
                    end
                else  # backward difference
                    if sub′[nw] == 1
                        sub′[nw] = Nw
                    else
                        sub′[nw] -= 1
                    end
                end

                ind′ = LinearIndices(N)[sub′...]
                ∂ws[ind, ind′] += ns * ∆wᵢ⁻¹  # off-diagonal entry
            end

            if !isbloch  # symmetry boundary
                # Initialize the masking array.
                Mask .= 1
                Mask[Base.setindex(axes(Mask), 1, nw)...] .= 0
                if ns == 1  # forward difference
                    # The input fields at the symmetry boundary should be zero, so apply the
                    # mask to the input field.
                    ∂ws = ∂ws * spdiagm(0=>Mask[:])
                else  # backward difference
                    # The output fields at the symmetry boundary should be zero, so apply
                    # mask to the output field.
                    ∂ws = spdiagm(0=>Mask[:]) * ∂ws
                end
            end

            # Test create_∂.
            @test create_∂(nw, ns==1, [N...], ∆w⁻¹, isbloch) == ∂ws

            # Test apply_∂!.
            fu = Fu[:]
            mul!(gv, ∂ws, fu)
            apply_∂!(Gv, Fu, Val(:(=)), nw, ns==1, ∆w⁻¹, isbloch)
            @test Gv[:] ≈ gv

            # print("matrix: "); @btime mul!($gv, $∂ws, $fu)
            # print("matrix-free: "); @btime apply_∂!($Gv, $Fu, $nw, $ns==1, $∆w⁻¹)
            # println()
        end
    end
end  # @testset "create_∂ and apply_∂!, 2D"

@testset "create_∂ and apply_∂!, 3D" begin
    N = (8,9,10)
    M = prod(N)

    Fu = rand(N...)
    Gv = similar(Fu)
    gv = zeros(M)

    Mask = similar(Fu)  # masking array
    for nw = 1:3
        sub = Vector{Int}(undef, 3)
        sub′ = Vector{Int}(undef, 3)

        Nw = N[nw]
        ∆w⁻¹ = rand(Nw)

        for ns = (-1,1), isbloch = (true,false)  # (backward,forward difference), (Bloch,symmetric)
            ∂ws = spzeros(M,M)

            for ind = 1:M
                sub .= CartesianIndices(N)[ind].I  # subscripts of diagonal entry
                indw = sub[nw]
                ∆wᵢ⁻¹ = ∆w⁻¹[indw]
                ∂ws[ind,ind] = -ns * ∆wᵢ⁻¹  # diagonal entries

                # Calculate the column index of the off-diagonal entry in the row `ind`.
                sub′ .= sub  # subscripts of off-diagonal entry
                if ns == 1  # forward difference
                    if sub′[nw] == Nw
                        sub′[nw] = 1
                    else
                        sub′[nw] += 1
                    end
                else  # backward difference
                    if sub′[nw] == 1
                        sub′[nw] = Nw
                    else
                        sub′[nw] -= 1
                    end
                end

                ind′ = LinearIndices(N)[sub′...]
                ∂ws[ind, ind′] += ns * ∆wᵢ⁻¹  # off-diagonal entry
            end

            if !isbloch  # symmetry boundary
                # Initialize the masking array.
                Mask .= 1
                Mask[Base.setindex(axes(Mask), 1, nw)...] .= 0
                if ns == 1  # forward difference
                    # The input fields at the symmetry boundary should be zero, so apply the
                    # mask to the input field.
                    ∂ws = ∂ws * spdiagm(0=>Mask[:])
                else  # backward difference
                    # The output fields at the symmetry boundary should be zero, so apply
                    # mask to the output field.
                    ∂ws = spdiagm(0=>Mask[:]) * ∂ws
                end
            end

            # Test create_∂.
            @test create_∂(nw, ns==1, [N...], ∆w⁻¹, isbloch) == ∂ws

            # Test apply_∂!.
            fu = Fu[:]
            mul!(gv, ∂ws, fu)
            apply_∂!(Gv, Fu, Val(:(=)), nw, ns==1, ∆w⁻¹, isbloch)
            @test Gv[:] ≈ gv

            # print("matrix: "); @btime mul!($gv, $∂ws, $fu)
            # print("matrix-free: "); @btime apply_∂!($Gv, $Fu, $nw, $ns==1, $∆w⁻¹)
            # println()
        end
    end
end  # @testset "create_∂ and apply_∂!, 3D"

end  # @testset "differential"


function fastModularExponentiation(m, e, n::Integer)
    n > 0  || error("n must be positive")
    e >= 0 || error("e must be non-negative")
    r = 1
    while e > 0
        if (e & 1) > 0
            r = (r*m)%n
        end
        e >>= 1
        m = (m^2)%n
    end
    return r>0 ? r : r+n
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# https://twitter.com/genkuroki/status/1459220790444920833

# %%
using Distributions, StatsPlots

# %%
mixnormal = MixtureModel([Normal(0, 1), Normal(1, 2)], [1/3, 2/3])

# %%
X = rand(mixnormal, 10^5)
f(x) = (1/3)*pdf(Normal(0, 1), x) + (2/3)*pdf(Normal(1, 2), x)

histogram(X; norm=true, alpha=0.3, label="sample of mixnormal")
plot!(f; label="(1/3)N(x|0,1)+(2/3)N(x|1,2))", lw=2)

# %%


# Quadratic Costs
n,m = 3,2
Q = Diagonal(@SVector fill(0.1, n))
R = Diagonal(@SVector fill(0.01, m))
Qf = Diagonal(@SVector fill(10., n))
xf = @SVector ones(n)

quadcost = QuadraticCost(Q,R)
@test quadcost.Q == Q
@test quadcost.q == zeros(n)
@test quadcost.R == R
@test quadcost.r == zeros(m)
@test quadcost.c == 0
@test state_dim(quadcost) == n
@test control_dim(quadcost) == m

quadcost = LQRCost(Q,R,xf)
@test quadcost.Q == Q
@test quadcost.q == -Q*xf
@test quadcost.R == R
@test quadcost.r == zeros(m)
@test quadcost.c == 0.5*xf'Q*xf

quadcost = TO.LQRCostTerminal(Q,xf)
@test quadcost.Q == Q
@test quadcost.q == -Q*xf
@test isempty(quadcost.R)
@test isempty(quadcost.r)
@test quadcost.c == 0.5*xf'Q*xf

q = @SVector rand(n)
quadcost = QuadraticCost(Q,q,0.5)
@test quadcost.Q == Q
@test quadcost.q == q
@test isempty(quadcost.R)
@test isempty(quadcost.r)
@test quadcost.c == 0.5

Q2 = Diagonal(@SVector [-1, 0.1, 0.2])
R2 = Diagonal(@SVector [0, 0.1])
@test_throws ArgumentError QuadraticCost(Q2,R)
@test_logs (:warn, "R is not positive definite") QuadraticCost(Q,R2)


using DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions,
      RecursiveArrayTools, Distributions, Test

println("One parameter case")
f1 = @ode_def begin
  dx = a*x - x*y
  dy = -3y + x*y
end a
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=50))
randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(1.5,0.1),1.0,1.8)]

bayesian_result = stan_inference(prob1,t,data,priors;num_samples=300,
                                 num_warmups=500,likelihood=Normal)

@test mean(get(bayesian_result.chains,:theta_1)[1]) ≈ 1.5 atol=3e-1

# Test norecompile
bayesian_result2 = stan_inference(prob1,t,data,priors,bayesian_result.model;
                                  num_samples=300,num_warmups=500,likelihood=Normal)

@test mean(get(bayesian_result2.chains,:theta_1)[1]) ≈ 1.5 atol=3e-1

priors = [truncated(Normal(1.,0.01),0.5,2.0),truncated(Normal(1.,0.01),0.5,2.0),truncated(Normal(1.5,0.01),1.0,2.0)]
bayesian_result = stan_inference(prob1,t,data,priors;num_samples=300,
                                 num_warmups=500,likelihood=Normal,sample_u0=true)

@test mean(get(bayesian_result.chains,:theta_1)[1]) ≈ 1. atol=3e-1
@test mean(get(bayesian_result.chains,:theta_2)[1]) ≈ 1. atol=3e-1
@test mean(get(bayesian_result.chains,:theta_3)[1]) ≈ 1.5 atol=3e-1

sol = solve(prob1,Tsit5(),save_idxs=[1])
randomized = VectorOfArray([(sol(t[i]) + .01 * randn(1)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(1.5,0.1),0.5,2)]
bayesian_result = stan_inference(prob1,t,data,priors;num_samples=300,
                                 num_warmups=500,likelihood=Normal,save_idxs=[1])

@test mean(get(bayesian_result.chains,:theta_1)[1]) ≈ 1.5 atol=3e-1

priors = [truncated(Normal(1.,0.01),0.5,2),truncated(Normal(1.5,0.01),0.5,2)]
bayesian_result = stan_inference(prob1,t,data,priors;num_samples=300,
                                 num_warmups=500,likelihood=Normal,save_idxs=[1],sample_u0=true)

@test mean(get(bayesian_result.chains,:theta_1)[1]) ≈ 1. atol=3e-1
@test mean(get(bayesian_result.chains,:theta_2)[1]) ≈ 1.5 atol=3e-1

println("Four parameter case")
f1 = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=50))
randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(1.5,0.01),0.5,2),truncated(Normal(1.0,0.01),0.5,1.5),
          truncated(Normal(3.0,0.01),0.5,4),truncated(Normal(1.0,0.01),0.5,2)]

bayesian_result = stan_inference(prob1,t,data,priors;num_samples=100,num_warmups=500,vars =(DiffEqBayes.StanODEData(),InverseGamma(4,1)))
@test mean(get(bayesian_result.chains,:theta_1)[1]) ≈ 1.5 atol=1e-1
@test mean(get(bayesian_result.chains,:theta_2)[1]) ≈ 1.0 atol=1e-1
@test mean(get(bayesian_result.chains,:theta_3)[1]) ≈ 3.0 atol=1e-1
@test mean(get(bayesian_result.chains,:theta_4)[1]) ≈ 1.0 atol=1e-1


"""
    eom_var1(x, p, n) -> Function

Equations of motion for a vector autoregressive system where X₁ → X₂ → X₃.
"""
function eom_var1(x, p, n)
    σ₁, σ₂, σ₃ = p[1], p[2], p[3]
    x₁, x₂, x₃ = x[1], x[2], x[3]
    θ = rand(Normal(0, σ₁))
    η = rand(Normal(0, σ₂))
    ϵ = rand(Normal(0, σ₃))

    dx₁ = θ
    dx₂ = x₁ * η
    dx₃ = 0.5*x₃ * x₂ + ϵ
    return SVector{3}(dx₁, dx₂, dx₃)
end

function var1(uᵢ, σ₁, σ₂, σ₃)
    p = [σ₁, σ₂, σ₃]
    DiscreteDynamicalSystem(eom_var1, uᵢ, p)
end

"""
    var1(x, p, n) -> DiscreteDynamicalSystem

Initialise a discrete vector autoregressive system where X₁ → X₂ → X₃.
"""
var1(;uᵢ = rand(3), σ₁ = 1.0, σ₂ = 0.2, σ₃ = 0.3) = var1(uᵢ, σ₁, σ₂, σ₃)


# Correlation Functions
# #+HTML: Source
# #+HTML:

# [[file:~/.julia/dev/MatrixProductStates/README.org::*Correlation%20Functions][Correlation Functions:1]]
"""
    two_point_correlator((i, op_i)::Pair{Int, Matrix}, (j, op_j)::Pair{Int, Matrix}, L)

Create an MPO on `L` sites (with bond dimension 1) representing identity operators everywhere except
sites `i` and `j` where `op_i` and `op_j` are inserted instead. ie.

    𝟙 ⊗ 𝟙 ⊗ ... ⊗ op_i ⊗ 𝟙 ⊗ ... ⊗ op_j ⊗ 𝟙 ⊗ ... ⊗ 𝟙

example: spin-spin correlation function

we can construct ⟨σᶻᵢσᶻⱼ⟩ on a 12 site lattice as
    σᶻ = [1 0; 0 -1]
    two_point_correlator(i=>σᶻ, j=>σᶻ, 12)
"""
function two_point_correlator((i, op_i), (j, op_j), L)
    d = size(op_i)[1]
    @assert (size(op_i) == (d, d)) && (size(op_j) == (d, d))
    @assert i in 1:L
    @assert j in 1:L
    id = diagm(0 => ones(Complex{Float64}, d))

    op_i_tnsr = reshape(convert(Matrix{Complex{Float64}}, op_i), 1, 1, d, d)
    op_j_tnsr = reshape(convert(Matrix{Complex{Float64}}, op_j), 1, 1, d, d)
    id_tnsr   = reshape(id, 1, 1, d, d)

    tensors = map(1:L) do l
        O_tnsr = (l == i ? op_i_tnsr :
                  l == j ? op_j_tnsr :
                  id_tnsr)
    end
    MPO{L,Complex{Float64}}(tensors)
end
# Correlation Functions:1 ends here


# ## Contour plot
# ### Contour plot

#md # ![FILE_NAME.png](images/FILE_NAME.png)

f(x,y) = x^2 + y^2
x = repeat(range(0, stop=2, length=50), 1, 50)
y = repeat(range(0, stop=2, length=50), 1, 50)'

contour( f.(x, y) )


using Random

# TODO: use Macros and Symbols

function divide_on_feature(x, idx, threshold)
    """
    Divide dataset based on if sample value on feature index is larger than
    the given threshold
    """

    x1, x2 = [], []

    for i in 1:size(x, 1)
        if x[i, idx] >= threshold
            push!(x1, x[i,:]')
        else
            push!(x2, x[i,:]')
        end
    end

    return vcat(x1...), vcat(x2...)

end

function sample(data::Array{T,2}, n::Int64; seed=123::Int64) where T <: Any
    # TODO: accept seed!
    n_samples = size(data, 1)

    # Shuffle data
    idx = shuffle(MersenneTwister(seed), 1:n_samples)
    data = data[idx, :]

    data[1:n, :]
end

function split_data(data, fraction)
    n_samples = size(data, 1)

    # Suffle data
    shuffle!(data)

    s = round(Integer, n_samples * fraction)

    return data[1:s, :], data[s:end, :]
end

# From https://bit.ly/2QnEacc

function entropy(y::AbstractArray)
    """ Calculate the entropy of label array y """
    unique_labels = unique(y)
    e = 0
    n = length(y)

    for label in unique_labels
        c = length(filter(x-> x==label, y))
        p = c / n
        e += -p * log2(p)
    end

    return e
end

function gini_index(y::AbstractArray)
    """ Calculate the Gini index of label array y """
    unique_labels = unique(y)
    e = 0
    n = length(y)

    for label in unique_labels
        c = length(filter(x-> x==label, y))
        p = c / n
        e += p^2
    end

    return 1 - e
end

function classification_error(y::AbstractArray)
    """ Calculate the Classification Error of label array y """
    unique_labels = unique(y)
    pp = []
    n = length(y)

    for label in unique_labels
        c = length(filter(x-> x==label, y))
        p = c / n
        push!(pp, p)
    end

    return 1 - max(pp)
end

function modes(values)
    # From https://rosettacode.org/wiki/Averages/Mode#Julia
    dict = Dict() # Values => Number of repetitions
    modesArray = typeof(values[1])[] # Array of the modes so far
    max = 0 # Max of repetitions so far

    for v in values
        # Add one to the dict[v] entry (create one if none)
        if v in keys(dict)
            dict[v] += 1
        else
            dict[v] = 1
        end

        # Update modesArray if the number of repetitions
        # of v reaches or surpasses the max value
        if dict[v] >= max
            if dict[v] > max
                empty!(modesArray)
                max += 1
            end
            append!(modesArray, [v])
        end
    end

    return modesArray
end

function count_distinct(seq)
    dict = Dict{Any, Int64}();
    for item in seq
        if !haskey(dict, item)
            dict[item] = 0;
        end
        dict[item] += 1
    end

    dict
end


function lin_interp(grid, vals)
    function func(x::Real)
        if x <= grid[1]
            index_1 = 1
            index_2 = 2
        elseif x >= grid[length(grid)]
            index_1 = length(grid) - 1
            index_2 = length(grid)
        else
            index_1 = searchsortedfirst(grid, x)
            index_2 = searchsortedlast(grid, x)
        end
        x_1 = grid[index_1]
        x_2 = grid[index_2]
        y_1 = vals[index_1]
        y_2 = vals[index_2]

        y = ((y_2- y_1)/(x_2 - x_1))*(x - x_1) + y_1
        return y
    end

    function func{T<:Real}(x::AbstractVector{T})
        y = zeros(length(x))
        for i in 1:length(x)
            y[i] = func(x[i])
        end
        return y
    end

    return func
end


# .................... cubic b-spline functions ...............................

# n / 0 == 0
#
function zero_divide(a, b)
    x = a / b
    return ifelse(isfinite(x), x, 0)
end

# uniform-knot p:degree n:number_of_points
#
function knot(p, n)
    base = [ifelse(k <= p, 0.0, 1.0) for k in (0:p + n)]
    if n - p == 1
        return base
    else
        for k in 0:n-p
            base[p + k + 1] = k / (n - p)
        end
    end
    return base
end

# b-spline basis p:degree i:number u:parameter k:knot
#
function bbasis(p, i, u, k)
    if p == 0
        return ifelse(k[i] <= u <= k[i+1], 1, 0)
    else
        m0 = zero_divide(u - k[i], k[i+p] - k[i])
        m1 = zero_divide(k[i+p+1] - u, k[i+p+1] - k[i+1])
        return m0 * bbasis(p-1, i, u, k) + m1 * bbasis(p-1, i+1, u, k)
    end
end

# b-spline basis first-derivative
#
function bbasis_fd(p, i, u, k)
    return zero_divide(p, k[i+p] - k[i]) * bbasis(p-1, i, u, k) -
        zero_divide(p, k[i+p+1] - k[i+1]) * bbasis(p-1, i+1, u, k)
end

# b-spline curve, points: as columns, t: 1d vector
#
function bcurve(points, u, bas=bbasis)
    knt = Ref(knot(3, size(points, 2)))
    tmp = zeros(size(points, 1), length(u), size(points, 2))
    for n in axes(points, 2)
        tmp[:,:,n] = points[:,n] .* bas.(3, n, u, knt)'
    end
    return dropdims(sum(tmp, dims=3), dims=3)
end

# b-spline surface, points: size(cp_u, cp_v, 3), uv: size(nu, nv, 2)
#
function bsurface(points, uv, bu=bbasis, bv=bbasis)
    ku, kv = knot.(3, size(points))
    uv_tmp = zeros(size(uv, 1), size(uv, 2), size(points, 3))

    for I in CartesianIndices(uv[:,:,1])
        for J in CartesianIndices(points[:,:,1])
            uv_tmp[I,:] += points[J,:] .*
                bu(3, J[1], uv[I, 1], ku) .*
                bv(3, J[2], uv[I, 2], kv)
        end
    end
    return uv_tmp
end

function mgrid(nu, nv, a=0:1, b=0:1)
    uv = zeros(nu, nv, 2)
    uv[:,:,1] .= first(a) : ((last(a)-first(a)) / (nu-1)) : last(a)
    uv[:,:,2] .= (first(b) : ((last(b)-first(b)) / (nv-1)) : last(b))'
    return uv
end

function rnd_grid(x, y)
    uv = mgrid(x, y, 0:x, 0:y)
    az = rand(Float32, x, y)
    return cat(uv, az, dims = 3)
end


using Test

@testset "Dirichlet0 Boundary Conditions" begin
    N = 100
    h_inv = 1/(N-1)
    d_order = 2
    approx_order = 2
    x = collect(1:h_inv:N).^2
    A = DerivativeOperator{Float64}(2,2,1/99,length(x),:Dirichlet0,:Dirichlet0,BC=(0.0,0.0))
    boundary_points = A.boundary_point_count
    res = A*x
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 2.0*ones(N - sum(boundary_points));

    N = 1000
    d_order = 4
    approx_order = 10
    y = collect(1:1.0:N).^4 .- 2*collect(1:1.0:N).^3 .+ collect(1:1.0:N).^2;
    A = DerivativeOperator{Float64}(d_order,approx_order,1.0,N,:Dirichlet0,:Dirichlet0,BC=(0.0,0.0))
    boundary_points = A.boundary_point_count
    res = A*y
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 24.0*ones(N - sum(boundary_points)) atol=10.0^-1; # Float64 is less stable

    A = DerivativeOperator{BigFloat}(d_order,approx_order,one(BigFloat),N,:Dirichlet0,:Dirichlet0,BC=(0.0,0.0))
    y = convert(Array{BigFloat, 1}, y)
    res = A*y
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 24.0*ones(N - sum(boundary_points)) atol=10.0^(-approx_order);
end

@testset "Dirichlet0 Boundary Conditions irregular grid" begin
    # tests as for DerivativeOperator, but with arbitrary grid (increasing step size):
    N = 1000
    d_order = 2
    approx_order = 2
    x = (1.:1.:N)*2.0 #erf.(linspace(-2,2,N))*N/2;
    dx = diff(x)
    y = x.^2

    A = FiniteDifference{Float64}(2,2,dx,length(x),:Dirichlet0,:Dirichlet0,BC=(0.0,0.0))
    boundary_points = A.boundary_point_count
    res = A*y
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 2.0*ones(N - sum(boundary_points));

    # N = 1000
    d_order = 4
    approx_order = 10
    # dx = rand(N)
    # x = [0.0; cumsum(dx)]
    y = @. x.^4 - 2*x.^3 + x.^2
    A = FiniteDifference{Float64}(d_order,approx_order,dx,N,:Dirichlet0,:Dirichlet0,BC=(0.0,0.0))
    boundary_points = A.boundary_point_count
    res = A*y
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 24.0*ones(N - sum(boundary_points)) atol=10.0^(-1); # Float64 is less stable

    A = FiniteDifference{BigFloat}(d_order,approx_order,BigFloat.(dx),N,:Dirichlet0,:Dirichlet0,BC=(0.0,0.0))
    y = convert(Array{BigFloat, 1}, y)
    res = A*y
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 24.0*ones(N - sum(boundary_points)) atol=10.0^(-approx_order);

end


"""
    SqExponentialKernel()

Squared exponential kernel.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the squared exponential kernel is defined as
```math
k(x, x') = \\exp\\bigg(- \\frac{\\|x - x'\\|_2^2}{2}\\bigg).
```

See also: [`GammaExponentialKernel`](@ref)
"""
struct SqExponentialKernel <: SimpleKernel end

kappa(::SqExponentialKernel, d²::Real) = exp(-d² / 2)

metric(::SqExponentialKernel) = SqEuclidean()

iskroncompatible(::SqExponentialKernel) = true

Base.show(io::IO, ::SqExponentialKernel) = print(io, "Squared Exponential Kernel")

## Aliases ##

"""
    RBFKernel()

Alias of [`SqExponentialKernel`](@ref).
"""
const RBFKernel = SqExponentialKernel

"""
    GaussianKernel()

Alias of [`SqExponentialKernel`](@ref).
"""
const GaussianKernel = SqExponentialKernel

"""
    SEKernel()

Alias of [`SqExponentialKernel`](@ref).
"""
const SEKernel = SqExponentialKernel

"""
    ExponentialKernel()

Exponential kernel.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the exponential kernel is defined as
```math
k(x, x') = \\exp\\big(- \\|x - x'\\|_2\\big).
```

See also: [`GammaExponentialKernel`](@ref)
"""
struct ExponentialKernel <: SimpleKernel end

kappa(::ExponentialKernel, d::Real) = exp(-d)

metric(::ExponentialKernel) = Euclidean()

iskroncompatible(::ExponentialKernel) = true

Base.show(io::IO, ::ExponentialKernel) = print(io, "Exponential Kernel")

## Aliases ##

"""
    LaplacianKernel()

Alias of [`ExponentialKernel`](@ref).
"""
const LaplacianKernel = ExponentialKernel

"""
    Matern12Kernel()

Alias of [`ExponentialKernel`](@ref).
"""
const Matern12Kernel = ExponentialKernel

"""
    GammaExponentialKernel(; γ::Real=2.0)

γ-exponential kernel with parameter `γ`.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the γ-exponential kernel[^RW] with parameter
``\\gamma \\in (0, 2]`` is defined as
```math
k(x, x'; \\gamma) = \\exp\\big(- \\|x - x'\\|_2^{\\gamma}\\big).
```

See also: [`ExponentialKernel`](@ref), [`SqExponentialKernel`](@ref)

[^RW]: C. E. Rasmussen & C. K. I. Williams (2006). Gaussian Processes for Machine Learning.
"""
struct GammaExponentialKernel{Tγ<:Real} <: SimpleKernel
    γ::Vector{Tγ}
    function GammaExponentialKernel(; gamma::Real=2.0, γ::Real=gamma)
        @check_args(GammaExponentialKernel, γ, zero(γ) < γ ≤ 2, "γ ∈ (0, 2]")
        return new{typeof(γ)}([γ])
    end
end

@functor GammaExponentialKernel

kappa(κ::GammaExponentialKernel, d::Real) = exp(-d^first(κ.γ))

metric(::GammaExponentialKernel) = Euclidean()

iskroncompatible(::GammaExponentialKernel) = true

function Base.show(io::IO, κ::GammaExponentialKernel)
    return print(io, "Gamma Exponential Kernel (γ = ", first(κ.γ), ")")
end


"""
    Link

An abstract type whose subtypes determine methods for [`linkfun`](@ref), [`linkinv`](@ref),
[`mueta`](@ref), and [`inverselink`](@ref).
"""
abstract type Link end

# Make links broadcast like a scalar
Base.Broadcast.broadcastable(l::Link) = Ref(l)

"""
    Link01

An abstract subtype of [`Link`](@ref) which are links defined on (0, 1)
"""
abstract type Link01 <: Link end

"""
    CauchitLink

A [`Link01`](@ref) corresponding to the standard Cauchy distribution,
[`Distributions.Cauchy`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Cauchy).
"""
mutable struct CauchitLink <: Link01 end

"""
    CloglogLink

A [`Link01`](@ref) corresponding to the extreme value (or log-Wiebull) distribution.  The
link is the complementary log-log transformation, `log(1 - log(-μ))`.
"""
mutable struct CloglogLink  <: Link01 end

"""
    IdentityLink

The canonical [`Link`](@ref) for the `Normal` distribution, defined as `η = μ`.
"""
mutable struct IdentityLink <: Link end

"""
    InverseLink

The canonical [`Link`](@ref) for [`Distributions.Gamma`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Gamma) distribution, defined as `η = inv(μ)`.
"""
mutable struct InverseLink  <: Link end

"""
    InverseSquareLink

The canonical [`Link`](@ref) for [`Distributions.InverseGaussian`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.InverseGaussian) distribution, defined as `η = inv(abs2(μ))`.
"""
mutable struct InverseSquareLink  <: Link end

"""
    LogitLink

The canonical [`Link01`](@ref) for [`Distributions.Bernoulli`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Bernoulli) and [`Distributions.Binomial`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Binomial).
The inverse link, [`linkinv`](@ref), is the c.d.f. of the standard logistic distribution,
[`Distributions.Logistic`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Logistic).
"""
mutable struct LogitLink <: Link01 end

"""
    LogLink

The canonical [`Link`](@ref) for [`Distributions.Poisson`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Poisson), defined as `η = log(μ)`.
"""
mutable struct LogLink <: Link end

"""
    NegativeBinomialLink

The canonical [`Link`](@ref) for [`Distributions.NegativeBinomial`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.NegativeBinomial) distribution, defined as `η = log(μ/(μ+θ))`.
The shape parameter θ has to be fixed for the distribution to belong to the exponential family.
"""
mutable struct NegativeBinomialLink  <: Link
    θ::Float64
end

"""
    ProbitLink

A [`Link01`](@ref) whose [`linkinv`](@ref) is the c.d.f. of the standard normal
distribution, [`Distributions.Normal()`](https://juliastats.github.io/Distributions.jl/stable/univariate.html#Distributions.Normal).
"""
mutable struct ProbitLink <: Link01 end

"""
    SqrtLink

A [`Link`](@ref) defined as `η = √μ`
"""
mutable struct SqrtLink <: Link end

"""
    linkfun(L::Link, μ)

Return `η`, the value of the linear predictor for link `L` at mean `μ`.

# Examples
```jldoctest
julia> μ = inv(10):inv(5):1
0.1:0.2:0.9

julia> show(linkfun.(LogitLink(), μ))
[-2.19722, -0.847298, 0.0, 0.847298, 2.19722]

```
"""
function linkfun end

"""
    linkinv(L::Link, η)

Return `μ`, the mean value, for link `L` at linear predictor value `η`.

# Examples
```jldoctest
julia> μ = 0.1:0.2:1
0.1:0.2:0.9

julia> η = logit.(μ);

julia> linkinv.(LogitLink(), η) ≈ μ
true
```
"""
function linkinv end

"""
    mueta(L::Link, η)

Return the derivative of [`linkinv`](@ref), `dμ/dη`, for link `L` at linear predictor value `η`.

# Examples
```jldoctest
julia> mueta(LogitLink(), 0.0)
0.25

julia> mueta(CloglogLink(), 0.0) ≈ 0.36787944117144233
true

julia> mueta(LogLink(), 2.0) ≈ 7.38905609893065
true
```
"""
function mueta end

"""
    inverselink(L::Link, η)

Return a 3-tuple of the inverse link, the derivative of the inverse link, and when appropriate, the variance function `μ*(1 - μ)`.

The variance function is returned as NaN unless the range of μ is (0, 1)

# Examples
```jldoctest
julia> inverselink(LogitLink(), 0.0)
(0.5, 0.25, 0.25)

julia> μ, oneminusμ, variance = inverselink(CloglogLink(), 0.0);

julia> μ + oneminusμ ≈ 1
true

julia> μ*(1 - μ) ≈ variance
true

julia> isnan(last(inverselink(LogLink(), 2.0)))
true
```
"""
function inverselink end

"""
    canonicallink(D::Distribution)

Return the canonical link for distribution `D`, which must be in the exponential family.

# Examples
```jldoctest
julia> canonicallink(Bernoulli())
LogitLink()
```
"""
function canonicallink end

linkfun(::CauchitLink, μ) = tan(pi * (μ - oftype(μ, 1/2)))
linkinv(::CauchitLink, η) = oftype(η, 1/2) + atan(η) / pi
mueta(::CauchitLink, η) = one(η) / (pi * (one(η) + abs2(η)))
function inverselink(::CauchitLink, η)
    μlower = atan(-abs(η)) / π
    μlower += oftype(μlower, 1/2)
    η > 0 ? 1 - μlower : μlower, inv(π * (1 + abs2(η))), μlower * (1 - μlower)
end

linkfun(::CloglogLink, μ) = log(-log1p(-μ))
function linkinv(::CloglogLink, η::T) where T<:Real
    clamp(-expm1(-exp(η)), eps(T), one(T) - eps(T))
end
function mueta(::CloglogLink, η::T) where T<:Real
    max(eps(T), exp(η) * exp(-exp(η)))
end
function inverselink(::CloglogLink, η)
    expη = exp(η)
    μ = -expm1(-expη)
    omμ = exp(-expη)   # the complement, 1 - μ
    μ, max(floatmin(μ), expη * omμ), max(floatmin(μ), μ * omμ)
end

linkfun(::IdentityLink, μ) = μ
linkinv(::IdentityLink, η) = η
mueta(::IdentityLink, η) = one(η)
inverselink(::IdentityLink, η) = η, one(η), oftype(η, NaN)

linkfun(::InverseLink, μ) = inv(μ)
linkinv(::InverseLink, η) = inv(η)
mueta(::InverseLink, η) = -inv(abs2(η))
function inverselink(::InverseLink, η)
    μ = inv(η)
    μ, -abs2(μ), oftype(μ, NaN)
end

linkfun(::InverseSquareLink, μ) = inv(abs2(μ))
linkinv(::InverseSquareLink, η) = inv(sqrt(η))
mueta(::InverseSquareLink, η) = -inv(2η*sqrt(η))
function inverselink(::InverseSquareLink, η)
    μ = inv(sqrt(η))
    μ, -μ / (2η), oftype(μ, NaN)
end

linkfun(::LogitLink, μ) = logit(μ)
linkinv(::LogitLink, η) = logistic(η)
function mueta(::LogitLink, η)
    expabs = exp(-abs(η))
    denom = 1 + expabs
    (expabs / denom) / denom
end
function inverselink(::LogitLink, η)
    expabs = exp(-abs(η))
    opexpabs = 1 + expabs
    deriv = (expabs / opexpabs) / opexpabs
    η ≤ 0 ? expabs / opexpabs : inv(opexpabs), deriv, deriv
end

linkfun(::LogLink, μ) = log(μ)
linkinv(::LogLink, η) = exp(η)
mueta(::LogLink, η) = exp(η)
function inverselink(::LogLink, η)
    μ = exp(η)
    μ, μ, oftype(μ, NaN)
end

linkfun(nbl::NegativeBinomialLink, μ) = log(μ / (μ + nbl.θ))
linkinv(nbl::NegativeBinomialLink, η) = ℯ^η * nbl.θ / (1-ℯ^η)
mueta(nbl::NegativeBinomialLink, η) = ℯ^η * nbl.θ / (1-ℯ^η)
function inverselink(nbl::NegativeBinomialLink, η)
    μ = ℯ^η * nbl.θ / (1-ℯ^η)
    deriv = μ * (1 + μ / nbl.θ)
    μ, deriv, oftype(μ, NaN)
end

linkfun(::ProbitLink, μ) = -sqrt2 * erfcinv(2μ)
linkinv(::ProbitLink, η) = erfc(-η / sqrt2) / 2
mueta(::ProbitLink, η) = exp(-abs2(η) / 2) / sqrt2π
function inverselink(::ProbitLink, η)
    μlower = erfc(abs(η) / sqrt2) / 2
    μupper = 1 - μlower
    η < 0 ? μlower : μupper, exp(-abs2(η) / 2 ) / sqrt2π, μlower * μupper
end

linkfun(::SqrtLink, μ) = sqrt(μ)
linkinv(::SqrtLink, η) = abs2(η)
mueta(::SqrtLink, η) = 2η
inverselink(::SqrtLink, η) = abs2(η), 2η, oftype(η, NaN)

canonicallink(::Bernoulli) = LogitLink()
canonicallink(::Binomial) = LogitLink()
canonicallink(::Gamma) = InverseLink()
canonicallink(::InverseGaussian) = InverseSquareLink()
canonicallink(d::NegativeBinomial) = NegativeBinomialLink(d.r)
canonicallink(::Normal) = IdentityLink()
canonicallink(::Poisson) = LogLink()

"""
    glmvar(D::Distribution, μ)

Return the value of the variance function for `D` at `μ`

The variance of `D` at `μ` is the product of the dispersion parameter, ϕ, which does not
depend on `μ` and the value of `glmvar`.  In other words `glmvar` returns the factor of the
variance that depends on `μ`.

# Examples
```jldoctest
julia> μ = 1/6:1/3:1;

julia> glmvar.(Normal(), μ)    # constant for Normal()
3-element Array{Float64,1}:
 1.0
 1.0
 1.0

julia> glmvar.(Bernoulli(), μ) ≈ μ .* (1 .- μ)
true

julia> glmvar.(Poisson(), μ) == μ
true
```
"""
function glmvar end

glmvar(::Union{Bernoulli,Binomial}, μ) = μ * (1 - μ)
glmvar(::Gamma, μ) = abs2(μ)
glmvar(::InverseGaussian, μ) = μ^3
glmvar(d::NegativeBinomial, μ) = μ * (1 + μ/d.r)
glmvar(::Normal, μ) = one(μ)
glmvar(::Poisson, μ) = μ

"""
    mustart(D::Distribution, y, wt)

Return a starting value for μ.

For some distributions it is appropriate to set `μ = y` to initialize the IRLS algorithm but
for others, notably the Bernoulli, the values of `y` are not allowed as values of `μ` and
must be modified.

# Examples
```jldoctest
julia> mustart(Bernoulli(), 0.0, 1) ≈ 1/4
true

julia> mustart(Bernoulli(), 1.0, 1) ≈ 3/4
true

julia> mustart(Binomial(), 0.0, 10) ≈ 1/22
true

julia> mustart(Normal(), 0.0, 1) ≈ 0
true
```
"""
function mustart end

mustart(::Bernoulli, y, wt) = (y + oftype(y, 1/2)) / 2
mustart(::Binomial, y, wt) = (wt * y + oftype(y, 1/2)) / (wt + one(y))
function mustart(::Union{Gamma, InverseGaussian}, y, wt)
    fy = float(y)
    iszero(y) ? oftype(y, 1/10) : fy
end
function mustart(::NegativeBinomial, y, wt)
    fy = float(y)
    iszero(y) ? fy + oftype(fy, 1/6) : fy
end
mustart(::Normal, y, wt) = y
function mustart(::Poisson, y, wt)
    fy = float(y)
    fy + oftype(fy, 1/10)
end

"""
    devresid(D, y, μ)

Return the squared deviance residual of `μ` from `y` for distribution `D`

The deviance of a GLM can be evaluated as the sum of the squared deviance residuals.  This
is the principal use for these values.  The actual deviance residual, say for plotting, is
the signed square root of this value
```julia
sign(y - μ) * sqrt(devresid(D, y, μ))
```

# Examples
```jldoctest
julia> devresid(Normal(), 0, 0.25) ≈ abs2(0.25)
true

julia> devresid(Bernoulli(), 1, 0.75) ≈ -2*log(0.75)
true

julia> devresid(Bernoulli(), 0, 0.25) ≈ -2*log1p(-0.25)
true
```
"""
function devresid end

function devresid(::Bernoulli, y, μ)
    if y == 1
        return -2 * log(μ)
    elseif y == 0
        return -2 * log1p(-μ)
    end
    throw(ArgumentError("y should be 0 or 1 (got $y)"))
end
function devresid(::Binomial, y, μ)
    if y == 1
        return -2 * log(μ)
    elseif y == 0
        return -2 * log1p(-μ)
    else
        return 2 * (y * (log(y) - log(μ)) + (1 - y)*(log1p(-y) - log1p(-μ)))
    end
end
devresid(::Gamma, y, μ) = -2 * (log(y / μ) - (y - μ) / μ)
devresid(::InverseGaussian, y, μ) = abs2(y - μ) / (y * abs2(μ))
function devresid(d::NegativeBinomial, y, μ)
    θ = d.r
    v = 2 * (xlogy(y, y / μ) + xlogy(y + θ, (μ + θ)/(y + θ)))
    return μ == 0 ? oftype(v, NaN) : v
end
devresid(::Normal, y, μ) = abs2(y - μ)
devresid(::Poisson, y, μ) = 2 * (xlogy(y, y / μ) - (y - μ))

"""
    dispersion_parameter(D)  # not exported

Does distribution `D` have a separate dispersion parameter, ϕ?

Returns `false` for the `Bernoulli`, `Binomial` and `Poisson` distributions, `true` otherwise.

# Examples
```jldoctest
julia> show(GLM.dispersion_parameter(Normal()))
true
julia> show(GLM.dispersion_parameter(Bernoulli()))
false
```
"""
dispersion_parameter(D) = true
dispersion_parameter(::Union{Bernoulli, Binomial, Poisson}) = false

"""
    loglik_obs(D, y, μ, wt, ϕ)  # not exported

Returns `wt * logpdf(D(μ, ϕ), y)` where the parameters of `D` are derived from `μ` and `ϕ`.

The `wt` argument is a multiplier of the result except in the case of the `Binomial` where
`wt` is the number of trials and `μ` is the proportion of successes.

The loglikelihood of a fitted model is the sum of these values over all the observations.
"""
function loglik_obs end

loglik_obs(::Bernoulli, y, μ, wt, ϕ) = wt*logpdf(Bernoulli(μ), y)
loglik_obs(::Binomial, y, μ, wt, ϕ) = logpdf(Binomial(Int(wt), μ), Int(y*wt))
loglik_obs(::Gamma, y, μ, wt, ϕ) = wt*logpdf(Gamma(inv(ϕ), μ*ϕ), y)
loglik_obs(::InverseGaussian, y, μ, wt, ϕ) = wt*logpdf(InverseGaussian(μ, inv(ϕ)), y)
loglik_obs(::Normal, y, μ, wt, ϕ) = wt*logpdf(Normal(μ, sqrt(ϕ)), y)
loglik_obs(::Poisson, y, μ, wt, ϕ) = wt*logpdf(Poisson(μ), y)
# We use the following parameterization for the Negative Binomial distribution:
#    (Γ(θ+y) / (Γ(θ) * y!)) * μ^y * θ^θ / (μ+θ)^{θ+y}
# The parameterization of NegativeBinomial(r=θ, p) in Distributions.jl is
#    Γ(θ+y) / (y! * Γ(θ)) * p^θ(1-p)^y
# Hence, p = θ/(μ+θ)
loglik_obs(d::NegativeBinomial, y, μ, wt, ϕ) = wt*logpdf(NegativeBinomial(d.r, d.r/(μ+d.r)), y)


function isprime(n)
    if n == 1
        return false
    end
    if n == 2 || n == 3
        return true
    end

    for i in 2:isqrt(n)
        if n%i == 0
            return false
        else
            if i == 2
                i += 1
            else
                i += 2
            end
        end
    end
    return true
end


using GeometricAlgebra
const ga = GeometricAlgebra
using Test

@basis "---+"
include("definitions.jl")
include("../operators.jl")


"""
     enkf(state_ens, d_matrix, hx_matrix)

Return updated states
"""
function enkf(state_ens, d_matrix, hx_matrix)

  # Variables (Mandel)

  X  = state_ens;
  D  = d_matrix;
  HX = hx_matrix;

  # Subtract ensemble mean (Mandel)

  (n, N) = size(X);
  (m, N) = size(D);

  A    = X - 1 / N * (X * ones(Float64, N, 1)) * ones(Float64, 1, N);
  HA   = HX - 1 / N * (HX * ones(Float64, N, 1)) * ones(Float64, 1, N);
  Dtmp = D - 1 / N * (D * ones(Float64, N, 1)) * ones(Float64, 1, N);

  # Observation error variance (Mandel-theoretic, Evensen-sample)

  R_sample = Dtmp * Dtmp' / (N-1);

  # Variance of predicted observations (DeChant and Mandel)

  C_YY = 1 / (N-1) * HA * HA';

  # Covariance between states ensemble and predicted observations (DeChant and Mandel)

  C_XY = 1 / (N-1) * A * HA';

  # Compute kalman gain (DeChant)

  K = C_XY / (C_YY + R_sample);

  if any(isnan(K))
    println("R_sample = $R_sample")
    println("C_YY = $C_YY")
    println("C_XY = $C_XY")
    error("Nans in Kalman gain")
  end

  # Update states (DeChant and Mandel)

  Xhat = X + K*(D-HX);

  return(Xhat);

end

"""
     enkf!(state_ens, d_matrix, hx_matrix)

Update states in place
"""
function enkf!(state_ens, d_matrix, hx_matrix)

  # Variables (Mandel)

  X  = state_ens;
  D  = d_matrix;
  HX = hx_matrix;

  # Subtract ensemble mean (Mandel)

  (n, N) = size(X);
  (m, N) = size(D);

  A    = X - 1 / N * (X * ones(Float64, N, 1)) * ones(Float64, 1, N);
  HA   = HX - 1 / N * (HX * ones(Float64, N, 1)) * ones(Float64, 1, N);
  Dtmp = D - 1 / N * (D * ones(Float64, N, 1)) * ones(Float64, 1, N);

  # Observation error variance (Mandel-theoretic, Evensen-sample)

  R_sample = Dtmp * Dtmp' / (N-1);

  # Variance of predicted observations (DeChant and Mandel)

  C_YY = 1 / (N-1) * HA * HA';

  # Covariance between states ensemble and predicted observations (DeChant and Mandel)

  C_XY = 1 / (N-1) * A * HA';

  # Compute kalman gain (DeChant)

  K = C_XY / (C_YY + R_sample);

  if any(isnan(K))
    println("R_sample = $R_sample")
    println("C_YY = $C_YY")
    println("C_XY = $C_XY")
    error("Nans in Kalman gain")
  end

  # Update states (DeChant and Mandel)

  X = X + K*(D-HX);

  nothing

end


#all kernel functions are normalized such that their support radius is h

@fastmath function pos(x::Float64)::Float64
	return (x > 0.0 ? x : 0.0)
end

"""
    spline23(h::Float64, r::Float64)::Float64

Returns ``w(r)``, the value of a 2d cubic spline ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function spline23(h::Float64, r::Float64)::Float64
	x = r/h
	if (x < 0.5)
		# 40/7pi = 1.8189136353359467
		return 1.8189136353359467*(1.0 - 6.0*x^2 + 6.0*x^3)/h^2
	elseif (x < 1.0)
		# 80/7pi = 3.6378272706718935
		return 3.6378272706718935*(1.0 - x)^3/h^2
	end
	return 0.0
end

"""
    Dspline23(h::Float64, r::Float64)::Float64

Returns ``\\frac{\\text{d}w}{\\text{d}r}(r)``, the first derivative of a 2d cubic spline ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function Dspline23(h::Float64, r::Float64)::Float64
	x = r/h
	if (x < 0.5)
		# 240/7pi = 10.91348181201568
		return -10.91348181201568*(2.0*x - 3.0*x^2)/h^3
	elseif (x < 1.0)
		return -10.91348181201568*(1.0 - x)^2/h^3
	end
	return 0.0
end

"""
    rDspline23(h::Float64, r::Float64)::Float64

Returns ``\\frac{1}{r}\\,\\frac{\\text{d}w}{\\text{d}r}(r)``, the reduced first derivative of a 2d cubic spline ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function rDspline23(h::Float64, r::Float64)::Float64
	x = r/h
	if (x < 0.5)
		# 240/7pi = 10.91348181201568
		return -10.91348181201568*(2.0 - 3.0*x)/h^4
	elseif (x < 1.0)
		return -10.91348181201568*(1.0 - x)^2/(x*h^4)
	end
	return 0.0
end

"""
    spline24(h::Float64, r::Float64)::Float64

Returns ``w(r)``, the value of a 2d quartic spline ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function spline24(h::Float64, r::Float64)::Float64
	x = r/h
	return 6.222175110452539*(pos(1.0 - x)^4 - 5*pos(0.6 - x)^4 + 10*pos(0.2 - x)^4)/h^2
end

"""
    Dspline24(h::Float64, r::Float64)::Float64

Returns ``\\frac{\\text{d}w}{\\text{d}r}(r)``, the first derivative of a 2d quartic spline ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function Dspline24(h::Float64, r::Float64)::Float64
	x = r/h
	return -24.888700441810155*(pos(1.0 - x)^3 - 5*pos(0.6 - x)^3 + 10*pos(0.2 - x)^3)/h^3
end

"""
    rDspline24(h::Float64, r::Float64)::Float64

Returns ``\\frac{1}{r}\\,\\frac{\\text{d}w}{\\text{d}r}(r)``, the reduced first derivative of a 2d quartic spline ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function rDspline24(h::Float64, r::Float64)::Float64
	x = r/h
	if x > 0.2
		return -24.888700441810155*(pos(1.0 - x)^3 - 5*pos(0.6 - x)^3)/(x*h^4)
	end
	return -24.888700441810155*(1.2 - 6.0*x^2)/h^4
end

"""
    wendland2(h::Float64, r::Float64)::Float64

Returns ``w(r)``, the value of a 2d quintic Wendland kernel ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function wendland2(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 7/pi ≐ 2.228169203286535
	return 2.228169203286535*((1.0 - x)^4)*(1.0 + 4.0*x)/h^2
end

"""
    Dwendland2(h::Float64, r::Float64)::Float64

Returns ``\\frac{\\text{d}w}{\\text{d}r}(r)``, the first derivative of a 2d quintic Wendland kernel ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function Dwendland2(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 140/pi ≐ 44.563384065730695
	return -44.563384065730695*x*((1.0 - x)^3)/h^3
end

"""
    rDwendland2(h::Float64, r::Float64)::Float64

Returns ``\\frac{1}{r}\\,\\frac{\\text{d}w}{\\text{d}r}(r)``, the reduced first derivative of a 2d quintic Wendland kernel ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function rDwendland2(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 140/pi ≐ 44.563384065730695
	return -44.563384065730695*((1.0 - x)^3)/h^4
end

"""
    wendland3(h::Float64, r::Float64)::Float64

Returns ``w(r)``, the value of a 3d quintic Wendland kernel ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function wendland3(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 21/2pi ≐ 3.3422538049298023
	return 3.3422538049298023*((1.0 - x)^4)*(1.0 + 4.0*x)/h^3
end

"""
    Dwendland3(h::Float64, r::Float64)::Float64

Returns ``\\frac{\\text{d}w}{\\text{d}r}(r)``, the first derivative of a 3d quintic Wendland kernel ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function Dwendland3(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 210/pi ≐ 66.84507609859604
	return -66.84507609859604*x*((1.0 - x)^3)/h^4
end

"""
    rDwendland3(h::Float64, r::Float64)::Float64

Returns ``\\frac{1}{r}\\,\\frac{\\text{d}w}{\\text{d}r}(r)``, the reduced first derivative of a 3d quintic Wendland kernel ``w`` with support radius `h`.
Integrates to unity.

"""
@fastmath function rDwendland3(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 210/pi ≐ 66.84507609859604
	return -66.84507609859604*((1.0 - x)^3)/h^5
end

@fastmath function DDwendland3(h::Float64, r::Float64)::Float64
	x = r/h
	if (x > 1.0)
		return 0.0
	end
	# 210/pi ≐ 66.84507609859604
	return -66.84507609859604*((1.0 - 4.0*x)*(1.0 - x)^2)/h^5
end

@fastmath function spiky3(h::Float64, r::Float64)::Float64
	x = r/h
	#15/2pi = 2.3873241463784303
	return 2.3873241463784303*(1.0 - x)^2/h^3
end

@fastmath function Dspiky3(h::Float64, r::Float64)::Float64
	x = r/h
	#15/pi = 4.7746482927568605
	return -(x < 1.0)*2.3873241463784303*(1.0 - x)/h^4
end

@fastmath function rDspiky3(h::Float64, r::Float64)::Float64
	x = r/h
	#15/pi = 4.7746482927568605
	return -(x < 1.0)*2.3873241463784303*(1.0 - x)/((x + 1e-6)*h^5)
end

@fastmath function DDspiky3(h::Float64, r::Float64)::Float64
	x = r/h
	#15/pi = 4.7746482927568605
	return (x < 1.0)*2.3873241463784303/h^5
end


using Combinatorics

function findMinWrapDist(f_true,f_hat)
# Assign each of the ground truth frequencies to closest
# (squared error) canditate of estimated frequencies.
# Metric is the wrap-aroud distance on the torus $\mathds{T}=[0,1]$.
#
# Inputs:
# f_true: ground truth frequencies
# f_hat: estimated frequencies
#
# Output:
# dists: wrap-around distances of the frequencies in f_true and f_hat

K = length(f_true)
WrapDist(x,y) = min(abs(x-y),abs(x-y+1))

dists = Array{Float64,1}(undef,K)
for ii = 1:K
    dists[ii] = minimum(WrapDist.(f_true[ii],f_hat))
end

return dists
end