# Themelis, Stella, Patrinos, "Forward-backward envelope for the sum of two
# nonconvex functions: Further properties and nonmonotone line-search
# algorithms", SIAM Journal on Optimization, vol. 28, no. 3, pp. 2274–2303
# (2018).

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalOperators: Zero
using LinearAlgebra
using Printf

"""
    ZeroFPRIteration(; <keyword-arguments>)

Instantiate the ZeroFPR algorithm (see [1]) for solving optimization problems
of the form

    minimize f(Ax) + g(x),

where `f` is smooth and `A` is a linear mapping (for example, a matrix).

# Arguments
- `x0`: initial point.
- `f=Zero()`: smooth objective term.
- `A=I`: linear operator (e.g. a matrix).
- `g=Zero()`: proximable objective term.
- `Lf=nothing`: Lipschitz constant of the gradient of x ↦ f(Ax).
- `gamma=nothing`: stepsize to use, defaults to `1/Lf` if not set (but `Lf` is).
- `adaptive=false`: forces the method stepsize to be adaptively adjusted.
- `minimum_gamma=1e-7`: lower bound to `gamma` in case `adaptive == true`.
- `max_backtracks=20`: maximum number of line-search backtracks.
- `directions=LBFGS(5)`: strategy to use to compute line-search directions.

# References
1. Themelis, Stella, Patrinos, "Forward-backward envelope for the sum of two nonconvex functions: Further properties and nonmonotone line-search algorithms", SIAM Journal on Optimization, vol. 28, no. 3, pp. 2274–2303 (2018).
"""
Base.@kwdef struct ZeroFPRIteration{R,Tx,Tf,TA,Tg,TLf,Tgamma,D}
    f::Tf = Zero()
    A::TA = I
    g::Tg = Zero()
    x0::Tx
    alpha::R = real(eltype(x0))(0.95)
    beta::R = real(eltype(x0))(0.5)
    Lf::TLf = nothing
    gamma::Tgamma = Lf === nothing ? nothing : (alpha / Lf)
    adaptive::Bool = gamma === nothing
    minimum_gamma::R = real(eltype(x0))(1e-7)
    max_backtracks::Int = 20
    directions::D = LBFGS(5)
end

Base.IteratorSize(::Type{<:ZeroFPRIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct ZeroFPRState{R,Tx,TAx,TH}
    x::Tx             # iterate
    Ax::TAx           # A times x
    f_Ax::R           # value of smooth term
    grad_f_Ax::TAx    # gradient of f at Ax
    At_grad_f_Ax::Tx  # gradient of smooth term
    gamma::R          # stepsize parameter of forward and backward steps
    y::Tx             # forward point
    xbar::Tx          # forward-backward point
    g_xbar::R         # value of nonsmooth term (at xbar)
    res::Tx           # fixed-point residual at iterate (= x - xbar)
    H::TH             # variable metric
    tau::R = zero(gamma)
    Axbar::TAx = similar(Ax)
    grad_f_Axbar::TAx = similar(Ax)
    At_grad_f_Axbar::Tx = similar(x)
    xbarbar::Tx = similar(x)
    res_xbar::Tx = similar(x)
    xbar_prev::Tx = similar(x)
    res_xbar_prev::Tx = similar(x)
    is_prev_set::Bool = false
    d::Tx = similar(x)
    Ad::TAx = similar(Ax)
end

f_model(iter::ZeroFPRIteration, state::ZeroFPRState) = f_model(state.f_Ax, state.At_grad_f_Ax, state.res, iter.alpha / state.gamma)

function Base.iterate(iter::ZeroFPRIteration{R}) where R
    x = copy(iter.x0)
    Ax = iter.A * x
    grad_f_Ax, f_Ax = gradient(iter.f, Ax)
    gamma = iter.gamma === nothing ? iter.alpha / lower_bound_smoothness_constant(iter.f, iter.A, x, grad_f_Ax) : iter.gamma
    At_grad_f_Ax = iter.A' * grad_f_Ax
    y = x - gamma .* At_grad_f_Ax
    xbar, g_xbar = prox(iter.g, y, gamma)
    state = ZeroFPRState(
        x=x, Ax=Ax, f_Ax=f_Ax, grad_f_Ax=grad_f_Ax, At_grad_f_Ax=At_grad_f_Ax,
        gamma=gamma, y=y, xbar=xbar, g_xbar=g_xbar, res=x - xbar, H=initialize(iter.directions, x),
    )
    return state, state
end

set_next_direction!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState) = mul!(state.d, state.H, -state.res_xbar)
set_next_direction!(::NoAccelerationStyle, ::ZeroFPRIteration, state::ZeroFPRState) = state.d .= .-state.res
set_next_direction!(iter::ZeroFPRIteration, state::ZeroFPRState) = set_next_direction!(acceleration_style(typeof(iter.directions)), iter, state)

update_direction_state!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState) = update!(state.H, state.xbar - state.xbar_prev, state.res_xbar - state.res_xbar_prev)
update_direction_state!(::NoAccelerationStyle, ::ZeroFPRIteration, state::ZeroFPRState) = return
update_direction_state!(iter::ZeroFPRIteration, state::ZeroFPRState) = update_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)

reset_direction_state!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState) = reset!(state.H)
reset_direction_state!(::NoAccelerationStyle, ::ZeroFPRIteration, state::ZeroFPRState) = return
reset_direction_state!(iter::ZeroFPRIteration, state::ZeroFPRState) = reset_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)

function Base.iterate(iter::ZeroFPRIteration{R}, state::ZeroFPRState) where R
    f_Axbar_upp, f_Axbar = if iter.adaptive == true
        gamma_prev = state.gamma
        state.gamma, state.g_xbar, f_Axbar, f_Axbar_upp = backtrack_stepsize!(
            state.gamma, iter.f, iter.A, iter.g,
            state.x, state.f_Ax, state.At_grad_f_Ax, state.y, state.xbar, state.g_xbar, state.res,
            state.Axbar, state.grad_f_Axbar,
            alpha = iter.alpha, minimum_gamma = iter.minimum_gamma,
        )
        if state.gamma != gamma_prev
            reset_direction_state!(iter, state)
        end
        f_Axbar_upp, f_Axbar
    else
        mul!(state.Axbar, iter.A, state.xbar)
        f_model(iter, state), gradient!(state.grad_f_Axbar, iter.f, state.Axbar)
    end

    # compute FBE
    FBE_x = f_Axbar_upp + state.g_xbar

    # compute residual at xbar
    mul!(state.At_grad_f_Axbar, iter.A', state.grad_f_Axbar)
    state.y .= state.xbar .- state.gamma .* state.At_grad_f_Axbar
    g_xbarbar = prox!(state.xbarbar, iter.g, state.y, state.gamma)
    state.res_xbar .= state.xbar .- state.xbarbar

    if state.is_prev_set == true
        update_direction_state!(iter, state)
    end

    copyto!(state.xbar_prev, state.xbar)
    copyto!(state.res_xbar_prev, state.res_xbar)
    state.is_prev_set = true

    set_next_direction!(iter, state)

    # Perform line-search over the FBE
    tau = R(1)
    mul!(state.Ad, iter.A, state.d)

    sigma = iter.beta * (0.5 / state.gamma) * (1 - iter.alpha)
    tol = 10 * eps(R) * (1 + abs(FBE_x))
    threshold = FBE_x - sigma * norm(state.res)^2 + tol

    for _ in 1:iter.max_backtracks
        state.x .= state.xbar_prev .+ tau .* state.d
        state.Ax .= state.Axbar .+ tau .* state.Ad
        # TODO: can precompute most of next line in case f is quadratic
        state.f_Ax = gradient!(state.grad_f_Ax, iter.f, state.Ax)
        mul!(state.At_grad_f_Ax, iter.A', state.grad_f_Ax)
        state.y .= state.x .- state.gamma .* state.At_grad_f_Ax
        state.g_xbar = prox!(state.xbar, iter.g, state.y, state.gamma)
        state.res .= state.x .- state.xbar
        FBE_x = f_model(iter, state) + state.g_xbar

        if FBE_x <= threshold
            state.tau = tau
            return state, state
        end

        tau *= 0.5
    end

    @warn "stepsize `tau` became too small ($(tau)), stopping the iterations"
    return nothing
end

# Solver

struct ZeroFPR{R, K}
    maxit::Int
    tol::R
    verbose::Bool
    freq::Int
    kwargs::K
end

function (solver::ZeroFPR)(x0; kwargs...)
    stop(state::ZeroFPRState) = norm(state.res, Inf) / state.gamma <= solver.tol
    disp((it, state)) = @printf(
        "%5d | %.3e | %.3e | %.3e\n",
        it,
        state.gamma,
        norm(state.res, Inf) / state.gamma,
        state.tau,
    )
    iter = ZeroFPRIteration(; x0=x0, solver.kwargs..., kwargs...)
    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)
    if solver.verbose
        iter = tee(sample(iter, solver.freq), disp)
    end
    num_iters, state_final = loop(iter)
    return state_final.xbar, num_iters
end

ZeroFPR(; maxit=1_000, tol=1e-8, verbose=false, freq=10, kwargs...) =
    ZeroFPR(maxit, tol, verbose, freq, kwargs)


module InterfacesTests

using Test
using TensorPolynomialBases
using StaticArrays

P = SVector{2,Float64}
V = SVector{3,Float64}
G = SMatrix{2,3,Float64,6}
filter(e,order) = sum(e) <= order
order = 1
basis = MonomialBasis{P,V}(filter,order)

x = SVector(2.0,3.0)

cache = ScratchData(basis)
v = zeros(V,length(basis))
g = zeros(G,length(basis))
evaluate!(v,basis,x,cache)
gradient!(g,basis,x,cache)
test_polynomial_basis(basis,x,v,g)

v = evaluate(basis,x)
g = gradient(basis,x)
test_polynomial_basis(basis,x,v,g)

@test P == point_type(basis)
@test V == value_type(basis)
@test G == gradient_type(basis)

end # module InterfacesTests


module UtilsTests

using TensorPolynomialBases: _gradient_type
using Test
using StaticArrays
using TensorValues

@test _gradient_type(SVector{3,Int},SVector{2,Int}) == SMatrix{2,3,Int,6}

@test _gradient_type(SMatrix{3,4,Int},SVector{2,Int}) == SArray{Tuple{2,3,4},Int,3,24}

@test _gradient_type(VectorValue{3,Int},VectorValue{2,Int}) == MultiValue{Tuple{2,3},Int,2,6}

@test _gradient_type(TensorValue{3,Int,9},VectorValue{2,Int}) == MultiValue{Tuple{2,3,3},Int,3,18}

end # module UtilsTests


"""Shift integer to push the mantissa in the right position. Used to determine
round up or down in the tie case. `keepbits` is the number of mantissa bits to
be kept (i.e. not zero-ed) after rounding."""
function get_shift(::Type{T},keepbits::Integer) where {T<:Base.IEEEFloat}
    return Base.significand_bits(T) - keepbits
end

"""Returns for a Float-type `T` and `keepbits`, the number of mantissa bits to be
kept/non-zeroed after rounding, half of the unit in the last place as unsigned integer.
Used in round (nearest) to add ulp/2 just before round down to achieve round nearest.
Technically ulp/2 here is just smaller than ulp/2 which rounds down the ties. For
a tie round up +1 is added in `round(T,keepbits)`."""
function get_ulp_half(  ::Type{T},
                        keepbits::Integer
                        ) where {T<:Base.IEEEFloat}
    # convert to signed for arithmetic bitshift
    # trick to push in 0s from the left and 1s from the right
    return ~unsigned(signed(~Base.significand_mask(T)) >> (keepbits+1))
end

"""Returns a mask that's 1 for all bits that are kept after rounding and 0 for the
discarded trailing bits. E.g.
```
julia> get_keep_mask(Float16,5)
0xffe0
```."""
function get_keep_mask( ::Type{T},
                        keepbits::Integer
                        ) where {T<:Base.IEEEFloat}
    # convert to signed for arithmetic bitshift
    # trick to push in 1s from the left and 0s from the right
    return unsigned(signed(~Base.significand_mask(T)) >> keepbits)
end

"""Returns a mask that's `1` for a given `mantissabit` and `0` else. Mantissa bits
are positive for the mantissa (`mantissabit = 1` is the first mantissa bit), `mantissa = 0`
is the last exponent bit, and negative for the other exponent bits."""
function get_bit_mask(  ::Type{T},
                        mantissabit::Integer
                        ) where {T<:Base.IEEEFloat}
    # push the sign mask (1000....) in the right position
    return Base.sign_mask(T) >> (Base.exponent_bits(T) + mantissabit)
end

"""IEEE's round to nearest tie to even for Float16/32/64."""
function Base.round(x::T,               # Float to be rounded
                    ulp_half::UIntT,    # obtain from get_ulp_half,
                    shift::Integer,     # get_shift, and
                    keepmask::UIntT     # get_keep_mask
                    ) where {T<:Base.IEEEFloat,UIntT<:Unsigned}
    ui = reinterpret(UIntT,x)                       # bitpattern as uint
    ui += ulp_half + ((ui >> shift) & one(UIntT))   # add ulp/2 with tie to even
    return reinterpret(T,ui & keepmask)             # set trailing bits to zero
end

"""Scalar version of `round(::Float,keepbits)` that first obtains
`shift, ulp_half, keep_mask` and then rounds."""
function Base.round(x::T,
                    keepbits::Integer
                    ) where {T<:Base.IEEEFloat}
    return round(x,get_ulp_half(T,keepbits),get_shift(T,keepbits),get_keep_mask(T,keepbits))
end

"""IEEE's round to nearest tie to even for a float array `X` in-place. Calculates from `keepbits`
only once the variables `ulp_half`, `shift` and `keep_mask` and loops over every element of the
array."""
function round!(X::AbstractArray{T},            # any array with element type T
                keepbits::Integer               # how many mantissa bits to keep
                ) where {T<:Base.IEEEFloat}     # constrain element type to Float16/32/64

    ulp_half = get_ulp_half(T,keepbits)         # half of unit in last place (ulp)
    shift = get_shift(T,keepbits)               # integer used for bitshift
    keep_mask = get_keep_mask(T,keepbits)       # mask to zero trailing mantissa bits

    @inbounds for i in eachindex(X)             # apply rounding to each element
        X[i] = round(X[i],ulp_half,shift,keep_mask)
    end

    return X
end

"""IEEE's round to nearest tie to even for a float array `X` which returns a rounded copy of `X`."""
function Base.round(X::AbstractArray{T},        # any array with element type T
                    keepbits::Integer           # mantissa bits to keep
                    ) where {T<:Base.IEEEFloat} # constrain element type to Float32/64

    Xcopy = copy(X)                             # copy array to avoid in-place changes
    round!(Xcopy,keepbits)                      # in-place round the copied array
    return Xcopy
end

"""Checks a given `mantissabit` of `x` for eveness. 1=odd, 0=even. Mantissa bits
are positive for the mantissa (`mantissabit = 1` is the first mantissa bit), `mantissa = 0`
is the last exponent bit, and negative for the other exponent bits."""
function Base.iseven(x::T,
                    mantissabit::Integer
                    ) where {T<:Base.IEEEFloat}

    bitmask = get_bit_mask(T,mantissabit)
    return 0x0 == reinterpret(typeof(bitmask),x) & bitmask
end

"""Checks a given `mantissabit` of `x` for oddness. 1=odd, 0=even. Mantissa bits
are positive for the mantissa (`mantissabit = 1` is the first mantissa bit), `mantissa = 0`
is the last exponent bit, and negative for the other exponent bits."""
Base.isodd(x::T,mantissabit::Integer) where {T<:Base.IEEEFloat} = ~iseven(x,mantissabit)


# -----------------------------------------------
# Types
# -----------------------------------------------

# using Optim

import Base.show

# Import types from Optim.jl
const AbsOptim = Optim.AbstractOptimizer
const AbsRes   = Optim.OptimizationResults

# #########   REGRESSION TYPES   #########

"""
## Regression Types

Every regression has a specific type inherited by the abstract type `AbstractRegressionType`.
"""
abstract type AbstractRegressionType end

# ---- LINEAR REGRESSION ----

"""
## Linear Regression

The abstract `LinearRegressionType` is the supertype for all *linear* regressions.
"""
abstract type LinearRegressionType<:AbstractRegressionType end

"""
### OLS (Ordinary Least Squares) Regression

Type that defines the simple OLS regression.
"""
struct OLS<:LinearRegressionType
    intercept::Bool
    robust::Any

    function OLS(; intercept::Bool=true, robust=false)
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(intercept, robust)
    end
end

"""
### GLS (Generalized Least Squares) Regression

Type that defines the GLS estimator. Since this is a theoretical estimator, the matrix Ω must be provided.
"""
struct GLS<:LinearRegressionType
    omega
    intercept::Bool
    robust::Any

    function GLS{T<:Number}(omega::Array{T}; intercept::Bool=true, robust=false)
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(omega, intercept, robust)
    end
end

"""
### FGLS (Feasible Generalized Least Squares) Regression

Type that defines the FGLS estimator: Ω is estimated by means of the residuals from an OLS regression.
"""
struct FGLS<:LinearRegressionType
    intercept::Bool
    robust::Any

    function FGLS(; intercept::Bool=true, robust=false)
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(intercept, robust)
    end
end

"""
### Iterated FGLS Regression

Type that defines the *Iterated* FGLS estimator: in this case, the process of estimation for Ω through residuals is iterated until convergence.
"""
struct IteratedFGLS<:LinearRegressionType
    iterations::Int
    intercept::Bool
    robust::Any

    function IteratedFGLS(iterations::Int=1; intercept::Bool=true, robust=false)
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        if iterations <= 0
            error("Number of iterations must be => 1")
        end
        new(iterations, intercept, robust)
    end
end

# ---- GENERAL LINEAR MODEL REGRESSION ----

abstract type GLMRegressionType<:AbstractRegressionType end

struct Logit<:GLMRegressionType
    intercept::Bool
    robust::Any

    function Logit(; intercept::Bool=true, robust=false)
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(intercept, robust)
    end
end

# ---- NONLINEAR REGRESSION ----

"""
## Non-Linear Regression

The abstract `NonLinearRegressionType` is the supertype for all *non-linear* regressions.
"""
abstract type NonLinearRegressionType<:AbstractRegressionType end

"""
### 'Optimize' Regression

Type that defines a regression where the relationship between response and regressors is a function. In this case, the package *Optim.jl* is used to carry out the minimization of the sum of squared residuals.
"""
struct Optimize<:NonLinearRegressionType
    method::AbsOptim
    modelfun
    beta0
    intercept::Bool
    robust::Any

    function Optimize{T<:Number}(modelfun, beta0::Array{T}, intercept::Bool=true, robust=false; method::AbsOptim=NewtonTrustRegion())
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(method, modelfun, beta0, intercept, robust)
    end
end

# #########   FITTING/RESULT TYPES   #########

"""
## Fitting Types

Every regression produces a collection of results/estimates.
Every output is thus packaged in a precise type, depending on the originating regression type, and inherits from the abstract type `AbstractFittingResult`.
"""
abstract type AbstractFittingResult end

"""
### Generic Fitting Result

Type that implements a neutral fitting result.
Should not be used in actual implementations.
"""
struct genericfitresult<:AbstractFittingResult
    callertype
    N
    beta
    residuals
    mse
    variance
    tstat
    tsig
    confint

    function genericfitresult{T<:Number}(callertype::Any, beta::Array{T}, residuals::Array{T}, variance::Array{T})
        N = length(residuals)
        mse = sum(residuals.^2)
        tstat = squeeze(beta ./ sqrt.(diag(variance)), 2)
        tsig = 2 * CDF(-abs.(tstat), (N - length(beta)), Student)
        confint = [beta - 2*sqrt.(diag(variance)), beta + 2*sqrt.(diag(variance))]
        new(callertype, N, beta, residuals, mse, variance, tstat, tsig, confint)
    end
end

"""
### Linear Fitting Result

Type that wraps results for outputs of a `LinearRegressionType` regression.
"""
struct linearfitresult<:AbstractFittingResult
    callertype
    N
    beta
    residuals
    mse
    variance
    tstat
    tsig
    confint
    robust::Any
    iterations::Int

    function linearfitresult{T<:Number}(callertype::Any, X::Array{T}, Y::Array{T}, beta::Array{T}, residuals::Array{T}, variance::Array{T}, robust::Any, iterations::Int=0)
        N = length(residuals)
        mse = sum(residuals.^2)
        tstat = squeeze(beta ./ sqrt.(diag(variance)), 2)
        tsig = 2 * CDF(-abs.(tstat), (N - length(beta)), Student)
        confint = [beta - 2*sqrt.(diag(variance)), beta + 2*sqrt.(diag(variance))]
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(callertype, N, beta, residuals, mse, variance, tstat, tsig, confint, robust, iterations)
    end
end

"""
### GLM Fitting Result

Type that wraps results for outputs of a `GLMRegressionType` regression.
"""

struct glmfitresult<:AbstractFittingResult
    callertype
    N
    beta
    residuals
    fit
    variance
    tstat
    tsig
    confint
    robust::Any
    converged::Bool
    iterations::Int

    function glmfitresult{T<:Number}(callertype::Logit, X::Array{T}, Y::Array{T}, beta::Array{T}, residuals::Array{T}, variance::Array{T}, robust::Any, converged::Bool, iterations::Int=0)
        N = length(residuals)
        fit = -2*logLike(X, Y, beta, Logistic)
        tstat = beta ./ sqrt.(diag(variance))
        if length(size(tstat)) > 1
            tstat = squeeze(tstat, 2)
        end
        tsig = 2 * CDF(-abs.(tstat), (N - length(beta)), Student)
        confint = [beta - 2*sqrt.(diag(variance)), beta + 2*sqrt.(diag(variance))]
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(callertype, N, beta, residuals, fit, variance, tstat, tsig, confint, robust, converged, iterations)
    end
end

"""
### Non-Linear Fitting Result

Type that wraps results for outputs of a `NonLinearRegressionType` regression.
"""
struct nlinfitresult<:AbstractFittingResult
    callertype
    N
    beta
    residuals
    mse
    variance
    tstat
    tsig
    confint
    robust::Any
    converged::Bool
    iterations::Int

    function nlinfitresult{T<:Number}(callertype::Any, X::Array{T}, Y::Array{T}, beta::Array{T}, residuals::Array{T}, variance::Array{T}, robust::Any, converged::Bool, iterations::Int=0)
        """
        beta = Optim.minimizer(optimresult)
        residuals = (Y - modelfun(X, beta))
        mse = Optim.minimum(optimresult)
        converged = Optim.converged(optimresult)
        """
        N = length(residuals)
        mse = sum(residuals.^2)
        tstat = beta ./ sqrt.(diag(variance))
        tsig = 2 * CDF(-abs.(tstat), (N - length(beta)), Student)
        confint = [beta - 2*sqrt.(diag(variance)), beta + 2*sqrt.(diag(variance))]
        if robust != false && !(typeof(robust) <: AbstractVarianceMatrix)
            error("Robust option is unrecognized")
        end
        new(callertype, N, beta, residuals, mse, variance, tstat, tsig, confint, robust, converged, iterations)
    end
end

# #########   ESTIMATION TYPES   #########

"""
## Estimation Types

Additional types that help with the estimation process of the regressions.

### Variance Matrices

Variance-covariance matrices are key elements in regression analysis, and require care.
Different types of *standard* and *robust* variance matrices are provided, and can be used standalone, if required.
All variance matrices types inherit from the abstract type `AbstractVarianceMatrix`.
"""
abstract type AbstractVarianceMatrix end

"""
### Basic-Naive Variance Matrix

Type that implements the naive variance-covariance matrix.
"""
struct BasicVariance<:AbstractVarianceMatrix

    function BasicVariance()
        new()
    end

    function BasicVariance{T<:Number}(X::Array{T}, residuals::Array{T}, k::Number)
        s2 = sum(residuals.^2) / (length(residuals) - k)
        return s2 .* inv(X' * X)
    end
end

"""
### HCE (Eicker-White) Heteroskedastic Consistent Variance Matrix

Type that implements the HCE variance-covariance matrix à la Eicker-White.
"""
struct HCEVariance<:AbstractVarianceMatrix

    function HCEVariance()
        new()
    end

    function HCEVariance{T<:Number}(X::Array{T}, residuals::Array{T})
        return inv(X' * X) * (X' * diagm(squeeze(residuals, 2).^2) * X) * inv(X' * X)
    end
end


function deferred_acceptance(apply_prefs::Array{Int64}, accept_prefs::Array{Int64}, caps=ones(Int64, length(accept_prefs[1, :])))
    num_of_apply = length(accept_prefs[:,1])-1
    num_of_accept = length(accept_prefs[1,:])

    ranking = zeros(Int64, num_of_apply+1, num_of_accept)

    da = zeros(Int64, num_of_apply+1, num_of_accept+1)

    num_of_da = zeros(Int64, num_of_accept)

    for i in 1:num_of_apply+1
        if i == num_of_apply+1
            search = 0
        else
            search = i
        end

        for j in 1:num_of_accept
            ranking[i, j] = findfirst(accept_prefs[:,j], search)
        end
    end

    rem = [1 for i in 1:length(apply_prefs[1,:])]
    depth = ones(Int64, length(accept_prefs[:,1])-1)

    b=1
    while true
        apply = findfirst(rem, 1)
        if apply == 0
            break
        end

        d = depth[apply]
        if d >= num_of_apply
            rem[apply] = 0
        else

            applied = apply_prefs[d, apply]

            if applied == 0
                rem[apply] = 0
            else
                my_rank = ranking[apply, applied]
                if num_of_da[applied] < caps[applied]
                    if my_rank < ranking[end, applied]
                        da[apply, applied] = 1
                        num_of_da[applied] += 1
                        rem[apply] = 0
                    else
                        depth[apply] += 1
                    end
                else
                    list_of_da = findn(da[:, applied])

                    da_rank = Array(Int64, length(list_of_da))

                    for i in 1:length(list_of_da)
                        da_rank[i] = ranking[list_of_da[i], applied]
                    end

                    worst_st_rank = maximum(da_rank)
                    worst_st = list_of_da[findfirst(da_rank, worst_st_rank)]

                    if my_rank < worst_st_rank
                        rem[apply] = 0
                        da[apply, applied] = 1
                        da[worst_st, applied] = 0
                        rem[worst_st] = 1
                        depth[worst_st] += 1
                    else
                        depth[apply] += 1
                    end

                end
            end
        end
    end

    accept_matched = Array(Int64, 0)
    apply_matched = zeros(Int64, num_of_apply)

    for i in 1:num_of_accept
        das = findn(da[:, i])
        append!(das, zeros(Int64, caps[i] - length(das)))
        append!(accept_matched, das)
        for d in das
            if d != 0
                apply_matched[d] = i
            end
        end

    end

    indptr = Array(Int64, num_of_accept+1)
    indptr[1] = 1
    for i in 1:num_of_accept
        indptr[i+1] = indptr[i] + caps[i]
    end
    return apply_matched, accept_matched, indptr
end


function hshoe_kernel(
	cx,
	cy,
	cz,
	x1,
	x2,
	x3,
	p1x,
	p1y,
	p1z,
	p2x,
	p2y,
	p2z;
	M∞::Real = 0.0
)

	β = sqrt(1.0 - M∞ ^ 2)

	ax = cx .- p1x'
	ay = (cy .- p1y') .* β
	az = (cz .- p1z') .* β

	bx = cx .- p2x'
	by = (cy .- p2y') .* β
	bz = (cz .- p2z') .* β

	na = @. sqrt(ax ^ 2 + ay ^ 2 + az ^ 2) + 1e-10
	nb = @. sqrt(bx ^ 2 + by ^ 2 + bz ^ 2) + 1e-10

	ab = @. ax * bx + ay * by + az * bz

	mult = @. (1.0 / na + 1.0 / nb) / (ab + na * nb)

	d2 = @. na - ax * x1 - ay * x2 - az * x3
	d3 = @. nb - bx * x1 - by * x2 - bz * x3

	mult2 = @. 1.0 / (na * d2)
	mult3 = @. 1.0 / (nb * d3)

	vx = @. (
		(ay * bz - az * by) * mult +
		(ay * x3 - az * x2) * mult2 -
		(by * x3 - bz * x2) * mult3
	) / (4 * π * β ^ 2)
	vy = @. (
		(az * bx - ax * bz) * mult +
		(az * x1 - ax * x3) * mult2 -
		(bz * x1 - bx * x3) * mult3
	) / (4 * π * β)
	vz = @. (
		(ax * by - ay * bx) * mult +
		(ax * x2 - ay * x1) * mult2 -
		(bx * x2 - by * x1) * mult3
	) / (4 * π * β)

	[vx, vy, vz]

end

function vlm_kernel(
	cx,
	cy,
	cz,
	p1x,
	p1y,
	p1z,
	p2x,
	p2y,
	p2z;
	M∞::Real = 0.0
)

	β = sqrt(1.0 - M∞ ^ 2)

	ax = cx .- p1x'
	ay = (cy .- p1y') .* β
	az = (cz .- p1z') .* β

	bx = cx .- p2x'
	by = (cy .- p2y') .* β
	bz = (cz .- p2z') .* β

	na = @. sqrt(ax ^ 2 + ay ^ 2 + az ^ 2) + 1e-10
	nb = @. sqrt(bx ^ 2 + by ^ 2 + bz ^ 2) + 1e-10

	ab = @. ax * bx + ay * by + az * bz

	mult = @. (1.0 / na + 1.0 / nb) / (ab + na * nb)

	vx = @. (
		ay * bz - az * by
	) * mult / (4 * π * β ^ 2)
	vy = @. (
		az * bx - ax * bz
	) * mult / (4 * π * β)
	vz = @. (
		ax * by - ay * bx
	) * mult / (4 * π * β)

	[vx, vy, vz]

end


# ------------------------------------------------------------------------------------------
# # Formatting The Plots
#
# Since the plotting functionality is implemented as a recipe to Plots.jl, [all of the
# options open to Plots.jl can be used in our
# plots](https://juliaplots.github.io/supported/). In addition, there are special features
# specifically for [differential equation
# plots](http://docs.juliadiffeq.org/latest/basics/plot.html). This tutorial will teach some
# of the most commonly used options. Let's first get the solution to some ODE. Here I will
# use one of the Lorenz ordinary differential equation. As with all commands in
# DifferentialEquations.jl, I got a plot of the solution by calling `solve` on the problem,
# and `plot` on the solution:
# ------------------------------------------------------------------------------------------

using DifferentialEquations, Plots
gr()
lorenz = @ode_def Lorenz begin
  dx = σ*(y-x)
  dy = ρ*x-y-x*z
  dz = x*y-β*z
end σ β ρ

p = [10.0,8/3,28]
u0 = [1., 5., 10.]
tspan = (0., 100.)
prob = ODEProblem(lorenz, u0, tspan, p)
sol = solve(prob)

plot(sol)

# ------------------------------------------------------------------------------------------
# Now let's change it to a phase plot. As discussed in the [plot functions
# page](http://docs.juliadiffeq.org/latest/basics/plot.html), we can use the `vars` command
# to choose the variables to plot. Let's plot variable `x` vs variable `y` vs variable `z`:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(:x,:y,:z))

# ------------------------------------------------------------------------------------------
# We can also choose to plot the timeseries for a single variable:
# ------------------------------------------------------------------------------------------

plot(sol,vars=[:x])

# ------------------------------------------------------------------------------------------
# Notice that we were able to use the variable names because we had defined the problem with
# the macro. But in general, we can use the indices. The previous plots would be:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3))
plot(sol,vars=[1])

# ------------------------------------------------------------------------------------------
# Common options are to add titles, axis, and labels. For example:
# ------------------------------------------------------------------------------------------

plot(sol,linewidth=5,title="Solution to the linear ODE with a thick line",
xaxis="Time (t)",yaxis="u(t) (in mm)",label=["X","Y","Z"])

# ------------------------------------------------------------------------------------------
# Notice that series recipes apply to the solution type as well. For example, we can use a
# scatter plot on the timeseries:
# ------------------------------------------------------------------------------------------

scatter(sol,vars=[:x])

# ------------------------------------------------------------------------------------------
# This shows that the recipe is using the interpolation to smooth the plot. It becomes
# abundantly clear when we turn it off using `denseplot=false`:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3),denseplot=false)

# ------------------------------------------------------------------------------------------
# When this is done, only the values the timestep hits are plotted. Using the interpolation
# usually results in a much nicer looking plot so it's recommended, and since the
# interpolations have similar orders to the numerical methods, their results are trustworthy
# on the full interval. We can control the number of points used in the interpolation's plot
# using the `plotdensity` command:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3),plotdensity=100)

# ------------------------------------------------------------------------------------------
# That's plotting the entire solution using 100 points spaced evenly in time.
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3),plotdensity=10000)

# ------------------------------------------------------------------------------------------
# That's more like it! By default it uses `100*length(sol)`, where the length is the number
# of internal steps it had to take. This heuristic usually does well, but unusually
# difficult equations it can be relaxed (since it will take small steps), and for equations
# with events / discontinuities raising the plot density can help resolve the discontinuity.
#
# Lastly notice that we can compose plots. Let's show where the 100 points are using a
# scatter plot:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3))
scatter!(sol,vars=(1,2,3),plotdensity=100)

# ------------------------------------------------------------------------------------------
# We can instead work with an explicit plot object. This form can be better for building a
# complex plot in a loop.
# ------------------------------------------------------------------------------------------

p = plot(sol,vars=(1,2,3))
scatter!(p,sol,vars=(1,2,3),plotdensity=100)
title!("I added a title")

# ------------------------------------------------------------------------------------------
# You can do all sorts of things. Have fun!
# ------------------------------------------------------------------------------------------


"""
    shiftedkernel = centered(kernel)

Shift the origin-of-coordinates to the center of `kernel`. The
center-element of `kernel` will be accessed by `shiftedkernel[0, 0,
...]`.

This function makes it easy to supply kernels using regular Arrays,
and provides compatibility with other languages that do not support
arbitrary axes.

See also: [`imfilter`](@ref).
"""
function centered(A::AbstractArray)
    offsetted = first.(axes(A)) .- 1
    total_offsets = .-((size(A) .+ 1 ) .÷ 2)
    OffsetArray(A, total_offsets .- offsetted)
end

"""
    kernfft = freqkernel([T::Type], kern, sz=size(kern); rfft=false)

Return a frequency-space representation of `kern`.
This embeds `kern` in an array of size `sz`, in a manner that implicitly imposes periodic boundary
conditions, and then returns the fourier transform. This is sometimes called the optical transfer
function, and known in some frameworks as `psf2otf`. If `rfft` is `true`, the fft for real-valued
arrays (`rfft`) is returned instead and the first dimension size will be approximately half of `sz[1]`.

`kern` should be zero-centered, i.e., `kern[0, 0]` should reference the center of your kernel.
See [`centered`](@ref). Optionally specify the numeric type `T` (which must be one of the
types supported by FFTW, either `Float32` or `Float64`).

The inverse of `freqkernel` is [`spacekernel`](@ref).
"""
function freqkernel(::Type{T}, kern::AbstractArray, sz::Dims=size(kern); rfft=false) where T<:Union{Float32,Float64}
    wrapindex(i, s) = i<1 ? i+s : i
    all(size(kern) .<= sz) || throw(DimensionMismatch("kernel size $(size(kern)) is bigger than supplied size $sz"))
    kernw = zeros(T, sz...)
    for I in CartesianIndices(kern)
        J = CartesianIndex(map(wrapindex, Tuple(I), sz))
        kernw[J] = kern[I]
    end
    return rfft ? FFTW.rfft(kernw) : fft(kernw)
end
freqkernel(kern::AbstractArray{T}, args...; rfft=false) where T =
    freqkernel(ffteltype(T), kern, args...; rfft=rfft)

"""
    kern = spacekernel(kernfft, axs; rfftsz=0)

Return a real-space representation of `kernfft`, a frequency-space representation of a kernel.
This performs an inverse fourier transform, implicitly imposes periodic boundary conditions,
and then trims & truncates axes of the output to `axs`. By default `kernfft` is assumed to have
been generated by `fft`; if it was instead generated by `rfft`, specify the original size
of the first dimension. (If `kernfft` was generated by [`freqkernel`](@ref), this is just `sz[1]`.)

The inverse of `spacekernel` is [`freqkernel`](@ref).
"""
function spacekernel(kernfft::AbstractArray, axs::Indices; rfftsz=0)
    wrapindex(i, s) = i<1 ? i+s : i
    kernw = rfftsz > 0 ? irfft(kernfft, rfftsz) : ifft(kernfft)
    kern = zeros(eltype(kernw), axs...)
    sz = size(kernw)
    for I in CartesianIndices(kern)
        J = CartesianIndex(map(wrapindex, Tuple(I), sz))
        kern[I] = kernw[J]
    end
    return kern
end
function spacekernel(kernfft::AbstractArray; rfftsz=0)
    sz = size(kernfft)
    if rfftsz > 0
        sz = (rfftsz, Base.tail(sz)...)
    end
    upper = map(s->s>>1, sz)
    axs = map((u, s)-> u-s+1:u, upper, sz)
    return spacekernel(kernfft, axs; rfftsz=rfftsz)
end

ffteltype(::Type{T}) where T<:Union{Float32,Float64} = T
ffteltype(::Type{Float16}) = Float32
ffteltype(::Type{T}) where T<:Normed = ffteltype(floattype(T))
ffteltype(::Type{T}) where T = Float64

dummyind(::Base.OneTo) = Base.OneTo(1)
dummyind(::AbstractUnitRange) = 0:0

dummykernel(inds::Indices{N}) where {N} = fill(1, map(dummyind, inds))

nextendeddims(inds::Indices) = sum(ind->length(ind)>1, inds)
nextendeddims(a::AbstractArray) = nextendeddims(axes(a))

function checkextended(inds::Indices, n)
    dimstr = n == 1 ? "dimension" : "dimensions"
    nextendeddims(inds) != n && throw(ArgumentError("need $n extended $dimstr, got axes $inds"))
    nothing
end
checkextended(a::AbstractArray, n) = checkextended(axes(a), n)

_reshape(A::OffsetArray{<:Any,N}, ::Val{N}) where N = A
_reshape(A::OffsetArray, ::Val{N}) where {N} = OffsetArray(reshape(parent(A), Val(N)), fill_to_length(A.offsets, -1, Val(N)))
_reshape(A::AbstractArray, ::Val{N}) where {N} = reshape(A, Val(N))

_vec(a::AbstractVector) = a
_vec(a::AbstractArray) = (checkextended(a, 1); a)
_vec(a::OffsetArray{<:Any,1}) = a
function _vec(a::OffsetArray)
    inds = axes(a)
    checkextended(inds, 1)
    i = findall(ind->length(ind)>1, inds)
    OffsetArray(vec(parent(a)), inds[i])
end

samedims(::Val{N}, kernel) where {N} = _reshape(kernel, Val(N))
samedims(::Val{N}, kernel::Tuple) where {N} = map(k->_reshape(k, Val(N)), kernel)
samedims(::AbstractArray{<:Any,N}, kernel) where {N} = samedims(Val(N), kernel)

_tail(R::CartesianIndices{0}) = R
_tail(R::CartesianIndices) = CartesianIndices(tail(axes(R)))

# ensure that overflow is detected, by ensuring that it doesn't happen
# at intermediate stages of the computation
accumfilter(pixelval, filterval) = pixelval * filterval
const SmallInts = Union{UInt8,Int8,UInt16,Int16}
accumfilter(pixelval::SmallInts, filterval::SmallInts) = Int(pixelval)*Int(filterval)
# advice: don't use FixedPoint for the kernel
accumfilter(pixelval::N0f8, filterval::N0f8) = Float32(pixelval)*Float32(filterval)
accumfilter(pixelval::Colorant{N0f8}, filterval::N0f8) = float32(c)*Float32(filterval)

# In theory, the following might need to be specialized. For safety, make it a
# standalone function call.
safe_for_prod(x, ref) = oftype(ref, x)


"""
    laplace(decisionMat)

    Apply Laplace method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.

# Output
- `::LaplaceResult`: LaplaceResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        3000 2750 2500 2250;
        1500 4750 8000 7750;
        2000 5250 8500 11750
]

julia> dm = makeDecisionMatrix(mat)

julia> result = laplace(dm)
```
"""
function laplace(decisionMatrix::DataFrame)::LaplaceResult

    n, p = size(decisionMatrix)

    probs = [1.0 / p for i in 1:p]

    expecteds = zeros(Float64, n)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    for i in 1:n
        expecteds[i] = (m[i,:] .* probs)  |> sum
    end

    bestIndex = sortperm(expecteds) |> last

    result = LaplaceResult(
        expecteds,
        bestIndex
    )

    return result
end

"""
    maximin(decisionMat)

    Apply Maximin method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.

# Output
- `::MaximinResult`: MaximinResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)

julia> result = maximin(dm)
```
"""
function maximin(decisionMatrix::DataFrame)::MaximinResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmins = rowmins(decisionMatrix)

    bestIndex = sortperm(rmins) |> last

    result = MaximinResult(
        rmins,
        bestIndex
    )

    return result
end

"""
    maximax(decisionMat)

    Apply Maximax method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.

# Output
- `::MaximaxResult`: MaximaxResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)

julia> result = maximax(dm)
```
"""
function maximax(decisionMatrix::DataFrame)::MaximaxResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmaxs = rowmaxs(decisionMatrix)

    bestIndex = sortperm(rmaxs) |> last

    result = MaximaxResult(
        rmaxs,
        bestIndex
    )

    return result
end

"""
    minimax(decisionMat)

    Apply Minimax method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.

# Output
- `::MinimaxResult`: MinimaxResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)

julia> result = minimax(dm)
```
"""
function minimax(decisionMatrix::DataFrame)::MinimaxResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmaxs = rowmaxs(decisionMatrix)

    bestIndex = sortperm(rmaxs) |> first

    result = MinimaxResult(
        rmaxs,
        bestIndex
    )

    return result
end

"""
    minimin(decisionMat)

    Apply Minimin method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.

# Output
- `::MiniminResult`: Minimin object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)

julia> result = minimin(dm)
```
"""
function minimin(decisionMatrix::DataFrame)::MiniminResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmins = rowmins(decisionMatrix)

    bestIndex = sortperm(rmins) |> first

    result = MiniminResult(
        rmins,
        bestIndex
    )

    return result
end

"""
    savage(decisionMat)

    Apply Savage method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.

# Output
- `::SavageResult`: SavageResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)

julia> result = savage(dm)

julia> result.bestIndex
4
```
"""
function savage(decisionMatrix::DataFrame)::SavageResult

    n, p = size(decisionMatrix)

    cmaxs = colmaxs(decisionMatrix)

    newDecisionMatrix = similar(decisionMatrix)

    @inbounds for i in 1:p
        newDecisionMatrix[:, i] = cmaxs[i] .- decisionMatrix[:, i]
    end

    rmaxs = scores = rowmaxs(newDecisionMatrix)

    bestIndex = rmaxs |> sortperm |> first

    result = SavageResult(
        newDecisionMatrix,
        scores,
        bestIndex
    )

    return result
end

"""
    hurwicz(decisionMat; alpha = 0.5)

    Apply Hurwicz method for a given decision matrix (for convenience, in type of DataFrame).

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.
 - `alpha::Float64`: The optional alpha value for the Hurwicz method. Default is 0.5.

# Output
- `::HurwiczResult`: HurwiczResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)

julia> result = hurwicz(dm)

julia> result.bestIndex
3
```
"""
function hurwicz(decisionMatrix::DataFrame; alpha::Float64=0.5)::HurwiczResult

    n, p = size(decisionMatrix)

    rmaxs = rowmaxs(decisionMatrix)
    rmins = rowmins(decisionMatrix)

    scores = alpha .* rmaxs .+ (1.0 - alpha) .* rmins

    bestIndex = scores |> sortperm |>  last

    result = HurwiczResult(
        scores,
        bestIndex
    )

    return result

end

"""
    mle(decisionMat, weights)

    Apply MLE (Maximum Likelihood) method for a given decision matrix (for convenience, in type of DataFrame) and weights.

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.
 - `weights::Array{Float64,1}`: Array of weights for each criterion that sums up to 1.0.

# Output
- `::MLEResult`: MLEResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)
julia> weights = [0.2, 0.5, 0.2, 0.1]
julia> result = mle(dm)

julia> result.bestIndex
2
```
"""
function mle(decisionMatrix::DataFrame, weights::Array{Float64,1})::MLEResult

    w = unitize(weights)

    weightedMatrix = w * decisionMatrix

    scores = rowsums(weightedMatrix)

    bestIndex = scores |> sortperm |> last

    result = MLEResult(
        scores,
        bestIndex
    )

    return result

end

"""
    expectedregret(decisionMat, weights)

    Apply Expected Regret method for a given decision matrix (for convenience, in type of DataFrame) and weights.

# Arguments:
 - `decisionMat::DataFrame`: Decision matrix with n alternatives and m criteria.
 - `weights::Array{Float64,1}`: Array of weights for each criterion that sums up to 1.0.

# Output
- `::ExpectedRegretResult`: ExpectedRegretResult object that holds multiple outputs including the best alternative.

# Examples
```julia-repl
julia> mat = [
        26 26 18 22;
        22 34 30 18;
        28 24 34 26;
        22 30 28 20
    ]

julia> dm = makeDecisionMatrix(mat)
julia> weights = [0.2, 0.5, 0.2, 0.1]
julia> result = expectedregret(dm)

julia> result.bestIndex
2
```
"""
function expectedregret(decisionMatrix::DataFrame, weights::Array{Float64,1})

    _, p = size(decisionMatrix)

    w = unitize(weights)

    cmaxs = colmaxs(decisionMatrix)

    regretmat = similar(decisionMatrix)

    for i in 1:p
        regretmat[:, i] = cmaxs[i] .- decisionMatrix[:, i]
    end

    weightedregret = w * regretmat

    scores = rowsums(weightedregret)

    bestIndex = scores |> sortperm |> first

    result = ExpectedRegretResult(
        scores,
        bestIndex
    )

    return result
end


"""
    countedinsertionsort!(a)

Sorts an array and returns the number of swaps required to sort.

This is useful when you want to know the number
of transpositions of vertices in a simplex. This determines the new
orientation of the simplex relative to the old one.
"""
function countedinsertionsort!(a)
    swaps = zero(Int64)
    for i in 2:length(a)
        key = a[i]
        j = i - 1
        while j >= 1 && key < a[j]
            a[j + 1] = a[j]
            j -= 1
            a[j + 1] = key
            swaps += 1
        end
    end
    swaps
end


function test_window_functions()
    try SPTK.Cwindow(0) catch @test false end
    try SPTK.Cwindow(5) catch @test false end
    @test_throws ArgumentError SPTK.Cwindow(-1)
    @test_throws ArgumentError SPTK.Cwindow(6)

    println("test windows functions")
    srand(98765)
    x = rand(1024)

    for f in [blackman, hamming, hanning, bartlett, trapezoid, rectangular]
        println("-- test_$f")
        y = f(length(x))
        @test length(y) == length(x)
        @test all(isfinite.(y))
    end

    # invalid normalize flag
    @test_throws ArgumentError blackman(512, normalize=-1)
    @test_throws ArgumentError blackman(512, normalize=3)
end

test_window_functions()


@testset "Cubic" begin
    for (constructor, copier) in ((interpolate, identity), (interpolate!, copy))
        isinplace = constructor == interpolate!
        f0(x) = sin((x-3)*2pi/9 - 1)
        f1(x) = 1.0 + 0.1*x + 0.01*x^2 + 0.001*x^3

        xmax = 10
        A0 = Float64[f0(x) for x in 1:xmax]
        A1 = Float64[f1(x) for x in 1:xmax]

        f2(x, y) = sin(x/10)*cos(y/6)
        xmax2, ymax2 = 30, 10
        A2 = Float64[f2(x, y) for x in 1:xmax2, y in 1:ymax2]

        for BC in (Line, Flat, Free, Periodic), GT in (OnGrid, OnCell)
            for (A, f) in ((A0, f0), (A1, f1))
                itp1 = @inferred(constructor(copier(A), BSpline(Cubic(BC(GT())))))
                ax1 = axes(itp1)[1]
                @test Interpolations.lbounds(itp1) == (GT == OnGrid ? (first(ax1),) : (first(ax1) - 0.5,))
                @test Interpolations.ubounds(itp1) == (GT == OnGrid ? (last(ax1),) : (last(ax1) + 0.5,))
                @test_throws ArgumentError parent(itp1)
                check_axes(itp1, A, isinplace)
                check_inbounds_values(itp1, A)
                check_oob(itp1)
                can_eval_near_boundaries(itp1)

                # test that inner region is close to data
                for x in 3.1:.2:8.1
                    @test f(x) ≈ itp1(x) atol=abs(0.1 * f(x))
                end
            end

            itp2 = @inferred(constructor(copier(A2), BSpline(Cubic(BC(GT())))))
            @test_throws ArgumentError parent(itp2)
            check_axes(itp2, A2, isinplace)
            check_inbounds_values(itp2, A2)
            check_oob(itp2)
            can_eval_near_boundaries(itp2)

            for x in 3.1:.2:xmax2-3, y in 3.1:2:ymax2-3
                @test f2(x,y) ≈ itp2(x,y) atol=abs(0.1 * f2(x,y))
            end
        end
        # Issue 419: Incorrect results with BSpline(Cubic(Periodic(OnCell()))))
        let BC = Periodic, GT = OnCell
            for (A, f) in ((A0, f0), (A1, f1))
                itp1 = @inferred(constructor(copier(A), BSpline(Cubic(BC(GT())))))
                @test itp1(Interpolations.lbounds(itp1)...) ≈ itp1(Interpolations.ubounds(itp1)...)
            end
        end
    end

    ix = 1:15
    k = length(ix) - 1
    f(x) = cos((x-1)*2pi/k)
    g(x) = -2pi/k * sin((x-1)*2pi/k)

    A = map(f, ix)

    for (constructor, copier) in ((interpolate, identity), (interpolate!, copy))

        for BC in (Line, Flat, Free, Periodic), GT in (OnGrid,OnCell)

            itp = constructor(copier(A), BSpline(Cubic(BC(GT()))))
            # test that inner region is close to data
            for x in range(ix[5], stop=ix[end-4], length=100)
                @test g(x) ≈ Interpolations.gradient1(itp,x) atol=cbrt(cbrt(eps(g(x))))
            end
        end
        # Issue 419: Incorrect results with BSpline(Cubic(Periodic(OnCell()))))
        let BC = Periodic, GT = OnCell
            itp = @inferred(constructor(A[1:end-1], BSpline(Cubic(BC(GT())))))
            @test itp(Interpolations.lbounds(itp)...) ≈ itp(Interpolations.ubounds(itp)...)
        end
    end
    itp_flat_g = interpolate(A, BSpline(Cubic(Flat(OnGrid()))))
    @test Interpolations.gradient(itp_flat_g,1)[1] ≈ 0 atol=eps()
    @test Interpolations.gradient(itp_flat_g,ix[end])[1] ≈ 0 atol=eps()

    itp_flat_c = interpolate(A, BSpline(Cubic(Flat(OnCell()))))
    @test Interpolations.gradient(itp_flat_c,0.5)[1] ≈ 0 atol=eps()
    @test Interpolations.gradient(itp_flat_c,ix[end] + 0.5)[1] ≈ 0 atol=eps()

    # Can construct from coefficients
    itp = interpolate(rand(5, 9), BSpline(Cubic(Flat(OnGrid()))))
    @test itp == Interpolations.BSplineInterpolation(itp.coefs, itp.it, itp.parentaxes)
    @test_throws ArgumentError Interpolations.BSplineInterpolation(itp.coefs, itp.it, (2:3, 2:10))
end


"""
    const FactorTable = Dict{NamedTuple,Float64}

Type alias of a Dictionary mapping tuples of Variable name=>value assignments to probabilities.
"""
const FactorTable = Dict{NamedTuple,Float64}

"""
    struct Variable

A Type representing a single discrete variable.

Params:

    name::Symbol
name of Variable

    m::Int:
number of possible (discrete) values this Variable can take
"""
struct Variable
    name::Symbol
    m::Int # number of possible values
end

"""
    struct Factor

A Type representing a joint or conditional probability distribution over the variables in vars.

Fields:

    vars:
The variables represented by this distribution

    table:
Probability table mapping Variable assignments (NamedTuple) => probability (Float64)
"""
struct Factor
    vars::Vector{Variable}
    table::FactorTable
end

"""
    function variablenames(φ::Factor)

Get a vector of variable names (Symbols) for the Variables in a Factor
"""
variablenames(φ::Factor) = [var.name for var in φ.vars]

"""
    function assignments(vars::AbstractVector{Variable})

Get an array of all possible variable=>value assignment tuples for Variables in vars.
"""
function assignments(vars::AbstractVector{Variable})

    # get a vector of variable names
    n = [var.name for var in vars]

    # product(1:v1.m, 1:v2.m, 1:v3.m) -> cartesian product of integer indicies
    #                                 -> iterator over tuples of indicies
    return [namedtuple(n, v) for v in product((1:v.m for v in vars)...)]
end

"""
    function normalize!(φ::Factor)

Update φ::Factor so all values sum to 1.0 while maintaining relative value (so values represent probabilities).
"""
function normalize!(ϕ::Factor)

    # sum all probabilities in Factor table
    z = sum(p for (a,p) in ϕ.table)

    # divide each value by total sum
    for (a,p) in ϕ.table
            ϕ.table[a] = p/z
    end

    return ϕ
end

"""
    struct BayesianNetwork

Represents a Bayesian Network.

Fields:

    vars::Vector{Variable}:
Vector of Variables contained in network

    factors::Vector{Factor}:
Vector of Factors, one for each Variable in vars, representing conditional probability distribution over that variable given parents of var in network.

    graph::SimpleDiGraph{Int64}:
Represents the structure of the network variables and implicitly defines independence relationships.

"""
struct BayesianNetwork

    """Vector of every Variable in the network"""
    vars::Vector{Variable}

    """
        Network conditional independence specified by a `Vector` of `Factor`s.
        Each Factor speifies the conditional joint PMF over its subset of variables, given the
    """
    factors::Vector{Factor}

    """
        DAG that specifies the independence relationships among variables.
        Nodes
    """
    graph::SimpleDiGraph{Int64}
end

"""
    function probability(bn::BayesianNetwork, assignment)

Calculate the probability of a given variable assignment given the bayesian network bn.

# Arguments

* bn: The BayesianNetwork describing the joint distribution
* assignment: a NamedTuple of :Symbol=>Integer pairs giving the
            categorical index value of each variable.

# Examples
```julia
julia> X = Variable(:x, 2);

julia> Y = Variable(:y, 2);

julia> Z = Variable(:z, 2);

julia> ϕx = Factor([X, Y, Z], FactorTable(
           (x=1, y=1, z=1) => 0.08,
           (x=1, y=1, z=2) => 0.31,
           (x=1, y=2, z=1) => 0.09,
           (x=1, y=2, z=2) => 0.37,
           (x=2, y=1, z=1) => 0.01,
           (x=2, y=1, z=2) => 0.05,
           (x=2, y=2, z=1) => 0.02,
           (x=2, y=2, z=2) => 0.07,
       ));

julia> ϕy = Factor([Y], FactorTable(
           (y=1,) => 0.20,
           (y=2,) => 0.80,
       ));

julia> ϕz = Factor([Z], FactorTable(
           (z=1,) => 0.40,
           (z=2,) => 0.60,
       ));

julia> g = SimpleDiGraph(3);

julia> add_edge!(g,1,3);

julia> add_edge!(g,2,3);

julia> bn = BayesianNetwork([Y, Z, X], [ϕy, ϕz, ϕ], g);

julia> assignment = (x=1, y=2, z=1);

julia> probability(bn, assignment)
0.028800000000000006
```
"""
function probability(bn::BayesianNetwork, assignment)

        # Function to map Factor to a NamedTuple containing only the name=>value pairs for names specified in the (conditionally independent) Factor φ and values for them (and all other variables in bn) specified in `assignment`
        __select(φ) = select(assignment, variablenames(φ))

        # Function mapping Factor to probability of factor assignment
        probability(φ) = φ.table[__select(φ)]

        # Final probability is product of individual factor's probabilities (since they are implicitly independent)
        return prod(probability(φ) for φ in bn.factors)

end


using StochasticGroundMotionSimulation
using Plots
using QuadGK
using BenchmarkTools
using Test
# using ForwardDiff
using Polynomials
using PolynomialRoots

fi = exp10.(range(-2.0, stop=2.0, length=2001))

m = 6.0
r = 100.0
fas = FASParams(100.0, [1.0, 50.0, Inf], [1.0, 0.5], 200.0, 0.4, 0.039)
sdof = Oscillator(1.0)

Afi = map(f -> fourier_spectral_ordinate(f, m, r, fas), fi)
Hfi = transfer(fi, sdof)

fac = 1.5

plot(fi, Afi.^2, lab="Af^2", box=:true, grid=:on, xscale=:log10)
plot!(fi, (Hfi .* Afi).^2, lab="Hf^2Af^2")
vline!([sdof.f_n, sdof.f_n/fac, sdof.f_n*fac])
xlims!(sdof.f_n/ceil(fac),sdof.f_n*ceil(fac))

fun(f) = squared_transfer(f, sdof)

fac = 1.5

fj = collect(range(sdof.f_n/fac, stop=sdof.f_n*fac, length=101))
H2j = map(f -> fun(f), fj)

@time igk = quadgk(fun, sdof.f_n/fac, sdof.f_n*fac)[1]
@time isr = simpsons_rule(fj, H2j)
@time itr = trapezoidal_rule(fj, H2j)

@test igk ≈ isr rtol=1e-5
@test igk ≈ itr rtol=1e-5

fun(f) = squared_transfer(f, sdof) * ( fourier_spectral_ordinate(f, m, r, fas)^2 )

fj = collect(range(sdof.f_n/fac, stop=sdof.f_n*fac, length=101))
H2j = map(f -> fun(f), fj)

@time igk = quadgk(fun, sdof.f_n/fac, sdof.f_n*fac)[1]
@time isr = simpsons_rule(fj, H2j)
@time itr = trapezoidal_rule(fj, H2j)

@test igk ≈ isr rtol=1e-5
@test igk ≈ itr rtol=1e-4

plot(fi, Afi.^2, lab="Af^2", box=:true, grid=:on)
plot!(fi, (Hfi .* Afi).^2, lab="Hf^2Af^2")
# vline!([sdof.f_n, sdof.f_n/fac, sdof.f_n*fac])
xlims!(1e-3, sdof.f_n/fac)

fj = collect(range(1e-3, stop=sdof.f_n/fac, length=101))
H2j = map(f -> fun(f), fj)

@time igk = quadgk(fun, 1e-3, sdof.f_n/fac)[1]
@time isr = simpsons_rule(fj, H2j)
@time itr = trapezoidal_rule(fj, H2j)

@test igk ≈ isr rtol=1e-5
@test igk ≈ itr rtol=1e-3

fj = collect(range(sdof.f_n*fac, stop=1e1, length=101))
H2j = map(f -> fun(f), fj)

@time igk = quadgk(fun, sdof.f_n*fac, 1e1)[1]
@time isr = simpsons_rule(fj, H2j)
@time itr = trapezoidal_rule(fj, H2j)

@test igk ≈ isr rtol=1e-3
@test igk ≈ itr rtol=1e-3

fun(f) = (2*π*f)^2 * squared_transfer(f, sdof) * ( fourier_spectral_ordinate(f, m, r, fas)^2 )

fj = collect(range(sdof.f_n/fac, stop=sdof.f_n*fac, length=101))
H2j = map(f -> fun(f), fj)
fk = collect(range(sdof.f_n/fac, stop=sdof.f_n*fac, length=71))
H2k = map(f -> fun(f), fk)

@time igk = quadgk(fun, sdof.f_n/fac, sdof.f_n*fac)[1]
@time isr = simpsons_rule(fj, H2j)
@time itr = trapezoidal_rule(fk, H2k)

@test igk ≈ isr rtol=1e-5
@test igk ≈ itr rtol=1e-4

function fourier_derivative(f, m, r, fas)
    βQ = π * r / ( fas.Q0 * fas.cQ )
    βκ = π * fas.κ0
    fc, ~, ~ = corner_frequency(m, fas)

    return -βκ*f^4 - βQ*(1.0-fas.η)*f^(4-fas.η) + 2*(1.0-fc^2)*f^3 - βκ*fc^2*f^2 - βQ*(1.0-fas.η)*fc^2*f^(2-fas.η) + 2*fc^2*f
end

fas = FASParams(100.0, [1.0, 50.0, Inf], [1.0, 0.5], 200.0, 0.0, 0.039)

fi = exp10.(range(-2, stop=2, length=101))
Afi = fourier_spectrum(fi, m, r, fas)
dAfi = map(f -> fourier_derivative(f, m, r, fas), fi)

# FAS(f) = fourier_spectral_ordinate(f, m, r, fas)
# dFAS(f) = ForwardDiff.derivative(FAS, f)
#
# gAfi = map(f -> dFAS(f), fi)

βQ = π * r / ( fas.Q0 * fas.cQ )
βκ = π * fas.κ0
fc, ~, ~ = corner_frequency(m, fas)

dAfp = Polynomial([ 0, 2*fc^2, -fc^2*(βQ+βκ), 2*(1.0-fc^2), -(βQ+βκ)])
ri = PolynomialRoots.roots([ 0, 2*fc^2, -fc^2*(βQ+βκ), 2*(1.0-fc^2), -(βQ+βκ)])

real.(ri)
imag.(ri)

plot(fi, Afi, lab="FAS")
vline!([abs(ri[1])])

plot(fi, dAfi, lab="dFAS")
ylims!(-1e2,1e2)


# =================
# Dependencies
# =================

using ReachabilityAnalysis, CarlemanLinearization
using Plots, LaTeXStrings, LinearAlgebra, SparseArrays
using Plots.PlotMeasures

using LazySets: center
using CarlemanLinearization: _error_bound_specabs_R

include("../utils.jl")

# =================
# Model definition
# =================

function quadratic(; a, b, N)
    dl = zeros(N-1)
    d = [a*i for i in 1:N]
    du = [b*i for i in 1:N-1]
    return Tridiagonal(dl, d, du)
end

function logistic(; r, K, N)
    return quadratic(a=r, b=-r/K, N=N)
end

# define model
r = -0.5
K = 0.8

# standard form
F1 = hcat(r)
F2 = hcat(-r/K)
a = r
b = -r/K
A = logistic(r=r, K=K, N=4);

x0 = interval(0.5)
R, Re_λ1 = _error_bound_specabs_R(x0, F1, F2; check=true)

using Revise
Tmax = 10.0

x0set = interval(0.47, 0.53)

NSTEPS = 1000

# analytic solution
u0 = mid(x0)
tsp = range(0.0, Tmax, length=NSTEPS+1)
f_analytic(t, u0) = u0 * a * exp(a * t) / (a + b*(1 - exp(a*t))*u0)
sol_a = f_analytic.(tsp, u0)

# solution with Carleman, discrete time
function sol_carlin_discrete_time(; N)
    A = logistic(r=r, K=K, N=N)
    y0 = center(kron_pow_stack(x0, N)) |> Singleton
    prob = @ivp(y' = A * y, y(0) ∈ y0, dim=N)

    return solve(prob, T=Tmax, alg=ORBIT(δ=0.5, approx_model=NoBloating()))
end

# solution with Carleman, continuous time
function sol_carlin_continuous_time(; N)
    F1 = hcat(r); F2 = hcat(-r/K)
    dirs = _template(n=1, N=N)
    alg = LGG09(δ=0.1, template=dirs, approx_model=Forward())
    return _solve_CARLIN(x0set, F1, F2; alg=alg, N=N, T=Tmax, bloat=false)
end

# ============
# Figure 1a
# ============

function figure_1a()
    sol_discr_N1 = sol_carlin_discrete_time(N=1)
    sol_discr_N2 = sol_carlin_discrete_time(N=2)
    sol_discr_N3 = sol_carlin_discrete_time(N=3)
    sol_discr_N4 = sol_carlin_discrete_time(N=4)
    sol_discr_N5 = sol_carlin_discrete_time(N=5)
    sol_discr_N6 = sol_carlin_discrete_time(N=6);

    fig = plot(legend=:topright, xlab=L"t", ylab=L"x(t)",
               legendfontsize=25,
               tickfont=font(25, "Times"),
               guidefontsize=25,
               xguidefont=font(35, "Times"),
               yguidefont=font(35, "Times"),
               xtick=([0.0, 2.5, 5.0, 7.5, 10.0], [L"0", L"2.5", L"5.0", L"7.5", L"10.0"]),
               ytick=([0.0, 0.1, 0.2, 0.3, 0.4, 0.5], [L"0", L"0.1", L"0.2", L"0.3", L"0.4", L"0.5"]),
               xlims=(0.0, 10.1),
               ylims=(-0.01, 0.5),
               bottom_margin=10mm, left_margin=6mm, right_margin=6mm, top_margin=3mm, size=(900, 600))

    plot!(fig, tsp, sol_a, c=:magenta, lab=L"\textrm{Analytic}", lw=3.0, alpha=1.)

    plot!(fig, sol_discr_N1, vars=(0, 1), lab=L"N=1", seriestype=:path, linestyle=:dash, c=:green, markersize=8, markershape=:circle, lw=2.0, alpha=1.)
    plot!(fig, sol_discr_N2, vars=(0, 1), lab=L"N=2", seriestype=:path, linestyle=:dash, c=:red, markersize=8, markershape=:utriangle, lw=2.0, alpha=1.)
    plot!(fig, sol_discr_N4, vars=(0, 1), lab=L"N=4", seriestype=:path, linestyle=:dash, c=:blue, markersize=8, markershape=:rect, lw=2.0, alpha=1.)

    plot!(fig, sol_discr_N6, vars=(0, 1), lab=L"N=6", seriestype=:path, linestyle=:dash, c=:orange, markersize=8, markershape=:star, lw=2.0, alpha=1.)

end

fig = figure_1a()
savefig(fig, joinpath(TARGET_FOLDER, "figure_1a.pdf"))

# ============
# Figure 1b
# ============

function figure_1b()
    sol_discr_N1 = sol_carlin_discrete_time(N=1)
    sol_discr_N2 = sol_carlin_discrete_time(N=2)
    sol_discr_N3 = sol_carlin_discrete_time(N=3)
    sol_discr_N4 = sol_carlin_discrete_time(N=4)
    sol_discr_N5 = sol_carlin_discrete_time(N=5)
    sol_discr_N6 = sol_carlin_discrete_time(N=6);

    fig = plot(legend=:topright, xlab=L"t", ylab=L"x(t)",
               legendfontsize=25,
               tickfont=font(25, "Times"),
               guidefontsize=25,
               xguidefont=font(35, "Times"),
               yguidefont=font(35, "Times"),
               xtick=([0.0, 2.5, 5.0, 7.5, 10.0], [L"0", L"2.5", L"5.0", L"7.5", L"10.0"]),
               ytick=([0.0, 0.1, 0.2, 0.3, 0.4, 0.5], [L"0", L"0.1", L"0.2", L"0.3", L"0.4", L"0.5"]),
               xlims=(0.0, 10.1),
               ylims=(-0.01, 0.6),
               bottom_margin=10mm, left_margin=6mm, right_margin=6mm, top_margin=3mm, size=(900, 600))

    tsp = range(0.0, Tmax, length=NSTEPS+1)

    sol_a_high = f_analytic.(tsp, x0set.hi)
    sol_a_low = f_analytic.(tsp, x0set.lo)

    # initial set
    plot!(Singleton([0.0]) × Interval(x0set), lw=8.0, alpha=1., c=:black, lab="", markershape=:none, seriestype=:shape, xlims=(-0.05, 10.1))

    plot!(fig, sol_carlin_continuous_time(N=1), vars=(0, 1), lab="", c=:green, lc=:green)
    plot!(fig, sol_carlin_continuous_time(N=2), vars=(0, 1), lab="", c=:red, lc=:red)
    plot!(fig, sol_carlin_continuous_time(N=6), vars=(0, 1), lab="", c=:orange, lc=:orange)

    plot!(fig, sol_discr_N1, vars=(0, 1), lab=L"N=1", seriestype=:path, linestyle=:dash, c=:green, markersize=8, markershape=:circle, lw=2.0, alpha=1.)
    plot!(fig, sol_discr_N2, vars=(0, 1), lab=L"N=2", seriestype=:path, linestyle=:dash, c=:red, markersize=8, markershape=:utriangle, lw=2.0, alpha=1.)
    plot!(fig, sol_discr_N6, vars=(0, 1), lab=L"N=6", seriestype=:path, linestyle=:dash, c=:orange, markersize=8, markershape=:star, lw=2.0, alpha=1.)

    plot!(fig, tsp, sol_a_high, c=:magenta, lab=L"\textrm{Analytic}", lw=3.0, alpha=1.)
    plot!(fig, tsp, sol_a_low, c=:magenta, lab="", lw=3.0, alpha=1.)

end

fig = figure_1b()
savefig(fig, joinpath(TARGET_FOLDER, "figure_1b.pdf"))


using  ApproxFunRational, AbstractIterativeSolvers
### Vector case
tol = 1.e-9
Ds1 = 0.5
Ds2 = 1.
Δs = Ds2-Ds1
Dn1 = 1.
Dn2 = 0.5
Δn = Dn2 - Dn1
rn = 1.
s0n = 1/sqrt(2.)
rs = 1.
s0s = 1/sqrt(2.)
pnn = 1
pmm = .2
L = 1.
D = 1. # overall delay

fPss = z -> 2*rs*s0s^2.0/(rs^2+z.^2)
#fPss = z -> .1*exp(-z^2/2)
fPnn = z -> pnn
fPmm = z -> pmm
Pss = Fun(zai(fPss),OscLaurent(0.0,L))
One = pad(Fun(1.,OscLaurent(0.0,L)),3)
Pnn = pnn*One
Pmm = pmm*One

H = zeros(Fun,2,2)
a = Pss + Pnn + Pmm - One
b = Pss*Fun(1.,OscLaurent(-Δs,L)) + Pnn*Fun(1.,OscLaurent(-Δn,L))
bt = Pss*Fun(1.,OscLaurent(Δs,L)) + Pnn*Fun(1.,OscLaurent(Δn,L))
H[1,1] = copy(a)
H[2,2] = copy(a)
H[1,2] = copy(b)
H[2,1] = copy(bt)
G = map(SumFun,copy(H))

# Set up right-hand side
b1 = Pss*Fun(1.,OscLaurent(- D + Ds1,L))
b2 = Pss*Fun(1.,OscLaurent(- D + Ds2,L))
h = map(SumFun,[b1,b2])

𝓒 = Cauchy(-1)
𝓢 = Cauchy(1)
simp(f) = combine!(chop!(f,tol))
inner(a,b) = ⋅(simp(a),simp(b))

function op(x::Array{T,1}) where T<:SumFun
    println("Apply Cauchy:")
    @time y = 𝓒*x
    println("Apply G:")
    @time y = G*y
    println("Subtract:")
    @time y = (x - y)
    return y
end

sop = x -> ( x - G*(𝓒*x))
out = GMRES(op,h,⋅,10*tol,30,simp)

u = sum([out[2][i]*out[1][i] for i=1:length(out[2])])

si_op(out[1][end])

𝓕 = FourierTransform(-1.0)
U = 𝓕*u

x = 0:.01:10
y1 = real(map(U[1],x))
y2 = real(map(U[2],x))

for i = 1:length(y1)-1
    if isinf(y1[i])
        y1[i] = y1[i+1]
    end
    if isinf(y2[i])
        y2[i] = y2[i+1]
    end
end

plot(x,y1)
plot!(x,y2)


using BenchmarkTools
using Distributed

function calc_pi_by_gcd(N)
    s = 0
    for a in 1:N
        for b in 1:N
            s += ifelse(gcd(a,b)==1, 1, 0)
        end
        s
    end
    println(sqrt(6N^2/s))
end

#@time calc_pi_by_gcd(20000)

function calc_pi_by_gcd_parallel(N)
    """
    Note that @parallel has been renamed into @difrom Julia 1.0
    https://discourse.julialang.org/t/loaderror-parallel/15126
    """
    c = @distributed (+) for a in 1:N
        s = 0
        for b in 1:N
            s += ifelse(gcd(a,b)==1, 1, 0)
        end
        s
    end
    println(sqrt(6N^2/c))
end

addprocs()

#@time calc_pi_by_gcd_parallel(20000)


push!(LOAD_PATH, "Modules")
using Colors
using JuliRay

φ=MathConstants.φ

## Tetrahedron
vertices=[[1.0,1.0,1.0],[1.0,-1.0,-1.0],[-1.0,1.0,-1.0],[-1.0,-1.0,1.0]]
edges=[[1,2],[1,3],[1,4],[2,3],[2,4],[3,4]]
faces=[[2,3,4],[1,3,4],[1,2,4],[1,2,3]]

V=rgbColor(csgUnion((p->Sphere(p,0.05)).(vertices)),RGB(0.1,0.1,0.1))
E=rgbColor(csgUnion([Cylinder(vertices[I...],vertices[J...],0.025) for (I,J) in edges]),RGB(0.2,0.2,0.2))
F=rgbftColor(csgUnion([Polygon((i->vertices[i]).(I)) for I in faces]),RGB(1,0,0),FT(0.2,0.2))
render(csgUnion(V,E,F),camera=LngLatCamera(lng=π/5,lat=π/5,pers=0.2,zoom=0.25),name="PlatonicSolid_4")

## Cube
vertices=[[[i,j,k] for i ∈ [1,-1], j ∈ [1,-1], k ∈ [1,-1]]...]
edges=[[[1+i+j,2+i+j] for i ∈ [0,2], j ∈ [0,4]]...]
edges=vcat(edges,[[[1+i+j,3+i+j] for i ∈ [0,1], j ∈ [0,4]]...])
edges=vcat(edges,[[[1+i+j,5+i+j] for i ∈ [0,1], j ∈ [0,2]]...])
faces=[[1,2,4,3],[1,2,4,3].+4,[1,3,7,5],[1,3,7,5].+1,[1,5,6,2],[1,5,6,2].+2]

V=rgbColor(csgUnion((p->Sphere(p,0.05)).(vertices)),RGB(0.1,0.1,0.1))
E=rgbColor(csgUnion([Cylinder(vertices[I...],vertices[J...],0.025) for (I,J) in edges]),RGB(0.2,0.2,0.2))
F=rgbftColor(csgUnion([Polygon((i->vertices[i]).(I)) for I in faces]),RGB(1,0,0),FT(0.2,0.2))
render(csgUnion(V,E,F),camera=LngLatCamera(lng=π/5,lat=π/5,pers=0.2,zoom=0.25),name="PlatonicSolid_6")

## Octahedron
vertices=[[1,0,0],[-1,0,0],[0,1,0],[0,-1,0],[0,0,1],[0,0,-1]]*sqrt(3)
edges=[[1,3],[3,2],[2,4],[4,1],[3,5],[5,4],[4,6],[6,3],[5,1],[1,6],[6,2],[2,5]]
faces=[[1,3,5],[2,3,5],[2,4,5],[1,4,5],[1,3,6],[2,3,6],[2,4,6],[1,4,6]]

V=rgbColor(csgUnion((p->Sphere(p,0.05)).(vertices)),RGB(0.1,0.1,0.1))
E=rgbColor(csgUnion([Cylinder(vertices[I...],vertices[J...],0.025) for (I,J) in edges]),RGB(0.2,0.2,0.2))
F=rgbftColor(csgUnion([Polygon((i->vertices[i]).(I)) for I in faces]),RGB(1,0,0),FT(0.2,0.2))
render(csgUnion(V,E,F),camera=LngLatCamera(lng=π/5,lat=π/5,pers=0.2,zoom=0.25),name="PlatonicSolid_8")

## Dodecahedron
vertices0=[[[i,j,k] for i ∈ [1,-1], j ∈ [1,-1], k ∈ [1,-1]]...]
vertices1=[[[0,i/φ,j*φ] for i ∈ [1,-1], j ∈ [1,-1]]...]
vertices2=[[[j*φ,0,i/φ] for i ∈ [1,-1], j ∈ [1,-1]]...]
vertices3=[[[i/φ,j*φ,0] for i ∈ [1,-1], j ∈ [1,-1]]...]
vertices=vcat(vertices0,vertices1,vertices2,vertices3)
edges1=[[1,9],[2,9],[3,10],[4,10],[9,10]]
edges2=[[5,11],[6,11],[7,12],[8,12],[11,12]]
edges3=[[1,13],[3,13],[5,14],[7,14],[13,14]]
edges4=[[2,15],[4,15],[6,16],[8,16],[15,16]]
edges5=[[1,17],[5,17],[2,18],[6,18],[17,18]]
edges6=[[3,19],[7,19],[4,20],[8,20],[19,20]]
edges=vcat(edges1,edges2,edges3,edges4,edges5,edges6)
faces=[[1,9,10,3,13],[2,9,10,4,15],[5,11,12,7,14],[6,11,12,8,16],[1,13,14,5,17],[3,13,14,7,19],[2,15,16,6,18],[4,15,16,8,20],[1,17,18,2,9],[5,17,18,6,11],[3,19,20,4,10],[7,19,20,8,12]]

V=rgbColor(csgUnion((p->Sphere(p,0.05)).(vertices)),RGB(0.1,0.1,0.1))
E=rgbColor(csgUnion([Cylinder(vertices[I...],vertices[J...],0.025) for (I,J) in edges]),RGB(0.2,0.2,0.2))
F=rgbftColor(csgUnion([Polygon((i->vertices[i]).(I)) for I in faces]),RGB(1,0,0),FT(0.2,0.2))
render(csgUnion(V,E,F),camera=LngLatCamera(lng=π/5,lat=π/5,pers=0.2,zoom=0.25),name="PlatonicSolid_12")

## Icosahedron
vertices1=[[[0,i,j*φ] for i ∈ [1,-1], j ∈ [1,-1]]...]
vertices2=[[[j*φ,0,i] for i ∈ [1,-1], j ∈ [1,-1]]...]
vertices3=[[[i,j*φ,0] for i ∈ [1,-1], j ∈ [1,-1]]...]
vertices=vcat(vertices1,vertices2,vertices3)*sqrt(3)/√(φ+2)
edges=[[2i-1,2i] for i ∈ 1:6]
edges=[edges...,[1,5],[5,9],[9,1],[2,5],[5,11],[11,2],[3,6],[6,9],[9,3],[4,6],[6,11],[11,4],[1,7],[7,10],[10,1],[2,7],[7,12],[12,2],[3,8],[8,10],[10,3],[4,8],[8,12],[12,4]]
faces=[[1,5,9],[2,5,11],[3,6,9],[4,6,11],[1,7,10],[2,7,12],[3,8,10],[4,8,12],[1,2,5],[1,2,7],[3,4,6],[3,4,8],[5,6,9],[5,6,11],[7,8,10],[7,8,12],[9,10,1],[9,10,3],[11,12,2],[11,12,4]]

V=rgbColor(csgUnion((p->Sphere(p,0.05)).(vertices)),RGB(0.1,0.1,0.1))
E=rgbColor(csgUnion([Cylinder(vertices[I...],vertices[J...],0.025) for (I,J) in edges]),RGB(0.2,0.2,0.2))
F=rgbftColor(csgUnion([Polygon((i->vertices[i]).(I)) for I in faces]),RGB(1,0,0),FT(0.2,0.2))
render(csgUnion(V,E,F),camera=LngLatCamera(lng=π/5,lat=π/5,pers=0.2,zoom=0.25),name="PlatonicSolid_20")


# https://fenicsproject.org/docs/dolfin/latest/python/demos/mixed-poisson/demo_mixed-poisson.py.html

module mixed_poisson

using FenicsPy

# Create mesh
mesh = UnitSquareMesh(32, 32)

# Define finite elements spaces and build mixed space
BDM = FiniteElement("BDM", mesh.ufl_cell(), 1)
DG  = FiniteElement("DG", mesh.ufl_cell(), 0)
W = FunctionSpace(mesh, BDM * DG)

# Define trial and test functions
(sigma, u) = TrialFunctions(W)
(tau, v) = TestFunctions(W)

# Define source function
f = Expression("10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
g = Expression("sin(5*x[0])", degree=2)

# Define variational form
a = (dot(sigma, tau) + div(tau)*u + div(sigma)*v)*dx
L = - f*v*dx

# Define function G such that G \cdot n = g
py"""
from math import sin
from dolfin import Cell, UserExpression
class BoundarySource(UserExpression):
    def __init__(self, mesh, **kwargs):
        self.mesh = mesh
        super().__init__(**kwargs)
    def eval_cell(self, values, x, ufc_cell):
        cell = Cell(self.mesh, ufc_cell.index)
        n = cell.normal(ufc_cell.local_facet)
        g = sin(5*x[0])
        values[0] = g*n[0]
        values[1] = g*n[1]
    def value_shape(self):
        return (2,)
"""

G = Expression(py"BoundarySource"(mesh.pyobject, degree=2))

# Define essential boundary
#def boundary(x):
#    return x[1] < DOLFIN_EPS or x[1] > 1.0 - DOLFIN_EPS
boundary = "x[1] < $DOLFIN_EPS || x[1] > 1.0 - $DOLFIN_EPS"

bc = DirichletBC(W.sub(0), G, boundary)

# Compute solution
w = FeFunction(W)
solve(a == L, w, bc)
(sigma, u) = w.split()

# Save solution to file in VTK format
vtkfile = File("mixed_poisson/sigma.pvd")
vtkfile << sigma

vtkfile = File("mixed_poisson/u.pvd")
vtkfile << u

# Plot sigma and u
#plot(sigma)
#plot(u)

end # module mixed_poisson


struct KernelRegression{R,T}
    x::R
    y::Vector{T}
    k::Function
    h::Float64
end

# Want a version for x as a vector (this) and X as a matrix (doesn't exist yet)
# function kernelregression{R <: Any,
#                           T <: Real}(x::Vector{R},
#                                      y::Vector{T};
#                                      kernel::Symbol = :epanechnikov,
#                                      bandwidth::Real = NaN,
#                                      getbandwidth::Function = kNN.bandwidth)
#     if isnan(bandwidth)
#         h_x = getbandwidth(x)
#         h_y = getbandwidth(y)
#         h = sqrt(h_x * h_y)
#     else
#         h = bandwidth
#     end
#     k = SmoothingKernels.kernels[kernel]
#     return KernelRegression(x, y, k, h)
# end

function StatsBase.predict(model::KernelRegression, x::Real)
    y, n = 0.0, length(model.x)
    h, k = model.h, model.k
    normalizer = 0.0
    for i in 1:n
        # TODO: Get definitions and usage of bandwidth right
        # w_i = (1 / h) * k((x - model.x[i]) / h)
        w_i = h * k(h * (x - model.x[i]))
        y += model.y[i] * w_i
        normalizer += w_i
    end
    return y / normalizer
end

function StatsBase.predict!(ys::Vector, model::KernelRegression, xs::AbstractVector{T}) where {T}
    n = length(xs)
    # @assert length(ys) == n
    for i in 1:n
        ys[i] = predict(model, xs[i])
    end
    return
end

function StatsBase.predict(model::KernelRegression, xs::AbstractVector{T}) where {T}
    ys = Array{T}(undef,length(xs))
    predict!(ys, model, xs)
    return ys
end


using JuliaPackageDemo
using MittagLeffler
using Test

const testx = 0.5

function _coolfunction1(x)
    (2*x + x^2) == coolfunction1(x)
end
@test _coolfunction1(testx)

function _coolfunction2(x)
    (50*mittleff(0.5, -x) - 10) == coolfunction2(x)
end
@test _coolfunction2(testx)

include(joinpath(@__DIR__, "..", "docs", "src", "simpleexamples.jl"))