"""
    function _rippa(A,a)

Estimate the Leave-One-Out (LOO) errors using rippas method. Complexity of 𝛰(3)
compared to calculating the exact LOO at a cost of 𝛰(4). `A` is the RBF matrix
and `a` is the weights of the RBF.
"""
function _rippa(A,a)
    N = size(A, 1)           #Number of Leave-One-Out (LOO) errors
    E = Array{Float64}(undef,N)

    _rippa!(E,A,a)
end

"""
    function _rippa!(E,A,a)

Same as _rippa but inplace.
"""
function _rippa!(E,A,a)
    N = size(A, 1)           #Number of Leave-One-Out (LOO) errors
    xₖ = A\I
    @inbounds for k = 1:N
        E[k] = a[k] / xₖ[k,k]   #Estimated error for k-th subset
    end

    return E
end

"""
    function _rmse_rippa!(A,a)

Same as _rippa! but calculates the RMSE value directly.
"""
function _rmse_rippa!(A,a)
    N = size(A, 1)           #Number of Leave-One-Out (LOO) errors
    xₖ = A\I
    RMSE = 0.0
    @inbounds for k = 1:N
        RMSE += (a[k] / xₖ[k,k])^2   #Accumulate square of estimated error for k-th subset
    end

    return sqrt(RMSE/N)
end

"""
    function RMSErrorLOO(interp,samples::Array{Float64,2},plan;
    cond_max = 1/eps(Float64)/10000, rippa = false)

Calculate the Leave-One-Out RMS error for a interpolation method.
`rippa` can be used to calculate the approximation of the LOO error
for Radial Basis Functions at a cost of 𝛰(3)
(compared to 𝛰(4)). `cond_max` sets the maximum allowed condition number for
matrix `A` used in the RBF calculation.
"""
function RMSErrorLOO(interp, samples, plan, smooth;
    cond_max::Float64=1e6, cond_check::Bool=false, rippa::Bool=false, rbf_dist_metric = Euclidean())

    #initiate arrays
    N = length(samples)
    E = Array{Float64}(undef,N)       #Error
    A = zeros(Float64, N, N)    #RBF matrix A
    ests = Array{Float64}(undef,N)    #Function estimate based on LOO
    LOOinds = Array{Int}(undef,N - 1, N)
    for i = 1:N
        LOOinds[:,i] = filter(x -> x != i, 1:N) # Get the leave one out sub indices
    end

    RMSErrorLOO!(A, ests, LOOinds, interp, samples, plan, smooth;
    cond_max=cond_max, rippa=rippa, rbf_dist_metric = rbf_dist_metric)
end

function RMSErrorLOO!(A, ests, LOOinds, interp::U,
 samples, plan::T, smooth; cond_max::Float64=1e6, cond_check::Bool=false,
  rippa::Bool=false, rbf_dist_metric = Euclidean()) where T <: AbstractArray where U <: AbstractArray

    N = length(samples)
    RMSE = Inf

    #perform Leave-One-Out estimation and calculate error
    if rippa
        @assert typeof(interp) <: Vector{ScatteredInterpolation.RadialBasisFunction} "Rippas
              algorithm only available for Radial Basis Functions"

        #interpolation object trained on the entire dataset
        itp, A = interpolate(interp, plan,
            samples, returnRBFmatrix=true, metric = rbf_dist_metric, smooth=smooth)

        #check the conditioning of matrix A for RBFs
        if cond_check && (cond(A) > cond_max)
            #RMSE is already Inf
        else
            #RBF error estimation based on Rippas algorithm
            RMSE = _rmse_rippa!(A,itp.w)
        end
    else
        for i = 1:N

            if typeof(smooth) <: AbstractVector
                loo_smooth = smooth[LOOinds[:,i]]
            else
                loo_smooth = smooth
            end

            #interpolation object trained on the LOO information
            itp, A = interpolate(interp[LOOinds[:,i]], plan[:,LOOinds[:,i]],
                samples[LOOinds[:,i]], returnRBFmatrix=true, metric = rbf_dist_metric, smooth=loo_smooth)

            #check the conditioning of matrix A for RBFs
            if cond_check && (cond(A) > cond_max)
                #RMSE is already Inf
            else
                #evaluate the interpolation object in the LOO position
                ests[i] = ScatteredInterpolation.evaluate(itp, plan[:,i])[1]
            end
        end
        RMSE = sqrt(mean((samples .- ests).^2))
    end

    return RMSE
end

function RMSErrorLOO!(A, ests, LOOinds, interp, samples, plan::T, smooth;
 cond_max::Float64=1e6, cond_check::Bool=false, rippa::Bool=false, rbf_dist_metric = Euclidean()) where T <: AbstractArray

    N = length(samples)
    RMSE = Inf

    #perform Leave-One-Out estimation and calculate error
    if rippa
        @assert typeof(interp) <: ScatteredInterpolation.RadialBasisFunction "Rippas algorithm only available for Radial Basis Functions"

        #interpolation object trained on the entire dataset
        itp, A = interpolate(interp, plan,
         samples, returnRBFmatrix=true, metric = rbf_dist_metric, smooth=smooth)

        #check the conditioning of matrix A for RBFs
        if cond_check && (cond(A) > cond_max)
            #RMSE is already Inf
        else
            #RBF error estimation based on Rippas algorithm
            RMSE = _rmse_rippa!(A,itp.w)
        end
    else
        for i = 1:N
            if typeof(smooth) <: AbstractVector
                loo_smooth = smooth[LOOinds[:,i]]
            else
                loo_smooth = smooth
            end

            #interpolation object trained on the LOO information
            itp, A = interpolate(interp, plan[:,LOOinds[:,i]],
                samples[LOOinds[:,i]], returnRBFmatrix=true, metric = rbf_dist_metric, smooth=loo_smooth)

            #check the conditioning of matrix A for RBFs
            if cond_check && (cond(A) > cond_max)
                #RMSE is already Inf
            else
                #evaluate the interpolation object in the LOO position
                ests[i] = ScatteredInterpolation.evaluate(itp, plan[:,i])[1]
            end
        end
        RMSE = sqrt(mean((samples .- ests).^2))
    end

    return RMSE
end

"""
    function scale(old_x::T,new_min,new_max;old_min=minimum(old_x),
    old_max=maximum(old_x)) where T <: Real

Scale a scalar to match a new range.
"""
function scale(old_x::T,new_min,new_max;old_min=minimum(old_x),old_max=maximum(old_x)) where T <: Real

    old_x = (((old_x - old_min) * (new_max - new_min)) / (old_max - old_min)) + new_min
end

"""
    function scale(old_x::Array{T,1},new_min,new_max;old_min=minimum(old_x),
    old_max=maximum(old_x)) where T <: Real

Scale a vector to match a new range.
"""
function scale(old_x::Array{T,1},new_min,new_max;old_min=minimum(old_x),old_max=maximum(old_x)) where T <: Real

    newX = (((old_x .- old_min) .* (new_max .- new_min)) ./ (old_max .- old_min)) .+ new_min
end

"""
    function scale(old_x::Array{T,2},direction::Int,new_min,new_max;
    old_min=minimum(old_x,direction),old_max=maximum(old_x,direction)) where T <: Real

Scale a 2D matrix to match a new range along the specified `direction`. Specify
the old extrema as scalars to use the same factor per row/column, or as a
vector to scale each row/column by a unique value.
"""
function scale(old_x::Array{T,2},direction::Int,new_min,new_max;old_min=minimum(old_x,dims=direction),old_max=maximum(old_x,dims=direction)) where T <: Real

    if length(old_min) == length(old_max) == 1
        newX = mapslices(x->scale(x,new_min,new_max;old_min=old_min,old_max=old_max), old_x, dims=direction)
    else
        old_min_itr = (x for x in old_min)
        old_max_itr = (x for x in old_max)

        mapfun = function (x)
            (omin, old_min_itr) = Iterators.peel(old_min_itr)
            (omax, old_max_itr) = Iterators.peel(old_max_itr)
            scale(x,new_min,new_max;old_min=omin,old_max=omax)
        end
        newX = mapslices(mapfun, old_x, dims=direction)
    end
    return newX
end

scale(x::Missing,min_val,max_val;old_min,old_max) = missing

"""
    preprocess_point(points,scaling;base_scale::Array{Float64,2})

Scales the input points in order to evaluate the adaptively scaled RBF correctly.
"""
function preprocess_point(points,scaling;base_scale::Array{Float64,2})

    min_max = extrema(base_scale, dims=2)

    preprocessed_point = similar(points)
    for i = 1:size(preprocessed_point,1)
        preprocessed_point[i,:] = scale(points[i,:],-1.0*scaling[i],1.0*scaling[i],
        old_min = min_max[i][1], old_max = min_max[i][2])
    end
    return preprocessed_point
end

"""
    preprocess_point!(preprocessed_point,points,scaling;old_min::Array{Float64,2},old_max::Array{Float64,2})

Scales the input points in order to evaluate the adaptively scaled RBF correctly.
"""
function preprocess_point!(preprocessed_point,points,scaling;old_min::Array{Float64,2},old_max::Array{Float64,2})

    for i = 1:size(preprocessed_point,1)
        preprocessed_point[i,:] = scale(points[i,:],-1.0*scaling[i],1.0*scaling[i],
        old_min = old_min[i], old_max = old_max[i])
    end
    return preprocessed_point
end

"""
    preprocess_point!(preprocessed_point,points,scaling::Bool;old_min::Array{Float64,2},old_max::Array{Float64,2}) where U

Scales the input points in order to evaluate the adaptively scaled RBF correctly.
"""
function preprocess_point!(preprocessed_point,points,scaling::Bool;old_min::Array{Float64,2},old_max::Array{Float64,2}) where U
    preprocessed_point = points
end

"""
    preprocess_point(points,scaling::Bool;base_scale::Array{Float64,2}) where U

Scales the input points in order to evaluate the adaptively scaled RBF correctly.
"""
function preprocess_point(points,scaling::Bool;base_scale::Array{Float64,2}) where U
    points
end

"""
    surrogate_evaluate(preprocessed_est_point,estimation_point,itp,optres,old_min,old_max)

Processes the input to scale correctly based on the options used.
"""
function surrogate_evaluate(preprocessed_est_point,estimation_point,itp,optres,old_min,old_max)
    preprocessed_est_point = preprocess_point!( preprocessed_est_point,
                                                estimation_point, optres.scaling;
                                                old_min=old_min,
                                                old_max=old_max,
                                                )

    # Evaluate the interpolation object
    res = ScatteredInterpolation.evaluate(itp, preprocessed_est_point)[1]
end

"""
    construct_search_range(plan::Array{Float64,2}, variable_kernel_width,
    min_rbf_width, max_rbf_width, variable_dim_scaling,
    min_scale, max_scale, smooth, max_smooth)

Creates array of tuples for the optimization of the RBF hyperparameters.
"""
function construct_search_range(plan::Array{Float64,2}, variable_kernel_width,
    min_rbf_width, max_rbf_width, variable_dim_scaling,
    min_scale, max_scale, smooth, max_smooth)

    n_dims, n_samples = size(plan)
    sr = Array{Tuple{Float64,Float64},1}()

    # Add the kernel width and type search range
    variable_kernel_width ? n_kerns = n_samples : n_kerns = 1
    push!(sr,   create_sr(  (min_range=min_rbf_width, max_range=max_rbf_width, n_times=n_kerns),
    (min_range=0.0, max_range=1.0, n_times=n_kerns))...)

    # Add the dimensional scaling search range
    variable_dim_scaling ? n_dim_scales = n_dims : n_dim_scales = 0
    push!(sr,   create_sr(  (min_range=min_scale, max_range=max_scale, n_times=n_dim_scales))...)

    # Add the ridge regression smoothing search range
    (smooth == :variable) && (n_smooth = n_samples)
    (smooth == :single)   && (n_smooth = 1)
    (smooth == :single_user)   && (n_smooth = 0)
    (smooth == false)     && (n_smooth = 0)
    push!(sr,   create_sr(  (min_range=0.0, max_range=max_smooth, n_times=n_smooth))...)

    return sr
end

"""
    create_sr(vargs::NamedTuple{(:min_range, :max_range, :n_times),Tuple{Float64,Float64,Int64}}...)

Facilitate the creation of array of tuples for the optimization of the RBF hyperparameters.
"""
function create_sr(vargs::NamedTuple{(:min_range, :max_range, :n_times),Tuple{Float64,Float64,Int64}}...)
    sr = Array{Tuple{Float64,Float64},1}()
    for varg in vargs
        for i = 1:varg[3]
            push!(sr,(varg[1],varg[2]))
        end
    end
    return sr
end

"""
    extract_vector_range(vargs::Int64...)

Supply length of subsets contained in a vector and receive a tuple containing
all the ranges needed to extract each subset in the vector.
"""
function extract_vector_range(vargs::Int...)
    output = Array{Union{UnitRange{Int}, Int, Bool},1}()

    count = 0
    for (i,varg) in enumerate(vargs)
        if varg == 0
            push!(output, false)
        else
            push!(output, (count+1):(count+varg))
            count += varg
            (length(output[end]) == 1) && (output[end] = output[end][1])
        end
    end
    return Tuple(output)
end

"""
    extract_bboptim_hypers(bboptim_fcall_vector,plan,kerns,
    variable_kernel_width,variable_dim_scaling,
    smooth,smooth_user)

Get the optimized RBF hyperparameters in useful format for further use.
"""
function extract_bboptim_hypers(bboptim_fcall_vector,plan,kerns,
    variable_kernel_width,variable_dim_scaling,
    smooth,smooth_user)

    n_dims, n_samples = size(plan)

    # Kernel width and type length
    variable_kernel_width ? n_kerns = n_samples : n_kerns = 1

    # Dimensional scaling length
    variable_dim_scaling ? n_dim_scales = n_dims : n_dim_scales = 0

    # Ridge regression smoothing length
    (smooth == :variable) && (n_smooth = n_samples)
    (smooth == :single)   && (n_smooth = 1)
    (smooth == :single_user)   && (n_smooth = 0)
    (smooth == false)     && (n_smooth = 0)

    width_inds, kernel_float_inds, scaling_inds, smooth_inds = extract_vector_range(n_kerns,
                                                            n_kerns,
                                                            n_dim_scales,
                                                            n_smooth
                                                            )

    # Arrange the width and smoothing results
    width = bboptim_fcall_vector[width_inds]
    kernel_float = bboptim_fcall_vector[kernel_float_inds]
    !(scaling_inds == false) ? scaling = bboptim_fcall_vector[scaling_inds] : scaling = scaling_inds
    (smooth == false) &&            (smooth = smooth_inds)
    (smooth == :single) &&          (smooth = bboptim_fcall_vector[smooth_inds])
    (smooth == :variable) &&        (smooth = bboptim_fcall_vector[smooth_inds])
    (smooth == :single_user) &&     (smooth = smooth_user)

    # Arrange the RBF kernel result
    kern_ind = round.(Int,scale(kernel_float,1,length(kerns),old_min=0,old_max=1))
    if variable_kernel_width
        kern = Vector{ScatteredInterpolation.RadialBasisFunction}(undef,size(plan,2))
        for i = 1:size(plan,2)
            kern[i] = kerns[kern_ind[i]](width[i])
        end
    elseif !variable_kernel_width
        kern = kerns[kern_ind](width)
    end

    return kern, scaling, smooth
end

"""
    interp_obj(inpt::Vector{Float64}, kerns, samples,
    plan::Array{Float64,2}; rippa::Bool = false,
    variable_kernel_width::Bool = true, variable_dim_scaling::Bool = true,
    smooth = false, cond_max=cond_max, rbf_dist_metric = Distances.Euclidean(),
    smooth_user::Float64 = 0.0)

Objective function for optimization of interpolation kernel function and width.
"""
function interp_obj(inpt::Vector{Float64}, A, ests, LOOinds, kerns, samples,
        plan::Array{Float64,2}; rippa::Bool = false,
        variable_kernel_width::Bool = true, variable_dim_scaling::Bool = true,
        smooth::Union{Bool, Symbol} = false, cond_max::Float64=Inf,  cond_check::Bool=false, rbf_dist_metric = Distances.Euclidean(),
        smooth_user::Float64 = 0.0)

    kern, scaling, smooth = extract_bboptim_hypers( inpt,plan,kerns,variable_kernel_width,
                                                    variable_dim_scaling,smooth,
                                                    smooth_user)

    #Preprocess the plan based on the settings used.
    preprocessed_plan = preprocess_point(plan,scaling,base_scale=plan)

    #Wrapped in try catch-block to catch failure to solve linear eq. system
    E = try
        E = RMSErrorLOO!(A, ests, LOOinds, kern, samples, preprocessed_plan, smooth; rippa = rippa,
        cond_max = cond_max, rbf_dist_metric = rbf_dist_metric, cond_check=cond_check)
    catch
        E = Inf
    end

    return E
end

"""
    rbf_hypers_opt(samples_org::Array{Float64,2}, plan::Array{Float64,2}, options::Options)

Optimization function of Radial Basis Function kernel and width.
"""
function rbf_hypers_opt(samples_org::Array{Float64,2}, plan::Array{Float64,2}, options::Options)

    @unpack rippa, variable_kernel_width, variable_dim_scaling, rbf_opt_method,
            min_rbf_width, max_rbf_width, min_scale, max_scale, cond_max,
            rbf_dist_metric, rbf_opt_gens, kerns, rbf_opt_pop,
            smooth, max_smooth, smooth_user, cond_check, constrained_seed_gens = options

    # Sample order for ScatteredInterpolation
    samples = vec(samples_org)

    # Pre-allocate for inplace functions
    N = length(samples) #Number of Leave-One-Out (LOO) errors
    A = zeros(Float64, N, N)    #RBF matrix A
    ests = Array{Float64}(undef,N)    #Function estimate based on LOO
    LOOinds = Array{Int}(undef,N - 1, N)
    for i = 1:N
        LOOinds[:,i] = filter(x -> x != i, 1:N) # Get the leave one out sub indices
    end

    # Create the hyperparameter search range based on the input options for a constrained space
    sr = construct_search_range(plan, false,
    min_rbf_width, max_rbf_width, variable_dim_scaling,
    min_scale, max_scale, smooth, max_smooth)

    #Improve the interpolant by first optimising the hyperparameters in a constrained search space
    if constrained_seed_gens > 0

        # RBF hyperparameter objective function
        itp_obj = function (x)
            interp_obj(x,A,ests,LOOinds,kerns,samples,plan;
                    rippa=rippa, variable_kernel_width=false,
                    variable_dim_scaling=variable_dim_scaling, smooth=smooth,
                    cond_max=cond_max,rbf_dist_metric=rbf_dist_metric,cond_check=cond_check)
        end

        @timeit_debug "hyper opt constrained" begin
        # Optimize the interpolant hyperparameters
        res = bboptimize(itp_obj;
                Method=rbf_opt_method,SearchRange=sr, MaxFuncEvals=rbf_opt_gens,
                TraceMode=:silent, rbf_dist_metric=rbf_dist_metric,
                TargetFitness = 1e-5, FitnessTolerance = 1e-6,
                PopulationSize = rbf_opt_pop,
                MaxStepsWithoutProgress=rbf_opt_gens,
                MaxNumStepsWithoutFuncEvals=rbf_opt_gens,
                );
        end
        #@show best_fitness(res)
        kern, scaling, smoothing = extract_bboptim_hypers( res.archive_output.best_candidate,
                                                        plan,kerns,false,
                                                        variable_dim_scaling,smooth,smooth_user)

        x0 = res.archive_output.best_candidate
    else
        x0 = BlackBoxOptim.Utils.latin_hypercube_sampling(first.(sr),last.(sr),1)
    end

    # Create the hyperparameter search range based on the input options
    sr = construct_search_range(plan, variable_kernel_width,
                                min_rbf_width, max_rbf_width, variable_dim_scaling,
                                min_scale, max_scale, smooth, max_smooth)

    # RBF hyperparameter objective function
    itp_obj = function (x)
        interp_obj(x,A,ests,LOOinds,kerns,samples,plan;
                rippa=rippa, variable_kernel_width=variable_kernel_width,
                variable_dim_scaling=variable_dim_scaling, smooth=smooth,
                cond_max=cond_max,rbf_dist_metric=rbf_dist_metric,cond_check=cond_check)
    end

    #First two values are the kern and kern_width, x0[3:end] is scaling and smooth
    #kern    kern_width    scaling    smooth
    if variable_kernel_width
        x1=[[x0[1] for _ in 1:size(plan,2)]; [x0[2] for _ in 1:size(plan,2)]; x0[3:end]]
    else
        x1 = x0
    end
    @assert length(x1)==length(sr)

    #Create population from LHC using the results from the constrained optimisation
    #of hyperparameters (if available), else supply random population.
    nvals=length(sr)
    if constrained_seed_gens > 0
        pop_constrained = [x1[i] for i in 1:nvals, _ in 1:1]
        pop = Array{Float64,2}(undef,nvals,rbf_opt_pop)
        pop[:,1] = pop_constrained
        for i in size(pop_constrained,2)+1:size(pop,2)
            pop[:,i] = BlackBoxOptim.Utils.latin_hypercube_sampling(first.(sr),last.(sr),1)
        end
    else
        pop = BlackBoxOptim.Utils.latin_hypercube_sampling(first.(sr),last.(sr),rbf_opt_pop)
    end

    @timeit_debug "hyper opt full" begin
    # Optimize the interpolant hyperparameters
    res = bboptimize(itp_obj;
            Method=rbf_opt_method,SearchRange=sr, MaxFuncEvals=rbf_opt_gens,
            TraceMode=:silent, rbf_dist_metric=rbf_dist_metric,
            TargetFitness = 1e-5, FitnessTolerance = 1e-6,
            PopulationSize = rbf_opt_pop,
            MaxStepsWithoutProgress=rbf_opt_gens,
            MaxNumStepsWithoutFuncEvals=rbf_opt_gens,
            Population=pop
            );
    end
    #@show best_fitness(res)
    kern, scaling, smoothing = extract_bboptim_hypers( res.archive_output.best_candidate,
                                                    plan,kerns,variable_kernel_width,
                                                    variable_dim_scaling,smooth,smooth_user)

    # Return the optimized hyperparameters in the correct type
    return RBFHypers(kern, scaling, smoothing), best_fitness(res)
end


function EvaluateFisherMatrixElement!(FisherMatrix::Fisherαβ, Cov::aₗₘCovariance,
    ∂Cℓα::∂Cℓ, ∂Cℓβ::∂Cℓ, Parα::String, Parβ::String)
    αMatrix = mymatmul(∂Cℓα.∂CℓArray,Cov.Covariance⁻¹)
    βMatrix = mymatmul(∂Cℓβ.∂CℓArray,Cov.Covariance⁻¹)
    Fisherℓ = mymatmul(αMatrix, βMatrix)
    fisherelement = SumℓAndTrace(Fisherℓ)
    FisherMatrix.FisherDict[Parα*"_"*Parβ] = fisherelement
end

function mymatmul(A::Array{Float64, 3}, B::Array{Float64, 3})
    C = zeros(size(A))
    @avx for l in 1:length(A[:,1,1])
        for i in 1:length(A[1,:,1])
            for j in 1:length(B[1,1,:])
                for k in 1:length(B[1,:,1])
                    C[l,i,j] += A[l,i,k]*B[l,k,j]
                end
            end
        end
    end
    return C
end

function SumℓAndTrace(Fisherℓ::Array{Float64, 3})
    fisherelement = 0
    @avx for ℓ in 1:size(Fisherℓ)[1]
        for i in 1:size(Fisherℓ)[2]
            fisherelement += Fisherℓ[ℓ, i, i]
        end
    end
    return fisherelement
end

function EvaluateFisherMatrixElement!(FisherMatrix::Fisherαβ, Cov::CℓCovariance,
    ∂Cℓα::∂Cℓ, ∂Cℓβ::∂Cℓ, Parα::String, Parβ::String)
    ℓnumber = length(∂Cℓα.∂CℓArray[:,1,1])
    inumber = length(∂Cℓα.∂CℓArray[1,:,1])
    L = EliminationMatrix(inumber)

    vecp∂CℓαᵀCov⁻¹ = zeros(1,floor(Int,inumber*0.5*(inumber+1)))
    fisher_temp = zeros(1,1)
    fisherelement = 0
    for ℓ in 1:ℓnumber
        vec∂Cℓα = vec(∂Cℓα.∂CℓArray[ℓ,:,:])
        vecp∂Cℓα = zeros(floor(Int,inumber*0.5*(inumber+1)))
        LinearAlgebra.mul!(vecp∂Cℓα, L, vec∂Cℓα)

        vec∂Cℓβ = vec(∂Cℓβ.∂CℓArray[ℓ,:,:])
        vecp∂Cℓβ = zeros(floor(Int,inumber*0.5*(inumber+1)))
        LinearAlgebra.mul!(vecp∂Cℓβ, L, vec∂Cℓβ)

        Covariance⁻¹ = Cov.Covariance⁻¹[ℓ,:,:]
        LinearAlgebra.mul!(vecp∂CℓαᵀCov⁻¹, transpose(vecp∂Cℓα), Covariance⁻¹)
        LinearAlgebra.mul!(fisher_temp, vecp∂CℓαᵀCov⁻¹, vecp∂Cℓβ)
        fisherelement += fisher_temp[1,1]
    end
    FisherMatrix.FisherDict[Parα*"_"*Parβ] = fisherelement
end


"""
    Born-Mayer Potential
"""
mutable struct BornMayer <: Potential
    A::Float64
    ρ::Float64
end

"""
    BornMayer(params::Dict)

Creates a BM potential.
"""
function BornMayer(params::Dict)
    A = params["A"]
    ρ = params["ρ"]
    return BornMayer(A, ρ)
end

"""
    potential_energy(p::BornMayer, r::Position, args...)

Calculates Born-Mayer potential energy.
"""
function potential_energy(p::BornMayer, r::Position, args...)
    return p.A * exp(-norm(r) / p.ρ)
end


solve(i::BudgetedSpanningTreeInstance{T}, ::LagrangianAlgorithm, ε; kwargs...) where T = st_prim_budgeted_lagrangian_search(i, ε; kwargs...)
solve(i::BudgetedSpanningTreeInstance{T}, ::LagrangianRefinementAlgorithm; kwargs...) where T = st_prim_budgeted_lagrangian_refinement(i; kwargs...)
solve(i::BudgetedSpanningTreeInstance{T}, ::IteratedLagrangianRefinementAlgorithm; kwargs...) where T = st_prim_budgeted_lagrangian_approx_half(i; kwargs...)

approximation_term(::BudgetedSpanningTreeInstance{T}, ::LagrangianAlgorithm) where T = NaN
approximation_ratio(::BudgetedSpanningTreeInstance{T}, ::LagrangianAlgorithm) where T = NaN

approximation_term(i::BudgetedSpanningTreeInstance{T}, ::LagrangianRefinementAlgorithm) where T = maximum(values(i.rewards))
approximation_ratio(::BudgetedSpanningTreeInstance{T}, ::LagrangianRefinementAlgorithm) where T = NaN

approximation_term(::BudgetedSpanningTreeInstance{T}, ::IteratedLagrangianRefinementAlgorithm) where T = NaN
approximation_ratio(::BudgetedSpanningTreeInstance{T}, ::IteratedLagrangianRefinementAlgorithm) where T = 0.5

function _budgeted_spanning_tree_compute_weight(i::BudgetedSpanningTreeInstance{T, U}, tree::Vector{Edge{T}}) where {T, U}
  return sum(i.weights[(e in keys(i.weights)) ? e : reverse(e)] for e in tree)
end

struct BudgetedSpanningTreeLagrangianSolution{T, U} <: BudgetedSpanningTreeSolution{T, U}
  # Used to store important temporary results from solving the Lagrangian dual.
  instance::BudgetedSpanningTreeInstance{T, U}
  tree::Vector{Edge{T}}
  λ::Float64 # Optimum dual multiplier.
  value::Float64 # Optimum value of the dual problem (i.e. with penalised constraint).
  λmax::Float64 # No dual value higher than this is useful (i.e. they all yield the same solution).
end

function st_prim_budgeted_lagrangian(i::BudgetedSpanningTreeInstance{T, U}, λ::Float64) where {T, U}
  # Solve the subproblem for one value of the dual multiplier λ:
  #     l(λ) = \max_{x spanning tree} (rewards + λ weights) x - λ budget.
  sti_rewards = Dict{Edge{T}, Float64}(e => i.rewards[e] + λ * i.weights[e] for e in keys(i.rewards))
  sti = SpanningTreeInstance(i.graph, sti_rewards)
  sti_sol = st_prim(sti)
  return _budgeted_spanning_tree_compute_value(sti, sti_sol.tree) - λ * i.budget, sti_sol.tree
end

function st_prim_budgeted_lagrangian_search(i::BudgetedSpanningTreeInstance{T, U}, ε::Float64) where {T, U}
  # Approximately solve the problem \min_{l ≥ 0} l(λ), where
  #     l(λ) = \max_{x spanning tree} (rewards + λ weights) x - λ budget.
  # This problem is the Lagrangian dual of the budgeted maximum spanning-tree problem:
  #     \max_{x spanning tree} rewards x  s.t.  weights x >= budget.
  # This algorithm provides no guarantee on the optimality of the solution.

  # Initial set of values for λ. The optimum is guaranteed to be contained in this interval.
  weights_norm_inf = maximum(values(i.weights)) # Maximum weight.
  m = nv(i.graph) - 1 # Maximum number of items in a solution. Easy to compute for a spanning tree!
  λmax = Float64(weights_norm_inf * m + 1)

  λlow = 0.0
  λhigh = λmax

  # Perform a golden-ratio search.
  while (λhigh - λlow) > ε
    λmidlow = λhigh - (λhigh - λlow) / MathConstants.φ
    λmidhigh = λlow + (λhigh - λlow) / MathConstants.φ
    vmidlow, _ = st_prim_budgeted_lagrangian(i, λmidlow)
    vmidhigh, _ = st_prim_budgeted_lagrangian(i, λmidhigh)

    if vmidlow < vmidhigh
      λhigh = λmidhigh
      vhigh = vmidhigh
    else
      λlow = λmidlow
      vlow = vmidlow
    end
  end

  vlow, stlow = st_prim_budgeted_lagrangian(i, λlow)
  vhigh, sthigh = st_prim_budgeted_lagrangian(i, λhigh)
  if vlow < vhigh
    return BudgetedSpanningTreeLagrangianSolution(i, stlow, λlow, vlow, λmax)
  else
    return BudgetedSpanningTreeLagrangianSolution(i, sthigh, λhigh, vhigh, λmax)
  end
end

function st_prim_budgeted_lagrangian_refinement(i::BudgetedSpanningTreeInstance{T, U};
                                                ε::Float64=1.0e-3, ζ⁻::Float64=0.2, ζ⁺::Float64=5.0,
                                                stalling⁻::Float64=1.0e-5) where {T, U}
  # Approximately solve the following problem:
  #     \max_{x spanning tree} rewards x  s.t.  weights x >= budget
  # This algorithm provides an additive approximation to this problem. If x* is the optimum solution and x~ the one
  # returned by this algorithm,
  #     weights x* >= budget   and   weights x~ >= budget                 (the returned solution is feasible)
  #     rewards x~ >= rewards x* - \max{e edge} reward[e]                 (additive approximation)

  # Check assumptions.
  if ζ⁻ >= 1.0
    error("ζ⁻ must be strictly less than 1.0: the dual multiplier λ is multiplied by ζ⁻ to reach an infeasible solution by less penalising the budget constraint.")
  end
  if ζ⁺ <= 1.0
    error("ζ⁺ must be strictly greater than 1.0: the dual multiplier λ is multiplied by ζ⁺ to reach a feasible solution by penalising more the budget constraint.")
  end

  # Ensure the problem is feasible by only considering the budget constraint.
  feasible_rewards = Dict{Edge{T}, Float64}(e => i.weights[e] for e in keys(i.rewards))
  feasible_instance = SpanningTreeInstance(i.graph, feasible_rewards)
  feasible_solution = st_prim(feasible_instance)
  if _budgeted_spanning_tree_compute_value(feasible_instance, feasible_solution.tree) < i.budget
    # By maximising the left-hand side of the budget constraint, impossible to reach the target budget. No solution!
    return SimpleBudgetedSpanningTreeSolution(i)
  end

  # Solve the Lagrangian relaxation to optimality.
  lagrangian = st_prim_budgeted_lagrangian_search(i, ε)
  λ0, v0, st0 = lagrangian.λ, lagrangian.value, lagrangian.tree
  λmax = lagrangian.λmax
  b0 = _budgeted_spanning_tree_compute_weight(i, st0) # Budget consumption of this first solution.

  # If already respecting the budget constraint exactly, done!
  if b0 == i.budget
    return SimpleBudgetedSpanningTreeSolution(i, st0)
  end

  # Find two solutions: one above the budget x⁺ (i.e. respecting the constraint), the other not x⁻.
  x⁺, x⁻ = nothing, nothing

  λi = λ0
  if b0 > i.budget
    x⁺ = st0
    @assert _budgeted_spanning_tree_compute_weight(i, x⁺) > i.budget

    stalling = false
    while true
      # Penalise less the constraint: it should no more be satisfied.
      λi *= ζ⁻
      _, sti = st_prim_budgeted_lagrangian(i, λi)
      if _budgeted_spanning_tree_compute_weight(i, sti) < i.budget
        x⁻ = sti
        break
      end

      # Is the process stalling?
      if λi <= stalling⁻
        stalling = true
        break
      end
    end

    # Specific handling of stallings.
    if stalling # First test: don't penalise the constraint at all.
      _, sti = st_prim_budgeted_lagrangian(i, 0.0)
      new_budget = _budgeted_spanning_tree_compute_weight(i, sti)
      if new_budget < i.budget
        x⁻ = sti
        stalling = false
      end
    end

    if stalling # Second test: minimise the left-hand side of the budget constraint, in hope of finding a feasible solution.
      # This process is highly similar to the computation of feasible_solution, but with a reverse objective function.
      infeasible_rewards = Dict{Edge{T}, Float64}(e => - i.weights[e] for e in keys(i.weights))
      infeasible_solution = st_prim(SpanningTreeInstance(i.graph, infeasible_rewards)).tree

      if _budgeted_spanning_tree_compute_weight(i, infeasible_solution) < i.budget
        x⁻ = infeasible_solution
        stalling = false
      end
    end

    if stalling # Third: decide there is no solution strictly below the budget. No refinement is possible.
      # As x⁺ is feasible, return it.
      return SimpleBudgetedSpanningTreeSolution(i, x⁺)
    end
  else
    x⁻ = st0
    @assert _budgeted_spanning_tree_compute_weight(i, x⁻) < i.budget

    while true
      # Penalise more the constraint: it should become satisfied at some point.
      λi *= ζ⁺
      _, sti = st_prim_budgeted_lagrangian(i, λi)
      if _budgeted_spanning_tree_compute_weight(i, sti) >= i.budget
        x⁺ = sti
        break
      end

      # Is the process stalling? If so, reuse feasible_solution, which is guaranteed to be feasible.
      if λi >= λmax
        x⁺ = feasible_solution.tree
        break
      end
    end
  end

  # Normalise the solutions: the input graph is undirected, the direction of the edges is not important.
  # In case one solution has the edge v -> w and the other one w -> v, make them equal. s
  sort_edge(e::Edge{T}) where T = (src(e) < dst(e)) ? e : reverse(e)
  x⁺ = [sort_edge(e) for e in x⁺]
  x⁻ = [sort_edge(e) for e in x⁻]

  # Iterative refinement. Stop as soon as there is a difference of at most one edge between the two solutions.
  while _solution_symmetric_difference_size(x⁺, x⁻) > 2
    # Enforce the loop invariant.
    @assert x⁺ !== nothing
    @assert x⁻ !== nothing
    @assert _budgeted_spanning_tree_compute_weight(i, x⁺) >= i.budget # Feasible.
    @assert _budgeted_spanning_tree_compute_weight(i, x⁻) < i.budget # Infeasible.

    # Switch elements from one solution to another.
    only_in_x⁺, only_in_x⁻ = _solution_symmetric_difference(x⁺, x⁻)
    e1 = first(only_in_x⁺)
    e2 = first(only_in_x⁻)

    # Create the new solution (don't erase x⁺ nor x⁻: only one of the two will be forgotten, the other will be kept).
    new_x = copy(x⁺)
    filter!(e -> e != e1, new_x)
    push!(new_x, e2)

    # Replace one of the two solutions, depending on whether this solution is feasible (x⁺) or not (x⁻).
    if _budgeted_spanning_tree_compute_weight(i, new_x) >= i.budget
      x⁺ = new_x
    else
      x⁻ = new_x
    end
  end

  # Done!
  return SimpleBudgetedSpanningTreeSolution(i, x⁺)
end

function st_prim_budgeted_lagrangian_approx_half(i::BudgetedSpanningTreeInstance{T, U}; kwargs...) where {T, U}
  # Approximately solve the following problem:
  #     \max_{x spanning tree} rewards x  s.t.  weights x >= budget
  # This algorithm provides a multiplicative approximation to this problem. If x* is the optimum solution and x~ the one
  # returned by this algorithm,
  #     weights x* >= budget   and   weights x~ >= budget                 (the returned solution is feasible)
  #     rewards x~ >= rewards x* / 2                                      (multiplicative approximation)

  # Based on the same ideas as https://doi.org/10.1007/s10107-009-0307-4

  # If there are too few edges, not much to do.
  if ne(i.graph) <= 2
    return st_prim_budgeted_lagrangian_refinement(i; kwargs...)
  end

  # For each pair of edges, force these two edges to be part of the solution and discard all edges with a higher value.
  best_sol = nothing
  for e1 in edges(i.graph)
    for e2 in edges(i.graph)
      # (e1, e2) must be a pair of distinct edges.
      if (src(e1) == src(e2) && dst(e1) == dst(e2)) || (src(e1) == dst(e2) && src(e1) == dst(e2))
        continue
      end

      # Filter out the edges that have a higher value than any of these two edges. Give a very large reward to them both.
      cutoff = min(i.rewards[e1], i.rewards[e2])
      rewards = filter(kv -> kv[2] < cutoff, i.rewards)
      rewards[e1] = rewards[e2] = prevfloat(Inf)

      graph = SimpleGraph(nv(i.graph))
      for e in keys(rewards)
        add_edge!(graph, e)
      end

      weights = Dict(e => i.weights[e] for e in keys(rewards))

      # No other simplification can be made, unfortunately: all other edges might participate in the optimum solution.
      # (Only exception: e1 and e2 are incident to the same vertex; any edge linking the other extremities of these edges
      # can be removed, but that's only one edge at most.)
      # As e1 and e2 have the best reward (as they are really bumped), they must be in any optimum solution.

      # Solve this subproblem.
      bsti = BudgetedSpanningTreeInstance(graph, rewards, weights, i.budget)
      sol = st_prim_budgeted_lagrangian_refinement(bsti; kwargs...)

      # This subproblem is infeasible. Maybe it's because the overall problem is infeasible or just because too many
      # edges were removed.
      if length(sol.tree) == 0
        continue
      end

      # Impossible to have a feasible solution with these two edges, probably because of the budget constraint.
      # Due to the direction of the budget constraint (>= budget), it is not possible to check for feasibility
      # before solving an instance.
      if ! (e1 in sol.tree) || ! (e2 in sol.tree)
        continue
      end

      # Only keep the best solution.
      if best_sol === nothing || sol.value > best_sol.value
        # sol's instance is the one used internally for the subproblems.
        best_sol = SimpleBudgetedSpanningTreeSolution(i, sol.tree)
      end
    end
  end

  if best_sol !== nothing
    return best_sol
  else
    return SimpleBudgetedSpanningTreeSolution(i)
  end
end


# module TestODEModel
using CovidModels
using CovidModels.Data
using OrdinaryDiffEq
import OrdinaryDiffEq: ODEProblem
using Petri
using Test
using Plots
using StatsPlots
using Colors

function fitplot(sol, df)
    p = plot(sol, vars=[2, 4, 5, 6], label=["IA" "I" "R" "D"])
    @df df scatter!(p,
            0:length(italy.Deaths)-1,
            [:Cases, :Deaths, :Recoveries],
            label=["cases" "deaths" "recoveries"],
            marker=true,
            color=[RGB(1,0,0) RGB(0.5,0,0.5) RGB(0,1,0)]
            )
    return p
end

function makeplots_seir(sol, prefix)
    mkpath(dirname(prefix))
    p1 = plot(sol,vars=[1,2,3,4], xlabel="", ylabel="people", linewidth=3,title="Cities", legend=false)
    p2 = plot(sol,vars=[5,6,7,8], xlabel="", ylabel="people", linewidth=3, legend=false)
    p3 = plot(sol,vars=[9,10,11,12], xlabel="time", ylabel="people", linewidth=3, legend=false)
    p4 = plot(sol,vars=[2,6,10], xlabel="", linewidth=3, labels=["i1" "i2" "i3"], legend=true)
    p5 = plot(sol,vars=[3,7,11], xlabel="", linewidth=3,title="Populations", labels=["e1" "e2" "e3"], legend=true)
    p6 = plot(sol,vars=[4,8,12], xlabel="time", linewidth=3, labels=["r1" "r2" "r3"], legend=true)
    p = plot(p1, p5, p2, p4, p3, p6, layout=(3,2), linewidth=3, link=:both)
    savefig(p, "$(prefix)combined.pdf")
    p
end

function makeplots_seird(sol, prefix)
    mkpath(dirname(prefix))
    p1 = plot(sol,vars=[1,2,3,4,5], xlabel="", ylabel="people", linewidth=3,title="Cities", legend=false)
    p2 = plot(sol,vars=[6,7,8,9,10], xlabel="", ylabel="people", linewidth=3, legend=false)
    p3 = plot(sol,vars=[11,12,13,14,15], xlabel="time", ylabel="people", linewidth=3, legend=false)
    p4 = plot(sol,vars=[2,7,12], xlabel="", linewidth=3, labels=["i1" "i2" "i3"], legend=true)
    p5 = plot(sol,vars=[3,8,13], xlabel="", linewidth=3,title="Populations", labels=["e1" "e2" "e3"], legend=true)
    p6 = plot(sol,vars=[5,10,15], xlabel="time", linewidth=3, labels=["d1" "d2" "d3"], legend=true)
    p = plot(p1, p5, p2, p4, p3, p6, layout=(3,2), linewidth=3, link=:both)
    savefig(p, "$(prefix)combined.pdf")
    p
end

T = [
    # City 1 SEIR
    ([1, 2], [3, 2]),
    ([3], [2]), # E→I
    ([2], [4]), # I→R
    # outflow 1→2
    ([1], [5]), # S→S′
    ([2], [6]), # I→I′
    ([3], [7]), # E→E′
    # City 2 SEIR
    ([5, 6], [7, 6]),
    ([7], [6]), # E→I
    ([6], [8]), # I→R
    # outflow 2→3
    ([5], [9]), # S→S′
    ([6], [10]),# I→I′
    ([7], [11]),# E→E′
    # City 3 SEIR
    ([9, 10], [11, 10]),
    ([11], [10]), # E→I
    ([10], [12]) # I→R
]

m = Petri.Model(1:12, T, missing, missing)
nS = length(m.S)
β = ones(Float64, length(T))

tspan = (0.0,60.0)
prob = ODEProblem(fluxes(m), u₀(m, 1,0), tspan, β)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)

u0 = zeros(Float64, length(m.S))
u0[1]  = 10000
u0[5]  = 10000
u0[9]  = 10000
u0[2]  = 1
u0
βseir = [10/sum(u0), 1/2, 1/5]
βtravel = [1/2, 1/2, 1/2]/1000
β = vcat(βseir, βtravel, βseir, βtravel, βseir)
@show β
prob = ODEProblem(fluxes(m), u0, tspan, β)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
@show sol.u[end]

using Plots
makeplots_seir(sol, "img/estimation/seir")

T = [
    # City 1 SEIₐIRD
    ([1, 2], [3, 2]), # S+Iₐ → E+Iₐ
    ([3], [2]),       # E  → Iₐ
    ([2], [4]),       # Iₐ → I
    ([2], [5]),       # Iₐ → R
    ([4], [5]),       # I  → R
    ([4], [6]),       # I  → D
]

m = Petri.Model(1:6, T, missing, missing)
nS = length(m.S)
β = ones(Float64, length(T))

tspan = (0.0,60.0)
prob = ODEProblem(fluxes(m), u₀(m, 1,0), tspan, β)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)

u0 = zeros(Float64, length(m.S))
u0[1]  = 60000000
u0[2]  = 1
u0

βseir = [10/5sum(u0), 1/2, 1/5, 1/3, 1/4, 1/7]
# βtravel = [1/2, 1/2, 1/2]/1000
# β = vcat(βseir, βtravel, βseir, βtravel, βseir)
β = βseir
@show β
prob = ODEProblem(fluxes(m), u0, tspan, β)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
@show sol.u[end]

# makeplots_seir(sol, "img/estimation/seiird")

plot(sol, vars=[3, 4, 5, 6], label=["IA" "I" "R" "D"])
# end

using CSV
using Query
using DataFrames

df = CSV.read("../eucovid.csv")

italy = @from i in df begin
    @where i.Country == "Italy" && i.Cases > 0
    @select i
    @collect DataFrame
end

caseplot(italy)

function lossfunc(df::DataFrame, m, u0)
    tspan = (0.0, length(df.Deaths)-1)
    mean(x) = sum(x)/length(x)
    t = tspan[1]:tspan[end]
    f(sol) = begin
        deaths = [u[end] for u in sol.(t)]
        recovr = [u[end-1] for u in sol.(t)]
        infsym = [u[end-2] for u in sol.(t)]
        dt = mean((deaths .- df.Deaths).^2) / mean(df.Deaths)
        rt = mean((recovr .- df.Recoveries).^2) / mean(df.Recoveries)
        it = mean((infsym .- df.Cases).^2) / mean(df.Cases)
        return dt + rt + it
    end

    g(θ) = begin
        prob = ODEProblem(fluxes(m), u0, tspan, θ)
        sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
        neg = sum(abs.(θ[θ .< 0]))
        λ = 10000000
        # if final deaths > final recoveries then penalty
        # deathpenalty = sol.u[end][end] > sol.u[end][end-1] ? 1 : 0
        # iapenalty = sum(tanh.([u[3] - u[2] for u in sol.(t)]).+1)
        iapenalty = 0
        lossval, regval = f(sol), λ*(neg + iapenalty)
        # lossval, regval = f(sol) , λ*(neg + deathpenalty)
        println("loss:$lossval\tregval:$regval")
        lossval+regval
    end
    return g
end

u0_italy = [60000000.0, 0, 1, 0, 0, 0]
ℓ = lossfunc(italy, m, u0_italy)

ℓ(β)

using Optim

function calibrate(k)
    β = [40/5sum(u0), 11/2, 1/1, 2/13, 1/18, 1/35]/k
    prob = ODEProblem(fluxes(m), u0_italy, (0.0,95), β)
    sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
    # plot(sol, vars=[2, 4, 5, 6, 3], label=["IA" "I" "R" "D" "E"])
    # p = fitplot(sol, italy)
    p = plot(sol, vars=[4, 5, 6], label=["I" "R" "D"],
            color=[RGB(1,0,0) RGB(0,1,0) RGB(0.5,0,0.5)])
    @df df scatter!(p,
            0:length(italy.Deaths)-1,
            [:Cases, :Deaths, :Recoveries],
            label=["cases" "deaths" "recoveries"],
            marker=true,
            color=[RGB(1,0,0) RGB(0.5,0,0.5) RGB(0,1,0)]
            )
    return β, sol, p
end

calibrate(15.5)[end]

β = [40/5sum(u0), 11/2, 1/1, 2/13, 1/18, 1/35]/15
# β = [60/5sum(u0), 11/2, 1/1, 2/13, 1/24, 1/39]/2
prob = ODEProblem(fluxes(m), u0_italy, (0.0,64), β)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
plot(sol, vars=[2, 4, 5, 6, 3], label=["IA" "I" "R" "D" "E"])
fitplot(sol, italy)

lb, ub = zeros(Float64, 6), 100*ones(Float64, 6)

res = optimize(ℓ, β)
# res = optimize(ℓ, lb, ub, β, Fminbox(NelderMead()))
@show res.minimizer
@show res.minimizer -β

# candidate minimizer
β̂ = [
 9.478916256341902e-9
 0.2915773795849941
 0.09434402982528367
 0.0061044262074341285
 0.016241417062038518
 0.02234216514573622
]

prob = ODEProblem(fluxes(m), u0_italy, (0.0,128), res.minimizer)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
fitplot(sol, italy)

prob = ODEProblem(fluxes(m), u0_italy, (0.0,70), res.minimizer)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
fitplot(sol, df)

# trying to add constraints to the parameter space
realparams(θ) = [θ[1:4]..., θ[5], θ[5]/5]
res = optimize(θ->ℓ(realparams(θ)), β[1:6])
@show res.minimizer
@show res.minimizer -β

prob = ODEProblem(fluxes(m), u0_italy, (0.0,128), realparams(res.minimizer))
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
fitplot(sol, df)


export bissect
function bissect(h       :: AbstractNLPModel,
                      nlpstop :: AbstractStopping;
                      t1 = h.meta.x0, t2 = 100.0,
                      verbose = true)

    (tₐ, tᵦ) = trouve_intervalle(h, t1, 1.0)
    gₐ = grad(h, tₐ)
    gᵦ = grad(h, tᵦ)
    iter = 0

    OK = update_and_start!(nlpstop, x = tₐ, gx = gₐ, g0 = copy(gₐ))

    verbose &&
        @printf(" iter        tₐ        tᵦ         gₐ        gᵦ        \n")
    verbose &&
        @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n", iter, tₐ[1], tᵦ[1], gₐ[1], gᵦ[1])
    while !OK
      tₚ = (tₐ .+ tᵦ) ./ 2
      gₚ = grad(h, tₚ)
      if true in (gₚ .<= 0.0)
        tₐ = tₚ
        gₐ = gₚ
        OK = update_and_stop!(nlpstop, x = tₐ, gx = gₐ)
      else
        tᵦ = tₚ
        gᵦ = gₚ
        OK = update_and_stop!(nlpstop, x = tᵦ, gx = gᵦ)
      end

      iter += 1
      verbose &&
        @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n", iter, tₐ[1], tᵦ[1], gₐ[1], gᵦ[1])
    end

    topt = (tₐ .+ tᵦ) ./ 2.0
    update!(nlpstop.current_state, x = topt, gx = grad(h, topt))
    optimal = nlpstop.meta.optimal
    return optimal, nlpstop
end


struct GLMResult{S, T}
    x::AbstractVector{S}
    y::AbstractVector{T}
    l::AbstractVector{T}
    u::AbstractVector{T}
end

function linear(x::AbstractVector{T1}, y::AbstractVector{T2}; n_points = 100) where {T1, T2}
    try
        lin_model = GLM.lm(@formula(Y ~ X), (X=x, Y=y))
        x_min, x_max = extrema(x)
        x_new = range(x_min, x_max, length = n_points)
        y_new, lower, upper = GLM.predict(lin_model,
         [ones(T1, n_points) x_new], interval=:confidence)
        # the GLM predictions always return matrices
       return GLMResult(x_new, vec(y_new), vec(lower), vec(upper))
    catch e
        error("Linear fit not possible for the given data")
    end
end

convert_arguments(P::PlotFunc, g::GLMResult) = PlotSpec{ShadedLine}(g.x, g.y, g.l, g.u)

struct Smooth{S, T}
    x::Vector{S}
    y::Vector{T}
end

convert_arguments(P::PlotFunc, l::Smooth) = PlotSpec{Lines}(Point2f0.(l.x,l.y))

function smooth(x, y; length = 100, kwargs...)
    model = loess(x, y; kwargs...)
    min, max = extrema(x)
    us = collect(range(min, stop = max, length = length))
    vs = Loess.predict(model, us)
    return Smooth(us, vs)
end

smooth(; kwargs...) = (args...) -> smooth(args...; kwargs...)


using QuasinormalModes
using SymEngine

# ------------------------------------------------------------------
# 1. Analytic harmonic oscillator
# ------------------------------------------------------------------

struct HarmonicOscilatorData{N,T} <: QuadraticEigenvalueProblem{N,T}
    nIter::N
    x0::T

    vars::Tuple{Basic,Basic}
    exprs::Tuple{Basic,Basic}
end

function HarmonicOscilatorData(nIter::N, x0::T) where {N,T}

    vars = @vars x ω

    λ0 = 2 * x
    S0 = 1 - ω

    return HarmonicOscilatorData{N,T}(nIter, x0, vars, (λ0, S0))
end

QuasinormalModes.λ0(d::HarmonicOscilatorData{N,T}) where {N,T} = d.exprs[1]
QuasinormalModes.S0(d::HarmonicOscilatorData{N,T}) where {N,T} = d.exprs[2]

QuasinormalModes.get_niter(d::HarmonicOscilatorData{N,T}) where {N,T} = d.nIter
QuasinormalModes.get_x0(d::HarmonicOscilatorData{N,T}) where {N,T} = d.x0

QuasinormalModes.get_ODEvar(d::HarmonicOscilatorData{N,T}) where {N,T} = d.vars[1]
QuasinormalModes.get_ODEeigen(d::HarmonicOscilatorData{N,T}) where {N,T} = d.vars[2]

# ------------------------------------------------------------------
# 2. Numeric harmonic oscillator
# ------------------------------------------------------------------

struct NHarmonicOscilatorData{N,T} <: NumericAIMProblem{N,T}
    nIter::N
    x0::T
end

function NHarmonicOscilatorData(nIter::N, x0::T) where {N,T}
    return NHarmonicOscilatorData{N,T}(nIter, x0)
end

QuasinormalModes.λ0(::NHarmonicOscilatorData{N,T}) where {N,T} = (x, ω) -> 2 * x
QuasinormalModes.S0(::NHarmonicOscilatorData{N,T}) where {N,T} = (x, ω) -> 1 - ω + x - x

QuasinormalModes.get_niter(d::NHarmonicOscilatorData{N,T}) where {N,T} = d.nIter
QuasinormalModes.get_x0(d::NHarmonicOscilatorData{N,T}) where {N,T} = d.x0

# ------------------------------------------------------------------
# 3. Constructing problems and caches
# ------------------------------------------------------------------

p_ana = HarmonicOscilatorData(0x0000A, 0.5);
p_num = NHarmonicOscilatorData(0x0000A, 0.5);

c_ana = AIMCache(p_ana)
c_num = AIMCache(p_num)

# ------------------------------------------------------------------
# 4. Computing quasinormal modes
# ------------------------------------------------------------------

ev_ana = computeEigenvalues(Serial(), p_ana, c_ana)
ev_num = eigenvaluesInGrid(Serial(), p_num, c_num, (0.0, 21.0))

function printEigen(eigenvalues)
    println("--------------------------------------")

    for i in eachindex(eigenvalues)
        println("n = $i, ω = $(eigenvalues[i])")
    end

    println("--------------------------------------")

    return nothing
end

println("Analytic results")
printEigen(reverse!(ev_ana))

println("Numeric results")
printEigen(ev_num)


# PowellSingular function
#
# source: Nielsen, Hans Bruun and others, Damping parameter in Marquardt's method, 1999, IMM
#
function PowellSingular()
NLS_ForwardDiff(4,4) do θ
        T = eltype(θ)

        T[θ[1]+10*θ[2],
          sqrt(5)*(θ[3]-θ[4]),
          (θ[2]-2*θ[3])^2,
          sqrt(10)*(θ[1]-θ[4])^2]

    end
end


x = randn(16,5)
wt = wavelet(WT.haar)
xw = cat([wpd(x[:,i], wt) for i in axes(x,2)]..., dims=3)
xsw = cat([swpd(x[:,i], wt) for i in axes(x,2)]..., dims=3)
xacw = cat([acwpd(x[:,i], wt) for i in axes(x,2)]..., dims=3)

# bb
@test isvalidtree(x[:,1], bestbasistree(xw[:,:,1], BB()))
@test isvalidtree(x[:,1], bestbasistree(xw, method=BB())[:,1])
@test isvalidtree(x[:,1],
    bestbasistree(xw, BB(LogEnergyEntropyCost(), false))[:,1])
@test isvalidtree(x[:,1], bestbasistree(xsw, BB(redundant=true))[:,1])
@test isvalidtree(x[:,1], bestbasistree(xacw, BB(redundant=true))[:,1])

# jbb
@test isvalidtree(x[:,1], bestbasistree(xw))
@test isvalidtree(x[:,1], bestbasistree(xw, JBB(NormCost(), false)))
@test isvalidtree(x[:,1], bestbasistree(xsw, JBB(redundant=true)))

# lsdb
@test isvalidtree(x[:,1], bestbasistree(xw, method=LSDB()))
@test isvalidtree(x[:,1], bestbasistree(xw, LSDB()))
@test isvalidtree(x[:,1], bestbasistree(xsw, LSDB(redundant=true)))

# siwpd_bb
x = randn(16)
xw0 = siwpd(x, wt)
xw4 = siwpd(circshift(x,4), wt)
bt0 = map(node -> sum(node), bestbasistree(xw0, 4, SIBB()))
bt4 = map(node -> sum(node), bestbasistree(xw4, 4, SIBB()))
@test bt0 == bt4

# misc
@test_throws ArgumentError bestbasis_treeselection(randn(15), 8, :fail)
@test_throws AssertionError bestbasis_treeselection(randn(7), 3, :fail) # true n=4


using ShapleyValues
using StatsBase
using Base.Test

include("utils.jl")
include("coalitions.jl")
include("samplespaces.jl")

# simple test of logistic and linear regression
srand(1)
K = 4
X = rand(K,100)
beta = randn(K)
f(x) = x'beta
p(x) = logistic(f(x))
x = 2*randn(K)

"A direct way to compute the expected value conditioned on the given inds."
function raw_exp(f, x, inds, X)
    Xtmp = copy(X)
    for i in inds
        Xtmp[i,:] = x[i]
    end
    mean(f(Xtmp))
end

fnull,φ,φVar = shapleyvalues(x, f, X) # linear regression
@test raw_exp(f, x, 1:K, X) - raw_exp(f, x, Int64[], X) ≈ sum(φ)
@test sum(abs(φVar)) < 1e-10
fnull,φ,φVar = shapleyvalues(x, f, X, nsamples=8)
fnull,φ,φVar = shapleyvalues(x, f, sparse(X), nsamples=8)

fnull,φ,φVar = shapleyvalues(x, p, X, logit) # logistic regression
@test logit(raw_exp(p, x, 1:K, X)) - logit(raw_exp(p, x, Int64[], X)) ≈ sum(φ)
@test sum(abs(φVar)) < 1e-10

# test with many features
srand(1)
K = 400
X = rand(K,10)
beta = randn(K)
f(x) = x'beta
p(x) = logistic(f(x))
x = 2*randn(K)
fnull,φ,φVar = shapleyvalues(x, f, X)
@test raw_exp(f, x, 1:K, X) - raw_exp(f, x, Int64[], X) ≈ sum(φ)

# test with many features
srand(1)
K = 400
X = rand(K,10)
beta = randn(K)
f(x) = x'beta
p(x) = logistic(f(x))
x = .02*randn(K)
fnull,φ,φVar = shapleyvalues(x, p, X, logit)
@test logit(raw_exp(p, x, 1:K, X)) - logit(raw_exp(p, x, Int64[], X)) ≈ sum(φ)


using StatsBase

export summarize

abstract type Summary end

summarize(nt::NamedTuple) = rmap(summarize, nt)

summarize(tv::TupleVector) = summarize(unwrap(tv))

import Base

###############################################################################
# FullSummary

struct FullSummary{T} <: Summary
    data :: T
end

summarize(x) = FullSummary(x)

Base.show(io, s::FullSummary) = print(io, s.data)

###############################################################################
# RealSummary
export RealSummary
struct RealSummary <: Summary
    μ :: Float64
    σ :: Float64

    function RealSummary(μ,σ)
        μ = isnan(μ) ? 0.0 : μ
        σ = isnan(σ) ? 0.0 : σ
        new(μ,σ)
    end
end

Base.typeinfo_prefix(io::IO, ::AbstractArray{<:Summary}) = ("", false)

function Base.show(io::IO, s::RealSummary)
    io = IOContext(io, :compact => true)
    σ = round(s.σ, sigdigits=2)
    if s.μ == 0 || σ == 0
        μdigits = 2
    else
        μdigits = max(0, ceil(Int, log10(2) * (exponent(s.μ) - exponent(σ))) + 2)
    end
    μ = round(s.μ, sigdigits = μdigits)
    print(io, μ, "±", σ)
end

summarize(x::AbstractVector{<:Real}) = RealSummary(mean_and_std(x)...)

function summarize(x::AbstractArray{<:AbstractArray{T}}) where {T <: Real}
    [RealSummary(μ,σ) for (μ,σ) in zip(mean_and_std(x)...)]
end


"""
    factorial(n::Int)

Compute the factorial of n.
"""
function factorial(n::Int)
    if n < 0
        throw(DomainError("n must be non-negative"))
    end
    result = 1
    for i in 1:n
        result *= i
    end
    return result
end


using ModelingToolkit, OrdinaryDiffEq, Test

@parameters t σ ρ β
@variables x(t) y(t) z(t)
D = Differential(t)

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]

@named lorenz1 = ODESystem(eqs)
@named lorenz2 = ODESystem(eqs)

@parameters γ
@variables a(t) α(t)
connections = [0 ~ lorenz1.x + lorenz2.y + a*γ,
               α ~ 2lorenz1.x + a*γ]
@named sys = ODESystem(connections,t,[a,α],[γ],systems=[lorenz1,lorenz2])
sys_simplified = structural_simplify(sys)

u0 = [lorenz1.x => 1.0,
      lorenz1.y => 0.0,
      lorenz1.z => 0.0,
      lorenz2.x => 0.0,
      lorenz2.y => 1.0,
      lorenz2.z => 0.0,
      a => 2.0]

p  = [lorenz1.σ => 10.0,
      lorenz1.ρ => 28.0,
      lorenz1.β => 8/3,
      lorenz2.σ => 10.0,
      lorenz2.ρ => 28.0,
      lorenz2.β => 8/3,
      γ => 2.0]

tspan = (0.0,100.0)
prob = ODEProblem(sys_simplified,u0,tspan,p)
sol = solve(prob,Rodas4())

@test_throws UndefVarError sol[b]
@test_throws UndefVarError sol[b, 1]
@test_throws UndefVarError sol[b, 1:5]
@test_throws UndefVarError sol[b, [1,2,3]]
@test_throws ErrorException sol['a']
@test_throws ErrorException sol['a',  1]
@test_throws ErrorException sol['a',  1:5]
@test_throws ErrorException sol['a',  [1,2,3]]

@test sol[a] isa AbstractVector
@test sol[a,  1] isa Real
@test sol[a,  1:5] isa AbstractVector
@test sol[a,  [1,2,3]] isa AbstractVector

@test sol[1] isa AbstractVector
@test sol[1:2] isa AbstractArray
@test sol[[1,2]] isa AbstractArray

@test sol[lorenz1.x] isa Vector
@test sol[lorenz1.x,2] isa Float64
@test sol[lorenz1.x,:] isa Vector
@test sol[t] isa Vector
@test sol[t,2] isa Float64
@test sol[t,:] isa Vector
@test length(sol[lorenz1.x,1:5]) == 5
@test sol[α] isa Vector
@test sol[α,3] isa Float64
@test length(sol[α,5:10]) == 6

# Check if indexing using variable names from interpolated solution works
interpolated_sol = sol(0.0:1.0:10.0)
@test interpolated_sol[α] isa Vector
@test interpolated_sol[α,:] isa Vector
@test interpolated_sol[α,2] isa Float64
@test length(interpolated_sol[α,1:5]) == 5
@test interpolated_sol[α] ≈ 2interpolated_sol[lorenz1.x] .+ interpolated_sol[a].*2.0
@test collect(interpolated_sol[t]) isa Vector
@test collect(interpolated_sol[t,:]) isa Vector
@test interpolated_sol[t,2] isa Float64
@test length(interpolated_sol[t,1:5]) == 5

sol1 = sol(0.0:1.0:10.0)
@test sol1.u isa Vector
@test first(sol1.u) isa Vector
@test length(sol1.u) == 11
@test length(sol1.t) == 11

sol2 = sol(0.1)
@test sol2 isa Vector
@test length(sol2) == length(states(sys_simplified))
@test first(sol2) isa Real

sol3 = sol(0.0:1.0:10.0, idxs=[lorenz1.x, lorenz2.x])
@test sol3.u isa Vector
@test first(sol3.u) isa Vector
@test length(sol3.u) == 11
@test length(sol3.t) == 11
@test collect(sol3[t]) ≈ sol3.t
@test collect(sol3[t, 1:5]) ≈ sol3.t[1:5]
@test_throws ErrorException sol(0.0:1.0:10.0, idxs=[lorenz1.x, 1])

sol4 = sol(0.1, idxs=[lorenz1.x, lorenz2.x])
@test sol4 isa Vector
@test length(sol4) == 2
@test first(sol4) isa Real
@test_throws ErrorException sol(0.1, idxs=[lorenz1.x, 1])

sol5 = sol(0.0:1.0:10.0, idxs=lorenz1.x)
@test sol5.u isa Vector
@test first(sol5.u) isa Real
@test length(sol5.u) == 11
@test length(sol5.t) == 11
@test collect(sol5[t]) ≈ sol3.t
@test collect(sol5[t, 1:5]) ≈ sol3.t[1:5]
@test_throws ErrorException sol(0.0:1.0:10.0, idxs=1.2)

sol6 = sol(0.1, idxs=lorenz1.x)
@test sol6 isa Real
@test_throws ErrorException sol(0.1, idxs=1.2)

sol7 = sol(0.0:1.0:10.0, idxs=[2,1])
@test sol7.u isa Vector
@test first(sol7.u) isa Vector
@test length(sol7.u) == 11
@test length(sol7.t) == 11
@test collect(sol7[t]) ≈ sol3.t
@test collect(sol7[t, 1:5]) ≈ sol3.t[1:5]

sol8 = sol(0.1, idxs=[2,1])
@test sol8 isa Vector
@test length(sol8) == 2
@test first(sol8) isa Real

sol9 = sol(0.0:1.0:10.0, idxs=2)
@test sol9.u isa Vector
@test first(sol9.u) isa Real
@test length(sol9.u) == 11
@test length(sol9.t) == 11
@test collect(sol9[t]) ≈ sol3.t
@test collect(sol9[t, 1:5]) ≈ sol3.t[1:5]

sol10 = sol(0.1, idxs=2)
@test sol10 isa Real

#=
using Plots
plot(sol,vars=(lorenz2.x,lorenz2.z))
plot(sol,vars=(α,lorenz2.z))
plot(sol,vars=(lorenz2.x,α))
plot(sol,vars=α)
plot(sol,vars=(t,α))
=#


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Desription
# ==============================================================================
#
#   Tests related to conversion from quaternions to Euler angle and axis.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File: ./src/conversions/quaternion_to_angleaxis.jl
# ==================================================

# Functions: quat_to_angleaxis
# ----------------------------

@testset "Quaternion => Euler angle and axis (Float64)" begin
    T = Float64

    q = Quaternion(cosd(T(75 / 2)), 0, sind(T(75 / 2)), 0)
    av = quat_to_angleaxis(q)
    @test eltype(av) === T
    @test av.a ≈ 75 * pi / 180
    @test av.v ≈ [0, 1, 0]

    q = Quaternion(cosd(T(225 / 2)), 0, sind(T(225 / 2)), 0)
    av = quat_to_angleaxis(q)
    @test eltype(av) === T
    @test av.a ≈ 135 * pi / 180
    @test av.v ≈ [0, -1, 0]
end

@testset "Quaternion => Euler angle and axis (Float32)" begin
    T = Float32

    q = Quaternion(cosd(T(75 / 2)), 0, sind(T(75 / 2)), 0)
    av = quat_to_angleaxis(q)
    @test eltype(av) === T
    @test av.a ≈ T(75 * pi / 180)
    @test av.v ≈ [0, 1, 0]

    q = Quaternion(cosd(T(225 / 2)), 0, sind(T(225 / 2)), 0)
    av = quat_to_angleaxis(q)
    @test eltype(av) === T
    @test av.a ≈ T(135 * pi / 180)
    @test av.v ≈ [0, -1, 0]
end

@testset "Quaternion => Euler angle and axis (Special cases)" begin
    q = Quaternion{Float64}(1, 0, 0, 0)
    av = quat_to_angleaxis(q)
    @test eltype(av) == Float64
    @test av.a == 0
    @test av.v == [0, 0, 0]

    q = Quaternion{Float32}(-1, 0, 0, 0)
    av = quat_to_angleaxis(q)
    @test eltype(av) == Float32
    @test av.a == 0
    @test av.v == [0, 0, 0]
end


"""
    nearestneighbor(x, y)

Nearest neighbor interpolation with the spline values in the center of intervals.

- aliases: `spline0`, `stepspline`
"""
function nearestneighbor(x::Vector, y::Vector)
    NearestNeighborSpline(y, x)
end

const spline0 = nearestneighbor
const stepspline = nearestneighbor
const nspline = nearestneighbor

"""
    linearspline(x, y)

Linear spline interpolation.

- aliases: `spline1`, `lspline`
"""
function linearspline(x::Vector, y::Vector)
    #n = length(x)
    b = Vector{Float64}(undef, length(x))

    @inbounds for i in 1:length(x)-1
        b[i] = (y[i+1] - y[i]) / (x[i+1] - x[i])
    end
    #yy = y[begin+1:end].-y[begin:end-1]
    #xx = x[begin+1:end].-x[begin:end-1]
    LinearSpline(b, y, x)
end

const spline1 = linearspline
const lspline = linearspline

"""
    quadraticspline(x::Vector, y::Vector, b1 = 0.0)

Quadratic spline interpolation.

- aliases: `spline2`, `qspline`

Optional argument `b1` sets the 1st derivative at the first knot. Default value = `0.0`.
"""
function quadraticspline(x::Vector, y::Vector, b1 = 0.0)
    n = length(x)
    b = Vector{Float64}(undef, n-1)
    c = Vector{Float64}(undef, n-1)
    b[1] = b1

    Δx = Vector{Float64}(undef, n-1)
    Δy = Vector{Float64}(undef, n-1)

    @inbounds for i in 1:n-1
        Δx[i] = x[i+1] - x[i]
        Δy[i] = y[i+1] - y[i]
    end

    @inbounds for i in 1:n-1
        b[i+1] = -b[i] + 2 * Δy[i] / Δx[i]
    end

    @inbounds for i in 1:n-1
         c[i] = (b[i+1] - b[i]) / (2 * Δx[i])
    end
    QuadraticSpline(c, b, y, x)
end

const spline2 = quadraticspline
const qspline = quadraticspline

"""
    cubicspline(x::Vector, y::Vector, style...)

- alias: `spline3`, `cspline`

Creates a cubic spline. Boundary conditions are defined by `style`.

`style...`: settings for boundary conditions.

- `:natural`: zero curvature at the start and end points of the spline.
    - aliases: `cubicspline(x, y)`, `naturalspline(x, y)`
- `:clamped, ds, de`: predefined 1st derivative of start (ds) and end points (de).
    - alias: `clampedspline(x, y, ds, de)`
- `:periodic`: periodic boundary conditions. `y` values at start and end points have to be equal,
conditions at start and end points are identical.
    - alias: `periodicspline(x, y)`
- `:parabolic`, `:notaknot`: parabolic runout spline, also known as not-a-knot spline
- `:cubic`: cubic runout spline
- `:constrained`: constrained cubic spline, see Kruger’s paper
- `:hermite`: hermite spline
    - `:kruger`: constrained cubic spline, see Kruger’s paper. No overshooting.
    - `:free`: hermite spline with free derivative conditions. No overshooting.
    - `:predef`: predefined derivatives at knots

# Examples:
```jldoctest
julia> x=[0,1,2,3,4]; y=[-1,1,5,2,-3]; spline3(x,y,:natural)
CubicSpline([1.0, -3.0, 2.0, 0.0], [0.0, 3.0, -6.0, 0.0], [1.0, 4.0, 1.0, -5.0], [-1.0, 1.0, 5.0, 2.0])
```
"""
function cubicspline(x::Vector, y::Vector, style...)
    if style[1] == :periodic
        y[1] != y[end] ? error("Periodic spline error: different y values at endpoints!") : nothing
    elseif style[1] == :constrained
        return hermitespline(x, y, :kruger)
    elseif style[1] == :hermite
        if length(style) == 1
            return hermitespline(x, y, :kruger)
        elseif length(style) == 2
            return hermitespline(x, y, style[2])
        else
            throw(DomainError(style[1], "Spline style does not exist."))
        end
    else
        n = length(x)
        #Δx= Vector{Float64}(undef, n-1)
        Δx = zeros(Float64, n-1)
        #Δy= Vector{Float64}(undef, n-1)
        Δy = zeros(Float64, n-1)
        #λ = Vector{Float64}(undef, n)    # upper diagonal
        λ = zeros(Float64, n)
        δ = ones(n) * 2 # main diagonal
        #μ = Vector{Float64}(undef, n)    # lower diagonal
        μ = zeros(Float64, n)
        #r = Vector{Float64}(undef, n)    # derivatives
        r = zeros(Float64, n)
        #M = Vector{Float64}(undef, n)
        #M = zeros(Float64, n)

        @inbounds for i in 1:n-1
            Δx[i] = x[i+1] - x[i]
            Δy[i] = y[i+1] - y[i]
        end
        if style == :notaknot
            A = zeros(n, n)
            #=@inbounds=# for i in 2:n-1
                A[i, i-1] = Δx[i-1]
                A[i, i] = 2.0 * (Δx[i-1] + Δx[i])
                A[i, i+1] = Δx[i+1]
                r[i] = 6.0 * (Δy[i] / Δx[i] - Δy[i-1] / Δx[i-1])
            end
            #A[begin, begin] = -Δx[begin + 2]
            #A[begin, begin + 1] = Δx[begin + 1] + Δx[begin + 2]
            #A[begin, begin + 2] = -Δx[begin + 1]
            A[1,1] = -(x[3] - x[2])
            A[1,2] = x[3] - x[1]
            A[1,3] = -(x[2] - x[1])
            #A[end, end - 2] = -Δx[end - 1]
            #A[end, end - 1] = Δx[end - 2] + Δx[end - 1]
            #A[end, end]   = Δx[end - 2]
            A[n-1, n-3] = -(x[n-1] - x[n-2])
            A[n-1, n-2] = x[n-1] - x[n-3]
            A[n-1, n-1] = -(x[n-2] - x[n-3])
            #A[n-1, n-1] = -(x[n-3] - x[n-2])
            #A[n-1, n-2] = x[n-3] - x[n-1]
            #A[n-1, n-3] = -(x[n-2] - x[n-1])
            #for i in 2:n-1
            #    r[i] = 6.0 / (Δx[i-1] + Δx[i]) * (Δy[i] / Δx[i] - Δy[i-1] / Δx[i-1])
            #    #r[i] = 6.0 * (Δy[i] / Δx[i] - Δy[i-1] / Δx[i-1])
            #end
        else
            #=@inbounds=# for i in 2:n-1
                λ[i] = Δx[i] / (Δx[i-1] + Δx[i])
                μ[i] = 1.0 - λ[i]
                r[i] = (6.0 / (Δx[i-1] + Δx[i])) * (Δy[i] / Δx[i] - Δy[i-1] / Δx[i-1])
            end
        end

        e = length(Δx)
        r = rbound(r, e, Δx, Δy, style...)
        if style ≠ :notaknot
            λ, μ = λµbound(λ, μ, e, Δx, style...)
            A = Tridiagonal(μ[2:n], δ, λ[1:n-1])
            M = A \ r
            M = Mbound(M,style...)
        else
            M = A \ r
        end
        b = zeros(Float64, n-1) # Vector{Float64}(undef, n-1)
        c = zeros(Float64, n-1) # Vector{Float64}(undef, n-1)
        d = zeros(Float64, n-1) # Vector{Float64}(undef, n-1)

        #=@inbounds=# for i in 1:n-1
            d[i] = (M[i+1] - M[i]) / (6 * Δx[i])
            c[i] = 0.5 * M[i] #
            b[i] = Δy[i] / Δx[i] - Δx[i] / 6.0 * (2.0 * M[i] + M[i+1])
        end
    end
    if style[1] in (:hermite, :constrained)
        nothing
    else
        return CubicSpline(d, c, b, y, x)
    end
end

const spline3 = cubicspline
const cspline = cubicspline

"""
    rbound(r,Δx,Δy,style...)

Returns vector `r` according to boundary conditions defined by `style...`
Helper function for `cubicspline`.

Returns: `r`
"""
function rbound(r, e, Δx, Δy, style...)
    if style[1] in (:natural, :notaknot)
        r[1]   = 0.0
        r[end] = 0.0
    elseif style[1] == :parabolic
        r[1] = r[2]
        r[end] = r[end-1]
    elseif style[1] == :clamped
        ds     = style[2]
        de     = style[3]
        r[1]   = 6.0 / Δx[1] * (Δy[1] / Δx[1] - ds)    #ds = start derivative
        r[end] = -6.0 / Δx[e] * (Δy[e] / Δx[e] - de)   #de = end derivative
    elseif style[1] == :periodic
        r[end] = 6.0 / (Δx[1] + Δx[e]) * (Δy[1] / Δx[1] - Δy[e] / Δx[e])
        r[1]   = r[end]
    end
    return r
end

"""
    Mbound(M,style...)

Returns vector `M` (moments) according to boundary conditions defined by `style...`

Helper function for `cubicspline`.
"""
function Mbound(M, style...)
    if style[1] == :natural
        M[1]   = 0.0
        M[end] = 0.0
    elseif style[1] == :periodic
        M[1] = M[end]
    elseif style[1] == :parabolic
        M[1] = M[2]
        M[end] = M[end-1]
    elseif style[1] == :notaknot
        #M[1] = M[2]
        #M[end] = M[end-1]
    end
    return M
end

function λµbound(λ,μ,e,Δx,style...)
    if style[1] == :natural
        λ[1]   = 0.0
        μ[end] = 0.0
    elseif style[1] == :clamped
        λ[1]   = 1.0
        μ[end] = 1.0
    elseif style[1] == :parabolic
        λ[1] = λ[2]
        μ[end] = λ[end-1]
    elseif style[1] == :periodic
        λ[end] = Δx[1] / (Δx[1] + Δx[e-1])
        μ[end] = 1 - Δx[1] / (Δx[1] + Δx[e])
    end
    return λ, µ
end

"""
    cubicspline(x, y)
Fall back function, setting the default boundary condition to `:natural`

- alias: `naturalspline`
"""
function cubicspline(x, y)
    cubicspline(x, y, :natural)
end

"""
    naturalspline(x::Vector, y::Vector)

Equivalent to `cubicspline(x::Vector, y::Vector, :natural)`.
"""
naturalspline(x, y) = cubicspline(x, y, :natural)

"""
    clampedspline(x::Vector, y::Vector, ds = 0.0, de = 0.0)

Equivalent to `cubicspline(x::Vector, y::Vector, ds, de)`.

Returns a clamped cubic spline.

Default arguments: 1st and last derivatives = `0.0`
"""
function clampedspline(x, y, ds = 0.0, de = 0.0)
    cubicspline(x, y, :clamped, ds, de)
end

"""
    periodicspline(x::Vector, y::Vector)

Equivalent to `cubicspline(x::Vector, y::Vector, :periodic)`.

`y` values of start and end points have to be equal.
"""
function periodicspline(x, y)
    cubicspline(x, y, :periodic)
end

notaknotspline(x, y) = cubicspline(x, y, :notaknot)

"""
    hermitespline(x::Vector, y::Vector, style::Symbol = :kruger, args...)

Creates cubic hermite interpolation spline with defined 1st derivatives at all knots.

Allowed `style`s:
- `:kruger`: Default value. Defines 1st derivatives using harmonic means to prevent overshoot at extremal knots.
- `:free`: Simple mean values for 1st derivatives.
- `:predef`: user defined 1st derivatives at all knots.
"""
function hermitespline(x::Vector, y::Vector, style::Symbol = :kruger, args...)
    n = length(x)
    # 1st derivatives
    b = Vector{Float64}(undef, n)
    # finite differences
    Δx = Vector{Float64}(undef, n-1)
    Δy = Vector{Float64}(undef, n-1)

    @inbounds for i in 1:n-1
        Δx[i] = x[i+1] - x[i]
        Δy[i] = y[i+1] - y[i]
    end

    #generate 1st derivatives at kntos

    # calculate derivatives based on C.J.C. Kruger,
    # [`Constrained Cubic Spline Interpolation for Chemical Engineering Applications`](https://pages.uoregon.edu/dgavin/software/spline.pdf), 2002
    # which utilizes the harmonic mean of the neighboring slopes of each knot.
    if style == :kruger
        @inbounds for i in 2:n-1
            if Δy[i-1] * Δy[i] > 0.0
                b[i] = 2.0 / (Δx[i] / Δy[i] + Δx[i-1] / Δy[i-1]) #harmonic mean
            end
        end
        # boundary knots
        b[1] = 1.5 * Δy[1]/Δx[1] - b[2] / 2
        b[n] = 1.5 * Δy[n-1]/Δx[n-1] - b[n-1] / 2

    # 1st derivatives at knots are arithmetic means of neighboring slopes
    elseif style == :free
        @inbounds for i in 2:n-1
            b[i] = Δy[i] / Δx[i]
        end
        # boundary knots
        b[1] = Δy[1] / Δx[1]
        b[n] = Δy[n-1] / Δx[n-1]
    elseif style == :predef
        if length(args) == 0
            throw(ArgumentError("Derivative vector missing."))
        end
        if length(args[1]) ≠ n-1
            throw(DimensionMismatch("x, y, and derivative vectors must have the same length"))
        end
        @inbounds for i in 1:length(args[1])
            b[i] = args[1][i]
        end
    else
        throw(DomainError(style, "Derivative constraint style does not exist!"))
    end

    HermiteSpline(b, y, x)
end

const hspline = hermitespline

#    for i in 1:n-1
#        Φ0 = 2 * Δx[i]^3 - 3 * Δx[i]^2 + 1
#        Φ1 = Δx[i]^3 - 2 * Δx[i]^2 + Δx[i]
#        Ψ0 = -2 * Δx[i]^3 + 3 * Δx[i]^3
#        Ψ1 = Δx[i]^3 - Δx[i]^2
#
#        Φ0_ = 6 * Δx[i]^2 - 6 * Δx[i]
#        Φ1_ = 3 * Δx[i]^2 - 4 * Δx[i] + 1
#        Ψ0_ = -6 * Δx[i]^2 + 6 * Δx[i]
#        Ψ1_ = 3 * Δx[i]^2 - 2 * Δx[i]
#
#        Φ0__ = 12 * Δx[i] - 6
#        Φ1__ = 6 * Δx[i] - 4
#        Ψ0__ = -12 * Δx[i] + 6
#        Ψ1__ = 6 * Δx[i] - 2
#
#        Φ0___ = 12
#        Φ1___ = 6
#        Ψ0___ = -12
#        Ψ1___ = 6
#        #xb = (xs[i] - spl.x[j]) / (spl.x[j+1] - spl.x[j])
#        c[i] = (y[i] * Φ0_ + Δx[i] * b[i] * Φ1_ + y[i+1] * Ψ0_ + Δx[i] * b[i+1] * Ψ1_) / Δx[i]
#        #d[i] = (y[i] * Φ0__ + Δx[i] * b[i] * Φ1__ + y[i+1] * Ψ0__ + Δx[i] * b[i+1] * Ψ1__) / Δx[i]
#    end

    #CubicSpline(d, c, b, y, x)

"""
    nakspline(x::Vector, y::Vector)

Creates not-a-knot spline.
"""
function nakspline(x::Vector, y::Vector)
    n = length(x)
    Δx = zeros(Float64, n)
    Δy = zeros(Float64, n)
    r = zeros(Float64, n-1)
    for i in 1:n-1
        Δx[i] = x[i+1] - x[i]
        Δy[i] = y[i+1] - y[i]
    end
    A = zeros(n-1, n-1)

    for i in 2:n-2
        A[i,i-1] = Δx[i]
        A[i,i]   = 2 * (Δx[i] + Δx[i+1])
        A[i,i+1] = Δx[i+1]
        r[i] = 3 * Δy[i+1] / Δx[i+1] - 3 * Δy[i] / Δx[i]
    end

    A[1,1] = Δx[2]
    A[1,2] = -(Δx[1] + Δx[2])
    A[1,3] = Δx[1]

    A[n-1, n-3] = Δx[n-1]
    A[n-1, n-2] = -(Δx[n-2] + Δx[n-1])
    A[n-1, n-1] = Δx[n-2]

    r[1] = r[2]
    r[n-1] = r[n-2]

    #return A,r

    M = A \ r

    b = zeros(Float64, n-1) # Vector{Float64}(undef, n-1)
    c = zeros(Float64, n-1) # Vector{Float64}(undef, n-1)
    d = zeros(Float64, n-1) # Vector{Float64}(undef, n-1)
    for i in 1:n-2
        #d[i] = (M[i+1] - M[i]) / (6 * Δx[i])
        d[i] = (M[i+1] - M[i]) / (3 * Δx[i])
        c[i] = 0.5 * M[i] #
        #b[i] = Δy[i] / Δx[i] - Δx[i] / 6.0 * (2.0 * M[i] + M[i+1])
        b[i] = Δy[i] / Δx[i] - (M[i+1] + 2 * M[i]) * Δx[i] / 3
    end
    return CubicSpline(d, c, b, y, x)
end


# Michelle Medina
# July 21, 2020
# Compute the cost vs. accuracy for
# Hierarchical basis vs. sparse grids

export cost1D, cost2D, cost3D, cost7D

using PyPlot
# ---------------------------------------------------------
function Base.:-(u::NodalBasis2D{T}, w::NodalBasis2D{T})::NodalBasis2D{T} where {T}
	value = NodalBasis2D(u.values - w.values)
	return value
end

using LinearAlgebra

function LinearAlgebra.norm(u::NodalBasis2D{T})::T where {T}
	return norm(u.values)
end

function cost1D(level::Int)::Int
	if level == 0
		return 2
	else
		return 2^(level-1)
	end
end

function cost2D(level::Int)::NTuple{2,Int}
    full = 0
	sparse = 0
	for l in 0:level, m in 0:level
        full = full + cost1D(l)*cost1D(m)
       	if l + m <= level
        	sparse = sparse + cost1D(l)*cost1D(m)
        end
	end
    return (full, sparse)
end

function cost3D(level::Int)::NTuple{2,Int}
    full   = 0
    sparse = 0
    for l in 0:level, m in 0:level, n in 0:level
        full = full + cost1D(l)*cost1D(m)*cost1D(n)
        if l + m + n <= level
            sparse = sparse + cost1D(l)*cost1D(m)*cost1D(n)
        end
    end
    return (full, sparse)
end

function cost7D(level::Int)::NTuple{2,Float64}
    full   = 0.0
    sparse = 0.0
    for l in 0:level, m in 0:level, n in 0:level, o in 0:level, p in 0:level, q in 0:level, r in 0:level
        full = full + 1.0 * cost1D(l)*cost1D(m)*cost1D(n)*cost1D(o)*cost1D(p)*cost1D(q)*cost1D(r)
        if l + m + n + o + p + q + r <= level
            sparse = sparse + 1.0 * cost1D(l)*cost1D(m)*cost1D(n)*cost1D(o)*cost1D(p)*cost1D(q)*cost1D(r)
        end
    end
    return (full, sparse)
end

#Cost (number of total points) vs. accurary (interpolation error)
#for both full basis and sparse basis
error_full = []
error_sparse = []
cost_full = []
cost_sparse = []

level = 3
for l in 0:level
	N = 2^l + 1

	xVal = collect(range(0.0, stop=1.0, length = N))
	yVal = collect(range(0.0, stop=1.0, length = N))
	nodalf = calcNodal2D((x,y)->exp(x)*exp(y), xVal, yVal)
	modalf = Nodal_2_H(nodalf)
	modalf_sparse = Nodal_2_H_sparse(nodalf,l)
	costF, costS = cost2D(l)

	# Compute the interpolation error
	le = 10 * (2^l+1)
	h = 1/le
	xNew = collect(range(0.0 + h/2, stop=1.0 - h/2, length = le))
	yNew = collect(range(0.0 + h/2, stop=1.0 - h/2, length = le))
	# xNew = xNew[2:end-1]
	# yNew = yNew[2:end-1]
	f_analytic = calcNodal2D((x,y)->exp(x)*exp(y), xNew, yNew)
	f_nodal_from_modalf = calcNodal2D((x,y)->evaluate(modalf, x, y), xNew, yNew)
	f_nodal_from_modalf_sparse = calcNodal2D((x,y)->evaluate(modalf_sparse, x, y), xNew, yNew)

	# error_f = sum(abs.((f_analytic - f_nodal_from_modalf).values))
	# error_s = sum(abs.((f_analytic - f_nodal_from_modalf_sparse).values))
	error_f = sqrt(sum((f_analytic.values - f_nodal_from_modalf.values).^2)/length(f_analytic.values))
	error_s = sqrt(sum((f_analytic.values - f_nodal_from_modalf_sparse.values).^2)/length(f_analytic.values))

	append!(error_full, error_f)
	append!(error_sparse, error_s)
	append!(cost_full, costF)
	append!(cost_sparse, costS)

	@show error_f, error_s, cost_full, cost_sparse

	# if l == 4
	# 	fig = figure()
	# 	subplot(1,2,1)
	# 	contourf((f_analytic - f_nodal_from_modalf).values)
	# 	colorbar()
	#
	# 	subplot(1,2,2)
	# 	contourf((f_analytic - f_nodal_from_modalf_sparse).values)
	# 	colorbar()
	# 	show()
	# end
	#
	#
end

# loglog(error_full, cost_full, "g-o", label = "Full")
# loglog(error_sparse, cost_sparse, "b-o", label = "Sparse")
# title("Cost vs Interpolation Error")
# legend()
# ylabel("Cost")
# xlabel("Interpolation Error")
# show()

# for l in 0:10
# 	@show l cost7D(l)
# end

# Heat Map of how coefficients fall off with level
# Tests how coefficients fall off with level
# n = 5
# xValue = collect(range(0.0, stop=1.0, length = 2^n + 1))
# yValue = collect(range(0.0, stop=1.0, length = 2^n + 1))
# nodalTest = calcNodal2D((x,y)->sin(pi*x)*sin(pi*y) + cos(pi*x)*cos(pi*y), xValue, yValue)
# modal = Nodal_2_H(nodalTest)
#
# A = zeros(6,6)
# for l in 0:5, m in 0:5
# 	A[l+1, m+1] = maximum(abs.(modal.levels[l+1, m+1].coefficients))
# end
#
# # use heat map
# using PyPlot
# imshow(A, vmin=0, vmax=0.5)
# title("f(x,y) = sin(pi*x)*sin(pi*y) + cos(pi*x)*cos(pi*y)")
# colorbar()
# show()


using CsminWel
using Base.Test

f(x) = sum(x.^2)
function grad(x)
    2*x, false
end

function g!(stor, x)
    stor[:] = 2*x
end

res0 = CsminWel.csminwel(f, grad, [0.1,0.1])
res1 = CsminWel.csminwel(f, [0.1,0.1])

res2 = CsminWel.csminwel(f, grad, [0.1,0.1])
res3 = CsminWel.csminwel(f, [0.1,0.1], H = eye(2))

res4 = optimize(f, g!, [.1, .1], Csminwel())
res5 = optimize(f, [.1, .1], Csminwel())   ## Default to finite difference
res6 = optimize(OnceDifferentiable(f, [.1, .1]; autodiff = :forward), [.1, .1], Csminwel())

@test Optim.minimum(res0) ≈ Optim.minimum(res1) atol=1.0e-9
@test Optim.minimum(res1) ≈ Optim.minimum(res2) atol=1.0e-9
@test Optim.minimum(res2) ≈ Optim.minimum(res3) atol=1.0e-9
@test Optim.minimum(res3) ≈ Optim.minimum(res4) atol=1.0e-9
@test Optim.minimum(res4) ≈ Optim.minimum(res5) atol=1.0e-9
@test Optim.minimum(res6) ≈ Optim.minimum(res6) atol=1.0e-9

#Maximizing loglikelihood logistic models
using StatsFuns
srand(1)
x = [ones(200) randn(200,4)]
y = [rand() < 0.5 ? 1. : 0. for j in 1:200]

function loglik(beta)
    xb = x*beta
    sum(-y.*xb + log1pexp.(xb))
end

function dloglik(beta)
    xb = x*beta
    px = logistic.(xb)
    -x'*(y.-px)
end

function g!(stor, beta)
    copy!(stor, dloglik(beta))
end

res1 = optimize(loglik, g!, zeros(5), BFGS())
res2 = optimize(loglik, g!, zeros(5), Csminwel())
res3 = optimize(loglik, zeros(5), Csminwel())
res4 = optimize(Optim.OnceDifferentiable(loglik, zeros(5), autodiff = :forward), zeros(5), Csminwel())
res5 = optimize(Optim.OnceDifferentiable(loglik, zeros(5), autodiff = :finite), zeros(5), Csminwel())

@test Optim.minimum(res1) ≈ Optim.minimum(res2) atol=1.0e-9
@test Optim.minimum(res2) ≈ Optim.minimum(res3) atol=1.0e-9
@test Optim.minimum(res3) ≈ Optim.minimum(res4) atol=1.0e-9
@test Optim.minimum(res4) ≈ Optim.minimum(res5) atol=1.0e-9

res = optimize(loglik, g!, zeros(5), Csminwel(), Optim.Options(extended_trace = true, show_trace = true, store_trace = true))

@test isa(res.trace, Array{Optim.OptimizationState{Float64,CsminWel.Csminwel},1})


using GlobalSensitivity, Test

f = X -> X[1] + 2 * X[2] + 3 * X[3] + 4 * X[7] * X[12]
res1 = gsa(f,FractionalFactorial(),num_params = 12,N=10)

@test res1[1][1:3] == [1.0, 2.0, 3.0]
@test all(res1[1][4:end] .== 0.0)


using LinearAlgebra: I, qr

push!(LOAD_PATH, @__DIR__)
using MultiprecisionFloats
using RKTK2
using DZMisc

setprecision(5000)
const AccurateReal = BigFloat

function log(args...)
    println(args...)
    flush(stdout)
end

function constraint_residual(x)
    [x[1]^2 + x[2]^2 + x[3]^2 - one(AccurateReal), x[1] + x[2]^2]
end

function constraint_jacobian(x)
    [dbl(x[1]) dbl(x[2]) dbl(x[3]);
     one(AccurateReal) dbl(x[2]) zero(AccurateReal)]
end

function constraint_objective(x)
    norm2(constraint_residual(x))
end

function objective(x)
    x[1] * x[3]
end

function gradient(x)
    [x[3], zero(AccurateReal), x[1]]
end

function constrained_gradient(x)
    constraint_jac = copy(transpose(constraint_jacobian(x)))
    orthonormalize_columns!(constraint_jac)
    grad = gradient(x)
    grad - constraint_jac * (transpose(constraint_jac) * grad)
end

function constrain_step(x, step)
    constraint_jac = copy(transpose(constraint_jacobian(x)))
    orthonormalize_columns!(constraint_jac)
    step - constraint_jac * (transpose(constraint_jac) * step)
end

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

function constrain(x)
    x_old, obj_old = x, constraint_objective(x)
    while true
        direction = qr(constraint_jacobian(x_old)) \ constraint_residual(x_old)
        x_new = x_old - direction
        obj_new = constraint_objective(x_new)
        if obj_new < obj_old
            x_old, obj_old = x_new, obj_new
        else
            break
        end
    end
    x_old, obj_old
end

function constrained_step_value(step_size, x, step_direction, step_norm)
    x_new, obj_new = constrain(x - (step_size / step_norm) * step_direction)
    if obj_new < AccurateReal(1e-75)
        objective(x_new)
    else
        AccurateReal(NaN)
    end
end

struct ConstrainedBFGSOptimizer
    objective_value::Ref{AccurateReal}
    last_step_size::Ref{AccurateReal}
end

const INPUT_POINT = AccurateReal[-0.5, 0.707107, -0.5]
const NUM_VARS = 3

const opt = ConstrainedBFGSOptimizer(
    objective(INPUT_POINT),
    AccurateReal(0.00001))

function main()

    x = copy(INPUT_POINT)
    inv_hess = Matrix{AccurateReal}(I, NUM_VARS, NUM_VARS)

    cons_grad = constrained_gradient(x)
    cons_grad_norm = norm(cons_grad)
    while true

        # log("Performing gradient descent step...")
        grad_step_size, obj_grad = quadratic_line_search(constrained_step_value,
            opt.objective_value[], opt.last_step_size[], x, cons_grad, cons_grad_norm)

        # log("Performing BFGS step...")
        bfgs_step = constrain_step(x, inv_hess * cons_grad)
        bfgs_step_norm = norm(bfgs_step)
        bfgs_step_size, obj_bfgs = quadratic_line_search(constrained_step_value,
            opt.objective_value[], opt.last_step_size[], x, bfgs_step, bfgs_step_norm)

        if obj_bfgs < opt.objective_value[] && obj_bfgs <= obj_grad
            x, _ = constrain(x - (bfgs_step_size / bfgs_step_norm) * bfgs_step)
            cons_grad_new = constrained_gradient(x)
            cons_grad_norm_new = norm(cons_grad_new)
            update_inverse_hessian!(inv_hess,
                -bfgs_step_size / bfgs_step_norm,
                bfgs_step,
                cons_grad_new - cons_grad,
                Vector{AccurateReal}(undef, NUM_VARS))
                opt.last_step_size[] = bfgs_step_size
                opt.objective_value[] = obj_bfgs
            cons_grad = cons_grad_new
            cons_grad_norm = cons_grad_norm_new
            log(Float64(opt.objective_value[]), " | ",
                Float64(cons_grad_norm), " | ",
                Float64(opt.last_step_size[]), " | BFGS")
        elseif obj_grad < opt.objective_value[]
            x, _ = constrain(x - (grad_step_size / cons_grad_norm) * cons_grad)
            inv_hess = Matrix{AccurateReal}(I, NUM_VARS, NUM_VARS)
            opt.last_step_size[] = grad_step_size
            opt.objective_value[] = obj_grad
            cons_grad = constrained_gradient(x)
            cons_grad_norm = norm(cons_grad)
            log(Float64(opt.objective_value[]), " | ",
                Float64(cons_grad_norm), " | ",
                Float64(opt.last_step_size[]), " | GRAD")
        else
            log("Neither BFGS step (", Float64(obj_bfgs),
                ") or gradient step (", Float64(obj_grad),
                ") are better than previous point (", Float64(opt.objective_value[]), ")")
            break
        end
    end
end

main()


const MOD = 20201227

function root(value)
    for i in 1:10^9
        if powermod(7, i, MOD) == value
            return i
        end
    end
end

# println(powermod(5764801, root(17807724), MOD) == 14897079) # test
println("P1: ", powermod(19774466, root(7290641), MOD))
println("P2: resort deposit")


export MvnModel, mvn_loss_4obj, mvn_loss_1obj

function mvn_logpx(θ::Vector, data::Matrix{Float64})
    n_vars = size(θ, 1) ÷ 2
    μ = θ[1:n_vars]
    Σ = Diagonal(θ[n_vars+1:end] .^ 2)
    dist = MvNormal(μ, Σ)
    logpx = logpdf(dist, data)
    logpx
end

function mvn_loss_4obj(data::Matrix{Float64})
    function p_loss(θ::Vector)
        loss_4obj(mvn_logpx(θ, data))
    end
    p_loss
end

function mvn_loss_1obj(data::Matrix{Float64})
    function p_loss(θ::Vector)
        loss_1obj(mvn_logpx(θ, data))
    end
    p_loss
end

MLJBase.@mlj_model mutable struct MvnModel <: PDFEstimator
    adtype::SciMLBase.AbstractADType = GalacticOptim.AutoZygote()

    optms::Vector{OptM} = long_optms

    loss::Function = mvn_loss_1obj
end

function MLJBase.fit(model::MvnModel, verbosity, X)
    x = collect(MLJBase.matrix(X)')

    n_vars = size(x, 1)
    p = rand(n_vars * 2)
    lss_f = model.loss(x)
    res = optimizeit(model, lss_f, p)

    μ = res.u[1:n_vars]
    Σ = Diagonal(res.u[n_vars+1:end] .^ 2)
    learned_dist = MvNormal(μ, Σ)

    fitresult = (learned_dist, res, n_vars)
    cache = nothing
    report = nothing
    fitresult, cache, report
end

function MLJBase.transform(model::MvnModel, fitresult, Xnew)
    xnew = collect(MLJBase.matrix(Xnew)')
    learned_dist, res, n_vars = fitresult

    ynew = pdf(learned_dist, xnew)
    ynew = reshape(ynew, size(ynew, 1), 1)
    ynew = DataFrame(ynew, :auto)
end

function MLJBase.fitted_params(model::MvnModel, fitresult)
    learned_dist, res, n_vars = fitresult

    (
        learned_params=res.u,
        learned_dist=learned_dist,
        res=res,
        n_vars=n_vars,
    )
end


@doc doc"""
    SpecialEuclidean(n)

Special Euclidean group $\mathrm{SE}(n)$, the group of rigid motions.

``\mathrm{SE}(n)`` is the semidirect product of the [`TranslationGroup`](@ref) on $ℝ^n$ and
[`SpecialOrthogonal(n)`](@ref)

````math
\mathrm{SE}(n) ≐ \mathrm{T}(n) ⋊_θ \mathrm{SO}(n),
````

where $θ$ is the canonical action of $\mathrm{SO}(n)$ on $\mathrm{T}(n)$ by vector rotation.

This constructor is equivalent to calling

```julia
Tn = TranslationGroup(n)
SOn = SpecialOrthogonal(n)
SemidirectProductGroup(Tn, SOn, RotationAction(Tn, SOn))
```

Points on $\mathrm{SE}(n)$ may be represented as points on the underlying product manifold
$\mathrm{T}(n) \times \mathrm{SO}(n)$ or as affine matrices with size `(n + 1, n + 1)`.
"""
const SpecialEuclidean{N} = SemidirectProductGroup{
    TranslationGroup{Tuple{N},ℝ},
    SpecialOrthogonal{N},
    RotationAction{TranslationGroup{Tuple{N},ℝ},SpecialOrthogonal{N},LeftAction},
}

function SpecialEuclidean(n)
    Tn = TranslationGroup(n)
    SOn = SpecialOrthogonal(n)
    A = RotationAction(Tn, SOn)
    return SemidirectProductGroup(Tn, SOn, A)
end

show(io::IO, ::SpecialEuclidean{n}) where {n} = print(io, "SpecialEuclidean($(n))")

Base.@propagate_inbounds function submanifold_component(
    ::SpecialEuclidean{n},
    x::AbstractMatrix,
    ::Val{1},
) where {n}
    return view(x, 1:n, n + 1)
end
Base.@propagate_inbounds function submanifold_component(
    ::SpecialEuclidean{n},
    x::AbstractMatrix,
    ::Val{2},
) where {n}
    return view(x, 1:n, 1:n)
end

function submanifold_components(G::SpecialEuclidean{n}, x::AbstractMatrix) where {n}
    @assert size(x) == (n + 1, n + 1)
    @inbounds t = submanifold_component(G, x, Val(1))
    @inbounds R = submanifold_component(G, x, Val(2))
    return (t, R)
end

Base.@propagate_inbounds function _padpoint!(
    ::SpecialEuclidean{n},
    y::AbstractMatrix,
) where {n}
    for i ∈ 1:n
        y[n+1, i] = 0
    end
    y[n+1, n+1] = 1
    return y
end

Base.@propagate_inbounds function _padvector!(
    ::SpecialEuclidean{n},
    v::AbstractMatrix,
) where {n}
    for i ∈ 1:n+1
        v[n+1, i] = 0
    end
    return v
end

compose(::SpecialEuclidean, x::AbstractMatrix, y::AbstractMatrix) = x * y

function compose!(
    ::SpecialEuclidean,
    z::AbstractMatrix,
    x::AbstractMatrix,
    y::AbstractMatrix,
)
    return mul!(z, x, y)
end


function approximate_Floquet_multipliers(filename::AbstractString)
    x = load(filename)["x"]
    return _approximate_Floquet_multipliers(x)
end

function _approximate_Floquet_multipliers(x)
    x_₆₄ = complex.(Float64.(mid.(real.(x))), Float64.(mid.(imag.(x))))
    return eigvals(monodromy(x_₆₄); sortby = abs2)
end

function monodromy(p)
    M = zeros(5, 5)
    tspan = (0.0, 2π)
    for i ∈ 1:5
        u0 = I[1:5,i]
        prob = ODEProblem(vector_field!, u0, tspan, p)
        sol = solve(prob, Tsit5())
        M[:,i] = sol(2π)
    end
    return M
end

function vector_field!(dx, x, p, t)
    u = real.(coefficients(component(p, 1)(t)))
    γ = real(component(p, 2)[1])
    sqrt3 = sqrt(3.0)
    dx[1] = 2γ*(-(4u[3] - u[4] - u[5]) * x[3] + (u[3] + 2(u[4] - u[5])) * x[4] + (u[3] - 2(u[4] - u[5])) * x[5])
    dx[2] = 2γ*sqrt3*((u[5] - u[4]) * x[3] + (2u[4] - u[3]) * x[4] + (u[3] - 2u[5]) * x[5])
    dx[3] = 4γ*(u[3] * x[1] + u[1] * x[3])
    dx[4] = -2γ*(u[4] * x[1] + sqrt3 * u[4] * x[2] + (u[1] + sqrt3 * u[2]) * x[4])
    dx[5] = -2γ*(u[5] * x[1] - sqrt3 * u[5] * x[2] + (u[1] - sqrt3 * u[2]) * x[5])
end


sigmoid(x::Real) = inv(one(x) + exp(-x))
binary_step(x::Real) = x < 0 ? 0 : 1

ActivationType = Union{typeof(sigmoid), typeof(binary_step)}


using LinearAlgebra
using PyCall
@pyimport ehtim as eh

function amp_debias(amp, σ; force_nonzero=false)
    amp = abs(amp)
    σ = abs(σ)
    if amp < σ
        if force_nonzero
            deb² = σ^2
        else
            deb² = 0
        end
    else
        deb² = amp^2 - σ^2
    end
    return √deb²
end

function logcamp_debias(log_camp, snr1, snr2, snr3, snr4)
    log_camp_debias = log_camp + 0.5*(1/snr1^2 + 1/snr2^2 - 1/snr3^2 -1/snr4^2)
    return log_camp_debias
end

"""
    Stealing this from eht-imaging, as long as
    we use the quadrangles from ehtim then
    the definitions of numerators and denominators
    will work out.
"""
function make_log_closure_amplitude(n1amp, n2amp, d1amp, d2amp, n1err, n2err, d1err, d2err; debias=true)
    p1 = abs(n1amp)
    p2 = abs(n2amp)
    p3 = abs(d1amp)
    p4 = abs(d2amp)

    if debias
        p1 = amp_debias(p1, n1err)
        p2 = amp_debias(p2, n2err)
        p3 = amp_debias(p3, d1err)
        p4 = amp_debias(p4, d2err)
    end
    snr1 = p1/n1err
    snr2 = p2/n2err
    snr3 = p3/d1err
    snr4 = p4/d2err

    logcamp = log(p1)+log(p2)-log(p3)-log(p4)
    logcamp_err = √(1/snr1^2 + 1/snr2^2 + 1/snr3^2 + 1/snr4^2)
    if debias
        logcamp = logcamp_debias(logcamp, snr1, snr2, snr3, snr4)
    end
    return logcamp, logcamp_err
end

function amp_add_syserr(amp, amp_error; fractional=0, additive=0)
    σ = √(amp_error^2+(fractional*amp)^2+additive^2)
    return amp, σ
end

function vis_add_syserr(vis, amp_error; fractional=0, additive=0)
    return amp_add_syserr(abs(vis), amp_error, fractional=fractional, additive=additive)
end

function logcamp_add_syserr(n1amp, n2amp, d1amp, d2amp, n1err, n2err, d1err, d2err; fractional=0, additive = 0, debias=true)
    n1amp, n1err = amp_add_syserr(n1amp, n1err, fractional=fractional, additive=additive)
    n2amp, n2err = amp_add_syserr(n2amp, n2err, fractional=fractional, additive=additive)
    d1amp, d1err = amp_add_syserr(d1amp, d1err, fractional=fractional, additive=additive)
    d2amp, d2err = amp_add_syserr(d2amp, d2err, fractional=fractional, additive=additive)
    return make_log_closure_amplitude(n1amp, n2amp, d1amp, d2amp, n1err, n2err, d1err, d2err, debias=debias)
end

function make_bispectrum(v1, v2, v3, v1err, v2err, v3err)
    bi = v1*v2*v3
    bisig = abs(bi)*√((v1err/abs(v1))^2+(v2err/abs(v2))^2+(v3err/abs(v3))^2)
    return bi, bisig
end

function closure_phase_from_bispectrum(bi, bisig)
    cphase = angle(bi)
    cphase_error = bisig / abs(bi)
    return cphase, cphase_error
end

function bispectrum_add_syserr(v1, v2, v3, v1err, v2err, v3err; fractional=0, additive=0)
    v1, v1err =vis_add_syserr(v1, v1err, fractional=fractional, additive=additive)
    v2, v2err =vis_add_syserr(v2, v2err, fractional=fractional, additive=additive)
    v3, v3err =vis_add_syserr(v3, v3err, fractional=fractional, additive=additive)
    return make_bispectrum(v1, v2, v3, v1err, v2err, v3err)
end

function cphase_add_syserr(v1, v2, v3, v1err, v2err, v3err; fractional=0, additive=0)
    bi, bisig = bispectrum_add_syserr(v1, v2, v3, v1err, v2err, v3err, fractional=fractional, additive=additive)
    return closure_phase_from_bispectrum(bi, bisig)
end


"""
complex bulk modulus K(ω)
then Q is given by Re(K(ω))/Im(K(ω))
* `nsls` : number of standard linear solids
*
"""
function get_relaxation_times(Q, nsls, fmin, fmax)
	tau_epsilon=zeros(nsls, size(Q)...)
	tau_sigma=zeros(nsls, size(Q)...)

	nom=100
	om=range(fmin, stop=fmax, length=nom);

	# unifrmly distribute tau_sigma over the pass band (eq. 27)
	τ_σ0=inv.(exp.(range(log(fmin), stop=log(fmax), length=nsls)))

	for ix in 1:size(Q,2)
		for iz in 1:size(Q,1)
			Q1=Q[iz,ix]

			for i in 1:nsls
				tau_sigma[i, iz, ix] = τ_σ0[i]
			end

			# see eq. 24
			A=hcat([om .* inv.(1. .+ om.*om .* tau_sigma[i,iz,ix].*tau_sigma[i,iz,ix]) for i in 1:nsls]...)

			# our objective is to optimize Q such that it is close to Q_0, for all ω
			τ_ϵ0 = A\fill(inv(Q1),nom)

			for i in 1:nsls
				tau_epsilon[i,iz,ix] = tau_sigma[i,iz,ix] .+ τ_ϵ0[i]
			end
		end
    end
    return tau_sigma, tau_epsilon
end

"""
Compute that factor that converts unrelaxed bulk modulus to relaxed bulk modulus.
"""
function relaxed_conversion_factor(tau_sigma, tau_epsilon)
	@assert size(tau_epsilon) == size(tau_sigma)
	nsls, nz, nx = size(tau_sigma)
	K_relaxation_factor=zeros(nz,nx)
	for ix in 1:nx
		for iz in 1:nz
			stau=0.0
			for i in 1:nsls
				stau  += (tau_epsilon[i,iz,ix]*inv(tau_sigma[i,iz,ix]) - 1)
			end
			K_relaxation_factor[iz,ix] = 1.0 - stau
		end
	end
	return K_relaxation_factor
end

# return complex bulk modulus, for a frequency domain solver (currently implemented for homogeneous media)
function complexK(Ku, freqs, tau_sigma, tau_epsilon)
    fac=relaxed_conversion_factor(tau_sigma, tau_epsilon)[1]
    nsls=size(tau_sigma,1)

    Kc=complex.(zeros(length(freqs)), zeros(length(freqs)))
    for (ifreq,freq) in enumerate(freqs)
        s=complex(0.0,0.0)
        for isls in 1:nsls
            s+=(1.0 + im*freq*tau_epsilon[isls,1,1])*inv(1.0 + im*freq*tau_sigma[isls,1,1])
        end
		Kc[ifreq]=(1.0 + inv(nsls)*(-nsls + s))*Ku#*inv(fac)
		s=0.0
		for isls in 1:nsls
			s+=tau_epsilon[isls,1,1]*inv(tau_sigma[isls,1,1])
		end
		Kc[ifreq] *= inv(1 + inv(nsls)*(-nsls + s))
	end
	return Kc

end


#============================================================================================
			Time Evolve
============================================================================================#

export euler, Solution, solveAdvection

#export euler
function euler(rhs::Function, dt::T,
               s0::AdvectionState{T})::AdvectionState{T} where {T}
    r0 = rhs(s0)
    s1 = s0 + dt * r0
    AdvectionState{T}(s0.time + dt, s1.u)
end

struct Solution{T}
    dx::T
    dt::T
    states::Vector{AdvectionState{T}}
end

function solveAdvection(tmax::T, nlines::Int, lambda::T)::Solution{T} where {T}
    dx = 1 / nlines
    dt = lambda * dx
    nsteps = round(Int, tmax / dt)
    sol = Solution{T}(dx, dt, AdvectionState{T}[])
    s = initialGaussian(nlines)
    push!(sol.states, s)
    for step in 1:nsteps
        s = euler(rhsAdvection, dt, s)
        push!(sol.states, s)
    end
    sol
end


# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using SparseArrays, Base.Threads, LinearAlgebra, JuMP

"""
    _injection_shift_factors(; buses, lines)

Returns a (B-1)xL matrix M, where B is the number of buses and L is the number
of transmission lines. For a given bus b and transmission line l, the entry
M[l.offset, b.offset] indicates the amount of power (in MW) that flows through
transmission line l when 1 MW of power is injected at the slack bus (the bus
that has offset zero) and withdrawn from b.
"""
function _injection_shift_factors(;
    buses::Array{Bus},
    lines::Array{TransmissionLine},
)
    susceptance = _susceptance_matrix(lines)
    incidence = _reduced_incidence_matrix(lines = lines, buses = buses)
    laplacian = transpose(incidence) * susceptance * incidence
    isf = susceptance * incidence * inv(Array(laplacian))
    return isf
end

"""
    _reduced_incidence_matrix(; buses::Array{Bus}, lines::Array{TransmissionLine})

Returns the incidence matrix for the network, with the column corresponding to
the slack bus is removed. More precisely, returns a (B-1) x L matrix, where B
is the number of buses and L is the number of lines. For each row, there is a 1
element and a -1 element, indicating the source and target buses, respectively,
for that line.
"""
function _reduced_incidence_matrix(;
    buses::Array{Bus},
    lines::Array{TransmissionLine},
)
    matrix = spzeros(Float64, length(lines), length(buses) - 1)
    for line in lines
        if line.source.offset > 0
            matrix[line.offset, line.source.offset] = 1
        end
        if line.target.offset > 0
            matrix[line.offset, line.target.offset] = -1
        end
    end
    return matrix
end

"""
    _susceptance_matrix(lines::Array{TransmissionLine})

Returns a LxL diagonal matrix, where each diagonal entry is the susceptance of
the corresponding transmission line.
"""
function _susceptance_matrix(lines::Array{TransmissionLine})
    return Diagonal([l.susceptance for l in lines])
end

"""

    _line_outage_factors(; buses, lines, isf)

Returns a LxL matrix containing the Line Outage Distribution Factors (LODFs)
for the given network. This matrix how does the pre-contingency flow change
when each individual transmission line is removed.
"""
function _line_outage_factors(;
    buses::Array{Bus,1},
    lines::Array{TransmissionLine,1},
    isf::Array{Float64,2},
)::Array{Float64,2}
    incidence = Array(_reduced_incidence_matrix(lines = lines, buses = buses))
    lodf::Array{Float64,2} = isf * transpose(incidence)
    _, n = size(lodf)
    for i in 1:n
        lodf[:, i] *= 1.0 / (1.0 - lodf[i, i])
        lodf[i, i] = -1
    end
    return lodf
end