"""
Description createSystem
"""
function createSystem(sys::System,G=6.673e-11::Float64)
    if      typeof(sys) == EMsystem
        return SystemInfo(EM.mu,EM.Rp,EM.Rs,EM.r12,EM.m1,Float64(G))
    elseif  typeof(sys) == SEsystem
        return SystemInfo(SE.mu,SE.Rp,SE.Rs,SE.r12,SE.m1,Float64(G))
    end
end
"""
Description createConstants
"""
function createConstants(sysData::SystemInfo)
    G       = sysData.G # Gravity constant      [m3/Kg/s2]
    ### Make sure that the arguments are Floats
    mu      = sysData.mu
    if mu > 1 || mu < 0
        error("mu parameter has to be 0<mu<1")
    end
    Rp      = sysData.Rp
    Rs      = sysData.Rs
    r12     = sysData.r12
    m1      = sysData.m1
    ###
    m2      = m1*mu/(1-mu)
    # Rotation of reference system [rad/s]
    w       =  sqrt(G*(m1+m2)/r12^3)
    # Gravity constant of Earth [m^3/s^2]
    mup     = G*m1
    return ProblemInfo(mu,Rp/r12,Rs/r12,r12,w,mup/r12^3/w^2,-mu,1-mu)
end
"""
Description createPar
"""
function createPar(prob_info;   tspan = -2.0, hgoal = 300, hmmin = 100 ,
                        hmmax = 20e3, sType = "global", points = 20, x0 = [],
                        maxt0 = -1.0, rest = "basic", lim = 2000,
                        cTol = 1e-4, tRest = "ineq", alg = "lowOrder",
                        atol = 1e-7, rtol = 1e-6, retorno = true, detMin = true, y = [], optTol = 1e-8, neval = 400)

    if length(x0) == 0
        println("There has to be an initial guess ('x0' argument)")
        return nothing
    elseif length(x0) == 3
        stop = true
    elseif length(x0) == 4 || length(x0) == 6
        detMin  = 1.0
        stop    = true
    elseif length(x0) == 5 || length(x0) == 7
        stop = false
    end

    if length(x0) == 4 || length(x0) == 5
        col = true
    else
        col = false
    end

    if sType == "global"
        type = true
    elseif sType == "local"
        type = false
    else
        println("The search type has to be either 'global' or 'local' ('sType' argument)")
        return nothing
    end

    if tRest == "ineq"
        restType = true
    elseif tRest == "mixed"
        restType = false
    else
        println("The restriction type has to be 'ineq' or 'mixed' ('tRest' argument)")
        return nothing
    end
    restriction = dictRest[rest]

    if alg == "lowOrder"
        odeAlg = true
    elseif alg == "highOrder"
        odeAlg = false
    else
        println("Wrong ODE algorithm, choose between 'lowOrder' or 'highOrder'")
        return nothing
    end

    if y == []
        println("There has to be an insertion point ('y' argument)")
        return nothing
    elseif length(y) > 6 || length(y) < 6
        println("The insertion point ('y' argument) has to be of length 6")
        return nothing
    end

    println("--- Problem parameters created correctly ---")
    return ProblemParameters(Float64(tspan),Float64(hgoal*1e3/prob_info.r12),Float64(hmmin*1e3/prob_info.r12),Float64(hmmax*1e3/prob_info.r12),type,Int64(points),stop,Float64(maxt0),
                            restriction,Float64(lim/prob_info.r12/prob_info.w),Float64(cTol),restType,odeAlg,Float64(atol),Float64(rtol),
                            retorno,Float64(detMin),convert.(Float64,x0),convert.(Float64,y),optTol,Int64(neval),col)
end
"""
Description dineq!
"""
function dineq!(dy,y,p,t,mu::Float64)
    #mu  = p[1]
    xp, yp, zp, vx, vy, vz = y
    ro1 = sqrt((xp+mu)^2      +yp^2 +zp^2)
    ro2 = sqrt((xp-(1-mu))^2  +yp^2 +zp^2)
    dy[1] = vx
    dy[2] = vy
    dy[3] = vz
    dy[4] = 2*vy+xp-(1-mu)*(xp+mu)/ro1^3-mu*(xp-(1-mu))/ro2^3
    dy[5] = -2*vx+yp*(1 - (1-mu)/ro1^3-mu/ro2^3)
    dy[6] = -zp*((1-mu)/ro1^3+mu/ro2^3)
end
"""
Description calc_data
"""
function calc_data(x::Array{Float64,1},par::ProblemParameters, prob_info::ProblemInfo, PROB::Union{Tuple{ODEProblem,ODEProblem},Tuple{ODEProblem}})
    x1      = prob_info.x1      ::Float64
    x2      = prob_info.x2      ::Float64
    rearth  = prob_info.rp      ::Float64
    rmoon   = prob_info.rs      ::Float64
    prob    = PROB[1]
    y       = prob.u0           ::Array{Float64,1}

    ny0 = y - vcat(zeros(3),x[1:3])

    if par.stop
        if length(x) == 3
            intProb          = remake(prob, u0 = ny0 )
        elseif length(x) == 4
            intProb          = remake(prob, u0 = ny0 , p = [x[4], 0.0, 0.0, 0.0, par.detMin])
        else
            intProb          = remake(prob, u0 = ny0 , p = [x[4], x[5], x[6], 0.0, par.detMin])
        end
        tstop = [0.0]
    else
        if length(x) == 5
            intProb          = remake(prob, u0 = ny0 , p = [x[4], 0.0, 0.0, x[5], par.detMin])
        elseif length(x) == 7
            intProb          = remake(prob, u0 = ny0 , p = [x[4], x[5], x[6], x[7], par.detMin])
        end
        tstop = [x[end]]
    end

    if par.odeAlg
        sol     = solve(intProb, DP5(), abstol = par.atol, reltol= par.rtol, tstops= tstop, dense=false, save_everystep=false)
    else
        sol     = solve(intProb, Vern7(), abstol = par.atol, reltol= par.rtol, tstops= tstop, dense=false, save_everystep=false)
    end

    ts, tp, ps, pp, ImpF, ImpS, ImpI, tImpF, tImpS, tImpI = calcOuts(x,sol,par,prob_info)

    if par.retorno
        probR = PROB[2]
        if length(x) == 3
            Y0 = ImpF
        elseif length(sol.u) > 2
            Y0 = ImpS
        else
            Y0 = ImpF
        end
        prob2   = remake(probR, u0 = Y0 , tspan = (0,2.5) )
        if par.odeAlg
            solr    = solve(prob2, DP5(), abstol = par.atol, reltol= par.rtol, dense=false, save_everystep=false)
        else
            solr    = solve(prob2, Vern7(), abstol = par.atol, reltol= par.rtol, dense=false, save_everystep=false)
        end
        eRet       = calcDists(solr.u[end], x1,rearth)
    else
        eRet       = 0.0
    end

    ds, dp   = calcDists(ps, x2, rmoon), calcDists(pp, x1, rearth) # distances to primaries in the minimum distance points

    fimp = fImp(pp,prob_info)

    return ds, dp, ts, tp, eRet, ImpF, ImpS, ImpI, tImpF, tImpS, tImpI, fimp
end
# Add method to calc_data
function calc_data(x::Vector{T},par::ProblemParameters, prob_info::ProblemInfo, PROB::Tuple{ODEProblem}) where {T<:ForwardDiff.Dual}
    x1      = prob_info.x1      ::Float64
    x2      = prob_info.x2      ::Float64
    rearth  = prob_info.rp      ::Float64
    rmoon   = prob_info.rs      ::Float64
    prob    = PROB[1]

    xType = eltype(x)
    if par.stop
        if length(x) == 3
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x), tspan = convert.(xType, prob.tspan) )
        elseif length(x) == 4
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], 0.0, 0.0, 0.0, par.detMin]) )
        else
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], x[5], x[6], 0.0, par.detMin]) )
        end
        tstop = convert.(xType,[0.0])
    else
        if length(x) == 5
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], 0.0, 0.0, x[5], par.detMin]) )
        elseif length(x) == 7
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], x[5], x[6], x[7], par.detMin]) )
        end
        tstop = x[end]
    end

    if par.odeAlg
        sol     = solve(intProb, DP5(), abstol = par.atol, reltol= par.rtol, tstops=tstop, dense=false, save_everystep=false)
    else
        sol     = solve(intProb, Vern7(), abstol = par.atol, reltol= par.rtol, tstops=tstop, dense=false, save_everystep=false)
    end

    ts, tp, ps, pp = calcOuts(x,sol,par,prob_info,PROB)

    #ds, dp   = calcDists(ps, x2, rmoon), calcDists(pp, x1, rearth) # distances to primaries in the minimum distance points
    ds = calcDists(ps, x2, rmoon)
    dp = calcDists(pp, x1, rearth)

    return ds, dp, ts, tp
end
function calc_data(x::Vector{T},par::ProblemParameters, prob_info::ProblemInfo, PROB::Tuple{ODEProblem,ODEProblem}) where {T<:ForwardDiff.Dual}
    x1      = prob_info.x1      ::Float64
    x2      = prob_info.x2      ::Float64
    rearth  = prob_info.rp      ::Float64
    rmoon   = prob_info.rs      ::Float64
    prob    = PROB[1]
    probR   = PROB[2]

    xType = eltype(x)
    if par.stop
        if length(x) == 3
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x), tspan = convert.(xType, prob.tspan) )
        elseif length(x) == 4
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], 0.0, 0.0, 0.0, par.detMin]) )
        else
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], x[5], x[6], 0.0, par.detMin]) )
        end
        tstop = convert.(xType,[0.0])
    else
        if length(x) == 5
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], 0.0, 0.0, x[5], par.detMin]) )
        elseif length(x) == 7
            intProb          = remake(prob, u0 = prob.u0 - vcat(zeros(3),x[1:3]), tspan = convert.(xType, prob.tspan), p = convert.(xType,[x[4], x[5], x[6], x[7], par.detMin]) )
        end
        tstop = x[end]
    end

    if par.odeAlg
        sol     = solve(intProb, DP5(), abstol = par.atol, reltol= par.rtol, tstops=tstop, dense=false, save_everystep=false)
    else
        sol     = solve(intProb, Vern7(), abstol = par.atol, reltol= par.rtol, tstops=tstop, dense=false, save_everystep=false)
    end

    ts, tp, ps, pp, IMP = calcOuts(x,sol,par,prob_info,PROB)

    prob2   = remake(probR, u0 = IMP , tspan = convert.(xType,(0.0,2.5)) )
    if par.odeAlg
        solr    = solve(prob2, DP5(), abstol = par.atol, reltol= par.rtol, dense=false, save_everystep=false, callback = callRet )
    else
        solr    = solve(prob2, Vern7(), abstol = par.atol, reltol= par.rtol, dense=false, save_everystep=false, callback = callRet )
    end
    eRet       = calcDists(solr.u[end], x1,rearth)

    #ds, dp   = calcDists(ps, x2, rmoon), calcDists(pp, x1, rearth) # distances to primaries in the minimum distance points
    ds = calcDists(ps, x2, rmoon)
    dp = calcDists(pp, x1, rearth)

    return ds, dp, ts, tp, eRet
end
"""
Description calcOuts
"""
function calcOuts(x::Array{Float64,1},sol::ODESolution,par::ProblemParameters,prob_info::ProblemInfo)
    ts = sol.t[1] # Instant of close approach to secondary body
    tp = sol.t[end] # Instant of close approach to primary body
    ps = sol.u[1] # State of close approach to secondary body
    pp = sol.u[end] # State of close approach to primary body

    ImpF    = sol.u[1]
    tImpF   = sol.t[1]
    ImpI    = sol.u[end]
    tImpI   = sol.t[end]

    ImpS    = sol.u[2]
    tImpS   = sol.t[2]
    if length(x) == 3
        if length(sol.t) == 3
            ts = sol.t[2]
            ps = sol.u[2]
        end
    else
        if par.stop
            if length(sol.t) == 3
                ts = sol.t[2]
                ps = sol.u[2]
            end
        else
            if length(sol.t) >= 4
                tImp        = x[end]

                indexImp = findfirst( x -> sol.t[x] == tImp, sol.t)
                ImpS  = sol.u[indexImp]
                tImpS   = sol.t[indexImp]

                Mclose = [ (calcDists(sol.u[i],prob_info.x2,prob_info.rs),i) for i in 2:(length(sol.t)-1) if i != indexImp ]
                sort!(Mclose,by = x -> x[1])
                ts = sol.t[Mclose[1][2]]
                ps = sol.u[Mclose[1][2]]
            end
        end
    end
    return ts, tp, ps, pp, ImpF, ImpS, ImpI, tImpF, tImpS, tImpI
end
function calcOuts(x::Vector{T},sol::ODESolution,par::ProblemParameters,prob_info::ProblemInfo,PROB::Tuple{ODEProblem}) where {T<:ForwardDiff.Dual}
    ts = sol.t[1] # Instant of close approach to secondary body
    tp = sol.t[end] # Instant of close approach to primary body
    ps = sol.u[1] # State of close approach to secondary body
    pp = sol.u[end] # State of close approach to primary body

    if length(x) == 3
        if length(sol.t) == 3
            ts = sol.t[2]
            ps = sol.u[2]
        end
    else
        if par.stop
            if length(sol.t) == 3
                ts = sol.t[2]
                ps = sol.u[2]
            end
        else
            if length(sol.t) >= 4
                tImp        = x[end]

                indexImp = findfirst( x -> sol.t[x] == tImp, sol.t)

                Mclose = [ (calcDists(sol.u[i],prob_info.x2,prob_info.rs),i) for i in 2:(length(sol.t)-1) if i != indexImp ]
                sort!(Mclose,by = x -> x[1])
                ts = sol.t[Mclose[1][2]]
                ps = sol.u[Mclose[1][2]]
            end
        end
    end
    return ts, tp, ps, pp
end
function calcOuts(x::Vector{T},sol::ODESolution,par::ProblemParameters,prob_info::ProblemInfo,PROB::Tuple{ODEProblem,ODEProblem}) where {T<:ForwardDiff.Dual}
    ts = sol.t[1] # Instant of close approach to secondary body
    tp = sol.t[end] # Instant of close approach to primary body
    ps = sol.u[1] # State of close approach to secondary body
    pp = sol.u[end] # State of close approach to primary body

    if length(sol.u) > 2
        IMP     = sol.u[2]
    else
        IMP     = sol.u[1]
    end

    if length(x) == 3
        IMP     = sol.u[1]
        if length(sol.t) == 3
            ts = sol.t[2]
            ps = sol.u[2]
        end
    else
        if par.stop
            if length(sol.t) == 3
                ts = sol.t[2]
                ps = sol.u[2]
            end
        else
            if length(sol.t) >= 4
                tImp        = x[end]

                indexImp = findfirst( x -> sol.t[x] == tImp, sol.t)
                IMP  = sol.u[indexImp]

                Mclose = [ (calcDists(sol.u[i],prob_info.x2,prob_info.rs),i) for i in 2:(length(sol.t)-1) if i != indexImp ]
                sort!(Mclose,by = x -> x[1])
                ts = sol.t[Mclose[1][2]]
                ps = sol.u[Mclose[1][2]]
            end
        end
    end
    return ts, tp, ps, pp, IMP
end
"""
Description myObj
"""
function myObj(X::Vector,grad::Vector,par::ProblemParameters)
    if par.stop == 0
        x = X[1:end-1]
    else
        x = X
    end
    #println(x)
    if length(grad) > 0
        if par.stop == 0
            grad[:] = vcat(2x,[0])
        else
            grad[:] = 2x
        end
    end
    return dot(x,x)
end
"""
Description fConstraint
"""
function fConstraint(x::Vector{T}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem}) where {T<:ForwardDiff.Dual}

    ds, dp, ts, tp  = calc_data(x, par, prob_info, PROB)

    n = par.rest[1]
    if      n == 1
        res = [ dp - par.hgoal ]
    elseif n == 2
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9 ]
    elseif n == 3
        res = [ dp - par.hgoal,
                -tp+par.tspan*0.95 ]
    elseif  n == 4
        res = [ dp - par.hgoal,
                -ds + par.hmmin ]
    elseif n == 8
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9,
                -ds + par.hmmin ]
    end
    return  res
end
function fConstraint(x::Vector{T}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem,ODEProblem}) where {T<:ForwardDiff.Dual}

    ds, dp, ts, tp, eRet  = calc_data(x, par, prob_info, PROB)

    n = par.rest[1]
    if      n == 1
        res = [ dp - par.hgoal ]
    elseif n == 2
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9 ]
    elseif n == 3
        res = [ dp - par.hgoal,
                -tp+par.tspan*0.95 ]
    elseif  n == 4
        res = [ dp - par.hgoal,
                -ds + par.hmmin ]
    elseif n == 5
        res = [ dp - par.hgoal,
                eRet-1.1*par.hgoal ]
    elseif n == 6
        res = [ dp - par.hgoal,
                -ds+par.hmmin,
                eRet-1.1*par.hgoal ]
    elseif n == 7
        res = [ dp - par.hgoal,
                -eRet,
                eRet-1.1*par.hgoal ]
    elseif n == 8
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9,
                -ds + par.hmmin ]
    end
    return  res
end
function fConstraint(x::Array{Float64,1}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem})

    ds, dp, ts, tp, eRet, ImpF, ImpS, ImpI, tImpF, tImpS, tImpI, fimp  = calc_data(x, par, prob_info, PROB)

    n = par.rest[1]
    if      n == 1
        res = [ dp - par.hgoal ]
    elseif n == 2
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9 ]
    elseif n == 3
        res = [ dp - par.hgoal,
                -tp+par.tspan*0.95 ]
    elseif  n == 4
        res = [ dp - par.hgoal,
                -ds + par.hmmin ]
    elseif n == 5
        res = [ dp - par.hgoal,
                eRet-1.1*par.hgoal ]
    elseif n == 6
        res = [ dp - par.hgoal,
                -ds+par.hmmin,
                eRet-1.1*par.hgoal ]
    elseif n == 7
        res = [ dp - par.hgoal,
                -eRet,
                eRet-1.1*par.hgoal ]
    elseif n == 8
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9,
                -ds + par.hmmin ]
    end
    return  res
end
"""
Description fConstraint
"""
function fAltitude(x::Vector{T}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem}) where {T<:ForwardDiff.Dual}
    ds, dp, ts, tp  = calc_data(x, par, prob_info, PROB)
    return  dp - par.hgoal
end
function fAltitude(x::Vector{T}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem,ODEProblem}) where {T<:ForwardDiff.Dual}

    ds, dp, ts, tp, eRet  = calc_data(x, par, prob_info, PROB)

    return  dp - par.hgoal
end
"""
Description fRest
"""
function fRest(x::Vector{T}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem}) where {T<:ForwardDiff.Dual}

    ds, dp, ts, tp  = calc_data(x, par, prob_info, PROB)

    n = par.rest[1]
    if n == 2
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9 ]
    elseif n == 3
        res = [ dp - par.hgoal,
                -tp+par.tspan*0.95 ]
    elseif  n == 4
        res = [ dp - par.hgoal,
                -ds + par.hmmin ]
    elseif n == 8
        res = [ dp - par.hgoal,
                -dp +  par.hgoal*0.9,
                -ds + par.hmmin ]
    end
    return  res
end
function fRest(x::Vector{T}, par::ProblemParameters,prob_info::ProblemInfo, PROB::Tuple{ODEProblem,ODEProblem}) where {T<:ForwardDiff.Dual}

    ds, dp, ts, tp, eRet  = calc_data(x, par, prob_info, PROB)

    if n == 2
        res = [ -dp +  par.hgoal*0.9 ]
    elseif n == 3
        res = [ -tp+par.tspan*0.95 ]
    elseif  n == 4
        res = [ -ds + par.hmmin ]
    elseif n == 5
        res = [ eRet-1.1*par.hgoal ]
    elseif n == 6
        res = [ -ds+par.hmmin,
                eRet-1.1*par.hgoal ]
    elseif n == 7
        res = [ -eRet,
                eRet-1.1*par.hgoal ]
    elseif n == 8
        res = [ -dp +  par.hgoal*0.9,
                -ds + par.hmmin ]
    end
    return  res
end
"""
Description myConstraint, myConstAlt, myConstRest
"""
function myConstraint(result::Vector,x::Vector,grad::Matrix, par::ProblemParameters,prob_info::ProblemInfo, PROB, results)
    ForwardDiff.jacobian!(results, (x0::Vector) -> fConstraint(x0, par,prob_info, PROB) , x)
    if length(grad) > 0
        grad[:] = DiffResults.jacobian(results)'
    end
    result[:] = DiffResults.value(results)
end
function myConstAlt(result::Vector,x::Vector,grad::Matrix, par::ProblemParameters,prob_info::ProblemInfo, PROB, results)
    ForwardDiff.gradient!(results, (x0::Vector) -> fAltitude(x0, par,prob_info, PROB) , x)
    if length(grad) > 0
        grad[:] = reshape( DiffResults.gradient(results) ,length(x),1)
    end
    result[:] = [DiffResults.value(results)]
end
function myConstRest(result::Vector,x::Vector,grad::Matrix, par::ProblemParameters,prob_info::ProblemInfo, PROB, results)
    ForwardDiff.jacobian!(results, (x0::Vector) -> fRest(x0, par,prob_info, PROB) , x)
    if length(grad) > 0
        grad[:] = DiffResults.jacobian(results)'
    end
    result[:] = DiffResults.value(results)
end
"""
Description fImp
"""
function fImp(p::Vector{<:Real},prob_info::ProblemInfo) # Computes the value of the collinear impulse to perform in p to circularize (adimensional velocity)
    xp, yp, zp, vx, vy, vz = p
    mu  = prob_info.mu  ::Float64
    mup = prob_info.mup ::Float64
    return sqrt((vx-yp)^2+(vy+xp+mu)^2+vz^2) - sqrt( mup / sqrt( (xp+mu)^2 + yp^2 + zp^2 ) )
end
"""
Description fImp
"""
calcDists(y::Vector{<:Real},x::Float64,r::Float64)       = sqrt((y[1]-x)^2+y[2]^2+y[3]^2)-r
"""
Description plotCC
"""
function plotCC(body::Int, n = 16 ::Int, d=prob_info::ProblemInfo) # Function to plot Earth and Moon
    if body == 1 # plot prymary body
        pos             = -d.mu
        rad             = d.rp
        colorSurface    = :blues
    elseif body == 0 # plot secondary
        pos             = 1-d.mu
        rad             = d.rs
        colorSurface    = :Greys
    end
    u = collect(range(0,stop=2,length=n));
    v = collect(range(0,stop=1,length=n));
    x = rad * cospi.(u) * sinpi.(v)' .+ pos;
    y = rad * sinpi.(u) * sinpi.(v)';
    z = rad * repeat(cospi.(v)',outer=[n, 1])
    surface!(x,y,z, color=colorSurface)
end
"""
Description fsObj
"""
function fsObj(x::Vector,grad::Vector, intF)
    if length(grad) > 0
        grad[:] = ForwardDiff.jacobian(intF, x)'
    end
    return intF(x)[1]
end
"""
Description point
"""
function point(i::Int,orbit::Array{Float64,2}) # Uses interpolation to extract the relevant points
    ORB     = interpolate(orbit, BSpline(Quadratic(Reflect(OnCell()))))
    Y(x)    = ORB(x,2)

    if i == 1 # Point of maximum Z coordinate
        imax = argmax(orbit[:,3])
        if (imax==0) || (imax==length(orbit[:,1]))
            ind = imax
        else
            ind = find_zero(Y, imax, Order2())
        end
    elseif i == 4 # Point of minimum Z coordinate
        imax = argmax(-orbit[:,3])
        if (imax==0) || (imax==length(orbit[:,1]))
            ind = imax
        else
            ind = find_zero(Y, imax, Order2())
        end
    elseif i == 6 # Point of most negative Y coordinate
        imax = argmax(-orbit[:,2])
        if (imax==0) || (imax==length(orbit[:,1]))
            ind = imax
        else
            OPT = Opt(:LD_SLSQP, 1)
            #LN_COBYLA
            #LD_SLSQP
            lower_bounds!(OPT, [1.1])
            upper_bounds!(OPT, [200.9])
            ftol_rel!(OPT::Opt, 1e-14::Float64)
            ftol_abs!(OPT::Opt, 1e-14::Float64)
            xtol_rel!(OPT::Opt, 1e-14::Float64)
            xtol_abs!(OPT::Opt, 1e-14::Float64)

            min_objective!(OPT::Opt, (x,grad) -> fsObj(x,grad,Y))
            (optf,optx,ret) = optimize(OPT, [imax])
            ind = optx[1]
        end
    elseif i == 2 # Point of most positive Y coordinate
        imax = argmax(orbit[:,2])
        if (imax==0) || (imax==length(orbit[:,1]))
            ind = imax
        else
            OPT = Opt(:LD_SLSQP, 1)
            lower_bounds!(OPT, [1.1])
            upper_bounds!(OPT, [200.9])
            ftol_rel!(OPT::Opt, 1e-14::Float64)
            ftol_abs!(OPT::Opt, 1e-14::Float64)
            xtol_rel!(OPT::Opt, 1e-14::Float64)
            xtol_abs!(OPT::Opt, 1e-14::Float64)

            max_objective!(OPT::Opt, (x::Vector,grad::Vector) -> fsObj(x,grad,Y))
            (optf,optx,ret) = optimize(OPT, [imax])
            ind = optx[1]

        end
    elseif (i == 3) || (i == 5) # Points of zero Z coordinate
        Z(x)    = ORB(x,3)
        indexes = [i for i in 1:length(orbit[:,1])-1 if orbit[i,3].*orbit[i+1,3] < 0 ]
        if      i == 5
            iz  = [indexes[i] for i in 1:2 if Y(indexes[i]) < 0 ]
            ind = find_zero(Z, iz[1], Order2())
        elseif  i == 3
            iz  = [indexes[i] for i in 1:2 if Y(indexes[i]) > 0 ]
            ind = find_zero(Z, iz[1], Order2())
        end
    end
    return ORB(ind,1:6), ind
end
"""
Description Pset
"""
function Pset(lbounds::Array{Float64,1},ubounds::Array{Float64,1}, x0::Array{Float64,1}, n = 50 ::Int64)::Array{Array{Float64,1},1} # Creates the starting points of the GlobalSearch like MultiStart
    cuts = Int(floor(n^(1/length(x0)))) # number of cuts in the grid
    C = ones(Int, length(x0))*(cuts + 1)
    p = prod(C); i = 1
    while p > n
        C[i]    -= 1
        p       = prod(C)
        i       += 1
        if i > length(x0)
            i = 1
        end
    end

    for i in 1:length(C)
        if C[i] == 1
            C[i] += 1
        end
    end

    rangeVar = [collect(range(lbounds[i],stop = ubounds[i], length = C[i])) for i in 1:length(x0)]
    if length(x0) == 3
        setOfPoints = [ [i, j, k] for i=rangeVar[1] for j=rangeVar[2] for k=rangeVar[3]  ]
    elseif length(x0) == 4
        setOfPoints = [ [i,j,k,l] for i=rangeVar[1] for j=rangeVar[2] for k=rangeVar[3] for l=rangeVar[4]  ]
    elseif length(x0) == 5
        setOfPoints = [ [i,j,k,l,m] for i=rangeVar[1] for j=rangeVar[2] for k=rangeVar[3] for l=rangeVar[4] for m=rangeVar[5]  ]
    elseif length(x0) == 6
        setOfPoints = [ [i,j,k,l,m,n] for i=rangeVar[1] for j=rangeVar[2] for k=rangeVar[3] for l=rangeVar[4] for m=rangeVar[5] for n=rangeVar[6]  ]
    elseif length(x0) == 7
        setOfPoints = [ [i,j,k,l,m,n,p] for i=rangeVar[1] for j=rangeVar[2] for k=rangeVar[3] for l=rangeVar[4] for m=rangeVar[5] for n=rangeVar[6] for p=rangeVar[7]  ]
    else
        println("Wrong input, x0 dimension greater than 7")
        return [[0.0]]
    end

    while length(setOfPoints) < n-1
        point = ubounds.*rand(length(x0))+lbounds.*rand(length(x0))
        setOfPoints = cat(setOfPoints, [point], dims = 1)
    end
    return cat(setOfPoints, [x0], dims = 1)
end
"""
Description setOptProblem
"""
function setOptProblem(par::ProblemParameters,prob_info::ProblemInfo,PROB)

    if par.stop == 1
        lbounds = -ones(length(par.x0))*par.lim
        ubounds =  ones(length(par.x0))*par.lim
    else
        if par.maxt0 < 0
            lbounds = vcat(-ones(length(par.x0)-1)*par.lim, [par.maxt0] )
            ubounds = vcat( ones(length(par.x0)-1)*par.lim, [-1e-3] )
        else
            lbounds = vcat(-ones(length(par.x0)-1)*par.lim, [-1e-3] )
            ubounds = vcat( ones(length(par.x0)-1)*par.lim, [par.maxt0] )
        end
    end

    # Optimization
    Nopt = length(par.x0) # Number of optimization Variables
    alg = :LD_SLSQP
    OPT = Opt(alg, Nopt)
    lower_bounds!(OPT::Opt, lbounds)
    upper_bounds!(OPT::Opt, ubounds)
    min_objective!(OPT::Opt, (X::Vector,grad::Vector) -> myObj(X,grad,par))
    ftol_rel!(OPT::Opt, par.optTol)
    ftol_abs!(OPT::Opt, par.optTol)
    xtol_rel!(OPT::Opt, par.optTol)
    xtol_abs!(OPT::Opt, par.optTol)
    maxeval!(OPT::Opt, par.neval)

    if par.tRest
        # Create DiffResults object
        results = DiffResults.JacobianResult(ones(par.rest[2]),par.x0)
        inequality_constraint!(OPT, (result::Vector,x::Vector,grad::Matrix) -> myConstraint(result,x,grad, par,prob_info, PROB, results), par.cTol*ones(par.rest[2]) )
    else
        if par.rest[2] > 1
            results = DiffResults.JacobianResult(ones(par.rest[2]-1),par.x0)
            inequality_constraint!(OPT, (result::Vector,x::Vector,grad::Matrix) -> myConstRest(result,x,grad, par,prob_info, PROB, results), par.cTol*ones(par.rest[2]-1) )
        end
        Eresults = DiffResults.GradientResult(par.x0)
        equality_constraint!(OPT, (result::Vector,x::Vector,grad::Matrix) -> myConstAlt(result,x,grad, par,prob_info, PROB, Eresults), par.cTol*ones(1) )
    end

    if par.sType
        return OPT, Pset(lbounds,ubounds,par.x0,par.points)
    else
        return OPT, [par.x0]
    end
end
"""
Description createPROB
"""
function createPROB(par::ProblemParameters, prob_info::ProblemInfo)
    call, cRet          = createCallback(par,prob_info)
    # create ODE problem
    if par.retorno
        prob    = ODEProblem((dy,y,p,t) -> dineq!(dy,y,p,t,prob_info.mu), par.y, par.tspan, [0.0, 0.0, 0.0, 0.0, 1.0], callback = call)
        probRet = ODEProblem((dy,y,p,t) -> dineq!(dy,y,p,t,prob_info.mu), par.y, par.tspan, [0.0, 0.0, 0.0, 0.0, 1.0], callback = cRet)
        PROB    = tuple(prob,probRet)
    else
        prob    = ODEProblem((dy,y,p,t) -> dineq!(dy,y,p,t,prob_info.mu), par.y, par.tspan, [0.0, 0.0, 0.0, 0.0, 1.0], callback = call)
        PROB    = tuple(prob)
    end
    return PROB
end
"""
Description createPROB
"""
function search(par::ProblemParameters, prob_info::ProblemInfo)
    # create ODE problem(s)
    PROB = createPROB(par, prob_info)
    #set optimization problem
    OPT, searchPoints   = setOptProblem(par,prob_info,PROB)
    @show OPT
    # Perform search
    sol = map((x) -> myOptimizer(x,OPT),searchPoints)
    return sol
end
"""
Description createCallback
"""
function createCallback(par::ProblemParameters,prob_info::ProblemInfo)
    lX0     = length(par.x0)
    t2      = par.tspan
    detMin  = par.detMin
    stop    = par.stop
    col     = par.col
    mu      = prob_info.mu

    cRet   = ContinuousCallback((y,t,integrator) -> eConditionRet(y,t,integrator,mu),terminate!, nothing, save_positions = (false,true), rootfind=true) # Fee Return stoping

    if lX0 == 3

        if t2 < 0
            cDist   = ContinuousCallback((y,t,integrator) -> eCondition(y,t,integrator,mu),nothing, terminate!, save_positions = (false,true), rootfind=true)
            mDist   = ContinuousCallback((y,t,integrator) -> mCondition(y,t,integrator,mu),nothing, saveEvent!, save_positions = (false,true), rootfind=true)
        else
            cDist   = ContinuousCallback((y,t,integrator) -> eCondition(y,t,integrator,mu),terminate!,nothing, save_positions = (false,true), rootfind=true)
            mDist   = ContinuousCallback((y,t,integrator) -> mCondition(y,t,integrator,mu),saveEvent!, nothing, save_positions = (false,true), rootfind=true)
        end

        if detMin == 1
            call = CallbackSet(cDist,mDist)
        else
            call = CallbackSet(cDist)
        end

    else

        if t2 < 0
            if stop == 1
                cDist   = ContinuousCallback((y,t,integrator) -> gCondition(y,t,integrator,mu),nothing, (integrator) -> impulse!(integrator,col),   save_positions = (false,true), rootfind=true)
            else
                mDist   = ContinuousCallback((y,t,integrator) -> mCondition(y,t,integrator,mu),nothing, saveEvent!, save_positions = (false,true), rootfind=true)
                eDist   = ContinuousCallback((y,t,integrator) -> eCondition(y,t,integrator,mu),nothing, terminate!, save_positions = (false,true), rootfind=true)
            end
        else
            if stop == 1
                cDist   = ContinuousCallback((y,t,integrator) -> gCondition(y,t,integrator,mu), (integrator) -> impulse!(integrator,col),   nothing, save_positions = (false,true), rootfind=true)
            else
                mDist   = ContinuousCallback((y,t,integrator) -> mCondition(y,t,integrator,mu),saveEvent!, nothing, save_positions = (false,true), rootfind=true)
                eDist   = ContinuousCallback((y,t,integrator) -> eCondition(y,t,integrator,mu),terminate!, nothing, save_positions = (false,true), rootfind=true)
            end
        end

        if stop == 1
            call = CallbackSet(cDist)
        else
            iDist   = DiscreteCallback(iCondition, (integrator) -> impulseInt!(integrator,col),save_positions = (false,true) )
            if detMin == 0
                call = CallbackSet(eDist,iDist)
            else
                call = CallbackSet(eDist,iDist,mDist)
            end
        end
    end
    return call, cRet
end
"""
Description conditions and affect!'s
"""
function gCondition(y,t,integrator,mu::Float64)
    xp, yp, zp, vx, vy, vz = y
    #mu      = integrator.p[1]
    # the last parameter indicates where the minimum is to be looked for
    if integrator.p[end] == 1.0  # the minimum is with the secondary body when p[end] is "true"
        return (xp - 1.0 + mu)*vx +  yp*vy + zp*vz
    else
        return (xp + mu)*vx +  yp*vy + zp*vz
    end
end
function iCondition(y,t,integrator)
    integrator.p[4] == t
end
function eCondition(y::Array{Float64,1},t::Float64,integrator,mu::Float64)  #Float version
    xp, yp, zp, vx, vy, vz = y
    #mu      = integrator.p[1]
    tint    = integrator.p[4]
    if xp > 0.6 || t > tint
        return 0.0
    else
        return (xp + mu)*vx +  yp*vy + zp*vz
    end
end
function eCondition(y,t,integrator,mu::Float64)
    xp, yp, zp, vx, vy, vz = y
    #mu      = integrator.p[1]
    tint    = integrator.p[4]
    if xp > 0.6 || t > tint
        return convert(eltype(y),0.0)
    else
        return (xp + mu)*vx +  yp*vy + zp*vz
    end
end
function eConditionRet(y::Array{Float64,1},t::Float64,integrator,mu::Float64) # Float version
    xp, yp, zp, vx, vy, vz = y
    #mu = integrator.p[1]
    if xp > 0.6
        return 0.0
    else
        return (xp + mu)*vx +  yp*vy + zp*vz
    end
end
function eConditionRet(y,t,integrator,mu::Float64)
    xp, yp, zp, vx, vy, vz = y
    #mu = integrator.p[1]
    if xp > 0.6
        return convert(eltype(y),0.0)
    else
        return (xp + mu)*vx +  yp*vy + zp*vz
    end
end
function mCondition(y,t,integrator,mu::Float64)
    xp, yp, zp, vx, vy, vz = y
    #mu = integrator.p[1]
    cond = (xp - 1.0 + mu)*vx +  yp*vy + zp*vz
    return cond
end
function impulse!(integrator,col::Bool)
    imp = integrator.p[1:3]
    if integrator.p[end] == 1.0
        if col
            integrator.u[4:6] = integrator.u[4:6]*(1.0-imp[1]/norm(integrator.u[4:6]))
        else
            integrator.u[4:6] = integrator.u[4:6] - imp
        end
        integrator.p[end] = convert(eltype(integrator.u),0.0) # change the minimum search to the other body
    else
        terminate!(integrator)
    end
end
function saveEvent!(integrator)
    nothing
    #println("Moon encounter at t: ",integrator.t)
end
function impulseInt!(integrator,col::Bool)
    imp = integrator.p[1:3]
    if col
        integrator.u[4:6] = integrator.u[4:6]*(1.0-imp[1]/norm(integrator.u[4:6]))
    else
        integrator.u[4:6] = integrator.u[4:6] - imp
    end
end
"""
Description myOptimizer
"""
function myOptimizer(x::Array{Float64,1}, OPT)
    f, xf, res = optimize(OPT,x)
    println("The optimizer reached the value: ",f)
    println(res)
    return f, xf, res, x
end
"""
Description myOptimizer
"""
function filterGlobalSol(Gsol, par::ProblemParameters,prob_info::ProblemInfo)
    # create callbacks
    PROB = createPROB(par, prob_info)
    #
    #Eresults = DiffResults.GradientResult(par.x0)
    #for i in 1:length(Gsol)
    #    ForwardDiff.gradient!(Eresults, (x0::Vector) -> fAltitude(x0, par,prob_info, PROB) , Gsol[i][2])
    #    @show abs.( [DiffResults.value(Eresults)] ), Gsol[i][1]

        #@show abs.( fConstraint( Gsol[i][2], par,prob_info, PROB)  ), Gsol[i][1]
    #end
    filterSol = [ Gsol[i] for i in 1:length(Gsol) if abs.( fConstraint( Gsol[i][2], par,prob_info, PROB) ) < par.cTol*ones(par.rest[2]) ]

    sfSol = [filterSol[1]] # filtrar por diferencias
    for i in 2:length(filterSol)
        #println(filterSol[i][1])
        dif = 1
        for j in 1:length(sfSol)
            difF = abs(sfSol[j][1] - filterSol[i][1])/maximum([1,sfSol[j][1]])
            difX = norm(sfSol[j][2] - filterSol[i][2])/maximum([1,norm(sfSol[j][2])])
            if (difF < 0.1) && (difX < 0.1)
                dif = 0
            end
        end
        if dif == 1
            #println("new sol")
            sfSol = vcat(sfSol, filterSol[i])
        end
    end
    println("Number of solutions: ",length(sfSol))
    sfSol = sort!(sfSol, by = x -> x[1])
    return filterSol, sfSol
end
"""
Description plotSols
"""
function plotSols(sfSol, par::ProblemParameters, prob_info::ProblemInfo, orbit::Matrix{<:Real},guiP::Bool)
    # create ODE problem
    PROB = createPROB(par, prob_info)

    pyplot()
    ploty = plot(orbit[:,1],orbit[:,2],orbit[:,3], yflip=true, leg=true, lab="Halo Orbit", aspect_ratio=:equal)
    scatter!([-prob_info.mu],[0],[0],c=:blue,label="Earth",markersize=2)
    scatter!([1-prob_info.mu],[0],[0],c=:grey,label="Moon",markersize=1)

    y0 = par.y
    prob = PROB[1]
    for i in 1:length(sfSol)

        x = sfSol[i][2]

        ds, dp, ts, tp, eRet, ImpF, ImpS, ImpI, tImpF, tImpS, tImpI, fimp = calc_data(x, par, prob_info, PROB)

        if par.stop
            if length(x) == 3
                intProb           = remake(prob, u0 = y0 - vcat(zeros(3),x[1:3]) )
            elseif length(x) == 4
                intProb           = remake(prob, u0 = y0 - vcat(zeros(3),x[1:3]) , p = [x[4], 0.0, 0.0, 0.0, par.detMin] )
            else
                intProb           = remake(prob, u0 = y0 - vcat(zeros(3),x[1:3]) , p = [x[4], x[5], x[6], 0.0, par.detMin] )
            end
            tstop = [0.0]
        else
            if length(x) == 5
                intProb           = remake(prob, u0 = y0 - vcat(zeros(3),x[1:3]) , p = [x[4], 0.0, 0.0, x[5], par.detMin] )
            elseif length(x) == 7
                intProb           = remake(prob, u0 = y0 - vcat(zeros(3),x[1:3]) , p = [x[4], x[5], x[6], x[7], par.detMin] )
            end
            tstop = [x[end]]
        end
        if par.odeAlg
            sol         = solve(intProb, DP5(), abstol= par.atol, reltol= par.rtol, tstops =  tstop , dense=true, save_everystep=true)
        else
            sol         = solve(intProb, Vern7(), abstol= par.atol, reltol= par.rtol, tstops = tstop , dense=true, save_everystep=true)
        end
        if length(x) == 3
            deltav = @sprintf("%.2f (m/s)",(fimp + norm(x))*prob_info.r12*prob_info.w)
        elseif par.stop
            deltav = @sprintf("%.2f (m/s)",(fimp + norm(x[1:3]) + norm(x[4:end]))*prob_info.r12*prob_info.w)
        else
            deltav = @sprintf("%.2f (m/s)",(fimp + norm(x[1:3]) + norm(x[4:end-1]))*prob_info.r12*prob_info.w)
        end
        plot!(sol, vars=(1,2,3), leg=true, label=string(deltav), xlim=:auto)
    end
    #xlims!(-0.05,1.1)
    if guiP
        gui()
        return ploty
    end
end


using ChainRulesCore

ChainRulesCore.debug_mode() = true

#####
##### `eigen`
#####

@Zygote.adjoint function LinearAlgebra.eigen(A::LinearAlgebra.RealHermSymComplexHerm)
    dU = eigen(A)
    return dU, function (Δ)
      d, U = dU
      d̄, Ū = Δ
      if Ū === nothing
        P = Diagonal(d̄)
      else
        F = inv.(d' .- d)
        P = F .* (U' * Ū)
        if d̄ === nothing
          P[diagind(P)] .= 0
        else
          P[diagind(P)] = d̄
        end
      end
      return (U * P * U',)
    end
  end

Zygote.refresh()

# function rrule(::typeof(eigen), X::AbstractMatrix{<:Real})
#     F = eigen(X)
#     function eigen_pullback(Ȳ::Composite)
#         # `getproperty` on `Composite`s ensures we have no thunks.
#         ∂X = eigen_rev(F, Ȳ.values, Ȳ.vectors)
#         return (NO_FIELDS, ∂X)
#     end
#     return F, eigen_pullback
# end

# function rrule(::typeof(getproperty), F::T, x::Symbol) where T <: Eigen
#     function getproperty_eigen_pullback(Ȳ)
#         C = Composite{T}
#         ∂F = if x === :values
#             C(values=Ȳ,)
#         elseif x === :vectors
#             C(vectors=Ȳ,)
#         # elseif x === :V
#         #     C(V=Ȳ,)
#         # elseif x === :Vt
#         #     # TODO: https://github.com/JuliaDiff/ChainRules.jl/issues/106
#         #     throw(ArgumentError("Vt is unsupported; use V and transpose the result"))
#         end
#         return NO_FIELDS, ∂F, DoesNotExist()
#     end
#     return getproperty(F, x), getproperty_eigen_pullback
# end

# # When not `Zero`s expect `λ̄ ::AbstractVector, Ū::AbstractMatrix`
# function eigen_rev(res::Eigen, λ̄ , Ū)
#     λ = res.values
#     U = res.vectors

#     k = length(λ)
#     T = eltype(λ)
#     F = T[i == j ? 0 : inv(@inbounds λ[j] - λ[i]) for i = 1:k, j = 1:k]
#     D = Diagonal(λ)
#     D̄ = λ̄  isa AbstractZero ? λ̄  : Diagonal(λ̄ )
#     Ut = U'
#     Ā = conj.(U) * ( D̄ + ( F .* ( Ut * Ū ) ) ) * Ut
#     return Ā
# end

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

# Manifolds.jl Forward mode eigen AD by mateuszbaran
# https://github.com/JuliaManifolds/Manifolds.jl/pull/27

##
using LinearAlgebra
using ForwardDiff

import LinearAlgebra.eigen
function eigen(A::StridedMatrix{<:ForwardDiff.Dual})
    A_values = map(d -> d.value, A)
    A_values_eig = eigen(A_values)
    UinvAU = A_values_eig.vectors \ A * A_values_eig.vectors
    vals_diff = diag(UinvAU)
    F = similar(A_values)
    for i ∈ axes(A_values, 1), j ∈ axes(A_values, 2)
        if i == j
            F[i, j] = 0
        else
            F[i, j] = inv(A_values_eig.values[j] - A_values_eig.values[i])
        end
    end
    vectors_diff = A_values_eig.vectors * (F .* UinvAU)
    for i ∈ eachindex(vectors_diff)
        vectors_diff[i] = ForwardDiff.Dual{ForwardDiff.tagtype(vectors_diff[i])}(A_values_eig.vectors[i], vectors_diff[i].partials)
    end
    return Eigen(vals_diff, vectors_diff)
end

function f(A)
    A_eig = eigen(A)
    return A_eig.vectors * Diagonal(A_eig.values) / A_eig.vectors
end
##
ForwardDiff.derivative(t -> f([t+1 t; t t+1]), 1.0)
##
ForwardDiff.jacobian(x -> f(x), [1.0 2.0; 4.0 1.0])

##

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

### Legume code

# def vjp_maker_eigsh(ans, x, **kwargs):
#     """Gradient for eigenvalues and vectors of a hermitian matrix."""
#     numeig = kwargs['k']
#     N = x.shape[-1]
#     w, v = ans              # Eigenvalues, eigenvectors.
#     vc = np.conj(v)

#     def vjp(g):
#         wg, vg = g          # Gradient w.r.t. eigenvalues, eigenvectors.
#         w_repeated = np.repeat(w[..., np.newaxis], numeig, axis=-1)

#         # Eigenvalue part
#         vjp_temp = np.dot(vc * wg[..., np.newaxis, :], T(v))

#         # Add eigenvector part only if non-zero backward signal is present.
#         # This can avoid NaN results for degenerate cases if the function
#         # depends on the eigenvalues only.
#         if np.any(vg):
#             off_diag = np.ones((numeig, numeig)) - np.eye(numeig)
#             F = off_diag / (T(w_repeated) - w_repeated + np.eye(numeig))
#             vjp_temp += np.dot(np.dot(vc, F * np.dot(T(v), vg)), T(v))

#         return vjp_temp

#     return vjp

# def vjp_maker_eigsh(ans, mat, **kwargs):
#     """Steven Johnson method extended to a Hermitian matrix
#     https://math.mit.edu/~stevenj/18.336/adjoint.pdf
#     """
#     numeig = kwargs['k']
#     N = mat.shape[0]

#     def vjp(g):
#         vjp_temp = np.zeros_like(mat)
#         for iv in range(numeig):
#             a = ans[0][iv]
#             v = ans[1][:, iv]
#             vc = np.conj(v)
#             ag = g[0][iv]
#             vg = g[1][:, iv]

#             # Eigenvalue part
#             vjp_temp += ag*np.outer(vc, v)

#             # Add eigenvector part only if non-zero backward signal is present.
#             # This can avoid NaN results for degenerate cases if the function
#             # depends on the eigenvalues only.
#             if np.any(vg):
#                 # Projection operator on space orthogonal to v
#                 P = np.eye(N, N) - np.outer(vc, v)
#                 Amat = T(mat - a*np.eye(N, N))
#                 b = P.dot(vg)

#                 # Initial guess orthogonal to v
#                 v0 = P.dot(np.random.randn(N))

#                 # Find a solution lambda_0 using conjugate gradient
#                 (l0, _) = sp.linalg.cg(Amat, b, x0=v0, atol=0)
#                 # Project to correct for round-off errors
#                 l0 = P.dot(l0)

#                 vjp_temp -= np.outer(l0, v)

#         return vjp_temp


using PureEOS
using Base.Test

# write your own tests here
println("Viscosity test")
@test waterVisc(298.15,998/18.015268) ≈ 889.735100 atol=0.001
@test waterVisc(298.15,1200/18.015268) ≈ 1437.649467 atol=0.001
@test waterVisc(373.15,1000/18.015268) ≈ 307.883622 atol=0.001
@test waterVisc(433.15,1/18.015268) ≈ 14.538324 atol=0.001
@test waterVisc(433.15,1000/18.015268) ≈ 217.685358 atol=0.001
@test waterVisc(873.15,1/18.015268) ≈ 32.619287 atol=0.001
@test waterVisc(873.15,100/18.015268) ≈ 35.802262 atol=0.001
@test waterVisc(873.15,600/18.015268) ≈ 77.430195 atol=0.001
@test waterVisc(1173.15,1/18.015268) ≈ 44.217245 atol=0.001
@test waterVisc(1173.15,100/18.015268) ≈ 47.640433 atol=0.001
@test waterVisc(1173.15,400/18.015268) ≈ 64.154608 atol=0.001
println("thermalconductivity test")
@test waterTherm(647.35,1/18.015268) ≈ 51.9298924 atol=0.001
@test waterTherm(647.35,122/18.015268) ≈ 130.922885 atol=0.1
@test waterTherm(647.35,222/18.015268) ≈ 367.787459 atol=20
@test waterTherm(647.35,272/18.015268) ≈ 757.959776 atol=40
@test waterTherm(647.35,322/18.015268) ≈ 1443.75556 atol=200
@test waterTherm(647.35,372/18.015268) ≈ 650.319402 atol=40
@test waterTherm(647.35,422/18.015268) ≈ 448.883487 atol=20
@test waterTherm(647.35,750/18.015268) ≈ 600.961346 atol=0.001

println("ethanol thermalconductivity test")
#@test ethanolTherm(300.0,850.0/46.0) ≈ 209.68 atol=2
#@test ethanolTherm(400.0,2.0/46.0) ≈ 26.108 atol=1
#@test ethanolTherm(400.0,690.0/46.0) ≈ 149.21 atol=1.5
#@test ethanolTherm(500.0,10.0/46.0) ≈ 39.594 atol=1
#@test ethanolTherm(500.0,10.0/46.0) ≈ 40.755 atol=1

#Data from NIST


function test_grad_desc()

    # Vectorize
    for i in 1:10
        target = [10.0 * rand() + 5.0, 10.0 * rand() + 5.0]
        @show target
        x0 = [0.0, 0.0]
        rexps = rand(2:10, 2)
        @show rexps
        maxΔx = 1.0 ./ rexps
        minΔx = 0.0 ./ rexps
        x1 = x0 .+ maxΔx .* 0.1
        gdth = 1e-5
        maxiter = 50000

        function up_fun(gdmodel)
            xs = SimTools.gd_value(gdmodel)
            xs .^ rexps
        end

        gdmodel = SimTools.grad_desc_vec(up_fun;
            target, x0, x1, minΔx, maxΔx, gdth, maxiter,
        )
        @show gdmodel.ϵi
        @show gdmodel.iter
        @show up_fun(gdmodel)
        @test all(isapprox.(up_fun(gdmodel), target; atol = 1e-3))
        println()
    end

    # Scalar
    @testset begin
        for i in 1:10
            target = 10.0 * rand() + 5.0
            @show target
            x0 = 0.0
            rexp = rand(1:10)
            @show rexp
            maxΔx = 1.0 / rexp
            x1 = x0 + maxΔx * 0.1
            maxiter = 1e4
            gdth = 1e-5

            function up_fun(gdmodel)
                x = SimTools.gd_value(gdmodel)
                x ^ rexp
            end
            gdmodel = SimTools.grad_desc(up_fun;
                target, x0, x1, maxΔx, gdth, maxiter
            )
            @show gdmodel.ϵi
            @show gdmodel.iter
            @show up_fun(gdmodel)
            @test all(isapprox.(up_fun(gdmodel), target; atol = 1e-3))
            println()
        end
    end

end
test_grad_desc()


using DifferentialDynamicsModels
using StaticArrays
using Test

@maintain_type struct RoboState{T} <: FieldVector{2,T}
    x::T
    y::T
end
x0 = RoboState(1., 1.)
xf = RoboState(4., 5.)
@test x0 + xf isa RoboState

SI = SingleIntegratorDynamics{2}()
@test state_dim(SI) == 2
@test control_dim(SI) == 2

bvp = SteeringBVP(SI, Time(), constraints=BoundedControlNorm())
@test issymmetric(bvp)
sol = bvp(x0, xf)
@test sol == bvp(x0, xf, Inf)
@test sol.cost ≈ 5
@test sol.controls isa StepControl{2,Float64,SVector{2,Float64}}
@test sol.controls.t ≈ 5
@test sol.controls.u ≈ SVector(.6, .8)

ctrl = sol.controls
@test propagate(SI, x0, ctrl) isa RoboState
@test propagate(SI, x0, ctrl) ≈ xf
@test propagate(SI, x0, fill(ctrl, 2)) isa RoboState
@test propagate(SI, x0, fill(ctrl, 2)) ≈ RoboState(7., 9.)

@test propagate(SI, x0, ctrl, -1:1:duration(ctrl)+1) ≈ [x0, [x0 + i*(xf-x0)/5 for i in 0:5]..., xf]
@test waypoints(SI, x0, fill(ctrl, 2), 1.0) ≈ waypoints(SI, x0, fill(ctrl, 2), 11)

c1 = StepControl(.5, rand(SVector{2,Float64}))
c2 = StepControl(.5, rand(SVector{2,Float64}))
@test instantaneous_control([c1, c2], .25) == c1.u
@test instantaneous_control([c1, c2], .75) == c2.u
@test propagate(SI, x0, [c1, c2]) == x0 + 0.5*c1.u + 0.5*c2.u

rc = RampControl(.5, rand(SVector{2,Float64}), rand(SVector{2,Float64}))
@test propagate(SI, x0, rc) ≈ DifferentialDynamicsModels.propagate_ode(SI, x0, rc)


import Blosc: compress, decompress

function parallel_density_grid{T}(C::Matrix{T}, width::T, d::Int, σ::T, ϵ::T = 1e-4, num_workers::Int = nworkers())
    num_points = size(C, 2)

    num_points_per_worker = fill(fld(num_points, num_workers), num_workers)
    num_points_per_worker[1:mod(num_points, num_workers)] += 1

    start_indices = unshift!(cumsum(num_points_per_worker), 0) .+ 1
    pop!(start_indices)
    params = map(zip(start_indices, num_points_per_worker)) do r
        C[:, range(r[1], r[2])], width, d, σ, ϵ
    end

    #ρs_compressed = pmap(params, pids = collect(take(workers(), num_workers))) do p
    ρs_compressed = pmap(params) do p
        compress(density_grid(p...), level = 9, shuffle = true, itemsize = sizeof(T))
    end

    ρ = mapreduce(.+, zeros(T, d^3), ρs_compressed) do ρ_compressed
        decompress(T, ρ_compressed)
    end

    return reshape(ρ, d, d, d)
end

function density_grid{T}(C::Matrix{T}, width::T, d::Int, σ::T, ϵ::T= 1e-4)
    hw = width / 2
    ρ = zeros(T, d, d, d)

    num_points = size(C, 2)

    r_max_sq = max(0, -σ^2 * log(ϵ))
    r_max = sqrt(r_max_sq)

    b = floor(Int, d * r_max / width) + 1

    g = collect(linspace(-hw, hw, d))

    @inbounds for i in 1:num_points
        ix = floor(Int, d * (C[1, i] + hw) / width) + 1
        iy = floor(Int, d * (C[2, i] + hw) / width) + 1
        iz = floor(Int, d * (C[3, i] + hw) / width) + 1

        x_min = max(1, ix - b)
        y_min = max(1, iy - b)
        z_min = max(1, iz - b)

        x_max = min(d, ix + b)
        y_max = min(d, iy + b)
        z_max = min(d, iz + b)

        for jx in x_min:x_max, jy in y_min:y_max, jz in z_min:z_max
            r_sq = (C[1, i] - g[jx])^2 + (C[2, i] - g[jy])^2 + (C[3, i] - g[jz])^2

            if r_sq < r_max_sq
                ρ[jx, jy, jz] += exp(-r_sq / σ^2)
            end
        end
    end

    return ρ
end

function density_grid{T}(x::Array{T, 3}, y::Array{T, 3}, z::Array{T, 3}, C::Matrix{T}, σ::T, ϵ::T = 1e-4)
    D = size(x)
    ρ = zeros(T, D)

    num_points = size(C, 2)

    R = CartesianRange(size(ρ))
    I_min, I_max = first(R), last(R)

    x_spread = x[I_max] - x[I_min]
    y_spread = y[I_max] - y[I_min]
    z_spread = z[I_max] - z[I_min]

    r_max_sq = max(0, -σ^2 * log(ϵ))
    r_max = sqrt(r_max_sq)

    dx = floor(Int, D[1] * r_max / x_spread) + 1
    dy = floor(Int, D[2] * r_max / y_spread) + 1
    dz = floor(Int, D[3] * r_max / z_spread) + 1

    B = CartesianIndex((dx, dy, dz))

    @inbounds for i in 1:num_points
        ix = floor(Int, D[1] * (C[1, i] - x[I_min]) / x_spread) + 1
        iy = floor(Int, D[2] * (C[2, i] - y[I_min]) / y_spread) + 1
        iz = floor(Int, D[3] * (C[3, i] - z[I_min]) / z_spread) + 1
        I = CartesianIndex((ix, iy, iz))

        for J in CartesianRange(max(I_min, I - B), min(I_max, I + B))
            r_sq = (C[1, i] - x[J])^2 + (C[2, i] - y[J])^2 + (C[3, i] - z[J])^2
            if r_sq < r_max_sq
                ρ[J] += exp(-r_sq / σ^2)
            end
        end
    end

    return ρ
end

function density_grid_stats{T}(width::T, d::Int, σ::T, ϵ::T)
    r_max = sqrt(max(0, -σ^2 * log(ϵ)))
    b = floor(Int, r_max / (width / d)) + 1
    num_blocks = (2b + 1)^3

    println("Max radius: $(r_max)")
    println("Blocks per dimension: $b")
    println("Total blocks: $(num_blocks)")
    flush(STDOUT)
end


# This implements the Adaptive Iterative Tresholding algorithm for
# compressed sensing / regularized regression. This implementation
# is based on the paper:
#  Yu Wang, Jinshan Zeng, Zhimin Peng, Xiangyu Chang, Zongben Xu,
#  "On Linear Convergence of Adaptively Iterative Thresholding
#  Algorithms for Compressed Sensing", arxiv, 2014, http://arxiv.org/abs/1408.6890
#
function adaptive_iterative_tresholding(A, b; kws...)
  M, N, k, xprev, x, An, acolnorms, sAnt, f, maxiterations, tolerance = unpack_ait_parameters(A, b; kws...)

  # Since we need one iteration before we can judge convergence
  # we check for convergence last in the loop.
  t = 0
  while true
    t += 1
    z = xprev .+ sAnt * (b .- An * xprev)
    zp = sortperm(abs(z), rev = true)
    tau = abs(z[zp[k+1]])

    # Treshold x
    for i in 1:N
      ip = zp[i]
      if i <= k
        x[ip] = f(tau, z[ip])
      else
        x[ip] = 0.0
      end
    end

    # Check convergence
    diff = norm(x .- xprev)
    if t >= maxiterations || diff < tolerance
      return (x ./ acolnorms, t)
    end

    # Switch their meaning => we need not allocate new memory in each iteration
    xprev, x = x, xprev
  end
end

function unpack_ait_parameters(A, b;
  k = -1,             # Sparsity level, this should be set as an upper bound on the number of non-zero coefficients of x. If set below the actual sparseness level the algorithm might not find right solution.
  s = 1.0,            # Step length, should be ~ in range [0.4, 1.3] but 1.0 seem to be generally good
  x0 = nothing,       # A starting point, will be randomized if not given
  f::Function = scad, # Defining function used in tresholding, by default SCAD
  tolerance = 1e-6,   # Convergence level, i.e. terminate algorithm if norm change is less than this between iterations
  maxIterations = 1e3 # Maximum number of iterations allowed. Often it converges in ~20-30 iterations if ever.
  )

  # Check and then initialize all parameters. Note that we use N for the num of covariates and
  # M for the number of cases. This is in line with the paper above but it is more common to use
  # N = M and P = N.
  M, N = size(A)
  @assert length(b) == M
  k = (k == -1) ? ifloor(sqrt(N)) : k
  @assert 1 <= k < N
  xprev = (x0 == nothing) ? zeros(N) : copy(x0)
  x = zeros(Float64, N)

  An, acolnorms = normalize_columns(A)
  sAnt = s * An'

  return (M, N, k, xprev, x, An, acolnorms, sAnt, f, maxIterations, tolerance)
end

function normalize_columns(A)
  acolnorms = map((j) -> norm(A[:,j]), 1:size(A, 2))
  An = broadcast(/, A, acolnorms')
  return (An, acolnorms)
end

# Inner AIT iteration loop. Assumes An has been normalized.
function ait_iteration_loop!(An, b, k, N, M, sAnt, tolerance, maxiterations, xprev, trainindices, Ntrain, f)
  t = 0
  x = zeros(Float64, N)

  # For now, later we will loop over the trainindices
  btrain = b[trainindices]
  Antrain = An[trainindices, :]
  sAnttrain = sAnt[:, trainindices]

  while true
    t += 1
    z = xprev .+ sAnttrain * (btrain .- Antrain * xprev)
    zp = sortperm(abs(z), rev = true)
    tau = abs(z[zp[k+1]])

    # Treshold x
    for i in 1:N
      ip = zp[i]
      if i <= k
        x[ip]::Float64 = f(tau::Float64, z[ip]::Float64)
      else
        x[ip]::Float64 = 0.0
      end
    end

    diff = norm(x .- xprev)
    if t >= maxiterations || diff < tolerance
      return (x, t)
    end

    xprev, x = x, xprev
  end
end

function ait2(A, b; kws...)
  M, N, k, xprev, x, An, acolnorms, sAnt, f, maxiterations, tolerance = unpack_ait_parameters(A, b; kws...)
  x, its = ait_iteration_loop!(An, b, k, N, M, sAnt, tolerance, maxiterations, xprev, collect(1:M), M, f)
  return (x ./ acolnorms, its)
end

function scad_penalty_tresholding(a, tau, u)
  absu = abs(u)
  if absu > a * tau
    u
  elseif absu > 2 * tau
    ((a-1)*u - sign(u)*a*tau)/(a-2)
  elseif absu > tau
    absu - sign(u)*tau
  else
    0.0
  end
end

scad(tau, u) = scad_penalty_tresholding(3.7, tau, u)
hard(tau, u) = abs(u) > tau ? u : 0.0

splitindices(n, ratio = 0.10) = begin
  randindices = shuffle(collect(1:n))
  numtest = ifloor(ratio*n)
  test = randindices[1:numtest]
  train = randindices[(numtest+1):end]
  return (train, test)
end

# Run AIt for multiple k-values from 2*sqrt(P) down to 1
# and then select the one with lowest MSE on a TestRatio
# holdout test set.
function ait_best_k(A, b, TestRatio = 0.20; kws...)
  N, P = size(A)

  # Split in train and test sets
  trainidxs, testidxs = splitindices(N, TestRatio)
  Atrain = A[trainidxs,:]
  btrain = b[trainidxs]
  Atest = A[testidxs,:]
  btest = b[testidxs]

  maxP = ifloor(2*sqrt(P))
  if maxP >= P
    maxP = P - 1
  end

  local bestcoefs = bestits = bestmse = Inf

  mses = Dict{Int64,Float64}()

  for k in maxP:-1:1
    coefs, its = adaptive_iterative_tresholding(Atrain, btrain; k = k, kws...)
    mses[k] = mse(btest, Atest * coefs)
    if mses[k] < bestmse
      bestcoefs = coefs
      bestits = its
      bestmse = mses[k]
    end
  end

  return bestcoefs, bestits, mses
end

# using PyPlot
# x = linspace(-5.0, 5.0, 200)
# y = map((x) -> scad_penalty_tresholding(3.7, 1.0, x), x)
# plot(x, y, color="red", linewidth=2.0, linestyle="--")

mse(y,yhat) = mean( (y .- yhat).^2 )

# N-fold cross-validation to select the sparseness level for AIT.
function nfoldCV_ait(A, b)
  M, N = size(A)

  num_ks = 5
  ks = ifloor(exp(log(N-1) / num_ks * reverse(1:num_ks)))
  xs = zeros(num_ks, N)
  x0 = zeros(N)
  its = zeros(Int64, num_ks)

  for i in 1:num_ks
    x0, its[i] = adaptive_iterative_tresholding(A, b; k=k, x0 = x0)
    xs[i,:] = x
  end
end

function ait_mse(A, b, k)
  betas, its = adaptive_iterative_tresholding(A, b; k = k)
  return (betas, mse(b, A * betas))
end

function adaptive_k_ait(A, b; kws...)
  N, P = size(A)

  mses = Dict{Int64, Float64}()
  betas = Dict{Int64, Vector{Float64}}()
  mse_for_k(k) = begin
    if !haskey(mses, k)
      betas[k], mses[k] = ait_mse(A, b, k)
    end
    mses[k]
  end

  binsearch(lowk, highk) = begin
    mselow = mse_for_k(lowk)
    msehigh = mse_for_k(highk)
    if mselow > 1e3 * msehigh
      if highk == lowk + 1
        return (highk, msehigh)
      end
      midk = lowk + ifloor((highk-lowk)/2)
      k, mseb = binsearch(lowk, midk)
      if k == -1
        return binsearch(midk, highk)
      else
        return (k, mseb)
      end
    else
      return (-1, nothing)
    end
  end

  k, mse = binsearch(1, ifloor(min(P/2, 2*sqrt(P))))

  return (k, mse, betas[k], mses)
end

function test_ait_on_random_problem(N, P, kstar)
  A = randn(N, P)
  nzw = map((i) -> rand(10.0:100.0) * (rand() < 0.5 ? 1.0 : (-1.0)), 1:kstar)
  betas = shuffle([nzw, zeros(P-kstar)])
  b = A * betas

  k, mse, betasp, mses = adaptive_k_ait(A, b)

  show(mses)
  println()

  return (N, P, kstar, k, mse, coeffquality(betas, betasp))
end

#for kstar in [3, 4, 5, 6]
#  for P in [10, 20, 50, 100]
#    println(test_ait_on_random_problem(100, P, kstar), "\n")
#  end
#end

# N, P = (rand(20:100), 20)
# kstar = ifloor(sqrt(P)) - 1
# A = randn(N, P)
# nzw = rand(1.0:10.0) * rand(-1.0:0.01:1.0, kstar)
# betas = shuffle([nzw, zeros(P-kstar)])
# b = A * betas
# k = -1.0
# s = 1.0
# x0 = nothing
# f = hard
# tolerance = 1e-5
# maxIterations = 1e4
#
# # Ensure compiled
# adaptive_iterative_tresholding(A[1:10,:], b[1:10])
# ait2(A[1:10,:], b[1:10])
#
# tic()
# betasp, its = adaptive_iterative_tresholding(A, b)
# t1 = toq()
# tic()
# betasp2, its2 = ait2(A, b)
# t2 = toq()
# nzwp = betasp[abs(betasp) .> 0.0]
# bhat = X' * betasp
# bhat2 = X' * betasp2
#
#
# show( [betas betasp betasp2] )
# println("\nN = ", N, ", P = ", P)
# println("MSE1: ", mse(b, bhat))
# println("MSE2: ", mse(b, bhat2))
# println("b1: ", coeffquality(betas, betasp))
# println("b2: ", coeffquality(betas, betasp2))
# println("Timediff: ", t1-t2)


include("basic_binary_tree.jl")
"""
array-based binary search tree
left tree values < root value < right tree values
"""
abstract type AbstractBinarySearchTree_arr<:AbstractBinaryTree_arr end
function Base.findfirst(value::T,tree::Tree,place::Int=tree.root)where {T,Tree<:AbstractBinarySearchTree_arr}
    if tree.val[place]==value
        return place
    elseif tree.val[place]>value
        if tree.lch[place]==0 return 0 end
        return findfirst(value,tree,tree.lch[place])
    else
        if tree.rch[place]==0 return 0 end
        return findfirst(value,tree,tree.rch[place])
    end
end


# Integrate with k-means better
function dp_means{T <: Real}(data::Matrix{T}, lambda::Real, opts::Options)
    @defaults opts tolerance => 10e-8
    @defaults opts max_iter => 1_000

    # Keep track of the number of data points
    n, p = size(data)

    # Initially assign all data points to a single cluster
    k = 1
    assignments = ones(Int64, n)
    centers = mean(data, 1)
    distances = Array(Float64, k)

    # Keep working until convergence or exceeding bound on iterations
    converged, iter = false, 0

    # Keep track of error in system
    ss_old, ss_new = Inf, Inf

    while !converged && iter < max_iter
        iter += 1

        # Update cluster assignments
        for i in 1:n
            for cluster in 1:k
                d = 0.0
                for j in 1:p
                    d += (data[i, j] - centers[cluster, j])^2
                end
                distances[cluster] = d
            end
            d_min, ind_min = findmin(distances)
            if d_min > lambda
                k += 1
                assignments[i] = k
                centers = vcat(centers, data[i, :])
                push!(distances, 0.0)
            else
                assignments[i] = ind_min
            end
        end

        # Update cluster centers for non-empty clusters
        for cluster in 1:k
            n_obs = 0
            for i in 1:n
                if assignments[i] == cluster
                    n_obs += 1
                end
            end
            if n_obs > 0
                for j in 1:p
                    centers[cluster, j] = 0.0
                end
            end
            for i in 1:n
                if assignments[i] == cluster
                    for j in 1:p
                        centers[cluster, j] += data[i, j]
                    end
                end
            end
            for j in 1:p
                centers[cluster, j] /= n_obs
            end
        end

        # Update convergence criterion
        ss_old, ss_new = ss_new, 0.0
        for i in 1:n
            for j in 1:p
                ss_new += (data[i, j] - centers[assignments[i], j])^2
            end
        end
        ss_change = ss_old - ss_new
        if !isnan(ss_change) && ss_change < tolerance
            converged = true
        end
    end

    results = DPMeansOutput(assignments,
                            centers,
                            k,
                            iter,
                            ss_new,
                            converged)

    return results
end

function dp_means{T <: Real}(data::Matrix{T}, lambda::Real)
    dp_means(data, lambda, Options())
end


# Copyright 2019 Tobias Frilling
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
    neighbourhoodOffsets(radius::Float64)

Return every (unique) `CartesianIndex{2}` in a `radius`-sized circle around `index`.
"""
function neighbourhood(index::CartesianIndex{2}, radius::Float64,
                       settings::Settings = defaultSettings)
    offsets = neighbourhoodOffsets(radius)
    neighbours = (i -> i + index).(offsets)
    return unique(correctCoords(neighbours, settings))
end

"""
    neighbourhoodOffsets(radius::Float64)

Return every `CartesianIndex{2}` in a `radius`-sized circle around (0,0).
"""
function neighbourhoodOffsets(radius::Float64)
    r_int = ceil(Int, radius)
    offsets = (CartesianIndices((r_int,r_int) .* 2 .+ 1) .-
               CartesianIndex(r_int+1, r_int+1))[:]
    filter!(i -> sum(i.I.^2) <= radius^2, offsets)
    return offsets
end

"""
    directNeighbours(ind::CartesianIndex{2})

Return the eight direct neighbours of index `ind`.
"""
function directNeighbours(ind::CartesianIndex{2}, settings::Settings = defaultSettings)
    neighbours = (CartesianIndices((3,3)) .- (CartesianIndex(2,2) - ind))[:]
    # remove the index itself
    deleteat!(neighbours, 5)
    return correctCoords(neighbours, settings)
end

"""
    correctCoords(coords::AbstractVector{CartesianIndex{2}})

Remove (if the lattice is a plane) or wrap (if the lattice is a toroid)
coordinates outside of the lattice size.
"""
function correctCoords(coords::AbstractVector{CartesianIndex{2}},
                       settings::Settings = defaultSettings)
    if settings.toroid
        return wrapCoordsOnToroid(coords, settings)
    else
        return removeCoordsOutsideBounds(coords, settings)
    end
end

"""
    removeCoordsOutsideBounds(coords::AbstractVector{CartesianIndex{2}})

Remove coordinates outside of the lattice size.
"""
function removeCoordsOutsideBounds(coords::AbstractVector{CartesianIndex{2}},
                                   settings::Settings = defaultSettings)
    filter(i -> all((1,1) .<= i.I .<= settings.latticeSize), coords)
end

"""
    wrapCoordsOnToroid(coords::AbstractVector{CartesianIndex{2}},

Wrap coordinates outside of the lattice size around the torus.
"""
function wrapCoordsOnToroid(coords::AbstractVector{CartesianIndex{2}},
                            settings::Settings = defaultSettings)
    # like mod, but returns y instead of 0
    mod_replace_zero(x,y) = if (m = mod(x,y)) == 0 y else m end
    map(i -> CartesianIndex(mod_replace_zero.(i.I,settings.latticeSize)), coords)
end

"""
    latticeDistance(a::CartesianIndex{2}, b::CartesianIndex{2},

Euclidean distance on the lattice.
"""
function latticeDistance(a::CartesianIndex{2}, b::CartesianIndex{2},
                         settings::Settings = defaultSettings)
    diff = abs.((a - b).I)
    if settings.toroid
        diff = min.(diff, settings.latticeSize .- diff)
    end
    return sqrt(sum(diff))
end

@inline neighbourhoodKernel(kernel::Symbol, radius::Float64) =
    neighbourhoodKernel(Val(kernel), radius)

"""
    neighbourhoodKernel(::Val{:cone}, radius::Float64)

Returns a linear neighbourhood decay function.
"""
function neighbourhoodKernel(::Val{:cone}, radius::Float64)
    (dist::Float64) -> (radius - dist)/radius
end

"""
    neighbourhoodKernel(::Val{:gauss}, radius::Float64)

Returns a gaussian neighbourhood decay function.
"""
function neighbourhoodKernel(::Val{:gauss}, radius::Float64)
    (dist::Float64) -> exp(-(dist/radius)^2/2)
end

"""
    neighbourhoodKernel(::Val{:mexhat}, radius::Float64)

Returns a "Mexican hat" (Ricker wavelet) neighbourhood decay function.
"""
function neighbourhoodKernel(::Val{:mexhat}, radius::Float64)
    (dist::Float64) -> begin
        square = (dist/radius)^2
        (1 - square) * exp(-square/2)
    end
end


using Flux, Flux.Tracker
using Distributions

struct VP
  b0
  b1
  log_sig

  VP() = new(param(randn(2)), param(randn(2)), param(randn(2)))
end

function lpdf_logx(dist, logx)
  return logpdf(dist, exp(logx)) + logx
end

function loglike(y, x, b0, b1, sig, N)
  return sum(logpdf.(Normal.(b0 .+ b1 .* x, sig), y)) * (N / length(y))
end

function log_p(b0, b1, log_sig)
  return logpdf(Normal(0, 1), b0) + logpdf(Normal(0, 1), b1) + lpdf_logx(Gamma(1, 1), log_sig)
end

lpdf_vp(p, x) = logpdf(Normal(p[1], exp(p[2])), x)
function log_q(b0, b1, log_sig, vp)
  return lpdf_vp(vp.b0, b0) + lpdf_vp(vp.b1, b1) + lpdf_vp(vp.log_sig, log_sig)
end

rsample(p) = randn() * exp(p[2]) + p[1]
function elbo(y, x, vp, N)
  b0 = rsample(vp.b0)
  b1 = rsample(vp.b1)
  log_sig = rsample(vp.log_sig)
  sig = exp(log_sig)

  return loglike(y, x, b0, b1, sig, N) + log_p(b0, b1, log_sig) - log_q(b0, b1, log_sig, vp)
end


#
# asian-asciiplot.jl
#
using ASCIIPlots;

S0  = 100;      # Spot price
K   = 102;      # Strike price
r   = 0.05;     # Risk free rate
q   = 0.0;      # Dividend yield
v   = 0.2;      # Volatility
tma = 0.25;     # Time to maturity
T = 100;        # Number of time steps
dt = tma/T;     # Time increment
S = zeros(Float64,T);
S[1] = S0;
dW = randn(T)*sqrt(dt)
[ S[t] = S[t-1] * (1 + (r - q - 0.5*v*v)*dt + v*dW[t] + 0.5*v*v*dW[t]*dW[t]) for t=2:T ]
x = linspace(1,T);
scatterplot(x,S,sym='*')


#=
Author: Stephan Scholz
Year: 2021

This file contains a Simulation of the Heat equation
Boundary conditions: heat transfer and heat radiation (nonlinear)
Dimension: 1D
Spatial Approximation: Finite Differences
Temporal Approximation: Forward-Euler method
=#

using LinearAlgebra, Plots, SparseArrays, DifferentialEquations

# Define the constants for the PDE
const λ = 45.0          # Conductivity
const ρ = 7800.0;       # Density
const cap = 480.0;      # Specific heat capacitivity
const α = λ/(cap * ρ)   # Diffusivity
const L = 0.1           # Length of rod

#Boundary conditions
const h = 10            # Conduction coefficient
const ϵ = 0.6;          # Emissivity
const sb = 5.67*10^(-8) # Stefan-Boltzmann constant
const k = ϵ * sb        # Radiation coefficient

# Discretization
const dx = L/100;                     # Spatial sampling: o--o--o--o--o--o
const N = round(Int, (L / dx + 1));   # Number of elements

const Tf = 300.0    # Final time
const dt = 10^(-2); # Time step width
const t_len = round(Int, (Tf/dt + 1))

#Laplace operator
const M = spdiagm(-1 => ones(N-1), 0 => -2*ones(N), 1 => ones(N-1))
M[1,2] = 2
M[end,end-1] = 2

const Mx = zeros(N)
const C = spzeros(N,2) # Boundary Conditions: Output / Outflow

C[1,1] = 1	  # Left side
C[end,2] = 1  # Right side

const q = α*dt/(dx^2);
if q > 0.5
  error("Numerical stability is NOT guaranteed! STOP program!")
end

# Temperatures
const ϑ₀   = 400.0; # Initial temperature
const ϑamb = 273.0; # Ambient temperature

# Define initial temperature
function init_temperature(ϑoffset, scale, dx, num_points)
	a = scale
	b = ϑoffset

	L = (num_points-1)*dx
	x = 0 : dx : L

	ϑinit = a*sin.(x*2*pi/L) .+ b

	return ϑinit
end

# Define the discretized PDE as an ODE function
function heat_eq(dθ,θ,p,t)

  λ = p[1]
  ρ = p[2]
  cap = p[3]

  mul!(Mx,M,θ[1:end]) # means: Mx = M * Θ[1:end]

  # Now: zero Neumann boundary conditions
  bc = zeros(2)

  # Later: Natural boundary condition: heat conduction + radiation
  bc[1] =  -1*(h*(θ[1] - ϑamb) + k*(θ[1]^4 - ϑamb^4));
  bc[2] =  -1*(h*(θ[end] - ϑamb) + k*(θ[end]^4 - ϑamb^4));

  flux_out =(2*dx/λ)*C*bc                     # Temperature input and output

  dθ .= λ/(ρ*cap) * 1/(dx^2) * (Mx + flux_out)   # Integration of temperatures

end

θinit = init_temperature(ϑ₀, 30.0, dx, N)

tspan = (0.0, Tf) # Time span
param = [λ, ρ, cap] # Parameter that shall be learned

prob = ODEProblem( heat_eq, θinit,tspan, param )
sol = solve(prob,Euler(),dt=dt,progress=true,save_everystep=true,save_start=true) # use 'save_everystep' only for small systems with few timesteps

xspan = 0 : dx : L

# Temperature at the beginning (t=0) vs at the end (t=Tf)
# plot(xspan, [sol[:,1], sol[:,end]], label=["t=0" "t=Tf"], xlabel="Position x", ylabel="Temperature")

# Temperature at left and right side
# plot(sol.t,[sol[1,:], sol[end,:]], label=["Left" "Right"], xlabel="Time t", ylabel="Temperature", legend=:bottomright)

# 2D Heatmap plot: xaxis = time ; yaxis = position
# heatmap(sol.t, xspan, sol[:,:], xlabel="Time t", ylabel="Position x")


"""
Output structure
"""
mutable struct Solution{T1<:AbstractFace, T2<:Vector{Float64}, T3<:Float64}
    F::T1
    Fnew::T1
    c::T2
    cnew::T2
    feas::T3
    feasnew::T3
    r::T2
    function Solution(M::AbstractLinearOp, b::Vector{Float64}, A::AbstractAtomicSet)
        z = M'*b
        F = face(A, z)
        Fnew = deepcopy(F)
        c = zeros( size(M*F, 2))
        cnew = deepcopy(c)
        feas = norm(b)^2/2
        feasnew = deepcopy(feas)
        r = deepcopy(b)
        new{AbstractFace, Vector{Float64},Float64}(F, Fnew, c, cnew, feas, feasnew, r)
    end
end

function update!(sol::Solution)
    sol.F = deepcopy(sol.Fnew)
    sol.c = deepcopy(sol.cnew)
    sol.feas = sol.feasnew
end

function constructPrimal(sol::Solution)
    if typeof(sol.F) <: SumFace
        rs = [rank(face) for face in sol.F.faces]
        cs = BlockArray(sol.c, rs)
        xs = [sol.F.faces[i]*cs[Block(i)] for i in 1:length(sol.F.faces)]
        return xs
    else
        return sol.F * sol.c
    end
end


using SkewDist
using Base.Test

function rand_pd(d::Int)
    A = randn(d,d)
    A'A
end

function rand_cor(d::Int)
    A = randn(d,d)
    Σ = A'A
    cov2cor(Σ)
end

include("test_MvSkewNormal.jl")
include("test_MvSkewTDist.jl")
include("test_operators.jl")
include("test_fit_skew.jl")
include("test_skewtdist.jl")


using StaticArrays
struct Particles{T,V <: AbstractVector{SVector{3,T}}}
    npar::Int
    charge::T
    mass::T
    positions::V
    momentums::V
    self_efields::V
    self_bfields::V
end

# general interface
function Particles(;pos, mom, charge=-1.0, mass=1.0)
    return Particles(pos, mom; charge=charge, mass=mass)
end

function Particles(pos::Vector{SVector{3,T}}, mom::Vector{SVector{3,T}}; charge=-1.0, mass=1.0) where {T}
    @assert length(pos) == length(mom)
    npar = length(pos)
    return Particles(npar, charge, mass, pos, mom, fill(SVector(0.0, 0.0, 0.0), npar), fill(SVector(0.0, 0.0, 0.0), npar))
end

function Particles(pos::Matrix{T}, mom::Matrix{T}; charge=-1.0, mass=1.0) where {T}
    @assert size(pos, 1) == 3
    @assert size(mom, 1) == 3
    @assert size(pos, 2) == size(mom, 2)
    npar = size(pos, 2)
    pos_new = reinterpret(reshape, SVector{3,T}, pos)
    mom_new = reinterpret(reshape, SVector{3,T}, mom)
    return Particles(npar, charge, mass, pos_new, mom_new, reinterpret(reshape, SVector{3,T}, zeros(3, npar)), reinterpret(reshape, SVector{3,T}, zeros(3, npar)))
end

function update_self_field!(beam::Particles; lambda)
    q = beam.charge
    @inbounds Threads.@threads for i in 1:beam.npar
        beam.self_efields[i] = SVector(0.0, 0.0, 0.0)
        beam.self_bfields[i] = SVector(0.0, 0.0, 0.0)
        xi = beam.positions[i]
        for j in 1:beam.npar
            xj = beam.positions[j]
            pj = beam.momentums[j]
            xij = xi - xj
            kernel = xij / sqrt(dot(xij, xij) + dot(pj, xij)^2 + eps())^3
            beam.self_efields[i] += 2.8179403699772166e-15 * q / lambda * sqrt(1.0 + dot(pj, pj)) * kernel
            beam.self_bfields[i] += 2.8179403699772166e-15 * q / lambda * cross(pj, kernel)
        end
    end
end

function output(beam::Particles, fname)
    pos = beam.positions
    mom = beam.momentums
    open(fname, "w") do io
        writedlm(io, pos)
    end
end


using LightGraphs, DifferentialEquations, Plots, LinearAlgebra

lil(a) = vcat(a, circshift(reverse(a),1))

function circulant_kura(a::Vector) # a is half vec
	n = 2length(a)
    g = SimpleGraph(reshape(vcat(map(x->circshift(lil(a), x), 0:(n-1))...), (n, n)))
    dgs = degree(g)
    @assert all(x->x==dgs[1], dgs)
    println("μ: $(dgs[1]/(n-1))")
    # println("fake connectivity: ≈$((round(ne / binomial(nv, 2), digits=2)))")
    g
end

# shitty wrapper
function kura_prob(u0, tspan, p)::ODEProblem
    ODEProblem(kuramoto!, u0, tspan, p)
end

function kuramoto!(du, u, p, t)
    n = length(u)
    A, ω = p # A is adjacency_matrix, and ω is
    u .+= ω * t
    for i in 1:n
        du[i] = sum(A[i, j] * sin(u[j] - u[i]) for j in 1:n)
    end
end

function anim_kura(u0, g)
    θn = θo = u0
    # p = scatter(θo)
    # p = plot(20)
    p = scatter(u0)
    anim = @animate for _ in 1:100
       θn = kuramoto(g, θo)
       println("$(θn - θo)")
       push!(p, 20, θn)
       θo = θn
    end
end

function demo(g)

    u0 = rand(nv(g))
    tspan = (0., 10.)
    p = (adjacency_matrix(g), 10.)
    ODEProblem(kuramoto!, u0, tspan, p)

end


@testset "setup_test.jl" begin

    fov = LaserLab.Fov(1.0mm,1.0mm)
    @test fov.a ≈ π * (fov.d/2.)^2
    @test fov.v ≈ fov.a * fov.z
    @test fov.a ≈ π * (fov.d/2.)^2

    luv = LaserLab.CLaser(375.0nm, 1.0μW)
    @test luv.P/μW ≈ 1
    @test luv.λ/nm ≈ 375.0

    obj = LaserLab.Objective("Generic", 0.5, 40.0)
    @test obj.NA ≈ 0.5
    @test obj.M ≈ 40.0

    obj2 = LaserLab.Objective("Generic", 5.0mm, 5.0mm, 40.0)
    @test obj2.NA ≈ 0.5
    @test obj2.M ≈ 40.0
    @test obj2.f/mm ≈ 5.0
    @test obj2.d/mm ≈ 5.0

    epl375 = LaserLab.PulsedLaser(375.0nm, 140.0mW, 1.0MHz, 65.0ps)
    @test epl375.P == epl375.Pk * epl375.w * epl375.f

    lmu_40x_nuv = LaserLab.Objective("LMU-40X-NUV", 5.3mm, 5.1mm, 40.0)
    @test lmu_40x_nuv.NA  ≈ lmu_40x_nuv.d / (2.0*lmu_40x_nuv.f)
    @test LaserLab.lmu_40x_nuv_transmission(0.500μm) ==LaserLab.lmu_40x_nuv_transmission(500nm)

    @test isapprox(LaserLab.photon_energy(500.0nm)/eV, hnu_eV(500.0), rtol=0.01)
    @test uconvert(J, LaserLab.delivered_energy(luv, 1*s)) ≈ 1.0μJ
    @test LaserLab.n_photons(luv) ≈ LaserLab.n_photons(375.0nm, 1.0μW)
    @test LaserLab.n_photons(luv) ≈ LaserLab.n_photons(375.0nm, 1.0μW)
    @test LaserLab.n_photons_int(luv, 1.0s) ≈ LaserLab.n_photons(luv)/Hz
    @test LaserLab.photon_density(375.0nm, 1.0μW, fov.a) ≈ LaserLab.photon_density(luv, fov)
end


# Sparsity-promoting least squares migration with source subsampling using the linearized Bregman algorithm
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: December 2016
#

export linbreg_lsrtm

function linbreg_lsrtm(J::TimeModeling.joJacobian, dD::TimeModeling.joData; iterations=10, lambda=[], shots_per_iteration=1, snapshot=[])

# Sparsity promoting LSRTM using the linearized Bregman method

	# Setup Curvelet transform
	C = joCurvelet2D(J.model.n[1], J.model.n[2]; zero_finest=true, DDT=Float32, RDT=Float64)

	# Main loop
	x = zeros(Float32,prod(J.model.n))
	z = zeros(Float32,prod(J.model.n))
	for j=1:iterations
		println("Iteration ", j, " of ", iterations)

		# Subsample shots
		shotIdx = randperm(J.info.nsrc)[1:shots_per_iteration]

		# Subsample observed data
		dDsub = subsample(dD, shotIdx)

		# Subsample Jacobian
		Jsub = subsample(J, shotIdx)

		# Residual
		dDpred = Jsub*x
		residual = dDpred - dDsub

		# Gradient
		gradient = Jsub'*residual

		# Step size
		step = norm(residual)^2/norm(gradient)^2

		# Automatic thresholding
		if isempty(lambda) && j==1
			lambda = convert(Float32,0.1*maximum(abs.(step*C*gradient)))
		end

		# Update dual variable
		z = z - step*gradient;

		# Update primal variable
		x = C'*softThresholding(C*z, lambda)

		# Save snapshots
		if ~isempty(snapshot)
			save(join([snapshot, "_", string(j), ".jld"]), "z", reshape(z,J.model.n), "x", reshape(x,J.model.n))
		end
	end

	return x
end


@testset "acquisition functions" begin
    for ac in [ProbabilityOfImprovement(), ExpectedImprovement(),
               UpperConfidenceBound(), ThompsonSamplingSimple(), MutualInformation()]
        model = GPE(rand(3, 4), rand(4), MeanZero(), SEIso(0., 0.))
        acfunc = BayesianOptimization.acquisitionfunction(ac, model)
        x = rand(3, 2)
        acvector = acfunc(x)
        @test length(acvector) == 2
        if typeof(ac) != ThompsonSamplingSimple
            @test acvector[1] == acfunc(x[:, 1])
        end
    end
end


module _XSpecEn
export XSpecEn
using FFTW: fft
using DSP: conv
    """
        XSpec, BandEn = XSpecEn(Sig)

    Returns the cross-spectral entropy estimate (`XSpec`) of the full cross-
    spectrum and the within-band entropy (`BandEn`) estimated between the data
    sequences contained in `Sig` using the default  parameters:
    N-point FFT = 2*len(`Sig`) + 1, normalised band edge frequencies = [0 1],
    logarithm = base 2, normalisation = w.r.t # of spectrum/band frequency
    values.

        XSpec, BandEn = XSpecEn(Sig::AbstractArray{T,2} where T<:Real; N::Int=1 + (2*size(Sig,1)), Freqs::Tuple{Real,Real}=(0,1), Logx::Real=exp(1), Norm::Bool=true)

    Returns the cross-spectral entropy (`XSpec`) and the within-band entropy
    (`BandEn`) estimate between the data sequences contained in `Sig` using the
    following specified 'keyword' arguments:

    # Arguments:
    `N`     - Resolution of spectrum (N-point FFT), an integer > 1    \n
    `Freqs` - Normalised band edge frequencies, a scalar in range [0 1]
              where 1 corresponds to the Nyquist frequency (Fs/2).
              **Note: When no band frequencies are entered, BandEn == SpecEn** \n
    `Logx`  - Logarithm base, a positive scalar     [default: base 2]
              ** enter 0 for natural log**  \n
    `Norm`  - Normalisation of `XSpec` value:
              [false]  no normalisation.
              [true]   normalises w.r.t # of spectrum/band frequency values [default] \n

    For more info, see the EntropyHub guide

    # See also `SpecEn`, `fft`, `XDistEn`, `periodogram`, `XSampEn`, `XApEn`

    # References:
        [1]  Matthew W. Flood,
            "XSpecEn - EntropyHub Project"
            (2021) https://github.com/MattWillFlood/EntropyHub

    """
    function XSpecEn(Sig::AbstractArray{T,2} where T<:Real; N::Int=1 + (2*size(Sig,1)),
        Freqs::Tuple{Real,Real}=(0,1), Logx::Real=exp(1), Norm::Bool=true)

    (size(Sig,2) > size(Sig,1)) ? Sig = transpose(Sig) : nothing

    (size(Sig,1) > 10 && size(Sig,2)==2) ? nothing :  error("Sig:   must be a 2-columns matrix")
    (N > 1) ? nothing :  error("N:     must be an integer > 1")
    (0<=Freqs[1]<1 && 0<Freqs[2]<=1 && Freqs[1]<Freqs[2]) ? nothing :
        error("Freq:    must be a two element tuple with values in range [0 1].
                The values must be in increasing order.")
    (Logx>0) ? nothing : error("Logx:     must be a positive number > 0")

    S1 = Sig[:,1]; S2 = Sig[:,2]
    Freqs = collect(Freqs)
    Fx = Int(ceil(N/2))
    Freqs = Int.(round.(Freqs.*Fx))
    Freqs[Freqs.==0] .= 1

    if Freqs[1] > Freqs[2]
        error("Lower band frequency must come first.")
    elseif Freqs[2]-Freqs[1]<1
        error("Spectrum resoution too low to determine bandwidth.")
    elseif minimum(Freqs)<0 || maximum(Freqs)>Fx
        error("Freqs must be normalized w.r.t sampling frequency [0 1].")
    end

    Temp = conv(S1,S2)
    N <= size(Temp,1) ? Temp = Temp[1:N] : Temp = vcat(Temp,zeros(N-size(Temp,1)))

    Pt = abs.(fft(Temp));
    Pxx = Pt[1:Fx]/sum(Pt[1:Fx]);
    XSpec = -transpose(Pxx)*log.(Logx, Pxx)
    Pband = (Pxx[Freqs[1]:Freqs[2]])/sum(Pxx[Freqs[1]:Freqs[2]]);
    BandEn = -transpose(Pband)*log.(Logx, Pband)

    if Norm
        XSpec = XSpec/log(Logx, Fx)
        BandEn = BandEn/log(Logx, Freqs[2]-Freqs[1]+1)
    end

    return XSpec, BandEn
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


function qualityMeasure(xN::Array{Int64,1}, yN::Array{Int64,1},
                        xU::Array{Int64,1}, yU::Array{Int64,1}
                       )

    # ---- Set the reference point
    xRef = xN[end] + 1
    yRef = yN[1] + 1

    # ---- Measure the area corresponding to the set of non-dominated points Y_N
    areaN = 0
    for i = 1:length(xN)-1
        areaN = areaN + (xN[i+1]-xN[i]) * (yRef - yN[i])
    end
    areaN = areaN + (xRef-xN[end]) * (yRef - yN[end])

    # ---- Measure the (restricted) area corresponding to the bound set U
    areaU = 0
    for i = 1:length(xU)-1
        areaU = areaU + ( min(xRef,xU[i+1]) - min(xRef,xU[i]) ) * (yRef - min(yRef,yU[i]) )
    end
    areaU = areaU + ( xRef - min(xRef,xU[end]) ) * (yRef - min(yRef,yU[end]) )

    # ---- Computer the ratio
    result = areaU/areaN
    return result

end


function transformcoords(a,b,c,A,B,C,X)

R = [a[:] b[:] c[:]]*inv([A[:] B[:] C[:]]);

d = size(X,1)
n = size(X,2)
x = R*reshape(X, (d, n));

return x, R

end


using ApproxFun, BandedMatrices, BlockBandedMatrices, SparseArrays, LinearAlgebra, LazyArrays, BlockArrays, MultivariateOrthogonalPolynomials, Test
using SparseVolterraExamples

#############################
## Various tests for integration Qy and expansion Ey operator consistency.
#############################
@testset "Testing that E_x from P01 works correctly" begin
    g = Fun(x -> exp(x+3x^2), Jacobi(0,1,0..1));
    N = ncoefficients(g);
    G = Fun(JacobiTriangle(), triExgenP01(N) * g.coefficients)
    @test G(0.1,0.2) ≈ g(0.1);
    #########################################################################
    f = Fun(x -> exp(x+3x^2), Legendre(0..1));
    N = ncoefficients(f);
    C = (I : Legendre(0..1) → Jacobi(0,1,0..1))[1:N,1:N];
    F = Fun(JacobiTriangle(), triExgenP01(N) * (C * f.coefficients));
    @test F(0.1,0.2) ≈ f(0.1);
end

@testset "Testing that Q and E_y from P01 work correctly" begin
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Jacobi(0,1,0..1));
    Ey=triEygenP01(length(f.coefficients));
    F = Fun(JacobiTriangle(), Ey*f.coefficients);
    @test F(0.1,0.2) ≈ f(0.2);
    @test F(0.1,0.3) ≈ f(0.3);
    #########################################################################
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Jacobi(0,1,0..1));
    N = ncoefficients(f);
    # option 1: use direct ApproxFun + Q generator. This does not use E.
    @test triNoKernelDirect(f,0.9) ≈ sum(Fun(x -> f(x), 0..0.9));
    # option 2: generate Q and E separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP01(length(f.coefficients));
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients)(0.1) ≈
                sum(Fun(x -> f(x), 0..0.9));
    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> f(y), JacobiTriangle()).coefficients,length(Ey*f.coefficients)) ≈ Ey*f.coefficients;
    #########################################################################
    # again for a different function from 0 to 0.5, expected result here is approx. 0.0917462
    g = Fun(y -> sin(8*y)+y^3-tan(y), Jacobi(0,1,0..1));
    N = ncoefficients(g);

    # option 1: use direct ApproxFun + Q generator. This does not use E.
    triNoKernelDirect(g,0.5);

    # option 2: generate Q and E separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP01(N);
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*g.coefficients)(0.5) ≈
                sum(Fun(x -> g(x), 0..0.5));

    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> g(y), JacobiTriangle()).coefficients,length(Ey*g.coefficients)) ≈ Ey*g.coefficients;
end

@testset "Testing that Q and E_y from P00 work correctly" begin
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Legendre(0..1));
    Ey=triEygenP00(length(f.coefficients));
    F = Fun(JacobiTriangle(), Ey*f.coefficients);
    @test F(0.1,0.2) ≈ f(0.2);
    @test F(0.1,0.3) ≈ f(0.3);
    #########################################################################
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Legendre(0..1));
    N = ncoefficients(f);
    # option 1: use direct ApproxFun + Q generator. This does not use R.
    @test triNoKernelDirect(f,0.9) ≈ sum(Fun(x -> f(x), 0..0.9));
    # option 2: generate Q and E separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP00(length(f.coefficients));
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients)(0.1) ≈
                sum(Fun(x -> f(x), 0..0.9));
    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> f(y), JacobiTriangle()).coefficients,length(Ey*f.coefficients)) ≈ Ey*f.coefficients;
    #########################################################################
    # again for a different function from 0 to 0.5, expected result here is approx. 0.0917462
    g = Fun(y -> sin(8*y)+y^3-tan(y), Legendre(0..1));
    N = ncoefficients(g);

    # option 1: use direct ApproxFun + Q generator. This does not use R.
    triNoKernelDirect(g,0.5);

    # option 2: generate Q and R separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP00(N);
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*g.coefficients)(0.5) ≈
                sum(Fun(x -> g(x), 0..0.5));

    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> g(y), JacobiTriangle()).coefficients,length(Ey*g.coefficients)) ≈ Ey*g.coefficients;
end

#############################
## Various tests for the interactions of Qy,Ey with the Jacobi operators
#############################

@testset "Testing that Q*E*f integrates correctly" begin
    S = Jacobi(0,1,0..1);
    f = Fun(x -> 1, S);
    N = length(f.coefficients);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
    #########################################################################
    S = Jacobi(0,1,0..1);
    f = Fun(x -> x, S);
    N = length(f.coefficients);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
    #########################################################################
    S = Jacobi(0,1,0..1);
    f = Fun(x -> 1,);
    xf = Fun(x -> x, S);
    N = length(f.coefficients);
    y = Fun(identity, S);
    Jy = Multiplication(y, S);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    QExf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*Jy[Block.(1:N), Block.(1:N)]*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
    #########################################################################
    S = Jacobi(0,1,0..1);
    f = Fun(x -> 1+exp(3*x^2), S);
    xf = Fun(x -> x*(1+exp(3*x^2)),S);
    N = length(f.coefficients);
    y = Fun(identity, S);
    Jy = Multiplication(y, S);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    QExf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*Jy[Block.(1:N), Block.(1:N)]*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
end

@testset "Testing if expected commutation relations hold" begin
    # we test if the direct Q*EyP01 diagonal matrix is computed correctly
    N=50;
    @test triQEygenP01(N) ≈ triQgen(N)*triEygenP01(N);
    #########################################################################
    # we test expected commutation relations of R and Y on the triangle
    # once with a function and one without function
    N = 50;
    S = Jacobi(0,1,0..1);
    T = JacobiTriangle();
    f = Fun(x->2+x+exp(x^5+1), S);
    E = sparse(triEygenP01(N));
    y = Fun(identity, S);
    J = sparse(Multiplication(y, S)[Block.(1:N), Block.(1:N)]);
    x, y = Fun(identity, T);
    Y = sparse(Multiplication(y, T)[Block.(1:N), Block.(1:N)]);
    @test Y*E ≈ E*J;
    @test Y*E*pad(f.coefficients,N) ≈ E*(J*pad(f.coefficients,N));
    #########################################################################
    # we test expected commutation relations of Q and X
    N = 50;
    S = Jacobi(0,1,0..1);
    T = JacobiTriangle();
    Q = sparse(triQgen(N));
    x = Fun(identity, S);
    J = sparse(Multiplication(x, S)[Block.(1:N), Block.(1:N)]);
    x, y = Fun(identity, T);
    X = sparse(Multiplication(x, T)[Block.(1:N), Block.(1:N)]);
    @test Q*X ≈ J*Q
end

@testset "Selective testing of monomial kernel basis consistency" begin
    k = 81;
    N = 30;
    basisvec = triJQEJbasislexP01(k,N,false);
    basisgrid = triJQEJbasisgridgenP01(k,k,N);
    # first couple
    @test basisvec[1] ≈ basisgrid[1][1];
    @test basisvec[2] ≈ basisgrid[1][2];
    @test basisvec[3] ≈ basisgrid[2][1];
    @test basisvec[4] ≈ basisgrid[2][2];
    # square number sequence elements
    @test basisvec[9] ≈ basisgrid[3][3];
    @test basisvec[16] ≈ basisgrid[4][4];
    @test basisvec[25] ≈ basisgrid[5][5];
    @test basisvec[36] ≈ basisgrid[6][6];
    @test basisvec[81] ≈ basisgrid[9][9];
    # some pseudo random picks
    @test basisvec[12] ≈ basisgrid[3][4];
    @test basisvec[27] ≈ basisgrid[2][6];
    @test basisvec[80] ≈ basisgrid[9][1];
end


using MittagLefflerFunctions
import MittagLefflerFunctions.Quadrature.MLQuad
using RationalApproximations
using PyPlot
using Printf

G(t) = (1-t)/(1+t)
big_α_vals = collect(1:5) ./ big(5)
α_vals = collect(1:5) ./ 5
β = big(1.0)
#big_Eα = [ MLQuad(α, β, 20, :hyperbola) for α in big_α_vals ]
Eα = [ MLQuad(α, 1.0, 14, :hyperbola) for α in α_vals ]
fα = Vector{Function}(undef, 5)
for col = 1:5
#    fα[col] = ( t -> convert(Float64, (big_Eα[col](-G(big(t))))) )
    fα[col] = ( t -> Eα[col](-G(t)) )
end
t = range(-1, 1, length=800)
max_m = 10
tol = 1e-10
max_iterations = 5

rows = 10
cols = length(α_vals)
errors = Array{Float64}(undef, rows, cols)
δ = similar(errors)
for col = 1:cols
    α = α_vals[col]
    println("α = ", α)
    F = fα[col].(t)
    r, err = aaa_v2(F, t, max_m, tol)
    M = length(r)
    for row = 1:rows
        if row ≤ M
            m = row
            resid = F - r[m].(t)
            idx = locate_extrema_idx(resid, m)
            if length(idx) ≠ 2m
                errors[row,col] = NaN
                δ[row,col] = NaN
            else
                supp_pt = r[m].supp_pt
                opt_r, zmin, zmax = minimax(fα[col], t[idx], supp_pt,
                                    (-1.0,1.0), max_iterations)
                errors[row,col] = zmax[end]
                δ[row,col] = zmax[end] - zmin[end]
            end
        else
            errors[row,col] = NaN
            δ[row,col] = NaN
        end
    end
end

@printf("%5s ", "m")
for col = 1:cols
    @printf("&    α = %3.1f", α_vals[col])
end
@printf("\n\n")

for row = 1:rows
    @printf("%5d ", row)
    for col = 1:cols
        @printf("& %10.2e", errors[row,col])
    end
    @printf("\\\\\n")
end