function KS_solve_Emin_sd( pw::PWGrid, V_ionic, Focc, Nstates::Int;
                         psi0=nothing, Potentials0 = nothing, E_NN = nothing,
                         α_t = 3e-5, NiterMax=1000, verbose=false )
    Ns = pw.Ns
    Npoints = prod(Ns)
    Ngwx = pw.gvecw.Ngwx

    if psi0 == nothing
        srand(2222)
        psi = randn(Ngwx,Nstates) + im*randn(Ngwx,Nstates)
        psi = ortho_gram_schmidt(psi)
    else
        psi = copy(psi0)
    end

    if Potentials0 == nothing
        Potentials = PotentialsT( V_ionic, zeros(Npoints), zeros(Npoints) )
        rho = calc_rho( pw, Focc, psi )
        Potentials.Hartree = real( G_to_R( Ns, Poisson_solve(pw, rho) ) )
        Potentials.XC = excVWN( rho ) + rho .* excpVWN( rho )
    else
        Potentials = PotentialsT( Potentials0.Ionic,
                                  Potentials0.Hartree,
                                  Potentials0.XC )
    end

    Etot = 0.0
    Etot_old = 0.0
    if( E_NN != nothing )
        Energies = EnergiesT( 0.0, 0.0, 0.0, 0.0, 0.0, E_NN )
    else
        Energies = EnergiesT( 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 )
    end

    for iter = 1:NiterMax
        psi = psi - α_t*calc_grad( pw, Potentials, Focc, psi )

        psi = ortho_gram_schmidt(psi)

        # Update potentials
        rho = calc_rho( pw, Focc, psi)
        Potentials.Hartree = real( G_to_R( Ns, Poisson_solve( pw, rho ) ) )
        Potentials.XC = excVWN( rho ) + rho .* excpVWN( rho )

        Energies = calc_Energies( pw, Potentials, Focc, psi, Energies.NN )
        Etot = Energies.Total

        conv = abs(Etot-Etot_old)
        @printf("Emin SD: (s, step, conv) %8d %18.10f %18.10e\n", iter, Etot, conv )
        if conv < 1e-6
          print("Convergence achieved\n")
          break
        end
        Etot_old = Etot
    end
    return psi, Energies, Potentials
end


function solve_tsp(dist_mtx::Matrix{Int})
    if !issymmetric(dist_mtx)
        error("Asymmetric TSP is not supported.")
    end

    n_nodes = size(dist_mtx, 1)

    lower_diag_row = Int[]
    for i in 1:n_nodes
        for j in 1:i
            push!(lower_diag_row, dist_mtx[i, j])
        end
    end
    buf = 10
    n_rows = length(lower_diag_row) / buf |> ceil |> Int
    rows = String[]
    for i in 1:n_rows
        s = buf * (i-1) + 1
        t = min(buf * i, length(lower_diag_row))
        push!(rows, join(lower_diag_row[s:t], " "))
    end

    name = randstring(10)
    filename = name * ".tsp"
    open(filename, "w") do io
        write(io, "NAME: $(name)\n")
        write(io, "TYPE: TSP\n")
        write(io, "COMMENT: $(name)\n")
        write(io, "DIMENSION: $(n_nodes)\n")
        write(io, "EDGE_WEIGHT_TYPE: EXPLICIT\n")
        write(io, "EDGE_WEIGHT_FORMAT: LOWER_DIAG_ROW \n")
        write(io, "EDGE_WEIGHT_SECTION\n")
        for r in rows
            write(io, "$r\n")
        end
        write(io, "EOF\n")
    end

    return __solve_tsp__(filename)
end

function solve_tsp(x::Vector{Float64}, y::Vector{Float64}; dist="EUC_2D")
    n_nodes = length(x)
    @assert length(x) == length(y)

    name = randstring(10)
    filename = name * ".tsp"

    open(filename, "w") do io
        write(io, "NAME: $(name)\n")
        write(io, "TYPE: TSP\n")
        write(io, "COMMENT: $(name)\n")
        write(io, "DIMENSION: $(n_nodes)\n")
        write(io, "EDGE_WEIGHT_TYPE: $(dist)\n")
        write(io, "EDGE_WEIGHT_FORMAT: FUNCTION \n")
        write(io, "NODE_COORD_TYPE: TWOD_COORDS \n")
        write(io, "NODE_COORD_SECTION\n")
        for i in 1:n_nodes
            write(io, "$i $(x[i]) $(y[i])\n")
        end
        write(io, "EOF\n")
    end

    return __solve_tsp__(filename)
end

function solve_tsp(org_tsp_file::String)
    if !isfile(org_tsp_file)
        error("$(org_tsp_file) is not a file.")
    end

    name = randstring(10)
    filename = name * ".tsp"
    cp(org_tsp_file, filename)

    return __solve_tsp__(filename)
end

function __solve_tsp__(tsp_file::String)
    status = run(`$(Concorde.CONCORDE_EXECUTABLE) $(tsp_file)`, wait=false)
    while !success(status)
        #
    end

    name = split(basename(tsp_file), ".")[1]
    sol_filepath =  name * ".sol"
    opt_tour = read_solution(sol_filepath)

    tsp = readTSP(tsp_file)
    M = Int.(tsp.weights)
    opt_len = tour_length(opt_tour, M)

    cleanup(name)

    return opt_tour, opt_len
end


"""
	clean_proba_(probability:Number, atol=ATOL)

Checks whether a (complex) number is close enough to a valid probability with tolerance
`ATOL`. If so, convert it to a positive real number.
"""
function clean_proba(probability::Number, atol=ATOL)

    """checks if a (complex) number is a close enough probability
    converts it to a positive real number if complex or just positive if real"""

    not_a_proba() = error(probability , " is not a probability")

    if real(probability) >= -ATOL && real(probability) <= 1 + ATOL
        if isa(probability,Complex)
            if abs(imag(probability)) <= ATOL
                return abs(probability)
            else
                not_a_proba()
            end
        elseif isa(probability,Real)
            return abs(probability)
        else
            error(typeof(probability), " probability type undefined")
        end
    else
        not_a_proba()
    end
end

"""
	clean_pdf(A::Array, atol=ATOL)

Checks wether an array is an acceptable discrete probability distribution with
tolerance `ATOL`. If so, converts its elements to normalized positive real numbers.
"""
function clean_pdf(A::Array, atol = ATOL)

	"""checks if an array has all elements as acceptable probabilities within atol
	and converts them to that and summing to one within length(A) * atol
	and renormalizes"""

	A .= clean_proba.(A)
	normalization = sum(A)
	if isapprox(normalization, 1, atol = length(A) * atol)
		A .= 1/normalization * A
		return convert(Vector{Real}, A)
	else
		error("A not normalized")
	end
end

"""
	isa_pdf(pdf)

Asserts if `pdf`	is a valid probability distribution.
"""
function isa_pdf(pdf)

	"""asserts if pdf is a probability distribution"""
	clean_pdf(pdf)

end

"""
	tvd(a,b)

Computes the total variation distance between two probability distributions.
"""
function tvd(a,b)
	"""total variation distance"""
	sum(abs.(a-b))
end

"""
	sqr(a,b)

Computes the euclidian distance between two probability distributions.
"""
function sqr(a,b)
	"""euclidian distance"""
	sqrt(sum((a-b).^2))
end


# Comparing with Matt Taddy's gamlr.R
# To rebuild the test cases source(gammalasso.R)
using Lasso
using CSV, GLM, DataFrames, Random, SparseArrays

# often path length is different because of different stopping rules...
function issimilarhead(a::AbstractVector,b::AbstractVector;rtol=1e-4)
    n = min(size(a,1),size(b,1))
    isapprox(a[1:n],b[1:n];rtol=rtol) ? true : [a[1:n] b[1:n]]
end

function issimilarhead(a::AbstractMatrix,b::AbstractMatrix;rtol=1e-4)
    n = min(size(a,1),size(b,1))
    m = min(size(a,2),size(b,2))
    isapprox(a[1:n,1:m],b[1:n,1:m];rtol=rtol) ? true : [a[1:n,1:m] b[1:n,1:m]]
end

datapath = joinpath(dirname(@__FILE__), "data")

penaltyfactors = readcsvmat(joinpath(datapath,"penaltyfactors.csv"))

rtol=1e-2
Random.seed!(243214)
@testset "GammaLassoPath" begin
    @testset "$family" for (family, dist, link) in (("gaussian", Normal(), IdentityLink()), ("binomial", Binomial(), LogitLink()), ("poisson", Poisson(), LogLink()))
        data = readcsvmat(joinpath(datapath,"gamlr.$family.data.csv"))
        y = convert(Vector{Float64},data[:,1])
        X = convert(Matrix{Float64},data[:,2:end])
        (n,p) = size(X)
        @testset "penalty_factor=$pf" for pf = 1:size(penaltyfactors,2)
            penalty_factor = penaltyfactors[:,pf]
            if penalty_factor == ones(p)
                penalty_factor = nothing
            end
            @testset "γ=$γ" for γ in [0 2 10]
                fitname = "gamma$γ.pf$pf"

                # get gamlr.R prms and estimates
                prms = CSV.read(joinpath(datapath,"gamlr.$family.$fitname.params.csv"))
                fittable = CSV.read(joinpath(datapath,"gamlr.$family.$fitname.fit.csv"))
                gcoefs = readcsvmat(joinpath(datapath,"gamlr.$family.$fitname.coefs.csv");types=[Float64 for i=1:100])
                family = prms[1,Symbol("fit.family")]
                γ = prms[1,Symbol("fit.gamma")]
                λ = nothing #convert(Vector{Float64},fittable[!, Symbol("fit.lambda")]) # should be set to nothing evenatually

                # fit julia version
                glp = fit(GammaLassoPath, X, y, dist, link; γ=γ, stopearly=false,
                    λminratio=0.001, penalty_factor=penalty_factor, λ=λ,
                    standardize=false, standardizeω=false)

                # compare
                @test true==issimilarhead(glp.λ,fittable[!, Symbol("fit.lambda")];rtol=rtol)
                @test true==issimilarhead(glp.b0,fittable[!, Symbol("fit.alpha")];rtol=rtol)
                @test true==issimilarhead(convert(Matrix{Float64},glp.coefs'),gcoefs';rtol=rtol)
                # we follow GLM.jl convention where deviance is scaled by nobs, while in gamlr it is not
                @test true==issimilarhead(deviance(glp),fittable[!, Symbol("fit.deviance")]/nobs(glp);rtol=rtol)
                @test true==issimilarhead(deviance(glp,X,y),fittable[!, Symbol("fit.deviance")]/nobs(glp);rtol=rtol)
                # @test true==issimilarhead(round(df(glp)[2:end]),round(fittable[2:end,Symbol("fit.df")]))
                @test true==issimilarhead(loglikelihood(glp),fittable[!, Symbol("fit.logLik")];rtol=rtol)
                @test true==issimilarhead(aicc(glp),fittable[!, Symbol("fit.AICc")];rtol=rtol)

                # TODO: figure out why these are so off, maybe because most are corner solutions
                # and stopping rules for lambda are different
                # # what we really need all these stats for is that the AICc identifies the same minima:
                # if argmin(aicc(glp)) != lastindex(aicc(glp)) && argmin(fittable[!, Symbol("fit.AICc")]) != lastindex(fittable[!, Symbol("fit.AICc")])
                #     # interior minima
                #     println("comparing intereior AICc")
                #     @test argmin(aicc(glp)) == argmin(fittable[!, Symbol("fit.AICc")])
                # end

                # comparse CV, NOTE: this involves a random choice of train subsamples
                gcoefs_CVmin = vec(readcsvmat(joinpath(datapath,"gamlr.$family.$fitname.coefs.CVmin.csv")))
                gcoefs_CV1se = vec(readcsvmat(joinpath(datapath,"gamlr.$family.$fitname.coefs.CV1se.csv")))

                glp_CVmin = coef(glp,MinCVmse(glp, 10))
                glp_CV1se = coef(glp,MinCV1se(glp, 10))

                @test glp_CVmin ≈ gcoefs_CVmin rtol=0.3
                @test glp_CV1se ≈ gcoefs_CV1se rtol=0.3

                if γ==0
                    # Compare with LassoPath
                    lp = fit(LassoPath, X, y, dist, link; stopearly=false,
                        λminratio=0.001, penalty_factor=penalty_factor, λ=λ,
                        standardize=false, standardizeω=false)
                    @test glp.λ == lp.λ
                    @test glp.b0 ≈ lp.b0
                    @test glp.coefs ≈ lp.coefs
                end
            end
        end
    end
end
# ## the following code is useful for understanding comparison failures
#
# rdist(x::Number,y::Number) = abs(x-y)/max(abs(x),abs(y))
# rdist(x::AbstractArray{T}, y::AbstractArray{S}; norm::Function=norm) where {T<:Number,S<:Number} = norm(x - y) / max(norm(x), norm(y))

# (family, dist, link) = (("gaussian", Normal(), IdentityLink()), ("binomial", Binomial(), LogitLink()), ("poisson", Poisson(), LogLink()))[2]
# data = readcsvmat(joinpath(datapath,"gamlr.$family.data.csv"))
# y = convert(Vector{Float64},data[:,1])
# X = convert(Matrix{Float64},data[:,2:end])
# (n,p) = size(X)
# pf = (1:size(penaltyfactors,2))[1]
# penalty_factor = penaltyfactors[:,pf]
# if penalty_factor == ones(p)
#     penalty_factor = nothing
# end
#
# γ = [0 2 10][1]
# fitname = "gamma$γ.pf$pf"
#
# # get gamlr.R prms and estimates
# prms = CSV.read(joinpath(datapath,"gamlr.$family.$fitname.params.csv"))
# fittable = CSV.read(joinpath(datapath,"gamlr.$family.$fitname.fit.csv"))
# gcoefs = readcsvmat(joinpath(datapath,"gamlr.$family.$fitname.coefs.csv");types=[Float64 for i=1:100])
# family = prms[1,Symbol("fit.family")]
# γ = prms[1,Symbol("fit.gamma")]
# λ = nothing #convert(Vector{Float64},fittable[!, Symbol("fit.lambda")]) # should be set to nothing evenatually
#
# # fit julia version
# glp = fit(GammaLassoPath, X, y, dist, link; γ=γ, stopearly=false,
#     λminratio=0.001, penalty_factor=penalty_factor, λ=λ,
#     standardize=false, standardizeω=false)
#
# using InteractiveUtils, BenchmarkTools
# @code_warntype coef(glp; select=:all)
# @code_warntype minAICc(glp)
# @code_warntype segselect(glp, Lasso.MinAIC())
#
# @code_warntype coef(glp)
# @code_warntype coef(glp, AllSeg())
# @code_warntype coef(glp; select=:AICc)
# @code_warntype coef(glp, MinAICc())
#
# @btime coef(glp)
# s = AllSeg()
# @btime coef(glp, s)
# coef(glp; select=:AICc)
# @btime coef(glp; select=:AICc)
# @btime (s = MinAICc())
# s = MinAICc()
# @btime coef(glp, s)
#
# m = fit(LassoModel, X, y, dist, link; stopearly=false,
#     λminratio=0.001, penalty_factor=penalty_factor, λ=λ,
#     standardize=false, standardizeω=false)
# show(coef(m))
# predict(m)
#
# γ = 1
# m = fit(GammaLassoModel, X, y, dist, link; γ=γ, stopearly=false,
#     λminratio=0.001, penalty_factor=penalty_factor, λ=λ,
#     standardize=false, standardizeω=false)
# show(coef(m))
# predict(m)
#
# # compare
# @test true==issimilarhead(glp.λ,fittable[!, Symbol("fit.lambda")];rtol=rtol)
# @test true==issimilarhead(glp.b0,fittable[!, Symbol("fit.alpha")];rtol=rtol)
# @test true==issimilarhead(convert(Matrix{Float64},glp.coefs'),gcoefs';rtol=rtol)
# # we follow GLM.jl convention where deviance is scaled by nobs, while in gamlr it is not
# @test true==issimilarhead(deviance(glp),fittable[!, Symbol("fit.deviance")]/nobs(glp);rtol=rtol)
# @test true==issimilarhead(deviance(glp,X,y),fittable[!, Symbol("fit.deviance")]/nobs(glp);rtol=rtol)
# # @test true==issimilarhead(round(df(glp)[2:end]),round(fittable[2:end,Symbol("fit.df")]))
# @test true==issimilarhead(loglikelihood(glp),fittable[!, Symbol("fit.logLik")];rtol=rtol)
# @test true==issimilarhead(aicc(glp),fittable[!, Symbol("fit.AICc")];rtol=rtol)
#
# # TODO: figure out why these are so off, maybe because most are corner solutions
# # and stopping rules for lambda are different
# # # what we really need all these stats for is that the AICc identifies the same minima:
# # if argmin(aicc(glp)) != lastindex(aicc(glp)) && argmin(fittable[!, Symbol("fit.AICc")]) != lastindex(fittable[!, Symbol("fit.AICc")])
# #     # interior minima
# #     println("comparing intereior AICc")
# #     @test argmin(aicc(glp)) == argmin(fittable[!, Symbol("fit.AICc")])
# # end
#
# # comparse CV, NOTE: this involves a random choice of train subsamples
# gcoefs_CVmin = vec(readcsvmat(joinpath(datapath,"gamlr.$family.$fitname.coefs.CVmin.csv")))
# gcoefs_CV1se = vec(readcsvmat(joinpath(datapath,"gamlr.$family.$fitname.coefs.CV1se.csv")))
#
# @btime coef(glp,select=:CVmin,nCVfolds=10)
# @btime coef(glp, MinCVmse(glp, 10))
# @btime coef(glp, MinCV1se(glp, 10))
# glp_CVmin = coef(glp,select=:CVmin,nCVfolds=10)
#
# glp_CV1se = coef(glp,select=:CV1se,nCVfolds=10)
#
# @test glp_CVmin ≈ gcoefs_CVmin rtol=0.3
# @test glp_CV1se ≈ gcoefs_CV1se rtol=0.3
#
# if γ==0
#     # Compare with LassoPath
#     lp = fit(LassoPath, X, y, dist, link; stopearly=false,
#         λminratio=0.001, penalty_factor=penalty_factor, λ=λ,
#         standardize=false, standardizeω=false)
#     @test glp.λ == lp.λ
#     @test glp.b0 ≈ lp.b0
#     @test glp.coefs ≈ lp.coefs
# end


# Example: Compare the performance of some math functions

using BenchmarkLite

## Define VecMath procedures
#
# The configuration is the vector length, represented by an integer
#

type VecMath{Op} <: Proc end

type Sqrt end
calc(::Sqrt, x) = sqrt(x)

type Exp end
calc(::Exp, x) = exp(x)

type Log end
calc(::Log, x) = log(x)

type Sin end
calc(::Sin, x) = sin(x)

# procedure name
Base.string{Op}(::VecMath{Op}) = string("vec-", lowercase("$Op"))

# pre-allocated arrays for running the procedure
typealias FVecPair (Vector{Float64},Vector{Float64})

# procedure codes
Base.length(p::VecMath, n::Int) = n

Base.isvalid(p::VecMath, n::Int) = (n > 0)

Base.start(p::VecMath, n::Int) = (rand(n), zeros(n))

function Base.run{Op}(p::VecMath{Op}, n::Int, s::FVecPair)
    x, y = s
    op = Op()
    for i = 1:n
        @inbounds y[i] = calc(op, x[i])
    end
end

Base.done(p::VecMath, n, s) = nothing

## perform the benchmark

procs = Proc[ VecMath{Sqrt}(),
              VecMath{Exp}(),
              VecMath{Log}(),
              VecMath{Sin}() ]

cfgs = 2 .^ (4:12)

println("Running log:")
println("--------------------")
rtable = run(procs, cfgs; duration=0.2)
@assert isa(rtable, BenchmarkTable)
println()

## show results

show(rtable; unit=:mps, cfghead="len")

# open("mathfuns.bench.csv", "w") do f
#   writecsv(f, rtable)
# end


# This file is part of the TaylorIntegration.jl package; MIT licensed

using TaylorIntegration
using Test
using LinearAlgebra: norm
using Logging
import Logging: Warn

@testset "Testing `one_ode.jl`" begin

    local _order = 28
    local _abstol = 1.0E-20
    local tT = Taylor1(_order)

    max_iters_reached() = "Maximum number of integration steps reached; exiting.\n"

    @testset "Tests: dot{x}=x^2, x(0) = 1" begin
        eqs_mov(x, p, t) = x^2
        t0 = 0.0
        x0 = 1.0
        x0T = Taylor1(x0, _order)
        tT[1] = t0
        TaylorIntegration.jetcoeffs!(eqs_mov, tT, x0T, nothing)
        @test x0T.coeffs[end] == 1.0
        δt = _abstol^inv(_order-1)
        @test TaylorIntegration.stepsize(x0T, _abstol) == δt

        tv, xv = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov, 1, 0.0, 1.0, _order, _abstol))
        @test length(tv) == 501
        @test length(xv) == 501
        @test xv[1] == x0
        @test tv[end] < 1.0

        tv, xv = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov, x0, 0.0, 1.0, _order, _abstol, nothing))
        @test length(tv) == 501
        @test length(xv) == 501
        @test xv[1] == x0
        @test tv[end] < 1.0

        trange = 0.0:1/8:1.0
        xv = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov, 1, trange, _order, _abstol))
        @test length(xv) == length(trange)
        @test typeof(xv) == Array{typeof(x0),1}
        @test xv[1] == x0
        @test isnan(xv[end])
        @test abs(xv[5] - 2.0) ≤ eps(2.0)
        tvr, xvr = (@test_logs min_level=Logging.Warn taylorinteg(
            eqs_mov, x0, trange[1], trange[end-1], _order, _abstol))
        @test tvr[1] == trange[1]
        @test tvr[end] == trange[end-1]
        @test xvr[1] == xv[1]
        @test xvr[end] == xv[end-1]

        trange = 0.0:1/8:1.0
        xv = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov, x0, trange, _order, _abstol, nothing))
        @test length(xv) == length(trange)
        @test typeof(xv) == Array{typeof(x0),1}
        @test xv[1] == x0
        @test isnan(xv[end])
        @test abs(xv[5] - 2.0) ≤ eps(2.0)
        tvr, xvr = (@test_logs min_level=Logging.Warn taylorinteg(
            eqs_mov, x0, trange[1], trange[end-1], _order, _abstol, nothing))
        @test tvr[1] == trange[1]
        @test tvr[end] == trange[end-1]
        @test xvr[1] == xv[1]
        @test xvr[end] == xv[end-1]

        tarray = vec(trange)
        xv2 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov, x0, tarray, _order, _abstol))
        @test xv[1:end-1] == xv2[1:end-1]
        @test xv2[1:end-1] ≈ xv[1:end-1] atol=eps() rtol=0.0
        @test length(xv2) == length(tarray)
        @test typeof(xv2) == Array{typeof(x0),1}
        @test xv2[1] == x0
        @test isnan(xv2[end])
        @test abs(xv2[5] - 2.0) ≤ eps(2.0)

        # Output includes Taylor polynomial solution
        tv, xv, polynV = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov, x0, 0, 0.5, _order, _abstol, Val(true), maxsteps=2))
        @test length(polynV) == 3
        @test xv[1] == x0
        @test polynV[1] == Taylor1(x0, _order)
        @test xv[2] == evaluate(polynV[2], tv[2]-tv[1])
        @test polynV[2] == Taylor1(ones(_order+1))
    end

    @testset "Tests: dot{x}=x^2, x(0) = 3; nsteps <= maxsteps" begin
        eqs_mov(x, p, t) = x.^2 #the ODE (i.e., the equations of motion)
        exactsol(t, x0) = x0/(1.0-x0*t) #the analytical solution
        t0 = 0.0
        tmax = 0.3
        x0 = 3.0
        x0T = Taylor1(x0, _order)
        tT[1] = t0
        TaylorIntegration.jetcoeffs!(eqs_mov, tT, x0T, nothing)
        @test x0T.coeffs[end] == 3.0^(_order+1)
        δt = (_abstol/x0T.coeffs[end-1])^inv(_order-1)
        @test TaylorIntegration.stepsize(x0T, _abstol) == δt

        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            eqs_mov, x0, 0, tmax, _order, _abstol))
        @test length(tv) < 501
        @test length(xv) < 501
        @test length(tv) == 14
        @test length(xv) == 14
        @test xv[1] == x0
        @test tv[end] < 1/3
        @test tv[end] == tmax
        @test abs(xv[end]-exactsol(tv[end], xv[1])) < 5e-14

        tmax = 0.33
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            eqs_mov, x0, t0, tmax, _order, _abstol))
        @test length(tv) < 501
        @test length(xv) < 501
        @test length(tv) == 28
        @test length(xv) == 28
        @test xv[1] == x0
        @test tv[end] < 1/3
        @test tv[end] == tmax
        @test abs(xv[end]-exactsol(tv[end], xv[1])) < 1e-11
    end

    @testset "Test non-autonomous ODE (1): dot{x}=cos(t)" begin
        let f(x, p, t) = cos(t)
            t0 = 0//1
            tmax = 10.25*(2pi)
            abstol = 1e-20
            order = 25
            x0 = 0.0 #initial conditions such that x(t)=sin(t)

            tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0, t0, tmax, order, abstol))
            @test length(tv) < 501
            @test length(xv) < 501
            @test xv[1] == x0
            @test tv[1] == t0
            @test abs(sin(tmax)-xv[end]) < 1e-14

            # Backward integration
            tb, xb = (@test_logs min_level=Logging.Warn taylorinteg(
                f, sin(tmax), tmax, t0, order, abstol))
            @test length(tb) < 501
            @test length(xb) < 501
            @test xb[1] == sin(tmax)
            @test tb[1] > tb[end]
            @test abs(sin(t0)-xb[end]) < 5e-14

            # Tests with a range, for comparison with backward integration
            tmax = 15*(2pi)
            Δt = (tmax-t0)/1024
            tspan = t0:Δt:tmax
            xv = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0, tspan, order, abstol))
            @test xv[1] == x0
            @test abs(sin(tmax)-xv[end]) < 1e-14

            # Backward integration
            xback = (@test_logs min_level=Logging.Warn taylorinteg(
                f, xv[end], reverse(tspan), order, abstol))
            @test xback[1] == xv[end]
            @test abs(sin(t0)-xback[end]) < 5e-14
            @test norm(xv[:]-xback[end:-1:1], Inf) < 5.0e-14

            # Tests if trange is properly sorted
            @test_throws AssertionError taylorinteg(f, x0, rand(t0:Δt:tmax, 100), order, abstol)
        end
    end

end


iscompact(d::TypedEndpointsInterval{:closed,:closed}) = true
iscompact(d::TypedEndpointsInterval) = false

isinterval(d::Domain) = false
isinterval(d::AbstractInterval) = true

hash(d::AbstractInterval, h::UInt) =
    hashrec(isleftopen(d), isrightopen(d), leftendpoint(d), rightendpoint(d), h)

Display.object_parentheses(::AbstractInterval) = true

approx_indomain(x, d::AbstractInterval, tolerance) =
    (x >= leftendpoint(d)-tolerance) && (x <= rightendpoint(d)+tolerance)

approx_indomain(x, d::TypedEndpointsInterval{:closed,:closed}, tolerance) =
    (x >= leftendpoint(d)-tolerance) && (x <= rightendpoint(d)+tolerance)
approx_indomain(x, d::TypedEndpointsInterval{:open,:closed}, tolerance) =
    (x > leftendpoint(d)-tolerance) && (x <= rightendpoint(d)+tolerance)
approx_indomain(x, d::TypedEndpointsInterval{:closed,:open}, tolerance) =
    (x >= leftendpoint(d)-tolerance) && (x < rightendpoint(d)+tolerance)
approx_indomain(x, d::TypedEndpointsInterval{:open,:open}, tolerance) =
    (x > leftendpoint(d)-tolerance) && (x < rightendpoint(d)+tolerance)

function point_in_domain(d::AbstractInterval)
    isempty(d) && throw(BoundsError())
    mean(d)
end

# For an interval of integers, try to find an integer point
function point_in_domain(d::AbstractInterval{T}) where {T<:Integer}
    isempty(d) && throw(BoundsError())
    x = round(T, mean(d))
    if x ∈ d
        x
    else
        # the mean is not inside the interval when rounded, this means
        # we have to choose an endpoint
        if isleftclosed(d)
            leftendpoint(d)
        elseif isrightclosed(d)
            rightendpoint(d)
        else
            # the interval is open and contains no integers in its interior
            throw(BoundsError())
        end
    end
end

isapprox(d1::AbstractInterval, d2::AbstractInterval) =
    isapprox(leftendpoint(d1), leftendpoint(d2); atol=default_tolerance(d1)) &&
    isapprox(rightendpoint(d1), rightendpoint(d2); atol=default_tolerance(d1))

boundary(d::AbstractInterval) = Point(leftendpoint(d)) ∪ Point(rightendpoint(d))
corners(d::AbstractInterval) = [leftendpoint(d), rightendpoint(d)]

normal(d::AbstractInterval, x) = (abs(minimum(d)-x) < abs(maximum(d)-x)) ? -one(eltype(d)) : one(eltype(d))

distance_to(d::AbstractInterval, x) = x ∈ d ? zero(eltype(d)) : min(abs(x-supremum(d)), abs(x-infimum(d)))

boundingbox(d::AbstractInterval) = d

volume(d::AbstractInterval) = width(d)

similar_interval(d::AbstractInterval, a, b) = similar_interval(d, promote(a, b)...)

## Some special intervals follow, e.g.:
# - the unit interval [0,1]
# - the 'Chebyshev' interval [-1,1]
# - ...
# Unlike a generic interval, these specific intervals have no data, and only one
# type parameter (T).

"""
The abstract type `FixedInterval` is the supertype of intervals with endpoints
determined by the type, rather than field values. Examples include `UnitInterval`
and `ChebyshevInterval`.
"""
abstract type FixedInterval{L,R,T} <: TypedEndpointsInterval{L,R,T} end
const ClosedFixedInterval{T} = FixedInterval{:closed,:closed,T}

closure(d::AbstractInterval) = ClosedInterval(endpoints(d)...)
closure(d::ClosedInterval) = d

interior(d::AbstractInterval) = OpenInterval(endpoints(d)...)

"""
Return an interval that is similar to the given interval, but with endpoints
`a` and `b` instead.
"""
similar_interval(d::ClosedFixedInterval{T}, a::S, b::S) where {S,T} =
    ClosedInterval{promote_type(float(T),S)}(a, b)

"The closed unit interval [0,1]."
struct UnitInterval{T} <: ClosedFixedInterval{T} end

UnitInterval() = UnitInterval{Float64}()

endpoints(d::UnitInterval{T}) where {T} = (zero(T), one(T))

similardomain(::UnitInterval, ::Type{T}) where {T} = UnitInterval{T}()

"The closed interval [-1,1]."
struct ChebyshevInterval{T} <: ClosedFixedInterval{T}
end

ChebyshevInterval() = ChebyshevInterval{Float64}()

endpoints(d::ChebyshevInterval{T}) where {T} = (-one(T),one(T))

similardomain(::ChebyshevInterval, ::Type{T}) where {T} = ChebyshevInterval{T}()

-(d::ChebyshevInterval) = d

"Map the interval [a,b] to the interval [c,d]."
interval_map(a, b, c, d) = AffineMap((d-c)/(b-a), c - a*(d-c)/(b-a))

canonicaldomain(d::ClosedInterval{T}) where {T} = ChebyshevInterval{float(T)}()
canonicaldomain(d::Interval{L,R,T}) where {L,R,T} = Interval{L,R,float(T)}(-1, 1)
canonicaldomain(d::FixedInterval) = d

mapfrom_canonical(d::Interval{L,R,T}) where {L,R,T} = interval_map(-one(T), one(T), endpoints(d)...)
mapfrom_canonical(d::FixedInterval{L,R,T}) where {L,R,T} = StaticIdentityMap{T}()

mapto(d1::D, d2::D) where {D <: FixedInterval} = identitymap(d1)
mapto(d1::D1, d2::D2) where {D1 <: FixedInterval,D2 <: FixedInterval} =
    interval_map(infimum(d1), supremum(d1), infimum(d2), supremum(d2))

"The half-open positive halfline `[0,∞)`."
struct HalfLine{T} <: FixedInterval{:closed,:open,T} end
HalfLine() = HalfLine{Float64}()

endpoints(d::HalfLine{T}) where {T} = (zero(T), T(Inf))

boundary(d::HalfLine) = Point(leftendpoint(d))

similardomain(::HalfLine, ::Type{T}) where {T} = HalfLine{T}()

# intercept and simplify the definition of IntervalSets
in(x, d::HalfLine) = x >= 0

approx_indomain(x, d::HalfLine, tolerance) = x >= -tolerance

function similar_interval(d::HalfLine{T}, a::S, b::S) where {T,S}
    @assert a == 0
    @assert isinf(b) && b > 0
    HalfLine{promote_type(float(T),S)}()
end

point_in_domain(d::HalfLine) = zero(eltype(d))

"The open negative halfline `(-∞,0)`."
struct NegativeHalfLine{T} <: FixedInterval{:open,:open,T} end
NegativeHalfLine() = NegativeHalfLine{Float64}()

similardomain(::NegativeHalfLine, ::Type{T}) where {T} = NegativeHalfLine{T}()

endpoints(d::NegativeHalfLine{T}) where {T} = (-T(Inf), zero(T))

boundary(d::NegativeHalfLine) = Point(rightendpoint(d))

# intercept and simplify the definition of IntervalSets
in(x, d::NegativeHalfLine) = x < 0

approx_indomain(x, d::NegativeHalfLine, tolerance) = x < tolerance

function similar_interval(d::NegativeHalfLine{T}, a::S, b::S) where {S,T}
    @assert isinf(a) && a < 0
    @assert b == 0
    NegativeHalfLine{promote_type(S,float(T))}()
end

point_in_domain(d::NegativeHalfLine) = -one(eltype(d))

similar_interval(d::Interval{L,R,T}, a::S, b::S) where {L,R,T,S} =
    Interval{L,R,promote_type(float(T),S)}(a, b)

#########################################
# A few set operations with known result
#########################################

# We define an exhaustive list of combinations of the four fixed intervals
# combined above in the routines 'intersect', 'union' and 'setdiff' where
# the output is known explicitly.

# Override the definition of Intervals.jl for FixedInterval's defined here
union(d1::FixedInterval, d2::FixedInterval) = uniondomain(d1, d2)
intersect(d1::FixedInterval, d2::FixedInterval) = intersectdomain(d1, d2)

uniondomain(d1::TypedEndpointsInterval, d2::TypedEndpointsInterval) =
    uniondomain(promote(d1,d2)...)

# type-unstable union of intervals. This function differs from `union` in
# IntervalSets.jl because that one throws an error if the intervals do not overlap
function uniondomain(d1::TypedEndpointsInterval{L1,R1,T}, d2::TypedEndpointsInterval{L2,R2,T}) where {L1,R1,L2,R2,T}
    a1 = leftendpoint(d1)
    b1 = rightendpoint(d1)
    a2 = leftendpoint(d2)
    b2 = rightendpoint(d2)

    # are they empty?
    isempty(d1) && return d2
    isempty(d2) && return d1
    # are they equal?
    d1 == d2 && return d1
    # does one lie within the other?
    a2 > a1 && b2 < b1 && return d1
    a1 > a2 && b1 < b2 && return d2
    # are they disjoint?
    b2 < a1 && return UnionDomain(d1, d2)   # return a UnionDomain for disjoint intervals
    b1 < a2 && return UnionDomain(d1, d2)
    # at this stage they must overlap
    if a1 < a2
        a = a1
        L = L1
    elseif a1 == a2
        a = a1
        L = (L1==L2==:open) ? :open : :closed
    else
        a = a2
        L = L2
    end
    if b1 > b2
        b = b1
        R = R1
    elseif b1 == b2
        b = b1
        R = (R1==R2==:open) ? :open : :closed
    else
        b = b2
        R = R2
    end
    Interval{L,R,T}(a, b)
end

function intersectdomain(d1::TypedEndpointsInterval, d2::TypedEndpointsInterval)
    # go back to the definition of IntervalSets.jl
    d = intersect(d1, d2)
    if isempty(d)
        # but avoid returning an interval like 2..1
        EmptySpace{eltype(d)}()
    elseif leftendpoint(d)==rightendpoint(d)
        # and avoid an interval like 1..1
        Point(leftendpoint(d))
    else
        d
    end
end

==(d1::TypedEndpointsInterval, d2::Point) =
    isclosedset(d1) && (leftendpoint(d1)==rightendpoint(d1)==d2.x)

# intersectdomain(d1::AbstractInterval, d2::AbstractInterval) = intersect(d1, d2)
# uniondomain(d1::AbstractInterval, d2::AbstractInterval) = union(d1, d2)

# Since fixed intervals are fully determined by their type,
# the result of intersect, union or setdiff is always known for two
# domains of the same type.
intersectdomain(d1::D, d2::D) where {D <: FixedInterval} = d1
uniondomain(d1::D, d2::D) where {D <: FixedInterval} = d1
setdiffdomain(d1::D, d2::D) where {D <: FixedInterval} = EmptySpace{eltype(D)}()

# [0,1] ∩ [-1,1] = [0,1]
intersectdomain(d1::UnitInterval{T}, d2::ChebyshevInterval{T}) where {T} = UnitInterval{T}()
intersectdomain(d1::ChebyshevInterval{T}, d2::UnitInterval{T}) where {T} = UnitInterval{T}()
# [0,1] ∩ [0,∞) = [0,1]
intersectdomain(d1::UnitInterval{T}, d2::HalfLine{T}) where {T} = UnitInterval{T}()
intersectdomain(d1::HalfLine{T}, d2::UnitInterval{T}) where {T} = UnitInterval{T}()
# [0,1] ∩ (-∞,0) = {}
intersectdomain(d1::UnitInterval{T}, d2::NegativeHalfLine{T}) where {T} = EmptySpace{T}()
intersectdomain(d1::NegativeHalfLine{T}, d2::UnitInterval{T}) where {T} = EmptySpace{T}()
# [-1,1] ∩ [0,∞) = [0,1]
intersectdomain(d1::ChebyshevInterval{T}, d2::HalfLine{T}) where {T} = UnitInterval{T}()
intersectdomain(d1::HalfLine{T}, d2::ChebyshevInterval{T}) where {T} = UnitInterval{T}()
# [0,∞) ∩ (-∞,0) = {}
intersectdomain(d1::HalfLine{T}, d2::NegativeHalfLine{T}) where {T} = EmptySpace{T}()
intersectdomain(d1::NegativeHalfLine{T}, d2::HalfLine{T}) where {T} = EmptySpace{T}()

# [0,1] ∪ [-1,1] = [-1,1]
uniondomain(d1::UnitInterval{T}, d2::ChebyshevInterval{T}) where {T} = ChebyshevInterval{T}()
uniondomain(d1::ChebyshevInterval{T}, d2::UnitInterval{T}) where {T} = ChebyshevInterval{T}()
# [0,1] ∪ [0,∞) = [0,∞)
uniondomain(d1::UnitInterval{T}, d2::HalfLine{T}) where {T} = HalfLine{T}()
uniondomain(d1::HalfLine{T}, d2::UnitInterval{T}) where {T} = HalfLine{T}()

# (-∞,0) ∪ [0,∞) = (-∞,∞)
# Note: T<:real to ensure that FullSpace{T} is not larger than intended.
uniondomain(d1::NegativeHalfLine{T}, d2::HalfLine{T}) where {T<:Real} = FullSpace{T}()
uniondomain(d1::HalfLine{T}, d2::NegativeHalfLine{T}) where {T<:Real} = FullSpace{T}()

# [0,1] ∖ [-1,1] = {}
setdiffdomain(d1::UnitInterval{T}, d2::ChebyshevInterval{T}) where {T} = EmptySpace{T}()
# [0,1] ∖ [0,∞) = {}
setdiffdomain(d1::UnitInterval{T}, d2::HalfLine{T}) where {T} = EmptySpace{T}()
# [0,1] ∖ (-∞,0) = [0,1]
setdiffdomain(d1::UnitInterval{T}, d2::NegativeHalfLine{T}) where {T} = UnitInterval{T}()
# [-1,1] ∖ (-∞,0) = [0,1]
setdiffdomain(d1::ChebyshevInterval{T}, d2::NegativeHalfLine{T}) where {T} = UnitInterval{T}()
# [0,∞) ∖ (-∞,0) = [0,∞)
setdiffdomain(d1::HalfLine{T}, d2::NegativeHalfLine{T}) where {T} = HalfLine{T}()
# (-∞,0) ∖ [0,1] = (-∞,0)
setdiffdomain(d1::NegativeHalfLine{T}, d2::UnitInterval{T}) where {T} = NegativeHalfLine{T}()
# (-∞,0) ∖ [0,∞) = (-∞,0)
setdiffdomain(d1::NegativeHalfLine{T}, d2::HalfLine{T}) where {T} = NegativeHalfLine{T}()

#################################
# Conversions between intervals
#################################
# Avoid depcrecated warning: Warning: Constructors no longer fall back to `convert`.
# example: A constructor `AbstractInterval{Float64}(::IntervalSets.ClosedInterval{Float64})` should be defined instead.

function convert(::Type{UnitInterval{T}}, d::AbstractInterval) where {T}
    endpoints(d) == (0,1) || throw(InexactError(:convert,UnitInterval,d))
    UnitInterval{T}()
end

function convert(::Type{ChebyshevInterval{T}}, d::AbstractInterval) where {T}
    endpoints(d) == (-1,1)|| throw(InexactError(:convert,ChebyshevInterval,d))
    ChebyshevInterval{T}()
end

UnitInterval{T}(d::AbstractInterval) where T = convert(UnitInterval{T}, d)
ChebyshevInterval{T}(d::AbstractInterval) where T = convert(ChebyshevInterval{T}, d)

function show(io::IO, d::ChebyshevInterval)
    print(io, Interval(d))
    print(io, " (Chebyshev)")
end
function show(io::IO, d::UnitInterval)
    print(io, Interval(d))
    print(io, " (Unit)")
end
function show(io::IO, d::HalfLine)
    print(io, Interval(d))
    print(io, " (HalfLine)")
end

########################
# Arithmetic operations
########################

function setdiffdomain(d1::AbstractInterval{T}, d2::AbstractInterval{T}) where T
    a1 = leftendpoint(d1)
    b1 = rightendpoint(d1)
    a2 = leftendpoint(d2)
    b2 = rightendpoint(d2)

    isempty(d1) && return d1
    isempty(d2) && return d1
    b1 < a2 && return d1
    a1 < a2 ≤ b1 ≤ b2 && return (a1..a2)
    a1 < a2 ≤ b2 < b1 && return uniondomain(a1..a2, b2..b1)
    a2 ≤ a1 < b2 < b1 && return (b2..b1)
    a2 ≤ a1 ≤ b1 ≤ b2 && return EmptySpace{T}()

    @assert b2 ≤ a1
    d1
end

setdiffdomain(d1::AbstractInterval, d2::AbstractInterval) =
    setdiffdomain(promote(d1,d2)...)

switch_open_closed(d::AbstractInterval) = d
switch_open_closed(d::Interval{L,R,T}) where {L,R,T} =
    Interval{R,L,T}(leftendpoint(d),rightendpoint(d))

function map_domain(map::AbstractAffineMap{<:Number}, domain::AbstractInterval)
    le = map(leftendpoint(domain))
    re = map(rightendpoint(domain))
    if le<re
        similar_interval(domain,le,re)
    else
        similar_interval(switch_open_closed(domain),re,le)
    end
end

mapped_domain(invmap::AbstractAffineMap{<:Number}, domain::AbstractInterval) =
    map_domain(inverse(invmap), domain)

# Preserve maps when acting on the half line and negative halfline, because
# the map may not always preserve the domain type
map_domain(map::AbstractAffineMap{<:Number}, domain::HalfLine) = MappedDomain(inverse(map), domain)
map_domain(map::AbstractAffineMap{<:Number}, domain::NegativeHalfLine) = MappedDomain(inverse(map), domain)


cd(@__DIR__)
using Pkg
Pkg.activate(".")

##

using Revise
using Statistics
using ProBeTune
using DiffEqFlux
using LaTeXStrings

##

##
include("PlotUtils.jl")
path = "../plots/diffusive_nl_system/" # for saving plots
mkpath(path)
##

#=
## Diffusive non-linear system example
As an example we want to tune a system of 10 nodes with one grid connection point to a specification of 2 nodes.
Our goal is to find such parameters for the 10-node system that it behaves close enough to a 2-node system under all possible inputs.
In this example we start with just 10 possible inputs.
=#
dim_sys = 10
diffusive_nl_example = PBTExample.create_diffusive_nl_example(dim_sys, 2, 0.:0.1:10., 10);

# => Can we make a large graph react like a small graph by tuning the non-linearity at the nodes?
scenarios = 1:3 # we can choose what scnearios to use for plots everywhere

p_initial = ones(2*10+dim_sys)
plot_callback(diffusive_nl_example, p_initial, 10., scenario_nums=scenarios)
savefig(path*"diffusive_nl_example10_initial_rand.png")
##
#=
Finished initialization
=#
d, p_tuned, dists = behavioural_distance(diffusive_nl_example, p_initial; abstol=1e-2, reltol=1e-2)

##
scenarios = 1:3 # we can choose what scnearios to use for plots everywhere
plot_callback(diffusive_nl_example, p_tuned, d, scenario_nums=scenarios)
savefig(path*"diffusive_nl_example10_initial_"*string(round(d, digits = 4))*".png")

##
res_10 = pbt_tuning(diffusive_nl_example, p_tuned; abstol=1e-4, reltol=1e-4,
                    optimizer = ADAM(0.5),
                    optimizer_options = (
                        :maxiters => 50,
                        :cb => PBTLibrary.basic_pbt_callback))

p_tuned = res_10.minimizer
##
res_10 = pbt_tuning(diffusive_nl_example, p_tuned; abstol=1e-6, reltol=1e-6,
                    optimizer = ADAM(0.1),
                    optimizer_options = (
                        :maxiters => 100,
                        :cb => PBTLibrary.basic_pbt_callback))

p_tuned = res_10.minimizer
##
res_10 = pbt_tuning(diffusive_nl_example, p_tuned; abstol=1e-6, reltol=1e-6,
                    optimizer = AMSGrad(0.01),
                    optimizer_options = (
                        :maxiters => 1000,
                        :cb => PBTLibrary.basic_pbt_callback))

p_tuned = res_10.minimizer
##
plot_callback(diffusive_nl_example, p_tuned, res_10.minimum, scenario_nums = scenarios)

savefig(path*"diffusive_nl_example10_final_"*string(round(res_10.minimum, digits = 4))*".png")
##

#=
After getting a good initial approximation, we can look at the minimizer.
Going further makes little sense as we would be overfitting to the small sample.
=#
print(relu.(p_tuned))

#= ## Resampling
We can check the quality of the resulting minimizer by optimizing the specs only (a much simpler repeated 2d optimization problem)
=#
d, p2 = behavioural_distance(diffusive_nl_example, p_tuned)
plot_callback(diffusive_nl_example, p2, d, scenario_nums = scenarios)

#=
In order to understand how much we were overfitting with respect to the concrete sample, we resample.
That is, we generate a new problem with a new sample from the same distribution:
This will give us information on the system tuning with a sample different from the one that the tuning was optimized for.
=#
resample!(diffusive_nl_example)
d_rs, p3 = behavioural_distance(diffusive_nl_example, p_tuned)

println(d_rs/d)

##
plot_callback(diffusive_nl_example, p3, d_rs, scenario_nums = scenarios)
savefig(path*"diffusive_nl_example10_resampled_"*string(round(d_rs, digits = 4))*".png")

#=
The median individual loss has gone up by a factor of 4.
This means that the system is somewhat overfit to the initial sample.
## Larger number of samples
We warmed up the optimization with a very small number of scenarios, we can now initialize a higher sampled optimization using the system parameters found in the lower one:
=#
p_100 = ones(2*100+dim_sys)
p_100[1:dim_sys] .= p2[1:dim_sys]
p_100[dim_sys+1:end] .= repeat(p2[dim_sys+1:dim_sys+2], 100)

resample!(diffusive_nl_example; n = 100);

#=
Optimizing only the specs is a task linear in the number of scenarios, the idea is that this will help with warming up the optimization.
We can also study the quality of the tuning found by the optimization based on a small number of samples.
=#
d100, p_100_initial = behavioural_distance(diffusive_nl_example, p_100;
                    optimizer = DiffEqFlux.ADAM(0.1),
                    optimizer_options = (:maxiters => 10,),
                    abstol = 1e-3, reltol=1e-3)

plot_callback(diffusive_nl_example, p_100_initial, d100, scenario_nums = scenarios)
savefig(path*"diffusive_nl_example100_initial_"*string(round(d100, digits = 4))*".png")

#= Now we can train the full system:
=#
res_100 = pbt_tuning(diffusive_nl_example, p_100_initial; abstol=1e-4, reltol=1e-4,
                    optimizer = BFGS(),
                    optimizer_options = (
                        :maxiters => 10,
                        :cb => PBTLibrary.basic_pbt_callback))

p_tuned = res_100.minimizer
p_tuned0 = res_100.minimizer
plot_callback(diffusive_nl_example, p_tuned, res_100.minimum, scenario_nums = scenarios)
#= Continue improving it for 150 Steps with some plotting in between:=#
for i in 1:10
    global res_100
    global p_tuned

    res_100 = pbt_tuning(diffusive_nl_example, p_tuned; abstol=1e-6, reltol=1e-6,
                        optimizer = ADAM(0.01),
                        optimizer_options = (
                            :maxiters => 50,
                            :cb => PBTLibrary.basic_pbt_callback))
    p_tuned = res_100.minimizer
    res_100 = pbt_tuning(diffusive_nl_example, p_tuned; abstol=1e-6, reltol=1e-6,
                        optimizer = AMSGrad(0.01),
                        optimizer_options = (
                            :maxiters => 200,
                            :cb => PBTLibrary.basic_pbt_callback))
    p_tuned = res_100.minimizer
    plot_callback(diffusive_nl_example, res_100.minimizer, res_100.minimum, scenario_nums = scenarios)
end

##
d, p, losses = behavioural_distance(diffusive_nl_example, res_100.minimizer)

plot_callback(diffusive_nl_example, p, d; scenario_nums=scenarios)
savefig(path*"diffusive_nl_example100_final_"*string(round(d, digits = 4))*".png")
##
resample!(diffusive_nl_example; n = 100);
d_rs, p_rs, losses = behavioural_distance(diffusive_nl_example, p)

plot_callback(diffusive_nl_example, p_rs, d_rs; scenario_nums=scenarios)
savefig(path*"diffusive_nl_example100_tuned_rs_"*string(round(d_rs, digits = 4))*".png")
##

x= exp10.(range(log10(.1),stop=log10(0.5), length = 50))
plot(x, (x)->1-PBTLibrary.confidence_interval(losses, x)[1],
    xlabel = "ε", ylabel=L"\hat d^{\rho,\varepsilon}",legend=:bottomright, label=false,c=:blue)
    #label="Fraction of scenarios within set distance from specification")
    plot!(x, (x)->1-PBTLibrary.confidence_interval(losses, x)[2],
        xlabel = "ε", ylabel=L"\hat d^{\rho,\varepsilon}",legend=:bottomright, label=false, linestyle=:dash, c=:blue)
        #label="Fraction of scenarios within set distance from specification")
    plot!(x, (x)->1-PBTLibrary.confidence_interval(losses, x)[3],
        xlabel = "ε", ylabel=L"\hat d^{\rho,\varepsilon}",legend=:bottomright, label=false, linestyle=:dash, c=:blue)
        #label="Fraction of scenarios within set distance from specification")

savefig(path*"diffusive_nl_example_confidence_interval.png")


cd(@__DIR__)
using Pkg
Pkg.activate(".")

##

using Revise
using Statistics
using ProBeTune
using DiffEqFlux
using LaTeXStrings

##
include("PlotUtils.jl")
path = "../plots/kuramoto/" # for saving plots
mkpath(path)
save_plots = true # set to true
##
#=
Define global system parameters to use in the tune_run function
=#
global t_steps = 0.:0.1:10.
global tspan = (0., 10.)
global N_osc = 10
global size_p_sys = N_osc * (N_osc - 1) ÷ 2 + N_osc
global size_p_spec = 2
global N_samples = 10
global p_sys_init = 6. * rand(size_p_sys) .+ 1.
global p_spec_init = rand(size_p_spec) .+ 1.
global p_initial = vcat(p_sys_init, repeat(p_spec_init, N_samples))
global K_av = 1.
##
#=
Create a set of intrinsic frequencies with mean(omega)=0
In tune_run we use omega*scale with variating scales, to explore systems with different synchronisity.
=#
omega = randn(N_osc)
omega[1] = 0.
omega .-= mean(omega[1:end])

##
#=
Tuning pipeline for the system of kuramoto oscillators.
We want to tune a system of 10 oscillators with one grid connection to behave like a single oscillator (specification) in terms of input-output behavior.
We seek to find such parameters for the system, that it behaves close to some single oscillator under all possible inputs.
In this example we use a sample of size 10.

First, we estimate the distance and optimize only specification parameters by running behavioural_distance.
Then we tune the system to specification using different optimizers.
We verify the tuning by resampling and tuning to the new sample.
All systems have the same set of inputs, as we fix the random seed.
=#
function tune_run(omega, scale; save_plots = false)
    Random.seed!(1)
    kur = PBTExample.create_kuramoto_example(omega*scale, N_osc, size_p_spec, K_av, t_steps, N_samples)

    d, p_dist, = behavioural_distance(kur, p_initial; abstol=1e-4, reltol=1e-4,
                                        optimizer=DiffEqFlux.BFGS(),
                                        optimizer_options=(
                                            :maxiters => 10,
                                            :cb => PBTLibrary.basic_pbt_callback))
    d_initial = d
    println("Initial distance to specified behaviour is lower than $d\n")
    ##
    scen = 1:3
    ##
    plot_callback(kur, p_dist, d, scenario_nums = scen, title = "Scenarios $scen, s = $scale", offset = 0.5, xlims = (kur.t_span[2]/2, kur.t_span[2]))
    if save_plots
        savefig(path*string(round(scale, digits = 4))*"_initial_"*string(round(d, digits = 4))*".png")
    end
    ##
    res = pbt_tuning(kur, p_dist; abstol=1e-4, reltol=1e-4)
    p_tuned = res.minimizer
    ##
    d2, p_dist_2, = behavioural_distance(kur, p_tuned; abstol=1e-4, reltol=1e-4)

    println("After a first tuning using ADAM(0.01) the distance to specified behaviour is lower than $d2\n")

    ##

    res = pbt_tuning(kur, p_dist_2; abstol=1e-4, reltol=1e-4,
                    optimizer=DiffEqFlux.BFGS(),
                    optimizer_options=(
                        :maxiters => 5,
                        :cb => PBTLibrary.basic_pbt_callback))

    p_t2 = res.minimizer
    ##
    plot_callback(kur, p_t2, res.minimum, scenario_nums = scen, title = "Scenarios $scen, s = $scale", offset = 0.5, xlims = (kur.t_span[2]/2, kur.t_span[2]))

    if save_plots
        savefig(path*string(round(scale, digits = 4))*"_tuned_"*string(round(res.minimum, digits = 4))*".png")
    end
    ##
    d3, p_dist_3, = behavioural_distance(kur, p_t2; abstol=1e-4, reltol=1e-4, optimizer_options=(:maxiters => 1000,))

    println("After a second tuning using BFGS the distance to specified behaviour is lower than $d3\n")
    ##
    plot_callback(kur, p_dist_3, d3, scenario_nums = scen, title = "Scenarios $scen, s = $scale", offset = 0.5, xlims = (kur.t_span[2]/2, kur.t_span[2]))
    if save_plots
        savefig(path*string(round(scale, digits = 4))*"_tuned2_"*string(round(d3, digits = 4))*".png")
    end
    resample!(kur)
    d4, p_dist_4, = behavioural_distance(kur, p_dist_3; abstol=1e-4, reltol=1e-4,
                    optimizer = DiffEqFlux.BFGS(),
                    optimizer_options=(
                        :maxiters => 10,
                        :cb => PBTLibrary.basic_pbt_callback))

    d5, p_dist_5, = behavioural_distance(kur, p_dist_4; abstol=1e-4, reltol=1e-4,
                    optimizer = DiffEqFlux.AMSGrad(0.01),
                    optimizer_options=(
                        :maxiters => 1000,))

    plot_callback(kur, p_dist_5, d5, scenario_nums = scen, title = "Scenarios $scen, s = $scale", offset = 0.5, xlims = (kur.t_span[2]/2, kur.t_span[2]))
    if save_plots
        savefig(path*string(round(scale, digits = 4))*"_final_rs_"*string(round(d5, digits = 4))*".png")
    end
    println("After resampling the distance to specified behaviour is lower than $d4\n")
    d_final = d5
    p_final = p_dist_5
    d_initial, d_final, p_final
end
##
#=
Different scales represent degree of synchronisity, the higher the number the less synchronous the system is.
Non-synchronous systems are more difficult to tune, so we chose a non-uniform distribution of scales.
Thus we can explore the large range of possible systems, while keeping the computation time not too high.
=#
scales = vcat(1.0:0.1:1.5,2.0:0.5:5.)

##
#=
Run tune_run for the set of systems, defined by scales.
res[i] contains initial and final distances to the specification, as well as the final parameters of the tuned system.
=#
res = [tune_run(omega, s, save_plots = save_plots) for s in scales]

##
#=
We want to evaluate performance of the tuning on all systems.
q = d_initial/d_final
First we plot the improvement depending on scale.
Secondly, we look at the absolute value of final distance.
=#
q = [r[1]/r[2] for r in res]

plot(scales, q)
if save_plots
    savefig(path*"improvement(scale).png")
end

plot(scales, [r[2] for r in res], xlabel = "s", ylabel = "d", label = "Final distance to specification")
if save_plots
    savefig(path*"final_d(omega).png")
end
##
#=
By running tune_run on all systems, we kept tuning parameters constant for all of them.
However, the chosen tuning parameters were limited by performance on systems with larger frequency spread.
In other words, we stop tuning when the systems with larger s start to become unstable.
At that point we could still continue tuning for more synchronous systems to show the best tuning result.
We chose to continue optimization for s = 1.2
=#
kur_new = PBTExample.create_kuramoto_example(omega*1.2, N_osc, size_p_spec, K_av, t_steps, N_samples)
p_initial_new = res[3][3] # for scale = 1.2

d, p_dist, = behavioural_distance(kur_new, p_initial_new; abstol=1e-4, reltol=1e-4,
                    optimizer = DiffEqFlux.BFGS(),
                    optimizer_options=(
                        :maxiters => 10,
                        :cb => PBTLibrary.basic_pbt_callback))
##
scen = 1:3
plot_callback(kur_new, p_dist, d, scenario_nums = scen, title = "Scenarios $scen, s = 1.2", offset = 0.5, xlims = (kur_new.t_span[2]/2, kur_new.t_span[2]))

##
d, p_dist, = behavioural_distance(kur_new, p_initial_new; abstol=1e-4, reltol=1e-4,
                    optimizer = DiffEqFlux.BFGS(),
                    optimizer_options=(
                        :maxiters => 10,
                        :cb => PBTLibrary.basic_pbt_callback))

plot_callback(kur_new, p_dist, d, scenario_nums = scen, title = "Scenarios $scen, s = 1.2", offset = 0.5, xlims = (kur_new.t_span[2]/2, kur_new.t_span[2]))
##
res = pbt_tuning(kur_new, p_dist)

plot_callback(kur_new, res.minimizer, res.minimum, scenario_nums = scen, title = "Scenarios $scen, s = 1.2", offset = 0.5, xlims = (kur_new.t_span[2]/2, kur_new.t_span[2]))
if save_plots
    savefig(path*"1.2_tuned_3_"*string(round(res.minimum, digits = 4))*".png")
end

resample!(kur_new)

d_rs, p_dist_rs = behavioural_distance(kur_new, res.minimizer; abstol=1e-4, reltol=1e-4,
optimizer = DiffEqFlux.BFGS(),
optimizer_options=(
    :maxiters => 10,
    :cb => PBTLibrary.basic_pbt_callback))

plot_callback(kur_new, p_dist_rs, d_rs, scenario_nums = scen, title = "Scenarios $scen, s = 1.2", offset = 0.5, xlims = (kur_new.t_span[2]/2, kur_new.t_span[2]))
if save_plots
    savefig(path*"1.2_rs_further_"*string(round(d_rs, digits = 4))*".png")
end
##
res2 = pbt_tuning(kur_new, p_dist_rs)

plot_callback(kur_new, res2.minimizer, res2.minimum, scenario_nums = scen, title = "Scenarios $scen, s = 1.2", offset = 0.5, xlims = (kur_new.t_span[2]/2, kur_new.t_span[2]))
if save_plots
    savefig(path*"1.2_rtuned4_"*string(round(res2.minimum, digits = 4))*".png")
end
##
resample!(kur_new)

d_rs, p_dist_rs, losses = behavioural_distance(kur_new, res.minimizer; abstol=1e-4, reltol=1e-4,
optimizer = DiffEqFlux.BFGS(),
optimizer_options=(
    :maxiters => 10,
    :cb => PBTLibrary.basic_pbt_callback))

plot_callback(kur_new, p_dist_rs, d_rs, scenario_nums = scen, title = "Scenarios $scen, s = 1.2", offset = 0.5, xlims = (kur_new.t_span[2]/2, kur_new.t_span[2]))
if save_plots
    savefig(path*"1.2_rs_final_"*string(round(d_rs, digits = 4))*".png")
end


module A5P1

#Use this to print
#=
weave(Pkg.dir("$(pwd())\\A5P1.jl"),
  informat="script",
  out_path = "$(pwd())\\A5P1.html",
  doctype = "md2html")
=#
if pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

using DataFrames, Distributions, StatsBase, GZip, JLD,
  Gadfly, CTModCopy, JuMP, NLopt, ForwardDiff, Roots

const DATA_PATH = pwd() * "\\data" #path where we store the data
const OUTPUT_PATH = pwd() * "\\output"
const DATE_FORMAT_STR = "m/d/y"
const SP500_NAME = "GSPC"
const TBILL_NAME = "t-bill"
const TERM_NAME = "termStructure"
const FOOTER_NAME = "footer.tex" #these are just headers for the tex tables
const HEADER_NAME = "header.tex"
const MAX_YEAR = 2016
const IN_SAMPLE_CUTOFF_YEAR = 1975

const DEFAULT_NUM_SIMS = 10000
const DEFAULT_LOGμ_TOL = 10.0^-10.0
const PlotContainer = Union{Plot,Gadfly.Compose.Context}
const CSRegSpec = Tuple{Symbol, Symbol, Int}
const CPRegSpec = Tuple{Symbol, Int}
const TOLERANCE = 10.0^-8.0

const MIN_VALS_FOR_EST = 10

abstract type AbstractARCH end

function preProcessSP500()::Void
  dateFormat::String = DATE_FORMAT_STR
  sp500DF::DataFrame = readtable("$DATA_PATH\\$SP500_NAME.csv")

  #fix the dates
  rename!(sp500DF, :Date, :DATE)
  sp500DF[:,:DATE] = ((s::String)->Date(s, dateFormat)::Date).(sp500DF[:,:DATE])::DataVector{Date}
  sort!(sp500DF, cols = [:DATE])
  sp500DF[:,:returns] = similar(sp500DF[:,:Adj_Close])
  sp500DF[2:end,:returns] = log.(sp500DF[2:end,:Adj_Close] ./  sp500DF[1:(end-1),:Adj_Close])

  sp500DF[:,:dayOfQuarter] = zeros(Int,size(sp500DF,1))
  for i::Int ∈ 2:size(sp500DF,1)
    if Dates.firstdayofquarter(sp500DF[i,:DATE]) ≠ Dates.firstdayofquarter(sp500DF[i-1,:DATE])
      sp500DF[i,:dayOfQuarter] = 1
    else
      sp500DF[i,:dayOfQuarter] = sp500DF[i-1,:dayOfQuarter] + 1
    end
  end
  #showcols(sp500DF)

  stream::IOStream = open("$DATA_PATH\\$SP500_NAME.jls", "w")
  serialize(stream, sp500DF)
  close(stream)

  return nothing
end

#holds a GARCH parameter object
mutable struct ΘGARCH <: AbstractARCH
  λ::Real
  β::Real
  α::Real
  μ::Real
  valid::Bool

end

#constructor
ΘGARCH(λ::Real, β::Real, α::Real, μ::Real)::ΘGARCH = ΘGARCH(λ, β, α, μ, true)

mutable struct ΘEGARCH <: AbstractARCH
  λ::Real
  β::Real
  α::Real
  μ::Real
  γ::Real
  valid::Bool

end

#constructor
ΘEGARCH(λ::Real, β::Real, α::Real, μ::Real, γ::Real)::ΘEGARCH = ΘEGARCH(λ, β, α, μ, γ, true)

function estimateGARCH(r::Vector{Float64}, verbose::Bool = false)::ΘGARCH

  #for convenience
  T::Int = length(r)

  function obj(λ::Real, β::Real, α::Real, μ::Real)

    #σ²::Vector{Real} = Vector{Real}(T-1)
    σ²tM1::Real = λ/(1.0-β)
    #l::Float64 = -(r[1]-μ)*(r[1]-μ)/(2.0*σ²₀)-0.5*log(σ²₀) #initial likelihood
    l::Real = 0.0

    #calculate the likelihood and get the next volatiltiy in the sequence
    for i::Int ∈ 1:T
      u²::Real = (r[i]-μ)*(r[i]-μ)
      l += -u²/σ²tM1 - 0.5*log(σ²tM1)
      σ²tM1 = λ+β*σ²tM1+α*u²
    end

    #println("λ=$λ, β=$β, α=$α, μ=$μ")

    return l
  end

  meanRet = mean(r)

  #set initial meta parameters
  minVal = 10.0^-10.0
  maxVal = 1.0-minVal

  #Algorithms: :LD_LBFGS, :LN_COBYLA, :LN_BOBYQA
  #  :LD_SLSQP, :LD_TNEWTON, :LD_VAR2/:LD_VAR1,
  #  :LD_MMA, :GN_ESCH, :GN_DIRECT, GN_DIRECT_L, :GD_STOGO
  # GN_ISRES, :LN_NELDERMEAD, :LN_SBPLX,
  #GN_CRS2_LM
  #:GN_CRS2_LM***, :GN_ISRES*** (slow), :LD_MMA**, LN_SBPLX**, LD_TNEWTON** best:62,848
  #MMA: λ=2.991604487710108e-6, β=0.8691809689019686, α=0.2447041114406467, μ=0.0005059861641042039
  #NEWTON: 3.7136249890797323e-6, β=0.8641416287811661, α=0.23581482405030202, μ=0.0005248887137138578
  #CRS2 λ=1.8304983359585568e-6, β=0.9091365756370056, α=0.1664779062477393, μ=0.00048124017464920696
  #SBPLX λ=1.7187359905806954e-6, β=0.9122999784465299, α=0.16106473023211287, μ=0.00029331108088162115
  #LR Vol>0 #:LN_COBYLA(61450), LD_MMA(61144), ISRES(62578),

  mod = Model(solver=NLoptSolver(algorithm=:LD_TNEWTON_PRECOND,maxtime=10))

  @variable(mod, minVal <= λ <= maxVal)
  @variable(mod, minVal <= β <= maxVal)
  @variable(mod, minVal <= α <= maxVal)
  @variable(mod, minVal <= μ <= maxVal)
  #@NLconstraint(mod, 1.0-α-β>=0.0)

  setvalue(λ, 0.5)
  setvalue(β, 0.5)
  setvalue(α, 0.5)
  setvalue(μ, mean(r)) #guess the unconditional mean

  JuMP.register(mod, :obj, 4, obj, autodiff=true)
  @NLobjective(mod, Max, obj(λ, β, α, μ))

  ΘBest::ΘGARCH  = ΘGARCH(0.0, 0.0, 0.0, 0.0, false)
  objBest::Float64 = -9999.0
  try
    output = solve(mod)
    ΘBest  = ΘGARCH(getvalue(λ), getvalue(β), getvalue(α), getvalue(μ))
    objBest= obj(getvalue(λ), getvalue(β), getvalue(α), getvalue(μ))
  end

  if verbose
    impliedAnnualLRVol::Float64 = (ΘBest.λ/(1.0 - ΘBest.β)*255)^0.5
    historicalLRVol::Float64 = std(r)*255.0^0.5

    impliedReturn::Float64 = exp(ΘBest.μ*255)-1.0
    historicalReturn::Float64 = exp(mean(r)*255)-1.0

    println("λ=$(ΘBest.λ), β=$(ΘBest.β), α=$(ΘBest.α), μ=$(ΘBest.μ)")
    println("Vol: Implied=$impliedAnnualLRVol, Hist=$historicalLRVol")
    println("Return: Implied=$impliedReturn, Hist=$historicalReturn")
    println("Objective: $objBest")
    println("For Reference: x̄=$(mean(r))")
  end

  return ΘBest

end

function estimateEGARCH(r::Vector{Float64}, verbose::Bool = false)::ΘEGARCH

  T::Int = length(r)

  function obj(λ::Real, β::Real, α::Real, μ::Real, γ::Real)

    #σ²::Vector{Real} = Vector{Real}(T-1)
    EAbsZ::Float64 = (2.0/π)^0.5
    lσ²tM1::Real = (λ-γ*EAbsZ)/(1.0-β)
    #lσ²tM1::Real = λ/(1.0-β)

    l::Real = 0.0

    #calculate the likelihood and get the next volatiltiy in the sequence
    for i::Int ∈ 1:T
      u²::Real = (r[i]-μ)*(r[i]-μ)
      σtM1::Real = (exp(lσ²tM1))^0.5
      l += -u²/(σtM1*σtM1) - 0.5*lσ²tM1
      lσ²tM1 = λ+β*lσ²tM1+α*u²+γ*(abs((r[i]-μ)/σtM1)-EAbsZ)
    end

    return l
  end

  meanRet = mean(r)

  #set initial meta parameters
  minVal = -10.0
  maxVal = 100.0

  #Algorithms: :LD_LBFGS, :LN_COBYLA, :LN_BOBYQA
  #  :LD_SLSQP, :LD_TNEWTON, :LD_VAR2/:LD_VAR1,
  #  :LD_MMA, :GN_ESCH, :GN_DIRECT, GN_DIRECT_L, :GD_STOGO
  # GN_ISRES, :LN_NELDERMEAD, :LN_SBPLX,
  #GN_CRS2_LM
  #:LD_LBFGS(62774), LD_VAR2(62774), LN_NELDERMEAD(61880), :LN_SBPLX (62791), :LD_TNEWTON(62800)

  mod = Model(solver=NLoptSolver(algorithm=:LD_TNEWTON_PRECOND, maxtime=20))

  @variable(mod, minVal <= λ <= maxVal)
  @variable(mod, minVal <= β <= maxVal)
  @variable(mod, minVal <= α <= maxVal)
  @variable(mod, 0.0 <= μ <= maxVal)
  @variable(mod, minVal <= γ <= maxVal)
  #@NLconstraint(mod, 1.0-α-β>=0.0)

  setvalue(λ, 0.5)
  setvalue(β, 0.5)
  setvalue(α, 0.5)
  setvalue(μ, mean(r)) #guess the unconditional mean
  setvalue(γ, 0.5) #guess the unconditional mean

  JuMP.register(mod, :obj, 5, obj, autodiff=true)
  @NLobjective(mod, Max, obj(λ, β, α, μ, γ))

  ΘBest::ΘEGARCH  = ΘEGARCH(getvalue(λ), getvalue(β), getvalue(α), getvalue(μ), getvalue(γ), false)
  objBest::Float64 = -9999.0
  try
    output = solve(mod)
    ΘBest  = ΘEGARCH(getvalue(λ), getvalue(β), getvalue(α), getvalue(μ), getvalue(γ))
    objBest= obj(getvalue(λ), getvalue(β), getvalue(α), getvalue(μ), getvalue(γ))
  end

  if verbose
    impliedAnnualLRVol::Float64 = (exp((ΘBest.λ-ΘBest.γ*(2/π)^0.5)/(1.0-ΘBest.β))*255)^0.5
    #impliedAnnualLRVol::Float64 = (exp(ΘBest.λ)/(1-ΘBest.β)*255)^0.5
    historicalLRVol::Float64 = std(r)*255.0^0.5

    impliedReturn::Float64 = exp(ΘBest.μ*255)-1.0
    historicalReturn::Float64 = exp(mean(r)*255)-1.0

    println("λ=$(ΘBest.λ), β=$(ΘBest.β), α=$(ΘBest.α), μ=$(ΘBest.μ), γ=$(ΘBest.γ)")
    println("Vol: Implied=$impliedAnnualLRVol, Hist=$historicalLRVol")
    println("Return: Implied=$impliedReturn, Hist=$historicalReturn")
    println("Objective: $objBest")
    println("For Reference: x̄=$(mean(r))")
  end

  return ΘBest

end

function estimateGARCHByQ(sp500DF::DataFrame, estimateModel::Function)#::T where T <: AbstractARCH

  dateDict::Dict = Dict(sp500DF[i,:DATE]=>i for i::Int ∈ 1:size(sp500DF,1))

  minValsForEst::Int = MIN_VALS_FOR_EST
  r::Vector{Float64} = sp500DF[2:end,:returns]

  #get a list of target indices
  quarterlyDates::Vector{Date} = sp500DF[sp500DF[:,:dayOfQuarter].==1,:DATE]
  terminalIndicesToEst::Vector{Int} = ((d::Date)->dateDict[d]).(quarterlyDates)
  #terminalIndicesToEst = terminalIndicesToEst[terminalIndicesToEst[:].>minValsForEst]

  #set up a container of contianers
  estimated::Vector{Vector{T where T<:AbstractARCH}} =
    [Vector{T where T<:AbstractARCH}() for i::Int ∈ 1:length(terminalIndicesToEst)]

  println("beginning $(length(estimated)) simulations")
  fails::Int = 0
  @fastmath for i ∈ 1:length(terminalIndicesToEst)

  #estimate the grabbed indices
    estimated[i] = [estimateModel(r[1:(terminalIndicesToEst[i]-1)], false)]
    if estimated[i][1].valid == false
      fails += 1
    end
  end
  println("simulations complete. $fails failures out of $(length(estimated)) simulations.")

  return [estimated[i][1] for i::Int ∈ 1:length(terminalIndicesToEst)]

end

#script function for running the GARCH problem
function A5P1Script(;refreshData::Bool=true, refreshGARCH::Bool=true, refreshEGARCH::Bool=true,)
  if !isfile("$DATA_PATH\\$SP500_NAME.jls") || refreshData
    preProcessSP500()
  end

  #read the binary file
  stream::IOStream = open("$DATA_PATH\\$SP500_NAME.jls")
  sp500DF::DataFrame = deserialize(stream)
  close(stream)

  r::Vector{Float64} = sp500DF[2:end,:returns]
  quarterlyDates::Vector{Date} = sp500DF[sp500DF[:,:dayOfQuarter].==1,:DATE]

  #estimate the GARCH model and print the output

  estimateGARCH(r, true)
  estimateEGARCH(r, true)

  #this setup is so that we don't ahve to solve the model each quarter repeatedly
  if refreshGARCH
    estimatedGARCH::Vector{ΘGARCH} = estimateGARCHByQ(sp500DF, estimateGARCH)

    #cache the results on disk
    stream = open("$OUTPUT_PATH\\GARCH_RESULTS.jls", "w")
    serialize(stream, estimatedGARCH)
    close(stream)
  else

    #access the resutls
    stream = open("$OUTPUT_PATH\\GARCH_RESULTS.jls")
    estimatedGARCH = deserialize(stream)
    close(stream)
  end

  if refreshEGARCH
    estimatedEGARCH::Vector{ΘEGARCH} = estimateGARCHByQ(sp500DF, estimateEGARCH)

    #cache the results on disk
    stream = open("$OUTPUT_PATH\\EGARCH_RESULTS.jls", "w")
    serialize(stream, estimatedEGARCH)
    close(stream)
  else

    #access the resutls
    stream = open("$OUTPUT_PATH\\EGARCH_RESULTS.jls")
    estimatedEGARCH = deserialize(stream)
    close(stream)
  end

  ##############Now work on the outputs. First make them into a dataframe
  garchDF::DataFrame = DataFrame(DATE=quarterlyDates, λ=[g.λ for g ∈ estimatedGARCH], β=[g.β for g ∈ estimatedGARCH],
    α=[g.α for g ∈ estimatedGARCH], μ=[g.μ for g ∈ estimatedGARCH],
    μLR=[exp(255*g.μ)-1.0 for g ∈ estimatedEGARCH], σ=[(g.λ/(1.0 - g.β)*255)^0.5 for g ∈ estimatedGARCH])

  egarchDF::DataFrame  = DataFrame(DATE=quarterlyDates, λ=[g.λ for g ∈ estimatedEGARCH], β=[g.β for g ∈ estimatedEGARCH],
    α=[g.α for g ∈ estimatedEGARCH], μ=[g.μ for g ∈ estimatedEGARCH], γ=[g.γ for g ∈ estimatedEGARCH],
    μLR=[exp(255*g.μ)-1.0 for g ∈ estimatedEGARCH], σ=[(exp((g.λ-g.γ*(2/π)^0.5)/(1.0-g.β))*255)^0.5 for g ∈ estimatedEGARCH])

  nQuarters::Int = length(quarterlyDates)

  for i::Int∈1:nQuarters
    if estimatedGARCH[i].valid == false
      garchDF[i,[:λ, :β, :α, :μ, :σ]] = NA
    end

    if estimatedEGARCH[i].valid == false
      egarchDF[i,[:λ, :β, :α, :μ,:γ, :σ]] = NA
    end
  end

  Gadfly.push_theme(:default)
  plots::Vector{PlotContainer} = Vector{PlotContainer}()
  plotNames::Vector{String} = Vector{String}()

  push!(plotNames, "GarchParameters")
  push!(plots, vstack(plot(x=garchDF[:,:DATE], y=garchDF[:,:λ],
    Geom.point, #Theme(line_width = 0.25pt),
    Guide.ylabel("Value"), Guide.xlabel("Date"),
    Guide.title("GARCH λ Param")),
    plot(x=garchDF[:,:DATE], y=garchDF[:,:β],
       Geom.point, #Theme(line_width = 0.25pt),
      Guide.ylabel("Value"), Guide.xlabel("Date"),
      Guide.title("GARCH β Param")),
    plot(x=garchDF[:,:DATE], y=garchDF[:,:α],
      Geom.point, #Theme(line_width = 0.25pt),
      Guide.ylabel("Value"), Guide.xlabel("Date"),
      Guide.title("GARCH α Param"))))

  egarchDF = egarchDF[((!isna).(egarchDF[:,:σ])),:]

  push!(plotNames, "EGarchParameters")
  push!(plots, vstack(plot(x=egarchDF[:,:DATE], y=egarchDF[:,:λ],
    Geom.point, #Theme(line_width = 0.25pt),
    Guide.ylabel("Value"), Guide.xlabel("Date"),
    Guide.title("GARCH λ Param")),
    plot(x=egarchDF[:,:DATE], y=egarchDF[:,:β],
       Geom.point, #Theme(line_width = 0.25pt),
      Guide.ylabel("Value"), Guide.xlabel("Date"),
      Guide.title("GARCH β Param")),
    plot(x=egarchDF[:,:DATE], y=egarchDF[:,:α],
      Geom.point, #Theme(line_width = 0.25pt),
      Guide.ylabel("Value"), Guide.xlabel("Date"),
      Guide.title("GARCH α Param"))))

  push!(plotNames, "GarchVol")
  push!(plots, vstack(plot(x=garchDF[:,:DATE], y=garchDF[:,:μLR],
    Geom.point, #Theme(line_width = 0.25pt),
    Guide.ylabel("Value"), Guide.xlabel("Date"), Coord.Cartesian(ymin=0.0,ymax=0.5),
    Guide.title("GARCH μ Param")),
    plot(x=garchDF[:,:DATE], y=garchDF[:,:σ],
       Geom.point, #Theme(line_width = 0.25pt),
      Guide.ylabel("Value"), Guide.xlabel("Date"),
      Guide.title("GARCH σ Param"))))

  push!(plotNames, "EGarchVol")
  push!(plots, vstack(plot(x=egarchDF[:,:DATE], y=egarchDF[:,:μLR],
    Geom.point, #Theme(line_width = 0.25pt),
    Guide.ylabel("Value"), Guide.xlabel("Date"), Coord.Cartesian(ymin=0.0,ymax=0.5),
    Guide.title("GARCH μ Param")),
  plot(x=egarchDF[:,:DATE], y=egarchDF[:,:σ],
     Geom.point, #Theme(line_width = 0.25pt),
    Guide.ylabel("Value"), Guide.xlabel("Date"),
    Guide.title("GARCH σ Param"))))

  for i ∈ 1:length(plots)
    draw(SVG("$OUTPUT_PATH\\$(plotNames[i]).svg", 7.5inch, 9inch),plots[i])
  end

  #graph the GARCH model

end

@time begin
  #A5P1Script(refreshData=false, refreshGARCH=false, refreshEGARCH=false)
end

end


"""

Plotting of symbolic objects.

The `Plots` package provide a uniform interface to many of `Julia`'s
plotting packages. `SymPy` plugs into `Plots`' "recipes."

The basic goal is that when `Plots` provides an interface for function
objects, this package extends the interface to symbolic expressions.

In particular:

* `plot(ex::Sym, a, b; kwargs...)` will plot a function evaluating `ex` over [a,b]

Example. Here we use the default backend for `Plots` to make a plot:

```
using Plots
@vars x
plot(x^2 - 2x, 0, 4)
```

* `plot(exs::Vector{Sym}, a, b; kwargs...)` will plot the functions evaluating `exs` over [a,b]

Example:

```
@vars x
plot([sin(x), cos(x)], 0, 2pi)
```

* `plot(ex1, ex2, a, b; kwargs...)` will plot the two expressions in a parametric plot over the interval `[a,b]`.

Example:

```
@vars x
plot(sin(2x), cos(3x), 0, 4pi) ## also
```

For a few backends (those that support `:path3d`) a third symbolic
expression may be added to have a 3d parametric plot rendered:

```
plot(sin(x), cos(x), x, 0, 4pi) # helix in 3d
```

* `plot(xs, ys, expression)` will make a contour plot (for many backends).

```
@vars x y
plot(linspace(0,5), linspace(0,5), x*y)
```

* To plot the surface  `z=ex(x,y)` over a region we have `Plots.surface`. For example,

```
@vars x y
surface(-5:5, -5:5, 25 - x^2 - y^2)
```

* a vectorfield plot can (inefficiently but directly) be produced following this example:

```
function vfieldplot(fx, fy; xlim=(-5,5), ylim=(-5,5), n=7)
    xs = linspace(xlim..., n)
    ys = linspace(ylim..., n)

    us = vec([x for x in xs, y in ys])
    vs = vec([y for x in xs, y in ys])
    fxs = vec([fx(x,y) for x in xs, y in ys])
    fys = vec([fy(x,y) for x in xs, y in ys])

    quiver(us, vs, quiver=(fxs, fys))
end
fx = (x + y) / sqrt(x^2 + y^2)
fy = (x - y) / sqrt(x^2 + y^2)
vfieldplot(fx, fy)
```

----

Some graphics provided by `SymPy` are available if `PyPlot` is installed.

* `plot_parametric_surface(exs::Tuple, (uvar,a0,b0), (vvar,a1,b1);
  kwargs)` will make parametrically defined surface plots.

Plot the parametrically defined surface `[exs[1](u,v), exs[2](u,v), exs[3](u,v)]` over `[a0,a1] x
[b0,b1]`. The specification of the variables uses a tuple of the form
`(Sym, Real, Real)` following the style of SymPy in `integrate`, say,
where disambiguation of variable names is needed.

```
@vars theta, phi
r = 1
plot_parametric_surface((r*sin(theta)*sin(phi), r*sin(theta)*cos(phi), r*cos(theta)),
                        (theta, 0, pi), (phi, 0, pi/2))
```

(The SymPy name for this function is `plot3d_parametric_surface`, we have dropped the "`3d`" part.)

* `plot_implicit(equation, (xvar, x0, x1), (yvar, y0, y1))` will plot implicitly the equation.

```
@syms x y
plot_implicit(Eq(x^2+ y^2,3), (x, -2, 2), (y, -2, 2))  # draw a circle
```

"""
sympy_plotting = nothing
export sympy_plotting

## Recipes for hooking into Plots

using RecipesBase

##
@recipe f{T<:Sym}(::Type{T}, v::T) = lambdify(v)

## for vectors of expressions
@recipe f{S<:AbstractVector{Sym}}(::Type{S}, ss::S) = Function[lambdify(s) for s in ss]

## vectorfieldplot is an addon as we avoid quiver
## Not ready yet!
## @recipe function f(::Type{Val{:vectorfieldplot}}, fx, fy, xlim=(-5,5), ylim=(-5,5); n=15)

##     x₀, x₁ = xlim
##     y₀, y₁ = ylim
##     Δx = (x₁ - x₀) / (n-1)
##     Δy = (y₁ - y₀) / (n-1)

##     xs = x₀:Δx:x₁
##     ys = y₀:Δy:y₁

##     mx, my = 0.0, 0.0

##     ## two loops, first to identify scaling factor so
##     ## lines stay within box of width Δx by Δy
##     for x in xs, y in ys
##         mx = max(mx, abs(fx(x,y)))
##         my = max(my, abs(fy(x,y)))
##     end

##     ## we want all lines to be in the box, so we scale
##     λ = .95 *  min(Δx/mx, Δy/my)

##     XS = Float64[]
##     YS = Float64[]

##     for x in xs, y in ys
##         append!(XS, [x, x + λ * fx(x,y), NaN])
##         append!(YS, [y, y + λ * fy(x,y), NaN])
##     end
##     pop!(XS); pop!(YS)

##     seriestype := :path
##     xlims := xlim
##     ylims := ylim
##     legend := false

##     x := XS
##     y := YS

##     ()
## end

## function vectorfieldplot(fx::Sym, fy::Sym, xlim=(-5,5), ylim=(-5,5); n::Int=15, kwargs...)
##     xs = free_symbols([fx, fy])
##     vectorfieldplot(lambdify(fx, xs) , lambdify(fy, xs), xlim, ylim; n=n, kwargs...)
## end

## function vectorfieldplot!(fx::Sym, fy::Sym, xlim=(-5,5), ylim=(-5,5); n::Int=15, kwargs...)
##     xs = free_symbols([fx, fy])
##     vectorfieldplot!(lambdify(fx, xs), lambdify(fy, xs), xlim, ylim; n=n, kwargs...)
## end

## function vectorfieldplot(fx, fy, xlim=(-5,5), ylim=(-5,5); n::Int=15, kwargs...)
##     try
##         eval(Expr(:using, :Plots)) # XXX Super hacky way to do this
##     catch err
##         warn("Can't plot without the `Plots` package.")
##     end
##     plot(Val{:vectorfieldplot}, fx, fy, xlim, ylim; n=n, kwargs...)
## end

## function vectorfieldplot!(fx, fy, xlim=(-5,5), ylim=(-5,5); n::Int=15, kwargs...)
##     try
##         eval(Expr(:using, :Plots)) # XXX Super hacky way to do this
##     catch err
##         warn("Can't plot without the `Plots` package.")
##     end
##     plot!(Val{:vectorfieldplot}, fx, fy, xlim, ylim; n=n, kwargs...)
## end
## export vectorfieldplot, vectorfieldplot!

## ---------------------

## These functions give acces to SymPy's plotting module. They will work if PyPlot is installed, but may otherwise cause an error

## surface plot xvar = Tuple(Sym, Real, Real)
##
"""

Render a parametrically defined surface plot.

Example:
```
@vars u, v
plot_parametric_surface((u*v,u-v,u+v), (u,0,1), (v,0,1))
```

This uses `PyPlot`, not `Plots` for now.
"""
function plot_parametric_surface(exs::(@compat Tuple{Sym,Sym,Sym}),
                                 xvar=(-5.0, 5.0),
                                 yvar=(-5.0, 5.0),
                                 args...;
                                 kwargs...)

    SymPy.call_sympy_fun(sympy[:plotting][:plot3d_parametric_surface], exs..., args...; kwargs...)

end
export plot_parametric_surface

"""
Plot an implicit equation

```
@syms x y
plot_implicit(Eq(x^2+ y^2,3), (x, -2, 2), (y, -2, 2))
```

"""
plot_implicit(ex, args...; kwargs...) = SymPy.call_sympy_fun(sympy[:plotting][:plot_implicit], ex, args...; kwargs...)
export plot_implicit


using Test
using CVChannel

"""
This script verifies that the entanglement-assisted communication value for
axisymmetric channels is bound as `eaCV(J_N) ≤ d * CV(J_N)`, where `d` is the
Hilbert space dimension of the channel's output. PPT axisymmetric states are also
separable, therefore, we can compute the communication value with `pptCVPrimal`.
"""

"""
    _verify_eaCV_bound_for_axisymmetric_channels(
        d::Int64,
        x_step::Float64,
        y_step::Float64
    )

Scans across the `d`-dimensional axisymmetric channel parameters evaluating both
the entanglement-assisted CV (`eaCVPrimal`) and the CV (`pptCVPrimal`).
The two axisymmetric parameters x and y are scanned across with step-sizes `x_step`
and `y_step` respectively.
For each channel, it is verified that the entanglement-assisted CV is bound by
`eaCV ≦ d * CV`.
"""
function _verify_eaCV_bound_for_axisymmetric_channels(d::Int64, x_step::Float64, y_step::Float64)
    y_bounds = CVChannel._axisymmetric_y_bounds(d)

    it = 0
    y_range = y_bounds[1]:y_step:y_bounds[2]
    for y in y_range
        it = it + 1
        println("(d = $d) verifying y-slice $it of $(length(y_range))")
        x_constraints = CVChannel._axisymmetric_x_constraints(d,y)
        for x in x_constraints[1]:x_step:x_constraints[2]
            @testset "(x,y) = ($x,$y)" begin
                ρ_axi = axisymmetricState(d,x,y)

                J_N = d*ρ_axi
                (ea_cv_N, ea_σAB_N) = eaCVPrimal(J_N, d, d)
                (ppt_cv_N, ppt_σAB_N) = pptCVPrimal(J_N, d, d)

                @test ea_cv_N ≤ d * ppt_cv_N || ea_cv_N ≈ d * ppt_cv_N
            end
        end
    end
end

println("Verifying bounds on entanglement-assisted CV")
@testset "Coarse-grained verification of entanglement-assisted CV bounds (d = $d)" for d in 2:8
    y_step = 0.1
    x_step = 0.2

    _verify_eaCV_bound_for_axisymmetric_channels(d,x_step,y_step)
end

@testset "Fine-grained verification of entanglement-assisted CV bounds (d = $d)" for d in 2:4
    y_step = 0.01
    x_step = 0.01

    _verify_eaCV_bound_for_axisymmetric_channels(d,x_step,y_step)
end


"A grid corresponding to an unstructured collection of points."
struct ScatteredGrid{T} <: AbstractGrid{T,1}
    points     ::  Vector{T}
    domain     ::  Domain
    ScatteredGrid(points::Vector{T}, domain=DomainSets.FullSpace{T}()) where T =
        new{T}(points, domain)
end

size(g::ScatteredGrid) = (length(g.points),)

unsafe_grid_getindex(g::ScatteredGrid, idx) = (@inbounds g.points[idx])

covering(g::ScatteredGrid) = g.domain

## Some functionality for random grids

"""
Compute a scattered grid of M points randomly distributed in `Ω`, using
the uniform probability measure on `Ω`.
"""
function randomgrid(Ω::Domain, M::Int)
    points = [randompoint(Ω, boundingbox(Ω)) for m in 1:M]
    ScatteredGrid(sort(points))
end

"Generate a single random point inside the given box, with `eltype` `T`."
randompoint(dom::ProductDomain) =
    DomainSets.toexternalpoint(dom, map(randompoint,components(dom)))

# Don't return an SVector in 1d, just a value
function randompoint(dom::AbstractInterval)
    T = prectype(dom)
    val = rand(T)
    convert(T,val * infimum(dom) + (1-val) * supremum(dom))
end

"""
Generate a single random point inside the given domain, with `eltype` `T`.
Random points are generated inside the given box, until one is inside the domain.
"""
function randompoint(Ω::Domain, box = boundingbox(Ω))
    local point
    indomain = false
    while ~indomain
        point = randompoint(box)
        indomain = point ∈ Ω
    end
    point
end


using SwitchTimeOpt
using MathOptInterface
const MOI = MathOptInterface

using Plots
using Ipopt

# Define Solver options
maxiter = 300
maxtime = 1000.0
verbose = 5
tolerance = 1e-06

#Define solver
solver = Ipopt.Optimizer()
MOI.set(solver, MOI.RawOptimizerAttribute("tol"), tolerance)
MOI.set(solver, MOI.RawOptimizerAttribute("print_level"), verbose)
MOI.set(solver, MOI.RawOptimizerAttribute("max_cpu_time"), maxtime)
MOI.set(solver, MOI.RawOptimizerAttribute("max_iter"), maxiter)

### Define system parameters
#  Time vector
t0 = 0.0
tf = 10.0

# Integer inputs
uvec = repeat([1.0; 2.0], 8, 1)
uvec = copy(uvec')

# Number of switches
N = size(uvec,2) -1

# Cost function matrix
K = 1.0
C = sqrt(K)*[0.0 1.0 -1.0]
Q = C'*C

# Initial state
x0 = [2.0; 2.0; 3.0]

### Define system dynamics
function nldyn(x::Array{Float64,1}, u::Array{Float64,1})
  n = length(x)
  f = zeros(n)

  f[1] = -sqrt(x[1]) + u[1]
  f[2] = sqrt(x[1]) - sqrt(x[2])
  f[3] = -0.05
  return f

end

function nldyn_deriv(x::Array{Float64,1}, u::Array{Float64,1})
  df = [-1/(2*sqrt(x[1]))        0                      0;
        1/(2*sqrt(x[1]))        -1/(2*sqrt(x[2]))       0;
        0                        0                      0]
end

# Generate and solve problem
ngrid = 100

# Initialize the problem
m = stoproblem(
  x0,                 # Initial state
  nldyn,              # Nonlinear dynamics
  nldyn_deriv,        # Nonlinear dynamics derivative
  uvec,               # Vector of integer inputs
  ngrid = ngrid,      # Number of grid points
  t0 = t0,            # InitialtTime
  tf = tf,            # Final time
  Q = Q,              # Cost matrix
  solver = solver)

solve!(m)
objlin = getobjval(m)
tauopt = gettau(m)
status = getstat(m)
xsim, xpts, objode45, t = simulate(m, tauopt)

println("Status: ", status)
println("Objective value: ", objlin)

showPlots = true
savePlots = false
plotSolution(m, "Tank", "Tank", showPlots, savePlots)


module ODE

export evolve_ode, Midpoint

"""    stepsize_adjust(y, yerr, h, order; rtol=1e-8, atol=1e-8)

Returns `(hnew, retry_step)` giving a stepsize that will try to achieve the
given absolute and relative tolerances on the next step (or a retry of this
one).

`retry_step` is `true` if the error exceeds the bound by more than a factor of
two, and `false` otherwise.
"""
function stepsize_adjust(y, yerr, h, order; rtol=1e-8, atol=1e-8)
    scale_max = zero(y[1])
    for i in 1:length(y)
        scale = yerr[i]/(atol + rtol*abs(y[i]))
        scale_max = max(scale_max, scale)
    end

    retry_step = (scale_max > 2 ? true : false)

    hnew = h/(scale_max)^(1/(order+1))

    if hnew > 2*h
        hnew = 2*h
    elseif hnew < 1/2*h
        hnew = 1/2*h
    end

    (hnew, retry_step)
end

struct Stepper
    order::Int
    step
end

"""    evolve_ode(stepper, f_rhs, x0, y0, x1; h_init, atol, rtol)

Returns the state of the ODE represented by `f_rhs` evolved from its initial
condition `(x0, y0)` to `x1` using the single-step algorithm in `stepper`.

`h_init` is used to guess the first step.

The evolution is adjusted to control the the absolute and relative error in the
solution according to the given absolute and relative tolerances; steps may be
re-tried until an acceptable error is achieved.

The final step will land *exactly* on `x1`.
"""
function evolve_ode(stepper, f_rhs, x0, y0, x1; h_init=(x1-x0)/100, atol=1e-8, rtol=1e-8)
    x = x0
    y = y0
    h = h_init
    not_done = true

    while not_done
        if x + h > x1
            h = x1-x
            not_done = false
        end

        y_new, yerr = stepper.step(f_rhs, x, y, h)
        h_new, retry = stepsize_adjust(y_new, yerr, h, stepper.order; rtol=rtol, atol=atol)

        if retry
            h = h_new
            not_done = true
        else
            y = y_new
            x = x + h
            h = h_new
        end
    end

    y
end

function step_midpoint(f_rhs, x, y, h)
    dydx0 = f_rhs(x, y)
    h_half = h/2
    y_half = y .+ dydx0 .* h_half

    dydx_half = f_rhs(x+h_half, y_half)

    y_new = y .+ dydx_half .* h
    yerr = h .* abs.(dydx_half .- dydx0)

    (y_new, yerr)
end

Midpoint = Stepper(2, step_midpoint)

end # module