@testset "Storage data misspecification" begin
    # See https://discourse.julialang.org/t/how-to-use-test-warn/15557/5 about testing for warning throwing
    warn_message = "The data doesn't include devices of type GenericBattery, consider changing the device models"
    model = DeviceModel(GenericBattery, BookKeeping)
    c_sys5 = build_system("c_sys5")
    c_sys14 = build_system("c_sys14")
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5)
    @test_logs (:warn, warn_message) construct_device!(op_problem, :Storage, model)
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys14)
    @test_logs (:warn, warn_message) construct_device!(op_problem, :Storage, model)
end

@testset "Storage Basic Storage With DC - PF" begin
    model = DeviceModel(GenericBattery, BookKeeping)
    c_sys5_bat = build_system("c_sys5_bat")
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 72, 0, 72, 72, 24, false)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Storage Basic Storage With AC - PF" begin
    model = DeviceModel(GenericBattery, BookKeeping)
    c_sys5_bat = build_system("c_sys5_bat")
    op_problem = OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 96, 0, 96, 96, 24, false)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Storage with Reservation DC - PF" begin
    model = DeviceModel(GenericBattery, BookKeepingwReservation)
    c_sys5_bat = build_system("c_sys5_bat")
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 96, 0, 72, 72, 24, true)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Storage with Reservation With AC - PF" begin
    model = DeviceModel(GenericBattery, BookKeepingwReservation)
    c_sys5_bat = build_system("c_sys5_bat")
    op_problem = OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 120, 0, 96, 96, 24, true)
    psi_checkobjfun_test(op_problem, GAEVF)
end


module QJuliaFields

#load path to qjulia home directory
push!(LOAD_PATH, joinpath(@__DIR__, "../..", "core"))
push!(LOAD_PATH, @__DIR__)

import QJuliaRegisters
import QJuliaEnums
import QJuliaGrid

const complex_length = 2

abstract type QJuliaGenericField_qj end

mutable struct QJuliaLatticeField_qj <: QJuliaGenericField_qj
  # Field Matrix{T}, e.g., a single color spinor, a block color spinor (an eigenvector set), an SU(N) field etc.:
  v::Any

  # Lattice field structure
  field_desc::QJuliaGrid.QJuliaLatticeFieldDescr_qj

  QJuliaLatticeField_qj() = new( missing )

  function QJuliaLatticeField_qj(fdesc::QJuliaGrid.QJuliaLatticeFieldDescr_qj)

    # Since the field descriptor is an external argument (and data are not protected in Julia) we need this check again:
    if fdesc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY && fdesc.nSpin == nothing ; error("Spin dof is not defined for the field with scalar geometry."); end

    if fdesc.grid.location == QJuliaEnums.QJULIA_CPU_FIELD_LOCATION
      # Set the field array total elements:
      tot_elems = complex_length*fdesc.siteSubset*(fdesc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY ? fdesc.nSpin : fdesc.nColor)*fdesc.nColor*fdesc.grid.grid_volumeCB

      v = Matrix{fdesc.grid.register_type}(undef, tot_elems, Int(fdesc.geom)*fdesc.nBlock)
    else
      # Ok, set the field as a "missing" one. It will be inititalized later.
      v = missing
    end

    # call constructor
    new(v, fdesc)

  end #QJuliaLatticeField_qj constructor

  function QJuliaLatticeField_qj(grid::QJuliaGrid.QJuliaGridDescr_qj, geom::QJuliaEnums.QJuliaFieldGeometry_qj, parity::QJuliaEnums.QJuliaParity_qj; NSpin = nothing, NColor = 3, NBlock = 1)

    fdesc = QJuliaLatticeFieldDescr_qj{NSpin, NColor, NBlock}(grid::QJuliaGridDescr_qj, geom::QJuliaEnums.QJuliaFieldGeometry_qj, parity::QJuliaEnums.QJuliaParity_qj)

    if fdesc.grid.location == QJuliaEnums.QJULIA_CPU_FIELD_LOCATION
      # Set the field array total elements:
      tot_elems = complex_length*fdesc.siteSubset*(fdesc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY ? fdesc.nSpin : fdesc.nColor)*fdesc.nColor*fdesc.grid.grid_volumeCB

      v = Matrix{fdesc.grid.register_type}(undef, tot_elems, Int(fdesc.geom)*fdesc.nBlock)
    else
      # Ok, set the field as a "missing" one. It will be inititalized later.
      v = missing
    end

    # call constructor
    new(v, fdesc)

  end #QJuliaLatticeField_qj constructor

  function QJuliaLatticeField_qj(field::QJuliaLatticeField_qj, parity::QJuliaEnums.QJuliaParity_qj)

     if field.v === missing; error("Cannot reference a missing field"); end
     if field.field_desc.siteSubset == 1; error("Cannot reference a parity field from the parity field"); end

     fdesc = deepcopy(field.field_desc)

     griddsc   = fdesc.grid
     griddsc.X = NTuple{QJuliaGrid.QJULIA_MAX_DIMS, Int}(ntuple(i->(i == 1 ? field.field_desc.grid.X[i] >> 1 : field.field_desc.grid.X[i]), QJuliaGrid.QJULIA_MAX_DIMS))

     griddsc.volume        = field.field_desc.grid.volumeCB
     griddsc.volumeCB      = field.field_desc.grid.volumeCB
     griddsc.grid_volume   = field.field_desc.grid.grid_volumeCB
     griddsc.grid_volumeCB = field.field_desc.grid.grid_volumeCB

     fdesc.parity     = parity
     fdesc.siteSubset = 1

     parity_volume = complex_length*griddsc.grid_volumeCB*(fdesc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY ? fdesc.nSpin : fdesc.nColor)*fdesc.nColor

     offset = parity == QJuliaEnums.QJULIA_EVEN_PARITY ? 1 : parity_volume+1

     v = view(field.v, offset:offset+parity_volume-1, :)

     new(v, fdesc)

  end

end #QJuliaLatticeField_qj

# Pointer helpers:
@inline function GetPtr(field::QJuliaGenericField_qj); return field.v; end

@inline function GetParityPtr(field::QJuliaGenericField_qj, parity::QJuliaEnums.QJuliaParity_qj)
   if(field.field_desc.siteSubset == 1); return field.v; end

   parity_volume = complex_length*field.field_desc.real_volumeCB*(field.field_desc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY ? field.field_desc.nSpin : field.field_desc.nColor)*field.field_desc.nColor
   offset        = parity == QJuliaEnums.QJULIA_EVEN_PARITY ? 1 : parity_volume+1

   return view(field.v, offset:offset+parity_volume-1, :)
end

function field_info(field::QJuliaGenericField_qj; vrbs = true)

  field_desc = field.field_desc
  println(" ")
  println("General field info for ", typeof(field), ": ")

  if vrbs == true; QJuliaGrid.field_descr_info(field_desc); end

  println("Field data reference type : ", typeof(field.v))
  println("Field object pointer: ", pointer_from_objref(field))

  if field.v === missing
    println("Warning: data array was not allocated.")
  elseif field_desc.grid.location == QJuliaEnums.QJULIA_CUDA_FIELD_LOCATION
    println("Warning: data array was allocated on GPU")
    println("Physical length : ", length(field.v))
  else
    println("Row data pointer: ", pointer(field.v))
    println("Bytes : ", length(field.v))
  end

  println(" ")

end

# Specialized outer constructors:

# WARNING: transfered into CUDA impl
#CreateGenericField(fdesc::QJuliaGrid.QJuliaLatticeFieldDescr_qj) = QJuliaLatticeField_qj(fdesc)

CreateColorSpinor(grid::QJuliaGrid.QJuliaGridDescr_qj; parity::QJuliaEnums.QJuliaParity_qj = QJuliaEnums.QJULIA_INVALID_PARITY) = QJuliaLatticeField_qj(grid, QJuliaEnums.QJULIA_SCALAR_GEOMETRY, parity)

CreateBlockColorSpinor(grid::QJuliaGrid.QJuliaGridDescr_qj, NBlock; parity::QJuliaEnums.QJuliaParity_qj = QJuliaEnums.QJULIA_INVALID_PARITY) = QJuliaLatticeField_qj(grid, QJuliaEnums.QJULIA_SCALAR_GEOMETRY, parity; nBlock = NBlock)

CreateGaugeField(grid::QJuliaGrid.QJuliaGridDescr_qj) = QJuliaLatticeField_qj(grid, QJuliaEnums.QJULIA_VECTOR_GEOMETRY, QJuliaEnums.QJULIA_INVALID_PARITY)

#Create references to parity fields
Even(field::QJuliaGenericField_qj) = QJuliaLatticeField_qj(field, QJuliaEnums.QJULIA_EVEN_PARITY)
Odd(field::QJuliaGenericField_qj)  = QJuliaLatticeField_qj(field, QJuliaEnums.QJULIA_ODD_PARITY)

function (field::QJuliaLatticeField_qj)(i::Int)
  return field.v[i]
end

# Field Accessors
# input : field object
# out   : tuple of data array and args (which is a tuple of dof and offsets)

struct AccessorArgs
  #
  nParity::Int
  #
  nSpin::Int
  #
  nColor::Int
  #
  nVec::Int
  #
  length::Int
  #
  offset_cb::Int
  #
  field_order_id::Int
  # 1 for matter fields, 4 for gauge fields
  geom::Int
end

function create_field_accessor(field::QJuliaGenericField_qj; verbosity = true)

  nParity = field.field_desc.siteSubset
  nSpin   = field.field_desc.nSpin == nothing ? 0 : field.field_desc.nSpin
  nColor  = field.field_desc.nColor
  nVecs   = field.field_desc.nBlock

  field_order_id = 0 # currently default field order id : color inside spin inside space-time for spinors, color cols inside color rows inside space-time for gauge fields

  stride  = complex_length*field.field_desc.grid.grid_volume #include padding

  if     field.field_desc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY
    field_order_shape  = (stride, nColor, nSpin, nVecs)
    geom = 1
  elseif field.field_desc.geom == QJuliaEnums.QJULIA_VECTOR_GEOMETRY
    field_order_shape  = (stride, nColor, nColor, 4) #4 directions
    geom = 4
  elseif field.field_desc.geom == QJuliaEnums.QJULIA_COARSE_GEOMETRY
    field_order_shape  = (stride, nColor, nColor, nSpin, nSpin, 4) #4 directions
    geom = 8
  else
    error("Field geometry is not supported.")
  end

  if verbosity == true ; println("Create field accessor with dims ", field_order_shape);end

  field_order = reshape(field.v, field_order_shape)
  #
  offset_cb   = Int(length(field.v) / field.field_desc.siteSubset)
  #
  args        = AccessorArgs(nParity, nSpin, nColor, nVecs, stride, offset_cb, field_order_id, geom)
  #
  return (field_order, args)
end

end # module QJuliaFields


using .MLJ

export deterministic, probabilistic

# deterministic model api
function deterministic(model, Xs::MLJ.Source, ys::MLJ.Source)
    model = scitype(model) <: MLJ.UnsupervisedScitype ?
            transform(machine(model, Xs), Xs) :
            predict(machine(model, Xs, ys), Xs)
    mach = machine(Deterministic(), Xs, ys, predict=model)
    return mach
end

# deterministic network api
function deterministic(model, Xs::MLJ.Source, ys::MLJ.Source, scores...)
    model = transform(machine(model), scores...)
    mach = machine(Deterministic(), Xs, ys, predict=model)
    return mach
end

# probabilistic model api
function probabilistic(model, Xs::MLJ.Source, ys::MLJ.Source)
    model = scitype(model) <: MLJ.UnsupervisedScitype ?
            transform(machine(model, Xs), Xs) :
            predict(machine(model, Xs, ys), Xs)
    mach = machine(Probabilistic(), Xs, ys, predict=model)
    return mach
end

# probabilistic network api
function probabilistic(model, Xs::MLJ.Source, ys::MLJ.Source, scores...)
    model = transform(machine(model), scores...)
    mach = machine(Probabilistic(), Xs, ys, predict=model)
    return mach
end


function position_calc(x...)
    _position_calc(Iterators.filter(x->!isa(x, Nothing), x)...)
end

function normal_calc(x::Bool, invert_normals::Bool = false)
    i = invert_normals ? "-" : ""
    if x
        return "$(i)getnormal(position, position_x, position_y, position_z, o_uv);"
    else
        return "vec3(0, 0, $(i)1);"
    end
end

function light_calc(x::Bool)
    if x
        """
        vec3 L      = normalize(o_lightdir);
        vec3 N      = normalize(o_normal);
        vec3 light1 = blinnphong(N, o_camdir, L, color.rgb);
        vec3 light2 = blinnphong(N, o_camdir, -L, color.rgb);
        color       = vec4(ambient * color.rgb + light1 + backlight * light2, color.a);
        """
    else
        ""
    end
end

function _position_calc(
        position_x::MatTypes{T}, position_y::MatTypes{T}, position_z::MatTypes{T}, target::Type{Texture}
    ) where T<:AbstractFloat
    """
    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
    ivec2 index2D = ind2sub(dims, index1D);
    vec2 index01 = vec2(index2D) / (vec2(dims)-1.0);
    pos = vec3(
        texelFetch(position_x, index2D, 0).x,
        texelFetch(position_y, index2D, 0).x,
        texelFetch(position_z, index2D, 0).x
    );
    """
end

function _position_calc(
        position_x::VectorTypes{T}, position_y::VectorTypes{T}, position_z::MatTypes{T},
        target::Type{Texture}
    ) where T<:AbstractFloat
    """
    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
    ivec2 index2D = ind2sub(dims, index1D);
    vec2 index01 = vec2(index2D) / (vec2(dims)-1.0);
    pos = vec3(
        texelFetch(position_x, index2D.x, 0).x,
        texelFetch(position_y, index2D.y, 0).x,
        texelFetch(position_z, index2D, 0).x
    );
    """
end

function _position_calc(
        position_xyz::VectorTypes{T}, target::Type{TextureBuffer}
    ) where T <: StaticVector
    "pos = texelFetch(position, index).xyz;"
end

function _position_calc(
        position_xyz::VectorTypes{T}, target::Type{GLBuffer}
    ) where T <: StaticVector
    len = length(T)
    filler = join(ntuple(x->0, 3-len), ", ")
    needs_comma = len != 3 ? ", " : ""
    "pos = vec3(position $needs_comma $filler);"
end

function _position_calc(
        grid::Grid{2}, position_z::MatTypes{T}, target::Type{Texture}
    ) where T<:AbstractFloat
    """
    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
    ivec2 index2D = ind2sub(dims, index1D);
    vec2 index01 = vec2(index2D) / (vec2(dims)-1.0);
    float height = texture(position_z, index01).x;
    pos = vec3(grid_pos(position, index01), height);
    """
end

@nospecialize
# surface(::Matrix, ::Matrix, ::Matrix)
function draw_surface(main::Tuple{MatTypes{T}, MatTypes{T}, MatTypes{T}}, data::Dict) where T <: AbstractFloat
    @gen_defaults! data begin
        position_x = main[1] => (Texture, "x position, must be a `Matrix{Float}`")
        position_y = main[2] => (Texture, "y position, must be a `Matrix{Float}`")
        position_z = main[3] => (Texture, "z position, must be a `Matrix{Float}`")
        scale = Vec3f(0) => "scale must be 0, for a surfacemesh"
    end
    return draw_surface(position_z, data)
end

# surface(Vector or Range, Vector or Range, ::Matrix)
function draw_surface(main::Tuple{VectorTypes{T}, VectorTypes{T}, MatTypes{T}}, data::Dict) where T <: AbstractFloat
    @gen_defaults! data begin
        position_x = main[1] => (Texture, "x position, must be a `Vector{Float}`")
        position_y = main[2] => (Texture, "y position, must be a `Vector{Float}`")
        position_z = main[3] => (Texture, "z position, must be a `Matrix{Float}`")
        scale = Vec3f(0) => "scale must be 0, for a surfacemesh"
    end
    return draw_surface(position_z, data)
end

function draw_surface(main, data::Dict)
    primitive = triangle_mesh(Rect2(0f0,0f0,1f0,1f0))
    to_opengl_mesh!(data, primitive)
    @gen_defaults! data begin
        scale = nothing
        position = nothing
        position_x = nothing => Texture
        position_y = nothing => Texture
        position_z = nothing => Texture
        image = nothing => Texture
        shading = true
        normal = shading
        invert_normals = false
        backlight = 0f0
    end
    @gen_defaults! data begin
        color = nothing => Texture
        color_map = nothing => Texture
        color_norm = nothing
        fetch_pixel = false
        matcap = nothing => Texture

        nan_color = RGBAf(1, 0, 0, 1)
        highclip = RGBAf(0, 0, 0, 0)
        lowclip = RGBAf(0, 0, 0, 0)

        uv_scale = Vec2f(1)
        instances = const_lift(x->(size(x,1)-1) * (size(x,2)-1), main) => "number of planes used to render the surface"
        transparency = false
        shader = GLVisualizeShader(
            "fragment_output.frag", "util.vert", "surface.vert",
            "mesh.frag",
            view = Dict(
                "position_calc" => position_calc(position, position_x, position_y, position_z, Texture),
                "normal_calc" => normal_calc(normal, to_value(invert_normals)),
                "light_calc" => light_calc(shading),
                "buffers" => output_buffers(to_value(transparency)),
                "buffer_writes" => output_buffer_writes(to_value(transparency))
            )
        )
    end
    return assemble_shader(data)
end
@specialize


module PowerSimulations

#################################################################################
# Exports

# Base Models
export Stage
export Simulation
export OperationModel
export ModelReference
export InitialCondition

#Network Relevant Exports
export StandardPTDFModel
export CopperPlatePowerModel

######## Device Models ########
export DeviceModel
######## Service Models ########
export ServiceModel
######## Branch Models ########
export StaticLine
export StaticTransformer
export TapControl
export StaticLineUnbounded
#export PhaseControl
export HVDCLossless
export HVDCDispatch
#export VoltageSourceDC
######## Load Models ########
export StaticPowerLoad
export InterruptiblePowerLoad
export DispatchablePowerLoad
######## Renewable Formulations ########
export RenewableFixed
export RenewableFullDispatch
export RenewableConstantPowerFactor
######## Hydro Formulations ########
export HydroFixed
######## Renewable Formulations ########
export BookKeeping
export BookKeepingwReservation
######## Thermal Formulations ########
export ThermalStandardUnitCommitment
export ThermalBasicUnitCommitment
export ThermalDispatch
export ThermalRampLimited
export ThermalDispatchNoMin

# feedforward sequences
export RecedingHorizon
export Synchronize
export Sequential

# feedforward models
export UpperBoundFF
export SemiContinuousFF
export RangeFF

# Initial Conditions Quantities
export DevicePower
export DeviceStatus
export TimeDurationON
export TimeDurationOFF
export DeviceEnergy

# cache_models
export TimeStatusChange

#operation_models
#export UnitCommitment
#export EconomicDispatch
#export OptimalPowerFlow

# Functions
## Construction Exports
export construct_device!
export construct_network!
## Op Model Exports
export solve_op_model!
export get_initial_conditions
export set_transmission_ref!
export set_devices_ref!
export set_branches_ref!
export set_services_ref!
export set_device_model!
export set_branch_model!
export set_device_model!
## Sim Model Exports
export run_sim_model!
## Utils Exports
export write_op_model
export write_model_results
export load_operation_results
export load_simulation_results
export get_all_constraint_index
export get_all_var_index
export get_con_index
export get_var_index
# Plotting Utils
export sort_data
export get_stacked_plot_data
export get_bar_plot_data
export get_stacked_generation_data
export bar_plot
export stack_plot
export report

#################################################################################
# Imports
#Modeling Imports
import JuMP
# so that users do not need to import JuMP to use a solver with PowerModels
import JuMP: with_optimizer
export with_optimizer
import MathOptInterface
import ParameterJuMP
import LinearAlgebra
import PowerSystems
import InfrastructureSystems
import PowerModels
import RecipesBase
import Requires

#TimeStamp Management Imports
import Dates
import TimeSeries

#I/O Imports
import MathOptFormat
import DataFrames
import Feather

include("core/definitions.jl")

#################################################################################
##### JuMP methods overloading
JuMP.Model(optimizer::Nothing; kwargs...) = JuMP.Model(kwargs...)

################################################################################
# Includes

#Abstract Models
include("network_models/networks.jl")
include("service_models/services.jl")

#Core Models and constructors
include("core/core_structs/aux_structs.jl")
include("core/core_structs/cache_models.jl")
include("core/core_structs/feedforward_model.jl")
include("core/core_structs/device_model.jl")
include("core/core_structs/initial_conditions.jl")
include("core/core_structs/canonical_model.jl")
include("core/core_structs/service_model.jl")
include("core/core_structs/operation_model.jl")
include("core/core_structs/chronology.jl")
include("core/core_structs/simulations_stages.jl")
include("core/core_structs/simulation_model.jl")
include("core/core_structs/results_model.jl")
include("core/build_cache.jl")
include("core/build_operations.jl")
include("core/build_simulations.jl")

#FeedForward Model Files
include("simulation_models/feedforward_models.jl")

#Device Modeling components
include("devices/device_models/common.jl")
include("devices/device_models/renewable_generation.jl")
include("devices/device_models/thermal_generation.jl")
include("devices/device_models/electric_loads.jl")
include("devices/device_models/AC_branches.jl")
include("devices/device_models/DC_branches.jl")
include("devices/device_models/storage.jl")
include("devices/device_models/hydro_generation.jl")

#Network models
include("network_models/copperplate_model.jl")
include("network_models/powermodels_interface.jl")
include("network_models/ptdf_model.jl")

#Device constructors
include("devices/device_constructors/device_constructors.jl")

#Network constructors
include("network_models/network_constructor.jl")

#Services Models
#include("service_models/reserves.jl")

#Services constructors
include("service_models/services_constructor.jl")

#Operational Model Constructors
include("operation_models/operation_models.jl")

#Simulations Model Files
include("simulation_models/stage_update.jl")

#Routines
include("routines/make_initial_conditions.jl")
include("routines/get_results.jl")
include("routines/solve_routines.jl")
include("routines/write_model.jl")

#Utils
include("utils/optimization_debugging.jl")
include("utils/printing.jl")
include("utils/plot_results.jl")
include("utils/plot_recipes.jl")

#Initialization

function __init__()
   Requires.@require Weave = "44d3d7a6-8a23-5bf8-98c5-b353f8df5ec9" include("utils/make_report.jl")
end

end


abstract type AbstractBranchFormulation <: AbstractDeviceFormulation end

abstract type AbstractLineForm <: AbstractBranchFormulation end

abstract type AbstractDCLineForm <: AbstractBranchFormulation end

abstract type AbstractTransformerForm <: AbstractBranchFormulation end

abstract type PiLine <: AbstractLineForm end

abstract type SeriesLine <: AbstractLineForm end

abstract type SimpleHVDC <: AbstractDCLineForm end

function line_flow_limit(ps_m::CanonicalModel,
                         devices::Array{Br,1},
                         device_formulation::Type{D},
                         system_formulation::Type{S},
                         time_range::UnitRange{Int64}) where {Br <: PSY.MonitoredLine,
                                                               D <: AbstractBranchFormulation,
                                                               S <: PM.AbstractPowerFormulation}

    #rate_data = [(h.name, (min = -1*h.rate, max = h.rate) for h in devices]

    device_range(ps_m, range_data, time_range, :dc_rate_const, :Fbr)

    return

end

function line_rate_constraints(ps_m::CanonicalModel,
                               devices::Array{Br,1},
                               device_formulation::Type{D},
                                system_formulation::Type{PSI.StandardPTDFForm},
                          time_range::UnitRange{Int64}) where {Br <: PSY.Branch, D <: AbstractBranchFormulation}

    range_data = [(h.name, (min = -1*h.rate, max = h.rate)) for h in devices]

    device_range(ps_m, range_data, time_range, :line_rate_limit, :Fbr)

    return

end

function line_flow_limit(ps_m::CanonicalModel,
                         devices::Array{Br,1},
                         device_formulation::Type{D},
                         system_formulation::Type{S},
                         time_range::UnitRange{Int64}) where {Br <: PSY.MonitoredLine,
                                                               D <: AbstractBranchFormulation,
                                                               S <: PM.AbstractActivePowerFormulation}

    #rate_data = [(h.name, (min = -1*h.rate, max = h.rate) for h in devices]

    device_range(ps_m, range_data, time_range, :dc_rate_const, :Fbr)

    return

end

function flow_variables(ps_m::CanonicalModel,
                        system_formulation::Type{S},
                        devices::Array{B,1},
                        time_range::UnitRange{Int64}) where {B <: PSY.Branch,
                                                            S <: PM.DCPlosslessForm}

    add_variable(ps_m, devices, time_range, :Fbr, false)

end

function flow_variables(ps_m::CanonicalModel,
                        system_formulation::Type{S},
                        devices::Array{B,1},
                        time_range::UnitRange{Int64}) where {B <: PSY.Branch,
                                                             S <: PM.AbstractDCPLLForm}

    add_variable(ps_m, devices, time_range, :Fbr_to, false)
    add_variable(ps_m, devices, time_range, :Fbr_fr, false)

end

function flow_variables(ps_m::CanonicalModel,
                        system_formulation::Type{S},
                        devices::Array{B,1},
                        time_range::UnitRange{Int64}) where {B <: PSY.Branch,
                                                             S <:PM.AbstractPowerFormulation}

    add_variable(ps_m, devices, time_range, :PFbr_to, false)
    add_variable(ps_m, devices, time_range, :PFbr_fr, false)

    add_variable(ps_m, devices, time_range, :QFbr_to, false)
    add_variable(ps_m, devices, time_range, :QFbr_fr, false)

    return

end


mutable struct GFRegenerateState
    prev_trace::GFTrace
    trace::GFTrace
    selection::AddressSet
    score::Float64
    weight::Float64
    visitor::AddressVisitor
    params::Dict{Symbol,Any}
    args_change::Any
    retchange::Union{Some{Any},Nothing}
    callee_output_changes::HomogenousTrie{Any,Any}
end

function GFRegenerateState(args_change, prev_trace, selection, params)
    visitor = AddressVisitor()
    GFRegenerateState(prev_trace, GFTrace(), selection, 0., 0., visitor,
                    params, args_change, nothing, HomogenousTrie{Any,Any}())
end

get_args_change(state::GFRegenerateState) = state.args_change

function set_ret_change!(state::GFRegenerateState, value)
    if state.retchange === nothing
        state.retchange = value
    else
        error("@retchange! was already used")
    end
end

function get_addr_change(state::GFRegenerateState, addr)
    get_leaf_node(state.callee_output_changes, addr)
end

function addr(state::GFRegenerateState, dist::Distribution{T}, args, addr, args_change) where {T}
    visit!(state.visitor, addr)
    has_previous = has_primitive_call(state.prev_trace, addr)
    if has_internal_node(state.selection, addr)
        error("Got internal node but expected leaf node in selection at $addr")
    end
    in_selection = has_leaf_node(state.selection, addr)
    local retval::T
    if has_previous && !in_selection
        # there was a previous value, and it was not in the selection
        # use the previous value
        # it contributes to the weight
        prev_call = get_primitive_call(state.prev_trace, addr)
        retval = prev_call.retval
        score = logpdf(dist, retval, args...)
        state.weight += score - prev_call.score
        retchange = NoChange()
    elseif has_previous
        # there was a previous value, and it was in the selection
        # simulate a new value
        # it does not contribute to the weight
        # retchange indicates that there was a change and gives the previous value
        retval = random(dist, args...)
        score = logpdf(dist, retval, args...)
        prev_call = get_primitive_call(state.prev_trace, addr)
        retchange = (true, prev_call.retval)
    else
        # there is no previous value
        # simulate a new value
        # it does not contribute to the weight
        # retchange is nothing, because the address is new
        retval = random(dist, args...)
        score = logpdf(dist, retval, args...)
        retchange = nothing
    end
    set_leaf_node!(state.callee_output_changes, addr, retchange)
    call = CallRecord(score, retval, args)
    state.trace = assoc_primitive_call(state.trace, addr, call)
    state.score += score
    retval
end

function addr(state::GFRegenerateState, gen::Generator{T}, args, addr, args_change) where {T}
    visit!(state.visitor, addr)
    has_previous = has_subtrace(state.prev_trace, addr)
    if has_previous
        if addr in state.selection
            error("Entire sub-traces cannot be selected, tried to select $addr")
        end
        prev_trace = get_subtrace(state.prev_trace, addr)
        prev_call = get_call_record(prev_trace)
        if has_internal_node(state.selection, addr)
            selection = get_internal_node(state.selection, addr)
        else
            selection = EmptyAddressSet()
        end
        (trace, weight, retchange) = regenerate(
            gen, args, args_change, prev_trace, selection)
        state.weight += weight
        set_leaf_node!(state.callee_output_changes, addr, retchange)
    else
        trace = simulate(gen, args)
        set_leaf_node!(state.callee_output_changes, addr, nothing)
    end
    call = get_call_record(trace)
    state.score += call.score
    state.trace = assoc_subtrace(state.trace, addr, trace)
    call.retval::T
end

function splice(state::GFRegenerateState, gf::GenFunction, args::Tuple)
    exec(gf, state, args)
end

function regenerate(gen::GenFunction, new_args, args_change, trace, selection)
    state = GFRegenerateState(args_change, trace, selection, gen.params)
    retval = exec(gen, state, new_args)
    new_call = CallRecord(state.score, retval, new_args)
    state.trace.call = new_call
    (state.trace, state.weight, state.retchange)
end


# -----------------------------------------------------------------------------
# Implementation of Atomistic MolecularDynamicsSimulator API
# -----------------------------------------------------------------------------

"""
    MollySimulator{S,T<:Unitful.Time,C} <: MolecularDynamicsSimulator

A wrapper around Molly to implement the Atomistic API.

**Type parameter descriptions**
- `S` the type of simulator
- `T<:Unitful.Time` the type for the time fields
- `C` the type for the coupling field

**Field descriptions**
- `Δt::T` the time between timesteps, assumed to be in atomic units
- `steps::Integer` the number of timesteps for the simulation
- `t₀::T` the starting time of the simulation, assumed to be in atomic units;
    defaults to 0
- `coupling::C` the coupling for the simulation;
    many options are defined by `Molly`, but a user could also define a custom thermostat;
    defaults to `NoCoupling()`
- `stride::Int` the number of timesteps between logger runs
    defaults to 1
"""
struct MollySimulator{S,T<:Unitful.Time,C} <: MolecularDynamicsSimulator
    Δt::T
    steps::Int
    t₀::T
    coupling::C
    stride::Int
end
function MollySimulator{S}(Δt::T, steps::Int; t₀::Real=zero(T), coupling::C=NoCoupling(), stride::Int=1) where {S,T<:Real,C}
    Δt, t₀ = promote(Δt * TIME_UNIT, t₀ * TIME_UNIT)
    MollySimulator{S,typeof(Δt),C}(Δt, steps, t₀, coupling, stride)
end
function MollySimulator{S}(Δt::T, steps::Int; t₀::Unitful.Time=zero(T), coupling::C=NoCoupling(), stride::Int=1) where {S,T<:Unitful.Time,C}
    Δt, t₀ = promote(Δt, t₀)
    MollySimulator{S,T,C}(Δt, steps, t₀, coupling, stride)
end
function MollySimulator(Δt::T, steps::Int; t₀::Real=zero(T), coupling::C=NoCoupling(), stride::Int=1) where {T<:Real,C}
    MollySimulator{Molly.VelocityVerlet}(Δt, steps; t₀=t₀, coupling=coupling, stride=stride)
end
function MollySimulator(Δt::T, steps::Int; t₀::Unitful.Time=zero(T), coupling::C=NoCoupling(), stride::Int=1) where {T<:Unitful.Time,C}
    MollySimulator{Molly.VelocityVerlet}(Δt, steps; t₀=t₀, coupling=coupling, stride=stride)
end

function simulate(system::AbstractSystem{3}, simulator::MollySimulator{S}, potential::AbstractPotential) where {S}
    wrapper = InteratomicPotentialInter(potential)
    loggers = Dict(
        "c" => CoordinateLogger(LENGTH_TYPE, simulator.stride),
        "v" => VelocityLogger(VELOCITY_TYPE, simulator.stride),
        "p" => PotentialEnergyLogger(ENERGY_TYPE, simulator.stride),
        "k" => KineticEnergyLogger(ENERGY_TYPE, simulator.stride),
        "t" => TemperatureLogger(TEMPERATURE_TYPE, simulator.stride)
    )
    system = System(system; general_inters=(wrapper,), loggers=loggers)
    simulate!(system, S(dt=simulator.Δt, coupling=simulator.coupling), simulator.steps)
    MollyResult(system, simulator)
end

# -----------------------------------------------------------------------------
# Integration with InteratomicPotentials
# -----------------------------------------------------------------------------

# Internal struct that represents a potential in the format accepted by Molly
# Wraps the underlying InteratomicPotential with a cache of the potential energy for the current timestep
struct InteratomicPotentialInter{P<:AbstractPotential,E<:Unitful.Energy}
    potential::P
    energy_cache::Ref{E}
    InteratomicPotentialInter(potential::AbstractPotential) = new{typeof(potential),ENERGY_TYPE}(potential, Ref(Inf * ENERGY_UNIT))
end

function Molly.forces(inter::InteratomicPotentialInter, sys, neighbors=nothing)
    eandf = energy_and_force(sys, inter.potential)
    inter.energy_cache[] = eandf.e
    eandf.f
end

function Molly.potential_energy(inter::InteratomicPotentialInter, sys, neighbors=nothing)
    isinf(inter.energy_cache[]) ? InteratomicPotentials.potential_energy(sys, inter.potential) : inter.energy_cache[]
end


#!/usr/bin/env julia
# Tim Sterne-Weiler 2015

using Pkg

const dir = abspath( splitdir(@__FILE__)[1] )
const ver = readline(open(dir * "/../VERSION"))

start = time_ns()
println( stderr, "Whippet $ver loading and compiling... " )

Pkg.activate(dir * "/..")

using ArgParse
using Libz

using Whippet

function parse_cmd()
   s = ArgParseSettings()
   @add_arg_table s begin
    "--tpm"
      help     = "File containing two columns, isoform names, and transcripts-per-million (TpM)"
      arg_type = String
    "--index", "-x"
      help     = "Output prefix for saving index 'dir/prefix' (default Whippet/index/graph[.jls])"
      arg_type = String
      default  = fixpath( "$dir/../../index/graph" )
    "--out", "-o"
      help     = "Where should the gzipped output go 'dir/prefix'?"
      arg_type = String
      default  = fixpath( "./output" )
   end
   return parse_args(s)
end

function main()

   args = parse_cmd()

   println(stderr, " $( round( (time_ns()-start)/1e9, digits=6 ) ) seconds" )

   indexpath = fixpath( args["index"] )
   println(stderr, "Loading splice graph index... $( indexpath ).jls")
   @timer const lib = open(deserialize, "$( indexpath ).jls")

   data = load_tpm_file( fixpath(args["tpm"]) )

   tpms = Vector{Vector{Float64}}(length(lib.graphs))
   allocate_tpms!( tpms, lib )
   fill_tpms!( tpms, lib, data )

   io = open( args["out"], "w" )
   stream = ZlibDeflateOutputStream( io )

   output_tpms( stream, tpms, lib )

   close(stream)
   close(io)

   println(stderr, "Whippet $ver done." )
end

function load_tpm_file( filename::String, header=true )
   tpmfile = Dict{String,Float64}()
   fh = open( filename )
   for l in eachline(fh)
      if header
         header = false
         continue
      end
      s = split(l)
      tpmfile[String(s[1])] = parse(Float64, s[2])
   end
   tpmfile
end

function allocate_tpms!( tpms, lib )
   for i in 1:length(lib.graphs)
      tpms[i] = zeros(length(lib.graphs[i].annoname))
   end
   tpms
end

function fill_tpms!( tpms, lib, data )
   for i in 1:length(lib.graphs)
      for j in 1:length(lib.graphs[i].annoname)
         if haskey( data, lib.graphs[i].annoname[j] )
            tpms[i][j] = data[ lib.graphs[i].annoname[j] ]
         end
      end
   end
   tpms
end

function flanking_spliced( edges, node, name )
   upstream   = false
   downstream = false
   for i in node:-1:1
      if edges[i] in (EDGETYPE_RR, EDGETYPE_LR, EDGETYPE_SR)
         upstream = true
         #print("$name $(edges[i]):$i for n:$node")
         break
      elseif edges[i] in (EDGETYPE_LS, EDGETYPE_SL, EDGETYPE_RS)
         return false
      end
   end
   for i in node+1:length(edges)
      if edges[i] in (EDGETYPE_LL, EDGETYPE_LR, EDGETYPE_LS)
         downstream = true
         #println(" to $(edges[i]):$i")
         break
      elseif edges[i] in (EDGETYPE_SR, EDGETYPE_RS, EDGETYPE_SL)
         return false
      end
   end
   return (upstream && downstream)
end

function output_tpms( stream, tpms, lib )
   for i in 1:length(lib.graphs)
      for n in 1:length(lib.graphs[i].nodelen)
         if lib.graphs[i].edgetype[n]   in (EDGETYPE_RR, EDGETYPE_LR, EDGETYPE_LL) &&
            lib.graphs[i].edgetype[n+1] in (EDGETYPE_LL, EDGETYPE_LR, EDGETYPE_RR) &&
            flanking_spliced(lib.graphs[i].edgetype, n, lib.info[i].gene)

             tpm_total = 0.0
             tpm_node  = 0.0
             for j in 1:length(lib.graphs[i].annopath)
                curpath = lib.graphs[i].annopath[j]
                if n in first(curpath):last(curpath) # is within bounds
                   tpm_total += tpms[i][j]
                   if n in curpath
                      tpm_node += tpms[i][j]
                   end
                end
             end
             if tpm_total > 0.0
                psi = tpm_node / tpm_total
                write( stream, lib.info[i].gene * "\t" )
                write( stream, lib.info[i].name * ":" * string(lib.graphs[i].nodecoord[n]) *
                                                  "-" * string(lib.graphs[i].nodecoord[n]+lib.graphs[i].nodelen[n]-1) * "\t" )
                write( stream, string(n) * "\t" )
                write( stream, string(psi) * "\n" )
             end
         end
      end
   end
end

@timer main()


const DevicesModelContainer = Dict{Symbol, DeviceModel}
const BranchModelContainer = Dict{Symbol, DeviceModelForBranches}
const ServicesModelContainer = Dict{Tuple{String, Symbol}, ServiceModel}

"""
    ProblemTemplate(::Type{T}) where {T<:PM.AbstractPowerFormulation}

Creates a model reference of the PowerSimulations Optimization Problem.

# Arguments

  - `model::Type{T<:PM.AbstractPowerFormulation}`:

# Example

template = ProblemTemplate(CopperPlatePowerModel)
"""
mutable struct ProblemTemplate
    network_model::NetworkModel{<:PM.AbstractPowerModel}
    devices::DevicesModelContainer
    branches::BranchModelContainer
    services::ServicesModelContainer
    function ProblemTemplate(network::NetworkModel{T}) where {T <: PM.AbstractPowerModel}
        new(
            network,
            DevicesModelContainer(),
            BranchModelContainer(),
            ServicesModelContainer(),
        )
    end
end

ProblemTemplate(::Type{T}) where {T <: PM.AbstractPowerModel} =
    ProblemTemplate(NetworkModel(T))
ProblemTemplate() = ProblemTemplate(CopperPlatePowerModel)

get_device_models(template::ProblemTemplate) = template.devices
get_branch_models(template::ProblemTemplate) = template.branches
get_service_models(template::ProblemTemplate) = template.services
get_network_model(template::ProblemTemplate) = template.network_model
get_network_formulation(template::ProblemTemplate) =
    get_network_formulation(get_network_model(template))

function get_model(template::ProblemTemplate, device_type)
    if device_type <: PSY.Device
        return get(template.devices, Symbol(device_type), nothing)
    elseif device_type <: PSY.Branch
        return get(template.branches, Symbol(device_type), nothing)
    elseif device_type <: PSY.Service
        return get(template.services, Symbol(device_type), nothing)
    else
        error("not supported: $device_type")
    end
end

# Note to devs. PSY exports set_model! these names are chosen to avoid name clashes

"""
Sets the network model in a template.
"""
function set_network_model!(
    template::ProblemTemplate,
    model::NetworkModel{<:PM.AbstractPowerModel},
)
    template.network_model = model
    return
end

"""
Sets the device model in a template using the component type and formulation.
Builds a default DeviceModel
"""
function set_device_model!(
    template::ProblemTemplate,
    component_type::Type{<:PSY.Device},
    formulation::Type{<:AbstractDeviceFormulation},
)
    set_device_model!(template, DeviceModel(component_type, formulation))
    return
end

"""
Sets the device model in a template using a DeviceModel instance
"""
function set_device_model!(
    template::ProblemTemplate,
    model::DeviceModel{<:PSY.Device, <:AbstractDeviceFormulation},
)
    _set_model!(template.devices, model)
    return
end

function set_device_model!(
    template::ProblemTemplate,
    model::DeviceModel{<:PSY.Branch, <:AbstractDeviceFormulation},
)
    _set_model!(template.branches, model)
    return
end

"""
Sets the service model in a template using a name and the service type and formulation.
Builds a default ServiceModel with use_service_name set to true.
"""
function set_service_model!(
    template::ProblemTemplate,
    service_name::String,
    service_type::Type{<:PSY.Service},
    formulation::Type{<:AbstractServiceFormulation},
)
    set_service_model!(
        template,
        service_name,
        ServiceModel(service_type, formulation, use_service_name=true),
    )
    return
end

"""
Sets the service model in a template using a ServiceModel instance.
"""
function set_service_model!(
    template::ProblemTemplate,
    service_type::Type{<:PSY.Service},
    formulation::Type{<:AbstractServiceFormulation},
)
    set_service_model!(template, ServiceModel(service_type, formulation))
    return
end

function set_service_model!(
    template::ProblemTemplate,
    service_name::String,
    model::ServiceModel{<:PSY.Service, <:AbstractServiceFormulation},
)
    _set_model!(template.services, service_name, model)
    return
end

function set_service_model!(
    template::ProblemTemplate,
    model::ServiceModel{<:PSY.Service, <:AbstractServiceFormulation},
)
    _set_model!(template.services, model)
    return
end


# In addition to illustrating transformations, we'll start loading some
# shaders from a file.  We'll also demo using Reactive for animations.
using ModernGL, GeometryTypes, GLAbstraction, GLWindow, Images, FileIO, Reactive

# Load our textures. See "downloads.jl" to get the images.
kitten = load(GLAbstraction.dir("tutorials", "images", "kitten.png"))
puppy  = load(GLAbstraction.dir("tutorials", "images", "puppy.png"))

# Create the window. This sets all the hints and makes the context current.
window = create_glcontext("Transformations 1", resolution=(600,600))

vao = glGenVertexArrays()
glBindVertexArray(vao)

# The positions of the vertices in our rectangle
vertex_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 texture coordinates of each vertex
vertex_texcoords = Vec2f0[(0, 0),
                          (1, 0),
                          (1, 1),
                          (0, 1)]

# Specify how vertices are arranged into faces
elements = Face{3,UInt32,-1}[(0,1,2),          # the first triangle
                             (2,3,0)]          # the second triangle

vertex_shader = vert"""
#version 150

in vec2 position;
in vec3 color;
in vec2 texcoord;

out vec3 Color;
out vec2 Texcoord;

uniform mat4 trans;

void main()
{
    Color = color;
    Texcoord = texcoord;
    gl_Position = trans*vec4(position, 0.0, 1.0);
}
"""
fragment_shader = load(joinpath(dirname(@__FILE__), "shaders", "puppykitten.frag"))

# Define the rotation matrix (could also use rotationmatrix_z)
# By wrapping it in a Signal, we can easily update it.
trans = Signal(rotate(0f0, Vec((0,0,1f0))))

# Link everything together, using the corresponding shader variable as
# the Dict key
bufferdict = Dict(:position=>Buffer(vertex_positions),
                  :texcoord=>Buffer(vertex_texcoords),
                  :texKitten=>Texture(kitten'),
                  :texPuppy=>Texture(puppy'),
                  :trans=>trans,
                  :indexes=>indexbuffer(elements)) # special for element buffers

ro = std_renderobject(bufferdict,
                      LazyShader(vertex_shader, fragment_shader))
# Do the rendering
glClearColor(0,0,0,1)
while !GLFW.WindowShouldClose(window)
    push!(trans, rotationmatrix_z(time()*deg2rad(180)))
    Reactive.run_till_now()
    glClear(GL_COLOR_BUFFER_BIT)
    GLAbstraction.render(ro)
    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


using Flux, Metalhead, Images, CUDA, NNlib, BSON, Random

CUDA.allowscalar(false)

include("..\\src\\utils.jl")
include("..\\src\\layers.jl")
include("..\\src\\vgg.jl")
include("..\\src\\transformer_net.jl")
include("..\\src\\neural_style.jl")

training_data = "C:\\Users\\michael\\Desktop\\mscoco\\train2017\\train2017"
style_folder = "C:\\Users\\michael\\Desktop\\Julia\\FastStyleTransfer.jl\\images"
style_image = "georges-lemmen_the-beach-at-heist_1891.jpg"
style_path = joinpath(style_folder, style_image)

model_path = "C:\\Users\\michael\\Desktop\\Julia\\FastStyleTransfer.jl\\models"
batch_size = 128
η = 0.01
epochs = 1

content_weight = 1.0
style_weight = 1.0

nimages = 100

train(training_data, batch_size, η, style_path, epochs, model_path, content_weight, style_weight, TransformerNet; images = nimages)


function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, D},
    ::Type{S},
) where {St <: PSY.Storage, D <: AbstractStorageFormulation, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerInVariable, devices, D())
    add_variables!(optimization_container, ActivePowerOutVariable, devices, D())
    add_variables!(optimization_container, ReactivePowerVariable, devices, D())
    add_variables!(optimization_container, EnergyVariable, devices, D())

    # Initial Conditions
    initial_conditions!(optimization_container, devices, D())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, D},
    ::Type{S},
) where {
    St <: PSY.Storage,
    D <: AbstractStorageFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerInVariable, devices, D())
    add_variables!(optimization_container, ActivePowerOutVariable, devices, D())
    add_variables!(optimization_container, EnergyVariable, devices, D())

    # Initial Conditions
    initial_conditions!(optimization_container, devices, D())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, BookKeepingwReservation},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerInVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        ActivePowerOutVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        EnergyVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        ReserveVariable,
        devices,
        BookKeepingwReservation(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, BookKeepingwReservation())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, BookKeepingwReservation},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerInVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        ActivePowerOutVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        EnergyVariable,
        devices,
        BookKeepingwReservation(),
    )
    add_variables!(
        optimization_container,
        ReserveVariable,
        devices,
        BookKeepingwReservation(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, BookKeepingwReservation())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, EnergyTarget},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerInVariable, devices, EnergyTarget())
    add_variables!(optimization_container, ActivePowerOutVariable, devices, EnergyTarget())
    add_variables!(optimization_container, ReactivePowerVariable, devices, EnergyTarget())
    add_variables!(optimization_container, EnergyVariable, devices, EnergyTarget())
    add_variables!(optimization_container, EnergyShortageVariable, devices, EnergyTarget())
    add_variables!(optimization_container, EnergySurplusVariable, devices, EnergyTarget())
    add_variables!(optimization_container, ReserveVariable, devices, EnergyTarget())
    # Initial Conditions
    initial_conditions!(optimization_container, devices, EnergyTarget())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_target_constraint!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, EnergyTarget},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerInVariable, devices, EnergyTarget())
    add_variables!(optimization_container, ActivePowerOutVariable, devices, EnergyTarget())
    add_variables!(optimization_container, EnergyVariable, devices, EnergyTarget())
    add_variables!(optimization_container, EnergyShortageVariable, devices, EnergyTarget())
    add_variables!(optimization_container, EnergySurplusVariable, devices, EnergyTarget())
    add_variables!(optimization_container, ReserveVariable, devices, EnergyTarget())
    # Initial Conditions
    initial_conditions!(optimization_container, devices, EnergyTarget())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_target_constraint!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, BatteryAncillaryServices},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerInVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        ActivePowerOutVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        EnergyVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        ReserveVariable,
        devices,
        BatteryAncillaryServices(),
    )
    # Initial Conditions
    initial_conditions!(optimization_container, devices, BatteryAncillaryServices())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    reserve_contribution_constraint!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{St, BatteryAncillaryServices},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(St, sys)

    if !validate_available_devices(St, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerInVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        ActivePowerOutVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        EnergyVariable,
        devices,
        BatteryAncillaryServices(),
    )
    add_variables!(
        optimization_container,
        ReserveVariable,
        devices,
        BatteryAncillaryServices(),
    )
    # Initial Conditions
    initial_conditions!(optimization_container, devices, BatteryAncillaryServices())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Energy Balanace limits
    add_constraints!(
        optimization_container,
        EnergyBalanceConstraint,
        EnergyVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    reserve_contribution_constraint!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )

    return
end


using LinearAlgebra
using Distributions
using Revise
using Random
using Printf
using StatsBase
using StaticArrays
using POMDPs
using BeliefUpdaters
using AutoViz
using AutomotiveDrivingModels
using AutomotivePOMDPs
using POMDPSimulators
using POMDPPolicies
using DeepQLearning
using Flux
using RLInterface
using POMDPModelTools
using TensorBoardLogger
using Distributions
# using Interact
# using Blink
using MCTS
using Reel

# LOGGER = Logger("log1", overwrite=true)
includet("environment.jl")
includet("generative_mdp.jl")
includet("masking.jl")
includet("cooperative_IDM.jl")
includet("belief_updater.jl")
includet("overlays.jl")
includet("rendering.jl")
includet("make_gif.jl")

rng = MersenneTwister(1)

# mdp = GenerativeMergingMDP(n_cars_main=5, observe_cooperation=false)

function test_state(s,v=0.0, acc=0.0,)
    ego = Vehicle(vehicle_state(35.0, merge_lane(mdp.env), 5.0, mdp.env.roadway), VehicleDef(), EGO_ID)
    veh1 = Vehicle(vehicle_state(s, main_lane(mdp.env), 4.9, mdp.env.roadway), VehicleDef(), EGO_ID + 1)
    veh2 = Vehicle(vehicle_state(s-15, main_lane(mdp.env), 4.9, mdp.env.roadway), VehicleDef(), EGO_ID + 2)
    scene = Scene()
    push!(scene, ego)
    push!(scene, veh1)
    push!(scene, veh2)
    return AugScene(scene, (acc=acc,))
end

rng = MersenneTwister(1)

mdp = GenerativeMergingMDP(random_n_cars=true, min_cars=10, max_cars=14, observe_cooperation=true)

mdp.driver_models[2] = CooperativeIDM(c=0.0)
set_desired_speed!(mdp.driver_models[2], 5.0)

s0 = test_state(85.0)

up = MergingUpdater(mdp)

s0 = initialstate(mdp, rng)
b0 = initialize_belief(up, s0)

super_render(mdp, s0, b0)
s = s0
b = b0
sp, r = generate_sr(mdp, s, 7, rng)
bp = update(up, b, 7, sp)
super_render(mdp, sp, bp)

s = sp
b = deepcopy(bp)
@show b.driver_types
a = 7
sp, r = generate_sr(mdp, s, a, rng)
bp = update(up, b, a, sp)
super_render(mdp, sp, bp)

pomdp = FullyObservablePOMDP(mdp)
up = MergingUpdater(mdp)
s0 = initialstate(mdp, rng)
b0 = initialize_belief(up, s0)
hr = HistoryRecorder(rng=rng, max_steps = 100)

hist= simulate(hr, pomdp, policy, up, b0, s0);

include("visualizer.jl")

includet("make_gif.jl")

make_gif(hist, mdp, hist.belief_hist)

function quick_evaluation(pomdp::POMDP, policy::Policy, rng::AbstractRNG, n_eval=1000)
    avg_r, avg_dr, c_rate, avg_steps, t_out = 0.0, 0.0, 0.0, 0.0, 0.0
    @showprogress for i=1:n_eval
        s0 = initialstate(pomdp, rng)
        up = MergingUpdater(pomdp.mdp)
        hr = HistoryRecorder(rng = rng, max_steps=100)
        hist = simulate(hr, pomdp, policy, up, b0, s0)
        avg_r += undiscounted_reward(hist)
        avg_dr += discounted_reward(hist)
        c_rate += undiscounted_reward(hist) <= pomdp.mdp.collision_cost
        t_out += n_steps(hist) >= hr.max_steps ? 1.0 : 0.0
        avg_steps += n_steps(hist)
    end
    avg_r /= n_eval
    avg_dr /= n_eval
    c_rate /= n_eval
    avg_steps /= n_eval
    t_out /= n_eval
    return avg_r, avg_dr, c_rate, avg_steps, t_out
end
policy = NNPolicy(mdp, policy.qnetwork, policy.action_map, policy.n_input_dims)
avg_r, avg_dr, c_rate, avg_steps, t_out = quick_evaluation(pomdp, policy, rng, 1000)

println("Collisions ", c_rate*100)
println("Avg steps ", avg_steps)
println("Time outs ", t_out*100)
println("avg disc reward ", avg_dr)
println("avg reward ", avg_r)

spbackup = deepcopy(sp)
sbackup = deepcopy(s)
bbackup = deepcopy(b)

neighbor = :fore
o = sp


"""
    usleep(us)

Sleep for `us` microseconds. A wrapper around the C `usleep` function in `unistd.h`.
"""
function usleep(us)
    ccall((:MicroSleep, libhg), Cvoid, (Cint,), us)
end

"""
    LoopRateLimiter

Runs a loop at a fixed rate. Works best for loops where the runtime is approximately
the same every iteration. The loop runtime is kept approximately constant by sleeping
for any time not used by the core computation. This is useful for situations where
the computation should take place a regular, predictable intervals.

To achieve better accuracy, the rate limiter records the error between the expected
runtime and actual runtime every `N_batch` iterations, and adjusts the sleep time
by the average. Unlike the standard sleep function in Julia, this limiter has a
minimum sleep time of 1 microsecond, and rates above 1000Hz can be achieved with
moderate accuracy.

# Example
```
lrl = LoopRateLimiter(100)  # Hz
reset!(lrl)
for i = 1:100
    startloop(lrl)          # start timing the loop
    myexpensivefunction()   # execute the core body of the loop
    sleep(lrl)              # sleep for the rest of the time
end
```
"""
mutable struct LoopRateLimiter{UseSleep}
    rate::Float64    # target rate (Hz)
    ns_target::Int   # target loop time (ns)
    N_batch::Int     # number of loops before estimating the timing offset
    us_offset::Int   # sleep offset, estimated online (us)
    t_batch::UInt64  # start time of last batch (ns)
    t_start::UInt64  # start time of the loop (ns)
    i_batch::Int
    function LoopRateLimiter{UseSleep}(rate; N_batch = 10) where {UseSleep}
        @assert UseSleep isa Bool
        ns_target = round(Int, 1 / rate * 1_000_000_000)
        new{UseSleep}(rate, ns_target, N_batch, 0, 0, 0, 0)
    end
end
function LoopRateLimiter(rate; N_batch = 10)
    use_sleep = rate < 50
    LoopRateLimiter{use_sleep}(rate, N_batch = N_batch)
end

doesblock(::LoopRateLimiter{UseSleep}) where {UseSleep} = !UseSleep

"""
    reset!(::LoopRateLimiter)

Reset the loop rate limiter before a loop. Not necessary if the object is created directly
before calling the loop.
"""
function reset!(lrl::LoopRateLimiter; all::Bool = true)
    lrl.i_batch = 0
end

"""
    startloop(::LoopRateLimiter)

Call this function at the beginning of a loop body to start timing the loop.
"""
function startloop(lrl::LoopRateLimiter)
    lrl.t_start = time_ns()
    us2ns = 1000
    if lrl.i_batch == 0  # very first time the loop is called
        lrl.t_batch = lrl.t_start
    end
    if lrl.i_batch == lrl.N_batch
        t_batch_expected = Int(lrl.ns_target * lrl.N_batch)
        t_batch_actual = Int(lrl.t_start - lrl.t_batch)
        us_error = (t_batch_actual - t_batch_expected) ÷ us2ns ÷ lrl.N_batch
        lrl.us_offset += us_error  # update the estimate of the offset
        lrl.t_batch = lrl.t_start
        lrl.i_batch = 1            # reset the batch loop count
    else
        lrl.i_batch += 1           # increment the batch loop count
    end
end

"""
    sleep(::LoopRateLimiter)

Sleep the OS for the amount of time needed to achieve the rate specified by the loop rate
limiter. Has a minimum sleep time of 1 microsecond (relies on the `usleep` C function).
"""
@generated function Base.sleep(lrl::LoopRateLimiter{UseSleep}) where {UseSleep}
    local sleep_code
    if UseSleep
        sleep_code = quote
            sleep(us_diff / 1e6)
            # usleep(us_diff)
        end
    else
        sleep_code = quote
            usleep(us_diff)
        end
    end
    quote
        us2ns = 1000
        t_stop = time_ns()
        ns_elapsed = round(Int, t_stop - lrl.t_start)
        ns_diff = lrl.ns_target - ns_elapsed                   # time error (ns)
        us_diff = round(Int, ns_diff ÷ us2ns) - lrl.us_offset  # sleep time (us)
        if (us_diff > 0)
            $sleep_code
            # if lrl.use_sleep
            #     sleep(us_diff / 1e6)
            # else
            # tsleep = @async usleep(us_diff)
            # while (!istaskdone(tsleep))
            # end
            # usleep(us_diff)
            # yield()
            # wait(tsleep)
            # end
            # sleep(us_diff / 1000.0)
        end
    end
end

"""
    @rate

Run a loop at a fixed rate, specified either by an integer literal or a
`LoopRateLimiter` object. It will run the loop so that it executes close
to `rate` iterations per second.

# Examples
```
@rate for i = 1:100
    myexpensivefunction()
end 200#Hz
```

```
lr = LoopRateLimiter(200, N_batch=10)
@rate while i < 100
    myexpensivefunction()
    i += 1
end lr
```

Note that the following does NOT work:
```
rate = 100
@rate for i = 1:100
    myexpensivefunction()
end rate
```
Since Julia macros dispatch on the compile-time types instead of the run-time types.
"""
macro rate(loop, lrl)

    # Make sure it's actually a loop
    if loop.head ∉ (:for, :while)
        error("@atrate macro can only be used before for or while loops.")
    end

    # Modify the loop body
    bodyblock = loop.args[2]
    newbody = quote
        startloop($(esc(lrl)))
        $(esc(bodyblock))
        sleep($(esc(lrl)))
    end
    loop.args[1] = esc(loop.args[1])
    loop.args[2] = newbody

    return loop
end

macro rate(loop, rate::Real)

    # Make sure it's actually a loop
    if loop.head ∉ (:for, :while)
        error("@atrate macro can only be used before for or while loops.")
    end

    # Modify the loop body
    bodyblock = loop.args[2]
    newbody = quote
        startloop(lrl)
        $(esc(bodyblock))
        sleep(lrl)
    end
    loop.args[1] = esc(loop.args[1])
    loop.args[2] = newbody

    expr = quote
        lrl = LoopRateLimiter($rate)
        $loop
    end

    return expr
end


__precompile__(true)

module Images

VERSION >= v"0.4.0" && import Base.take
import Base.Order: Ordering, ForwardOrdering, ReverseOrdering
import Base: ==, .==, +, -, *, /, .+, .-, .*, ./, .^, .<, .>
import Base: atan2, clamp, convert, copy, copy!, ctranspose, delete!, done,
             eltype, fft, float32, float64, get, getindex, haskey, hypot,
             ifft, length, linearindexing, map, map!, maximum, mimewritable,
             minimum, next, ndims, one, parent, permutedims, reinterpret, resize!,
             setindex!, show, showcompact, similar, size, slice, sqrt, squeeze,
             start, strides, sub, sum, write, writemime, zero
# "deprecated imports" are below

using Colors, ColorVectorSpace, FixedPointNumbers, FileIO
import Colors: Fractional, red, green, blue
typealias AbstractGray{T}                    Color{T,1}
typealias TransparentRGB{C<:AbstractRGB,T}   TransparentColor{C,T,4}
typealias TransparentGray{C<:AbstractGray,T} TransparentColor{C,T,2}
using Graphics
import Graphics: width, height, Point
import FixedPointNumbers: ufixed8, ufixed10, ufixed12, ufixed14, ufixed16

using Base.Cartesian
include("compatibility/forcartesian.jl")

# if isdefined(module_parent(Images), :Grid)
#     import ..Grid.restrict
# end

const is_little_endian = ENDIAN_BOM == 0x04030201
immutable TypeConst{N} end  # for passing compile-time constants to functions

include("core.jl")
include("map.jl")
include("overlays.jl")
include("labeledarrays.jl")
include("algorithms.jl")
include("connected.jl")
include("edge.jl")
include("writemime.jl")
include("corner.jl")
include("distances.jl")

function precompile()
    for T in (UInt8, UInt16, Int, Float32, Float64)
        Tdiv = typeof(one(T)/2)
        for N = 2:3
            precompile(restrict!, (Array{Tdiv, N}, Array{T,N}, Int))
            precompile(imfilter, (Array{T,N}, Array{Float64,N}))
            precompile(imfilter, (Array{T,N}, Array{Float64,2}))
            precompile(imfilter, (Array{T,N}, Array{Float32,N}))
            precompile(imfilter, (Array{T,N}, Array{Float32,2}))
        end
    end
    for T in (Float32, Float64)
        for N = 2:3
            precompile(_imfilter_gaussian!, (Array{T,N}, Vector{Float64}))
            precompile(_imfilter_gaussian!, (Array{T,N}, Vector{Int}))
            precompile(imfilter_gaussian_no_nans!, (Array{T,N}, Vector{Float64}))
            precompile(imfilter_gaussian_no_nans!, (Array{T,N}, Vector{Int}))
            precompile(fft, (Array{T,N},))
        end
    end
    for T in (Complex{Float32}, Complex{Float64})
        for N = 2:3
            precompile(ifft, (Array{T,N},))
        end
    end
end

export # types
    AbstractImage,
    AbstractImageDirect,
    AbstractImageIndexed,
    Image,
    ImageCmap,
    BitShift,
    ClampMin,
    ClampMax,
    ClampMinMax,
    Clamp,
    Clamp01NaN,
    LabeledArray,
    MapInfo,
    MapNone,
    Overlay,
    OverlayImage,
    ScaleAutoMinMax,
    ScaleMinMax,
    ScaleMinMaxNaN,
    ScaleSigned,
    SliceData,

    # macros
    @test_approx_eq_sigma_eps,

    # constants
    palette_fire,
    palette_gray32,
    palette_gray64,
    palette_rainbow,

    # core functions
    assert2d,
    assert_scalar_color,
    assert_timedim_last,
    assert_xfirst,
    assert_yfirst,
    colordim,
    colorspace,
    coords,
    coords_spatial,
    copyproperties,
    data,
    dimindex,
    dimindexes,
    getindexim,
    grayim,
    colorim,
    height,
    isdirect,
    isxfirst,
    isyfirst,
    limits,
    maxabsfinite,
    maxfinite,
    minfinite,
    nchannels,
    ncolorelem,
    nimages,
    pixelspacing,
    properties,
    rerange!,
    reslice!,
    restrict,
    sdims,
    size_spatial,
    shareproperties,
    sliceim,
    spacedirections,
    spatialorder,
    spatialpermutation,
    spatialproperties,
    storageorder,
    subim,
    timedim,
    width,
    widthheight,
    raw,

    # iterator functions
    first_index,
    iterate_spatial,
    parent,

    # color-related functions
    imadjustintensity,
    indexedcolor,
    lut,
    separate,
    uint32color,
    uint32color!,

    # Scaling of intensity
    sc,
    scale,
    mapinfo,
    uint8sc,
    uint16sc,
    uint32sc,
    ufixed8sc,
    ufixedsc,

    # flip dimensions
    flipx,
    flipy,
    flipz,

    # algorithms
    ando3,
    ando4,
    ando5,
    backdiffx,
    backdiffy,
    dilate,
    erode,
    extrema_filter,
    opening,
    closing,
    tophat,
    bothat,
    morphogradient,
    morpholaplace,
    forwarddiffx,
    forwarddiffy,
    gaussian2d,
    imaverage,
    imcomplement,
    imcorner,
    imdog,
    imedge,
    imfilter,
    imfilter!,
    imfilter_fft,
    imfilter_gaussian,
    imfilter_gaussian!,
    imfilter_LoG,
    blob_LoG,
    findlocalmaxima,
    findlocalminima,
    imgaussiannoise,
    imgradients,
    imlaplacian,
    imlineardiffusion,
    imlog,
    imROF,
    imstretch,
#     imthresh,
    label_components,
    label_components!,
    component_boxes,
    component_lengths,
    component_indices,
    component_subscripts,
    component_centroids,
    magnitude,
    magnitude_phase,
    meanfinite,
    entropy,
    ncc,
    orientation,
    padarray,
    phase,
    prewitt,
    sad,
    sadn,
    sobel,
    ssd,
    ssdn,
    thin_edges,
    thin_edges_subpix,
    thin_edges_nonmaxsup,
    thin_edges_nonmaxsup_subpix,

    # distances
    hausdorff_distance,

    # phantoms
    shepp_logan

"""
`Images` is a package for representing and processing images.

Constructors, conversions, and traits:

    - Construction: `Image`, `ImageCmap`, `grayim`, `colorim`, `convert`, `copyproperties`, `shareproperties`
    - Traits: `colordim`, `colorspace`, `coords_spatial`, `data`, `isdirect`, `isxfirst`, `isyfirst`, `pixelspacing`, `properties`, `sdims`, `spacedirections`, `spatialorder`, `storageorder`, `timedim`
    - Size-related traits: `height`, `nchannels`, `ncolorelem`, `nimages`, `size_spatial`, `width`, `widthheight`
    - Trait assertions: `assert_2d`, `assert_scalar_color`, `assert_timedim_last`, `assert_xfirst`, `assert_yfirst`
    - Indexing operations: `getindexim`, `sliceim`, `subim`
    - Conversions: `convert`, `raw`, `reinterpret`, `separate`

Contrast/coloration:

    - `MapInfo`: `MapNone`, `BitShift`, `ClampMinMax`, `ScaleMinMax`, `ScaleAutoMinMax`, etc.
    - `imadjustintensity`, `sc`, `imstretch`, `imcomplement`

Algorithms:

    - Reductions: `maxfinite`, `maxabsfinite`, `minfinite`, `meanfinite`, `sad`, `ssd`
    - Resizing: `restrict`, `imresize` (not yet exported)
    - Filtering: `imfilter`, `imfilter_fft`, `imfilter_gaussian`, `imfilter_LoG`, `imROF`, `ncc`, `padarray`
    - Filtering kernels: `ando[345]`, `guassian2d`, `imaverage`, `imdog`, `imlaplacian`, `prewitt`, `sobel`
    - Gradients: `backdiffx`, `backdiffy`, `forwarddiffx`, `forwarddiffy`, `imgradients`
    - Edge detection: `imedge`, `imgradients`, `thin_edges`, `magnitude`, `phase`, `magnitudephase`, `orientation`
    - Corner detection: `imcorner`
    - Blob detection: `blob_LoG`, `findlocalmaxima`, `findlocalminima`
    - Morphological operations: `dilate`, `erode`, `closing`, `opening`, `tophat`, `bothat`, `morphogradient`, `morpholaplace`
    - Connected components: `label_components`, `component_boxes`, `component_lengths`, `component_indices`, `component_subscripts`, `component_centroids`

Test images and phantoms (see also TestImages.jl):

    - `shepp_logan`
"""
Images

import FileIO: load, save
@deprecate imread(filename; kwargs...) load(filename; kwargs...)
@deprecate imwrite(img, filename; kwargs...) save(filename, img; kwargs...)
export load, save

function limits(img)
    Base.depwarn("limits is deprecated, all limits are (0,1)", :limits)
    oldlimits(img)
end

end


using InstrumentOperator
using NCDatasets
using OrderedCollections
using YAML

# Load L1 file (replace with your example!):
L1File = "/net/fluo/data1/group/oco2/L1bSc/oco2_L1bScND_26780a_190715_B10003r_200429212407.h5"
ocoData = Dataset(L1File);

# Load corresponding met file:
metFile = "/net/fluo/data1/group/oco2/L2Met/oco2_L2MetND_26780a_190715_B10003r_200429212406.h5"
metData = Dataset(metFile);

# Load dictionary:
dictOCO2 = YAML.load_file("json/oco2.yaml"; dicttype=OrderedDict{String,Any});

# Load L1 file (could just use filenames here as well)
oco = InstrumentOperator.load_L1(dictOCO2,ocoData, metData);

# Pick some bands as tuple (or just one)
bands = (1,2);
# Indices within that band:
indices = (1:1016,20:1000);
# Geo Index (footprint,sounding):
GeoInd = [5,3000];

# Get data for that sounding:
oco_sounding = InstrumentOperator.getMeasurement(oco, bands, indices, GeoInd)


@inline cos(x::Float32) = ccall("extern @air.cos.f32", llvmcall, Cfloat, (Cfloat,), x)
@inline cos_fast(x::Float32) = ccall("extern @air.fast_cos.f32", llvmcall, Cfloat, (Cfloat,), x)

@inline sin(x::Float32) = ccall("extern @air.sin.f32", llvmcall, Cfloat, (Cfloat,), x)
@inline sin_fast(x::Float32) = ccall("extern @air.fast_sin.f32", llvmcall, Cfloat, (Cfloat,), x)

@inline exp(x::Float32) = ccall("extern @air.exp.f32", llvmcall, Cfloat, (Cfloat,), x)
@inline exp_fast(x::Float32) = ccall("extern @air.fast_exp.f32", llvmcall, Cfloat, (Cfloat,), x)


# sample script to reproduce
# ERROR: fatal error in type inference (type bound)

# it seems related to Union{mseSplit,sseSplit}
# (search the code for that expression)
# I am quite sure that this did not happen on 0.7beta2

using Distributed
using Pkg
pkg"develop https://github.com/kafisatz/DecisionTrees.jl#master"
pkg"activate DecisionTrees.jl"
Distributed.@everywhere using DecisionTrees

# if the above does not work, maybe you need to clone the git
# git clone --depth=50 --branch=master https://github.com/kafisatz/DecisionTrees.jl.git
# cd kafisatz/DecisionTrees.jl

@assert VERSION >= v"0.7.0-rc2.0"

Distributed.@everywhere import CSV
Distributed.@everywhere import DataFrames
Distributed.@everywhere import DataFrames: DataFrame

##############################
# Read the data
##############################

dr = "C:\\Users\\bernhard.konig\\Documents\\ASync\\home\\Code\\Julia\\DecisionTrees.jl"
isdir(dr) && cd(dr)
pwd()

@warn("""This assumes pwd() == Pkg.dir("DecisionTrees.jl")""")
datafile = joinpath("data", "freMTPL2", "freMTPL2.csv")
if Sys.islinux()
    # trying to unzip the file...
    pathToZip = string(splitext(datafile)[1], ".zip")
    unzipLoc = splitdir(pathToZip)[1]
    cmd = `unzip $(pathToZip) -d $(unzipLoc)`
    run(cmd)
end

@assert isfile(datafile) "You need to unzip the file freMTPL2.zip in the folder data/freMTPL2"

@time fullData = CSV.read(datafile, rows_for_type_detect=100000, allowmissing=:none, categorical=false);

# some data manipulation
areasSorted = sort(unique(fullData[:Area]))
AreaInteger = map(x->findall((in)([x]), areasSorted)[1], fullData[:Area])
fullData[:AreaInteger] = AreaInteger
    # correct for unreasonable observations
for i = 1:size(fullData, 1)
    fullData[:ClaimNb][i] = min(4, fullData[:ClaimNb][i])
    fullData[:Exposure][i] = min(1.0, fullData[:Exposure][i])
end

# set independent variables
selected_explanatory_vars = ["Area","AreaInteger","VehPower","VehAge","DrivAge","BonusMalus","VehBrand","VehGas","Density","Region"]

##############################
# Prepare the data
##############################
dtmtable, sett, dfprepped = prepare_dataframe_for_dtm!(fullData, keycol="IDpol", trnvalcol="trnTest", numcol="ClaimNb", denomcol="Exposure", weightcol="Exposure", independent_vars=selected_explanatory_vars);

# run a few models

# this runs as expected
updateSettings!(sett, minWeight=-0.03, model_type="build_tree", calculatePoissonError=true)
resultingFiles, resM = dtm(dtmtable, sett)

# this runs as expected
updateSettings!(sett, crit="sse", model_type="build_tree")
resultingFiles, resM = dtm(dtmtable, sett)

# this one fails!
updateSettings!(sett, crit="sse", model_type="boosted_tree")
resultingFiles, resM = dtm(dtmtable, sett)

#=
# this is the error I get

"""
ERROR: fatal error in type inference (type bound)
Stacktrace:
 [1] _typed_hcat(::Type{Float64}, ::Tuple{Array{Int64,1},Array{Int64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1}}) at .\abstractarray.jl:1255
 [2] typed_hcat at .\abstractarray.jl:1243 [inlined]
 [3] hcat at C:\cygwin\home\Administrator\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v0.7\SparseArrays\src\sparsevector.jl:1063 [inlined]
 [4] macro expansion at .\logging.jl:307 [inlined]
 [5] macro expansion at .\util.jl:156 [inlined]
 [6] macro expansion at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\boosting.jl:89 [inlined]
 [7] macro expansion at .\util.jl:156 [inlined]
 [8] macro expansion at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\helper_functions.jl:1469 [inlined]
 [9] boosted_tree(::DecisionTrees.DTMTable, ::ModelSettings) at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\boosting.jl:82
 [10] run_model_actual(::DecisionTrees.DTMTable, ::ModelSettings, ::String) at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\runmodel.jl:467
 [11] #dtm#110 at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\runmodel.jl:798 [inlined]
 [12] dtm(::DecisionTrees.DTMTable, ::ModelSettings) at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\runmodel.jl:798
 [13] top-level scope at none:0
 """ =#


@testset "controller" begin
    @testset "options" begin
        @system SControllerOptions(Controller) begin
            a ~ preserve(extern)
            b ~ ::int(override)
        end
        o = (a=1, b=2)
        s = instance(SControllerOptions; options=o)
        @test s.a' == 1
        @test s.b == 2
    end

    @testset "config placeholder" begin
        @system SControllerConfigPlaceholder(Controller) begin
            a => 1 ~ preserve(parameter)
        end
        o = 0 => :a => 2
        s = instance(SControllerConfigPlaceholder; config=o)
        @test s.a' == 2
    end

    @testset "seed" begin
        @system SControllerSeed(Controller) begin
            a => 0 ± 1 ~ preserve
        end
        s1 = instance(SControllerSeed; seed=0)
        @test s1.a' == 0.6791074260357777
        s2 = instance(SControllerSeed; seed=0)
        @test s1.a' == s2.a'
    end

    @testset "no seed" begin
        @system SControllerNoSeed(Controller) begin
            a => 0 ± 1 ~ preserve
        end
        s1 = instance(SControllerNoSeed; seed=nothing)
        s2 = instance(SControllerNoSeed; seed=nothing)
        @test s1.a' != s2.a'
    end
end


@testset "DeviceModel Tests" begin
    @test_throws ArgumentError DeviceModel(PSY.ThermalGen, PSI.ThermalStandardUnitCommitment)
    @test_throws ArgumentError DeviceModel(PSY.ThermalStandard, PSI.AbstractDeviceFormulation)
    @test_throws ArgumentError PSI.CanonicalModel(JuMP.Model(),
                                                nothing,
                                                false,
                                                true,
                                                1:24,
                                                Dates.Hour(1),
                                                Dates.now(),
                                                Dict{Symbol, JuMP.Containers.DenseAxisArray}(),
                                                Dict{Symbol, JuMP.Containers.DenseAxisArray}(),
                                                JuMP.AffExpr(0),
                                                Dict{Symbol, JuMP.Containers.DenseAxisArray}(),
                                                nothing,
                                                Dict{PSI.ICKey, Array{InitialCondition}}(),
                                                nothing)
end

@testset "OperationModel Tests" begin
    for p in [true, false]
        t = OperationModel(TestOptModel, CopperPlatePowerModel, c_sys5; parameters = p)
        @test PSI.model_has_parameters(t.canonical) == p
    end
end


# export PCCBS
#
# module PCCBS
#
# using CRCBS
# using Parameters, LightGraphs, DataStructures
# using GraphUtils
# using ..PlanningPredicates

################################################################################
############################### ENVIRONMENT DEF ################################
################################################################################
const State = CRCBS.GraphState
# struct State{N} <: CRCBS.GraphState
#     vtx::Int
#     t::Int
#     node::N
# end
const Action = CRCBS.GraphAction

export PCCBSEnv

# @with_kw_noshow struct PCCBSEnv{E,N,I,T,C<:AbstractCostModel,H<:AbstractCostModel} <: GraphEnv{State{N},Action,C}
@with_kw_noshow struct PCCBSEnv{E,N,I,T,C<:AbstractCostModel,H<:AbstractCostModel} <: GraphEnv{State,Action,C}
    search_env::E                   = nothing
    schedule_node::N                = nothing
    node_id::I                      = nothing
    agent_idx::Int                  = -1
    # graph::G                        = get_graph(search_env)
    constraints::T                  = discrete_constraint_table(search_env,agent_idx)
    goal::State                     = State()
    cost_model::C                   = get_cost_model(search_env)
    heuristic::H                    = get_heuristic_model(search_env)
end
CRCBS.get_cost_model(env::PCCBSEnv)       = get_cost_model(env.search_env)
CRCBS.get_heuristic_model(env::PCCBSEnv)  = get_heuristic_model(env.search_env)
GraphUtils.get_node(env::PCCBSEnv)        = env.schedule_node
GraphUtils.get_graph(env::PCCBSEnv)       = get_graph(env.search_env,graph_key(get_node(env))) #graph

function Base.show(io::IO, env::PCCBSEnv)
    print(io,"PCCBSEnv: \n",
        "\t","schedule_node: ",string(get_node(env)),"\n",
        "\t","agent_idx:     ",get_agent_id(env),"\n",
        "\t","goal:          ",string(get_goal(env)),"\n")
end

CRCBS.get_next_state(s::State,a::Action)    = State(get_e(a).dst,get_t(s)+get_dt(a))
CRCBS.get_next_state(env::PCCBSEnv,s,a)  = get_next_state(s,a)
# CRCBS.get_next_state(s::State,a::Action)    = State(get_e(a).dst,get_t(s)+get_dt(a),s.node)
# CRCBS.get_next_state(env::PCCBSEnv,s,a)  = State(get_e(a).dst,get_t(s)+get_dt(a),env.schedule_node)
CRCBS.wait(env::PCCBSEnv,s)              = Action(e=Edge(get_vtx(s),get_vtx(s)))

CRCBS.get_possible_actions(env::PCCBSEnv,s::State) = get_possible_actions(get_node(env),env,s)
function CRCBS.get_possible_actions(node,env::PCCBSEnv,s::State)
    if CRCBS.is_valid(env,s)
        return map(v2->Action(e=Edge(s.vtx,v2)), outneighbors(get_graph(env),s.vtx))
    end
    return Action[]
end
function CRCBS.get_possible_actions(node::Union{COLLECT,DEPOSIT},env::PCCBSEnv,s::State)
    # @assert get_vtx(s) == get_id(get_initial_location_id(node)) "Robot $(get_agent_id(env)) should not be at s=$(string(s)) when task=$(string(node))"
    if !(get_vtx(s) == get_id(get_initial_location_id(node)))
        @log_info(-1,0,"Robot $(get_agent_id(env)) should not be at s=$(string(s)) when task=$(string(node))")
    end
    if CRCBS.is_valid(env,s)
        return [CRCBS.wait(env,s)]
    end
    return Action[]
end
function CRCBS.get_possible_actions(env::MetaAgentCBS.TeamMetaEnv,s::MetaAgentCBS.State{State})
    d_set = Set{Tuple}([(0,0),(-1,0),(0,1),(1,0),(0,-1)])
    for (e,s) in zip(env.envs, s.states)
        intersect!(d_set, CRCBS.get_graph(e).edge_cache[s.vtx])
    end
    meta_actions = Vector{MetaAgentCBS.Action{Action}}()
    for d in d_set
        meta_action = MetaAgentCBS.Action{Action}()
        valid_action = true
        for (e,s) in zip(env.envs, s.states)
            vtx = get_graph(e).vtxs[s.vtx]
            next_vtx = (vtx[1] + d[1], vtx[2] + d[2])
            edge = Edge(get_graph(e).vtx_map[vtx...],get_graph(e).vtx_map[next_vtx...])
            if has_edge(get_graph(e),edge)
                push!(meta_action.actions, Action(e = edge))
            else
                valid_action = false
                break
            end
        end
        if valid_action
            push!(meta_actions, meta_action)
        end
    end
    meta_actions
end

function CRCBS.is_goal(env::PCCBSEnv,s)
    if get_t(s) >= get_t(get_goal(env))
        if states_match(s, get_goal(env))
            #########################################
            # updated goal time
            v = get_vtx(get_schedule(env.search_env),env.node_id)
            path_spec = get_path_spec(get_schedule(env.search_env), v)
            # goal_time = get_tF(env.search_env,v) # time after which goal can be satisfied
            goal_time = get_t(get_goal(env))
            if path_spec.tight == true
                goal_time = get_tF(env.search_env,v) + minimum(get_local_slack(env.search_env,v))
                # goal_time += minimum(get_local_slack(env.search_env,v))
            end
            if (path_spec.free == true) && is_terminal_node(get_graph(get_schedule(env.search_env)),v)
                goal_time = maximum(get_cache(env.search_env).tF)
            end
            # check that all predecessor nodes have been completed.
            if get_t(s) < goal_time
                return false
            end
            #########################################
            return true
        elseif !CRCBS.is_valid(get_goal(env))
            return true
        end
    end
    return false
end

# end # end module PCCBS


export multiHeadAttnForward
import Base: unsafe_convert
using Knet.KnetArrays: DevArray
using AutoGrad: AutoGrad, @primitive1

using CUDA.CUDNN: handle,
    cudnnMultiHeadAttnForward,
    cudnnMultiHeadAttnBackwardData,
    cudnnMultiHeadAttnBackwardWeights,
    cudnnGetMultiHeadAttnBuffers,
    cudnnGetMultiHeadAttnWeights,
    cudnnAttnDescriptor_t,
        cudnnCreateAttnDescriptor,
        cudnnDestroyAttnDescriptor,
        cudnnSetAttnDescriptor,
        cudnnGetAttnDescriptor,
        cudnnDataType_t,
        cudnnDropoutDescriptor_t,
    cudnnSeqDataDescriptor_t,
        cudnnCreateSeqDataDescriptor,
        cudnnDestroySeqDataDescriptor,
        cudnnSetSeqDataDescriptor,
        cudnnGetSeqDataDescriptor,
    cudnnSeqDataAxis_t,
        CUDNN_SEQDATA_TIME_DIM,  # 0, /* index in time */
        CUDNN_SEQDATA_BATCH_DIM, # 1, /* index in batch */
        CUDNN_SEQDATA_BEAM_DIM,  # 2, /* index in beam */
        CUDNN_SEQDATA_VECT_DIM,  # 3  /* index in vector */
    cudnnAttnQueryMap_t,
        CUDNN_ATTN_QUERYMAP_ALL_TO_ONE, # 0         /* multiple Q-s map to a single (K,V) set when beam size > 1 */
        CUDNN_ATTN_QUERYMAP_ONE_TO_ONE, # (1U << 0) /* multiple Q-s map to multiple (K,V) sets when beam size > 1 */
        CUDNN_ATTN_DISABLE_PROJ_BIASES, # 0         /* no biases in attention input and output projections */
        CUDNN_ATTN_ENABLE_PROJ_BIASES,  # (1U << 1) /* use biases in attention input and output projections */
    cudnnMultiHeadAttnWeightKind_t,
        CUDNN_MH_ATTN_Q_WEIGHTS, # 0, /* input projection weights for 'queries' */
        CUDNN_MH_ATTN_K_WEIGHTS, # 1, /* input projection weights for 'keys' */
        CUDNN_MH_ATTN_V_WEIGHTS, # 2, /* input projection weights for 'values' */
        CUDNN_MH_ATTN_O_WEIGHTS, # 3, /* output projection weights */
        CUDNN_MH_ATTN_Q_BIASES,  # 4, /* input projection bias tensor for 'queries' */
        CUDNN_MH_ATTN_K_BIASES,  # 5, /* input projection bias for 'keys' */
        CUDNN_MH_ATTN_V_BIASES,  # 6, /* input projection bias for 'values' */
        CUDNN_MH_ATTN_O_BIASES,  # 7, /* output projection biases */
    cudnnMathType_t,
        CUDNN_DEFAULT_MATH,                    # 0,
        CUDNN_TENSOR_OP_MATH,                  # 1,
        CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION, # 2,
        CUDNN_FMA_MATH                         # 3,

@primitive1 multiHeadAttnForward(x; o...),dy,y  multiHeadAttnBackwardWeights(x,y,dy; o...)  multiHeadAttnBackwardData(x,y,dy; o...)
@primitive1 multiHeadAttnBackwardData(x,y...;o...)  throw(MethodError(back,multiHeadAttnBackwardData))
@primitive1 multiHeadAttnBackwardWeights(x,y...;o...)  throw(MethodError(back,multiHeadAttnBackwardWeights))

function multiHeadAttnForward
end

function multiHeadAttnBackwardData
end

function multiHeadAttnBackwardWeights
end

mutable struct AttnDescriptor; ptr; end
unsafe_convert(::Type{<:Ptr}, mha::AttnDescriptor)=mha.ptr
function MHA(; # Defaults:
             attnMode::Unsigned = CUDNN_ATTN_QUERYMAP_ALL_TO_ONE | CUDNN_ATTN_DISABLE_PROJ_BIASES,
             nHeads::Integer,
             smScaler::Real = 1.0,
             dataType::DataType,
             computePrec::DataType = dataType, # There doesn't seem to be any other option in cudnn 8.0.2 docs
             mathType::cudnnMathType_t = (dataType === Float16 ? CUDNN_TENSOR_OP_MATH :
                                          dataType === Float32 ? CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION :
                                          CUDNN_DEFAULT_MATH),
             attnDropout::Real = 0.0,
             postDropout::Real = 0.0,

             )
    ptr = cudnnAttnDescriptor_t[C_NULL]
    cudnnCreataAttnDescriptor(ptr)
    cudnnSetAttnDescriptor(ptr[1], mode, nHeads, DT(dataType), DT(computePrec), mathType, DD(attnDropout), DD(postDropout)
                           )
    mha = AttnDescriptor(ptr[1])
    finalizer(x->cudnnDestroyAttnDescriptor(x.ptr), mha)
    return mha
end