export FiniteDifferenceMethod, fdm, backward_fdm, forward_fdm, central_fdm, extrapolate_fdm

"""
    estimate_magitude(f, x::T) where T<:AbstractFloat

Estimate the magnitude of `f` in a neighbourhood of `x`, assuming that the outputs of `f`
have a "typical" order of magnitude. The result should be interpreted as a very rough
estimate. This function deals with the case that `f(x) = 0`.
"""
function estimate_magitude(f, x::T) where T<:AbstractFloat
    M = float(maximum(abs, f(x)))
    M > 0 && (return M)
    # Ouch, `f(x) = 0`. But it may not be zero around `x`. We conclude that `x` is likely a
    # pathological input for `f`. Perturb `x`. Assume that the pertubed value for `x` is
    # highly unlikely also a pathological value for `f`.
    Δ = convert(T, 0.1) * max(abs(x), one(x))
    return float(maximum(abs, f(x + Δ)))
end

"""
    estimate_roundoff_error(f, x::T) where T<:AbstractFloat

Estimate the round-off error of `f(x)`. This function deals with the case that `f(x) = 0`.
"""
function estimate_roundoff_error(f, x::T) where T<:AbstractFloat
    # Estimate the round-off error. It can happen that the function is zero around `x`, in
    # which case we cannot take `eps(f(x))`. Therefore, we assume a lower bound that is
    # equal to `eps(T) / 1000`, which gives `f` four orders of magnitude wiggle room.
    return max(eps(estimate_magitude(f, x)), eps(T) / 1000)
end

"""
    FiniteDifferences.DEFAULT_CONDITION

The default [condition number](https://en.wikipedia.org/wiki/Condition_number) used when
computing bounds. It provides amplification of the ∞-norm when passed to the function's
derivatives.
"""
const DEFAULT_CONDITION = 100

"""
    FiniteDifferenceMethod{G<:AbstractVector, C<:AbstractVector, E<:Function}

A finite difference method.

# Fields
- `grid::G`: Multiples of the step size that the function will be evaluated at.
- `q::Int`: Order of derivative to estimate.
- `coefs::C`: Coefficients corresponding to the grid functions that the function evaluations
    will be weighted by.
- `bound_estimator::Function`: A function that takes in the function and the evaluation
    point and returns a bound on the magnitude of the `length(grid)`th derivative.
"""
struct FiniteDifferenceMethod{G<:AbstractVector, C<:AbstractVector, E<:Function}
    grid::G
    q::Int
    coefs::C
    bound_estimator::E
end

"""
    FiniteDifferenceMethod(
        grid::AbstractVector,
        q::Int;
        condition::Real=DEFAULT_CONDITION
    )

Construct a finite difference method.

# Arguments
- `grid::Abstract`: The grid. See [`FiniteDifferenceMethod`](@ref).
- `q::Int`: Order of the derivative to estimate.
- `condition::Real`: Condition number. See [`DEFAULT_CONDITION`](@ref).

# Returns
- `FiniteDifferenceMethod`: Specified finite difference method.
"""
function FiniteDifferenceMethod(
    grid::AbstractVector,
    q::Int;
    condition::Real=DEFAULT_CONDITION
)
    p = length(grid)
    _check_p_q(p, q)
    return FiniteDifferenceMethod(
        grid,
        q,
        _coefs(grid, q),
        _make_default_bound_estimator(condition=condition)
    )
end

"""
    (m::FiniteDifferenceMethod)(
        f::Function,
        x::T;
        factor::Real=1,
        max_step::Real=0.1 * max(abs(x), one(x))
    ) where T<:AbstractFloat

Estimate the derivative of `f` at `x` using the finite differencing method `m` and an
automatically determined step size.

# Arguments
- `f::Function`: Function to estimate derivative of.
- `x::T`: Input to estimate derivative at.

# Keywords
- `factor::Real=1`: Factor to amplify the estimated round-off error by. This can be used
    to force a more conservative step size.
- `max_step::Real=0.1 * max(abs(x), one(x))`: Maximum step size.

# Returns
- Estimate of the derivative.

# Examples

```julia-repl
julia> fdm = central_fdm(5, 1)
FiniteDifferenceMethod:
  order of method:       5
  order of derivative:   1
  grid:                  [-2, -1, 0, 1, 2]
  coefficients:          [0.08333333333333333, -0.6666666666666666, 0.0, 0.6666666666666666, -0.08333333333333333]

julia> fdm(sin, 1)
0.5403023058681155

julia> fdm(sin, 1) - cos(1)  # Check the error.
-2.4313884239290928e-14

julia> FiniteDifferences.estimate_step(fdm, sin, 1.0)  # Computes step size and estimates the error.
(0.0010632902144695163, 1.9577610541734626e-13)
```
"""
@inline function (m::FiniteDifferenceMethod)(f::Function, x::Real; kw_args...)
    # Assume that converting to float is desired.
    return _call_method(m, f, float(x); kw_args...)
end
@inline function _call_method(
    m::FiniteDifferenceMethod,
    f::Function,
    x::T;
    factor::Real=1,
    max_step::Real=0.1 * max(abs(x), one(x))
) where T<:AbstractFloat
    # The automatic step size calculation fails if `m.q == 0`, so handle that edge case.
    iszero(m.q) && return f(x)
    h, _ = estimate_step(m, f, x, factor=factor, max_step=max_step)
    return _eval_method(m, f, x, h)
end

"""
    (m::FiniteDifferenceMethod)(f::Function, x::T, h::Real) where T<:AbstractFloat

Estimate the derivative of `f` at `x` using the finite differencing method `m` and a given
step size.

# Arguments
- `f::Function`: Function to estimate derivative of.
- `x::T`: Input to estimate derivative at.
- `h::Real`: Step size.

# Returns
- Estimate of the derivative.

# Examples

```julia-repl
julia> fdm = central_fdm(5, 1)
FiniteDifferenceMethod:
  order of method:       5
  order of derivative:   1
  grid:                  [-2, -1, 0, 1, 2]
  coefficients:          [0.08333333333333333, -0.6666666666666666, 0.0, 0.6666666666666666, -0.08333333333333333]

julia> fdm(sin, 1, 1e-3)
0.5403023058679624

julia> fdm(sin, 1, 1e-3) - cos(1)  # Check the error.
-1.7741363933510002e-13
```
"""
@inline function (m::FiniteDifferenceMethod)(f::Function, x::Real, h::Real)
    # Assume that converting to float is desired.
    return _eval_method(m, f, float(x), h)
end
@inline function _eval_method(
    m::FiniteDifferenceMethod,
    f::Function,
    x::T,
    h::Real
) where T<:AbstractFloat
    return sum(
        i -> convert(T, m.coefs[i]) * f(T(x + h * m.grid[i])),
        eachindex(m.grid)
    ) / h^m.q
end

# Check the method and derivative orders for consistency.
function _check_p_q(p::Integer, q::Integer)
    q >= 0 || throw(DomainError(q, "order of derivative (`q`) must be non-negative"))
    q < p || throw(DomainError(
        (q, p),
        "order of the method (q) must be strictly greater than that of the derivative (p)",
    ))
    # Check whether the method can be computed. We require the factorial of the method order
    # to be computable with regular `Int`s, but `factorial` will after 20, so 20 is the
    # largest we can allow.
    p > 20 && throw(DomainError(p, "order of the method (`p`) is too large to be computed"))
    return
end

const _COEFFS_CACHE = Dict{Tuple{AbstractVector{<:Real}, Integer}, Vector{Float64}}()

# Compute coefficients for the method and cache the result.
function _coefs(grid::AbstractVector{<:Real}, q::Integer)
    return get!(_COEFFS_CACHE, (grid, q)) do
        p = length(grid)
        # For high precision on the `\`, we use `Rational`, and to prevent overflows we use
        # `Int128`. At the end we go to `Float64` for fast floating point math, rather than
        # rational math.
        C = [Rational{Int128}(g^i) for i in 0:(p - 1), g in grid]
        x = zeros(Rational{Int128}, p)
        x[q + 1] = factorial(q)
        return Float64.(C \ x)
    end
end

# Estimate the bound on the derivative by amplifying the ∞-norm.
function _make_default_bound_estimator(; condition::Real=DEFAULT_CONDITION)
    default_bound_estimator(f, x) = condition * estimate_magitude(f, x)
    return default_bound_estimator
end

function Base.show(io::IO, m::MIME"text/plain", x::FiniteDifferenceMethod)
    @printf io "FiniteDifferenceMethod:\n"
    @printf io "  order of method:       %d\n" length(x.grid)
    @printf io "  order of derivative:   %d\n" x.q
    @printf io "  grid:                  %s\n" x.grid
    @printf io "  coefficients:          %s\n" x.coefs
end

"""
    function estimate_step(
        m::FiniteDifferenceMethod,
        f::Function,
        x::T;
        factor::Real=1,
        max_step::Real=0.1 * max(abs(x), one(x))
    ) where T<:AbstractFloat

Estimate the step size for a finite difference method `m`. Also estimates the error of the
estimate of the derivative.

# Arguments
- `m::FiniteDifferenceMethod`: Finite difference method to estimate the step size for.
- `f::Function`: Function to evaluate the derivative of.
- `x::T`: Point to estimate the derivative at.

# Keywords
- `factor::Real=1`. Factor to amplify the estimated round-off error by. This can be used
    to force a more conservative step size.
- `max_step::Real=0.1 * max(abs(x), one(x))`: Maximum step size.

# Returns
- `Tuple{T, <:AbstractFloat}`: Estimated step size and an estimate of the error of the
    finite difference estimate.
"""
function estimate_step(
    m::FiniteDifferenceMethod,
    f::Function,
    x::T;
    factor::Real=1,
    max_step::Real=0.1 * max(abs(x), one(x))
) where T<:AbstractFloat
    p = length(m.coefs)
    q = m.q

    # Estimate the round-off error.
    ε = estimate_roundoff_error(f, x) * factor

    # Estimate the bound on the derivatives.
    M = m.bound_estimator(f, x)

    # Set the step size by minimising an upper bound on the error of the estimate.
    C₁ = ε * sum(abs, m.coefs)
    C₂ = M * sum(n -> abs(m.coefs[n] * m.grid[n]^p), eachindex(m.coefs)) / factorial(p)
    # Type inference fails on this, so we annotate it, which gives big performance benefits.
    h::T = convert(T, min((q / (p - q) * (C₁ / C₂))^(1 / p), max_step))

    # Estimate the accuracy of the method.
    accuracy = h^(-q) * C₁ + h^(p - q) * C₂

    return h, accuracy
end

for direction in [:forward, :central, :backward]
    fdm_fun = Symbol(direction, "_fdm")
    grid_fun = Symbol("_", direction, "_grid")
    @eval begin function $fdm_fun(
            p::Int,
            q::Int;
            adapt::Int=1,
            condition::Real=DEFAULT_CONDITION,
            geom::Bool=false
        )
            _check_p_q(p, q)
            grid = $grid_fun(p)
            geom && (grid = _exponentiate_grid(grid))
            coefs = _coefs(grid, q)
            return FiniteDifferenceMethod(
                grid,
                q,
                coefs,
                _make_adaptive_bound_estimator($fdm_fun, p, q, adapt, condition, geom=geom),
            )
        end

        @doc """
    $($(Meta.quot(fdm_fun)))(
        p::Int,
        q::Int;
        adapt::Int=1,
        condition::Real=DEFAULT_CONDITION,
        geom::Bool=false
    )

Contruct a finite difference method at a $($(Meta.quot(direction))) grid of `p` linearly
spaced points.

# Arguments
- `p::Int`: Number of grid points.
- `q::Int`: Order of the derivative to estimate.

# Keywords
- `adapt::Int=1`: Use another finite difference method to estimate the magnitude of the
    `p`th order derivative, which is important for the step size computation. Recurse
    this procedure `adapt` times.
- `condition::Real`: Condition number. See [`DEFAULT_CONDITION`](@ref).
- `geom::Bool`: Use geometrically spaced points instead of linearly spaced points.

# Returns
- `FiniteDifferenceMethod`: The specified finite difference method.
        """ $fdm_fun
    end
end

function _make_adaptive_bound_estimator(
    constructor::Function,
    p::Int,
    q::Int,
    adapt::Int,
    condition::Int;
    kw_args...
)
    if adapt >= 1
        estimate_derivative = constructor(
            p + 1, p, adapt=adapt - 1, condition=condition; kw_args...
        )
        return (f, x) -> estimate_magitude(x′ -> estimate_derivative(f, x′), x)
    else
        return _make_default_bound_estimator(condition=condition)
    end
end

_forward_grid(p::Int) = collect(0:(p - 1))

_backward_grid(p::Int) = collect((1 - p):0)

function _central_grid(p::Int)
    if isodd(p)
        return collect(div(1 - p, 2):div(p - 1, 2))
    else
        return vcat(div(-p, 2):-1, 1:div(p, 2))
    end
end

_exponentiate_grid(grid::Vector, base::Int=3) = sign.(grid) .* base .^ abs.(grid) ./ base

function _is_symmetric(vec::Vector; centre_zero::Bool=false, negate_half::Bool=false)
    half_sign = negate_half ? -1 : 1
    if isodd(length(vec))
        centre_zero && vec[end ÷ 2 + 1] != 0 && return false
        return vec[1:end ÷ 2] == half_sign .* reverse(vec[(end ÷ 2 + 2):end])
    else
        return vec[1:end ÷ 2] == half_sign .* reverse(vec[(end ÷ 2 + 1):end])
    end
end

function _is_symmetric(m::FiniteDifferenceMethod)
    grid_symmetric = _is_symmetric(m.grid, centre_zero=true, negate_half=true)
    coefs_symmetric =_is_symmetric(m.coefs, negate_half=true)
    return grid_symmetric && coefs_symmetric
end

"""
    extrapolate_fdm(
        m::FiniteDifferenceMethod,
        f::Function,
        x::T,
        h::Real=0.1 * max(abs(x), one(x));
        power=nothing,
        breaktol=Inf,
        kw_args...
    ) where T<:AbstractFloat

Use Richardson extrapolation to refine a finite difference method.

Takes further in keyword arguments for `Richardson.extrapolate`. This method
automatically sets `power = 2` if `m` is symmetric and `power = 1`. Moreover, it defaults
`breaktol = Inf`.

# Arguments
- `m::FiniteDifferenceMethod`: Finite difference method to estimate the step size for.
- `f::Function`: Function to evaluate the derivative of.
- `x::T`: Point to estimate the derivative at.
- `h::Real=0.1 * max(abs(x), one(x))`: Initial step size.

# Returns
- `Tuple{<:AbstractFloat, <:AbstractFloat}`: Estimate of the derivative and error.
"""
function extrapolate_fdm(
    m::FiniteDifferenceMethod,
    f::Function,
    x::T,
    h::Real=0.1 * max(abs(x), one(x));
    power::Int=1,
    breaktol::Real=Inf,
    kw_args...
) where T<:AbstractFloat
    (power == 1 && _is_symmetric(m)) && (power = 2)
    return extrapolate(h -> m(f, x, h), h; power=power, breaktol=breaktol, kw_args...)
end


module TopologicalDataAnalysis

# Write your package code here.
using LinearAlgebra
using Combinatorics

export getpowerset

"""
    getpowerset(L)

Returns the power set of the list L.

Input:
    L (list).

Output:
    L2 (list): power set of L.

Example:
    L = collect(1:3)
    getpowerset(L)
    [Int64[], [1], [2], [3], [1, 2], [1, 3], [2, 3], [1, 2, 3]]
"""
function getpowerset(L)
    return collect(powerset(L))
end

end


## Day 20: Trench Map ##
########################

day20_input = (10, 3)
day20_test = (4, 8)

# Given two player positions, simulates the dirac dice game
function simulate(a, b)
    dice = 10
    score_a = 0
    score_b = 0
    rolls = 0
    while true
        for _ = 1:3
            rolls += 1
            dice = wrap(dice + 1)
            a = wrap(a + dice)
        end
        score_a += a
        if score_a >= 1000
            return score_b, rolls
        end
        # @show score_a
        for _ = 1:3
            rolls += 1
            dice = wrap(dice + 1)
            b = wrap(b + dice)
        end
        score_b += b
        if score_b >= 1000
            return score_a, rolls
        end
        # @show score_b
    end
end

wrap(position) = position > 10 ? position - 10 : position

# Part 2, this looks more tricky. I thing the solution is to group universes
# that are identical and reduce them to a number.

# What is the state of the universe? (a, b, score_a, score_b). There are
# 10 * 10 * 31 * 31 = 96100 possible states (including 20 + 10 = 30 and 0 score)

function dirac(a, b)
    # To make indexing easier, players start with score 1 and win with 22 points.
    states = zeros(Int, (10, 10, 31, 31))
    states[a, b, 1, 1] = 1
    a_wins = 0
    b_wins = 0
    while sum(states) > 0
        @show "loop", sum(states), a_wins, b_wins
        # Evolve the metaverse
        new_states = zeros(Int, (10, 10, 31, 31))
        for a = 1:10, b = 1:10, score_a = 1:21, score_b = 1:21
            if states[a, b, score_a, score_b] > 0
                new_states .+= dirac_one_a(a, b, score_a, score_b) .* states[a, b, score_a, score_b]
            end
        end
        states = new_states
        # Filter out states that are already won and sum them up.
        for a = 1:10, b = 1:10, score_a = 22:31, score_b = 1:31
            a_wins += states[a, b, score_a, score_b]
            states[a, b, score_a, score_b] = 0
        end
        # Evolve the metaverse
        new_states = zeros(Int, (10, 10, 31, 31))
        for a = 1:10, b = 1:10, score_a = 1:21, score_b = 1:21
            if states[a, b, score_a, score_b] > 0
                new_states .+= dirac_one_b(a, b, score_a, score_b) .* states[a, b, score_a, score_b]
            end
        end
        states = new_states
        # Filter out states that are already won and sum them up.
        for a = 1:10, b = 1:10, score_a = 1:31, score_b = 22:31
            b_wins += states[a, b, score_a, score_b]
            states[a, b, score_a, score_b] = 0
        end
    end
    (a_wins, b_wins)
end

# You can roll 3*3*3 = 27 different universes which collapse into
roll_count = Dict(3 => 1, 4 => 3, 5 => 6, 6 => 7, 7 => 6, 8 => 3, 9 => 1)
@assert sum(values(roll_count)) == 27

function dirac_one_a(a, b, score_a, score_b)::Array
    states = zeros(Int, (10, 10, 31, 31))
    for roll in keys(roll_count)
        new_a = wrap(a + roll)
        states[new_a, b, score_a+new_a, score_b] = roll_count[roll]
    end
    states
end

function dirac_one_b(a, b, score_a, score_b)::Array
    states = zeros(Int, (10, 10, 31, 31))
    for roll in keys(roll_count)
        new_b = wrap(b + roll)
        states[a, new_b, score_a, score_b+new_b] = roll_count[roll]
    end
    states
end


function plotfft(signal,Fs)
           freqs = fftfreq(length(signal),Fs)
           freqs = freqs[freqs .>=0]
           F = fft(signal)[1:length(freqs)]
           sig = 20 .* log10.((2/length(freqs)).*abs.(F))
           replace!(sig,-Inf=>-300)
           semilogx(freqs ,sig)
       end


function acoeffs(ly::Float64,ny::Int,Ξ::Float64,τ::Float64=0.0;jw::Float64=0.05)
    ζjet = zeros(Float64,2*ny-1)
    Δθ::Float64 = jw
    κ::Float64 = τ == 0.0 ? 0.0 : 1.0/τ
    ζjet = [-κ*Ξ*tanh(-y/Δθ) for y in LinRange(-ly/2.0,ly/2.0,2*ny-1)]
    fftshift(fft(ζjet))
end

function bcoeffs(lx::Float64,ly::Float64,nx::Int,ny::Int,β::Float64,τ::Float64=0.0,νn::Float64=0.0)
    B = zeros(ComplexF64,2*ny-1,nx)
    α::Int = 2
    kxmax::Float64 = 2.0*Float64(pi)/lx*Float64(nx-1)
    kymax::Float64 = 2.0*Float64(pi)/ly*Float64(ny-1)
    γ::Float64 = τ == 0.0 ? 0.0 : 1.0/τ
    for m = 0:1:nx-1
        nmin = m == 0 ? 1 : -(ny-1)
        for n=nmin:1:ny-1
            kx::Float64 = 2.0*Float64(pi)*Float64(m)/lx
            ky::Float64 = 2.0*Float64(pi)*Float64(n)/ly
            B[n+ny,m+1] = -γ + im*β*kx/(kx^2 + ky^2) - νn*((kx^2 + ky^2)/(kxmax^2 + kymax^2))^(2*α)
        end
    end
    B
end

function ccoeffs(lx::Float64,ly::Float64,nx::Int,ny::Int)

    M::Int = nx - 1
    N::Int = ny - 1

    Cp = zeros(Float64,2*ny-1,nx,2*ny-1,nx)
    Cm = zeros(Float64,2*ny-1,nx,2*ny-1,nx)

    # ++ interactions note: +0 has only (0,+n)
    for m1=1:1:M
        for n1=-N:1:N
            for m2=0:1:min(m1,M-m1)

                n2min = m2 == 0 ? 1 : -N
                for n2=max(n2min,-N-n1):1:min(N,N-n1)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    if m1 == m2
                        Cp[n2+ny,m2+1,n1+ny,m1+1] = -(px*qy - qx*py)/(px^2 + py^2)
                    else
                        Cp[n2+ny,m2+1,n1+ny,m1+1] = -(px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))
                    end

                end
            end
        end
    end

    # +- interactions note: - includes (0,-n) because it is conj(0,n)
    for m1=1:1:M
        for n1=-N:1:N
            for m2=0:1:m1

                n2min = m2 == 0 ? 1 : -N
                n2max = m2 == m1 ? n1 - 1 : N
                for n2=max(n2min,n1-N):1:min(n2max,n1+N)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    Cm[n2+ny,m2+1,n1+ny,m1+1] = (px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))

                end
            end
        end
    end
    Cp,Cm
end

function ccoeffs(lx::Float64,ly::Float64,nx::Int,ny::Int,Λ::Int)

    M::Int = nx - 1
    N::Int = ny - 1

    Cp = zeros(Float64,2*ny-1,nx,2*ny-1,nx)
    Cm = zeros(Float64,2*ny-1,nx,2*ny-1,nx)

    # L + L = L
    for m1=0:1:Λ
        n1min = m1 == 0 ? 1 : -N
        for n1=n1min:1:N
            for m2=0:1:min(m1,Λ-m1)

                n2min = m2 == 0 ? 1 : -N
                for n2=max(n2min,-N-n1):1:min(N,N-n1)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    if m1 == m2
                        Cp[n2+ny,m2+1,n1+ny,m1+1] = -(px*qy - qx*py)/(px^2 + py^2)
                    else
                        Cp[n2+ny,m2+1,n1+ny,m1+1] = -(px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))
                    end

                end
            end
        end
    end

    # L - L = L
    # note: -L should always include (0,-n)
    for m1=0:1:Λ
        n1min = m1 == 0 ? 1 : -N
        for n1=n1min:1:N

            for m2=0:1:m1

                n2min = m2 == 0 ? 1 : -N
                n2max = m2 == m1 ? n1 - 1 : N
                for n2=max(n2min,n1-N):1:min(n2max,n1+N)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    Cm[n2+ny,m2+1,n1+ny,m1+1] = (px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))

                end
            end
        end
    end

    # H - H = L
    for m1=Λ+1:1:M
        for n1=-N:1:N
            for m2=max(Λ+1,m1-Λ):1:m1

                n2max = m2 == m1 ? n1 - 1 : N
                for n2=max(-N,n1-N):1:min(n2max,n1+N)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    Cm[n2+ny,m2+1,n1+ny,m1+1] = (px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))

                end
            end
        end
    end

    # H + L = H
    for m1=Λ+1:1:M
        for n1=-N:1:N
            for m2=0:1:min(M-m1,Λ)

                n2min = m2 == 0 ? 1 : -N
                for n2=max(n2min,-N-n1):1:min(N,N-n1)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    Cp[n2+ny,m2+1,n1+ny,m1+1] = -(px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))

                end
            end
        end
    end

    # H - L = H
    # note: -L should always include (0,-n)
    for m1=Λ+1:1:M
        for n1=-N:1:N
            for m2=0:1:min(Λ,m1 - Λ - 1)

                n2min = m2 == 0 ? 1 : -N
                for n2=max(n2min,n1-N):1:min(N,n1+N)

                    px::Float64 = 2.0*Float64(pi)/lx*Float64(m1)
                    py::Float64 = 2.0*Float64(pi)/ly*Float64(n1)
                    qx::Float64 = 2.0*Float64(pi)/lx*Float64(m2)
                    qy::Float64 = 2.0*Float64(pi)/ly*Float64(n2)

                    Cm[n2+ny,m2+1,n1+ny,m1+1] = (px*qy - qx*py)*(1.0/(px^2 + py^2) - 1.0/(qx^2 + qy^2))

                end
            end
        end
    end
    Cp,Cm
end


using SparseArrays, Test
using CDCS

# Example page 4 of SeDuMi_Guide_105R5.pdf
@testset "Linear Programming example" begin
    c = [1.0, -1.0, 0.0, 0.0]
    A = [10.0 -7.0 -1.0 0.0
          1.0  0.5  0.0 1.0]
    b = [5.0, 3.0]
    sol, dual, z, status = cdcs(Matrix(A'), b, c, CDCS.Cone(0, 4), verbose=0)
    tol = 1e-2
    @test sol ≈ [47/24, 25//12, 0, 0] atol=tol rtol=tol
    @test dual ≈ [1/8, -1/4] atol=tol rtol=tol
    @test z ≈ [0, 0, 1/8, 1/4] atol=tol rtol=tol
    @test status["cost"] ≈ -1/8 atol=tol rtol=tol
end


# write your own tests here
N = 10
ph = PortHamiltonian.weak_phs1(N, 0, 1);
eigval, eigvec = eigen(ph);
num_freq = frequencies(eigval)

@test abs(num_freq[2]-pi) < 1e-13

N = 10
ph = PortHamiltonian.weak_phs2(N, 0, 1);
eigval, eigvec = eigen(ph);
num_freq = frequencies(eigval)

@test abs(num_freq[2]-pi) < 1e-13

ph = PortHamiltonian.finelem(N+1,5, 0, 1)
eigval, eigvec = eigen(ph);
num_freq = frequencies(eigval)
@test abs(num_freq[2]-pi) < 1e-13

ph = PortHamiltonian.finelem(N,5, 0, 1)
eigval, eigvec = eigen(ph);
num_freq = frequencies(eigval)
@test abs(num_freq[2]-pi*3/2) < 1e-13


using Images, TestImages
using Deconvolution
using FFTW
using Colors
using ZernikePolynomials
using FFTW
using NumberTheoreticTransforms

trim = 225:288 # need to limit size due to O(N^4) complexity
image_float = channelview(testimage("cameraman"))[trim.-110, trim.+40]
image_int = map(x -> x.:i, image_float) .|> Int64

blur_float = evaluateZernike(LinRange(-41,41,512), [12, 4, 0], [1.0, -1.0, 2.0], index=:OSA)
blur_float ./= (sum(blur_float))
blur_float = blur_float[trim, trim]
blur_int = blur_float .|> Normed{UInt8, 8} .|> x -> x.:i .|> Int64

blur_int = fftshift(blur_int)

@assert sum(blur_int) == 256 # to not change brightness

blurred_image_float = ifft(fft(image_float) .* fft(fftshift(blur_float))) |> real

(g, q, n) = (4, 274177, 64)
X = ntt(image_int, g, q)
@assert intt(X, g, q) == image_int
H = ntt( blur_int, g, q)
@assert intt(H, g, q) == blur_int
Y = mod.((X .* H), q)
y = intt(Y, g, q)
blurred_image_int = y / 2^16

using Statistics
@show maximum(abs.(blurred_image_float .- blurred_image_int))
@show mean(abs.(blurred_image_float .- blurred_image_int))
@show var(abs.(blurred_image_float .- blurred_image_int))

save("doc/src/ntt/original.jpg", Images.clamp01nan.(image_float))
save("doc/src/ntt/blurred_fft.jpg", Images.clamp01nan.(Gray.(Normed{UInt8, 8}.(blurred_image_float))))
save("doc/src/ntt/blurred_ntt.jpg", Images.clamp01nan.(Gray.(Normed{UInt8, 8}.(blurred_image_int))))


using Calculus
using QuadGK

f(k) = QuadGK.quadgk(η -> 1 / √(1 - k^2 * cos(η)^2), 0, π / 2 )[1]
g(k) = QuadGK.quadgk(η -> 1 / √(1 - k^2 * sin(η)^2), 0, π / 2 )[1]
δ(x) = f(x) - g(x)

for x = 0.01 : 0.02 : 0.99
    del = δ(x)
    println("δ($x) = $del")
end


"""
    swrfstar(data::Array{<:Real,2}, target::Array{<:Integer, 1}, m::Signed=-1,
                  dist_func::Any=(e1, e2) -> sum(abs.(e1 .- e2), dims=2);
                  f_type::String="continuous")::Array{Float64,1}

Compute feature weights using SWRF* algorithm. The f_type argument specifies whether the features are continuous or discrete
and can either have the value of "continuous" or "discrete".

---
# Reference:
- Matthew E. Stokes and Shyam Visweswaran.
Application of a spatially-weighted Relief algorithm for ranking genetic predictors of disease.
BioData mining, 5(1):20–20, Dec 2012. 23198930[pmid].
"""
function swrfstar(data::Array{<:Real,2}, target::Array{<:Integer, 1}, m::Signed=-1,
                  dist_func::Any=(e1, e2) -> sum(abs.(e1 .- e2), dims=2);
                  f_type::String="continuous")::Array{Float64,1}

    # Initialize feature weights vector.
    weights = zeros(Float64, 1, size(data, 2))

    # Compute vectors of maximum and minimum feature values.
    max_f_vals = maximum(data, dims=1)
    min_f_vals = minimum(data, dims=1)

    # Sample m examples without replacement.
    sample_idxs = StatsBase.sample(1:size(data, 1), if (m==-1) size(data,1) else m end, replace=false)
    if (m == -1) m = size(data, 1) end

    # Compute probabilities of class values in training set.
    classes_map = countmap(target)
    num_samples = length(target)
    p_classes = Array{Float64}(undef, length(keys(classes_map)), 2)
    @inbounds for (idx, p) = enumerate(classes_map)  # Compute matrix representation of probabilities.
        p_classes[idx, :] = [p[1], p[2]/num_samples]
    end

    # Go over sampled indices.
    @inbounds for idx = sample_idxs

        # Get samples from same class and different classes.
        samples_same_class = data[(target .== target[idx]) .& (1:length(target) .!= idx), :]
        samples_other_class = data[target .!= target[idx], :]

        # Compute distances to examples with same class value.
        distances_same = dist_func(data[idx:idx, :], samples_same_class)

        # Get class values of examples with different class value.
        target_other = target[target .!= target[idx]]

        # Compute t and u parameter values.
        t_same = Statistics.mean(distances_same)
        u_same = Statistics.std(distances_same)

        # Compute distances to examples with different class value.
        distances_other = dist_func(data[idx:idx, :], samples_other_class)

        # Compute t and u parameter values.
        t_other = Statistics.mean(distances_other)
        u_other = Statistics.std(distances_other)

        # Compute weights for examples from same class.
        neigh_weights_same = 2.0 ./ (1 .+ exp.(-(t_same.-distances_same)/(u_same/4.0 + eps(Float64))) .+ eps(Float64)) .- 1

        # Compute weights for examples from different class.
        neigh_weights_other = 2.0 ./ (1 .+ exp.(-(t_other.-distances_other)/(u_other/4.0 + eps(Float64))) .+ eps(Float64)) .- 1

        # Get probabilities of classes not equal to class of sampled example.
        p_classes_other = p_classes[p_classes[:, 1] .!= target[idx], 2]

        # Get other classes and compute diff sum weights for examples from different classes.
        classes_other = p_classes[p_classes[:, 1] .!= target[idx], 1]
        p_weights = p_classes_other./(1 .- p_classes[p_classes[:, 1] .== target[idx], 2] .+ eps(Float64))

        # Map weights vector of classes of samples with different class values to construct weights vector.
        weights_map = hcat(classes_other, p_weights)
        weights_mult = [weights_map[findall(weights_map[:, 1] .== t), 2][1] for t in target_other]

        ### Weights Update ###

        if f_type == "continuous"
            # If features continuous.

            # Penalty term
            penalty = sum(neigh_weights_same.*(abs.(data[idx:idx, :] .- samples_same_class)./(max_f_vals .- min_f_vals .+ eps(Float64))), dims=1)

            # Reward term
            reward = sum(weights_mult.*(neigh_weights_other.*(abs.(data[idx:idx, :] .- samples_other_class)./(max_f_vals .- min_f_vals .+ eps(Float64)))), dims=1)

            # Weights update
            weights = weights .- penalty./(m*size(samples_same_class, 1) + eps(Float64)) .+ reward./(m*size(samples_other_class, 1) + eps(Float64))

        elseif f_type == "discrete"
            # If features discrete.

            # Penalty term
            penalty = sum(neigh_weights_same.*Int64.(data[idx:idx, :] .!= samples_same_class), dims=1)

            # Reward term
            reward = sum(neigh_weights_other.*weights_mult .* Int64.(data[idx:idx, :] .!= samples_other_class), dims=1)

            # Weights update
            weights = weights .- penalty./(m*size(samples_same_class, 1) + eps(Float64)) .+ reward./(m*size(samples_other_class, 1) + eps(Float64))

        else
            throw(DomainError(f_type, "f_type can only be equal to \"continuous\" or \"discrete\"."))
        end

        ######################

    end

    # Return computed feature weights.
    return vec(weights)

end


using Distributed

# Number of matrices in the sum
nmats = 2
addprocs(nmats)

@everywhere begin
  using LinearAlgebra
  using KrylovKit

  BLAS.set_num_threads(1)

  include(joinpath("src", "DistributedSums.jl"))
  include(joinpath("src", "MatrixMaps.jl"))

  # Matrix dimension
  n = 100

  # Starting vector that we will apply
  # the matrix map to
  v0 = randn(n)
end

# Make a distributed sum of MatrixMaps, which is
# just a Matrix with special multiplication
# notation that can work with KrylovKit
A = DistributedSum(_ -> MatrixMap(randn(n, n)), nmats)

Av0 = A(v0)

@show norm(Av0 - sum(fetch.(A))(v0))

# Solve eigenvector equation
λ⃗, v⃗ = eigsolve(A, v0, 1, :LM)
λ = λ⃗[1]
v = v⃗[1]

@show norm(A(v) - λ * v)


using DiscreteEvents, Printf, Distributions, Random

const μ = 1/3       # service rate
const λ = 0.9       # arrival rate
count = [0]         # a job counter

# describe the server process
function serve(clk::Clock, id::Int, input::Channel, output::Channel, X::Distribution)
    job = take!(input)
    print(clk, @sprintf("%6.3f: server %d serving customer %d\n", tau(clk), id, job))
    @delay clk X
    print(clk, @sprintf("%6.3f: server %d finished serving %d\n", tau(clk), id, job))
    put!(output, job)
end

# model the arrivals
function arrive(c::Clock, input::Channel, cust::Vector{Int})
    cust[1] += 1
    @printf("%6.3f: customer %d arrived\n", tau(c), cust[1])
    put!(input, cust[1])
end

Random.seed!(123)  # set random number seed
clock = Clock()    # create a clock
input = Channel{Int}(Inf)
output = Channel{Int}(Inf)
for i in 1:3       # start three server processes
    @process serve(clock, i, input, output, Exponential(1/μ))
end
# create a repeating event for 10 arrivals
@event arrive(clock, input, count) every Exponential(1/λ) 10
@run! clock 20     # run the clock


include(joinpath(dirname(@__DIR__), "src", "BestApproximation.jl")); using .BestApproximation
using Test

@testset "BestApproximation.jl" begin
    @test best_approx(factorial(big(26)), 1:4) == (2, 88)
    @test best_approx(factorial(big(26)), 4) == (2, 88)
    @test best_approx(factorial(big(26)), 10) == (5, 38)
    @test best_approx(factorial(big(26)), 100) == (46, 16)
    @test best_approx(factorial(big(26)), 1000) == (904, 9)
    @test best_approx(factorial(big(26)), 10000) == (904, 9)
    @test best_approx(factorial(big(26)), 100000) == (27181, 6)
    @test best_approx(factorial(big(26)), 1000000) == (209470, 5)
    @test best_approx(1234567890, 94) == (5, 13)
    @test best_approx(12345678901234567890, 94) == (3, 40)
    @test best_approx(123456789012345678901234567890, 94878) == (70564, 6)
    @test best_approx(factorial(big(26)), 1:4, 1:100) == (2, 88)
    @test best_approx(factorial(big(26)), 4, 100) == (2, 88)
    @test best_approx(factorial(big(26)), 10, 100) == (5, 38)
end


"""
    AmoebaFiber2D(f::AbstractPolynomial, w=(0., 0.))

Construct the fiber `AmoebaFiber2D` of the bivariate polynomial `f` at `w`.
"""
function AmoebaFiber2D(p::MP.AbstractPolynomialLike{S}, w=(0.0, 0.0)) where S
    vars = MP.variables(p)
    @assert (length(vars) == 2) "Expected a bivariate polynomial, but got $(p)."

    T = realified_coefficient_type(S)
    f_re, f_im = SP.Polynomial.(1.0 .* realify(p))
    AmoebaFiber2D(f_re, f_im, w)
end

function AmoebaFiber2D(f_re::SP.Polynomial{T}, f_im::SP.Polynomial, w=(0.0, 0.0)) where T
    v = exp.(w)
    U = zeros(T, 2, 4)
    x = SVector{4, Float64}(0.0, 0.0, 0.0, 0.0)

    exponents = SP.exponents.((f_re, f_im))
    original_coefficients = copy.(SP.coefficients.((f_re, f_im)))

    AmoebaFiber2D(f_re, f_im, v, U, x, exponents, original_coefficients)
end

function update_fiber!(F::AmoebaFiber2D, w)
    w1, w2 = w
    F.v = (exp(w1), exp(w2))
    condition!(F)
    F
end

function condition!(fiber::AmoebaFiber2D)
    fiber.f_re.coefficients .= fiber.original_coefficients[1]
    fiber.f_im.coefficients .= fiber.original_coefficients[2]

    n1 = normalizer_amoeba_2d(fiber.exponents[1], fiber.f_re.coefficients, fiber.v)
    n2 = normalizer_amoeba_2d(fiber.exponents[2], fiber.f_im.coefficients, fiber.v)
    normalizer = inv(max(n1, n2))
    SP.scale_coefficients!(fiber.f_re, normalizer)
    SP.scale_coefficients!(fiber.f_im, normalizer)
    nothing
end

function normalizer_amoeba_2d(E, c, v)
    v1, v2 = v
    normalizer = sum(1:size(E, 2)) do j
        m1 = E[1,j]+E[3,j]
        m2 = E[2,j]+E[4,j]
        abs(c[j]) * Base.FastMath.pow_fast(v1, m1) *
        Base.FastMath.pow_fast(v2, m2)
    end
end

function phi!(fiber::AmoebaFiber2D, θ)
    v1, v2 = fiber.v
    x = SVector(v1 * cos(θ[1]), v2 * cos(θ[2]), v1 * sin(θ[1]), v2 * sin(θ[2]))
    fiber.x = x
    x
end

function evaluate(fiber::AmoebaFiber2D{T}, θ) where T
    x = phi!(fiber, θ)
    SVector(SP.evaluate(fiber.f_re, x), SP.evaluate(fiber.f_im, x))
end

function jacobian(fiber::AmoebaFiber2D, θ, precomputed=false)
    x = precomputed ? fiber.x : phi!(fiber, θ)
    U = fiber.U
    @inbounds begin
        ∇f_re = SP.gradient(fiber.f_re, x)
        for i=1:4
            U[1, i] = ∇f_re[i]
        end
        U[2, 3] = U[1, 1]
        U[2, 4] = U[1, 2]
        U[2, 1] = -U[1, 3]
        U[2, 2] = -U[1, 4]
    end

    @inbounds out = begin
        x3, x4 = -x[3], -x[4]
        ∂re∂θ1 = U[1, 1] * x3
        ∂re∂θ2 = U[1, 2] * x4
        ∂im∂θ1 = U[2, 1] * x3
        ∂im∂θ2 = U[2, 2] * x4

        ∂re∂θ1 = muladd(U[1, 3], x[1], ∂re∂θ1)
        ∂re∂θ2 = muladd(U[1, 4], x[2], ∂re∂θ2)
        ∂im∂θ1 = muladd(U[2, 3], x[1], ∂im∂θ1)
        ∂im∂θ2 = muladd(U[2, 4], x[2], ∂im∂θ2)

        # SMatrix has column-major storage
        SMatrix{2,2}(∂re∂θ1, ∂im∂θ1, ∂re∂θ2, ∂im∂θ2)
    end

    out
end


function gpdps_potts!(u,v,data,params;
                      uref=0*u,
                      vref=0*v)
    σ,τ,ω,maxit     = params
    α,γ,f,isotropic = data

    fᵧ = t -> 2t^2/(2t^2 + γ)
    if isotropic
        Fᵧ = (u₁,u₂) -> sum(@. fᵧ(sqrt(u₁^2+u₂^2)))
        κ₁ = (u₁,u₂,v₁,v₂) -> 2(1-u₁*v₁-u₂*v₂)
        κ₂ = (u₁,u₂,v₁,v₂) -> 2(1-u₁*v₁-u₂*v₂)
    else
        Fᵧ = (u₁,u₂) -> sum(@. fᵧ(u₁)+fᵧ(u₂))
        κ₁ = (u₁,u₂,v₁,v₂) -> 2(1-u₁*v₁)
        κ₂ = (u₁,u₂,v₁,v₂) -> 2(1-u₂*v₂)
    end

    # preallocate
    uold = similar(u)
    ξ    = similar(u)
    v₁,v₂ = view(v,:,:,1), view(v,:,:,2)
    u₁,u₂ = similar(u), similar(u)
    z₁,z₂ = similar(u), similar(u)
    Kᵤ,Kᵥ = similar(u), similar(v)
    for it = 1:maxit
        # Kx = ∇ᵀ(κ(∇u⋅v)⋅v)
        u₁,u₂ = ∇!(u₁,u₂,u)
        @. z₁ = κ₁(u₁,u₂,v₁,v₂)*v₁
        @. z₂ = κ₂(u₁,u₂,v₁,v₂)*v₂
        Kᵤ = ∇ᵀ!(Kᵤ,z₁,z₂)

        # print output
        if mod(it,1000) == 0
            J   = norm(u.-f)^2/(2α) + Fᵧ(u₁,u₂)
            err = norm(u.-uref)^2+norm(v.-vref)^2
            @printf("it %6d: J = %1.6e, err = %1.6e\n", it, J, err)
            imshow(u', colormap=2)
        end

        # x = prox_τG(x-τKx)
        @. uold = u
        @. u = (u-τ*Kᵤ+τ/α*f)/(1+τ/α)
        @. ξ = u + ω*(u-uold)

        # Ky = κ(∇ξ⋅v)⋅∇ξ
        u₁,u₂ = ∇!(u₁,u₂,ξ)
        @. Kᵥ[:,:,1] = κ₁(u₁,u₂,v₁,v₂)*u₁
        @. Kᵥ[:,:,2] = κ₂(u₁,u₂,v₁,v₂)*u₂

        # y = prox_σF*(y+σKy)
        @. v = (v+σ*Kᵥ)/(1+γ*σ)
    end
end

function ∇!(u₁,u₂,u)
    @. @views begin
        u₁[1:(end-1), :] = u[2:end, :] - u[1:(end-1), :]
        u₁[end, :] = 0

        u₂[:, 1:(end-1)] = u[:, 2:end] - u[:, 1:(end-1)]
        u₂[:, end] = 0
    end
    return u₁, u₂
end

function ∇ᵀ!(v,v₁,v₂)
    @. @views begin
        v[2:(end-1), :] = v₁[1:(end-2), :] - v₁[2:(end-1), :]
        v[1, :] = -v₁[1, :]
        v[end, :] = v₁[end-1, :]

        v[:, 2:(end-1)] += v₂[:, 1:(end-2)] - v₂[:, 2:(end-1)]
        v[:, 1] += -v₂[:, 1]
        v[:, end] += v₂[:, end-1]
    end
    return v
end

function estimate_steplengths(f,α,γ,isotropic;expected_max_jump,γ_for_my)
    dyn_range = maximum(abs.(f))-minimum(abs.(f))
    expected_max_jump = 1.0
    γ_for_my = 10

    L = √8
    R_K = 2*L

    if isotropic
        mx = √2*expected_max_jump*dyn_range
    else
        mx = expected_max_jump*dyn_range
    end

    fmy = t -> 2t/(2t^2+γ_for_my)
    # fmy has its maximum in the interval [0, mx] either at mx or at √(γ_for_my/2)
    my = max(fmy(mx), fmy(√(γ_for_my/2)))

    # This estimate ignores the difficult-to-estimate ρ_y term, essentially
    # assuming ρ_y=0. To offset this, we set δ≪1.
    δ = 0.1
    μ = (1+δ)/2
    γ̃ = γ/2
    α̃ = α*2
    ξx = 1/α-1/α̃
    @assert(ξx>2*L*my^2)
    λx = 2*L^2*my^4/(ξx-2*L*my^2)
    λy = mx^2
    τ₀ = δ/λx
    τ₁ = 2*γ̃*α̃/(λy+√(λy^2+4*γ̃*α̃*(R_K^2/(1-μ)+2*λy/α̃)))
    τ  = min(τ₀, τ₁)
    σ = τ/(γ̃*α̃)
    ω = 1/(1+τ/α̃)

    return τ,σ,ω
end


# The SteinMaternPowerKernel is given by
#
# k(x, y) = ((1 + ||x||^alpha)*(1 + ||y||^alpha))^-1 * km(x,y)
#
# where km is the matern kernel.

mutable struct SteinMaternPowerKernel <: SteinMaternWeightedKernel
    # alpha is the exponent on the norms
    alpha::Float64
    # gamma is the rate parameter for the matern kernel
    gamma::Float64
end

# default is use gamma = sqrt{3} and alpha = 2.5
SteinMaternPowerKernel() = SteinMaternPowerKernel(2.5, sqrt(3.0))

# default is use gamma sqrt{3}
SteinMaternPowerKernel(alpha::Float64) = SteinMaternPowerKernel(alpha, sqrt(3.0))

# The term (1 + ||z||^alpha)^-1
function _Qz(ker::SteinMaternPowerKernel, z::Array{Float64, 1})
    Qz = 1.0 + norm(z,2)^ker.alpha
    1.0 / Qz
end

function Q(ker::SteinMaternPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    _Qz(ker, x) * _Qz(ker, y)
end

function gradxQ(ker::SteinMaternPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    fx = _Qz(ker, x)
    fy = _Qz(ker, y)
    mx = norm(x, 2)^(ker.alpha - 2)
    -x * ker.alpha * fx^2 * fy * mx
end

function gradxyQ(ker::SteinMaternPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    mx = norm(x, 2)^(ker.alpha - 2.0)
    my = norm(y, 2)^(ker.alpha - 2.0)
    dot(x, y) * ker.alpha^2 * mx * my * Q(ker, x, y)^2
end


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

lower = -100.
upper = 100.

max_iter = 50

# example 0
true_d = MixtureModel(Normal, [(5.0, 2), (10.0, 2), (20.0, 1), (-2, 1)], [1.0/3, 1.0/4, 1.0/4, 1.0/6])
ll(x) = logpdf(true_d, x)

init_approx_d =  Normal(0, 40)
x = linspace(-10,30,500)

include("run_laplacian_gradient_boosting_exclusive.jl")
show()

# example 1
true_d = MixtureModel(Normal, [(-2.0, 1), (2.0, 1)], [0.5, 0.5])
ll(x) = logpdf(true_d, x)

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

include("run_laplacian_gradient_boosting_exclusive.jl")
show()

println("Done.")

# example 2
true_d = MixtureModel(Normal, [(-3.5, 1), (-1.0, 0.4), (1.0, 0.5), (6.0, 1)], [0.2, 0.3, 0.3, 0.2])
ll(x) = logpdf(true_d, x)
x = linspace(-10,10,500)

include("run_laplacian_gradient_boosting_exclusive.jl")
show()

println("Done.")


@testset "3 stages with $solver" for solver in lp_solvers
    include("prob5.2_data.jl")

    numScen = 2
    M = StructuredModel(num_scenarios=numScen)

    x = Vector{Vector{JuMP.Variable}}(undef, 2)
    v = Vector{Vector{JuMP.Variable}}(undef, 2)
    y = Vector{Matrix{JuMP.Variable}}(undef, 2)

    x[1] = @variable(M, [1:n], lowerbound=0)
    v[1] = @variable(M, [1:n], lowerbound=0)
    @constraints M begin
        x[1] .== v[1]
    end
    @objective(M, Min, dot(ic, v[1]))

    for s in 1:numScen
        M2 = StructuredModel(parent=M, prob=p2[s], id=s)
        y[1] = @variable(M2, [1:n, 1:m], lowerbound=0)
        x[2] = @variable(M2, [1:n], lowerbound=0)
        v[2] = @variable(M2, [1:n], lowerbound=0)
        @constraints M2 begin
            x[2] .== x[1] + v[2]
            demand[j=1:m], sum(y[1][:,j]) == D2[j,s]
            ylim[i=1:n], sum(y[1][i,:]) <= x[1][i]
        end
        @objective(M2, Min, dot(ic, v[2]) + dot(C, y[1] * T))
        for S in 1:numScen
            M3 = StructuredModel(parent=M2, prob=p2[S], id=S)
            y[2] = @variable(M3, [1:n, 1:m], lowerbound=0)
            @constraints M3 begin
                demand[j=1:m], sum(y[2][:,j]) == D2[j,S]
                ylim[i=1:n], sum(y[2][i,:]) <= x[2][i]
            end
            @objective(M3, Min, dot(C, y[2] * T))
        end
    end

    if !isdefined(@__MODULE__, :testniter_3stages)
        include("prob5.2_3stages_niter.jl")
    end
    fulltest(M, 3, 406712.49, [2986,0,7329,854], 402593.71614, 17319.095064, testniter_3stages, solver)
    # root = model2lattice(M, 3, solver, cutmode)
    # sol = SDDP(root, 3, cutmode, :All, verbose)
    #
    # v11value = sol.sol[1:4]
    # @test sol.status == :Optimal
    # @test abs(sol.objval - 406712.49) < 0.1
    # @test norm(v11value - [2986,0,7329,854]) < 0.1
    # SDDPclear(M)
end


for N in [Float64, Float32, Rational{Int}]
    H1 = BallInf(N[0, 0], N(1))
    H2 = BallInf(N[-1//10, 0], N(1))
    H3 = Hyperrectangle(N[-16//10, -2], N[1//2, 1//2])

    for p in N[1, 2, Inf]
        @test distance(H1, H2; p=p) == distance(H2, H1; p=p) == N(0)
        @test distance(H1, H3; p=p) == distance(H3, H1; p=p) ≈ norm(N[1//10, 1//2], p)
        @test distance(H2, H3; p=p) == distance(H3, H2; p=p) ≈ norm(N[0, 1//2], p)
    end
end


numbers = readlines("day3/numbers.txt")
numbers = permutedims(hcat(collect.(numbers)...))
numbers = parse.(Int, numbers)

function find_rates(numbers)
    γ_bits = sum(numbers, dims = 1) .> size(numbers, 1) / 2
    ε_bits = map(~, γ_bits)  # bitwise not
    γ = parse(Int, join(convert(Array{Int}, γ_bits)), base = 2)
    ε = parse(Int, join(convert(Array{Int}, ε_bits)), base = 2)
    γ, ε
end

γ, ε = find_rates(numbers)
p1 = γ * ε

@show p1

function find_rating(numbers, keep_bit)
    more_numbers = true
    bit_idx = 1
    while more_numbers
        bits = numbers[:, bit_idx]
        numbers = numbers[bits .== keep_bit(bits), :]
        more_numbers = size(numbers, 1) > 1
        bit_idx += 1
    end
    parse(Int, join(numbers), base = 2)
end

most_common_bit(bits) = sum(bits) >= length(bits) / 2 ? 1 : 0

least_common_bit(bits) = sum(bits) >= length(bits) / 2 ? 0 : 1

oxy = find_rating(numbers, most_common_bit)
co2 = find_rating(numbers, least_common_bit)
p2 = oxy * co2

@show p2


####
#### internal methods, not part of the API
####

####
#### scaled problem
####

"""
Given `Ê`, solves
```math
Ê = \\log(∑ᵢ \\exp(zᵢ + ϵᵢ Ĉ))
```
for `Ĉ`. It is assumed that called does all the sanity checks for parameters.

# Important

The original problem needs to be mapped into this one with

| original problem    | scaled problem |
|---------------------|----------------|
| ``(1 - σ) Ê``       | ``Ê``          |
| ``(1 - σ) Ĉ``       | ``Ĉ``          |
| ``Ω̂ᵢ + (1 - σ) p̂ᵢ`` | ``zᵢ``         |
| ``ϵᵢ``              | ``ϵᵢ`` (same)  |
"""
struct ScaledProblem{N,Tz<:Real,Tϵ<:Real}
    zs::SVector{N,Tz}
    ϵs::SVector{N,Tϵ}
end

Broadcast.broadcastable(problem::ScaledProblem) = Ref(problem)

# workaround for https://github.com/JuliaStats/LogExpFunctions.jl/issues/33
function _logsumexp(x::SVector)
    A = maximum(x)
    log(sum(exp.(x .- A))) + A
end

"""
$(SIGNATURES)

Solve the scaled `problem` using Newton's method. `atol` is the absolute tolerance (should
be positive) between the left and right hand sides, and `M` is the maximum number of
iterations.
"""
function calculate_Ĉ_newton(problem::ScaledProblem, Ê::Real, atol, M = 100)
    @unpack zs, ϵs = problem
    Ĉ = (Ê - _logsumexp(zs)) / mean(ϵs)
    for _ in 1:M
        x = zs .+ ϵs .* Ĉ
        # calculate logsumexp and “weighted softmax” for f′ at the same time
        # it works because all inputs are finite
        A = maximum(x)
        X = exp.(x .- A)        # numerical stability correction
        f = log(sum(X)) + A - Ê
        abs(f) ≤ atol && return Ĉ
        f′ = dot(X, ϵs) / sum(X)
        Δ = f / f′
        Ĉ′ = Ĉ - Δ
        Δ ≠ 0 && Ĉ′ == Ĉ && return Ĉ # no change because abs(Δ) < eps(Ĉ)
        Ĉ = Ĉ′
    end
    error("no solution after $M iterations")
end

###
### consumption aggregator
###

"""
$(SIGNATURES)

Check parameters, throw an error message for invalid ones.
"""
function check_σ_ϵs(σ, ϵs)
    @argcheck σ > 0 DomainError
    @argcheck σ ≠ 1 DomainError
    @argcheck all(ϵs .> 0) DomainError
    nothing
end

"""
Default tolerance for the residual in `Ê = ...` for calculating `Ĉ`.
"""
const DEFAULT_TOL = 1e-10

# `T` is for a separate, conditional dispatch path to catch `ForwardDiff.Dual`.
function calculate_Ĉ(::Type{T}, Ê, σ, Ω̂s, ϵs, p̂s;
                     tol = DEFAULT_TOL,
                     skipcheck::Bool = false) where {T <: Real}
    skipcheck || check_σ_ϵs(σ, ϵs)
    scaled_problem = ScaledProblem(Ω̂s .+ (1 - σ) .* p̂s, ϵs)
    atol = (1 + abs(Ê)) * abs(tol)
    calculate_Ĉ_newton(scaled_problem, (1 - σ) * Ê, atol) / (1 - σ)
end

"""
$(SIGNATURES)

Calculate `Ĉ` from parameters.

# Arguments

- `skipcheck`: skip sanity checks

- `tol`: (relative) tolerance, translated to absolute tolerance `(1 + abs(Ê)) * abs(tol)`.
"""
function calculate_Ĉ(Ê::TÊ, σ::Tσ, Ω̂s::AbstractVector{TΩ̂}, ϵs::AbstractVector{Tϵ},
                     p̂s::AbstractVector{Tp̂};
                     tol = DEFAULT_TOL,
                     skipcheck::Bool = false) where {Tσ,TΩ̂,Tϵ,TÊ,Tp̂}
    T = promote_type(Tσ,TΩ̂,Tϵ,TÊ,Tp̂)
    calculate_Ĉ(T, Ê, σ, Ω̂s, ϵs, p̂s; tol, skipcheck)
end

###
### partial derivatives
###
### Coded here internally so that they can be reused for various AD frameworks.

"""
$(SIGNATURES)

Calculate and return

1. `Zs`, where each element is ``Zᵢ ∝ exp(zᵢ)``, scaled to protect from under- and overflow,
where ``zᵢ = (Ĉ ϵᵢ + pᵢ) (1 - σ) + Ωᵢ``,

2. `sum(Zs .* ϵs)`.

These are useful for partial derivatives.
"""
function calculate_Zs_∑Zϵ(Ĉ, Ê, σ, Ω̂s, ϵs, p̂s)
    zs = (Ĉ .* ϵs .+ p̂s) .* (1 - σ) .+ Ω̂s
    Zs = exp.(zs .- maximum(zs))
    Zs, dot(Zs, ϵs)
end

"""
$(SIGNATURES)

Calculate `∂Ĉ/∂Ê`.
"""
@inline calculate_∂Ê(Zs, ∑Zϵ) = sum(Zs) / ∑Zϵ

"""
$(SIGNATURES)

Calculate `∂Ĉ/∂p̂ᵢ` for each price in `p̂s`.

Also note that the result can be reused, as ``∂Ĉ/∂ϵᵢ=Ĉ⋅∂Ĉ/∂p̂ᵢ`` and
``∂Ĉ/∂Ω̂ᵢ=∂p̂ᵢ/(1-σ)``.
"""
@inline calculate_∂p̂s(Zs, ∑Zϵ) = Zs ./ -∑Zϵ

"""
$(SIGNATURES)

Calculate `∂Ĉ/∂σ`.
"""
@inline function calculate_∂σ(Zs, ∑Zϵ, Ĉ, Ê, σ, ϵs, p̂s)
    (Ĉ + sum(Zs .* (p̂s .- Ê)) / ∑Zϵ) / (1 - σ)
end


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ b17843a0-05ae-11ec-1825-8b2b03bd5811
begin
    using Pkg
    Pkg.activate("..")
    using Revise
    using Javis
    using JUtils
    using Colors
    using Random
    using ROCAnalysis
    using GLM
    using DataFrames
    using StatsPlots
    # import Luxor
    import MLJ: make_regression
end

# ╔═╡ 5e3a6b7d-f15d-4541-9b48-7c3c3ce5ddd4
begin
    X, y = make_regression(
        100,
        1,
        noise = 0.0,
        sparse = 0.0,
        outliers = 0.0,
        binary = true,
        as_table = false,
    )
    df = DataFrame(x = (X[:, 1]), y = y)
    log = glm(@formula(y ~ x), df, Binomial())
    preds = predict(log)
    tar = preds[df.y .== 1.0]
    nontar = preds[df.y .== 0.0]
    myroc = roc(tar, nontar)

    myplot = plot(myroc.pmiss, 1 .- myroc.pfa, label = "")
    plot!(myplot, [0, 1], [0, 1], label = "")
end

# ╔═╡ e0bd9a55-5e4b-4d13-adf6-24ccd11a0384
@draw begin
    draw_roc(Point(100, 0), tar, nontar, 0.05)
end 700 500

# ╔═╡ f9061057-61fd-492f-9624-4aa44db95cdd
begin
    function ground(args...)
        background("white")
        sethue("black")
    end
end

# ╔═╡ 61a37185-92fe-4b31-b609-43167b9c3f76
begin
    scatter(sort(nontar), bins = 100)
    scatter!(sort(tar), bins = 100)
end

# ╔═╡ 54de45e1-f204-4521-a31c-3ce5e2e76a7d
begin
    n_frames = 200
    n_steps = 20
    step_size = n_frames ÷ n_steps

    size = 300
    global_shift = Point(100, 50)

    my_video = Video(800, 800)

    Background(1:n_frames, ground)

    max_pred = maximum(preds)

    for (frame_n, p) in zip(1:step_size:n_frames, range(0.0, 1.0, length = n_steps))
        Object(
            frame_n:(frame_n + step_size - 1),
            @JShape begin
                draw_roc(global_shift, tar, nontar, p)
            end
        )
    end

    margin = 50

    Object(
        1:n_frames,
        @JShape begin
            rocpoints = Point.(myroc.pmiss .* size, myroc.pfa .* size)
            Javis.translate(-midpoint(rocpoints[1], rocpoints[end]))
            bottomleft = Point(rocpoints[1].x, rocpoints[end].y - margin)
            topright = Point(rocpoints[end].x, bottomleft.y - margin - max_pred * size / 2)
            Javis.translate(global_shift)
            newidxs = sortperm(preds)
            barchart(
                preds[newidxs],
                boundingbox = Luxor.BoundingBox([bottomleft, topright]),
                margin = 0,
                barfunction = (values, i, lowpos, highpos, barwidth, scaledvalue) -> begin
                    @layer begin
                        height = (topright.y - bottomleft.y)
                        y[newidxs][i] ? sethue("blue") : sethue("red")
                        circle(Point(lowpos.x, lowpos.y + values[i] * height), 3, :fill)
                    end
                end,
            )
        end
    )

    for (frame_n, p) in zip(1:step_size:n_frames, range(0.0, 1.0, length = n_steps))
        Object(
            frame_n:(frame_n + step_size - 1),
            @JShape begin
                rocpoints = Point.(myroc.pmiss .* size, myroc.pfa .* size)
                Javis.translate(-midpoint(rocpoints[1], rocpoints[end]))
                Javis.translate(global_shift)
                bottomleft = Point(rocpoints[1].x, rocpoints[end].y - margin)
                topright =
                    Point(rocpoints[end].x, bottomleft.y - max_pred * size / 2 - margin)

                bottomright = Point(topright.x, bottomleft.y)

                line(bottomleft, Point(bottomleft.x, topright.y - margin))
                line(bottomleft, Point(topright.x + margin, bottomleft.y))

                height = (topright.y - bottomleft.y)
                shift = p * height
                line(bottomleft + (0, shift), bottomright + (0, shift), :stroke)
                setcolor(sethue("blue")..., 0.4)
                box(bottomleft, bottomright + (0, shift), :fill)
                setcolor(sethue("red")..., 0.4)
                box(bottomleft + (0, shift), bottomright + (0, height), :fill)
            end
        )
    end

    render(my_video, pathname = "../output/roc_animation.gif")
end

# ╔═╡ Cell order:
# ╠═b17843a0-05ae-11ec-1825-8b2b03bd5811
# ╠═5e3a6b7d-f15d-4541-9b48-7c3c3ce5ddd4
# ╠═e0bd9a55-5e4b-4d13-adf6-24ccd11a0384
# ╟─f9061057-61fd-492f-9624-4aa44db95cdd
# ╠═61a37185-92fe-4b31-b609-43167b9c3f76
# ╠═54de45e1-f204-4521-a31c-3ce5e2e76a7d


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

estimate a covariance matrix that satisfies some given prior information and
minimizes a given convex spectral function

p ∈ 𝕊ᵈ is the covariance variable
minimize    f(p)                    (note: enforces p ⪰ 0)
subject to  tr(p) = 1               (normalize)
            B p = b                 (prior info as equalities)
            C p ≤ c                 (prior info as inequalities)
where f is a convex spectral function
=#

struct CovarianceEstJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    d::Int
    ext::MatSpecExt # formulation specifier
end

function build(inst::CovarianceEstJuMP{T}) where {T <: Float64}
    d = inst.d
    @assert d >= 1
    @assert is_domain_pos(inst.ext)

    p0 = randn(T, d, d)
    p0 = p0 * p0' + I / 2
    p0 ./= tr(p0)
    vec_dim = Cones.svec_length(d)
    p0_vec = zeros(T, vec_dim)
    Cones.smat_to_svec!(p0_vec, p0, sqrt(T(2)))

    model = JuMP.Model()
    JuMP.@variable(model, p[1:d, 1:d], Symmetric)
    JuMP.@constraint(model, tr(p) == 1)
    p_vec = zeros(JuMP.AffExpr, vec_dim)
    Cones.smat_to_svec!(p_vec, one(T) * p, sqrt(T(2)))

    # convex objective
    JuMP.@variable(model, epi)
    JuMP.@objective(model, Min, epi)
    add_homog_spectral(inst.ext, d, vcat(1.0 * epi, p_vec), model)

    # linear prior constraints
    lin_dim = round(Int, sqrt(d - 1))
    B = randn(T, lin_dim, vec_dim)
    b = B * p0_vec
    JuMP.@constraint(model, B * p_vec .== b)
    C = randn(T, lin_dim, vec_dim)
    c = C * p0_vec
    JuMP.@constraint(model, C * p_vec .<= c)

    # save for use in tests
    model.ext[:p_var] = p

    return model
end

function test_extra(inst::CovarianceEstJuMP{T}, model::JuMP.Model) where T
    stat = JuMP.termination_status(model)
    @test stat == MOI.OPTIMAL
    (stat == MOI.OPTIMAL) || return

    # check objective
    tol = eps(T)^0.2
    p_opt = JuMP.value.(model.ext[:p_var])
    λ = eigvals(Symmetric(p_opt, :U))
    @test minimum(λ) >= -tol
    obj_result = get_val(pos_only(λ), inst.ext)
    @test JuMP.objective_value(model) ≈ obj_result atol=tol rtol=tol
    return
end


#
# blas.jl -
#
# Code based on BLAS (Basic Linear Algebra Subroutines).
#
#-------------------------------------------------------------------------------
#
# This file is part of LazyAlgebra (https://github.com/emmt/LazyAlgebra.jl)
# released under the MIT "Expat" license.
#
# Copyright (c) 2017-2019 Éric Thiébaut.
#

# The idea is to generalize the dot product as follows:
#
#   `vdot(x,y)` yields the sum of `conj(x[i])*y[i]` for each `i` in
#               `eachindex(x,y)` providing `x` and `y` have the same dimensions
#               (i.e., same `indices`).
#
#   `A*x` yields the matrix-vector product providing that the trailing
#                dimensions of `A` match the dimensions of `x`.  The result has
#                the same dimensions as the leading dimensions of `A`.
#
# We may want to use fast BLAS routines.
#
# According to the following timings (for n = 96 and 4 threads), the fastest
# method is the BLAS version of `apply!(,Adjoint,,)`.  When looking at the
# loops, this is understandable as `apply!(,Adjoint,,)` is easier to
# parallelize than `apply!(,Direct,,)`.  Note that Julia implementations are
# with SIMD and no bounds checking.
#
#             A⋅x       A'.x      x'⋅y
#     ---------------------------------
#     BLAS   3.4 µs     2.0 µs   65 ns
#     Julia  4.5 µs    24.2 μs   65 ns

const BlasVec{T} = Union{DenseVector{T},StridedVector{T}}
const BlasArr{T,N} = DenseArray{T,N}

for T in (Float32, Float64)

    @eval begin

        vdot(::Type{$T}, x::BlasVec{$T}, y::BlasVec{$T}) =
            __call_blas_dot(BLAS.dot, x, y)

        vdot(::Type{$T}, x::BlasArr{$T,N}, y::BlasArr{$T,N}) where {N} =
            __call_blas_dot(BLAS.dot, x, y)

        vdot(::Type{Complex{$T}}, x::BlasVec{Complex{$T}}, y::BlasVec{Complex{$T}}) =
            __call_blas_dot(BLAS.dotc, x, y)

        vdot(::Type{Complex{$T}}, x::BlasArr{Complex{$T},N}, y::BlasArr{Complex{$T},N}) where {N} =
            __call_blas_dot(BLAS.dotc, x, y)

    end

end

@inline function __call_blas_dot(f, x, y)
    size(x) == size(y) || throw_dimensions_mismatch()
    return f(length(x), pointer(x), stride(x, 1), pointer(y), stride(y, 1))
end

function vupdate!(y::BlasVec{T}, alpha::Real,
                  x::BlasVec{T}) where {T<:BlasFloat}
    size(x) == size(y) || throw_dimensions_mismatch()
    BLAS.axpy!(length(x), T(alpha),
               pointer(x), stride(x, 1),
               pointer(y), stride(y, 1))
    return y
end


@testset "TestZTest" begin

    # OneSampleZTest tests
    # Loop over n, d, p, s, two_sided values
    # TODO: Test one-sided vs two-sided
    params = (
        (33.36712895331925, 0.5, 0.80, 0.05, false),
        (43.995480906737704, 0.5, 0.90, 0.05, false),
        (53.940614296961286, 0.5, 0.95, 0.05, false),
        (16, 0.7494656445134399, 0.8, 0.05, false),
        (32, 0.5112738234562735, 0.8, 0.05, false),
        (64, 0.35565012796825096, 0.8, 0.05, false),
        (16, 0.5, 0.46487, 0.05, false),
        (32, 0.5, 0.782276, 0.05, false),
        (64, 0.5, 0.9760556, 0.05, false),
    )

    for (n, d, p, s, one_sided) in params
        p_z = power(OneSampleZTest, n, d, 1.0, s, one_sided)
        p_t = power(OneSampleTTest, n, d, 1.0, s, one_sided)
        n_z = sample_size(OneSampleZTest, d, 1.0, p, s, one_sided)
        n_t = sample_size(OneSampleTTest, d, 1.0, p, s, one_sided)
        d_z = effect_size(OneSampleZTest, n, p, s, one_sided)
        d_t = effect_size(OneSampleTTest, n, p, s, one_sided)
        # TODO: Get non-stupid bounds here on gap between z-test and t-test
        @test 0.9 * p_z < p_t < p_z
        @test 0.8 * n_t < n_z < n_t
        @test 0.9 * d_t < d_z < d_t
    end

    # EqualVarianceZTest tests
    # Loop over n, d, p, s, two_sided values
    # TODO: Test one-sided vs two-sided
    params = (
        (63.76561019093564, 0.5, 0.80, 0.05, false),
        (85.03128413735288, 0.5, 0.90, 0.05, false),
        (104.92794298436287, 0.5, 0.95, 0.05, false),
        (16, 1.0236637440012688, 0.8, 0.05, false),
        (32, 0.7114799360863416, 0.8, 0.05, false),
        (64, 0.4990691779657755, 0.8, 0.05, false),
        (16, 0.5, 0.2777445, 0.05, false),
        (32, 0.5, 0.5036382, 0.05, false),
        (64, 0.5, 0.8014596, 0.05, false),
    )

    for (n, d, p, s, one_sided) in params
        p_z = power(EqualVarianceZTest, n, d, 1.0, s, one_sided)
        p_t = power(EqualVarianceTTest, n, d, 1.0, s, one_sided)
        n_z = sample_size(EqualVarianceZTest, d, 1.0, p, s, one_sided)
        n_t = sample_size(EqualVarianceTTest, d, 1.0, p, s, one_sided)
        d_z = effect_size(EqualVarianceZTest, n, p, s, one_sided)
        d_t = effect_size(EqualVarianceTTest, n, p, s, one_sided)
        # TODO: Get non-stupid bounds here on gap between z-test and t-test
        @test 0.9 * p_z < p_t < p_z
        @test 0.8 * n_t < n_z < n_t
        @test 0.9 * d_t < d_z < d_t
    end
end


module AR1

export AR1Model, residuals, simulate

@doc """
Construct an AR(1) model with mean `mu`, standard deviation (at zero
lag) `sigma`, correlation time `tau`, and measurement uncertainty scale
`nu`.
""" ->
immutable AR1Model
    mu::Float64
    sigma::Float64
    tau::Float64
    nu::Float64
end

""" Returns the uncorrelated residuals and associated uncertainties
for the AR(1) model applied to the given data. """
function residuals(model::AR1Model, ts::Array{Float64, 1}, ys::Array{Float64, 1}, dys::Array{Float64, 1})
    n = size(ts, 1)

    s2 = model.sigma*model.sigma
    nu2 = model.nu*model.nu

    rs = zeros(n)
    ss = zeros(n)

    t = ts[1]
    x = 0.0
    dx2 = s2
    k = 1.0

    for i in 1:n
        # Advance the filter
        dt = ts[i] - t
        t = ts[i]
        lam = exp(-dt/model.tau)
        x = lam*x
        dx2 = lam*lam*(dx2 - s2) + s2

        # Predict the observation
        rs[i] = ys[i] - x - model.mu
        ss2 = dx2 + nu2*dys[i]*dys[i]
        ss[i] = sqrt(ss2)

        # Update gain
        k = dx2 / ss2

        # Update state
        x += rs[i]*k
        dx2 -= ss2*k*k
    end

    rs, ss
end

""" Returns a sample from the given AR(1) model at times `ts` with
observational uncertainties `dxs` """
function simulate(model::AR1Model, ts::Array{Float64, 1}, dxs::Array{Float64, 1})
    xs = zeros(ts)
    xsout = zeros(ts)

    xs[1] = model.sigma*randn()
    xsout[1] = xs[1] + model.nu*dxs[1]*randn() + model.mu

    for i in 2:size(xs,1)
        dt = ts[i] - ts[i-1]
        alpha = exp(-dt/model.tau)

        beta = model.sigma*sqrt(1.0 - alpha*alpha)

        xs[i] = alpha*xs[i-1] + beta*randn()
        xsout[i] = xs[i] + model.nu*dxs[i]*randn() + model.mu
    end

    xsout
end

end


using Test, Random
using MCMCChains
using Distributions

## Test Chain
# Define the experiment
n_iter = 4000
n_name = 3
n_chain = 2

# observations
Random.seed!(1234)
x = map(i -> i + 2e-1*randn(), 1:n_name)

# some sample experiment results
Random.seed!(1234)
val1 = 0.5*randn(n_iter, n_name, n_chain) .+ x'
val2 = 2*randn(n_iter, n_name, n_chain) .+ x'

# construct a Chains object
chn1 = Chains(val1)
chn2 = Chains(val2)

lpfun = function f(chain::Chains)

    p1 = Array(chain[:,1,:])
    p2 = Array(chain[:,2,:])
    p3 = Array(chain[:,3,:])

    lp = map(p -> logpdf(Normal(p), x[1]), p1)
    lp += map(p -> logpdf(Normal(p), x[2]), p2)
    lp += map(p -> logpdf(Normal(p), x[3]), p3)

    return lp
end

@testset "deviance information criterion" begin

    DIC1, pD1 = dic(chn1, lpfun)
    DIC2, pD2 = dic(chn2, lpfun)

    @test DIC1 < DIC2
    @test pD1 < pD2
end


###############################################################
# Part of BDisposal.jl package
###############################################################

"""
    prodIndex(gI,gO,bO,bI;retToScale,prodStructure,convexAssumption,startθ,startμ,startλ)

Compute productivity indexes

Given a set of measures of inputs, "good" ("desiderable") and "bad" ("undesiderable") outputs for different decision making units, compute their
productivity indexes improvements (or declines) between consecutive time periods.

## Parameters:
- Positional
  *  `gI`: "Good" inputs (3D matrix by DMUs, input items and periods)
  *  `gO`: "Good" outputs (3D matrix by DMUs, input items and periods)
  *  `bO`: "Bad" outputs (3D matrix by DMUs, input items and periods)
  *  `bI`: "Bad" inputs (optional 3D matrix by DMUs, input items and periods) [default: empty array]
- Keyword
  *  `retToScale`: Wheter the return to scales should be assumed "constant" (default) or "variable"
  *  `prodStructure`: Wheter the production structure should be assumed "additive" (default) or "multiplicative"
  *  `convexAssumption`: Wheter a convex production frontier should be assumed [default: `true`]

## Returns:
* A named touple where each element is a matrix of production indexes by DMUs and period passages (e.g. "year2 on year1" and "year3 on year2") or one of their decompositions.
* Currently the value reported are:
    - `prodIndexes`:          Overall production indexes
    - `prodIndexes_G`:        Decomposition for "good" inputs and outputs
    - `prodIndexes_B`:        Decomposition for "bad" inputs and outputs
    - `prodIndexes_T`:        Decomposition for the technological component, overall
    - `prodIndexes_T_O`:      Decomposition for the technological component, outputs
    - `prodIndexes_T_G_O`:    Decomposition for the technological component, good outputs
    - `prodIndexes_T_B_O`:    Decomposition for the technological component, bad outputs
    - `prodIndexes_T_I`:      Decomposition for the technological component, inputs
    - `prodIndexes_T_G_I`:    Decomposition for the technological component, good inputs
    - `prodIndexes_T_B_I`:    Decomposition for the technological component, bad inputs
    - `prodIndexes_E`:        Decomposition for the efficiency component, overall
    - `prodIndexes_E_O`:      Decomposition for the efficiency component, outputs
    - `prodIndexes_E_G_O`:    Decomposition for the efficiency component, good output
    - `prodIndexes_E_B_O`:    Decomposition for the efficiency component, bad outputs
    - `prodIndexes_E_I`:      Decomposition for the efficiency component, inputs
    - `prodIndexes_E_G_I`:    Decomposition for the efficiency component, good inputs
    - `prodIndexes_E_B_I`:    Decomposition for the efficiency component, bad inputs
    - `prodIndexes_S`:        Decomposition for the scale (residual) component, overall
    - `prodIndexes_S_O`:      Decomposition for the scale (residual) component, outputs
    - `prodIndexes_S_G_O`:    Decomposition for the scale (residual) component, good output
    - `prodIndexes_S_B_O`:    Decomposition for the scale (residual) component, bad outputs
    - `prodIndexes_S_I`:      Decomposition for the scale (residual) component, inputs
    - `prodIndexes_S_G_I`:    Decomposition for the scale (residual) component, good inputs
    - `prodIndexes_S_B_I`:    Decomposition for the scale (residual) component, bad inputs

The second and third element of the tuple are respectively the "good inputs/outputs" and "bad/inputs/outputs" components.
The first matrix can be retrieved from the two components by multiplying them (for `multiplicative` production structure)
or summing them (for `additive` production strucure).

## Description of the function

`prodIndex()` displays environmental productivity indices. These
productivity measures are implemented for different time periods (eg., years, months
etc.) or spatial units (eg., countries, cities etc.), based on the  environmental efficiency
indicators described in `efficiencyScores`. Hence, the environmental productivity indices inherit the
structure of additive and multiplicative productivity measures.

## Interpretation of the results

The additive productivity indicator shows combined desirable and undesirable outputs
productivity improvement (respectively decline) when it takes positive (respectively negative) values.
In the multiplicative context, if the productivity measure is greater (respectively lesser) than
1 then, desirable and undesirable outputs productivity increase (respectively decrease)
arises. The BDisposal package underscores the prominent sources of environmental productivity
change. The main drivers of productivity variation are technological change,
technical efficiency variation and scale efficiency change. For the additive background,
when the technological change is greater (respectively, lesser) than 0 then, technolog-
ical improvement (respectively, deterioration) occurs. A similar reasoning applies for
the additive technical and scale efficiency components. In the multiplicative context, if
the technological change is greater (respectively, lesser) than 1 then, technological in-
crease (respectively, decrrease) arises. A similar reasonnng applies for the multiplicative
technical and scale efficiency components.

## Notes:
* The decomposition by technological, efficiency and scale components is done only for the convex assumption,
as in the non-convex case the individual distance components used to compute these disaggregations are infinite.

"""
function prodIndex(gI::Array{Float64,3},gO::Array{Float64,3},bO::Array{Float64,3},bI::Array{Float64,3}=Array{Float64}(missing, (size(gI,1),0,size(gI,3)));
                   retToScale="constant",prodStructure="multiplicative",convexAssumption=true)

    nDMUs = size(gI,1)
    nPer = size(gI,3)
    nBI  = size(bI,2)

    # Output containers

    # - Global
    prodIndexes = Array{Union{Float64,Missing}}(missing, (nDMUs,nPer-1))

    # - decomposition by type of the inputs or outputs (Good vs Bad)
    prodIndexes_G = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_B = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))

    # - decomposition by tecnological/efficiency/scale (residual)
    prodIndexes_T     = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1)) # technological
    prodIndexes_T_O   = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_T_G_O = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_T_B_O = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_T_I   = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_T_G_I = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_T_B_I = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))

    prodIndexes_E     = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1)) # Efficiency
    prodIndexes_E_O   = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_E_G_O = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_E_B_O = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_E_I   = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_E_G_I = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_E_B_I = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))

    #prodIndexes_S     = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1)) # Scale (residual)
    prodIndexes_S_O   = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_S_G_O = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_S_B_O = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_S_I   = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_S_G_I = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_S_B_I = Array{Union{Float64,Missing}}(missing, (size(gI,1),nPer-1))
    prodIndexes_S     = Array{Union{Float64,Missing}}(missing, (nDMUs,nPer-1))

    noBadIndexDefault = prodStructure == "multiplicative" ? 1.0 : 0.0

    for t in 1:nPer-1
        #println(t)
        gIₜ = gI[:,:,t]
        gIᵤ = gI[:,:,t+1]
        bIₜ = bI[:,:,t]
        bIᵤ = bI[:,:,t+1]
        gOₜ = gO[:,:,t]
        gOᵤ = gO[:,:,t+1]
        bOₜ = bO[:,:,t]
        bOᵤ = bO[:,:,t+1]
        #if convexAssumption == true
            dirGI = prodStructure == "multiplicative" ? (-1,0,0,0) : (1,0,0,0)
            dirBI = prodStructure == "multiplicative" ? (0,-1,0,0) : (0,1,0,0)
            dirGO = prodStructure == "multiplicative" ? (0,0,1,0)  : (0,0,1,0)
            dirBO = prodStructure == "multiplicative" ? (0,0,0,-1) : (0,0,0,1)
        #else
        #    dirGI = prodStructure == "multiplicative" ? (1,0,0,0)  : (1,0,0,0)
        #    dirBI = prodStructure == "multiplicative" ? (0,1,0,0)  : (0,1,0,0)
        #    dirGO = prodStructure == "multiplicative" ? (0,0,-1,0) : (0,0,1,0)
        #    dirBO = prodStructure == "multiplicative" ? (0,0,0,1)  : (0,0,0,-1)
        #end
        for z in 1:nDMUs
            gIₜ₀ = gIₜ[z,:]
            bIₜ₀ = bIₜ[z,:]
            gOₜ₀ = gOₜ[z,:]
            bOₜ₀ = bOₜ[z,:]
            gIᵤ₀ = gIᵤ[z,:]
            bIᵤ₀ = bIᵤ[z,:]
            gOᵤ₀ = gOᵤ[z,:]
            bOᵤ₀ = bOᵤ[z,:]

           #if z == 1 && t == nPer-1
            #   println("here we stop and check")
           #end

            forceLinearModel = (convexAssumption == true) ?  true : false

            # t...
            idx_gi_t̃ = problem(gIᵤ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGI,crossTime=true)
            idx_gi_t = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGI)
            idx_bi_t̃ = (nBI == 0) ? noBadIndexDefault : problem(gIₜ₀,bIᵤ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBI,crossTime=true)
            idx_bi_t = (nBI == 0) ? noBadIndexDefault : problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBI)
            idx_go_t̃ = problem(gIₜ₀,bIₜ₀,gOᵤ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGO,crossTime=true)
            idx_go_t = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGO)
            idx_bo_t̃ = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOᵤ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBO,crossTime=true)
            idx_bo_t = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBO)

            # u...
            idx_gi_u = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGI)
            idx_gi_ũ = problem(gIₜ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGI,crossTime=true)
            idx_bi_u = (nBI == 0) ? noBadIndexDefault : problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBI)
            idx_bi_ũ = (nBI == 0) ? noBadIndexDefault : problem(gIᵤ₀,bIₜ₀,gOᵤ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBI,crossTime=true)
            idx_go_u = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGO)
            idx_go_ũ = problem(gIᵤ₀,bIᵤ₀,gOₜ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGO,crossTime=true)
            idx_bo_u = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBO)
            idx_bo_ũ = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOₜ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBO,crossTime=true)

            # DMU at t, frontier at t+1...
            idx_gi_tu = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGI,crossTime=true)
            idx_bi_tu = (nBI == 0) ? noBadIndexDefault : problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBI,crossTime=true)
            idx_go_tu = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGO,crossTime=true)
            idx_bo_tu = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIᵤ,bIᵤ,gOᵤ,bOᵤ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBO,crossTime=true)

            # DMU at t+1, frontier at t...
            idx_gi_ut = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGI,crossTime=true)
            idx_bi_ut = (nBI == 0) ? noBadIndexDefault : problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBI,crossTime=true)
            idx_go_ut = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirGO,crossTime=true)
            idx_bo_ut = problem(gIᵤ₀,bIᵤ₀,gOᵤ₀,bOᵤ₀,gIₜ,bIₜ,gOₜ,bOₜ,
                        retToScale=retToScale,prodStructure=prodStructure,
                        convexAssumption=convexAssumption,forceLinearModel=forceLinearModel,
                        directions=dirBO,crossTime=true)

            # Computing aggregated for the full (gI,Bi,bO,gO) partition...
            if prodStructure == "multiplicative"
                if convexAssumption == true
                    idx_i_t = (idx_gi_t̃/idx_gi_t) * (idx_bi_t̃/idx_bi_t)
                    idx_o_t = (idx_go_t/idx_go_t̃) * (idx_bo_t/idx_bo_t̃)
                    idx_t   = idx_o_t/idx_i_t
                    idx_i_u = (idx_gi_u/idx_gi_ũ) * (idx_bi_u/idx_bi_ũ)
                    idx_o_u = (idx_go_ũ/idx_go_u) * (idx_bo_ũ/idx_bo_u)
                    idx_u   = idx_o_u/idx_i_u
                    idx     = (idx_t * idx_u)^(1/2)
                else
                    idx_i_t = (idx_gi_t/idx_gi_t̃) * (idx_bi_t/idx_bi_t̃)
                    idx_o_t = (idx_go_t/idx_go_t̃) / (idx_bo_t̃/idx_bo_t)
                    idx_t   = idx_o_t/idx_i_t
                    idx_i_u = (idx_gi_ũ/idx_gi_u) * (idx_bi_ũ/idx_bi_u)
                    idx_o_u = (idx_go_ũ/idx_go_u) / (idx_bo_u/idx_bo_ũ)
                    idx_u   = idx_o_u/idx_i_u
                    idx     = (idx_t * idx_u)^(1/2)
                end

            else # addittive
                if convexAssumption == true
                    idx_i_t = (idx_gi_t̃ - idx_gi_t) + (idx_bi_t̃ - idx_bi_t)
                    idx_o_t = (idx_go_t - idx_go_t̃) - (idx_bo_t̃ - idx_bo_t)
                    idx_t   = idx_o_t - idx_i_t
                    idx_i_u = (idx_gi_u - idx_gi_ũ) + (idx_bi_u - idx_bi_ũ)
                    idx_o_u = (idx_go_ũ - idx_go_u) - (idx_bo_u - idx_bo_ũ)
                    idx_u   = idx_o_u - idx_i_u
                    idx     = (idx_t + idx_u) / 2
                else
                    idx_i_t = -(idx_gi_t-idx_gi_t̃) - (idx_bi_t-idx_bi_t̃)
                    idx_o_t = (idx_go_t-idx_go_t̃) - (idx_bo_t̃-idx_bo_t)
                    idx_t   = idx_o_t-idx_i_t
                    idx_i_u = -(idx_gi_ũ-idx_gi_u) - (idx_bi_ũ-idx_bi_u)
                    idx_o_u = (idx_go_ũ-idx_go_u) - (idx_bo_u-idx_bo_ũ)
                    idx_u   = idx_o_u-idx_i_u
                    idx     = (idx_t + idx_u)/2
                end
            end
            prodIndexes[z,t] = idx

            if convexAssumption == true
                # Computing aggregated for the good inputs/outputs only...
                if prodStructure == "multiplicative"
                   idx_Gi_t = (idx_gi_t̃/idx_gi_t)
                   idx_Go_t = (idx_go_t/idx_go_t̃)
                   idx_Gt   = idx_Go_t/idx_Gi_t
                   idx_Gi_u = (idx_gi_u/idx_gi_ũ)
                   idx_Go_u = (idx_go_ũ/idx_go_u)
                   idx_Gu   = idx_Go_u/idx_Gi_u
                   idx_G    = (idx_Gt * idx_Gu)^(1/2)
                else
                    idx_Gi_t = (idx_gi_t̃ - idx_gi_t)
                    idx_Go_t = (idx_go_t - idx_go_t̃)
                    idx_Gt   = idx_Go_t - idx_Gi_t
                    idx_Gi_u = (idx_gi_u - idx_gi_ũ)
                    idx_Go_u = (idx_go_ũ - idx_go_u)
                    idx_Gu   = idx_Go_u - idx_Gi_u
                    idx_G     = (idx_Gt + idx_Gu) / 2
                end
                prodIndexes_G[z,t] = idx_G

                # Computing aggregated for the bad inputs/outputs only...
                if prodStructure == "multiplicative"
                   idx_Bi_t = (idx_bi_t̃/idx_bi_t)
                   idx_Bo_t = (idx_bo_t̃/idx_bo_t)
                   idx_Bt   = 1/(idx_Bo_t * idx_Bi_t)
                   idx_Bi_u = (idx_bi_u/idx_bi_ũ)
                   idx_Bo_u = (idx_bo_u/idx_bo_ũ)
                   idx_Bu   = 1/(idx_Bo_u*idx_Bi_u)
                   idx_B    = (idx_Bt * idx_Bu)^(1/2)
                else
                    idx_Bi_t = (idx_bi_t̃ - idx_bi_t)
                    idx_Bo_t = (idx_bo_t̃ - idx_bo_t)
                    idx_Bt   = - idx_Bo_t - idx_Bi_t
                    idx_Bi_u = (idx_bi_u - idx_bi_ũ)
                    idx_Bo_u = (idx_bo_u - idx_bo_ũ)
                    idx_Bu   = - idx_Bo_u - idx_Bi_u
                    idx_B    = (idx_Bt + idx_Bu) / 2
                end
                prodIndexes_B[z,t] = idx_B

                # Computing tecnological/efficiency/residual decomposition
                if prodStructure == "multiplicative"
                    # Disaggregation T/E/S...
                    idx_T_G_O = ((idx_go_t/idx_go_tu) * (idx_go_ut/idx_go_u) )^-(1/2)
                    idx_T_B_O = ((idx_bo_t/idx_bo_tu) * (idx_bo_ut/idx_bo_u) )^-(1/2)
                    idx_T_O   = idx_T_G_O * idx_T_B_O
                    idx_T_G_I = ((idx_gi_t/idx_gi_tu) * (idx_gi_ut/idx_gi_u) )^-(1/2)
                    idx_T_B_I = ((idx_bi_t/idx_bi_tu) * (idx_bi_ut/idx_bi_u) )^-(1/2)
                    idx_T_I   = idx_T_G_I * idx_T_B_I
                    idx_T     = (idx_T_O * idx_T_I)

                    idx_E_G_O = (idx_go_u/idx_go_t)^-1
                    idx_E_B_O = (idx_bo_u/idx_bo_t)^-1
                    idx_E_O   = idx_E_G_O * idx_E_B_O
                    idx_E_G_I = (idx_gi_u/idx_gi_t)^-1
                    idx_E_B_I = (idx_bi_u/idx_bi_t)^-1
                    idx_E_I   = idx_E_G_I * idx_E_B_I
                    idx_E     = (idx_E_O * idx_E_I)

                    idx_S_I   = idx / (idx_T_I * idx_E_I)
                    idx_S_O   = idx / (idx_T_O * idx_E_O)
                    idx_S_G_O = ((idx_go_t̃ / idx_go_ut) * (idx_gi_t̃ / idx_gi_t) * (idx_go_tu / idx_go_ũ) * (idx_gi_u/idx_gi_ũ)  )^-(1/2)
                    idx_S_B_O = ((idx_bo_t̃ / idx_bo_ut) * (idx_bi_t̃ / idx_bi_t) * (idx_bo_tu / idx_bo_ũ) * (idx_bi_u/idx_bi_ũ)  )^-(1/2)
                    idx_S_G_I = ((idx_gi_t̃ / idx_gi_ut) * (idx_go_t̃ / idx_go_t) * (idx_gi_tu / idx_gi_ũ) * (idx_go_u/idx_go_ũ)  )^-(1/2)
                    idx_S_B_I = ((idx_bi_t̃ / idx_bi_ut) * (idx_bo_t̃ / idx_bo_t) * (idx_bi_tu / idx_bi_ũ) * (idx_bo_u/idx_bo_ũ)  )^-(1/2)

                    idx_S     = idx / (idx_T * idx_E)

                else # additive production structure
                    # Disaggregation T/E/S...
                    idx_T_G_O = ((-idx_go_t+idx_go_tu) + (-idx_go_ut+idx_go_u) )/2
                    idx_T_B_O = ((-idx_bo_t-idx_bo_tu) + (+idx_bo_ut+idx_bo_u) )/2
                    idx_T_O   = idx_T_G_O + idx_T_B_O
                    idx_T_G_I = ((-idx_gi_t+idx_gi_tu) + (-idx_gi_ut+idx_gi_u) )/2
                    idx_T_B_I = ((-idx_bi_t+idx_bi_tu) + (-idx_bi_ut+idx_bi_u) )/2
                    idx_T_I   = idx_T_G_I + idx_T_B_I
                    idx_T     = (idx_T_O + idx_T_I)

                    idx_E_G_O = -(idx_go_u-idx_go_t)
                    idx_E_B_O = -(idx_bo_u-idx_bo_t)
                    idx_E_O   = idx_E_G_O + idx_E_B_O
                    idx_E_G_I = -(idx_gi_u-idx_gi_t)
                    idx_E_B_I = -(idx_bi_u-idx_bi_t)
                    idx_E_I   = idx_E_G_I + idx_E_B_I
                    idx_E     = (idx_E_O + idx_E_I)

                    idx_S_I   = idx - (idx_T_I + idx_E_I)
                    idx_S_O   = idx - (idx_T_O + idx_E_O)
                    idx_S_G_O = -((idx_go_t̃ - idx_go_ut) + (idx_gi_t̃ - idx_gi_t) + (idx_go_tu - idx_go_ũ) + (idx_gi_u-idx_gi_ũ)  )/2
                    idx_S_B_O = -((idx_bo_t̃ + idx_bo_ut) + (idx_bi_t̃ - idx_bi_t) + (-idx_bo_tu - idx_bo_ũ) + (idx_bi_u-idx_bi_ũ)  )/2
                    idx_S_G_I = -((idx_gi_t̃ - idx_gi_ut) + (idx_go_t̃ - idx_go_t) + (idx_gi_tu - idx_gi_ũ) + (idx_go_u-idx_go_ũ)  )/2
                    idx_S_B_I = -((idx_bi_t̃ - idx_bi_ut) + (idx_bo_t̃ - idx_bo_t) + (idx_bi_tu - idx_bi_ũ) + (idx_bo_u-idx_bo_ũ)  )/2
                    idx_S     = idx - (idx_T + idx_E)

                    #=
                    idx_T_G_O = ((idx_go_t-idx_go_ũ) + (idx_go_t̃-idx_go_u) )/2
                    idx_T_B_O = ((idx_bo_t-idx_bo_ũ) + (idx_bo_t̃-idx_bo_u) )/2
                    idx_T_O   = idx_T_G_O + idx_T_B_O
                    idx_T_G_I = ((idx_gi_t-idx_gi_ũ) + (idx_gi_t̃-idx_gi_u) )/2
                    idx_T_B_I = ((idx_bi_t-idx_bi_ũ) + (idx_bi_t̃-idx_bi_u) )/2
                    idx_T_I   = idx_T_G_I + idx_T_B_I
                    idx_T     = (idx_T_O + idx_T_I)/2

                    idx_E_G_O = idx_go_u-idx_go_t
                    idx_E_B_O = idx_bo_u-idx_bo_t
                    idx_E_O   = idx_E_G_O + idx_E_B_O
                    idx_E_G_I = idx_gi_u-idx_gi_t
                    idx_E_B_I = idx_bi_u-idx_bi_t
                    idx_E_I   = idx_E_G_I + idx_E_B_I
                    idx_E     = (idx_E_O + idx_E_I)/2

                    idx_S_G_O = idx_G -  (idx_T_G_O + idx_E_G_O)
                    idx_S_B_O = idx_B -  (idx_T_B_O + idx_E_B_O)
                    idx_S_O   = idx_S_G_O + idx_S_B_O
                    idx_S_G_I = idx_G -  (idx_T_G_I + idx_E_G_I)
                    idx_S_B_I = idx_B -  (idx_T_B_I + idx_E_B_I)
                    idx_S_I   = idx_S_G_I + idx_S_B_I
                    #idx_S     = (idx_S_O + idx_S_I)/2
                    =#
                end

            prodIndexes_T[z,t]      =  idx_T
            prodIndexes_T_O[z,t]    =  idx_T_O
            prodIndexes_T_G_O[z,t]  =  idx_T_G_O
            prodIndexes_T_B_O[z,t]  =  idx_T_B_O
            prodIndexes_T_I[z,t]    =  idx_T_I
            prodIndexes_T_G_I[z,t]  =  idx_T_G_I
            prodIndexes_T_B_I[z,t]  =  idx_T_B_I
            prodIndexes_E[z,t]      =  idx_E
            prodIndexes_E_O[z,t]    =  idx_E_O
            prodIndexes_E_G_O[z,t]  =  idx_E_G_O
            prodIndexes_E_B_O[z,t]  =  idx_E_B_O
            prodIndexes_E_I[z,t]    =  idx_E_I
            prodIndexes_E_G_I[z,t]  =  idx_E_G_I
            prodIndexes_E_B_I[z,t]  =  idx_E_B_I
            prodIndexes_S[z,t]      =  idx_S
            prodIndexes_S_O[z,t]    =  idx_S_O
            prodIndexes_S_G_O[z,t]  =  idx_S_G_O
            prodIndexes_S_B_O[z,t]  =  idx_S_B_O
            prodIndexes_S_I[z,t]    =  idx_S_I
            prodIndexes_S_G_I[z,t]  =  idx_S_G_I
            prodIndexes_S_B_I[z,t]  =  idx_S_B_I

        else # non-convex case
            # Computing aggregated for the good inputs/outputs only...
            if prodStructure == "multiplicative"
                idx_Gi_t = 1/(idx_gi_t̃*idx_gi_t)
                idx_Go_t = (idx_go_t/idx_go_t̃)
                idx_Gt   = idx_Go_t/idx_Gi_t
                idx_Gi_u = 1/(idx_gi_u/idx_gi_ũ)
                idx_Go_u = (idx_go_ũ/idx_go_u)
                idx_Gu   = idx_Go_u/idx_Gi_u
                idx_G    = (idx_Gt * idx_Gu)^(1/2)

                idx_Bi_t = 1/(idx_bi_t̃/idx_bi_t)
                idx_Bo_t = 1/(idx_bo_t̃/idx_bo_t)
                idx_Bt   = (idx_Bo_t / idx_Bi_t)
                idx_Bi_u = 1/(idx_bi_u/idx_bi_ũ)
                idx_Bo_u = 1/(idx_bo_u/idx_bo_ũ)
                idx_Bu   = (idx_Bo_u/idx_Bi_u)
                idx_B    = (idx_Bt * idx_Bu)^(1/2)

            else

                idx_Gi_t = (idx_gi_t̃-idx_gi_t)
                idx_Go_t = (idx_go_t-idx_go_t̃)
                idx_Gt   = idx_Go_t-idx_Gi_t
                idx_Gi_u = (idx_gi_u-idx_gi_ũ)
                idx_Go_u = (idx_go_ũ-idx_go_u)
                idx_Gu   = idx_Go_u-idx_Gi_u
                idx_G    = (idx_Gt + idx_Gu)/2

                idx_Bi_t = (idx_bi_t̃-idx_bi_t)
                idx_Bo_t = (idx_bo_t̃-idx_bo_t)
                idx_Bt   = (-idx_Bo_t - idx_Bi_t)
                idx_Bi_u = (idx_bi_u-idx_bi_ũ)
                idx_Bo_u = (idx_bo_u-idx_bo_ũ)
                idx_Bu   = (-idx_Bo_u-idx_Bi_u)
                idx_B    = (idx_Bt + idx_Bu)/2

            end
            prodIndexes_G[z,t] = idx_G
            prodIndexes_B[z,t] = idx_B

        end

        end # end for each DMU
    end # end for each period

    return (prodIndexes        =  prodIndexes       ,
            prodIndexes_G      =  prodIndexes_G     ,
            prodIndexes_B      =  prodIndexes_B     ,
            prodIndexes_T      =  prodIndexes_T     ,
            prodIndexes_T_O    =  prodIndexes_T_O   ,
            prodIndexes_T_G_O  =  prodIndexes_T_G_O ,
            prodIndexes_T_B_O  =  prodIndexes_T_B_O ,
            prodIndexes_T_I    =  prodIndexes_T_I   ,
            prodIndexes_T_G_I  =  prodIndexes_T_G_I ,
            prodIndexes_T_B_I  =  prodIndexes_T_B_I ,
            prodIndexes_E      =  prodIndexes_E     ,
            prodIndexes_E_O    =  prodIndexes_E_O   ,
            prodIndexes_E_G_O  =  prodIndexes_E_G_O ,
            prodIndexes_E_B_O  =  prodIndexes_E_B_O ,
            prodIndexes_E_I    =  prodIndexes_E_I   ,
            prodIndexes_E_G_I  =  prodIndexes_E_G_I ,
            prodIndexes_E_B_I  =  prodIndexes_E_B_I ,
            prodIndexes_S      =  prodIndexes_S     ,
            prodIndexes_S_O    =  prodIndexes_S_O   ,
            prodIndexes_S_G_O  =  prodIndexes_S_G_O ,
            prodIndexes_S_B_O  =  prodIndexes_S_B_O ,
            prodIndexes_S_I    =  prodIndexes_S_I   ,
            prodIndexes_S_G_I  =  prodIndexes_S_G_I ,
            prodIndexes_S_B_I  =  prodIndexes_S_B_I )
end


using EffectiveWaves, Test
using LinearAlgebra, Statistics

@testset "Check material definitions" begin
    include("../src/acoustics/acoustic_mediums.jl")
    @test true
end

include("complex.jl")

# Single effective wavenumber tests
    @time include("effective-2D/strong_low_freq_effective.jl")
    @time include("effective-2D/high_frequency_effective.jl")

    @time include("effective-2D/large_vol_low_freq_effective.jl")
    @time include("effective-2D/weak_scatterers_effective.jl")

# Eigensystems for different symmetries (i.e. a sphere fille with particles or a halfspace filled with particles)
    @time include("effective-3D/equivalent-symmetries.jl")
    @time include("effective-3D/low_volume_fraction.jl")
    @time include("effective-3D/low_frequency.jl")
    @time include("effective-3D/planar-symmetry.jl")

# test equivalence between methods for finding wavenumbers
    # test does not run on Julia version < 0.7 due to differences in Optim versions
    # include("path_mesh_wavenumbers.jl")

# Test functions used for the discretisation
    include("discretisation/numerical_integration.jl")
    include("discretisation/integrated_reflection.jl")
    include("discretisation/average_integrand_kernel.jl")

# Test discretisation for 3D
    include("effective-3D/discrete_solver.jl")

# Test matching method
    include("match-wave/match_low_volumefraction.jl")
    include("match-wave/match_low_frequency.jl")
    include("match-wave/match_discrete_wave.jl")

# Test wiener hopf method for monopole scatterers
    # include("wiener-hopf-reflection.jl")
    # include("wiener-hopf-wave.jl")