# random functions that dispatch either to CURAND or GPUArrays' generic RNG

using Random

export rand_logn!, rand_poisson!

# the interface is split in two levels:
# - functions that extend the Random standard library, and take an RNG as first argument,
#   will only ever dispatch to CURAND and as a result are limited in the types they support.
# - functions that take an array will dispatch to either CURAND or GPUArrays
# - non-exported functions are provided for constructing GPU arrays from only an eltype

const curand_rng = CURAND.default_rng
gpuarrays_rng() = GPUArrays.default_rng(CuArray)

function seed!(seed=Base.rand(UInt64))
    Random.seed!(curand_rng(), seed)
    Random.seed!(gpuarrays_rng(), seed)
end

# CURAND in-place
Random.rand!(A::CURAND.UniformArray) = Random.rand!(curand_rng(), A)
Random.randn!(A::CURAND.NormalArray; kwargs...) = Random.randn!(curand_rng(), A; kwargs...)
rand_logn!(A::CURAND.LognormalArray; kwargs...) = CURAND.rand_logn!(curand_rng(), A; kwargs...)
rand_poisson!(A::CURAND.PoissonArray; kwargs...) = CURAND.rand_poisson!(curand_rng(), A; kwargs...)

# CURAND out-of-place
rand(T::CURAND.UniformType, dims::Dims) = Random.rand(curand_rng(), T, dims)
randn(T::CURAND.NormalType, dims::Dims; kwargs...) = Random.randn(curand_rng(), T, dims; kwargs...)
rand_logn(T::CURAND.LognormalType, dims::Dims; kwargs...) = CURAND.rand_logn(curand_rng(), T, dims; kwargs...)
rand_poisson(T::CURAND.PoissonType, dims::Dims; kwargs...) = CURAND.rand_poisson(curand_rng(), T, dims; kwargs...)

# support all dimension specifications
rand(T::CURAND.UniformType, dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), T, Dims((dim1, dims...)))
randn(T::CURAND.NormalType, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_logn(T::CURAND.LognormalType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_poisson(T::CURAND.PoissonType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), T, Dims((dim1, dims...)); kwargs...)

# GPUArrays in-place
Random.rand!(A::CuArray) = Random.rand!(gpuarrays_rng(), A)
Random.randn!(A::CuArray; kwargs...) =
    error("CUDA.jl does not support generating normally-distributed random numbers of type $(eltype(A))")
# FIXME: GPUArrays.jl has a randn! nowadays, but it doesn't work with e.g. Cuint
rand_logn!(A::CuArray; kwargs...) =
    error("CUDA.jl does not support generating lognormally-distributed random numbers of type $(eltype(A))")
rand_poisson!(A::CuArray; kwargs...) =
    error("CUDA.jl does not support generating Poisson-distributed random numbers of type $(eltype(A))")

# GPUArrays out-of-place
rand(T::Type, dims::Dims) = Random.rand!(CuArray{T}(undef, dims...))
randn(T::Type, dims::Dims; kwargs...) = Random.randn!(CuArray{T}(undef, dims...); kwargs...)
rand_logn(T::Type, dims::Dims; kwargs...) = rand_logn!(CuArray{T}(undef, dims...); kwargs...)
rand_poisson(T::Type, dims::Dims; kwargs...) = rand_poisson!(CuArray{T}(undef, dims...); kwargs...)

# support all dimension specifications
rand(T::Type, dim1::Integer, dims::Integer...) =
    Random.rand!(CuArray{T}(undef, dim1, dims...))
randn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_logn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_logn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_poisson(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_poisson!(CuArray{T}(undef, dim1, dims...); kwargs...)

# untyped out-of-place
rand(dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), Dims((dim1, dims...)))
randn(dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_logn(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_poisson(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), Dims((dim1, dims...)); kwargs...)


# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

"""
Contains all complex data used for testing and benchmarking.
"""
module TestData

using OpticSim
using OpticSim: tobeziersegments # try to use only exported functions so this list should stay short
using OpticSim.Geometry
using OpticSim.Emitters
using OpticSim.GlassCat

using StaticArrays
using DataFrames: DataFrame
using Unitful
using Images
using Base: @.
using LinearAlgebra

#define glasses locally rather than relying on GlassCat. These tests used to rely on downloaded glasses which made the tests brittle.
#If a glass didn't download the test would fail.

const Examples_N_BK7 = GlassCat.Glass(GlassCat.GlassID(GlassCat.AGF, 885), 2, 1.03961212, 0.00600069867, 0.231792344, 0.0200179144, 1.01046945, 103.560653, 0.0, 0.0, NaN, NaN, 0.3, 2.5, 1.86e-6, 1.31e-8, -1.37e-11, 4.34e-7, 6.27e-10, 0.17, 20.0, -0.0009, 2.3, 1.0, 7.1, 1.0, 1, 1.0, [(0.3, 0.05, 25.0), (0.31, 0.25, 25.0), (0.32, 0.52, 25.0), (0.334, 0.78, 25.0), (0.35, 0.92, 25.0), (0.365, 0.971, 25.0), (0.37, 0.977, 25.0), (0.38, 0.983, 25.0), (0.39, 0.989, 25.0), (0.4, 0.992, 25.0), (0.405, 0.993, 25.0), (0.42, 0.993, 25.0), (0.436, 0.992, 25.0), (0.46, 0.993, 25.0), (0.5, 0.994, 25.0), (0.546, 0.996, 25.0), (0.58, 0.995, 25.0), (0.62, 0.994, 25.0), (0.66, 0.994, 25.0), (0.7, 0.996, 25.0), (1.06, 0.997, 25.0), (1.53, 0.98, 25.0), (1.97, 0.84, 25.0), (2.325, 0.56, 25.0), (2.5, 0.36, 25.0)], 1.5168, 2.3, 0.0, 0, 64.17, 0, 2.51, 0.0)

const Examples_N_SK16 =GlassCat.Glass(GlassCat.GlassID(GlassCat.AGF, 828), 2, 1.34317774, 0.00704687339, 0.241144399, 0.0229005, 0.994317969, 92.7508526, 0.0, 0.0, NaN, NaN, 0.31, 2.5, -2.37e-8, 1.32e-8, -1.29e-11, 4.09e-7, 5.17e-10, 0.17, 20.0, -0.0011, 3.2, 1.4, 6.3, 4.0, 1, 53.3, [(0.31, 0.02, 25.0), (0.32, 0.11, 25.0), (0.334, 0.4, 25.0), (0.35, 0.7, 25.0), (0.365, 0.86, 25.0), (0.37, 0.89, 25.0), (0.38, 0.93, 25.0), (0.39, 0.956, 25.0), (0.4, 0.97, 25.0), (0.405, 0.974, 25.0), (0.42, 0.979, 25.0), (0.436, 0.981, 25.0), (0.46, 0.984, 25.0), (0.5, 0.991, 25.0), (0.546, 0.994, 25.0), (0.58, 0.994, 25.0), (0.62, 0.993, 25.0), (0.66, 0.994, 25.0), (0.7, 0.996, 25.0), (1.06, 0.995, 25.0), (1.53, 0.973, 25.0), (1.97, 0.88, 25.0), (2.325, 0.54, 25.0), (2.5, 0.26, 25.0)], 1.62041, 3.3, 4.0, 0, 60.32, 0, 3.58, 0.0)

const Examples_N_SF2 = GlassCat.Glass(GlassCat.GlassID(GlassCat.AGF, 815), 2, 1.47343127, 0.0109019098, 0.163681849, 0.0585683687, 1.36920899, 127.404933, 0.0, 0.0, NaN, NaN, 0.365, 2.5, 3.1e-6, 1.75e-8, 6.62e-11, 7.51e-7, 8.99e-10, 0.277, 20.0, 0.0081, 1.0, 1.4, 6.68, 1.0, 1, 1.0, [(0.365, 0.007, 25.0), (0.37, 0.06, 25.0), (0.38, 0.4, 25.0), (0.39, 0.68, 25.0), (0.4, 0.83, 25.0), (0.405, 0.865, 25.0), (0.42, 0.926, 25.0), (0.436, 0.949, 25.0), (0.46, 0.961, 25.0), (0.5, 0.975, 25.0), (0.546, 0.986, 25.0), (0.58, 0.987, 25.0), (0.62, 0.984, 25.0), (0.66, 0.984, 25.0), (0.7, 0.987, 25.0), (1.06, 0.997, 25.0), (1.53, 0.984, 25.0), (1.97, 0.93, 25.0), (2.325, 0.76, 25.0), (2.5, 0.67, 25.0)], 1.64769, 1.2, 0.0, 0, 33.82, 0, 2.718, 0.0)

const Examples_N_SF14 = GlassCat.Glass(GlassCat.GlassID(GlassCat.AGF, 896), 2, 1.69022361, 0.0130512113, 0.288870052, 0.061369188, 1.7045187, 149.517689, 0.0, 0.0, NaN, NaN, 0.365, 2.5, -5.56e-6, 7.09e-9, -1.09e-11, 9.85e-7, 1.39e-9, 0.287, 20.0, 0.013, 1.0, 2.2, 9.41, 1.0, 1, 1.0, [(0.365, 0.004, 25.0), (0.37, 0.04, 25.0), (0.38, 0.33, 25.0), (0.39, 0.61, 25.0), (0.4, 0.75, 25.0), (0.405, 0.79, 25.0), (0.42, 0.87, 25.0), (0.436, 0.91, 25.0), (0.46, 0.938, 25.0), (0.5, 0.964, 25.0), (0.546, 0.983, 25.0), (0.58, 0.987, 25.0), (0.62, 0.987, 25.0), (0.66, 0.987, 25.0), (0.7, 0.985, 25.0), (1.06, 0.998, 25.0), (1.53, 0.98, 25.0), (1.97, 0.88, 25.0), (2.325, 0.64, 25.0), (2.5, 0.57, 25.0)], 1.76182, 1.0, 0.0, 0, 26.53, 0, 3.118, 0.0)

include("curves.jl")
include("surfaces.jl")
include("lenses.jl")
include("systems.jl")
include("other.jl")

end # module TestData
export TestData


# default kernel function of difference types
const KERNEL_TYPE = Dict(ContinuousDim => gaussian, CategoricalDim => wang_ryzin, UnorderedCategoricalDim => aitchison_aitken)

function map_dim(candidates::Vector)
    unordered_to_index, index_to_unordered = Dict(), Dict()
    for (i, elem) in enumerate(candidates)
        unordered_to_index[elem] = i
        index_to_unordered[i] = elem
    end
    (unordered_to_index, index_to_unordered)
end

# Candidate need to be mapped from UnorderedCategoricalDim
function candidate_need_map(dim_type, candidate)
    if dim_type isa UnorderedCategoricalDim
        for _candidate in candidate
            if !(_candidate isa Real)
                return true
            end
        end
    end
    return false
end

# Observation need to be mapped from UnorderedCategoricalDim
function observation_need_map(dim_type, observation)
    if dim_type isa UnorderedCategoricalDim
        for _observation in observation
            if !(_observation isa Real)
                return true
            end
        end
    end
    return false
end

# Lazy evaluation univariate KDE
Base.@kwdef mutable struct KDEUniv
    type::DimensionType
    bandwidth::Real
    data::Vector
    kernel::Function
end
# Using default kernel for dim_type
function KDEUniv(dim_type::DimensionType, bandwidth, observation)
    KDEUniv(dim_type, bandwidth, observation, KERNEL_TYPE[typeof(dim_type)])
end

function pdf(kde::KDEUniv, x::Real; keep_all=true)
    densities = kde.kernel(kde.bandwidth, kde.data, x)
    if keep_all
        return densities
    else
        return mean(densities) / kde.bandwidth
    end
end

# Multivariate KDE based on KDEUniv
Base.@kwdef mutable struct KDEMulti
    # Type of every dimension, continuous or discrete
    dims::Vector
    # KDE from different dimensions
    KDEs::Vector
    # observations: An k*n matrix, where k is dimension of KDEs and n is number of observations
    observations::Vector
    mat_observations::Matrix
    # In UnorderedContinuous case, we assign an index 1:N to every value, where N=|D_i|, then we pass the 1:N to KDE
    # dimension i is number or not
    mapped::Vector
    unordered_to_index::Dict{KDEUniv, Dict{Any, Real}}
    index_to_unordered::Dict{KDEUniv, Dict{Real, Any}}
end

# Constructor without candidates
function KDEMulti(dims::Vector, observations::Vector)
    KDEMulti(dims, nothing, observations)
end
function KDEMulti(dims::Vector, bws::Union{Vector, Nothing}, observations::Vector)
    for (dim_type, observation) in zip(dims, observations)
        if observation_need_map(dim_type, observation)
            error("If there is Uncatogorical dimension and not a number, should specify its candidate value.")
        end
    end
    KDEMulti(dims, bws, observations, nothing)
end

# Constructor with candidates
## Convert tuple candidates to dict candidates
function get_dict_candidates(dim_types::Vector{DimensionType}, candidates::Tuple)
    dict_candidates = Dict{Int, Vector}()
    for (i, dim_type, candidate) in zip(1:length(dim_types), dim_types, candidates)
        if candidate_need_map(dim_type, candidate)
            dict_candidates[i] = candidate
        end
    end
    dict_candidates
end

## Tuple candidates
function KDEMulti(dims::Vector, observations::Vector, candidates::Tuple)
    KDEMulti(dims, nothing, observations, get_dict_candidates(dims, candidates))
end
## Dict candidates
function KDEMulti(dims::Vector, observations::Vector, candidates::Union{Dict{Int, Vector}, Nothing})
    KDEMulti(dims, nothing, observations, candidates)
end
function KDEMulti(dims::Vector, bws::Union{Vector, Nothing}, observations::Vector, candidates::Union{Dict{Int, Vector}, Nothing})
    mat_observations = hcat(observations...)
    KDEMulti(dims, bws, mat_observations, candidates)
end

# Main constructor
function KDEMulti(dims::Vector, bws::Union{Vector, Nothing}, mat_observations::Matrix, candidates::Union{Dict{Int, Vector}, Nothing})
    _KDEs, _observations, _unordered_to_index, _index_to_unordered = Vector(), Vector(), Dict{Int, Dict{Any, Real}}(),
                                                                            Dict{Int, Dict{Real, Any}}()
    mapped = Vector{Bool}()
    for (i, dim_type_i) in zip(1:size(mat_observations)[1], dims)
        _observations_i = mat_observations[i, :]
        if observation_need_map(dim_type_i, _observations_i) && ((candidates isa Nothing) || !haskey(candidates, i))
            error("No corresponding candidate at dimension. ")
        end
        if !(candidates isa Nothing) && haskey(candidates, i) && candidate_need_map(dim_type_i, candidates[i])
            _unordered_to_index_i, _index_to_unordered_i = map_dim(candidates[i])
            _unordered_to_index[i], _index_to_unordered[i] = _unordered_to_index_i, _index_to_unordered_i
            _observations_i = [_unordered_to_index_i[elem] for elem in _observations_i]
            push!(mapped, true)
        else
            push!(mapped, false)
        end
        push!(_observations, _observations_i)
    end
    if bws isa Nothing
        bws = default_bandwidth(_observations)
    end
    for (_observations_i, bandwidth_i, dim_type_i) in zip(_observations, bws, dims)
        # kde_i = KDEUniv(_observations_i, bandwidth_i, dim_type_i, type_kernel[typeof(dim_type_i)], is_number)
        kde_i = KDEUniv(dim_type_i, bandwidth_i, _observations_i)
        push!(_KDEs, kde_i)
    end
    unordered_to_index, index_to_unordered = Dict{KDEUniv, Dict{Any, Real}}(), Dict{KDEUniv, Dict{Real, Any}}()
    for i in 1:size(mat_observations)[1]
        if haskey(_index_to_unordered, i)
            unordered_to_index[_KDEs[i]] = _unordered_to_index[i]
            index_to_unordered[_KDEs[i]] = _index_to_unordered[i]
        end
    end
    KDEMulti(dims, _KDEs, _observations, mat_observations, mapped, unordered_to_index, index_to_unordered)
end

# GPKE and kernel code refers to Nonparametric part of ['statsmodels'](https://github.com/statsmodels/statsmodels)
#
# Copyright of statsmodels:
#
#     Copyright (C) 2006, Jonathan E. Taylor
#     All rights reserved.
#
#     Copyright (c) 2006-2008 Scipy Developers.
#     All rights reserved.
#
#     Copyright (c) 2009-2018 statsmodels Developers.
#     All rights reserved.

# pdf of KDEMulti, using (unnormalized) GPKE(Generalized Product Kernel Estimator)
function gpke(multi_kde::KDEMulti, x::Vector)
    Kval = Array{Real}(undef, (length(multi_kde.observations[1]), length(multi_kde.observations)))
    for (i, _kde, _x) in zip(1:length(x), multi_kde.KDEs, x)
        if multi_kde.mapped[i]
            _x = multi_kde.unordered_to_index[_kde][_x]
        end
        Kval[:, i] = pdf(_kde, _x)
    end
    iscontinuous = [_dim isa ContinuousDim ? true : false for _dim in multi_kde.dims]
    dens = prod(Kval, dims=2) / prod([kde.bandwidth for kde in multi_kde.KDEs][iscontinuous])
    sum(dens)
end

# Alias of gpke
function pdf(multi_kde::KDEMulti, x::Vector)
    gpke(multi_kde, x)
end

# Scott's normal reference rule of thumb bandwidth parameter
function default_bandwidth(observations::Vector)
    X = std.(observations)
    1.06 * X * length(observations).^(-1 / (4 + length(observations[1])))
end


#! format: off
struct RangeReserve <: AbstractReservesFormulation end
struct StepwiseCostReserve <: AbstractReservesFormulation end
############################### Reserve Variables #########################################

get_variable_sign(_, ::Type{<:PSY.Reserve}, ::AbstractReservesFormulation) = NaN
############################### ActiveServiceVariable, Reserve #########################################

get_variable_binary(::ActiveServiceVariable, ::Type{<:PSY.Reserve}, ::AbstractReservesFormulation) = false
get_variable_upper_bound(::ActiveServiceVariable, ::PSY.Reserve, ::PSY.Component, _) = nothing
get_variable_lower_bound(::ActiveServiceVariable, ::PSY.Reserve, ::PSY.Component, _) = 0.0

############################### ServiceRequirementVariable, ReserveDemandCurve ################################

get_variable_binary(::ServiceRequirementVariable, ::Type{<:PSY.ReserveDemandCurve}, ::AbstractReservesFormulation) = false
get_variable_upper_bound(::ServiceRequirementVariable, ::PSY.ReserveDemandCurve, ::PSY.Component, ::AbstractReservesFormulation) = nothing
get_variable_lower_bound(::ServiceRequirementVariable, ::PSY.ReserveDemandCurve, ::PSY.Component, ::AbstractReservesFormulation) = 0.0

#! format: on
################################## Reserve Requirement Constraint ##########################
function service_requirement_constraint!(
    optimization_container::OptimizationContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.Reserve, T <: AbstractReservesFormulation}
    parameters = model_has_parameters(optimization_container)
    use_forecast_data = model_uses_forecasts(optimization_container)
    initial_time = model_initial_time(optimization_container)
    @debug initial_time
    time_steps = model_time_steps(optimization_container)
    name = PSY.get_name(service)
    constraint =
        get_constraint(optimization_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(optimization_container, name, SR)
    use_slacks = get_services_slack_variables(optimization_container.settings)

    ts_vector = get_time_series(optimization_container, service, "requirement")

    use_slacks && (slack_vars = reserve_slacks(optimization_container, name))

    requirement = PSY.get_requirement(service)
    if parameters
        container = get_parameter_container(
            optimization_container,
            UpdateRef{SR}(SERVICE_REQUIREMENT, "requirement"),
        )
        param = get_parameter_array(container)
        multiplier = get_multiplier_array(container)
        for t in time_steps
            param[name, t] = add_parameter(optimization_container.JuMPmodel, ts_vector[t])
            multiplier[name, t] = 1.0
            if use_slacks
                resource_expression = sum(reserve_variable[:, t]) + slack_vars[t]
            else
                resource_expression = sum(reserve_variable[:, t])
            end
            mul = (requirement * multiplier[name, t])
            constraint[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                resource_expression >= param[name, t] * mul
            )
        end
    else
        for t in time_steps
            constraint[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                sum(reserve_variable[:, t]) >= ts_vector[t] * requirement
            )
        end
    end
    return
end

function service_requirement_constraint!(
    optimization_container::OptimizationContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.StaticReserve, T <: AbstractReservesFormulation}
    parameters = model_has_parameters(optimization_container)
    initial_time = model_initial_time(optimization_container)
    @debug initial_time
    time_steps = model_time_steps(optimization_container)
    name = PSY.get_name(service)
    constraint =
        get_constraint(optimization_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(optimization_container, name, SR)
    use_slacks = get_services_slack_variables(optimization_container.settings)

    use_slacks && (slack_vars = reserve_slacks(optimization_container, name))

    requirement = PSY.get_requirement(service)
    for t in time_steps
        resource_expression = JuMP.GenericAffExpr{Float64, JuMP.VariableRef}()
        JuMP.add_to_expression!(resource_expression, sum(reserve_variable[:, t]))
        if use_slacks
            resource_expression += slack_vars[t]
        end
        constraint[name, t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            resource_expression >= requirement
        )
    end

    return
end

function cost_function!(
    optimization_container::OptimizationContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.Reserve, T <: AbstractReservesFormulation}
    reserve = get_variable(optimization_container, PSY.get_name(service), SR)
    for r in reserve
        JuMP.add_to_expression!(
            optimization_container.cost_function,
            r,
            DEFAULT_RESERVE_COST,
        )
    end
    return
end

function service_requirement_constraint!(
    optimization_container::OptimizationContainer,
    service::SR,
    ::ServiceModel{SR, StepwiseCostReserve},
) where {SR <: PSY.ReserveDemandCurve}
    initial_time = model_initial_time(optimization_container)
    @debug initial_time
    time_steps = model_time_steps(optimization_container)
    name = PSY.get_name(service)
    constraint =
        get_constraint(optimization_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(optimization_container, name, SR)
    requirement_variable = get_variable(optimization_container, SERVICE_REQUIREMENT, SR)

    for t in time_steps
        constraint[name, t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            sum(reserve_variable[:, t]) >= requirement_variable[name, t]
        )
    end

    return
end

function AddCostSpec(
    ::Type{T},
    ::Type{StepwiseCostReserve},
    optimization_container::OptimizationContainer,
) where {T <: PSY.Reserve}
    return AddCostSpec(;
        variable_type = ServiceRequirementVariable,
        component_type = T,
        has_status_variable = false,
        has_status_parameter = false,
        variable_cost = PSY.get_variable,
        start_up_cost = nothing,
        shut_down_cost = nothing,
        fixed_cost = nothing,
        sos_status = SOSStatusVariable.NO_VARIABLE,
    )
end

function add_to_cost!(
    optimization_container::OptimizationContainer,
    spec::AddCostSpec,
    service::SR,
    component_name::String,
) where {SR <: PSY.Reserve}
    time_steps = model_time_steps(optimization_container)
    use_forecast_data = model_uses_forecasts(optimization_container)
    if !use_forecast_data
        error("StepwiseCostReserve is only supported with forecast")
    end
    variable_cost_forecast =
        get_time_series(optimization_container, service, "variable_cost")
    variable_cost_forecast = map(PSY.VariableCost, variable_cost_forecast)
    for t in time_steps
        variable_cost!(
            optimization_container,
            spec,
            component_name,
            variable_cost_forecast[t],
            t,
        )
    end
    return
end

function cost_function!(
    optimization_container::OptimizationContainer,
    service::SR,
    model::ServiceModel{SR, StepwiseCostReserve},
) where {SR <: PSY.ReserveDemandCurve}
    spec = AddCostSpec(SR, model.formulation, optimization_container)
    @debug SR, spec
    add_to_cost!(optimization_container, spec, service, PSY.get_name(service))
    return
end

function modify_device_model!(
    devices_template::Dict{Symbol, DeviceModel},
    service_model::ServiceModel{<:PSY.Reserve, <:AbstractReservesFormulation},
    contributing_devices::Vector{<:PSY.Device},
)
    device_types = unique(typeof.(contributing_devices))
    for dt in device_types
        for (device_model_name, device_model) in devices_template
            # add message here when it exists
            get_component_type(device_model) != dt && continue
            service_model in device_model.services && continue
            push!(device_model.services, service_model)
        end
    end

    return
end

function include_service!(
    constraint_info::T,
    services,
    ::ServiceModel{SR, <:AbstractReservesFormulation},
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    SR <: PSY.Reserve{PSY.ReserveUp},
}
    for (ix, service) in enumerate(services)
        push!(
            constraint_info.additional_terms_ub,
            make_constraint_name(PSY.get_name(service), SR),
        )
    end
    return
end

function include_service!(
    constraint_info::T,
    services,
    ::ServiceModel{SR, <:AbstractReservesFormulation},
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    SR <: PSY.Reserve{PSY.ReserveDown},
}
    for (ix, service) in enumerate(services)
        push!(
            constraint_info.additional_terms_lb,
            make_constraint_name(PSY.get_name(service), SR),
        )
    end
    return
end

function add_device_services!(
    constraint_info::T,
    device::D,
    model::DeviceModel,
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    D <: PSY.Device,
}
    for service_model in get_services(model)
        if PSY.has_service(device, service_model.component_type)
            services = (
                s for s in PSY.get_services(device) if isa(s, service_model.component_type)
            )
            @assert !isempty(services)
            include_service!(constraint_info, services, service_model)
        end
    end
    return
end

function add_device_services!(
    constraint_data_in::AbstractRangeConstraintInfo,
    constraint_data_out::AbstractRangeConstraintInfo,
    device::D,
    model::DeviceModel{D, <:AbstractStorageFormulation},
) where {D <: PSY.Storage}
    for service_model in get_services(model)
        if PSY.has_service(device, service_model.component_type)
            services = (
                s for s in PSY.get_services(device) if isa(s, service_model.component_type)
            )
            @assert !isempty(services)
            if service_model.component_type <: PSY.Reserve{PSY.ReserveDown}
                for service in services
                    push!(
                        constraint_data_in.additional_terms_ub,
                        make_constraint_name(
                            PSY.get_name(service),
                            service_model.component_type,
                        ),
                    )
                end
            elseif service_model.component_type <: PSY.Reserve{PSY.ReserveUp}
                for service in services
                    push!(
                        constraint_data_out.additional_terms_ub,
                        make_constraint_name(
                            PSY.get_name(service),
                            service_model.component_type,
                        ),
                    )
                end
            end
        end
    end
    return
end


if VERSION > v"1.6"
    exit() # Don't want to build the system image, since we will use LBT
end

using PackageCompiler
using MKL_jll

# if no environment variable ENV["USE_BLAS64"] is set install.jl
# tries to change USE_BLAS64 = false
const USEBLAS64 = parse(Bool,get(ENV, "USE_BLAS64","false"))

include("../src/install.jl")
enable_mkl_startup()


function plot_LES_3D(frame, FILE_DIR, OUTPUT_PATH, axis_images; title="", colorscheme=:ice, rev=false)
    xy_file = jldopen(joinpath(FILE_DIR, "xy_slice.jld2"))
    xz_file = jldopen(joinpath(FILE_DIR, "xz_slice.jld2"))
    yz_file = jldopen(joinpath(FILE_DIR, "yz_slice.jld2"))
    instantaneous_statistics = jldopen(joinpath(FILE_DIR, "instantaneous_statistics.jld2"))

    frames = keys(xy_file["timeseries/T"])
    iteration = frames[frame]

    color_range_max = maximum(maximum.([xz_file["timeseries/T/$iteration"], yz_file["timeseries/T/$iteration"], xy_file["timeseries/T/$iteration"]]))
    color_range_min = minimum(minimum.([xz_file["timeseries/T/$iteration"], yz_file["timeseries/T/$iteration"], xy_file["timeseries/T/$iteration"]]))
    color_range = (color_range_min, color_range_max)

    xC = xz_file["grid/xC"][4:end-3]
    yC = xz_file["grid/yC"][4:end-3]
    zC = xz_file["grid/zC"][4:end-3]

    colormap = cgrad(colorscheme, scale=:log10, rev=rev)

    fig = Figure(resolution=(1920, 1080))
    ax = fig[1,1] = CairoMakie.Axis3(fig, aspect=(1, 1, 0.5), xlabel="x /m", ylabel="y /m", zlabel="z /m")

    x_xz = fill(xC[1], 128)
    y_xz = yC
    z_xz = zeros(length(x_xz), length(y_xz))

    for i in 1:size(z_xz, 1)
        z_xz[i,:] .= zC[i]
    end

    z_xz
    T_xz = transpose(hcat([xz_file["timeseries/T/$iteration"][:, :, i] for i in 1:length(zC)]... ))

    xz_surface = CairoMakie.surface!(ax, x_xz, y_xz, z_xz, color=T_xz, colormap=colormap)

    x_yz = xC
    y_yz = fill(yC[1], 128)
    z_yz = zeros(length(x_yz), length(y_yz))

    for i in 1:size(z_yz, 2)
        z_yz[:,i] .= zC[i]
    end

    z_yz

    T_yz = similar(z_yz)

    for i in 1:size(T_yz, 2)
        T_yz[:,i] = yz_file["timeseries/T/$iteration"][:, :, i]
    end
    T_yz

    yz_surface = CairoMakie.surface!(ax, x_yz, y_yz, z_yz, color=T_yz, colormap=colormap)

    x_xy = xC
    y_xy = yC
    z_xy = zeros(length(y_xy), length(x_xy))

    T_xy = xy_file["timeseries/T/$iteration"][:,:,1]

    xy_surface = CairoMakie.surface!(ax, x_xy, y_xy, z_xy, color=T_xy, colormap=colormap)

    xz_surface.colorrange = color_range
    yz_surface.colorrange = color_range
    xy_surface.colorrange = color_range

    ax_T_3D = fig[2,1] = CairoMakie.Axis(fig, aspect=DataAspect())

    rel_size = 40
    aspect = 1 / 4

    hidedecorations!(ax_T_3D)
    hidespines!(ax_T_3D)
    image!(ax_T_3D, axis_images.T_3D)
    rowsize!(fig.layout, 2, CairoMakie.Relative(1 / rel_size))

    plots_sublayout = fig[:,2] = GridLayout()

    colsize!(fig.layout, 2, CairoMakie.Relative(aspect))
    rowgap!(fig.layout, Relative(1 / 50 / 3))

    colorbar = CairoMakie.Colorbar(fig[3, 1], xy_surface, vertical=false)

    y_ax_u = plots_sublayout[1,1] = CairoMakie.Axis(fig, aspect=DataAspect())
    y_ax_v = plots_sublayout[3,1] = CairoMakie.Axis(fig, aspect=DataAspect())
    y_ax_T = plots_sublayout[5,1] = CairoMakie.Axis(fig, aspect=DataAspect())

    x_ax_u = plots_sublayout[2,2] = CairoMakie.Axis(fig, aspect=DataAspect())
    x_ax_v = plots_sublayout[4,2] = CairoMakie.Axis(fig, aspect=DataAspect())
    x_ax_T = plots_sublayout[6,2] = CairoMakie.Axis(fig, aspect=DataAspect())

    ax_u = plots_sublayout[1,2] = CairoMakie.Axis(fig)
    ax_v = plots_sublayout[3,2] = CairoMakie.Axis(fig)
    ax_T = plots_sublayout[5,2] = CairoMakie.Axis(fig)

    hidedecorations!(y_ax_u)
    hidedecorations!(y_ax_v)
    hidedecorations!(y_ax_T)

    hidedecorations!(x_ax_u)
    hidedecorations!(x_ax_v)
    hidedecorations!(x_ax_T)

    hidespines!(y_ax_u)
    hidespines!(y_ax_v)
    hidespines!(y_ax_T)

    hidespines!(x_ax_u)
    hidespines!(x_ax_v)
    hidespines!(x_ax_T)

    image!(x_ax_u, axis_images.u)
    image!(x_ax_v, axis_images.v)
    image!(x_ax_T, axis_images.T)

    image!(y_ax_u, axis_images.z)
    image!(y_ax_v, axis_images.z)
    image!(y_ax_T, axis_images.z)

    rowsize!(plots_sublayout, 2, CairoMakie.Relative(1 / rel_size))
    rowsize!(plots_sublayout, 4, CairoMakie.Relative(1 / rel_size))
    rowsize!(plots_sublayout, 6, CairoMakie.Relative(1 / rel_size))
    colsize!(plots_sublayout, 1, CairoMakie.Aspect(2, 1))

    colgap!(plots_sublayout, 1 / 50)
    rowgap!(plots_sublayout, Relative(1 / 50 / 3))

    u_line = CairoMakie.lines!(ax_u, instantaneous_statistics["timeseries/u/$iteration"][:], zC)
    v_line = CairoMakie.lines!(ax_v, instantaneous_statistics["timeseries/v/$iteration"][:], zC)
    T_line = CairoMakie.lines!(ax_T, instantaneous_statistics["timeseries/T/$iteration"][:], zC)

    if title !== ""
        supertitle = fig[0, :] = Label(fig, title, textsize=25)
    end

    trim!(fig.layout)

    close(xy_file)
    close(xz_file)
    close(yz_file)

    save(OUTPUT_PATH, fig, px_per_unit=4, pt_per_unit=2)
end


# Here, we illustrate a pure ModernGL implementation of some polygon drawing
using ModernGL, GeometryTypes, GLFW

# Create the window. This sets all the hints and makes the context current.
window = GLFW.Window(name="Drawing polygons 5", resolution=(800,600))

# The shaders. Here we do everything manually, but life will get
# easier with GLAbstraction. See drawing_polygons5.jl for such an
# implementation.

# The vertex shader
vertex_source = Vector{UInt8}("""
#version 150

in vec2 position;
in vec3 color;

out vec3 Color;

void main()
{
    Color = color;
    gl_Position = vec4(position, 0.0, 1.0);
}
""")

# The fragment shader
fragment_source = Vector{UInt8}("""
# version 150

in vec3 Color;

out vec4 outColor;

void main()
{
    outColor = vec4(Color, 1.0);
}
""")

# Compile the vertex shader
vertex_shader = glCreateShader(GL_VERTEX_SHADER)
glShaderSource(vertex_shader, 1, Ptr{UInt8}[pointer(vertex_source)], Ref{GLint}(length(vertex_source)))  # nicer thanks to GLAbstraction
glCompileShader(vertex_shader)
# Check that it compiled correctly
status = Ref(GLint(0))
glGetShaderiv(vertex_shader, GL_COMPILE_STATUS, status)
if status[] != GL_TRUE
    buffer = Array(UInt8, 512)
    glGetShaderInfoLog(vertex_shader, 512, C_NULL, buffer)
    @error "$(unsafe_string(pointer(buffer), 512))"
end

# Compile the fragment shader
fragment_shader = glCreateShader(GL_FRAGMENT_SHADER)
glShaderSource(fragment_shader, 1, Ptr{UInt8}[pointer(fragment_source)], Ref{GLint}(length(fragment_source)))  # nicer thanks to GLAbstraction
glCompileShader(fragment_shader)
# Check that it compiled correctly
status = Ref(GLint(0))
glGetShaderiv(fragment_shader, GL_COMPILE_STATUS, status)
if status[] != GL_TRUE
    buffer = Array(UInt8, 512)
    glGetShaderInfoLog(fragment_shader, 512, C_NULL, buffer)
    @error "$(unsafe_string(pointer(buffer), 512))"
end

# Connect the shaders by combining them into a program
shader_program = glCreateProgram()
glAttachShader(shader_program, vertex_shader)
glAttachShader(shader_program, fragment_shader)
glBindFragDataLocation(shader_program, 0, "outColor") # optional

glLinkProgram(shader_program)
glUseProgram(shader_program)

# The vertices of our triangle, with color
positions = Point{2,Float32}[(0,    0.5),   # red vertex
                            ( 0.5, -0.5),   # green vertex
                            (-0.5, -0.5)]   # blue vertex

colors = Vec3f0[(1, 0, 0),                     # top-left
                (0, 1, 0),                     # top-right
                (0, 0, 1)]

vao = Ref(GLuint(0))
glGenVertexArrays(1, vao)
glBindVertexArray(vao[])

# Create the Vertex Buffer Objects (VBO)
vbo = Ref(GLuint(0))   # initial value is irrelevant, just allocate space
glGenBuffers(1, vbo)
glBindBuffer(GL_ARRAY_BUFFER, vbo[])
glBufferData(GL_ARRAY_BUFFER, sizeof(positions), positions, GL_STATIC_DRAW)
# Link vertex data to attributes
pos_attribute = glGetAttribLocation(shader_program, "position")
glEnableVertexAttribArray(pos_attribute)
glVertexAttribPointer(pos_attribute, 2,
                      GL_FLOAT, GL_FALSE, 0, C_NULL)

# Color VBO
vbo1 = Ref(GLuint(0))
glGenBuffers(1, vbo1)
glBindBuffer(GL_ARRAY_BUFFER, vbo1[])
glBufferData(GL_ARRAY_BUFFER, sizeof(colors), colors, GL_STATIC_DRAW)

# Link color data to attributes
col_attribute = glGetAttribLocation(shader_program, "color")
glBindBuffer(GL_ARRAY_BUFFER, vbo1[])
glEnableVertexAttribArray(col_attribute)
glVertexAttribPointer(col_attribute, 3,
                      GL_FLOAT, GL_FALSE, 0, C_NULL)

# Draw while waiting for a close event
glClearColor(0,0,0,0)
while !GLFW.WindowShouldClose(window)
    glClear(GL_COLOR_BUFFER_BIT)
    glDrawArrays(GL_TRIANGLES, 0, length(positions))
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS
        GLFW.SetWindowShouldClose(window, true)
    end
end
GLFW.SetWindowShouldClose(window, false)

# Now we define another geometry that we will render, a rectangle, this one with an index buffer
# The positions of the vertices in our rectangle
positions = Point{2,Float32}[(-0.5,  0.5),     # top-left
                             ( 0.5,  0.5),     # top-right
                             ( 0.5, -0.5),     # bottom-right
                             (-0.5, -0.5)]     # bottom-left

# The colors assigned to each vertex
colors = Vec3f0[(1, 0, 0),                     # top-left
                (0, 1, 0),                     # top-right
                (0, 0, 1),                     # bottom-right
                (1, 1, 1)]                     # bottom-left

# Specify how vertices are arranged into faces
# Face{N,T} type specifies a face with N vertices, with index type
# T (you should choose UInt32), and index-offset O. If you're
# specifying faces in terms of julia's 1-based indexing, you should set
# O=0. (If you instead number the vertices starting with 0, set
# O=-1.)
elements = Face{3,UInt32}[(0,1,2),          # the first triangle
                          (2,3,0)]          # the second triangle

vao = Ref(GLuint(0))
glGenVertexArrays(1, vao)
glBindVertexArray(vao[])

# Create the Vertex Buffer Objects (VBO)
vbo = Ref(GLuint(0))   # initial value is irrelevant, just allocate space
glGenBuffers(1, vbo)
glBindBuffer(GL_ARRAY_BUFFER, vbo[])
glBufferData(GL_ARRAY_BUFFER, sizeof(positions), positions, GL_STATIC_DRAW)
# Link vertex data to attributes
pos_attribute = glGetAttribLocation(shader_program, "position")
glEnableVertexAttribArray(pos_attribute)
glVertexAttribPointer(pos_attribute, 2,
                      GL_FLOAT, GL_FALSE, 0, C_NULL)

# Color VBO
vbo1 = Ref(GLuint(0))
glGenBuffers(1, vbo1)
glBindBuffer(GL_ARRAY_BUFFER, vbo1[])
glBufferData(GL_ARRAY_BUFFER, sizeof(colors), colors, GL_STATIC_DRAW)

# Link color data to attributes
col_attribute = glGetAttribLocation(shader_program, "color")
glBindBuffer(GL_ARRAY_BUFFER, vbo1[])
glEnableVertexAttribArray(col_attribute)
glVertexAttribPointer(col_attribute, 3,
                      GL_FLOAT, GL_FALSE, 0, C_NULL)

# Create the Element Buffer Object (EBO)
ebo = Ref(GLuint(0))
glGenBuffers(1, ebo)
glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo[])
glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(elements), elements, GL_STATIC_DRAW)

# Draw while waiting for a close event
glClearColor(0,0,0,0)
while !GLFW.WindowShouldClose(window)
    glClear(GL_COLOR_BUFFER_BIT)
    glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, C_NULL)
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS
        GLFW.SetWindowShouldClose(window, true)
    end
end

GLFW.DestroyWindow(window)  # needed if you're running this from the REPL