mutable struct UnivariateKalmanState{Fl<:AbstractFloat}
    v::Fl
    F::Fl
    att::Vector{Fl}
    a::Vector{Fl}
    Ptt::Matrix{Fl}
    P::Matrix{Fl}
    llk::Fl
    steady_state::Bool
    # Auxiliary matrices
    P_to_check_steady_state::Matrix{Fl}
    ZP::Vector{Fl}
    TPtt::Matrix{Fl}
    function UnivariateKalmanState(a1::Vector{Fl}, P1::Matrix{Fl}) where Fl
        m = length(a1)
        P_to_check_steady_state = zeros(Fl, m, m)
        ZP = zeros(Fl, m)
        TPtt = zeros(Fl, m, m)

        return new{Fl}(
            zero(Fl),
            zero(Fl),
            zeros(Fl, m),
            a1,
            zeros(Fl, m, m),
            P1,
            zero(Fl),
            false,
            P_to_check_steady_state,
            ZP,
            TPtt,
        )
    end
end

function save_a1_P1_in_filter_output!(
    filter_output::FilterOutput{Fl}, kalman_state::UnivariateKalmanState{Fl}
) where Fl
    filter_output.a[1] = deepcopy(kalman_state.a)
    filter_output.P[1] = deepcopy(kalman_state.P)
    return filter_output
end

function save_kalman_state_in_filter_output!(
    filter_output::FilterOutput{Fl}, kalman_state::UnivariateKalmanState{Fl}, t::Int
) where Fl
    filter_output.v[t] = copy(fill(kalman_state.v, 1))
    filter_output.F[t] = copy(fill(kalman_state.F, 1, 1))
    filter_output.a[t + 1] = copy(kalman_state.a)
    filter_output.att[t] = copy(kalman_state.att)
    filter_output.P[t + 1] = copy(kalman_state.P)
    filter_output.Ptt[t] = copy(kalman_state.Ptt)
    # There is no diffuse part in UnivariateKalmanState
    filter_output.Pinf[t] = copy(fill(zero(Fl), size(kalman_state.Ptt)))
    return filter_output
end

# Univariate Kalman filter with the recursions as described
# in Koopman's book TODO
"""
    UnivariateKalmanFilter{Fl <: AbstractFloat}

A Kalman filter that is tailored to univariate systems, exploiting the fact that the
dimension of the observations at any time period is 1.

# TODO equations and descriptions of a1 and P1
"""
mutable struct UnivariateKalmanFilter{Fl<:AbstractFloat} <: KalmanFilter
    steadystate_tol::Fl
    a1::Vector{Fl}
    P1::Matrix{Fl}
    skip_llk_instants::Int
    kalman_state::UnivariateKalmanState

    function UnivariateKalmanFilter(
        a1::Vector{Fl},
        P1::Matrix{Fl},
        skip_llk_instants::Int=length(a1),
        steadystate_tol::Fl=Fl(1e-5),
    ) where Fl
        kalman_state = UnivariateKalmanState(copy(a1), copy(P1))
        return new{Fl}(steadystate_tol, a1, P1, skip_llk_instants, kalman_state)
    end
end

function reset_filter!(kf::UnivariateKalmanFilter{Fl}) where Fl
    copyto!(kf.kalman_state.a, kf.a1)
    copyto!(kf.kalman_state.P, kf.P1)
    set_state_llk_to_zero!(kf.kalman_state)
    fill!(kf.kalman_state.P_to_check_steady_state, zero(Fl))
    fill!(kf.kalman_state.ZP, zero(Fl))
    fill!(kf.kalman_state.TPtt, zero(Fl))
    kf.kalman_state.steady_state = false
    return kf
end

function set_state_llk_to_zero!(kalman_state::UnivariateKalmanState{Fl}) where Fl
    kalman_state.llk = zero(Fl)
    return kalman_state
end

# TODO this should either have ! or return something
function check_steady_state(kalman_state::UnivariateKalmanState{Fl}, tol::Fl) where Fl
    @inbounds for j in axes(kalman_state.P, 2), i in axes(kalman_state.P, 1)
        if abs(
            (kalman_state.P[i, j] - kalman_state.P_to_check_steady_state[i, j]) /
            kalman_state.P[i, j],
        ) > tol
            # Update the P_to_check_steady_state matrix
            copyto!(kalman_state.P_to_check_steady_state, kalman_state.P)
            return nothing
        end
    end
    kalman_state.steady_state = true
    return nothing
end

function update_v!(
    kalman_state::UnivariateKalmanState{Fl}, y::Fl, Z::Vector{Fl}, d::Fl
) where Fl
    kalman_state.v = y - dot(Z, kalman_state.a) - d
    return kalman_state
end

function update_F!(kalman_state::UnivariateKalmanState{Fl}, Z::Vector{Fl}, H::Fl) where Fl
    LinearAlgebra.BLAS.gemv!('N', one(Fl), kalman_state.P, Z, zero(Fl), kalman_state.ZP)
    kalman_state.F = H
    kalman_state.F += dot(kalman_state.ZP, Z)
    return kalman_state
end

function update_att!(kalman_state::UnivariateKalmanState{Fl}, Z::Vector{Fl}) where Fl
    copyto!(kalman_state.att, kalman_state.a)
    @inbounds for i in axes(kalman_state.P, 1), j in axes(kalman_state.P, 2)
        kalman_state.att[i] +=
            (kalman_state.v / kalman_state.F) * kalman_state.P[i, j] * Z[j]
    end
    return kalman_state
end

function repeat_a_in_att!(kalman_state::UnivariateKalmanState{Fl}) where Fl
    for i in eachindex(kalman_state.att)
        kalman_state.att[i] = kalman_state.a[i]
    end
    return kalman_state
end

function update_a!(
    kalman_state::UnivariateKalmanState{Fl}, T::Matrix{Fl}, c::Vector{Fl}
) where Fl
    LinearAlgebra.BLAS.gemv!('N', one(Fl), T, kalman_state.att, zero(Fl), kalman_state.a)
    kalman_state.a .+= c
    return kalman_state
end

function update_Ptt!(kalman_state::UnivariateKalmanState{Fl}) where Fl
    LinearAlgebra.BLAS.gemm!(
        'N', 'T', one(Fl), kalman_state.ZP, kalman_state.ZP, zero(Fl), kalman_state.Ptt
    )
    @. kalman_state.Ptt = kalman_state.P - (kalman_state.Ptt / kalman_state.F)
    return kalman_state
end

function repeat_P_in_Ptt!(kalman_state::UnivariateKalmanState{Fl}) where Fl
    for i in axes(kalman_state.P, 1), j in axes(kalman_state.P, 2)
        kalman_state.Ptt[i, j] = kalman_state.P[i, j]
    end
    return kalman_state
end

function update_P!(
    kalman_state::UnivariateKalmanState{Fl}, T::Matrix{Fl}, RQR::Matrix{Fl}
) where Fl
    LinearAlgebra.BLAS.gemm!(
        'N', 'N', one(Fl), T, kalman_state.Ptt, zero(Fl), kalman_state.TPtt
    )
    LinearAlgebra.BLAS.gemm!(
        'N', 'T', one(Fl), kalman_state.TPtt, T, zero(Fl), kalman_state.P
    )
    kalman_state.P .+= RQR
    return kalman_state
end

function update_P!(
    kalman_state::UnivariateKalmanState{Fl}, T::Matrix{Fl}, R::Matrix{Fl}, Q::Matrix{Fl}
) where Fl
    LinearAlgebra.BLAS.gemm!(
        'N', 'N', one(Fl), T, kalman_state.Ptt, zero(Fl), kalman_state.TPtt
    )
    LinearAlgebra.BLAS.gemm!(
        'N', 'T', one(Fl), kalman_state.TPtt, T, zero(Fl), kalman_state.P
    )
    kalman_state.P .+= R * Q * R'
    return kalman_state
end

function update_llk!(kalman_state::UnivariateKalmanState{Fl}) where Fl
    kalman_state.llk -= (
        HALF_LOG_2_PI + (log(kalman_state.F) + kalman_state.v^2 / kalman_state.F) / 2
    )
    return kalman_state
end

function update_kalman_state!(
    kalman_state::UnivariateKalmanState{Fl},
    y::Fl,
    Z::Vector{Fl},
    T::Matrix{Fl},
    H::Fl,
    RQR::Matrix{Fl},
    d::Fl,
    c::Vector{Fl},
    skip_llk_instants::Int,
    tol::Fl,
    t::Int,
) where Fl
    if isnan(y)
        kalman_state.v = NaN
        update_F!(kalman_state, Z, H)
        repeat_a_in_att!(kalman_state)
        update_a!(kalman_state, T, c)
        repeat_P_in_Ptt!(kalman_state)
        update_P!(kalman_state, T, RQR)
        kalman_state.steady_state = false # Not on steadystate anymore
    elseif kalman_state.steady_state
        update_v!(kalman_state, y, Z, d)
        update_att!(kalman_state, Z)
        update_a!(kalman_state, T, c)
        if t > skip_llk_instants
            update_llk!(kalman_state)
        end
    else
        update_v!(kalman_state, y, Z, d)
        update_F!(kalman_state, Z, H)
        update_att!(kalman_state, Z)
        update_a!(kalman_state, T, c)
        update_Ptt!(kalman_state)
        update_P!(kalman_state, T, RQR)
        check_steady_state(kalman_state, tol)
        if t > skip_llk_instants
            update_llk!(kalman_state)
        end
    end
    return kalman_state
end

function update_kalman_state!(
    kalman_state::UnivariateKalmanState{Fl},
    y::Fl,
    Z::Vector{Fl},
    T::Matrix{Fl},
    H::Fl,
    R::Matrix{Fl},
    Q::Matrix{Fl},
    d::Fl,
    c::Vector{Fl},
    skip_llk_instants::Int,
    t::Int,
) where Fl
    if kalman_state.steady_state
        update_v!(kalman_state, y, Z, d)
        update_att!(kalman_state, Z)
        update_a!(kalman_state, T, c)
    elseif isnan(y)
        kalman_state.v = NaN
        update_F!(kalman_state, Z, H)
        repeat_a_in_att!(kalman_state)
        update_a!(kalman_state, T, c)
        repeat_P_in_Ptt!(kalman_state)
        update_P!(kalman_state, T, R, Q)
    else
        update_v!(kalman_state, y, Z, d)
        update_F!(kalman_state, Z, H)
        update_att!(kalman_state, Z)
        update_a!(kalman_state, T, c)
        update_Ptt!(kalman_state)
        update_P!(kalman_state, T, R, Q)
    end
    if t > skip_llk_instants
        update_llk!(kalman_state)
    end
    return kalman_state
end

function optim_kalman_filter(
    sys::StateSpaceSystem, filter::UnivariateKalmanFilter{Fl}
) where Fl
    return filter_recursions!(
        filter.kalman_state, sys, filter.steadystate_tol, filter.skip_llk_instants
    )
end

function kalman_filter!(
    filter_output::FilterOutput, sys::StateSpaceSystem, filter::UnivariateKalmanFilter{Fl}
) where Fl
    filter_recursions!(
        filter_output,
        filter.kalman_state,
        sys,
        filter.steadystate_tol,
        filter.skip_llk_instants,
    )
    return filter_output
end

function filter_recursions!(
    kalman_state::UnivariateKalmanState{Fl},
    sys::LinearUnivariateTimeInvariant,
    steadystate_tol::Fl,
    skip_llk_instants::Int,
) where Fl
    RQR = sys.R * sys.Q * sys.R'
    @inbounds for t in eachindex(sys.y)
        update_kalman_state!(
            kalman_state,
            sys.y[t],
            sys.Z,
            sys.T,
            sys.H,
            RQR,
            sys.d,
            sys.c,
            skip_llk_instants,
            steadystate_tol,
            t,
        )
    end
    return kalman_state.llk
end

# TODO this doesn't use steadystate_tol
function filter_recursions!(
    kalman_state::UnivariateKalmanState{Fl},
    sys::LinearUnivariateTimeVariant,
    steadystate_tol::Fl,
    skip_llk_instants::Int,
) where Fl
    @inbounds for t in eachindex(sys.y)
        update_kalman_state!(
            kalman_state,
            sys.y[t],
            sys.Z[t],
            sys.T[t],
            sys.H[t],
            sys.R[t],
            sys.Q[t],
            sys.d[t],
            sys.c[t],
            skip_llk_instants,
            t,
        )
    end
    return kalman_state.llk
end

function filter_recursions!(
    filter_output::FilterOutput,
    kalman_state::UnivariateKalmanState{Fl},
    sys::LinearUnivariateTimeInvariant,
    steadystate_tol::Fl,
    skip_llk_instants::Int,
) where Fl
    RQR = sys.R * sys.Q * sys.R'
    save_a1_P1_in_filter_output!(filter_output, kalman_state)
    @inbounds for t in eachindex(sys.y)
        update_kalman_state!(
            kalman_state,
            sys.y[t],
            sys.Z,
            sys.T,
            sys.H,
            RQR,
            sys.d,
            sys.c,
            skip_llk_instants,
            steadystate_tol,
            t,
        )
        save_kalman_state_in_filter_output!(filter_output, kalman_state, t)
    end
    return filter_output
end

# TODO this doesn't use steadystate_tol
function filter_recursions!(
    filter_output::FilterOutput,
    kalman_state::UnivariateKalmanState{Fl},
    sys::LinearUnivariateTimeVariant,
    steadystate_tol::Fl,
    skip_llk_instants::Int,
) where Fl
    save_a1_P1_in_filter_output!(filter_output, kalman_state)
    @inbounds for t in eachindex(sys.y)
        update_kalman_state!(
            kalman_state,
            sys.y[t],
            sys.Z[t],
            sys.T[t],
            sys.H[t],
            sys.R[t],
            sys.Q[t],
            sys.d[t],
            sys.c[t],
            skip_llk_instants,
            t,
        )
        save_kalman_state_in_filter_output!(filter_output, kalman_state, t)
    end
    return filter_output
end


# Create example data
angles = LinRange(0, 360, 40)
radii = LinRange(0, 2, 40)
# Draw the polar plot in degrees
polar(angles, radii, radians=false)


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

_sec2rad{T<:AbstractFloat}(sec::T) = deg2rad(sec/3600.0)

"""
    sec2rad(sec) -> radians

### Purpose ###

Convert from seconds to radians.

### Argument ###

* `sec`: number of seconds.  It can be either a scalar or a vector.

### Output ###

The number of radians corresponding to `sec`.  If `sec` is an array, an array of
the same length is returned.

### Example ###

``` julia
sec2rad(3600*30)
# => 0.5235987755982988
```

### Notes ###

Use `rad2sec` to convert radians to seconds.
"""
sec2rad(sec::Real) = _sec2rad(float(sec))

@vectorize_1arg Real sec2rad


# This script is the translation of run_example.m in
# Dylan Mikesell's DynamicWarping repo (https://github.com/dylanmikesell/DynamicWarping.)
# run this script in the terminal: '> julia DTWDT_run_example.jl'
# 2019/06 Kurama Okubo

using PlotlyJS, ORCA, FileIO, JLD2
# This script runs an example of dynamic time warping as an introduction to
# the method.
push!(LOAD_PATH,"../src")

include("../src/functions.jl")
using .DTWDTfunctions

example = 2; # shift function (1=step, 2=sine)
# you can try two test cases. Watch the 'maxLag' parameter. The two shift
# functions have different magnitudes so if you don't go to large enough
# lags in the sine model, you won't converge to the correction solution.

# in the case of the sine wave, you must have b=1 because the strains are
# on the order of dt!! Play with 'b' to see how you get stuck in the
# wrong minimum because you're not allowing large enough jumps.

# in the case of the step function, we break the DTW because the step shift
# is larger than dt. We would need a different step pattern to correctly
# solve for this step. It's possible, just not implemented in this version
# of the dynamic warping because don't expect strains > 100#.

# (you can also change the input examples using the scripts in exampleData/
# folder.)

maxLag = 80; # max nuber of points to search forward and backward (can be npts, just takes longer and is unrealistic)
b      = 1; # b-value to limit strain
# impose a strain limit: 1 dt is b=1, half dt is b=2, etc.
# if you mess with b you will see the results in the estimated shifts

## load the data file and plot

if example == 1
    #@load "../exampledata/stepShiftData.jld2"; # data with constant shifts and step in the middle
    error("stepShiftData is not implemented.")
elseif example == 2
    @load "../exampledata/sineShiftData.jld2"; # data with sine wave shifts of one cycle
end

lvec   = (-maxLag:maxLag).*dt; # lag array for plotting below
npts   = length(u0);            # number of samples
tvec   = ( 0 : npts-1 ) .* dt; # make the time axis
stTime = st;               # shift vector in time

# plot shift function
#time shift curve

function lineplot1()
    tr = scatter(;x=tvec, y=stTime, mode="lines", name="Timeshift")
    layout = Layout(
        title="Shift applied to u_{0}(t)",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
    )
    plot(tr, layout)
end

function lineplot2()
    tr1 = scatter(;x=tvec, y=u0, mode="lines+markers", name="u0(t)")
    tr2 = scatter(;x=tvec, y=u1, mode="lines+markers", name="u(t)")
    layout = Layout(
        title = "Traces",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="Amplitude [a.u.]"),
    )
    plot([tr1, tr2], layout)
end

#p = [lineplot1(), lineplot2()]

## compute error function and plot

err = computeErrorFunction(u1, u0, npts, maxLag); # cpmpute error function over Lags
# note the error function is independent of strain limit 'b'.

function ploterrfunc()
    tr = heatmap(;x=tvec, y=lvec, z=log10.(err), colorscale="Greys",
    zauto=false, zmin=-4, zmax=0, colorbar=attr(title="L2 error", titleside="right",))
    layout = Layout(
        title = "Error function",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="Lag"),
    )
    plot(tr, layout)
end

#ploterrfunc()

## accumuluate error in FORWARD direction

direction = 1; # direction to accumulate errors (1=forward, -1=backward)
# it is instructive to flip the sign of +/-1 here to see how the function
# changes as we start the backtracking on different sides of the traces.
# Also change 'b' to see how this influences the solution for stbar. You
# want to make sure you're doing things in the proper directions in each
# step!!!

dist  = accumulateErrorFunction(direction, err, npts, maxLag, b); # forward accumulation to make distance function

function plotdistfunc()
    tr = heatmap(;x=tvec, y=lvec, z=dist, colorscale="Jet",
    zauto=true, colorbar=attr(title="Dist", titleside="right",))
    layout = Layout(
        title = "Distance function dir = +1",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
        width=800,
        height=400,
    )
    plot(tr, layout)
end
#plotdistfunc()

stbar = backtrackDistanceFunction( -1*direction, dist, err, -maxLag, b ); # find shifts

## plot the results

stbarTime = stbar .* dt;      # convert from samples to time
tvec2     = tvec + stbarTime; # make the warped time axis

# make figure
function lineplot3()
    tr1 = scatter(;x=tvec, y=stTime, mode="lines", name="Actual")
    tr2 = scatter(;x=tvec[1:10:end], y=stbarTime[1:10:end], mode="markers", name="Estimated")
    layout = Layout(
        title = "Estimated shifts",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
        width=800,
        height=400,
    )
    plot([tr1, tr2], layout)
end

p = plotdistfunc()
display(p)
savefig(p, "./SINEdistarray.png")
p = lineplot3()
display(p)
savefig(p, "./SINEcomparison.png")

## An example of computing the residual misfit

error = computeDTWerror( err, stbar, maxLag );

## Apply dynamic time warping in the backward direction
direction = -1; # direction to accumulate errors (1=forward, -1=backward)

dist  = accumulateErrorFunction(direction, err, npts, maxLag, b ); # backward accumulation to make distance function
stbar = backtrackDistanceFunction( -1*direction, dist, err, -maxLag, b ); # find shifts

## plot the results
stbarTime = stbar .* dt;      # convert from samples to time
tvec2     = tvec + stbarTime; # make the warped time axis

function plotdistfunc2()
    tr = heatmap(;x=tvec, y=lvec, z=dist, colorscale="Jet",
    zauto=true, colorbar=attr(title="Dist", titleside="right",))
    layout = Layout(
        title = "Distance function dir = -1",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
    )
    plot(tr, layout)
end

function lineplot4()
    tr1 = scatter(;x=tvec, y=stTime, mode="lines", name="Actual")
    tr2 = scatter(;x=tvec, y=stbarTime, mode="markers", name="Estimated")
    layout = Layout(
        title = "Estimated shifts",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
    )
    plot([tr1, tr2], layout)
end

p = ([plotdistfunc2(), lineplot4()])
display(p)

# plot input traces
function lineplot5()
    tr1 = scatter(;x=tvec, y=u0, mode="lines", name="Raw")
    tr2 = scatter(;x=tvec, y=u1, mode="lines", name="Shifted",
    line = attr(dash="dash"))

    layout = Layout(
        title = "Input traces for dynamic time warping dir = -1",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="Amplitude (a.u.)"),
    )
    plot([tr1, tr2], layout)
end

function lineplot6()
    tr1 = scatter(;x=tvec, y=u0, mode="lines", name="Raw")
    tr2 = scatter(;x=tvec2, y=u1, mode="lines", name="Warped",
    line = attr(dash="dash"))

    layout = Layout(
        title = "Output traces for dynamic time warping",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="Amplitude (a.u.)"),
    )
    plot([tr1, tr2], layout)
end

p = ([lineplot5(), lineplot6()])
display(p)

## Apply dynamic time warping in both directions to smooth. (Follwoing example in Hale 2013)

dist1 = accumulateErrorFunction( -1, err, npts, maxLag, b ); # forward accumulation to make distance function
dist2 = accumulateErrorFunction( 1, err, npts, maxLag, b ); # backwward accumulation to make distance function

dist  = dist1 .+ dist2 .- err; # add them and remove 'err' to not count twice (see Hale's paper)

stbar = backtrackDistanceFunction( -1, dist, err, -maxLag, b ); # find shifts
# !! Notice now that you can backtrack in either direction and get the same
# result after you smooth the distance function in this way.

## plot the results

stbarTime = stbar .* dt;      # convert from samples to time
tvec2     = tvec + stbarTime; # make the warped time axis

# plot the distance function
function plotdistfunc_s()
    tr = heatmap(;x=tvec, y=lvec, z=dist, colorscale="Jet",
    zauto=true, colorbar=attr(title="Dist", titleside="right",))
    layout = Layout(
        title = "Distance function smoothed",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
    )
    plot(tr, layout)
end

function lineplot_s1()
    tr1 = scatter(;x=tvec, y=stTime, mode="lines", name="Actual")
    tr2 = scatter(;x=tvec, y=stbarTime, mode="markers", name="Estimated")
    layout = Layout(
        title = "Estimated shifts",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="τ [s]"),
    )
    plot([tr1, tr2], layout)
end

p = ([plotdistfunc_s(), lineplot_s1()])
display(p)

# plot input traces
function lineplot_s2()
    tr1 = scatter(;x=tvec, y=u0, mode="lines", name="Raw")
    tr2 = scatter(;x=tvec, y=u1, mode="lines", name="Shifted",
    line = attr(dash="dash"))

    layout = Layout(
        title = "Input traces for dynamic time warping smoothed",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="Amplitude (a.u.)"),
    )
    plot([tr1, tr2], layout)
end

function lineplot_s3()
    tr1 = scatter(;x=tvec, y=u0, mode="lines", name="Raw")
    tr2 = scatter(;x=tvec2, y=u1, mode="lines", name="Warped",
    line = attr(dash="dash"))

    layout = Layout(
        title = "Output traces for dynamic time warping",
        xaxis=attr(title="Time [s]"),
        yaxis=attr(title="Amplitude (a.u.)"),
    )
    plot([tr1, tr2], layout)
end

p = ([lineplot_s2(), lineplot_s3()])
display(p)
savefig(p, "./SINEtraces.png")

println("Type something to exit this example.");
readline();


module TestAsInterval

using Base.Test
using ArbFloats

afIvl1 = midpoint_radius(ArbFloat(1.0), ArbFloat(1.0e-8))
afIvl2 = midpoint_radius(ArbFloat(1.0), ArbFloat(1.0e-10))
afIvl3 = nextfloat(afIvl2)

@test contains(afIvl1, afIvl2)
@test !contains(afIvl2, afIvl1)
@test overlap(afIvl2, afIvl3)

@test midpoint(afIvl3) > midpoint(afIvl2)
@test radius(afIvl3) == radius(afIvl2)

end # module


facts("Single objective functions") do

context("Sphere") do
  p = BlackBoxOptim.example_problems["Sphere"]
  sphere = p.funcs[1]

  @fact sphere([0]) => 0

  @fact sphere([1]) => 1

  @fact sphere([1, 2]) => 5

  @fact sphere([1, 2, 3]) => 14

  @fact sphere([-1, 2, -3]) => 14

  @fact_throws sphere([])

  p2 = BlackBoxOptim.as_fixed_dim_problem(p, 3)
  @fact numdims(p2) => 3
  @fact ranges(search_space(p2)) => [(-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0)]
end

context("Schwefel2.22") do
  p = BlackBoxOptim.example_problems["Schwefel2.22"]
  schwefel2_22 = p.funcs[1]

  @fact schwefel2_22([0]) => 0

  @fact schwefel2_22([1]) => 2

  @fact schwefel2_22([1, 2]) => (1+2)+(1*2)

  @fact schwefel2_22([1, 2, 3]) => (1+2+3)+(1*2*3)

  @fact schwefel2_22([-1, 2, -3]) => (1+2+3)+(1*2*3)

  @fact_throws schwefel2_22([])

  p2 = BlackBoxOptim.as_fixed_dim_problem(p, 4)
  @fact numdims(p2) => 4
  @fact ranges(search_space(p2)) => [(-10.0, 10.0), (-10.0, 10.0), (-10.0, 10.0), (-10.0, 10.0)]
end

context("Schwefel1.2") do
  p = BlackBoxOptim.example_problems["Schwefel1.2"]
  schwefel1_2 = p.funcs[1]

  @fact schwefel1_2([0]) => 0

  @fact schwefel1_2([1]) => 1

  @fact schwefel1_2([1, 2]) => 1+9

  @fact schwefel1_2([1, 2, 3]) => 1+9+36

  @fact schwefel1_2([-1, 2, -3]) => 1+1+4

  @fact schwefel1_2([]) => 0
end

context("Schwefel2.21") do
  p = BlackBoxOptim.example_problems["Schwefel2.21"]
  schwefel2_21 = p.funcs[1]

  @fact schwefel2_21([0]) => 0

  @fact schwefel2_21([1]) => 1

  @fact schwefel2_21([1, 2]) => 2

  @fact schwefel2_21([1, 2, 3]) => 3

  @fact schwefel2_21([-1, 2, -3]) => 3

  @fact_throws schwefel2_21([])
end

context("Rosenbrock") do
  p = BlackBoxOptim.example_problems["Rosenbrock"]
  rosenbrock = p.funcs[1]

  @fact rosenbrock([1, 2]) => 100

  @fact rosenbrock([1, 2, 3]) => 201

  @fact rosenbrock([-1, 2, -3]) => 5005

  @fact_throws rosenbrock([])
end

end


function modes(values)
    dict = Dict() # Values => Number of repetitions
    modesArray = typeof(values[1])[] # Array of the modes so far
    max = 0 # Max of repetitions so far

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

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

    return modesArray
end

println(modes([1,3,6,6,6,6,7,7,12,12,17]))
println(modes((1,1,2,4,4)))


"""
    KT_residuals

Compute Kuhn-Turck residuals

# Arguments
- `df`: gradient
- `lambda`: Lagrangian multipliers with respect to the constraints
- `mult_x_U`: reduced cost with respect to the upper bounds
- `mult_x_L`: reduced cost with respect to the lower bounds
- `Jac`: Jacobian matrix
- `norm`: whether the residual is normalized or not
"""
function KT_residuals(
    df::Tv, lambda::Tv, mult_x_U::Tv, mult_x_L::Tv, Jac::Tm
) where {T, Tv<:AbstractArray{T}, Tm<:AbstractMatrix{T}}
    KT_res = norm(df - Jac' * lambda - mult_x_U - mult_x_L)
    scalar = max(1.0, norm(df))
    for i = 1:size(Jac,1)
        scalar = max(scalar, abs(lambda[i]) * norm(Jac[i,:]))
    end
    return KT_res / scalar
end

"""
    norm_complementarity

Compute the normalized complementeraity
"""
function norm_complementarity(
    E::Tv, g_L::Tv, g_U::Tv, x::Tv, x_L::Tv, x_U::Tv,
    lambda::Tv, mult_x_U::Tv, mult_x_L::Tv,
    p = Inf
) where {T, Tv <: AbstractArray{T}}
    m = length(E)
    compl = Tv(undef, m)
    denom = 0.0
    for i = 1:m
        if g_L[i] == g_U[i]
            compl[i] = 0.0
        else
            compl[i] = min(E[i] - g_L[i], g_U[i] - E[i]) * lambda[i]
            denom += lambda[i]^2
        end
    end
    return norm(compl, p) / (1 + sqrt(denom))
end

"""
    norm_violations

Compute the normalized constraint violation
"""
function norm_violations(
    E::Tv, g_L::Tv, g_U::Tv, x::Tv, x_L::Tv, x_U::Tv, p = 1
) where {T, Tv <: AbstractArray{T}}

    m = length(E)
    n = length(x)
    viol = Tv(undef, m+n)
    fill!(viol, 0.0)
    for i = 1:m
        if E[i] > g_U[i]
            viol[i] = E[i] - g_U[i]
        elseif E[i] < g_L[i]
            viol[i] = g_L[i] - E[i]
        end
    end
    for j = 1:n
        if x[j] > x_U[j]
            viol[m+j] = x[j] - x_U[j]
        elseif x[j] < x_L[j]
            viol[m+j] = x_L[j] - x[j]
        end
    end
    return norm(viol, p)
end


export DimensionDiagrams, EstimateDimension, EstimateDimensionMLE, EstimateDimensionANOVA, EstimateDimensionNPCA, EstimateDimensionPCA, EstimateDimensionCorrSum, EstimateDimensionBoxCounting

######################
# Dimension Diagrams #
######################
"""
DimensionDiagrams(
    data::Array{T,2},
    projective::Bool;
    methods  = [:CorrSum, :BoxCounting, :PHCurve, :NPCA, :MLE, :ANOVA, :PHCurve],
    eps_ticks = 25,
    fontsize = 16,
    lw = 4
    ) where {T <: Number}

Produces Dimension Diagrams.
* data is a matrix whose colums are the data points in Ω.
* projective = false makes diagrams in euclidean space.
* projective = true makes diagrams in projective space.
There are some optional arguments.
* methods lists the dimension estimators to be plotted.
* eps_ticks = k puts into k evenly spaces ϵ on [0,1] at which the dimension is computed.
* fontsize sets the fontsize of the axes.
* lw sets the linewidth.
"""
function DimensionDiagrams(
    data::Array{T,2},
    projective::Bool;
    diagrams  = [:CorrSum, :BoxCounting, :NPCA, :MLE, :ANOVA, :PHCurve],
    eps_ticks = 25,
    fontsize = 16,
    lw = 4
    ) where {T <: Number}

    n = size(data,1)

    cols = Colors.colormap("RdBu", mid = 0.5)
    colors = Dict("CorrSum" => cols[5], "BoxCounting"=> cols[20], "PHCurve"=> cols[35], "NPCA" => cols[55], "MLE" => cols[70], "ANOVA" => cols[85], "PHCurve" => cols[100])

    ϵ = Array{Float64}(range(0.1, length = eps_ticks, stop = 0.9))

    p = Plots.plot()
    dimension_estimates = map(diagrams) do m
        d = EstimateDimension(m, data, ϵ, projective)
        Plots.plot!(p, ϵ, d, label = string(m), lw = lw, linecolor = colors["$m"])

        d
    end

    if !projective
        y_upper = n + 1 + 0.1
    else
        y_upper = n + 0.1
    end
    y_upper = max(y_upper, maximum(vcat(dimension_estimates...)))

    Plots.plot(p,
              xlims = [0,1],
              ylims = [0,y_upper],
              titlefont = fontsize,
              tickfont = fontsize,
              labelfontsize = fontsize,
              xlabel = "epsilon",
              ylabel = "d(epsilon)",
              xscale = :none,
              yscale = :none
      )
end

#################
# Wrapper #
#################
function EstimateDimension(method::Symbol, data::Array{T,2}, ϵ::Vector{S}, projective::Bool) where {S,T<:Number}
    if method == :CorrSum
        return EstimateDimensionCorrSum(data, ϵ, projective)
    elseif method == :BoxCounting
        return EstimateDimensionBoxCounting(data, ϵ, projective)
    elseif method == :NPCA
        return EstimateDimensionNPCA(data, ϵ, projective)
    elseif method == :MLE
        return EstimateDimensionMLE(data, ϵ, projective)
    elseif method == :ANOVA
        return EstimateDimensionANOVA(data, ϵ, projective)
    elseif method == :PHCurve
        return EstimateDimensionPHCurve(data, ϵ, projective)
    else
        println("Method not supported")
        return zeros(length(ϵ))
    end
end

#################
# PCA #
#################
function EstimateDimensionPCA(data::Array{T,2}, projective::Bool) where {T <: Number}
    if !projective
        n = size(data, 1)
        data = data .- Statistics.mean.([data[i,:] for i in 1:n])
    else
        data = SafeDehomogenization(data)
        n = size(data, 1)
        data = data .- Statistics.mean.([data[i,:] for i in 1:n])
    end
    if size(data, 2) > 1
        L = log10.(LinearAlgebra.svdvals(data))
        if length(L) == 1
            if L[1] < 1e-15
                return 0
            else
                return 1
            end
        else
            i = findmin(diff(L))
            # return rank(data)
            return i[2]
        end
    else
        return 0
    end
end

#################
# Nonlinear PCA #
#################
function EstimateDimensionNPCA(data::Array{T,2}, ϵ_array::Vector{S}, projective::Bool) where {S,T<:Number}

    if !projective
        dist = ScaledEuclidean(data)
    else
        dist = ScaledFubiniStudy(data)
    end

    H = Clustering.hclust(dist, linkage = :single)
    m = size(data,2)
    return map(ϵ_array) do ϵ
        myclust = Clustering.cutree(H, h = ϵ)
        groups = unique(myclust)
        r = map(groups) do group
            index = findall(myclust .== group)
            return EstimateDimensionPCA(data[:,index], projective) * length(index)
        end
        return sum(r) / m
    end
end
function EstimateDimensionNPCA(data::Array{T,2}, projective::Bool; eps_ticks = 25) where {T<:Number}
    ϵ = Array{Float64}(range(0.1, length =eps_ticks, stop = 0.9))
    EstimateDimensionNPCA(data, ϵ, projective)
end

#########################
# Correlation Dimension #
#########################
function EstimateDimensionCorrSum(data::Array{T,2}, ϵ_array::Vector{S}, projective::Bool) where {S,T<:Number}
    n = size(data, 1)
    m = size(data, 2)
    if !projective
        dist = ScaledEuclidean(data)
    else
        dist = ScaledFubiniStudy(data)
    end

    output =  map(ϵ_array) do ϵ
        count = (sum(dist .<  ϵ) - m) / 2
        if count > 0
            return abs( log(count / (m * (m-1))))
        else
            return 0.0
        end
    end

    if !projective
        output = abs.(diff(output) ./ diff(log.(ϵ_array)))
        return [0.0; output]
    else
        output = abs.(diff(output) ./ diff(log.(sin.(ϵ_array))))
        return [0.0; output]
    end
end
function EstimateDimensionCorrSum(data::Array{T,2}, projective::Bool; eps_ticks = 25) where {T<:Number}
    ϵ = Array{Float64}(range(0.1, length = eps_ticks, stop = 0.9))
    EstimateDimensionCorrSum(data, ϵ, projective)
end

##########################
# Box Counting Dimension #
##########################
function EstimateDimensionBoxCounting(data::Array{T,2}, ϵ_array::Vector{S}, projective::Bool) where {S,T<:Number}
    n = size(data, 1)
    m = size(data, 2)

    if !projective
        u_min = minimum.([data[i,:] for i in 1:n])
        u_max = maximum.([data[i,:] for i in 1:n])

        d = abs.(u_max - u_min)
        λ = maximum(d) * √(n)

        return map(ϵ_array) do ϵ
            R = floor(Int64, λ/ϵ + 1)
            if R == 1
                return 0.0
            else
                qs = [floor.((abs.(data[:,i] - u_min) ./ d) .* R) for i in 1:m]
                return log(length(unique(qs)))/abs(log(R))
            end
        end
    else
        data = SafeDehomogenization(data)
        n = n-1
        u_min = minimum.([data[i,:] for i in 1:n])
        u_max = maximum.([data[i,:] for i in 1:n])

        i = 1
        #compute spherical angle for the box, not projective
        d = map(1:n) do i
            a = [u_min[i]; 1.0]
            b = [u_max[i]; 1.0]
            c = abs(LinearAlgebra.dot(a,b) / (LinearAlgebra.norm(a) * LinearAlgebra.norm(b)))
            if c < 1.0
                return acos(c)
            else
                return 0.0
            end
        end
        λ = maximum(d)
        return map(ϵ_array) do ϵ
            R = floor(Int64, λ/ϵ + 1)

            if R == 1
                return 0.0
            else
                D = map(1:m) do j
                    return map(1:n) do i
                        a = [u_min[i]; 1.0]
                        b = [data[i,j]; 1.0]
                        c = abs(LinearAlgebra.dot(a,b) / (LinearAlgebra.norm(a) * LinearAlgebra.norm(b)))
                        if c < 1.0
                            return acos(c)
                        else
                            return 0.0
                        end
                    end
                end
                qs = [floor.(( D[i] ./ d) .* R) for i in 1:m]
                return log(length(unique(qs)))/abs(log(R))
            end
        end
    end
end
function EstimateDimensionBoxCounting(data::Array{T,2}, projective::Bool; eps_ticks = 25) where {T<:Number}
    ϵ = Array{Float64}(range(0.1, length = 25, stop = 0.9))
    EstimateDimensionBoxCounting(data, ϵ, projective)
end

###########################
# Bickel-Levina-Algorithm #
###########################
function BL_MLE(dist::Array{T,1}, ϵ::Number, projective::Bool) where {T<:Number}
    k = length(dist)
    if !projective
        return float(abs(inv(sum(log.(ϵ .* inv.(dist))) / k)))

    else
        return float(abs(inv(sum(log.(sin(ϵ) .* inv.(sin.(dist)))) / k)))

    end
end
function EstimateDimensionMLE(data::Array{T,2}, ϵ_array::Vector{S}, projective::Bool) where {S,T<:Number}
    m = size(data, 2)
    if !projective
        dist = ScaledEuclidean(data)
    else
        dist = ScaledFubiniStudy(data)
    end
    estimators = map(ϵ_array) do ϵ
            estimator = zeros(Float64, m, 2)
            for i = 1:m
                dist_i = dist[i,[1:i-1;i+1:m]]

                if !projective
                    index = findall(dist_i .<  ϵ)
                else
                    index = findall(dist_i .<  ϵ)
                end

                k = length(index)
                if k > 1
                    estimator[i,:] = [BL_MLE(dist_i[index], ϵ, projective) * k, k]
                else
                    estimator[i,:] = [0.0, 0]
                end
            end
            return estimator
        end

    return map(estimators) do entry
        return sum(entry[:,1]) / sum(entry[:,2])
    end
end
function EstimateDimensionMLE(data::Array{T,2}, projective::Bool; eps_ticks = 25) where {T<:Number}
    ϵ = Array{Float64}(range(0.1, length = 25, stop = 0.9))
    EstimateDimensionMLE(data, ϵ, projective)
end

#################################
# Velasco-Quiroz-Diaz-Algorithm #
#################################
# Main function to compute the estimate of the dimension
# It is the algorithm from Sec. 2.1
function DQV_Estimator(dist::Array{T,2}) where {T<:Number}
    # Compute the U-statistic from Equation (3)
    m = size(dist,1)
    v = vec(dist) .- pi/2
    S = LinearAlgebra.dot(v, v)
    S /= (2 * binomial(m,2))

    # Now we have to compute the variances β from Lemma 2.5
    # Ininitialize the βs.
    # Treat d even and d odd separately
    β_even = pi^2 / 12
    β_odd = pi^2 / 4
    n_even = 0
    n_odd = 0

    # Compute the βs recursively. Since they are decreasing,
    # I only need to check for when U>β happens
    while S < β_even
        n_even = n_even + 1
        β_even = β_even - 2 / (2*n_even)^2
    end

    while S < β_odd
        n_odd = n_odd + 1
        β_odd = β_odd - 2 / (2*n_odd-1)^2
    end

    # Now I have four canditates for the closest β.
    index = findmin([abs(S-β_odd) abs(S-β_odd-2 / (2*n_odd-1)^2) abs(S-β_even) abs(S-β_even-2 / (2*n_even)^2) ])

    # return the corresponding dimension
    if index[2][2] == 1
        return 2 * n_odd + 1
    elseif index[2][2] == 2
        return maximum([2 * n_odd - 1,0])
    elseif index[2][2] == 3
        return 2 * n_even + 2
    elseif index[2][2] == 4
        return 2 * n_even
    end
end
function EstimateDimensionANOVA(data::Array{T,2}, ϵ_array::Vector{U}, projective) where {T,U <: Number}
    n = size(data, 1)
    m = size(data, 2)
    dist = zeros(T, m, m)
    D = zeros(Float64, m, m)

    if !projective
        dist = ScaledEuclidean(data)
    else
        data = hcat([LinearAlgebra.normalize(data[:,i]) for i in 1:m]...)
        dist = ScaledFubiniStudy(data)
    end

    estimators = map(1:m) do i
        P = zeros(T, n, n)
        if projective
            P = LinearAlgebra.diagm(0 => ones(n)) - data[:,i] * LinearAlgebra.transpose(data[:,i])
            data_x = (P*data)[:,[1:i-1;i+1:m]]
        else
            data_x = data[:,[1:i-1;i+1:m]] .- data[:,i]
        end
        cos_dist = FubiniStudyDistances!(D, data_x)
        return map(ϵ_array) do ϵ
            index = findall(dist[i,[1:i-1;i+1:m]] .< ϵ)
            k = length(index)
                if k > 1
                    return [DQV_Estimator(cos_dist[index,index]) * k, k]
                else
                    return [0, 0]
                end
        end
    end

    estimators = map(1:length(ϵ_array)) do i
        return [map(entry -> entry[i][1], estimators), map(entry -> entry[i][2], estimators)]
    end

    return map(1:length(ϵ_array)) do i
        sum(estimators[i][1]) / sum(estimators[i][2])
    end
end
function EstimateDimensionANOVA(data::Array{T,2}, projective::Bool; eps_ticks = 25) where {T<:Number}
    ϵ = Array{Float64}(range(0.1, length = 25, stop = 0.9))
    EstimateDimensionANOVA(data, ϵ, projective)
end

#################################
# PH curve dimension #
#################################
function EstimateDimensionPHCurve(dist::Array{T,2}) where {T <: Number}

    if sum(dist .> 0.0) > 0
        C = Eirene.eirene(dist, maxdim = 0)
        B = Eirene.barcode(C, dim = 0)
        m = size(B,1)

        if m > 1
            integral = sum([abs(B[i,2] - B[i,1]) for i in 1:(m-1)]) / (m-1)
            return abs(log(m) / log(integral))
        else
            return 0.0
        end
    else
        return 0.0
    end
end
function EstimateDimensionPHCurve(data::Array{T,2}, ϵ_array::Vector{S}, projective) where {T,S <: Number}
    if !projective
        dist = ScaledEuclidean(data)
    else
        dist = ScaledFubiniStudy(data)
    end

    H = Clustering.hclust(dist, :single)
    m = size(data,2)
    return map(ϵ_array) do ϵ

        myclust = Clustering.cutree(H, h = ϵ)
        groups = unique(myclust)
        r = map(groups) do group
            index = findall(myclust .== group)
            return EstimateDimensionPHCurve(dist[index,index]) * length(index)
        end
        return sum(r) / m
    end
end

function EstimateDimensionPHCurve(data::Array{T,2}, projective::Bool; eps_ticks = 25) where {T<:Number}
    ϵ = Array{Float64}(range(0.1, length =eps_ticks, stop = 0.9))
    EstimateDimensionPHCurve(data, ϵ, projective)
end


#=
Created on Thursday 16 Jan 2020
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Solution for the color transfer.
=#

using STMO
using STMO.OptimalTransport
using Plots, LaTeXStrings

using Images, Colors

# function to subsample image
subsample(image, every=8) = image[1:every:size(image,1), 1:every:size(image,2)]

# load images
image1 = load("chapters/06.OptimalTransport/Figures/butterfly3.jpg")
image2 = load("chapters/06.OptimalTransport/Figures/butterfly2.jpg")

# subsample
image1 = subsample(image1, 8)
image2 = subsample(image2, 5)

# get colors, can also using sampling
colors1 = vec(image1)
colors2 = vec(image2)

# distance matrix
C = [colordiff(c1, c2) for c1 in colors1, c2 in colors2]

n, m = size(C)

# optimal transport
P = sinkhorn(C, ones(n)/n, ones(m)/m, λ=10)

mapdistr(X, P) = Diagonal(sum(P, dims=2)[:].^-1) * P * X

# map colors
image1transf = reshape(mapdistr(colors2, P), size(image1)...)
image2transf = reshape(mapdistr(colors1, P'), size(image2)...)

plot(
plot(image1, title="image 1", aspect_ratio=:equal),
plot(image2, title="image 2", aspect_ratio=:equal),
plot(image1transf, title="image 1 transfered", aspect_ratio=:equal),
plot(image2transf, title="image 2 transfered", aspect_ratio=:equal)
)

savefig("figures/colortransfer")


# # Parallel evaluation and simple statistics
#
# In this example we will be looking at a way to evaluate several times a given problem
# by exploiting a distributed computing strategy known as **embarrasingly parallel**
# execution.
#
# The idea is very simple: compute the same problem several times, in parallel, and
# collect the results at the end.
# This should work well with metaheuristics due to the following properties:
#
# - Metaheuristics are _self-contained,_ meaning that they possess their own RNG and their own search space.
# - Due to the fact that solutions are searched with a random component, independent runs should be able to give independent results.
#
# With this in mind, we can use the fantastic parallel capabilities of `Julia` and
# show how easy it is to obtain some simple, but meaningful statistics for a given
# optimization problem.
#
# With this we want to also show that this is the **right way** to use this package.
# We should _always_ run several independent solvers to obtain good statistics.
# We must **never** settle with just one run from the solver.
#
# ## Setting up
#
# In order to use the multithreading capabilities of `Julia`, we need to launch the
# REPL using all the threads we can afford.
#
# For instance, the computer where this example is currently being executed has
# 8 physical cores, so I should launch my `Julia` REPL
# like so
#
# ```shell
# JULIA_NUM_THREADS=8 julia
# ```
#
# This will launch the REPL with 8 cores ready to be used. We can check this by entering
# the following in the REPL,
#
# ```julia
# julia> using Base.Threads
# julia> nthreads()
# ```
#
# making sure that this matches the same number introduced with the variable
# `JULIA_NUM_THREADS`.
#
# You should replace the number of threads, in this case 8, with the number of threads
# that you wish.
#
# ## Set up
#
# First, let us import all the necessary modules we will need for this example.

using BioMetaheuristics
using BioMetaheuristics.TestFunctions
using Base.Threads
using Statistics
using BenchmarkTools

# Now, let us create some setup functions.
# First, we will be using the `optimize` interface to solve the [`Levy`](@ref)
# optimization problem. This function just
# wraps the `optimize` function and returns the solution vector from the problem.
#
# We will be using Particle Swarm Optimization (PSO), with 20,000 iterations.
# It is very important to note that no specific RNG is being used.
# This tells the `optimize` function to always use a new, truly randomly seeded RNG.

function solve_problem(dim)
    val = optimize(Levy(), zeros(dim), [-35.0, 35.0], PSO(); iters=20_000)

    return val.x
end;

# ### Some benchmarks
#
# Before we continue, I would like to show a benchmark of the PSO implementation.
# We will use the fantastic [`BenchmarkTools.jl`](https://github.com/JuliaCI/BenchmarkTools.jl)
# to check the time it takes to finish a single optimization problem.
#
# Here, we will solve a 50-dimensional Levy optimization problem. We are not really
# interested if the solution is accurate, just the timing. We will deal with accuracy
# in the next subsection.

b = @benchmark solve_problem(50);
b

# Not bad at all. If it takes between 3-4 seconds to finish I think this is fairly
# good.
# Furthermore, it means that we can take advantage of how little it takes to run
# in order to really push the number of solvers to execute.
#
# ## Parallel evaluation
#
# We now turn to the main event. We will launch several solvers using the previously
# created wrapper function in parallel. First off, we need another function to do so.
#
# This function will use the `@threads` macro that enables a for-loop to run in parallel.
# The function has a matrix called `solutions` where we will be storing the solution
# vectors from each independent run.
#
# We will call the solver in parallel and each time we will save the solutions vector.
# We will then compute the mean, median and standard deviation from all the independent
# runs and report that as the _true_ value, and compare it to the _ground truth._

function parallel_compute(;dim=50, total_iterations=50)
    solutions = Matrix{Float64}(undef, dim, total_iterations)

    @threads for i in 1:total_iterations
        solutions[:, i] = solve_problem(dim)
    end

    mean_result = mean(solutions; dims=2)
    median_result = median(solutions; dims=2)
    std_result = std(solutions; dims=2)

    return (mean_result, median_result, std_result)
end;

# We are ready to go. We will do 200 iterations in order to have good statistics.
#
# !!! warning
#     Depending on your hardware this might take a long time. If that is the case
#     stop the process and reduce the number of iterations.
#
# !!! note
#     Every time you do this exercise, new results will appear. This is normal for
#     metaheuristics, and it is desired. We don't want reproducibility here, we aim
#     for _statistical significance._

μ, med, σ = parallel_compute(;total_iterations=200);

# We can now check the results

# Here is the mean value

μ'

# Here is the median obtained

med'

# And finally, the standard deviation of the solution

σ'

# Let us print the evaluation of the function on these results.
# First, the mean value.

println(evaluate(Levy(), μ))

# Then, the median obtained.

println(evaluate(Levy(), med))

# ## Conclusions
#
# Recall that for the Levy function the solution is
#
# ```math
# \mathbf{x^*} = (1, \dots, 1)
# ```
#
# and it should evaluate to
#
# ```math
# f(\mathbf{x^*}) = 0
# ```
#
# It it safe to say that we have obtained those results, but we are _statistically sure_
# that these results are the _true_ results, because we have several independent runs
# for the answer.
#
# ### On the use of the median
#
# Notice that most of the time, the median result actually performs better than the mean
# result.
# This is somewhat expected, as the median is not really affected by poorly performant runs.
# Sometimes, independent runs will not give good results and when taken into account,
# the mean value will suffer greatly from this.
# The median, however, as a measure of the "value in the middle", will not be so affected
# and a better estimation is expected.
#
# We should always report both results, and conclude that whenever the final function
# evaluation is approximately the same with both the mean and the median, the
# solution is statistically acceptable for that particular problem.
#


using StatsBase
function dstats(x, rnames="")
    k = size(x,2)
    if rnames==""
        rnames = 1:k
        rnames = rnames'
    end
    m = mean(x,1)
    s = std(x,1)
    sk = m-m
    k = m-m
    for i = 1:size(x,2)
        sk[:,i] = skewness(x[:,i])
        k[:,i] = kurtosis(x[:,i])
    end
    mn = minimum(x,1)
    mx = maximum(x,1)
    q05 = zeros(k)
    q95 = zeros(k)
    for i = 1:size(x,2) q05[i] = quantile(x[:,i], 0.05) end
    for i = 1:size(x,2) q95[i] = quantile(x[:,i], 0.95) end
    cnames = ["  mean", "  std", "skew", "kurt", "min", "max", "q05", "q95"]
    stats = [m' s' sk' k' mn' mx' q05' q95']
    prettyprint(stats, cnames, rnames);
    return stats
end


module problem108

description = """
Diophantine reciprocals I
Problem 108

In the following equation x, y, and n are positive integers.
1/x + 1/y = 1/n

For n = 4 there are exactly three distinct solutions:
1/5 + 1/20 = 1/4
1/6 + 1/12 = 1/4
1/8 + 1/8 = 1/4

What is the least value of n for which the number of distinct solutions exceeds one-thousand?

NOTE: This problem is an easier version of problem 110; it is strongly advised that you solve this one first.
"""

# Brute force check for 1/y + 1/x = 1/n for x in the range n+1 to 2n
# Integer y exists when (x-n) divides x*n, so no need to actually calulate the fractions.
function solutionCount(n :: Int)
  c = 0
  xn = n*(n+1)
  xmn = 1
  for x = n+1 : n * 2
    if xn % xmn == 0
      c += 1
    end
    xn += n
    xmn += 1
  end
  return c
end

using Base.Test

@test solutionCount(4) == 3

# all solutions take the form:
# a/(n*k) + (k-a)/(n*k) = 1/n,
# where a is a factor of n*k, and (k-a) is a factor of n*k
# I arrived at this solution by trial and error playing with the solutionCount above.
# Originally, I had only considered the case where a=1, reducing the problem to N having >1000 divisors.
# The solution to that problem requires a large value for Sigma0(n) which implies any prime factors.
@test solutionCount(2*2*3*3*5*7*11*13) > 1000

end


abstract AbstractPoint{T <: Real}

immutable Point{T} <: AbstractPoint
	x::T
	y::T
end

x(p::Point) = p.x
y(p::Point) = p.y
Base.eltype{T}(::Type{Point{T}}) = T

Base.start(p::Point) = 1
Base.next(p::Point, state) = (getfield(p,fieldname(Point,state)), state + 1)
Base.done(p::Point, state) = state > nfields(Point)
Base.length(p::Point) = nfields(Point)
Base.size(p::Point) = (nfields(Point),)

Base.getindex(p::Point,i) = getfield(p,fieldname(Point,i))
Base.endof(p::Point)  = length(p)

immutable Line{T <: AbstractPoint}
	p1::T
	p2::T
end

function perpendicular(l1::Line, l2::Line)
	m1 = (y(l1.p2) - y(l1.p1)) / (x(l1.p2) - x(l1.p1))
	m2 = (y(l2.p2) - y(l2.p1)) / (x(l2.p2) - x(l2.p1))

	return isapprox(m1,-(1/m2))
end

⟂(l1,l2) = perpendicular(l1,l2)

isleft(p::AbstractPoint, l::Line) = ((x(l.p2) - x(l.p1)) * (y(p) - y(l.p1) ) - (x(p) - x(l.p1)) * (y(l.p2) - y(l.p1)))

Base.norm(l::Line) = norm([x(l.p2) - x(l.p1), y(l.p2) - y(l.p1)])

abstract AbstractPolygon{T <: AbstractPoint}

immutable Polygon{T} <: AbstractPolygon{T}
	vertexs::Vector{T}
	function Polygon(v::Vector)
		if v[end] != v[1]
			push!(v,v[1])
		end
		return new(v)
	end
end

Polygon{T}(v::Vector{T}) = Polygon{T}(v)

Base.start(p::Polygon) = start(p.vertexs)
Base.next(p::Polygon, state) = next(p.vertexs, state)
Base.done(p::Polygon, state) = done(p.vertexs, state)
Base.length(p::Polygon) = length(p.vertexs)
Base.size(p::Polygon) = size(p.vertexs)

Base.getindex(p::Polygon,i) = p.vertexs[i]
Base.endof(p::Polygon)  = endof(p.vertexs)
Base.eltype(p::Polygon) = eltype(p.vertexs)

function boundingbox(poly::AbstractPolygon)
	minx = mapreduce(x,min,poly.vertexs)
	miny = mapreduce(y,min,poly.vertexs)
	maxx = mapreduce(x,max,poly.vertexs)
	maxy = mapreduce(y,max,poly.vertexs)

	(Point(minx,miny), Point(maxx,maxy))
end

function Base.in(p::AbstractPoint, poly::AbstractPolygon)
	wn = 0
	for idx in 1:(length(poly) - 1)
		V1 = poly[idx]
		V2 = poly[idx + 1]

		if y(V1) <= y(p)
			if y(V2) > y(p)
				if isleft(p, Line(V1,V2)) > 0
					wn+=1
				end
			end
		else
			if y(V2) <= y(p)
				if isleft(p, Line(V1,V2)) < 0
					wn-=1
				end
			end
		end
	end
	return wn != 0
end

function Base.rand{T}(rng, S::AbstractPolygon{T}, n::Int64)
	pT = eltype(T)
	ret = Point{pT}[]

	(minp,maxp) = boundingbox(S)

	bw = x(maxp) - x(minp)
	bh = y(maxp) - y(minp)

	while length(ret) < n
		rnd = [x(minp),y(minp)] .+ (rand(rng,pT,2) .* [bw,bh])

		newPoint = Point{pT}(rnd...)

		if newPoint in S
			push!(ret,newPoint)
		end
	end
	ret
end

Base.rand(rng,S::AbstractPolygon) = rand(rng,S,1)

Base.rand(S::AbstractPolygon, n) = rand(Base.Random.GLOBAL_RNG,S,n)

Base.rand(S::AbstractPolygon) = rand(Base.Random.GLOBAL_RNG,S,1)[1]


using Markdown
using InteractiveUtils
using ReinforcementLearning
using Plots
using Flux
using Statistics

begin
	const NX = 4
	const NY = 12
	const Start = CartesianIndex(1, 1)
	const Goal = CartesianIndex(1, 12)
	const UDZ = [
        CartesianIndex(-1, 0),  # up
        CartesianIndex(1, 0),   # down
        CartesianIndex(0, 0),  # zero
	]
	const LinearInds = LinearIndices((NX, NY))
end

function iscliff(p::CartesianIndex{2})
    x, y = Tuple(p)
    x == 4 && y > 1 && y < NY
end

heatmap((!iscliff).(CartesianIndices((NX, NY))); yflip = true)

begin
	Base.@kwdef mutable struct CliffWalkingEnv <: AbstractEnv
		position::CartesianIndex{2} = Start
	end
	function (env::CliffWalkingEnv)(a::Int)
		x, y = Tuple(env.position + UDZ[a])
		env.position = CartesianIndex(min(max(x, 1), NX), min(max(y, 1), NY))
	end
end

RLBase.state(env::CliffWalkingEnv) = LinearInds[env.position]

RLBase.state_space(env::CliffWalkingEnv) = Base.OneTo(length(LinearInds))

RLBase.action_space(env::CliffWalkingEnv) = Base.OneTo(length(UDZ))

RLBase.reward(env::CliffWalkingEnv) = env.position == Goal ? 0.0 : (iscliff(env.position) ? -100.0 : -1.0)

RLBase.is_terminated(env::CliffWalkingEnv) = env.position == Goal || iscliff(env.position)

RLBase.reset!(env::CliffWalkingEnv) = env.position = Start

world = CliffWalkingEnv()

begin
	NS = length(state_space(world))
	NA = length(action_space(world))
end

create_agent(α, method) = Agent(
	policy = QBasedPolicy(
		learner=TDLearner(
			approximator=TabularQApproximator(
				;n_state=NS,
				n_action=NA,
				opt=Descent(α),
			),
			method=method,
			γ=1.0,
			n=0
		),
		explorer=EpsilonGreedyExplorer(0.1)
	),
	trajectory=VectorSARTTrajectory()
)

function repeated_run(α, method, N, n_episode, is_mean=true)
	env = CliffWalkingEnv()
	rewards = []
	for _ in 1:N
		h = TotalRewardPerEpisode(;is_display_on_exit=false)
		run(
			create_agent(α, method),
			env,
			StopAfterEpisode(n_episode;is_show_progress=false),
			h
		)
		push!(rewards, is_mean ? mean(h.rewards) : h.rewards)
	end
	mean(rewards)
end

α = 0.001
N = 500
n_episode = 500

begin
	p = plot(legend=:bottomright, xlabel="Episodes", ylabel="Sum of rewards during episode")
	plot!(p, repeated_run(α, :SARS, N, n_episode, false), label="QLearning")
	plot!(p, repeated_run(α, :SARSA, N, n_episode, false), label="SARSA")
	p
end

# begin
# 	A = 0.1:0.05:0.95
# 	fig_9_4 = plot(;legend=:bottomright, xlabel="α", ylabel="Sum of rewards per episode")

# 	plot!(fig_9_4, A, [repeated_run(α, :SARS, 100, 100) for α in A], linestyle=:dash ,markershape=:rect, label="Interim Q")
# 	plot!(fig_9_4, A, [repeated_run(α, :SARSA, 100, 100) for α in A], linestyle=:dash, markershape=:dtriangle, label="Interim SARSA")
# 	plot!(fig_9_4, A, [repeated_run(α, :ExpectedSARSA, 100, 100) for α in A], linestyle=:dash, markershape=:cross, label="Interim ExpectedSARSA")

# 	plot!(fig_9_4, A, [repeated_run(α, :SARS, 10, 5000) for α in A], linestyle=:solid ,markershape=:rect, label="Asymptotic interim Q")
# 	plot!(fig_9_4, A, [repeated_run(α, :SARSA, 10, 5000) for α in A], linestyle=:solid, markershape=:dtriangle, label="Asymptotic SARSA")
# 	plot!(fig_9_4, A, [repeated_run(α, :ExpectedSARSA, 10, 5000) for α in A], linestyle=:solid, markershape=:cross, label="Asymptotic ExpectedSARSA")
# 	fig_9_4
# end


"""
```julia
    coldstartdelta( G::AbstractMatrix{Float64}, g::AbstractVector{Float64}; H::AbstractMatrix{Float64} = zeros( Float64, (1,1) ), L::Float64 = eigmax(H) )
```

Compute the Δ parameter needed when computing the upper iteration bound for a cold-started
fast gradient method.

Either the Hessian `H` or its maximum eigenvalue `L` must be passed as a parameter to the function.
"""
function coldstartdelta( G::AbstractMatrix{Float64}, g::AbstractVector{Float64}; H::AbstractMatrix{Float64} = zeros( Float64, (1,1) ), L::Float64 = eigmax(H) )
    (n, m) = size( G )

    model = OSQP.Model()

    options = Dict( :verbose => false )

    OSQP.setup!( model; P = sparse( triu( 1.0*I(m) ) ),
                        q = zeros( Float64, (m) ),
                        A = sparse( G ),
                        u = Vector( g ),
                        options... )

    results = OSQP.solve!( model )

    r = norm( results.x )
    Δ  = L/2 * r^2

    return Δ
end

"""
```julia
    upperiterationbound( ϵ::Float64, G::AbstractMatrix{Float64}, g::AbstractVector{Float64}; H::AbstractMatrix{Float64} = zeros( Float64, (1,1) ), L::Float64 = eigmax(H), κ::Float64 = cond(H) )
```

Compute the maximum number of iterations needed for the fast gradient method to reach `ϵ` sub-optimality level when projecting
into the constraint set ``G x ≦ g``.

To compute the iteration bound, either the Hessian `H` must be specified, or its largest eigenvalue `L`
and condition number `κ` must be supplied instead.
"""
function upperiterationbound( ϵ::Float64, G::AbstractMatrix{Float64}, g::AbstractVector{Float64}; H::AbstractMatrix{Float64} = zeros( Float64, (1,1) ), L::Float64 = eigmax(H), κ::Float64 = cond(H) )
    Δ = coldstartdelta( G, g, L = L )

    a₁ = ceil( ( log(ϵ) - log(Δ) ) / ( log( 1- sqrt(1/κ) ) ) )
    a₂ = ceil( 2 * sqrt(Δ/ϵ) - 2 )

    return max( 0, min( a₁, a₂ ) )
end


using SchumakerSpline
using Base.Test

# Run tests

tic()
println("Second Derivative Test")
@time @test include("SecondDerivativeTest.jl")
toc()


using WSMP
using Base.Test

const ia = Int32[1,5,9,13,17,21,25,27,29,30];
const ja = Int32[1,3,7,8,2,3,8,9,3,7,8,9,4,6,7,8,5,6,8,9,6,7,8,9,7,8,8,9,9];
const avals = [14.,-1.,-1.,-3.,14.,-1.,-3.,-1.,16.,-2.,-4.,-2.,14.,-1.,-1.,
               -3.,14.,-1.,-3.,-1.,16.,-2.,-4.,-2.,16.,-4.,71.,-4.,16.];
w = Wssmp(ia,ja,avals)
x = w\ones(9)
@test_approx_eq x [0.10094824180165943,0.10094824180165943,0.11951797708415644,
                   0.10094824180165943,0.10094824180165943,0.11951797708415644,
                   0.11951797708415646,0.058079810351639664,0.11951797708415648]
V,R,M = wsversion()
@test V == 14


function each_cons(n, collection)
    [collection[i:(i + n - 1)] for i = 1:(length(collection) - n + 1)]
end

n = ["73167176531330624919225119674426574742355349194934",
     "96983520312774506326239578318016984801869478851843",
     "85861560789112949495459501737958331952853208805511",
     "12540698747158523863050715693290963295227443043557",
     "66896648950445244523161731856403098711121722383113",
     "62229893423380308135336276614282806444486645238749",
     "30358907296290491560440772390713810515859307960866",
     "70172427121883998797908792274921901699720888093776",
     "65727333001053367881220235421809751254540594752243",
     "52584907711670556013604839586446706324415722155397",
     "53697817977846174064955149290862569321978468622482",
     "83972241375657056057490261407972968652414535100474",
     "82166370484403199890008895243450658541227588666881",
     "16427171479924442928230863465674813919123162824586",
     "17866458359124566529476545682848912883142607690042",
     "24219022671055626321111109370544217506941658960408",
     "07198403850962455444362981230987879927244284909188",
     "84580156166097919133875499200524063689912560717606",
     "05886116467109405077541002256983155200055935729725",
     "71636269561882670428252483600823257530420752963450"]

println(maximum(map(digits -> prod(map(digit -> parse(Int, digit), split(digits, ""))), each_cons(5, reduce(*, n)))))


@testset "monomial generation" begin
    @polyvar x y z

    @test TypedPolynomials.monomial_powers(Val{3}(), 3) == [
        (3, 0, 0),
        (2, 1, 0),
        (2, 0, 1),
        (1, 2, 0),
        (1, 1, 1),
        (1, 0, 2),
        (0, 3, 0),
        (0, 2, 1),
        (0, 1, 2),
        (0, 0, 3)
    ]
    @test @inferred(monomials((x, y), 2)) == [x^2, x * y, y^2]
    @test @inferred(monomials((x, y), 0:2)) == [x^2, x * y, y^2, x, y, 1]
    @test @inferred(monomials((x, y), [1, 0, 2])) == [x^2, x * y, y^2, x, y, 1]

    Random.seed!(1)
    for vars in ((x,), (x, y), (x, y, z), (x, z))
        for degree in 0:5
            @test issorted(@inferred(monomials(vars, degree)), rev=true)
            @test issorted(@inferred(monomials(vars, 0:degree)), rev=true)
        end
    end
    @test_throws ArgumentError monomials((z, y), 2)
    @test_throws ArgumentError monomials((z, y), [1, 2])
end


function sift!(Eav, decomp, d_osf, nsifts=5)
    N = length(decomp)
    e1 = zeros(N)
    e2 = zeros(N)
    avg = zeros(N)
    w = max(div(d_osf-1, 2), 3)
    if iseven(w)
        w += 1
    end
    for j in 1:nsifts
        stream_minmax(e1, e2, decomp, d_osf)
        e1 .= moving_average(e1, d_osf)
        e2 .= moving_average(e2, d_osf)
        avg .= (e1 .+ e2)./2
        avg .= moving_average(avg, d_osf)
        Eav .= Eav .+ avg
        decomp .-= avg
    end
    return nothing
end

function sEMD(signal;maximfs=10,nsifts=2)
    N = length(signal)
    a_new = copy(signal)
    imfs = ElasticArray{Float64}(undef,N,0)
    Eav = zeros(N)
    for i in 1:maximfs
        exs = find_extrema_count(a_new)
        d_osf =  div(N, div(exs,2))
        iseven(d_osf) && (d_osf+=1)

        if d_osf >= div(N,3) || exs < 5
            append!(imfs, a_new)
            break
        end
        fill!(Eav,0)
        sift!(Eav, a_new, d_osf, nsifts)
        append!(imfs, a_new)
        a_new .= Eav
    end
        return convert(Array,imfs)
end


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5856fdbe-c5e9-11eb-33d0-2d17b97b572c
begin
	using CSV
	using DataFrames
	using GLM
	using PlutoUI
	using Plots
end

# ╔═╡ b3b94edc-fd96-4403-bdc7-a1f5af651d6c
data = CSV.read("D:\\Workspace\\Machine Learning\\Indian GDP Analysis and Prediction\\Data\\Indian_gdp.csv",DataFrame)

# ╔═╡ e1c7d649-9922-4897-998f-f401ecb88158
scatter(data.year,data.GDP)

# ╔═╡ b2c5ec62-314f-4cd4-a7bc-b6e353c64398
reg = fit(LinearModel, @formula(GDP ~ year), data)

# ╔═╡ 97eb7071-cf68-4801-918d-53f9157a9f44
linear_regression = predict(reg)

# ╔═╡ a5de798c-aae2-452e-930b-6716c739b108
begin
	plot(data.year,linear_regression)
	scatter!(data.year,data.GDP)
end

# ╔═╡ 2fa71a79-93fa-43a8-9ee7-12fddeb5f971
md"Enter an year to predict GDP : $@bind y TextField()"

# ╔═╡ 6b671101-053c-4136-b3b7-37674655cccf
pyear = parse(Int64,y)

# ╔═╡ 22fa2415-d9d7-4c91-a235-0c94340d7f51
pGDP = coef(reg)[2]*pyear+coef(reg)[1]

# ╔═╡ Cell order:
# ╠═5856fdbe-c5e9-11eb-33d0-2d17b97b572c
# ╠═b3b94edc-fd96-4403-bdc7-a1f5af651d6c
# ╠═e1c7d649-9922-4897-998f-f401ecb88158
# ╠═b2c5ec62-314f-4cd4-a7bc-b6e353c64398
# ╠═97eb7071-cf68-4801-918d-53f9157a9f44
# ╠═a5de798c-aae2-452e-930b-6716c739b108
# ╠═2fa71a79-93fa-43a8-9ee7-12fddeb5f971
# ╠═6b671101-053c-4136-b3b7-37674655cccf
# ╠═22fa2415-d9d7-4c91-a235-0c94340d7f51

Indian GDP Analysis and Prediction/main.jl


@testset "SSOR" begin
	@testset "small full system" begin
		# CG: test sparse Laplacian
		A = getDivGrad(32,32,32)
		rhs = randn(size(A,1))
		tol = 1e-2

		# tests with A being matrix
		x1,flag1,relres1,iter1,resvec1       = ssor(A,rhs,tol=tol,maxIter=100)
		x2,flag2,relres2,iter2,resvec2       = ssor(A,rhs,tol=tol,maxIter=100,omega=2/3,out=1)
		x3,flag3,relres3,iter3,resvec3     = ssor(A,rhs,tol=tol,maxIter=100,omega=1/2,out=1,storeInterm=true)

		# test relative residuals
		@test norm(A*x1-rhs)/norm(rhs) <= tol
		@test norm(A*x2-rhs)/norm(rhs) <= tol
		@test norm(A*x3[:,end]-rhs)/norm(rhs) <= tol
	end

	@testset "preconditioner" begin
		omega = 1.2;
		A = getDivGrad(32,32,32)
		rhs = randn(size(A,1))
		d = omega./Array(diag(A));
		x = zeros(length(rhs)) # pre allocation for the preconditioner result.
		PC = r -> (x[:].=0.0; return ssorPrecTrans!(A,x,r,d));
		y = KrylovMethods.cg(A,rhs,tol=1e-12,maxIter=200,M=PC,out=1)[1]
		@test norm(A*y-rhs)/norm(rhs) <= 1e-12
	end

	@testset "complex " begin

		# CG: test sparse Laplacian
		A = getDivGrad(32,32,32)
		A += 1im*I
		rhs = randn(size(A,1)) + 1im * randn(size(A,1))
		tol = 1e-2

		#
		x0,flag0,relres0,iter0,resvec0       = ssor(A,rhs,tol=tol,maxIter=3,out=2)
		@test flag0==-1
		@test iter0==3
		@test length(resvec0)==3

		# tests with A being matrix
		x1,flag1,relres1,iter1,resvec1       = ssor(A,rhs,tol=tol,maxIter=100)
		x2,flag2,relres2,iter2,resvec2       = ssor(A,rhs,tol=tol,maxIter=100,omega=2/3,out=1)
		x3,flag3,relres3,iter3,resvec3       = ssor(A,rhs,tol=tol,maxIter=100,omega=1/2,out=1)

		# test relative residuals
		@test norm(A*x1-rhs)/norm(rhs) <= tol
		@test norm(A*x2-rhs)/norm(rhs) <= tol
		@test all(diff(resvec1).<0)
		@test all(diff(resvec2).<0)
		@test all(diff(resvec3).<0)

	end
end


using JuMP
using MathOptInterface
const MOI = MathOptInterface

using Gurobi
model = Model(Gurobi.Optimizer)
# set_optimizer_attribute(model, "Cuts",0)
set_optimizer_attribute(model, "TimeLimit", 600)

# using Cbc
# model = Model(Cbc.Optimizer)
# set_optimizer_attribute(model, "cuts","off")
# set_optimizer_attribute(model, "seconds", 60)
# set_optimizer_attribute(model, "thread", 16)

n_i = 12                # 产品种类个数
n_week = 15                # 时间周期数
n_day = 7
n_machine = 3           # 机器个数
machine = 1:n_machine   # 机器编号
product = 1:n_i         # 产品编号
weeks = 1:n_week           # 时间编号
days = 1:n_day          # 星期编号
paras = [   2 83 10 1 30;
            2 31 10 1 20;
            2 11 10 1 30;
            2 93 10 1 40;
            2 82 10 1 40;
            2 72 10 1 10;
            3 23 20 1 30;
            3 91 20 1 20;
            3 83 20 1 10;
            3 34 20 1 50;
            3 61 20 1 30;
            3 82 20 1 20]

f = paras[:,1] # 产品包装使用的机器编号、产品类别

f2 = product[f .== 2] # 谷物类产品
f3 = product[f .== 3] # 水果类产品

s_init = paras[:,2] # 初始库存
ss = paras[:, 3] # 安全库存
alpha = repeat(paras[:, 4], 1, n_machine) # 生产速度
beta = paras[:, 5] # 清洗时间
# 需求
d  = [   0 95 110 96 86 124 83 108 114 121 110 124 104 86 87;
        98 96 96 98 103 104 122 101 89 108 101 109 106 108 76;
        106 0 89 123 96 105 83 82 112 109 119 85 99 80 123;
        98 121 0 105 98 96 101 81 117 76 103 81 95 105 102;
        0 124 113 123 123 79 111 98 97 80 98 124 78 108 109;
        103 102 0 95 107 105 107 105 75 93 115 113 111 105 85;
        110 93 0 112 84 124 98 101 83 87 105 118 115 106 78;
        85 92 101 110 93 96 120 109 121 87 92 85 91 93 109;
        122 116 109 0 105 108 88 98 77 90 110 102 107 99 96;
        120 124 94 105 92 86 101 106 75 109 83 95 79 108 100;
        117 96 78 0 108 87 114 107 110 94 104 101 108 110 80;
        125 112 75 0 116 103 122 88 85 84 76 102 84 88 82]
L = [1400, 700, 700] # 机器上限

@variable(model, x[i in product, t in weeks, day in days] >= 0) # 产品生产计划
@variable(model, s[i in product, t in weeks] >= 0) # 库存
@variable(model, y[i in product, t in weeks,day in days], Bin) # 是否开工生产

# 库存平衡
@constraint(model ,dem_sat1[i in product], s_init[i] + sum(x[i, 1, day] for day in days) == d[i, 1] + s[i, 1])
@constraint(model, dem_sat[i in product, t in weeks[2:end]], s[i, t - 1] + sum(x[i, t, day] for day in days) == d[i, t] + s[i, t] )
# 安全库存
@constraint(model, ss_limit[i in product, t in weeks], s[i, t] >= ss[i])

# 最大产量约束
@constraint(model, vub[i in product, t in weeks, day in days], x[i, t, day] <= (sum(d[i, k]  for k in t:n_week) + ss[i]) * y[i, t, day])

# 混合机器上限
@constraint(model, mix_cap[i in product, t in weeks, day in days], alpha[i, 1] * x[i, t, day]  + beta[i] * y[i, t, day]  <= L[1] / 7)

# 包装机器上限
@constraint(model, pack_cap2[i in f2, t in weeks, day in days], alpha[i, 2] * x[i, t, day]  <= L[2] / 7)
@constraint(model, pack_cap3[i in f3, t in weeks, day in days], alpha[i, 3] * x[i, t, day]  <= L[3] / 7)

# 包装机器只能开启一台
@constraint(model, pcak_work2[t in weeks, day in days], sum(y[i, t, day] for i in f2) <= 1)
@constraint(model, pcak_work3[t in weeks, day in days], sum(y[i, t, day] for i in f3) <= 1)
@expression(model, inventory, sum(s[i, t] for i in product, t in weeks))
@objective(model, Min, inventory)

optimize!(model)


#-----------------------------------------------------------------------------# Losses
l2regloss(y, yhat) = .5 * abs2(y - yhat)
deriv(::typeof(l2regloss), y, yhat) = yhat - y

l1regloss(y, yhat) = abs(y - yhat)
deriv(::typeof(l1regloss), y, yhat) = sign(yhat - y)

logisticloss(y, yhat) = log1p(exp(-y * yhat))
deriv(::typeof(logisticloss), y, yhat) = -y / (1 + exp(y * yhat))

l1hingeloss(y, yhat) = max(1 - y*yhat, zero(yhat))
deriv(::typeof(l1hingeloss), y, yhat) = (u = y*yhat; u ≥ 1 ? zero(y) : -y)

struct DWDLoss{T<:Number}
    q::T
end
function (o::DWDLoss)(y, yhat)
    agreement = y * yhat
    q = o.q
    if agreement ≤ q / (q + 1)
        1 - agreement
    else
        (q ^ q / (q + 1) ^ (q + 1)) / agreement ^ q
    end
end
function deriv(loss::DWDLoss, y, yhat)
    u = y * yhat
    q = loss.q
    if u ≤ q / (q + 1)
        -y
    else
       - y * ( q / (q + 1)) ^ (q + 1) / u ^ (q + 1)
    end
end

#-----------------------------------------------------------------------------# Penalties
prox(::typeof(zero), x, s) = x
prox(::typeof(abs2), x, s) = x / (1 + s)
prox(::typeof(abs), x::T, s::Number) where {T} = sign(x) * max(T(0), abs(x) - s)

struct ElasticNet{T}
    α::T
end
(p::ElasticNet)(x) = smooth(abs2(x), abs(x), p.α)
function prox(p::ElasticNet, x, s)
    αs = p.α * s
    prox(abs, x, αs) / (1 + s - αs)
end

#-----------------------------------------------------------------------------# StatLearn
"""
    StatLearn(args...; penalty=zero, rate=LearningRate())

Fit a model (via stochastic approximation) that is linear in the parameters.  The (offline)
objective function that StatLearn approximately minimizes is

``(1/n) ∑ᵢ f(yᵢ, xᵢ'β) + ∑ⱼ λⱼ g(βⱼ),``

where ``fᵢ`` are loss functions of a response variable and linear predictor, ``λⱼ``s are
nonnegative regularization parameters, and ``g`` is a penalty function.

Use `StatLearn` with caution, as stochastic approximation algorithms are inherently noisy.

# Arguments

- `loss = OnlineStats.l2regloss`: The loss function to be (approximately) minimized.
    - Regression Losses:
        - `l2regloss`: Squared error loss
        - `l1regloss`: Absolute error loss
    - Classification (y ∈ {-1, 1}) Losses:
        - `logisticloss`: Logistic regression
        - `l1hingeloss`: Loss function used in Support Vector Machines.
        - `DWDLoss(q)`: Generalized Distance Weighted Discrimination (smoothed `l1hingeloss`)
- `algorithm = SGD()`: The stochastic approximation method to be used.
    - Algorithms based on Stochastic gradient:
        - `SGD()`: Stochastic Gradient Descent
        - `ADAGRAD()`: AdaGrad (adaptive version of SGD)
        - `RMSPROP()`: RMSProp (adaptive version of SGD)
    - Algorithms based on Majorization-Minimization Principle:
        - `MSPI()`: Majorized Stochastic Proximal Iteration
        - `OMAS()`: Online MM via Averaged Surrogate
        - `OMAP()`: Online MM via Averaged Parameter
- `λ = 0.0`: The hyperparameter(s) used for the penalty function
    - User can provide elementwise penalty hyperparameters (`Vector{Float64}`) or single hyperparameter (`Float64`).

## Keyword Arguments

- `penalty`: The regularization function used.
    - `zero`: no penalty (default)
    - `abs`: (LASSO) parameters penalized by their absolute value
    - `abs2`: (Ridge) parameters penalized by their squared value
    - `ElasticNet(α)`: α * (abs penalty) + (1-α) * (abs2 penalty)
- `rate = LearningRate(.6)`

# Example

    x = randn(1000, 5)
    y = x * range(-1, stop=1, length=5) + randn(1000)

    o = fit!(StatLearn(MSPI()), zip(eachrow(x), y))
    coef(o)

    o = fit!(StatLearn(OnlineStats.l1regloss, ADAGRAD()), zip(eachrow(x), y))
    coef(o)
"""
mutable struct StatLearn{A<:Algorithm, L, P, W} <: OnlineStat{XY}
    β::Vector{Float64}
    λ::Vector{Float64}
    gx::Vector{Float64}
    loss::L
    penalty::P
    alg::A
    rate::W
    n::Int
end
function StatLearn(args...; penalty=zero, rate=LearningRate())
    p = 0
    λ = zeros(1)
    loss = l2regloss
    alg = SGD()
    for a in args
        if a isa AbstractVector
            λ = a
        elseif a isa Float64
            λ = fill(a, 1)
        elseif a isa Algorithm
            alg = a
        elseif a isa Integer
            p = a
        elseif hasmethod(a, (Number, Number)) && hasmethod(deriv, (typeof(a), Number, Number))
            loss = a
        else
            @warn """
            Arguments of type $(typeof(a)) are not recognized by StatLearn.  You may be seeing this as
            a result of OnlineStats removing LearnBase, LossFunctions, and PenaltyFunctions as
            dependencies.  See the `StatLearn` docstring for details.

            - Previous behavior for PenaltyFunctions can be restored e.g.

            ```
            OnlineStats.prox(p::Penalty, x, s) = PenaltyFunctions.prox(p, x, s)
            StatLearn(penalty = L1Penalty())
            ```

            - Previous behavior for LossFunctions can be restored e.g.:

            ```
            f(y, yhat) = LossFunctions.value(L2DistLoss(), y, yhat)
            OnlineStats.deriv(::typeof(f), y, yhat) = LossFunctions.deriv(L2DistLoss(), y, yhat)
            StatLearn(f)
            ```
            """
        end
    end
    init!(alg, p)
    StatLearn(zeros(p), λ, zeros(p), loss, penalty, alg, rate, 0)
end

function Base.show(io::IO, o::StatLearn)
    print(io, "StatLearn: ")
    print(io, name(o.alg, false, false))
    o.penalty != zero && print(io, " | mean(λ)=", mean(o.λ))
    print(io, " | ", o.loss)
    print(io, " | ", o.penalty)
    print(io, " | nobs=", nobs(o))
    print(io, " | nvars=", length(o.β))
end
coef(o::StatLearn) = value(o)

function gradient!(o::StatLearn, x, y)
    d_dη = deriv(o.loss, y, predict(o, x))
    for j in eachindex(o.gx)
        o.gx[j] = x[j] * d_dη
    end
end
function _fit!(o::StatLearn{<:SGAlgorithm}, xy)
    x, y = xy
    (o.n += 1) == 1  && init!(o, length(x))
    gradient!(o, x, y)
    update!(o.alg, o.gx)
    updateβ!(o, o.rate(o.n))
end
function init!(o::StatLearn, p)
    o.β = zeros(p)
    o.gx = zeros(p)
    init!(o.alg, p)
    if length(o.λ) == 1
        o.λ = fill(o.λ[1], p)
    end
end

function _merge!(o::StatLearn, o2::StatLearn)
    o.n += o2.n
    γ = nobs(o2) / nobs(o)
    smooth!(o.β, o2.β, γ)
    merge!(o.alg, o2.alg, γ)
    smooth!(o.λ, o2.λ, γ)
end

predict(o::StatLearn, x::VectorOb) = dot(x, o.β)
predict(o::StatLearn, x::AbstractMatrix) = x * o.β
classify(o::StatLearn, x) = sign.(predict(o, x))

function objective(o::StatLearn, x::AbstractMatrix, y::VectorOb)
    mean(o.loss.(y, predict(o,x))) + sum(o.λ .* o.penalty.(o.β))
end

#-----------------------------------------------------------------------# updateβ!
function updateβ!(o::StatLearn{SGD}, γ)
    for j in eachindex(o.β)
        o.β[j] = prox(o.penalty, o.β[j] - γ * o.gx[j], γ * o.λ[j])
    end
end
function updateβ!(o::StatLearn{T}, γ) where {T<:Union{ADAGRAD, RMSPROP}}
    for j in eachindex(o.β)
        s = γ / sqrt(o.alg.h[j] + ϵ)
        o.β[j] = prox(o.penalty, o.β[j] - s * o.gx[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADAM}, γ)
    for j in eachindex(o.β)
        s = γ / sqrt(o.alg.v[j] + ϵ)
        o.β[j] = prox(o.penalty, o.β[j] - s * o.alg.m[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADAMAX}, γ)
    for j in eachindex(o.β)
        s = γ / ((1 - o.alg.β1^nobs(o)) * o.alg.v[j])
        o.β[j] = prox(o.penalty, o.β[j] - s * o.alg.m[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADADELTA}, γ)
    for j in eachindex(o.β)
        s = o.alg.δ[j]
        o.β[j] = prox(o.penalty, o.β[j] - s * o.gx[j], s * o.λ[j])
    end
end

#------------------------------------------------------------------# Majorization-based
# lipschitz_constant (L): f(θ) ≤ f(x) + ∇f(x)'(θ - x) + (L / 2) ||θ - x||^2
lconst(loss, x, y) = error("Unkown Lipschitz constant for loss $loss")

lconst(o::typeof(l2regloss), x, y) = dot(x, x)

lconst(o::typeof(logisticloss), x, y) = .25 * dot(x, x)

lconst(o::DWDLoss, x, y) = (o.q + 1)^2 / o.q * dot(x, x)

#-----------------------------------------------------------------------# OMAS
# L stored in o.alg.a[1]
function _fit!(o::StatLearn{OMAS}, xy)
    x, y = xy
    (o.n += 1) == 1  && init!(o, length(x))
    γ = o.rate(o.n)
    b = o.alg.b
    gradient!(o, x, y)
    ht = lconst(o.loss, x, y)
    L = (o.alg.a[1] = smooth(o.alg.a[1], ht, γ))
    for j in eachindex(o.β)
        b[j] = smooth(b[j], ht * o.β[j] - o.gx[j], γ)
        o.β[j] = prox(o.penalty, b[j] / L, o.λ[j] / L)
    end
end
#-----------------------------------------------------------------------# OMAP
function _fit!(o::StatLearn{OMAP}, xy)
    x, y = xy
    (o.n += 1) == 1 && init!(o, length(x))
    γ = o.rate(o.n)
    gradient!(o, x, y)
    h_inv = inv(lconst(o.loss, x, y))
    for j in eachindex(o.β)
        o.β[j] -= γ * h_inv * o.gx[j]
    end
end
#-----------------------------------------------------------------------# MSPI
function _fit!(o::StatLearn{MSPI}, xy)
    x, y = xy
    (o.n += 1) == 1  && init!(o, length(x))
    γ = o.rate(o.n)
    gradient!(o, x, y)
    γ2 = γ / (1 + γ * lconst(o.loss, x, y))
    for j in eachindex(o.β)
        @inbounds o.β[j] = prox(o.penalty, o.β[j] - γ2 * o.gx[j], γ2 * o.λ[j])
    end
end


normal_kurtosis(x) = kurtosis(x, false)

function ks_test(a,b)
    x = ApproximateTwoSampleKSTest(a[2:end],b[2:end])
    return sqrt(x.n_x*x.n_y/(x.n_x+x.n_y))*x.δ
end

function generalized_hurts_exp(log_prices)
    q = 1
    max_τ = 19

    k = 0
    H = zeros(15)
    L = size(log_prices, 1)

    for iter_τ = 5:max_τ
        k += 1

        x = collect(1:iter_τ)
        k_q_t = zeros(iter_τ)
        for τ = 1:iter_τ
            numer = log_prices[(τ+1):τ:L] - log_prices[1:τ:(L-τ)]
            denom = log_prices[1:τ:L]

            # Determine Drift
            N = size(numer, 1) + 1
            X = collect(1:N)
            Y = denom
            mx = sum(X)/N
            SSxx = sum(X.^2) - N * mx^2
            my = sum(Y) / N
            SSxy = sum(X .* Y) - N * mx * my
            cc = [SSxy/SSxx, my - SSxy/SSxx*mx]

            # Subtract Drift
            numer = numer .- cc[1]
            denom = denom .- cc[1] * collect(1:N) .- cc[2]

            k_q_t[τ] = mean(abs.(numer).^q) / mean(abs.(denom).^q)
        end

        # Calculate Hurst Exponent for current iteration
        log_10_x = log10.(x)
        mx = mean(log_10_x)
        SSxx = sum(log_10_x.^2) - iter_τ * mx^2
        log_10_k_q_t = log10.(k_q_t)
        my = mean(log_10_k_q_t)
        SSxy = sum(log_10_x .* log_10_k_q_t) - iter_τ * mx * my
        H[k] = SSxy/SSxx
    end

    H = mean(H) / q
    return H
end

function get_summary_stats(simualted_log_prices, log_prices)

    num_replications = size(simualted_log_prices,2)
    stats_mat = zeros(num_replications,5)

    stats_mat[:,1] = mean(simualted_log_prices, dims=1)
    stats_mat[:,2] = std(simualted_log_prices, dims=1)
    stats_mat[:,3] = mapslices(normal_kurtosis, simualted_log_prices, dims=1)

    ks_test_stat(x) = ks_test(log_prices, x)
    stats_mat[:,4] = mapslices(ks_test_stat, simualted_log_prices, dims=1)
    stats_mat[:,5] = mapslices(generalized_hurts_exp, simualted_log_prices,
        dims=1)
    return stats_mat
end


@testset "Rosenbrock" begin

    function test_result(result::Vector, fitness::Float64, N::Int, tol::Float64)
        @test length(result) == N
    	@test ≈(result, ones(N), atol=tol)
        @test ≈(fitness, 0.0, atol=tol)
    end

	# Objective function
	rosenbrock(x::Vector{Float64}) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

	# Parameters
	N = 2
	terminate(σ) = σ < 1e-10

	# Testing: (15/5+100)-ES
	# with isotropic mutation operator y' := y + σ(N_1(0, 1), ..., N_N(0, 1))
	result, fitness, cnt = es(rosenbrock, N;
        initStrategy = strategy(σ = 1.0),
        recombination = average, mutation = isotropic,
        μ = 15, ρ = 5, λ = 100, iterations = 1000)
	println("(15/5+100)-ES => F: $(fitness), C: $(cnt), OBJ: $(result)")
	test_result(result, fitness, N, 1e-1)

	# Testing: (15/15+100)-σ-Self-Adaptation-ES
	# with isotropic mutation operator y' := y + σ(N_1(0, 1), ..., N_N(0, 1))
	result, fitness, cnt = es(rosenbrock, N;
		initPopulation = [.5, .5],
        initStrategy = strategy(σ = 1.0, τ = 1/sqrt(2*N)),
		recombination = average, srecombination = averageSigma1,
		mutation = isotropic, smutation = isotropicSigma,
		μ = 15, λ = 100, iterations = 1000)
	println("(15/15+100)-σ-SA-ES-IS => F: $(fitness), C: $(cnt), OBJ: $(result)")
	test_result(result, fitness, N, 1e-1)

	# Testing: (15/15+100)-σ-Self-Adaptation-ES
	# with non-isotropic mutation operator y' := y + (σ_1 N_1(0, 1), ..., σ_N N_N(0, 1))
	result, fitness, cnt = es(rosenbrock, N;
		initPopulation = rand(N,25),
        initStrategy = strategy(σ = .5ones(N), τ = 1/sqrt(2*N), τ0 = 1/sqrt(N)),
		recombination = average, srecombination = averageSigmaN,
		mutation = anisotropic, smutation = anisotropicSigma,
		μ = 15, λ = 100, iterations = 1000)
	println("(15/15+100)-σ-SA-ES-AS => F: $(fitness), C: $(cnt), OBJ: $(result)")
	test_result(result, fitness, N, 1e-1)

	# Testing: CMA-ES
	result, fitness, cnt = cmaes(rosenbrock, N;
    	μ = 3, λ = 12, iterations = 100_000, tol = 1e-3)
    println("(3/3,12)-CMA-ES => F: $(fitness), C: $(cnt), OBJ: $(result)")
    test_result(result, fitness, N, 1e-1)

    # Testing: GA
    result, fitness, cnt = ga(rosenbrock, N;
    	initPopulation = (n -> rand(n)),
    	populationSize = 100,
    	ɛ = 0.1,
        selection = sus,
        crossover = intermediate(0.25),
        mutation = domainrange(fill(0.5,N)))
    println("GA(p=50,x=0.8,μ=0.1) => F: $(fitness), C: $(cnt), OBJ: $(result)")
    test_result(result, fitness, N, 3e-1)

end


using Zygote, LinearAlgebra
using KrylovKit: eigsolve

# helper for `herm`
function f(x, y, i, j)
    if i == j
        return x/2 + im*zero(y)
    elseif i > j
        return x + im * y
    elseif i < j
        return zero(x) + im*zero(y)
    end
end

# helper for invherm
function finv(z, i, j)
    if i <= j
        return real(z)
    elseif i > j
        return imag(z)
    end
end

# Create a complex `d` by `d` Hermitian matrix from a `d^2`-dimensional real vector `v`
function herm(v::AbstractVector)
    d = isqrt(length(v))
    @assert d^2 == length(v)
    m = reshape(v, d, d)
    m = [ f(m[i,j], m[j,i], i, j) for i = 1:d, j = 1:d]
    m = m + m'
    return m
end

# Recover the vector representation
function invherm(A::AbstractMatrix)
    d = size(A,1)
    m = [ finv(A[i,j], i, j) for i = 1:d for j = 1:d ]
    return vec(m)
end

# Find the minimal eigenvalue using KrylovKit (from a vector representation)
function eigmin_vec(v::AbstractVector)
    A = herm(v)
    eigsolve(A, 1, :SR)[1][1] |> real
end

# Give an adjoint for `Zygote`, which technically only works in the non-degenerate case
Zygote.@adjoint function eigmin_vec(v::AbstractVector)
    A = herm(v)
    eigs = eigsolve(A, 1, :SR)
    eval = eigs[1][1]
    evec = eigs[2][1]
    adj = (evec * evec') |> invherm
    return eval |> real, c->(c * adj,)
end

A = rand(4, 4); A = A + A';
v = invherm(A)

@show eigmin(A)
@show eigmin_vec(v)
@show Zygote.gradient(eigmin_vec, v)

Methods/@adjoint (real rep of Hermitian).jl


function sieve(limit::Integer)
    limit <= 0 && error("Invalid limit")
    nums = fill(true, limit)
    nums[1] = false
    primes = []
    while (i = findfirst(nums)) > 0
        push!(primes, i)
        for j = find(nums)
            if j % i == 0
                nums[j] = false
            end
        end
    end
    primes
end


@doc raw"""
    zernike(x, y, maxord)

Calculate  unit-normalized Zernike polynomials and their x,y derivatives

Translation of the code from
[1] T. B. Andersen, “Efficient and robust recurrence relations for the Zernike circle polynomials
and their derivatives in Cartesian coordinates,” _Opt. Express_, vol. 26, no. 15, p. 18878, Jul. 2018.

     Numbering scheme:
     Within a radial order, sine terms come first
             ...
           sin((n-2m)*theta)   for m = 0,..., [(n+1)/2]-1
             ...
              1                for n even, m = n/2
             ...
           cos((n-2m)*theta)   for m = [n/2]+1,...,n
             ...

     INPUT:
     x, y normalized (x,y) coordinates in unit circle
     MaxOrd: Maximum Zernike radial order

     OUTPUT:
     dUdx[...]   array to receive each derivative dU/dx at (x,y)
     dUdy[...]   array to receive each derivative dU/dy at (x,y)

"""
function zernike(x, y, maxord::Int)

    ztotnumber = Int((maxord + 2) * (maxord + 1) / 2)
    # print(ztotnumber)
    zern = zeros(ztotnumber)
    dudx = zeros(ztotnumber)
    dudy = zeros(ztotnumber)

    zern[1] = 1
    dudx[1] = 0
    dudy[1] = 0

    zern[2] = y
    zern[3] = x
    dudx[2] = 0
    dudx[3] = 1
    dudy[2] = 1
    dudy[3] = 0

    kndx = 1                # index for term from 2 orders down
    jbeg = 2                # start index for current radial order
    jend = 3                # end index for current radial order
    jndx = 3                # running index for current Zern
    even = -1

    #  Outer loop in radial order index
    for nn = 2:maxord

        even   = -even          # parity of radial index
        jndx1  = jbeg           # index for 1st ascending series in x
        jndx2  = jend           # index for 1st descending series in y
        jndx11 = jndx1 - 1      # index for 2nd ascending series in x
        jndx21 = jndx2 + 1      # index for 2nd descending series in y
        jbeg   = jend + 1       # end of previous radial order +1
        nn2    = nn / 2
        nn1    = (nn - 1) / 2

      # Inner loop in azimuthal index
        for mm = 0:nn
            jndx += 1                  # increment running index for current Zern

            if mm == 0
                zern[jndx] = x * zern[jndx1] + y * zern[jndx2]
                dudx[jndx] = zern[jndx1] * nn
                dudy[jndx] = zern[jndx2] * nn

            elseif mm == nn
                zern[jndx] = x * zern[jndx11] - y * zern[jndx21]
                dudx[jndx] = zern[jndx11] * nn
                dudy[jndx] = -zern[jndx21] * nn

            elseif even > 0 && mm == nn2
                zern[jndx] = 2 * (x * zern[jndx1] + y * zern[jndx2]) - zern[kndx]
                dudx[jndx] = 2 * nn * zern[jndx1] + dudx[kndx]
                dudy[jndx] = 2 * nn * zern[jndx2] + dudy[kndx]
                kndx += 1

            elseif even < 0 && mm == nn1
                qval = zern[jndx2] - zern[jndx21]
                zern[jndx] = x * zern[jndx11] + y * qval - zern[kndx]
                dudx[jndx] = zern[jndx11] * nn + dudx[kndx]
                dudy[jndx] = qval * nn + dudy[kndx]
                kndx += 1

            elseif even < 0 && mm == nn1 + 1
                pval = zern[jndx1] + zern[jndx11]
                zern[jndx] = x * pval + y * zern[jndx2] - zern[kndx]
                dudx[jndx] = pval * nn + dudx[kndx]
                dudy[jndx] = zern[jndx2] * nn + dudy[kndx]
                kndx += 1

            else
                pval = zern[jndx1] + zern[jndx11]
                qval = zern[jndx2] - zern[jndx21]
                zern[jndx] = x * pval + y * qval - zern[kndx]
                dudx[jndx] = pval * nn + dudx[kndx]
                dudy[jndx] = qval * nn + dudy[kndx]
                kndx += 1

            end

            jndx11 = jndx1                   # update indices
            jndx1 += 1
            jndx21 = jndx2
            jndx2 += -1
        end

        jend = jndx
    end

    (z = zern, zx = dudx, zy = dudy)

end

"""
    ZernikeBW(elements, ap, mask, norms)

Contains Zernike basis (in Born&Wolf norming) with the aperture, plotting mask, and element norms. Can be constructed as
    `ZernikeBW(dom::CartesianDomain2D, d::Real, maxorder::Integer) `
and
    `ZernikeBW(gridsize::Integer, maxorder::Integer)`

"""
struct ZernikeBW <: OrthogonalBasis
    elements::VectorOfArray
    ap::Array
    mask::Array
    norms::Vector
    ZernikeBW(elements, ap, mask) = new(elements, ap, mask, [sqrt.(inner(f, ap .* f)) for f in elements])
end

ZernikeBW(gridsize::Integer, maxorder::Integer)  = ZernikeBW(makezerniketable(gridsize, maxorder),makeaperture(gridsize)...)
ZernikeBW(dom::CartesianDomain2D, d::Real, maxorder::Integer)  = ZernikeBW(makezerniketable(dom, maxorder,d/2.),aperture(dom, d)...)

# function makezerniketable(gridsize::Integer, maxorder::Integer)
#     x = range(-1, 1, length=gridsize)
#     y = range(-1, 1, length=gridsize)
#     totalznum = Int((maxorder + 2) * (maxorder + 1) / 2)
#     ztable = [zernike(xc, yc, maxorder)[:z] for xc ∈ x,  yc ∈ y ]
#     zvec = VectorOfArray([zeros(gridsize, gridsize) for i = 1:totalznum])
#     [zvec[i,:] = ztable[i] for i = eachindex(ztable) ]
#     return zvec
# end

function makezerniketable(dom::CartesianDomain2D,  maxorder::Integer, scale = 1)
    x= dom.xrange/scale
    y = dom.yrange/scale
    totalznum = Int((maxorder + 2) * (maxorder + 1) / 2)
    ztable = [zernike(xc, yc, maxorder)[:z] for yc ∈ y, xc ∈ x ]
    zvec = VectorOfArray([zeros(length(y), length(x)) for i = 1:totalznum])
    [zvec[i,:] = ztable[i] for i = eachindex(ztable) ]
    return zvec
end

makezerniketable(gridsize::Integer, maxorder::Integer) =
    makezerniketable(CartesianDomain2D(range(-1, 1, length=gridsize), range(-1, 1, length=gridsize)), maxorder)

function makeaperture(gridsize::Integer, δ = 0.)
    x = range(-1, 1, length=gridsize)
    y = range(-1, 1, length=gridsize)
    # δ = 0. # tuning of the aperture size
    r =1 + δ  /gridsize
    ap = [ (xc^2 + yc^2) <= r^2 ? 1 : 0 for xc ∈ x, yc ∈ y]
    # area = +(ap[:]...)
    phmask = [ (xc^2 + yc^2) <= r^2 ? 1 : NaN for xc ∈ x, yc ∈ y]
    return(ap, phmask)
end


using ConvDiffMIPDECO
using Test
using jInv.Mesh

domain = [0. 3. 0 1. 0 2.]
n      = 3*[7 9 12] .- 1
M      = getRegularMesh(domain,n)
Mass, Mass_const, SM = getFEMMatrices3D(M)

e = ones(prod(M.n.+1))
@test abs(prod((domain[2:2:end]-domain[1:2:end])) - dot(e,Mass*e))/dot(e,Mass*e) < 1e-2

f = getFEMsource3D(M)
v = Mass_const*f
u = SM\v


#############
# Exp & Log #
#############

struct Exp{N} <: Bijector{N} end
struct Log{N} <: Bijector{N} end

Exp() = Exp{0}()
Log() = Log{0}()

(b::Exp)(y::Real) = exp(y)
(b::Log)(x::Real) = log(x)

(b::Exp)(y) = exp.(y)
(b::Log)(x) = log.(x)

inv(b::Exp{N}) where {N} = Log{N}()
inv(b::Log{N}) where {N} = Exp{N}()

logabsdetjac(b::Exp{0}, x::Real) = x
logabsdetjac(b::Exp{0}, x::AbstractVector) = x
logabsdetjac(b::Exp{1}, x::AbstractVector) = sum(x)
logabsdetjac(b::Exp{1}, x::AbstractMatrix) = vec(sum(x; dims = 1))

logabsdetjac(b::Log{0}, x::Real) = -log(x)
logabsdetjac(b::Log{0}, x::AbstractVector) = -log.(x)
logabsdetjac(b::Log{1}, x::AbstractVector) = - sum(log.(x))
logabsdetjac(b::Log{1}, x::AbstractMatrix) = - vec(sum(log.(x); dims = 1))


using MatrixCompletion
import Distributions

@testset "$(format("Estimator: MOM[NegativeBinomial]]"))"  begin
  let
    for i in 1:10
      p_test = rand()
      r_test = rand() * 20
      data = rand(Distributions.NegativeBinomial(r_test, p_test), 1000 * 200)
      tc = estimator(MOM{NegativeBinomial}(), data)
      @test abs(tc[:p] - p_test) / p_test < 0.05
      @test abs(tc[:r] - r_test) / r_test < 0.05
    end
  end
end


@testset "Tracker" begin
    @testset "tracking - AD : $AD" for AD = 0:3
        @var x a y b
        F = System([x^2 - a, x * y - a + b], [x, y], [a, b])

        tracker = Tracker(
            ParameterHomotopy(F, [1, 0], [2, 4]),
            options = TrackerOptions(automatic_differentiation = AD),
        )
        s = [1, 1]
        res = track(tracker, s, 1, 0)
        @test is_success(res)
        @test steps(res) ≤ 5
        @test isa(solution(res), Vector{ComplexF64})
        @test solution(res) ≈ [sqrt(2), -sqrt(2)]
        @test length(solution(res)) == 2
        @test is_success(track(tracker, res, 0, 1))

        code = track!(tracker, s, 1, 0)
        @test is_success(code)
        s0 = copy(tracker.state.x)
        @unpack μ, ω = tracker.state
        @test is_success(track(tracker, s0, 0, 1))
        @test is_success(track(tracker, s0, 0, 1, μ = μ, ω = ω))

        s = @SVector [1, 1]
        @test is_success(track(tracker, s, 1, 0))
    end

    @testset "projective tracking- AD: $AD" for AD = 0:3
        @var x a y b z
        F = System([x^2 - a * z^2, x * y + (b - a) * z^2], [x, y, z], [a, b])
        H = ParameterHomotopy(F, [1, 0], [2, 4])
        tracker = Tracker(
            on_affine_chart(H, (2,)),
            options = TrackerOptions(automatic_differentiation = AD),
        )

        s = [1, 1, 1]
        res = track(tracker, s, 1, 0)
        @test is_success(res)
        @test isa(solution(res), Vector{ComplexF64})
        x₀ = abs(solution(res)[end])
        @test affine_chart(PVector(solution(res))) ≈ [sqrt(2), -sqrt(2)] rtol = 1e-12 / x₀
    end

    @testset "iterator" begin
        @var x a y b
        F = System([x^2 - a, x * y - a + b], [x, y], [a, b])
        tracker = Tracker(ParameterHomotopy(F, [1, 0], [2, 4]))
        s = [1, 1]

        # path iterator
        typeof(first(iterator(tracker, s, 1.0, 0.0))) == Tuple{Vector{ComplexF64},Float64}

        tracker.options.max_step_size = 0.01
        @test length(collect(iterator(tracker, s, 1.0, 0.0))) ≥ 101

        F = System([x - a], [x], [a])
        ct = Tracker(
            ParameterHomotopy(F, [1], [2]),
            options = TrackerOptions(max_step_size = 0.015625),
        )
        Xs = Vector{ComplexF64}[]
        for (x, t) in iterator(ct, [1.0], 1.0, 0.0)
            push!(Xs, x)
        end

        @test length(Xs) ≥ length(1:0.015625:2)
    end

    @testset "path info" begin
        @var x a y b
        F = System([x^2 - a, x * y - a + b], [x, y], [a, b])
        tracker = Tracker(ParameterHomotopy(F, [1, 0], [2, 4]))
        info = path_info(tracker, [1, 1], 1, 0)
        @test !isempty(sprint(show, info))
    end

    @testset "Change parameters" begin
        @var x a y b
        F = System([x^2 - a, x * y - a + b]; parameters = [a, b])
        s = [1.0, 1.0 + 0im]
        tracker = Tracker(ParameterHomotopy(F, [2.2, 3.2], [2.2, 3.2]))
        start_parameters!(tracker, [1, 0])
        target_parameters!(tracker, [2, 4])
        res = track(tracker, s, 1.0, 0.0)
        @test is_success(res)
    end

    @testset "Straight Line Homotopy" begin
        @var x y
        F = System([x^2 + y^2 - 2.3, 2x + 3y - 4], [x, y])
        G = System((0.2 + 0.4im) .* [x^2 - 1, y - 1], [x, y])
        H = StraightLineHomotopy(G, F)
        S = [[1, 1], [-1, 1]]
        tracker = Tracker(H, options = (automatic_differentiation = 3,))

        @test is_success(track(tracker, S[1], 1, 0))
        @test is_success(track(tracker, S[2], 1, 0))
    end

    include("test_cases/steiner_higher_prec.jl")
    include("test_cases/four_bar.jl")
end


using Distributions
using Base.Test

# Test the consistency between the recursive and nonrecursive computation of the pdf
# of the Binomial distribution
srand(1234)
for (p, n) in [(0.6, 10), (0.8, 6), (0.5, 40), (0.04, 20), (1., 100), (0., 10), (0.999999, 1000), (1e-7, 1000)]

    d = Binomial(n, p)

    a = pdf(d, 0:n)
    for t=0:n
        @test pdf(d, t) ≈ a[1+t]
    end

    li = rand(0:n, 2)
    rng = minimum(li):maximum(li)
    b = pdf(d, rng)
    for t in rng
        @test pdf(d, t) ≈ b[t - first(rng) + 1]
    end

end


struct Projection{S, I}
    set::S
    indices::I
end
Polyhedra.project(set, indices) = Projection(set, indices)
function space_variables(p::Projection)
    vars = space_variables(p.set)
    if vars === nothing
        return nothing
    else
        return vars[p.indices]
    end
end
dimension(p::Projection) = length(p.indices)

# A^{-1} * S
struct LinearImage{S, T, MT <: AbstractMatrix{T}} <: AbstractSet{T}
    set::S
    A::MT
end
perspective_variable(li::LinearImage) = perspective_variable(li.set)
space_variables(li::LinearImage) = nothing
dimension(li::LinearImage) = size(li.A, 1)

# A^{-1} * S
struct LinearPreImage{S, T, MT <: AbstractMatrix{T}} <: AbstractSet{T}
    set::S
    A::MT
end

# S + c
struct Translation{S, T, VT <: AbstractVector{T}} <: AbstractSet{T}
    set::S
    c::VT
end
dimension(t::Translation) = length(t.c)
space_variables(t::Translation) = space_variables(t.set)
function Polyhedra.project(t::Translation, I)
    return Translation(Polyhedra.project(t.set, I), t.c[I])
end

"""
    householder(x)

Householder reflection
```math
I - 2 v v^T / (v^T v)
```
It is symmetric and orthogonal.
"""
function householder(x)
    y = copy(x)
    t = LinearAlgebra.reflector!(y)
    v = [1; y[2:end]]
    I - t * v * v'
end
_householder(h) = householder([1.0; h]) # We add 1, for perspective variable z

struct Householder{T, S <: AbstractSet{T}, U} <: AbstractSet{T}
    set::S
    p::DynamicPolynomials.Polynomial{true, U}
    h::Vector{Float64}
    z::SpaceVariable
    x::Vector{SpaceVariable}
end
perspective_gauge0(set::Householder) = set.p
perspective_variable(set::Householder) = set.z
space_variables(set::Householder) = set.x
convexity_proof(set::Householder) = convexity_proof(set.set)

const HouseDualOf{S, T, U} = PerspectiveDualOf{Householder{T, S, U}}

function Polyhedra.project(set::Polar{T}, I) where T
    return Polyhedra.polar(zero_eliminate(Polyhedra.polar(set), setdiff(1:dimension(set), I)))
end

function zero_eliminate(set::Householder{T}, I) where T
    J = setdiff(1:dimension(set), I)
    p = subs(set.p, set.x[I] => zeros(T, length(I)))
    return Householder(UnknownSet{T}(), p, set.h[J], set.z, set.x[J])
end
function Polyhedra.project(set::PerspectiveDual, I)
    return perspective_dual(zero_eliminate(perspective_dual(set),
                                           setdiff(1:dimension(set), I)))
end

function _HPH(set::Householder)
    H = _householder(set.h)
    return H * _HPH(set.set) * H
end