abstract type AbstractModelSet end
# using CSV
# using DataFrames

"""
    WindFarmModelSet(wakedeficitmodel, wake_deflection_model, wake_combination_model, local_ti_model)

Container for objects defining models to use in wind farm calculations

# Arguments
- `wake_defiict_model::AbstractWakeDeficitModel`: contains a struct defining the desired wake deficit model
- `wake_deflection_model::AbstractWakeDeflectionModel`: contains a struct defining the desired wake deflection model
- `wake_combination_model::AbstractWakeCombinationModel`: contains a struct defining the desired wake combination model
- `local_ti_model::AbstractTurbulenceIntensityModel`: contains a struct defining the desired turbulence intensity model
"""
struct WindFarmModelSet{DTM,DNM,CM,TIM} <: AbstractModelSet

    wake_deficit_model::DTM
    wake_deflection_model::DNM
    wake_combination_model::CM
    local_ti_model::TIM

end

"""
    point_velocity(loc, turbine_x, turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
    rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, wtvelocities,
    wind_resource, model_set::AbstractModelSet;
    wind_farm_state_id=1, downwind_turbine_id=0)

Calculates the wind speed at a given point for a given state

# Arguments
- `loc::Array{TF,3}`: Location of interest
- `turbine_x::Array{TF,nTurbines}`: turbine east-west locations in the state
    reference frame
- `turbine_y::Array{TF,nTurbines}`: turbine north-south locations in the state
    reference frame
- `turbine_z::Array{TF,nTurbines}`: turbine base height in the state reference frame
- `turbine_yaw::Array{TF,nTurbines}`: turbine yaw for the given wind direction in
    radians
- `turbine_ct::Array{TF,nTurbines}`: turbine thrust coefficients for the given state
- `turbine_ai::Array{TF,nTurbines}`: turbine axial induction for the given state
- `rotor_diameter::Array{TF,nTurbines}`: turbine rotor diameters
- `hub_height::Array{TF,nTurbines}`: turbine hub heights
- `turbine_local_ti::Array{TF,nTurbines}`: turbine local turbulence intensity for
    the given state
- `sorted_turbine_index::Array{TF,nTurbines}`: array containing indices of wind turbines
    from most upwind to most downwind turbine in the given state
- `wtvelocities::Array{TF,nTurbines}`: effective inflow wind speed for given state
- `wind_resource::DiscretizedWindResource`: contains wind resource discreption (directions,
    speeds, frequencies, etc)
- `wind_farm_state_id::Int`: index to correct state to use from wind resource provided.
    Defaults to 1
- `downwind_turbine_id::Int`: index of wind turbine of interest (if any). If not a point for
    calculating effective wind speed of a turbine, then provide 0 (default)
"""
function point_velocity(locx, locy, locz, turbine_x, turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
                    rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, wtvelocities,
                    wind_resource, model_set::AbstractModelSet;
                    wind_farm_state_id=1, downwind_turbine_id=0)

    wakedeficitmodel = model_set.wake_deficit_model
    wakedeflectionmodel = model_set.wake_deflection_model
    wakecombinationmodel = model_set.wake_combination_model

    # extract flow information
    wind_speed = wind_resource.wind_speeds[wind_farm_state_id]
    reference_height = wind_resource.measurement_heights[wind_farm_state_id]

    # set ground height
    ground_height = wind_resource.wind_shear_model.ground_height    # TODO: allow topology to be given

    # find order for wind shear and deficit calculations
    shear_order = wind_resource.wind_shear_model.shear_order

    # adjust wind speed for wind shear
    if shear_order == "nothing"
        wind_speed_internal = wind_speed
    elseif shear_order == "first"
        wind_speed_internal = adjust_for_wind_shear(locz, wind_speed, reference_height, ground_height, wind_resource.wind_shear_model)
    else
        wind_speed_internal = wind_speed
    end

    # get number of turbines
    nturbines = length(turbine_x)

    # initialize deficit summation term to zero
    deficit_sum = 0.0

    # loop through all turbines
    for u=1:nturbines

        # get index of upstream turbine
        upwind_turb_id = Int(sorted_turbine_index[u])

        # don't allow turbine to impact itself
        if upwind_turb_id == downwind_turbine_id; continue; end

        # downstream distance between upstream turbine and point
        x = locx - turbine_x[upwind_turb_id]

        # check turbine relative locations
        if x > 1E-6
            # skip this loop if it would include a turbine's impact on itself)
            if upwind_turb_id==downwind_turbine_id; continue; end

            # calculate wake deflection of the current wake at the point of interest
            horizontal_deflection = wake_deflection_model(locx, locy, locz, turbine_x, turbine_yaw, turbine_ct,
                            upwind_turb_id, rotor_diameter, turbine_local_ti, wakedeflectionmodel)

            vertical_deflection = 0.0

            # velocity difference in the wake
            deltav = wake_deficit_model(locx, locy, locz, turbine_x, turbine_y, turbine_z, horizontal_deflection, vertical_deflection,
                            upwind_turb_id, downwind_turbine_id, hub_height, rotor_diameter, turbine_ai,
                            turbine_local_ti, turbine_ct, turbine_yaw, wakedeficitmodel)

            # combine deficits according to selected wake combination method
            deficit_sum = wake_combination_model(deltav, wind_speed_internal, wtvelocities[upwind_turb_id], deficit_sum, wakecombinationmodel)
            # println(deficit_sum, " ", downwind_turbine_id, " ", upwind_turb_id)

        end
    end

    # find velocity at point without shear
    point_velocity = wind_speed_internal - deficit_sum

    if shear_order == "nothing"
        point_velocity_out = point_velocity
    elseif shear_order == "first"
        point_velocity_out = point_velocity
    else
        point_velocity_out = adjust_for_wind_shear(locz, point_velocity, reference_height, ground_height, wind_resource.wind_shear_model)
    end

    return point_velocity_out

end

"""
    point_velocity(turbine_x, turbine_y, turbine_z, rotor_diameter, hub_height, turbine_yaw,
    sorted_turbine_index, ct_model, rotor_sample_points_y, rotor_sample_points_z, wind_resource,
    model_set::AbstractModelSet; wind_farm_state_id=1)

Calculates the wind speed at a given point for a given state

# Arguments
- `turbine_x::Array{TF,nTurbines}`: turbine east-west locations in the state
    reference frame
- `turbine_y::Array{TF,nTurbines}`: turbine north-south locations in the state
    reference frame
- `turbine_z::Array{TF,nTurbines}`: turbine base height in the state reference frame
- `rotor_diameter::Array{TF,nTurbines}`: turbine rotor diameters
- `hub_height::Array{TF,nTurbines}`: turbine hub heights
- `turbine_yaw::Array{TF,nTurbines}`: turbine yaw for the given wind direction in
    radians
- `sorted_turbine_index::Array{TF,nTurbines}`: turbine sorted order upstream to downstream
    for given state
- `ct_model::AbstractThrustCoefficientModel`: defines how the thrust coefficient changes
    with state etc
- rotor_sample_points_y::Array{TF,N}`: horizontal wind location of points to sample across
    the rotor swept area when calculating the effective wind speed for the wind turbine.
    Points are centered at the hub (0,0) and scaled by the radius (1=tip of blades)
- rotor_sample_points_z::Array{TF,N}`: vertical wind location of points to sample across the
    rotor swept area when calculating the effective wind speed for the wind turbine. Points
    are centered at the hub (0,0) and scaled by the radius (1=tip of blades)
- `wind_resource::DiscretizedWindResource`: wind resource discreption (directions, speeds,
    frequencies, etc)
- `model_set::AbstractModelSet`: defines wake-realated models to be used in analysis
- `wind_farm_state_id::Int`: index to correct state to use from wind resource provided.
    Defaults to 1
"""
function turbine_velocities_one_direction(turbine_x, turbine_y, turbine_z, rotor_diameter, hub_height, turbine_yaw,
                    sorted_turbine_index, ct_model, rotor_sample_points_y, rotor_sample_points_z, wind_resource,
                    model_set::AbstractModelSet; wind_farm_state_id=1, velocity_only=true)

    # get number of turbines and rotor sample point
    n_turbines = length(turbine_x)
    n_rotor_sample_points = length(rotor_sample_points_y)

    # initialize correct array types
    arr_type = promote_type(typeof(turbine_x[1]),typeof(turbine_y[1]),typeof(turbine_z[1]),typeof(rotor_diameter[1]),
                            typeof(hub_height[1]),typeof(turbine_yaw[1]))

    # initialize arrays
    turbine_velocities = zeros(arr_type, n_turbines)
    turbine_ct = zeros(arr_type, n_turbines)
    turbine_ai = zeros(arr_type, n_turbines)
    turbine_local_ti = zeros(arr_type, n_turbines)

    # loop over all turbines
    for d=1:n_turbines

        # get index of downstream turbine
        downwind_turbine_id = Int(sorted_turbine_index[d])

        # initialize downstream wind turbine velocity to zero
        # println("start array: ", turbine_velocities[downwind_turbine_id])
        # wind_turbine_velocity = typeof(turbine_velocities[downwind_turbine_id])(0.0)
        wind_turbine_velocity = 0.0
        # turbine_velocities[downwind_turbine_id] = 0.0

        # loop over all rotor sample points to approximate the effective inflow velocity
        for p=1:n_rotor_sample_points

            # scale rotor sample point coordinate by rotor diameter (in rotor hub ref. frame)
            local_rotor_sample_point_y = rotor_sample_points_y[p]*0.5*rotor_diameter[downwind_turbine_id]
            local_rotor_sample_point_z = rotor_sample_points_z[p]*0.5*rotor_diameter[downwind_turbine_id]

            # put rotor sample points in wind direction coordinate system, and account for yaw
            locx = turbine_x[downwind_turbine_id] .+ local_rotor_sample_point_y*sin(turbine_yaw[downwind_turbine_id])
            locy = turbine_y[downwind_turbine_id] .+ local_rotor_sample_point_y*cos(turbine_yaw[downwind_turbine_id])
            locz = turbine_z[downwind_turbine_id] .+ hub_height[downwind_turbine_id] + local_rotor_sample_point_z

            # calculate the velocity at given point
            point_velocity_with_shear = point_velocity(locx, locy, locz, turbine_x, turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
                                    rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, turbine_velocities,
                                    wind_resource, model_set,
                                    wind_farm_state_id=wind_farm_state_id, downwind_turbine_id=downwind_turbine_id)

            # add sample point velocity to turbine velocity to be averaged later
            wind_turbine_velocity += point_velocity_with_shear

        end

        # final velocity calculation for downstream turbine (average equally across all points)
        wind_turbine_velocity /= n_rotor_sample_points

        turbine_velocities[downwind_turbine_id] = deepcopy(wind_turbine_velocity)

        # update thrust coefficient for downstream turbine
        turbine_ct[downwind_turbine_id] = calculate_ct(turbine_velocities[downwind_turbine_id], ct_model[downwind_turbine_id])

        # update axial induction for downstream turbine
        turbine_ai[downwind_turbine_id] = _ct_to_axial_ind_func(turbine_ct[downwind_turbine_id])

        # get local turbulence intensity for this wind state
        ambient_ti = wind_resource.ambient_tis[wind_farm_state_id]

        # update local turbulence intensity for downstream turbine
        turbine_local_ti[downwind_turbine_id] = calculate_local_ti(turbine_x, turbine_y, ambient_ti, rotor_diameter, hub_height, turbine_yaw, turbine_local_ti, sorted_turbine_index,
                            turbine_velocities, turbine_ct, model_set.local_ti_model; turbine_id=downwind_turbine_id, tol=1E-6)

        # println("local ti turb 9: ", turbine_local_ti[downwind_turbine_id])
    end

    # df = DataFrame(ID=1:n_turbines, V=turbine_velocities, TI=turbine_local_ti, CT=turbine_ct)
    # CSV.write("internaldata.txt", df)

    if velocity_only
        return turbine_velocities
    else
        return turbine_velocities, turbine_ct, turbine_ai, turbine_local_ti
    end
end

function turbine_velocities_one_direction(x, turbine_z, rotor_diameter, hub_height, turbine_yaw,
    sorted_turbine_index, ct_model, rotor_sample_points_y, rotor_sample_points_z, wind_resource,
    model_set::AbstractModelSet; wind_farm_state_id=1, velocity_only=true)

    n_turbines = Int(length(x)/2)
    # println(typeof(x), n_turbines)
    turbine_x = x[1:n_turbines]
    turbine_y = x[n_turbines+1:end]
    # println(turbine_x)
    # println("turbine_x type ", typeof(turbine_x))
    # println("type of x ", typeof(x))

    # get number of turbines and rotor sample point
    # n_turbines = length(turbine_x)
    n_rotor_sample_points = length(rotor_sample_points_y)

    arr_type = promote_type(typeof(turbine_x[1]),typeof(turbine_y[1]),typeof(turbine_z[1]),typeof(rotor_diameter[1]),
                typeof(hub_height[1]),typeof(turbine_yaw[1]))
    turbine_velocities = zeros(arr_type, n_turbines)
    turbine_ct = zeros(arr_type, n_turbines)
    turbine_ai = zeros(arr_type, n_turbines)
    turbine_local_ti = zeros(arr_type, n_turbines)

    for d=1:n_turbines

        # get index of downstream turbine
        downwind_turbine_id = Int(sorted_turbine_index[d])

        # initialize downstream wind turbine velocity to zero
        # println("start array: ", turbine_velocities[downwind_turbine_id])
        # wind_turbine_velocity = typeof(turbine_velocities[downwind_turbine_id])(0.0)
        wind_turbine_velocity = 0.0
        # turbine_velocities[downwind_turbine_id] = 0.0

        for p=1:n_rotor_sample_points

            # scale rotor sample point coordinate by rotor diameter (in rotor hub ref. frame)
            local_rotor_sample_point_y = rotor_sample_points_y[p]*0.5*rotor_diameter[downwind_turbine_id]
            local_rotor_sample_point_z = rotor_sample_points_z[p]*0.5*rotor_diameter[downwind_turbine_id]

            locx = turbine_x[downwind_turbine_id] .+ local_rotor_sample_point_y*sin(turbine_yaw[downwind_turbine_id])
            locy = turbine_y[downwind_turbine_id] .+ local_rotor_sample_point_y*cos(turbine_yaw[downwind_turbine_id])
            locz = turbine_z[downwind_turbine_id] .+ hub_height[downwind_turbine_id] + local_rotor_sample_point_z

            # calculate the velocity at given point
            point_velocity_with_shear = point_velocity(locx, locy, locz, turbine_x, turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
                                rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, turbine_velocities,
                                wind_resource, model_set,
                                wind_farm_state_id=wind_farm_state_id, downwind_turbine_id=downwind_turbine_id)

            # add sample point velocity to turbine velocity to be averaged later
            wind_turbine_velocity += point_velocity_with_shear

        end

        # final velocity calculation for downstream turbine (average equally across all points)
        wind_turbine_velocity /= n_rotor_sample_points

        turbine_velocities[downwind_turbine_id] = deepcopy(wind_turbine_velocity)

        # update thrust coefficient for downstream turbine
        turbine_ct[downwind_turbine_id] = calculate_ct(turbine_velocities[downwind_turbine_id], ct_model[downwind_turbine_id])

        # update axial induction for downstream turbine
        turbine_ai[downwind_turbine_id] = _ct_to_axial_ind_func(turbine_ct[downwind_turbine_id])

        # update local turbulence intensity for downstream turbine
        ambient_ti = wind_resource.ambient_tis[wind_farm_state_id]
        turbine_local_ti[downwind_turbine_id] = calculate_local_ti(turbine_x, turbine_y, ambient_ti, rotor_diameter, hub_height, turbine_yaw, turbine_local_ti, sorted_turbine_index,
                    turbine_velocities, turbine_ct, model_set.local_ti_model; turbine_id=downwind_turbine_id, tol=1E-6)

    end

    if velocity_only
        return turbine_velocities
    else
        return turbine_velocities, turbine_ct, turbine_ai, turbine_local_ti
    end

end

# turbine_velocities_one_direction!(model_set::AbstractModelSet, problem_description::AbstractWindFarmProblem; wind_farm_state_id=1) = turbine_velocities_one_direction!([0.0], [0.0],
# model_set::AbstractModelSet, problem_description::AbstractWindFarmProblem; wind_farm_state_id=1)

"""
calculate_flow_field(xrange, yrange, zrange, model_set::AbstractModelSet, turbine_x,
    turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai, rotor_diameter, hub_height,
    turbine_local_ti, sorted_turbine_index, wtvelocities, wind_resource; wind_farm_state_id=1)

Generates a flow field for a given state and cross section

# Arguments
- `xrange::Range`: range defining east-west locations to sample in global reference frame
- `yrange::Range`: range defining north-west locations to sample in global reference frame
- `zrange::Range`: range defining vertical locations to sample in global reference frame
- `model_set::AbstractModelSet`: defines wake-realated models to be used in analysis
- `turbine_x::Array{TF,nTurbines}`: turbine east-west locations in the global
    reference frame
- `turbine_y::Array{TF,nTurbines}`: turbine north-south locations in the global
    reference frame
- `turbine_z::Array{TF,nTurbines}`: turbine base height in the global reference frame
- `turbine_yaw::Array{TF,nTurbines}`: turbine yaw for the given wind direction in
    radians
- `turbine_ct::Array{TF,nTurbines}`: thrust coefficient of each turbine for the given state
- `turbine_ai::Array{TF,nTurbines}`: turbine axial induction for the given state
- `rotor_diameter::Array{TF,nTurbines}`: turbine rotor diameters
- `hub_height::Array{TF,nTurbines}`: turbine hub heights
- `turbine_local_ti::Array{TF,nTurbines}`: turbine local turbulence intensity for
    the given state
- `sorted_turbine_index::Array{TF,nTurbines}`: turbine north-south locations in the
    global reference frame
- `wtvelocities::Array{TF,nTurbines}`: effective inflow wind speed for given state
- `wind_resource::DiscretizedWindResource`: wind resource discreption (directions, speeds,
    frequencies, etc)
- `wind_farm_state_id::Int`: index to correct state to use from wind resource provided.
    Defaults to 1
"""
function calculate_flow_field(xrange, yrange, zrange,
    model_set::AbstractModelSet, turbine_x, turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
    rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, wtvelocities,
    wind_resource; wind_farm_state_id=1)

    xlen = length(xrange)
    ylen = length(yrange)
    zlen = length(zrange)
    npoints = xlen*ylen*zlen
    point_velocities = zeros(npoints)
    point_velocities = reshape(point_velocities, (zlen, ylen, xlen))

    # rotate to direction frame for velocity calculations
    rot_tx, rot_ty = rotate_to_wind_direction(turbine_x, turbine_y, wind_resource.wind_directions[wind_farm_state_id])

    # sort the turbines
    sorted_turbine_index = sortperm(rot_tx)

    for zi in 1:zlen
        for yi in 1:ylen
            for xi in 1:xlen
                locx = xrange[xi]
                locy = yrange[yi]
                locz = zrange[zi]
                locx, locy = rotate_to_wind_direction(locx, locy, wind_resource.wind_directions[wind_farm_state_id])

                point_velocities[zi, yi, xi] = point_velocity(locx, locy, locz, rot_tx, rot_ty, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
                    rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, wtvelocities,
                    wind_resource, model_set,
                    wind_farm_state_id=wind_farm_state_id, downwind_turbine_id=0)

            end
        end
    end

    # if zlen == 1
    #     return point_velocities[1,1:ylen,1:xlen]
    # elseif ylen == 1
    #     return point_velocities[1:zlen,1,1:xlen]
    # elseif xlen == 1
    #     return point_velocities[1:zlen,1:ylen,1]
    # else
    return point_velocities[1:zlen,1:ylen,1:xlen]
    # end

end

function calculate_flow_field(xrange, yrange, zrange,
    model_set::AbstractModelSet, turbine_x, turbine_y, turbine_z, turbine_yaw,
    rotor_diameter, hub_height, ct_models, rotor_sample_points_y, rotor_sample_points_z,
    wind_resource; wind_farm_state_id=1)

    # rotate to direction frame for velocity calculations
    rot_tx, rot_ty = rotate_to_wind_direction(turbine_x, turbine_y, wind_resource.wind_directions[wind_farm_state_id])

    # sort the turbines
    sorted_turbine_index = sortperm(rot_tx)

    turbine_velocities, turbine_ct, turbine_ai, turbine_local_ti = turbine_velocities_one_direction(rot_tx, rot_ty, turbine_z, rotor_diameter, hub_height, turbine_yaw,
    sorted_turbine_index, ct_models, rotor_sample_points_y, rotor_sample_points_z, wind_resource,
    model_set, wind_farm_state_id=wind_farm_state_id, velocity_only=false)

    return calculate_flow_field(xrange, yrange, zrange,
        model_set, turbine_x, turbine_y, turbine_z, turbine_yaw, turbine_ct, turbine_ai,
        rotor_diameter, hub_height, turbine_local_ti, sorted_turbine_index, turbine_velocities,
        wind_resource, wind_farm_state_id=wind_farm_state_id)

end


# Convenience data structure for holding all global parameters
# in a single container.
mutable struct Parameters

  # The minimum value produced by a call to FLOAT.RAND.
  min_random_float::Float32

  # The maximum value produced by a call to FLOAT.RAND.
  max_random_float::Float32

  # The maximum value produced by a call to INTEGER.RAND.
  max_random_integer::Int32

  # The minimum value produced by a call to INTEGER.RAND.
  min_random_integer::Int32

  # The maximum number of points that will be executed in a single
  # top-level call to the interpreter.
  eval_push_limit::Int32

  # The probability that the selection of the ephemeral random NAME constant
  # for inclusion in randomly generated code will yield a new name, rather than
  # an existing name.
  new_erc_name_probability::Float32

  # The maximum number of points in an expression produced by CODE.RAND.
  max_points_in_random_expressions::Int32

  # The maximum number of points that may occur in any program on the CODE
  # stack. Instructions which would violate this limit act as NOOPs.
  # If this parameter is set to -1, then no such limit is enforced.
  max_points_in_program::Int32

  # The minimum length of a string produced by NAME.RAND.
  min_random_name_length::Int32

  # The maximum length of a string produced by NAME.RAND.
  max_random_name_length::Int32

  # When this flag is set to true, all code passed to the top level of the
  # interpreter will be pushed onto the CODE stack prior to execution.
  top_level_push_code::Bool

  # When this flag is set to true, the CODE stack will be popped at the
  # end of top level calls to the interpreter.
  top_level_pop_code::Bool

  # Creates a new set of default parameters.
  Parameters() =
    new(typemin(Float32), typemax(Float32),
      typemin(Int32), typemax(Int32),
      500, 0.001, 10, 10, 1, 10, true, false)

end

# Adjusts the parameters of a given parameter object using a parameter table.
set!(p::Parameters, params::Dict{String, String}) = for (k, v) in params
  set!(p, k, v)
end

# Adjusts the value of a single parameter, k, to a given value, v.
set!(p::Parameters, k::String, v::String) =
  if k == "MAX-RANDOM-FLOAT"
    p.max_random_float = parse(Float32, v)
  elseif k == "MIN-RANDOM-FLOAT"
    p.min_random_float = parse(Float32, v)
  elseif k == "MAX-RANDOM-INTEGER"
    p.max_random_integer = parse(Int32, v)
  elseif k == "MIN-RANDOM-INTEGER"
    p.min_random_integer = parse(Int32, v)
  elseif k == "EVAL-PUSH-LIMIT"
    p.eval_push_limit = parse(Int32, v)
  elseif k == "NEW-ERC-NAME-PROBABILITY"
    p.new_erc_name_probability = parse(Float32, v)
  elseif k == "MAX-POINTS-IN-RANDOM-EXPRESSIONS"
    p.max_points_in_random_expressions = parse(Int32, v)
  elseif k == "MAX-POINTS-IN-PROGRAM"
    p.max_points_in_program = parse(Int32, v)
  elseif k == "TOP-LEVEL-PUSH-CODE"
    p.top_level_push_code = v == "TRUE"
  elseif k == "MIN-RANDOM-NAME-LENGTH"
    p.min_random_string_length = parse(Int32, v)
  elseif k == "MAX-RANDOM-NAME-LENGTH"
    p.max_random_string_length = parse(Int32, v)
  elseif k == "TOP-LEVEL-POP-CODE"
    p.top_level_pop_code = v == "TRUE"
  end


import ..Tagging: hastag, traithastag, tag, mergetag
import ..Util: mergef
import ..Traits: traits
export AbstractΩ, defΩ, defω, resolve, idtype, replacetags, like

# # Sample Space
# A sample space represents a set of possible values.
# Sample spaces are structured in the sense that that are composed of parts.
# Different parts are indicated by an id (see IDS)

"Abstract Sample Space"
abstract type AbstractΩ end

# # Tags
"""
`tag(ω::AbstractΩ, tags)`

tag `ω` with `tags`.
"""
function tag(ω::AbstractΩ, tags, mergefunc=mergetag)
  replacetags(ω, mergef(mergefunc, ω.tags, tags))
end

rmtag(ω::AbstractΩ, tag) =
  replacetags(ω, rmkey(ω.tags, tag))

updatetag(ω::AbstractΩ, tag, val) =
  replacetags(ω, update(ω.tags, tag, val))

traithastag(t::AbstractΩ, tag) = traithastag(t.tags, tag)
hastag(ω::AbstractΩ, tag) = hastag(ω.tags, tag)

traits(ω::AbstractΩ) = traits(ω.tags) # FIXME: Do this at the type level

# # Defaults
"Default sample space"
function defΩ end

"Default sample space object"
function defω end

# FIXME MOve this somewhere (shouldnt really be in AbstractΩ)

function replacetags end

function resolve end

"`ids(ω)` Collection of ids in `ω`, i.e. domain of ω"
function ids end

function idtype end

function like end


module FinancialAssets

using Base.Dates

import Base: show

export FinancialAsset,
       LongStock, ShortStock, Stock,
       LongFuture, ShortFuture, Future,
       LongPut, ShortPut, LongCall, ShortCall,
       LongPutVertical, LongCallVertical, ShortPutVertical, ShortCallVertical,
       CoveredCall,
       #Ticker,
       OptionExercise,
       CUSIP, BloombergID, ReutersID, FinancialID,
       American, European,
       Currency, CurrencyPair,
       USD, EUR, GBP, AUD, JPY,
       EURUSD, EURJPY, USDGBP, USDAUD, USDNZD, USDJPY,
       F, G, H, J, K, M, N, Q, U, V, X, Z

include("id.jl")
include("assets.jl")
include("const.jl")
include("spread.jl")
include("apply.jl")
include("risk.jl")

end


using Serialization, DataStructures, Random

export computePredictionMetricsSparseVectorData

# Returns a priority queue of ordered item counts, for a dataset in sparse vector format
function getOrderedItemCountsSparseVectorData(trainingBasketsDictFileName,
                                              trainingBasketsDictObjectName)

  # Load training data
  trainingUsersBasketsDict = open(deserialize, trainingBasketsDictFileName);
  println("Loaded $trainingBasketsDictFileName")

  # Build set of training instances
  numTrainingInstances = length(collect(keys(trainingUsersBasketsDict)))
  trainingInstances = fill(Vector{Int}(), numTrainingInstances)
  trainingInstanceIndex = 1
  numItems = 0
  for trainingInstanceUserId in collect(keys(trainingUsersBasketsDict))
    trainingInstanceItems = deepcopy(trainingUsersBasketsDict[trainingInstanceUserId].basketItems)
    heldOutItem = trainingUsersBasketsDict[trainingInstanceUserId].heldOutItem
    push!(trainingInstanceItems, heldOutItem)

    trainingInstances[trainingInstanceIndex] = trainingInstanceItems

    trainingInstanceIndex += 1

    numItems = trainingUsersBasketsDict[trainingInstanceUserId].numItemsInCatalog
  end

  # Compute ordered item counts, by descending popularity
  orderedItemCounts = PriorityQueue(Dict{Int, Int}(), Base.Order.Reverse)
  for i = 1:numTrainingInstances
    trainingInstance = trainingInstances[i]

    for itemId in trainingInstance
      if haskey(orderedItemCounts, itemId)
        orderedItemCounts[itemId] += 1
      else
        # Create a new entry for this item
        orderedItemCounts[itemId] = 1
      end
    end
  end

  return orderedItemCounts
end

# Compute MPR and precision metrics from DPP prediction results for a dataset in sparse vector format
function computePredictionMetricsSparseVectorData(trainingBasketsDictFileName, trainingBasketsDictObjectName,
                                                  testBasketsDictFileName, testBasketsDictObjectName,
                                                  resultsForTestInstancesDictFileName)
  resultsForTestInstancesDictObjectName = "resultsForTestInstancesDict"
  Random.seed!(1234)

  # Load test data
  testUsersBasketsDict = open(deserialize, testBasketsDictFileName);
  println("Loaded $testBasketsDictFileName")

  # Build set of test instances
  numTestInstances = length(collect(keys(testUsersBasketsDict)))
  testInstances = fill(Vector{Int}(), numTestInstances)
  testInstanceIndex = 1
  numItems = 0
  for testInstanceUserId in collect(keys(testUsersBasketsDict))
    testInstanceItems = deepcopy(testUsersBasketsDict[testInstanceUserId].basketItems)
    heldOutItem = testUsersBasketsDict[testInstanceUserId].heldOutItem
    push!(testInstanceItems, heldOutItem)

    testInstances[testInstanceIndex] = testInstanceItems

    testInstanceIndex += 1

    numItems = testUsersBasketsDict[testInstanceUserId].numItemsInCatalog
  end

  # Compute popularity-stratification weights for each item, where the weight is
  # proportional to 1 / (itemCount ^ beta), for beta = 0.5
  orderedItemCounts = getOrderedItemCountsSparseVectorData(trainingBasketsDictFileName,
                                                           trainingBasketsDictObjectName)
  sumItemCounts = sum(values(orderedItemCounts))
  itemPopWeights = PriorityQueue(Dict{Int, Float64}(), Base.Order.Reverse)
  beta = 0.5
  startingItemId = 1
  for i = startingItemId:numItems
    if haskey(orderedItemCounts, i)
      itemPopWeights[i] = 1 / (orderedItemCounts[i] ^ beta)
    else
      # If an item id is not available in orderedItemCounts, then its popularity is 0.
      itemPopWeights[i] = 0
      println("itemPopWeights[$i]: 0")
    end
  end
  # Normalize itemPopWeights, so they sum to 1
  sumItemPopWeights = sum(collect(values(itemPopWeights)))
  for i = startingItemId:numItems
    itemPopWeights[i] = itemPopWeights[i] / sumItemPopWeights
  end

  # Load DPP prediction results from disk
  resultsForTestInstancesDict = open(deserialize, resultsForTestInstancesDictFileName);
  println("Loaded $resultsForTestInstancesDictFileName")

  # Compute mean percentile rank and precision across predictions for all test instances
  shuffle!(testInstances)
  numTestInstances = length(testInstances)
  sumPredictedRanks = 0
  basketSizeResultsDict = Dict{Int, BasketSizeResult}()
  predictedRanks = Int[]
  for i = 1:numTestInstances
    testInstance = testInstances[i]

    # Get predicted rank for test instance
    testResult = resultsForTestInstancesDict[testInstance]
    sumPredictedRanks += testResult.predictedRankActualNextItem
    push!(predictedRanks, testResult.predictedRankActualNextItem)

    # Last item in the test instance is not part of the basket
    basketSize = length(testInstance) - 1
    # Update metrics for the size of the current basket
    if basketSize >= 2
      if haskey(basketSizeResultsDict, basketSize)
        basketSizeResult = basketSizeResultsDict[basketSize]
        basketSizeResult.numBasketInstances += 1
        basketSizeResult.sumPredictedRanks += testResult.predictedRankActualNextItem
        push!(basketSizeResult.predictedRanks, testResult.predictedRankActualNextItem)
        basketSizeResultsDict[basketSize] = basketSizeResult
      else
        # Create a new entry for this basket size
        basketSizeResult = BasketSizeResult(basketSize, 1, testResult.predictedRankActualNextItem,
                                            [testResult.predictedRankActualNextItem])
        basketSizeResultsDict[basketSize] = basketSizeResult
      end

      if i % 100 == 0
        println("Processed $i test instances")
      end
    end
  end

  # Compute overall mean percentile rank
  meanRank = sumPredictedRanks / numTestInstances
  meanRankPercent = meanRank / numItems
  meanPercentileRank = 100.0 - (100.0 * meanRankPercent)
  println("meanPercentileRank for all baskets: $meanPercentileRank")

  # Compute mean percentile rank for each basket size
  for basketSize in sort(collect(keys(basketSizeResultsDict)))
    basketSizeResult = basketSizeResultsDict[basketSize]
    meanRank = basketSizeResult.sumPredictedRanks / basketSizeResult.numBasketInstances
    meanRankPercent = meanRank / numItems
    meanPercentileRank = 100.0 - (100.0 * meanRankPercent)
    println("For basket size $basketSize, meanPercentileRank: $meanPercentileRank")
  end

  # Compute overall precision@1 (accuracy)
  hits = length(findall(x -> x < 1, predictedRanks))
  precisionAtOnePercent = hits / numTestInstances * 100
  println("")
  println("precision@1 for all baskets: $precisionAtOnePercent")

  # Compute overall popularity-weighted precision@1
  popWeightedPrecisionAtOnePercent =
    computePopularityWeightedPrecisionAtK(itemPopWeights, resultsForTestInstancesDict,
                                          testInstances, predictedRanks, 1)
  println("")
  println("Popularity-weighted precision@1 for all baskets: $popWeightedPrecisionAtOnePercent")

  # Compute precision@1 (accuracy) for each basket size
  for basketSize in sort(collect(keys(basketSizeResultsDict)))
    basketSizeResult = basketSizeResultsDict[basketSize]
    hits = length(findall(x -> x < 1, basketSizeResult.predictedRanks))
    precisionAtOnePercent = hits / basketSizeResult.numBasketInstances * 100
    println("For basket size $basketSize, precision@1: $precisionAtOnePercent")
  end

  # Compute overall precision@5
  hits = length(findall(x -> x < 5, predictedRanks))
  precisionAtFivePercent = hits / numTestInstances * 100
  println("")
  println("precision@5 for all baskets: $precisionAtFivePercent")

  # Compute overall popularity-weighted precision@5
  popWeightedPrecisionAtFivePercent =
    computePopularityWeightedPrecisionAtK(itemPopWeights, resultsForTestInstancesDict,
                                          testInstances, predictedRanks, 5)
  println("")
  println("Popularity-weighted precision@5 for all baskets: $popWeightedPrecisionAtFivePercent")

  # Compute precision@5 for each basket size
  for basketSize in sort(collect(keys(basketSizeResultsDict)))
    basketSizeResult = basketSizeResultsDict[basketSize]
    hits = length(findall(x -> x < 5, basketSizeResult.predictedRanks))
    precisionAtFivePercent = hits / basketSizeResult.numBasketInstances * 100
    println("For basket size $basketSize, precision@5: $precisionAtFivePercent")
  end

  # Compute overall precision@10
  hits = length(findall(x -> x < 10, predictedRanks))
  precisionAtTenPercent = hits / numTestInstances * 100
  println("")
  println("precision@10 for all baskets: $precisionAtTenPercent")

  # Compute overall popularity-weighted precision@10
  popWeightedPrecisionAtTenPercent =
    computePopularityWeightedPrecisionAtK(itemPopWeights, resultsForTestInstancesDict,
                                          testInstances, predictedRanks, 10)
  println("")
  println("Popularity-weighted precision@10 for all baskets: $popWeightedPrecisionAtTenPercent")

  # Compute precision@10 for each basket size
  for basketSize in sort(collect(keys(basketSizeResultsDict)))
    basketSizeResult = basketSizeResultsDict[basketSize]
    hits = length(findall(x -> x < 10, basketSizeResult.predictedRanks))
    precisionAtTenPercent = hits / basketSizeResult.numBasketInstances * 100
    println("For basket size $basketSize, precision@10: $precisionAtTenPercent")
  end

  # Compute overall precision@20
  hits = length(findall(x -> x < 20, predictedRanks))
  precisionAtTwentyPercent = hits / numTestInstances * 100
  println("")
  println("precision@20 for all baskets: $precisionAtTwentyPercent")

  # Compute overall popularity-weighted precision@20
  popWeightedPrecisionAtTwentyPercent =
    computePopularityWeightedPrecisionAtK(itemPopWeights, resultsForTestInstancesDict,
                                          testInstances, predictedRanks, 20)
  println("")
  println("Popularity-weighted precision@20 for all baskets: $popWeightedPrecisionAtTwentyPercent")

  # Compute precision@20 for each basket size
  for basketSize in sort(collect(keys(basketSizeResultsDict)))
    basketSizeResult = basketSizeResultsDict[basketSize]
    hits = length(findall(x -> x < 20, basketSizeResult.predictedRanks))
    precisionAtTwentyPercent = hits / basketSizeResult.numBasketInstances * 100
    println("For basket size $basketSize, precision@20: $precisionAtTwentyPercent")
  end
end

# Computes the popularity weighted (stratified) precision@k
function computePopularityWeightedPrecisionAtK(itemPopWeights, resultsForTestInstancesDict,
                                               testInstances, predictedRanks, k)
  popWeightedHits = 0
  numTestInstances = length(testInstances)
  sumPopWeightsForTestRatings = 0

  startingHeldOutItemId = 0

  for i = 1:numTestInstances
    testInstance = testInstances[i]
    heldOutTestItemId = resultsForTestInstancesDict[testInstance].actualNextItemId - startingHeldOutItemId

    sumPopWeightsForTestRatings += itemPopWeights[heldOutTestItemId]

    if predictedRanks[i] < k
      popWeightedHits += itemPopWeights[heldOutTestItemId]
    end
  end

  popWeightedPrecisionAtKPercent = popWeightedHits / sumPopWeightsForTestRatings * 100
  return popWeightedPrecisionAtKPercent
end


#=
This file is part of the replication code for: Hasenzagl, T., Pellegrino, F., Reichlin, L., & Ricco, G. (2020). A Model of the Fed's View on Inflation.
Please cite the paper if you are using any part of the code for academic work (including, but not limited to, conference and peer-reviewed papers).
=#

function quarterly2monthly(a,n)

    #=
    Filename: monthly2quarterly.jl
    a: quarterly data
    n: length of resulting monthly array
    =#

    if ndims(a) == 1
        nQ = 1;
    else
        nQ = size(a,2);
    end

    a[ismissing.(a)] .= NaN;
    b = hcat([kron(a[:,i], [NaN; NaN; 1]) for i=1:nQ]...) |> Array{Union{Float64, Missing}};
    b[isnan.(b)] .= missing;

    if size(b,1) < n
        b = [b; missing.*ones(n-size(b,1), nQ)];
    end

    return b;
end


#######################
## sorting utilities ##
#######################

@doc doc"""
Lexicographical sorting of paths given as Array{Array{Int, 1}, 1}.
This is implicitly called by `CTreePaths` contructors.
"""->
function sortPaths(paths::IntArrays)
    pathsArr = _cTreePaths2arr(paths)
    sortedArr = sortcols(pathsArr)
    return paths = _arr2CTreePaths(sortedArr)
end

## write tree paths as matrix
##---------------------------

@doc doc"""
`Array{Int, 2}` rectangular matrix with paths as columns, filled to
maximum depth with zeros. This is not a conversion to parent notation,
but simply exists to make use of built-in lexicographic sorting
functions for arrays.
"""->
function _cTreePaths2arr(paths::IntArrays)
    ## transform array of array into matrix, filling empty fields with
    ## zero
    nRows = maxDepth(paths)
    vals = zeros(Int, nRows, length(paths))
    for ii=1:length(paths)
        nVals = length(paths[ii])
        vals[1:nVals, ii] = paths[ii]
    end
    return vals
end

@doc doc"""
`Array{Array{Int, 1}, 1}` of tree paths, obtained from a rectangular
matrix `Array{Int, 2}` with trailing zeros. Trailing zeros are simply
cut off, in order to get paths of different length.
"""->
function _arr2CTreePaths(arr::Array{Int, 2})
    ## remove trailing zeros
    maxDepth, nPaths = size(arr)
    paths = Array(Array{Int, 1}, nPaths)
    for ii=1:nPaths
        currPath = arr[:, ii]
        noZeros = (currPath .!= 0)
        if !isa(noZeros, Array)
            paths[ii] = [currPath[noZeros]]
        else
            paths[ii] = currPath[noZeros]
        end
    end
    return paths
end

##########################
## conversion functions ##
##########################

## @doc doc"""
## Convert `Array{Int, 1}` of tree information given in parent notation
## into a `Tree` instance. Parent notation specifies for each variable
## its single parent in an acyclic tree.

## The algorithm starts by finding all final nodes as those that are
## listed as parent node for any other variable. For the case of
## incomplete trees, nodes not yet connected are encoded with a value of
## -1.
## """->
## function par2tree(parNot::Array{Int, 1})
##     ## transform single parent notation column to tree
##     tr = CTreeParRef(parNot)
##     return convert(CTreePaths, tr)
## end

## function par2tree(parNot::Array{Int, 2})
##     ## transform parent notation matrix to array of trees
##     nVars = size(parNot, 2)
##     trees = Array(CTreePaths, nVars)
##     for ii=1:nVars
##         trees[ii] = par2tree(parNot[:, ii])
##     end
##     return trees
## end

## @doc doc"""
## Transform `Tree` to `Array{Int, 1}` parent notation. Root nodes get
## parent value of 0. Not yet connected nodes are encoded as value -1.
## """->
## function tree2par(tP::Tree, nVars::Int)
##     return convert(CTreeParRef, tP, nVars)
## end

## @doc doc"""
## Transform `Array` of `Tree`s into parent notation square matrix.
## """->
## function tree2par(tPs::Array{Tree, 1}, nVars::Int)
##     ## transform array of trees to parent notation matrix
##     parNot = Array(Int, nVars, nVars)
##     for ii=1:nVars
##         tPar = convert(CTreeParRef, tPs[ii], nVars)
##         parNot[:, ii] = tPar.tree
##     end
##     return parNot
## end

## function vine2trees(vn::Vine)
##     return par2tree(vn.trees)
## end

## function trees2vine(trs::Array{Tree, 1}, nVars::Int)
##     return Vine(tree2par(trs, nVars))
## end


using NCDatasets
using Test

fname = tempname()

for define_standard_name in [false,true]

    NCDataset(fname,"c") do ds
        # Dimensions

        ds.dim["xi_rho"] = 6
        ds.dim["xi_u"] = 5
        ds.dim["xi_v"] = 6
        ds.dim["xi_psi"] = 5
        ds.dim["eta_rho"] = 6
        ds.dim["eta_u"] = 6
        ds.dim["eta_v"] = 5
        ds.dim["eta_psi"] = 5
        ds.dim["N"] = 30
        ds.dim["s_rho"] = 30
        ds.dim["s_w"] = 31
        ds.dim["tracer"] = 63
        ds.dim["ocean_time"] = Inf # unlimited dimension

        # Declare variables

        nclon_rho = defVar(ds,"lon_rho", Float64, ("xi_rho", "eta_rho"))
        nclon_rho.attrib["long_name"] = "longitude of RHO-points"
        nclon_rho.attrib["units"] = "degree_east"
        nclon_rho.attrib["field"] = "lon_rho, scalar"
        if define_standard_name
            nclon_rho.attrib["standard_name"] = "longitude"
        end

        nclat_rho = defVar(ds,"lat_rho", Float64, ("xi_rho", "eta_rho"))
        nclat_rho.attrib["long_name"] = "latitude of RHO-points"
        nclat_rho.attrib["units"] = "degree_north"
        nclat_rho.attrib["field"] = "lat_rho, scalar"
        if define_standard_name
            nclon_rho.attrib["standard_name"] = "latitude"
        end

        nclon_u = defVar(ds,"lon_u", Float64, ("xi_u", "eta_u"))
        nclon_u.attrib["long_name"] = "longitude of U-points"
        nclon_u.attrib["units"] = "degree_east"
        nclon_u.attrib["field"] = "lon_u, scalar"

        nclat_u = defVar(ds,"lat_u", Float64, ("xi_u", "eta_u"))
        nclat_u.attrib["long_name"] = "latitude of U-points"
        nclat_u.attrib["units"] = "degree_north"
        nclat_u.attrib["field"] = "lat_u, scalar"

        nczeta = defVar(ds,"zeta", Float32, ("xi_rho", "eta_rho", "ocean_time"))
        nczeta.attrib["long_name"] = "free-surface"
        nczeta.attrib["units"] = "meter"
        nczeta.attrib["time"] = "ocean_time"
        nczeta.attrib["coordinates"] = "lat_rho lon_rho"
        nczeta.attrib["field"] = "free-surface, scalar, series"

        ncubar = defVar(ds,"ubar", Float32, ("xi_u", "eta_u", "ocean_time"))
        ncubar.attrib["long_name"] = "vertically integrated u-momentum component"
        ncubar.attrib["units"] = "meter second-1"
        ncubar.attrib["time"] = "ocean_time"
        ncubar.attrib["coordinates"] = "lat_u lon_u"
        ncubar.attrib["field"] = "ubar-velocity, scalar, series"
    end

    local ds
    ds = NCDataset(fname)
    @test name(coord(ds["zeta"],"longitude")) == "lon_rho"
    @test name(coord(ds["ubar"],"longitude")) == "lon_u"
    @test coord(ds["ubar"],"foobar") == nothing
    close(ds)
end


layernames(::Type{WorldClim}, ::Type{Elevation}) = ("Elevation",)

function layernames(::Type{WorldClim}, ::Type{BioClim})
    return (
        "Annual Mean Temperature",
        "Mean Diurnal Range",
        "Isothermality",
        "Temperature Seasonality",
        "Max Temperature of Warmest Month",
        "Min Temperature of Coldest Month",
        "Temperature Annual Range",
        "Mean Temperature of Wettest Quarter",
        "Mean Temperature of Driest Quarter",
        "Mean Temperature of Warmest Quarter",
        "Mean Temperature of Coldest Quarter",
        "Annual Precipitation",
        "Precipitation of Wettest Month",
        "Precipitation of Driest Month",
        "Precipitation Seasonality",
        "Precipitation of Wettest Quarter",
        "Precipitation of Driest Quarter",
        "Precipitation of Warmest Quarter",
        "Precipitation of Coldest Quarter",
    )
end

layernames(::Type{CHELSA}, ::Type{BioClim}) = layernames(WorldClim, BioClim)

function layernames(::Type{EarthEnv}, ::Type{HabitatHeterogeneity})
    return (
        "Coefficient of variation",
        "Evenness",
        "Range",
        "Shannon",
        "Simpson",
        "Standard deviation",
        "Contrast",
        "Correlation",
        "Dissimilarity",
        "Entropy",
        "Homogeneity",
        "Maximum",
        "Uniformity",
        "Variance",
    )
end

function layernames(::Type{EarthEnv}, ::Type{LandCover})
    return (
        "Evergreen/Deciduous Needleleaf Trees",
        "Evergreen Broadleaf Trees",
        "Deciduous Broadleaf Trees",
        "Mixed/Other Trees",
        "Shrubs",
        "Herbaceous Vegetation",
        "Cultivated and Managed Vegetation",
        "Regularly Flooded Vegetation",
        "Urban/Built-up",
        "Snow/Ice",
        "Barren",
        "Open Water",
    )
end

function layernames(::Type{EarthEnv}, ::Type{Topography})
    return (
        "Elevation",
        "Slope",
        "Aspect Cosine",
        "Aspect Sine",
        "Aspect Eastness",
        "Aspect Northness",
        "Roughness",
        "Topographic Position Index",
        "Terrain Ruggedness Index",
        "Vector Ruggedness Measure",
        "∂(E-W slope)",
        "∂²(E-W slope)",
        "∂(N-S slope)",
        "∂²(N-S slope)",
        "Profile curvature",
        "Tangential curvature"
    )
end

layernames(p::Type{<:LayerProvider}, d::Type{<:LayerDataset}, i::Int) = layernames(p, d)[i]
layernames(p::Type{<:LayerProvider}, d::Type{<:LayerDataset}, i::AbstractArray) = layernames(p, d)[i]


# check that all interace functions are provided

# --- LazySet ---

# support vector
@test check_method_implementation(LazySet, σ,
                                  Function[S -> (Vector{Float64}, S{Float64})])
# dimension
check_method_implementation(LazySet, dim, Function[S -> (S{Float64},)])

# --- AbstractPolyhedron ---

# constraints list
@test check_method_implementation(AbstractPolyhedron, constraints_list,
                                  Function[S -> (S{Float64},)])

# --- AbstractPolytope ---

# vertices list
@test check_method_implementation(AbstractPolytope, vertices_list,
                                  Function[S -> (S{Float64},)])

# --- AbstractCentrallySymmetric ---

# center
@test check_method_implementation(AbstractCentrallySymmetric, center,
                                  Function[S -> (S{Float64},)])

# --- AbstractCentrallySymmetricPolytope ---
# (copies methods from AbstractCentrallySymmetric)

# center (from AbstractCentrallySymmetric)
@test check_method_implementation(AbstractCentrallySymmetricPolytope, center,
                                  Function[S -> (S{Float64},)])

# --- AbstractZonotope ---

# genmat
@test check_method_implementation(AbstractZonotope, genmat,
                                  Function[S -> (S{Float64},)])
# generators
@test check_method_implementation(AbstractZonotope, generators,
                                  Function[S -> (S{Float64},)])

# --- AbstractHyperrectangle ---

# radius_hyperrectangle (2x)
@test check_method_implementation(AbstractHyperrectangle, radius_hyperrectangle,
                                  Function[S -> (S{Float64},)])
@test check_method_implementation(AbstractHyperrectangle, radius_hyperrectangle,
                                  Function[S -> (S{Float64}, Int)])
# isflat
@test check_method_implementation(AbstractHyperrectangle, isflat,
                                  Function[S -> (S{Float64},)])

# --- AbstractSingleton ---

# element (2x)
@test check_method_implementation(AbstractSingleton, element,
                                  Function[S -> (S{Float64},)])
@test check_method_implementation(AbstractSingleton, element,
                                  Function[S -> (S{Float64}, Int)])

# --- AbstractPolygon ---

# tohrep
@test check_method_implementation(AbstractPolygon, tohrep,
                                  Function[S -> (S{Float64},)])
# tovrep
@test check_method_implementation(AbstractPolygon, tovrep,
                                  Function[S -> (S{Float64},)])

# --- AbstractHPolygon ---

# no functions yet


@compat for (func, elty) in [(:CLBlastStrmv, Float32), (:CLBlastDtrmv, Float64),
                     (:CLBlastCtrmv, ComplexF32), (:CLBlastZtrmv, ComplexF64)]
    #TODO: (:CLBlastHtrmv, Float16)

    @eval function $func(layout::CLBlastLayout, triangle::CLBlastTriangle,
                         a_transpose::CLBlastTranspose, diagonal::CLBlastDiagonal,
                         n::Integer,
                         a_buffer::cl.CL_mem, a_offset::Integer, a_ld::Integer,
                         x_buffer::cl.CL_mem, x_offset::Integer, x_inc::Integer,
                         queue::cl.CmdQueue, event::cl.Event)
        err = ccall(
            ($(string(func)), libCLBlast),
            cl.CL_int,
            (Cint, Cint, Cint, Cint, Csize_t, Ptr{Cvoid}, Csize_t, Csize_t,
              Ptr{Cvoid}, Csize_t, Csize_t, Ptr{Cvoid}, Ptr{Cvoid}),
            Cint(layout), Cint(triangle), Cint(a_transpose), Cint(diagonal), n,
              a_buffer, a_offset, a_ld, x_buffer, x_offset, x_inc, Ref(queue), Ref(event)
        )
        if err != cl.CL_SUCCESS
            println(stderr, "Calling function $(string($func)) failed!")
            throw(CLBlastError(err))
        end
        return err
    end

    @eval function trmv!(uplo::Char, trans::Char, diag::Char,
                         A::cl.CLArray{$elty,2}, x::cl.CLArray{$elty};
                         queue::cl.CmdQueue=cl.queue(x))
        # check and convert arguments
        m, n = size(A)
        if m != n
            throw(DimensionMismatch("`A` has dimensions $(size(A)) but must be square."))
        end
        if length(x) != n
            throw(DimensionMismatch("`x` has length $(length(x)) while $n is required."))
        end
        if uplo == 'U'
            triangle = CLBlastTriangleUpper
        elseif uplo == 'L'
            triangle = CLBlastTriangleLower
        else
            throw(ArgumentError("Upper/lower marker `uplo` is $(uplo) but only 'U' and 'L' are allowed."))
        end
        if trans == 'N'
            a_transpose = CLBlastTransposeNo
        elseif trans == 'T' || (trans == 'C' && $elty <: Real)
            a_transpose = CLBlastTransposeYes
        elseif trans == 'C' && $elty <: Complex
            a_transpose = CLBlastTransposeConjugate
        else
            throw(ArgumentError("Transpose marker `trans` is $(uplo) but only 'N', 'T', and 'C' are allowed."))
        end
        if diag == 'N'
            diagonal = CLBlastDiagonalNonUnit
        elseif diag == 'U'
            diagonal = CLBlastDiagonalUnit
        else
            throw(ArgumentError("Diagonal marker `diag` is $(diag) but only 'N' and 'U' are allowed."))
        end
        layout = CLBlastLayoutColMajor

        # output event
        event::cl.Event = cl.Event(C_NULL)

        $func(layout, triangle, a_transpose, diagonal,
              n,
              pointer(A), 0, size(A,1),
              pointer(x), 0, 1,
              queue, event)

        # wait for kernel
        cl.wait(event)

        x
    end

end


"""
    GridTopology(mesh)

A generic tensor-product topology defined on either an equispaced, regular rectangular
structured grid or irregular one.
"""
struct GridTopology{M <: AbstractMesh, B} <: AbstractTopology
    mesh::M
    boundaries::B
end
function GridTopology(mesh::M) where {M <: AbstractMesh}
    x1boundary = mesh.domain.x1boundary
    x2boundary = mesh.domain.x2boundary
    boundaries = if isnothing(x1boundary)
        if isnothing(x2boundary)
            NamedTuple()
        else
            NamedTuple{x2boundary}((3, 4))
        end
    else
        if isnothing(x2boundary)
            NamedTuple{x1boundary}((1, 2))
        else
            NamedTuple{(x1boundary..., x2boundary...)}((1, 2, 3, 4))
        end
    end
    GridTopology(mesh, boundaries)
end

function Base.show(io::IO, topology::GridTopology{M}) where {M <: AbstractMesh}
    print(io, "GridTopology on ", topology.mesh)
end
domain(topology::GridTopology{M}) where {M <: AbstractMesh} =
    topology.mesh.domain

function nlocalelems(topology::GridTopology{M}) where {M <: AbstractMesh}
    n1 = topology.mesh.n1
    n2 = topology.mesh.n2
    return n1 * n2
end

eachslabindex(topology::GridTopology{M}) where {M <: AbstractMesh} =
    1:nlocalelems(topology)

# InteriorFaceIterator
function Base.length(
    fiter::InteriorFaceIterator{T},
) where {M <: AbstractMesh, T <: GridTopology{M}}
    topology = fiter.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)
    return (x1periodic ? n1 : n1 - 1) * n2 + n1 * (x2periodic ? n2 : n2 - 1)
end

function Base.iterate(
    fiter::InteriorFaceIterator{T},
    (d, z1, z2) = (1, 0, 0),
) where {M <: AbstractMesh, T <: GridTopology{M}}
    # iteration state (major first)
    #  - d ∈ (1,2): face direction
    #  - z1 ∈ 0:n1-1: 0-based face index in direction 1
    #  - z2 ∈ 0:n2-1: 0-based face index in direction 2

    topology = fiter.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)

    # skip boundary faces
    if d == 1 && z1 == 0 && !x1periodic
        d = 2
    end
    if d == 2 && z2 == 0 && !x2periodic
        d = 1
        z1 += 1
        if z1 >= n1
            z1 = 0
            z2 += 1
            if !x1periodic
                d = 2
            end
        end
    end

    if z2 >= n2
        return nothing
    end

    if d == 1
        y1 = z1 == 0 ? n1 - 1 : z1 - 1
        y2 = z2
    else
        y1 = z1
        y2 = z2 == 0 ? n2 - 1 : z2 - 1
    end

    elem1 = z2 * n1 + z1 + 1
    elem2 = y2 * n1 + y1 + 1
    if d == 1
        nextstate = (2, z1, z2)
    else
        z1 += 1
        if z1 == n1
            z1 = 0
            z2 += 1
        end
        nextstate = (1, z1, z2)
    end
    if d == 1
        return (elem1, 1, elem2, 2, false), nextstate
    else
        return (elem1, 3, elem2, 4, false), nextstate
    end
end

# BoundaryFaceIterator
function boundary_names(topology::GridTopology{M}) where {M <: AbstractMesh}
    x1boundary = topology.mesh.domain.x1boundary
    x2boundary = topology.mesh.domain.x2boundary
    if isnothing(x1boundary)
        isnothing(x2boundary) ? () : x2boundary
    else
        isnothing(x2boundary) ? x1boundary : (x1boundary..., x2boundary...)
    end
end

function boundary_tag(
    topology::GridTopology{M},
    name::Symbol,
) where {M <: AbstractMesh}
    x1boundary = topology.mesh.domain.x1boundary
    x2boundary = topology.mesh.domain.x2boundary
    if !isnothing(x1boundary)
        x1boundary[1] == name && return 1
        x1boundary[2] == name && return 2
    end
    if !isnothing(x2boundary)
        x2boundary[1] == name && return 3
        x2boundary[2] == name && return 4
    end
    error("Invalid boundary name")
end

function boundaries(topology::GridTopology{M}) where {M <: AbstractMesh}
    return topology.boundaries
end

function Base.length(
    bfiter::BoundaryFaceIterator{T},
) where {M <: AbstractMesh, T <: GridTopology{M}}
    boundary = bfiter.boundary
    topology = bfiter.topology
    if boundary in (1, 2)
        if isnothing(topology.mesh.domain.x1boundary)
            return 0
        else
            return topology.mesh.n2
        end
    end
    if boundary in (3, 4)
        if isnothing(topology.mesh.domain.x2boundary)
            return 0
        else
            return topology.mesh.n1
        end
    end
end

function Base.iterate(
    bfiter::BoundaryFaceIterator{T},
) where {M <: AbstractMesh, T <: GridTopology{M}}
    boundary = bfiter.boundary
    topology = bfiter.topology
    if boundary in (1, 2) && isnothing(topology.mesh.domain.x1boundary)
        return nothing
    end
    if boundary in (3, 4) && isnothing(topology.mesh.domain.x2boundary)
        return nothing
    end
    Base.iterate(bfiter, 0)
end

function Base.iterate(
    bfiter::BoundaryFaceIterator{T},
    z,
) where {M <: AbstractMesh, T <: GridTopology{M}}
    boundary = bfiter.boundary
    topology = bfiter.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    if boundary == 1
        z >= n2 && return nothing
        elem = z * n1 + 1
    elseif boundary == 2
        z >= n2 && return nothing
        elem = z * n1 + n1
    elseif boundary == 3
        z >= n1 && return nothing
        elem = z + 1
    elseif boundary == 4
        z >= n1 && return nothing
        elem = (n2 - 1) * n1 + z + 1
    end
    return (elem, boundary), z + 1
end

# VertexIterator
function Base.length(
    viter::VertexIterator{T},
) where {M <: AbstractMesh, T <: GridTopology{M}}
    topology = viter.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)
    nv1 = x1periodic ? n1 : n1 + 1
    nv2 = x2periodic ? n2 : n2 + 1
    return nv1 * nv2
end

function Base.iterate(
    viter::VertexIterator{T},
    (z1, z2) = (0, 0),
) where {M <: AbstractMesh, T <: GridTopology{M}}
    topology = viter.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)
    nv1 = x1periodic ? n1 : n1 + 1 # unique vertices in x1 direction
    nv2 = x2periodic ? n2 : n2 + 1 # unique vertices in x2 direction

    if z2 >= nv2
        return nothing
    end
    vertex = Vertex(topology, (z1, z2))
    z1 += 1
    if z1 >= nv1
        nextstate = (0, z2 + 1)
    else
        nextstate = (z1, z2)
    end
    return vertex, nextstate
end

# Vertex
function Base.length(
    vertex::Vertex{T},
) where {M <: AbstractMesh, T <: GridTopology{M}}
    topology = vertex.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)

    z1, z2 = vertex.num

    k1 = !x1periodic && (z1 == 0 || z1 == n1) ? 1 : 2
    k2 = !x2periodic && (z2 == 0 || z2 == n2) ? 1 : 2
    return k1 * k2
end

function Base.iterate(
    vertex::Vertex{T},
    vert = 0,
) where {M <: AbstractMesh, T <: GridTopology{M}}
    # iterator of (element, vertnum) that share global vertex
    topology = vertex.topology
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)
    nv1 = x1periodic ? n1 : n1 + 1
    nv2 = x2periodic ? n2 : n2 + 1
    z1, z2 = vertex.num

    vert += 1
    # at the boundary, skip non-existent elements
    if !x1periodic
        if z1 == 0 && (vert == 2 || vert == 4)
            vert += 1
        end
        if z1 == n1 && (vert == 1 || vert == 3)
            vert += 1
        end
    end
    if !x2periodic
        if z2 == 0 && (vert == 3 || vert == 4)
            vert += 2
        end
        if z2 == n2 && (vert == 1 || vert == 2)
            vert += 2
        end
    end

    if vert > 4
        return nothing
    end

    if vert == 2 || vert == 4
        z1 = mod(z1 - 1, nv1)
    end
    if vert == 3 || vert == 4
        z2 = mod(z2 - 1, nv2)
    end
    elem = z2 * n1 + z1 + 1
    return (elem, vert), vert
end

# GridTopology1D
"""
    GridTopology1D(mesh)

A line topology defined on an equispaced linear mesh.
"""

function GridTopology1D(mesh::Meshes.EquispacedLineMesh)
    x1boundary = mesh.domain.x3boundary
    x2boundary = nothing
    boundaries =
        isnothing(x1boundary) ? NamedTuple() : NamedTuple{x1boundary}((1, 2))
    return GridTopology(mesh, boundaries)
end

# Uniform grid implementations, dispatching on EquispacedRectangleMesh
function vertex_coordinates(
    topology::GridTopology{M},
    elem::Integer,
) where {M <: EquispacedRectangleMesh}
    @assert 1 <= elem <= nlocalelems(topology)
    CT = Topologies.coordinate_type(topology)

    # convert to 0-based indices
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    range1 = mesh.range1
    range2 = mesh.range2

    z2, z1 = fldmod(elem - 1, n1)

    c1 = CT(range1[z1 + 1], range2[z2 + 1])
    c2 = CT(range1[z1 + 2], range2[z2 + 1])
    c3 = CT(range1[z1 + 1], range2[z2 + 2])
    c4 = CT(range1[z1 + 2], range2[z2 + 2])
    return (c1, c2, c3, c4)
end

function opposing_face(
    topology::GridTopology{M},
    elem::Integer,
    face::Integer,
) where {M <: EquispacedRectangleMesh}
    @assert 1 <= elem <= nlocalelems(topology)
    @assert 1 <= face <= 4

    # convert to 0-based indices
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    x1periodic = isnothing(mesh.domain.x1boundary)
    x2periodic = isnothing(mesh.domain.x2boundary)

    z2, z1 = fldmod(elem - 1, n1)
    if face == 1
        z1 -= 1
        if z1 < 0
            if !x1periodic
                return (0, 1, false)
            end
            z1 += n1
        end
        opface = 2
    elseif face == 2
        z1 += 1
        if z1 == n1
            if !x1periodic
                return (0, 2, false)
            end
            z1 -= n1
        end
        opface = 1
    elseif face == 3
        z2 -= 1
        if z2 < 0
            if !x2periodic
                return (0, 3, false)
            end
            z2 += n2
        end
        opface = 4
    elseif face == 4
        z2 += 1
        if z2 == n2
            if !x2periodic
                return (0, 4, false)
            end
            z2 -= n2
        end
        opface = 3
    end
    opelem = z2 * n1 + z1 + 1
    return opelem, opface, false
end

# Non-uniform grid implementations, dispatching on TensorProductMesh

function vertex_coordinates(
    topology::GridTopology{M},
    elem::Integer,
) where {M <: TensorProductMesh}
    @assert 1 <= elem <= nlocalelems(topology)

    # convert to 0-based indices
    mesh = topology.mesh
    n1 = mesh.n1
    n2 = mesh.n2
    coordinates = mesh.coordinates

    z2, z1 = fldmod(elem - 1, n1)

    c1 = coordinates[z1 * (n2 + 1) + (z2 + 1)]
    c2 = coordinates[(z1 + 1) * (n2 + 1) + (z2 + 1)]
    c3 = coordinates[z1 * (n2 + 1) + (z2 + 2)]
    c4 = coordinates[(z1 + 1) * (n2 + 1) + (z2 + 2)]

    return (c1, c2, c3, c4)
end

function opposing_face(
    topology::GridTopology{M},
    elem::Integer,
    face::Integer,
) where {M <: TensorProductMesh}
    @assert 1 <= elem <= nlocalelems(topology)
    @assert 1 <= face <= 4

    return topology.mesh.faces[(elem - 1) * 4 + face][3:5]
end


using Distributed
using Random
using Revise
# using ProgressMeter
# using ArgParse

const IN_SLURM = "SLURM_JOBID" in keys(ENV)
IN_SLURM && using ClusterManagers

ARGS_KWARGS = [
    Dict(:ARGS=>(0.5/8, 0.1, 1.0, 32, 8),   :KWARGS=>Dict([:n_episodes=>1000])),
    Dict(:ARGS=>(0.5/4, 0.1, 1.0, 32, 4),   :KWARGS=>Dict([:n_episodes=>1000])),
    Dict(:ARGS=>(0.5/2, 0.1, 1.0, 32, 2),   :KWARGS=>Dict([:n_episodes=>1000])),
    Dict(:ARGS=>(0.5/16, 0.1, 1.0, 32, 16), :KWARGS=>Dict([:n_episodes=>1000]))
]

function parallel_experiment_runs_args_episodic(experiment_file, args_kwargs_list, n_runs=10; exp_func_name=:run, num_workers=5, make_results=nothing)

    pids = nothing
    if IN_SLURM
        pids = addprocs(SlurmManager(parse(Int, ENV["SLURM_NTASKS"])))
        print("\n")
    else
        println(num_workers, " ", nworkers())
        if nworkers() == 1
            pids = addprocs(num_workers)
        elseif nworkers() < num_workers
            pids = addprocs((num_workers) - nworkers())
        end
    end

    println(nworkers(), " ", pids)

    @everywhere global exp_file=$experiment_file
    @everywhere include(exp_file)
    exp_func = getfield(Main, Symbol(exp_func_name))
    experiment(r, ak) = exp_func(ak[:ARGS]...; seed=1002349+r, ak[:KWARGS]...)

    res_dict = Dict{Int64, Any}()
    futures = Array{Array{Future,1}, 1}()

    for (args_kwargs_idx, args_kwargs) in enumerate(args_kwargs_list)

        in_futures = Array{Future,1}()

        for run in 1:n_runs
            append!(in_futures, [@spawn experiment(run, args_kwargs)])
        end
        if make_results == nothing
            results = zeros(n_runs, args_kwargs[:KWARGS][:n_episodes])
        else
            results = make_results(n_runs, args_kwargs)
        end

        append!(futures, [in_futures])

        res_dict[args_kwargs_idx] = Dict([
            "args"=>args_kwargs,
            "results"=>results
        ])
    end

    for (args_kwargs_idx, args_kwargs) in enumerate(args_kwargs_list)
        for run in 1:n_runs
            res_dict[args_kwargs_idx]["results"][run,:] .= fetch(futures[args_kwargs_idx][run])
        end
    end

    return res_dict
end

function parallel_experiment_runs(experiment_file, args_kwargs_list, id, n_runs=10)

end

function experiment(experiment_file, arg_file, id, r=1)

end


function generate_posit_h(io, posit_defs)
  #generates "posit.h" based on the posit_definitions
  write(io, """
    #ifndef __POSIT_H
    #define __POSIT_H

    #include <stdint.h>
    #include <math.h>
    #include <stdbool.h>
    """)

  #define a data structure that holds the environment settings and extern it.
  write(io, """

  typedef struct {
    bool nanmode;
    bool underflows;
  } environment_t;

  extern "C" environment_t POSIT_ENV;

  extern "C" void set_nanmode(bool nanmode);
  extern "C" void set_underflow(bool underflows);
  extern "C" bool get_nanmode();
  extern "C" bool get_underflow();

  """)

  #create literal constants.
  for n in sort(collect(keys(posit_defs)))
    T = top_bits(n)
    write(io,"""
    #define P$(n)INF      $(c_literal(T))
    #define P$(n)MAXREAL  $(c_literal(T - one(T)))
    #define P$(n)MINREAL  $(c_literal(T | one(T)))
    #define P$(n)POSSMALL $(c_literal(one(T)))
    #define P$(n)NEGSMALL $(c_literal(-one(T)))
    #define P$(n)ZER      $(c_literal(zero(T)))
    """)
  end
  write(io, "\n")

  for n in sort(collect(keys(posit_defs)))
    for es in posit_defs[n]
      write(io, "typedef union {uint$(n)_t udata; int$(n)_t sdata;} p$(n)e$(es)_t ;\n")
    end
    write(io, "\n")
  end

  for n in sort(collect(keys(posit_defs)))
    for es in posit_defs[n]
      write(io, "typedef struct {uint64_t *cache; uint64_t exponent; uint64_t leading_ones;} fdp_cache$(n)e$(es)_t;\n")
    end
    write(io, "\n")
  end

  write(io, "#endif")  #to close out the header guard #define

end


mutable struct SPPointMassConstraintOut <: SumProductRule{PointMassConstraint} end
outboundType(::Type{SPPointMassConstraintOut}) = Message{PointMass}
isApplicable(::Type{SPPointMassConstraintOut}, input_types::Vector{<:Type}) = (length(input_types) == 1) && (input_types[1] == Nothing)


include("ismv.jl")
include("lbv.jl")
include("pdf.jl")
include("sharp.jl")


"""
    Gibbs(n_iters, alg_1, alg_2)

Compositional MCMC interface.

Example:
```julia
alg = Gibbs(1000, HMC(1, 0.2, 3, :v1), PG(20, 1, :v2))
```
"""
mutable struct Gibbs{A} <: InferenceAlgorithm
    n_iters   ::  Int     # number of Gibbs iterations
    algs      ::  A   # component sampling algorithms
    thin      ::  Bool    # if thinning to output only after a whole Gibbs sweep
    gid       ::  Int
end
Gibbs(n_iters::Int, algs...; thin=true) = Gibbs(n_iters, algs, thin, 0)
Gibbs(alg::Gibbs, new_gid) = Gibbs(alg.n_iters, alg.algs, alg.thin, new_gid)

const GibbsComponent = Union{Hamiltonian,MH,PG}

function Sampler(alg::Gibbs)
    n_samplers = length(alg.algs)
    samplers = Array{Sampler}(undef, n_samplers)

    space = Set{Symbol}()

    for i in 1:n_samplers
        sub_alg = alg.algs[i]
        if isa(sub_alg, GibbsComponent)
            samplers[i] = Sampler(typeof(sub_alg)(sub_alg, i))
        else
            @error("[Gibbs] unsupport base sampling algorithm $alg")
        end
        space = union(space, sub_alg.space)
    end

    # Sanity check for space
    @assert issubset(Turing._compiler_[:pvars], space) "[Gibbs] symbols specified to samplers ($space) doesn't cover the model parameters ($(Turing._compiler_[:pvars]))"

    if Turing._compiler_[:pvars] != space
        @warn("[Gibbs] extra parameters specified by samplers don't exist in model: $(setdiff(space, Turing._compiler_[:pvars]))")
    end

    info = Dict{Symbol, Any}()
    info[:samplers] = samplers

    Sampler(alg, info)
end

function sample(
                model::Function,
                alg::Gibbs;
                save_state=false,         # flag for state saving
                resume_from=nothing,      # chain to continue
                reuse_spl_n=0             # flag for spl re-using
                )

    # Init the (master) Gibbs sampler
    spl = reuse_spl_n > 0 ? resume_from.info[:spl] : Sampler(alg)

    @assert typeof(spl.alg) == typeof(alg) "[Turing] alg type mismatch; please use resume() to re-use spl"

    # Initialize samples
    sub_sample_n = []
    for sub_alg in alg.algs
        if isa(sub_alg, GibbsComponent)
            push!(sub_sample_n, sub_alg.n_iters)
        else
            @error("[Gibbs] unsupport base sampling algorithm $alg")
        end
    end

    # Compute the number of samples to store
    n = reuse_spl_n > 0 ? reuse_spl_n : alg.n_iters
    sample_n = n * (alg.thin ? 1 : sum(sub_sample_n))

    # Init samples
    time_total = zero(Float64)
    samples = Array{Sample}(undef, sample_n)
    weight = 1 / sample_n
    for i = 1:sample_n
        samples[i] = Sample(weight, Dict{Symbol, Any}())
    end

    # Init parameters
    varInfo = if resume_from == nothing
        vi_ = VarInfo()
        Base.invokelatest(model, vi_, HamiltonianRobustInit())
        vi_
    else
        resume_from.info[:vi]
    end

    n = spl.alg.n_iters; i_thin = 1

    # Gibbs steps
    PROGRESS[] && (spl.info[:progress] = ProgressMeter.Progress(n, 1, "[Gibbs] Sampling...", 0))
    for i = 1:n
        @debug "Gibbs stepping..."

        time_elapsed = zero(Float64)
        lp = nothing; epsilon = nothing; lf_num = nothing

        for local_spl in spl.info[:samplers]
            last_spl = local_spl
      # PROGRESS[] && haskey(spl.info, :progress) && (local_spl.info[:progress] = spl.info[:progress])

            @debug "$(typeof(local_spl)) stepping..."

            if isa(local_spl.alg, GibbsComponent)
                if isa(local_spl.alg, Hamiltonian)  # clean cache
                    local_spl.info[:grad_cache] = Dict{UInt64,Vector}()
                    local_spl.info[:reverse_diff_cache] = Dict()
                end

                for _ = 1:local_spl.alg.n_iters
                    @debug "recording old θ..."
                    time_elapsed_thin = @elapsed varInfo = step(model, local_spl, varInfo, i==1)

                    if ~spl.alg.thin
                        samples[i_thin].value = Sample(varInfo).value
                        samples[i_thin].value[:elapsed] = time_elapsed_thin
                        if ~isa(local_spl.alg, Hamiltonian)
                            # If statement below is true if there is a HMC component which provides lp and epsilon
                            if lp != nothing samples[i_thin].value[:lp] = lp end
                            if epsilon != nothing samples[i_thin].value[:epsilon] = epsilon end
                            if lf_num != nothing samples[i_thin].value[:lf_num] = lf_num end
                        end
                        i_thin += 1
                    end
                    time_elapsed += time_elapsed_thin
                end

                if isa(local_spl.alg, Hamiltonian)
                    lp = getlogp(varInfo)
                    epsilon = local_spl.info[:wum][:ϵ][end]
                    lf_num = local_spl.info[:lf_num]
                end
            else
                @error("[Gibbs] unsupport base sampler $local_spl")
            end
        end

        time_total += time_elapsed

        if spl.alg.thin
            samples[i].value = Sample(varInfo).value
            samples[i].value[:elapsed] = time_elapsed
            # If statement below is true if there is a HMC component which provides lp and epsilon
            if lp != nothing samples[i].value[:lp] = lp end
            if epsilon != nothing samples[i].value[:epsilon] = epsilon end
            if lf_num != nothing samples[i].value[:lf_num] = lf_num end
        end

        if PROGRESS[]
            if haskey(spl.info, :progress)
                ProgressMeter.update!(spl.info[:progress], spl.info[:progress].counter + 1)
            end
        end
    end

    @info("[Gibbs] Finished with")
    @info("  Running time    = $time_total;")

    if resume_from != nothing   # concat samples
        pushfirst!(samples, resume_from.value2...)
    end
    c = Chain(0, samples)       # wrap the result by Chain

    if save_state               # save state
        save!(c, spl, model, varInfo)
    end

    return c
end


# depends on abstract type SurrogateConfig & combinable( :: SurrogateConfig )
using Lazy: @forward

Broadcast.broadcastable( objf::AbstractObjective ) = Ref( objf );

include("BatchObjectiveFunction.jl");

# MANDATORY METHODS
"A general constructor."
function _wrap_func( :: T, fn :: Function, model_cfg :: SurrogateConfig,
    n_vars :: Int, n_out :: Int ) :: T where {T<:Type{<:AbstractObjective}}
    nothing
end

"Return a function that evaluates an objective at an **unscaled** site."
eval_handle(::AbstractObjective) :: Function = nothing

num_vars( :: AbstractObjective ) :: Int = nothing

"Number of calls to the original objective function."
num_evals( :: AbstractObjective ) :: Int = 0
"Set evaluation counter to `N`."
num_evals!( :: AbstractObjective, N :: Int) :: Nothing = nothing

num_outputs( objf :: AbstractObjective ) :: Int = nothing

model_cfg( objf :: AbstractObjective ) :: SurrogateConfig = nothing

# DERIVED methods and defaults

"Evaluate the objective at unscaled site `x`. and increase counter."
function eval_objf(objf :: AbstractObjective, x :: Vec )
    inc_evals!(objf)
    eval_handle(objf)(x)
end

#=
function Broadcast.broadcasted( ::typeof(eval_objf), objf :: AbstractObjective, X̂ :: VecVec)
    inc_evals!(objf, length(X̂))
    eval_handle(objf).(X̂)
end
=#

"(Soft) upper bound on the number of function calls. "
max_evals( objf :: AbstractObjective) = max_evals( model_cfg(objf) );

"Increase evaluation count by `N`"
function inc_evals!( objf :: AbstractObjective, N :: Int = 1 )
    num_evals!( objf, num_evals(objf) + N )
end

#Does the model type allow for combination with models of same type?
combinable( objf :: AbstractObjective ) = combinable( model_cfg(objf) );

#Can objectives of different types be combined?
combinable( ::Type{<:AbstractObjective}, ::Type{<:AbstractObjective} ) = false;

function combinable( objf1 :: T, objf2 :: F ) where {T<:AbstractObjective, F<:AbstractObjective}
    return ( combinable( T,F ) &&
        combinable( objf1 ) && combinable( objf2 ) &&
        isequal(model_cfg( objf1 ), model_cfg( objf2 ))
    )
end

# generic combine function for abstract objectives
"Combine two objectives. Only needed if `combinable` can return true."
function combine( objf1 :: T, objf2 :: F ) where{T<:AbstractObjective, F<:AbstractObjective}
    new_fn = combine( eval_handle( objf1 ), eval_handle( objf2 ) );
    n_out = num_outputs( objf1 ) + num_outputs( objf2 );
    new_config = combine( model_cfg(objf1), model_cfg(objf2) )
    return _wrap_func( T, new_fn, new_config, num_vars(objf1), n_out );
end

################
# A wrapper around `AbstractObjective`s that
# ensures a certain output type `T` when evaluating
struct OutTypeWrapper{T, O <: AbstractObjective} <: AbstractObjective
    objf :: O
end

function OutTypeWrapper( inner_objf :: O, T :: Type{<:Vec} = Vector{Float64}) where O<:AbstractObjective
    return OutTypeWrapper{T,O}(inner_objf)
end

# force a new outtype if input is already wrapped
function OutTypeWrapper( wrapped :: OutTypeWrapper{T,O}, F :: Type{<:Vec} = Vector{Float64}) where{T,O}
    return OutTypeWrapper{F,O}(wrapped.objf)
end

@forward OutTypeWrapper.objf num_vars, num_evals, num_evals!, num_outputs
@forward OutTypeWrapper.objf model_cfg, max_evals, inc_evals!, combinable

function eval_handle( wrapped_objf :: OutTypeWrapper{T,O} ) where {T,O}
    return x -> convert(T, eval_handle(wrapped_objf.objf(x)))
end

"A general constructor."
function _wrap_func( :: Type{<:OutTypeWrapper{T,O}}, fn :: Function, model_cfg :: SurrogateConfig,
    n_vars :: Int, n_out :: Int ) where {T, O}
    inner = _wrap_func( O, fn, model_cfg, n_vars, n_out )
    return OutTypeWrapper{T, typeof(inner)}(inner)
end

function combinable( ::Type{OutTypeWrapper{T1,O1}}, :: Type{OutTypeWrapper{T2,O2}} ) where{T1,T2,O1,O2}
    return combinable(O1,O2) && promote_type(T1,T2) <: Vec
end

# generic combine function for abstract objectives
function combine( wrapped1 :: OutTypeWrapper{T1,O1}, wrapped2 :: OutTypeWrapper{T2,O2} ) where {T1,T2,O1,O2}
    new_fn = combine( eval_handle( wrapped1.objf ), eval_handle( wrapped2.objf ) );
    n_out = num_outputs( wrapped1.objf ) + num_outputs( wrapped2.objf );
    new_config = combine( model_cfg(wrapped1.objf), model_cfg(wrapped2.objf) )
    T = promote_type(T1,T2)
    return _wrap_func( T, new_fn, new_config, num_vars(wrapped1.objf), n_out );
end


export residual, residual!, jac_residual, jac_structure_residual, jac_structure_residual!
export jac_coord_residual!, jac_coord_residual, jprod_residual, jprod_residual!
export jtprod_residual, jtprod_residual!, jac_op_residual, jac_op_residual!
export hess_residual, hess_structure_residual, hess_structure_residual!
export hess_coord_residual!, hess_coord_residual, jth_hess_residual
export hprod_residual, hprod_residual!, hess_op_residual, hess_op_residual!

mutable struct NLSProjector{T, S} <: AbstractNLSModel{T, S}
  model::AbstractNLSModel{T, S}
  meta::NLPModelMeta{T, S}
  nls_meta::NLSMeta{T, S}
  counters::NLSCounters
  d::S
  function NLSProjector(
    model::AbstractNLSModel{T, S},
    d::S;
    name = "Projection over $(model.meta.name)",
  ) where {T, S}
    nvar = length(d)
    x0 = d #fill!(S(undef, nvar), zero(T))
    meta = NLPModelMeta{T, S}(nvar, x0 = x0, name = name)
    nls_meta = NLSMeta{T, S}(nvar, nvar, nnzj = nvar, nnzh = 0)
    return new{T, S}(model, meta, nls_meta, NLSCounters(), d)
  end
end

function NLPModels.residual!(model::NLSProjector, x, Fx)
  Fx .= x - model.d
  return Fx
end

function NLPModels.jac_structure_residual!(
  model::NLSProjector,
  rows::AbstractVector{<:Integer},
  cols::AbstractVector{<:Integer},
)
  @lencheck model.nls_meta.nnzj rows cols
  for i = 1:(model.meta.nvar)
    rows[i] = i
    cols[i] = i
  end
  return rows, cols
end

function NLPModels.jac_coord_residual!(
  model::NLSProjector{T, S},
  x::AbstractVector,
  vals::AbstractVector,
) where {T, S}
  @lencheck model.meta.nvar x
  @lencheck model.nls_meta.nnzj vals
  increment!(model, :neval_jac_residual)
  vals .= one(T)
  return vals
end

function NLPModels.jprod_residual!(
  model::NLSProjector,
  x::AbstractVector,
  v::AbstractVector,
  Jv::AbstractVector,
)
  @lencheck model.meta.nvar x v
  @lencheck model.nls_meta.nequ Jv
  increment!(model, :neval_jprod_residual)
  Jv .= v
  return Jv
end

function NLPModels.jtprod_residual!(
  model::NLSProjector,
  x::AbstractVector,
  v::AbstractVector,
  Jtv::AbstractVector,
)
  @lencheck model.meta.nvar x Jtv
  @lencheck model.nls_meta.nequ v
  increment!(model, :neval_jtprod_residual)
  Jtv .= v
  return Jtv
end

function NLPModels.hess_structure_residual!(::NLSProjector, rows, cols)
  return rows, cols
end

function NLPModels.hess_coord_residual!(::NLSProjector, x, v, vals)
  return vals
end

function NLPModels.hprod_residual!(::NLSProjector, x, i, v, Hiv)
  return Hiv
end

for meth in [
  :cons!,
  :jth_congrad!,
  :jth_sparse_congrad,
  :jac_structure!,
  :jac_coord!,
  :jprod!,
  :jtprod!,
  :jth_hess_coord!,
  :jth_hprod!,
  :ghjvprod!,
  :hess_structure!,
  :hess_coord!,
  :hprod!,
]
  premeth = Symbol("NLPModels.", meth)
  @eval begin
    $premeth(model::NLSProjector, args...; kwargs...) = $meth(model.model, args...; kwargs...)
  end
end


"""
    module Simulate_FallingBall4

Models and simulates a sphere that is falling downwards as 6 dof body
(described by quaternions), **with** collision handling.

See also models

- [`Simulate_FallingBall1`](@ref) (falling ball constrained with prismatic joint),
- [`Simulate_FallingBall2`](@ref) (falling ball constrained with prismatic joint and collision handling).
"""
module Simulate_FallingBall4

using  Modia3D
import Modia3D.ModiaMath

# Solids
groundMaterial  = Modia3D.Material(color="DarkGreen", transparency=0.5)
ballMaterial    = Modia3D.Material(color="LightBlue", transparency=0.5)
contactMaterial = Modia3D.ContactMaterialElastic(mu1=0.5)

@assembly FallingBall4(;h_start = 1.0, D=0.1) begin
   world  = Object3D(Modia3D.CoordinateSystem(0.5))
   ground = Object3D(world, Modia3D.Solid(Modia3D.SolidBox(2.0,0.1,2.0), nothing, groundMaterial,
                                          contactMaterial=contactMaterial), r = [1.0, -0.05, 0.0])
   ball   = Object3D(world, Modia3D.Solid(Modia3D.SolidSphere(D), "Aluminium", ballMaterial,
                                          contactMaterial=contactMaterial),
                     fixed=false, r=[0.5,h_start,0.0], v_start=[0.2,0,0], visualizeFrame=true)
end

fallingBall = FallingBall4()

# Modia3D.visualizeAssembly!( fallingBall )

simulationModel = Modia3D.SimulationModel( fallingBall, stopTime=2.5, scaleConstraintsAtEvents = true)
#Modia3D.print_ModelVariables(simulationModel)
result = ModiaMath.simulate!(simulationModel, log=true)

ModiaMath.plot(result, ["ball.r[2]"; "ball.v[1]"; "ball.v[2]"; "ball.v[3]"])

println("... success of Simulate_FallingBall4.jl!")

end


"Connection points at which energy balance is upheld, has `axis` and `addto!`"
struct Bus <: Component
    model::AbstractEnergyModel
    class::Symbol
end

"Represents a synchronuous zone. Induced by determine_subnetworks!(model)"
struct SubNetwork <: Component
    model::AbstractEnergyModel
    class::Symbol
    buses::Axis
end

mutable struct EnergyModel{MT <: ModelType, TF <: PM.AbstractPowerFormulation} <: AbstractEnergyModel
    devices::Dict{Symbol,Device}
    subnetworks::Dict{Symbol,SubNetwork}
    buses::Dict{Symbol,Bus}
    axes::Dict{Symbol,Axis}
    data::AbstractData
    parent::Union{AbstractEnergyModel,Nothing}
    jumpmodel::Union{JuMP.AbstractModel,Nothing}
    jumpnames::Bool
    jumpobjects::Dict{Symbol,Dict{Symbol}{Any}}
end

ensure_optimizerfactory(opt::JuMP.OptimizerFactory) = opt
ensure_optimizerfactory(::Type{T}) where T <: MOI.AbstractOptimizer = with_optimizer(T)

function EnergyModel(::Type{MT}, ::Type{TF}, data::AbstractData;
                     parent=nothing, jumpmodel=nothing, jumpnames=true, optimizer=nothing) where
        {MT <: ModelType, TF <: PM.AbstractPowerFormulation}

    if isnothing(jumpmodel)
        jumpmodel = isnothing(optimizer) ? JuMP.Model() : JuMP.Model(ensure_optimizerfactory(optimizer))
    end

    EnergyModel{MT,TF}(Dict{Symbol,Device}(),
                       Dict{Symbol,SubNetwork}(),
                       Dict{Symbol,Bus}(),
                       Dict{Symbol,Axis}(),
                       data,
                       parent,
                       jumpmodel,
                       jumpnames,
                       Dict{Symbol, Dict{Symbol}{Any}}())
end

function EnergyModel(other::EnergyModel{MT,TF};
                     modeltype=nothing, transmissionform=nothing,
                     jumpnames=nothing, kwargs...) where
        {MT <: ModelType, TF <: PM.AbstractPowerFormulation}

    if isnothing(modeltype) modeltype = MT end
    if isnothing(transmissionform) transmissionform = TF end

    kwargs = Dict{Symbol,Any}(kwargs)
    get!(kwargs, :jumpnames, other.jumpnames)

    model = EnergyModel(modeltype, transmissionform, other.data; kwargs...)

    merge!(model.buses, Dict(k=>withmodel(v, model) for (k,v) in other.buses))
    merge!(model.devices, Dict(k=>withmodel(v, model) for (k,v) in other.devices))
    merge!(model.axes, other.axes)

    determine_subnetworks!(model)

    model
end

EnergyModel(; kwargs...) = EnergyModel(Data(); kwargs...)
EnergyModel(filename; kwargs...) = EnergyModel(load(filename); kwargs...)
EnergyModel(data::AbstractData; modeltype=ExpansionModel, transmissionform=PM.DCPlosslessForm, kwargs...) = load(EnergyModel(modeltype, transmissionform, data; kwargs...))

function load(m::EnergyModel)
    for (class, T) in components(m.data)
        push!(m, T(m, class))
    end
    determine_subnetworks!(m)
    m
end

Base.push!(m::EnergyModel, d::Device) = (m.devices[d.class] = d)
Base.push!(m::EnergyModel, c::SubNetwork) = (m.subnetworks[c.class] = c)
Base.push!(m::EnergyModel, c::Bus) = (m.buses[c.class] = c)
Base.push!(m::EnergyModel, ax::Axis{name}) where name = (m.axes[name] = ax)

set_snapshots!(m::EnergyModel, snapshots::AbstractArray) =
    push!(m, Axis{:snapshots}(collect(snapshots)))

with_model(d::Device, DF::DeviceFormulation) = with_formulation(d, DF)
with_model(d::Device, ::Type{T}) where T <: Device = model()
function set_formulation!(m::EnergyModel, changes::Pair...)
    for (selector, newmodel) in changes
        for d in devices(m, selector)
            push!(m, with_formulation(d, newmodel))
        end
    end
end

store_results!(m::EnergyModel) = store_results!(m.data, m)

devices(m::EnergyModel) = values(m.devices)
devices(sn::SubNetwork) = devices(model(sn))

devices(m::EnergyModel, T::Type{<:Device}) = (d for d = devices(m) if isa(d, T))
devices(sn::SubNetwork, T::Type{<:Device}) = (d for d = devices(sn) if isa(d, T))

buses(m::EnergyModel) = values(m.buses)
subnetworks(m::EnergyModel) = values(m.subnetworks)

JuMP.optimize!(m::EnergyModel; kwargs...) = optimize!(m.jumpmodel; kwargs...)

Base.show(io::IO, m::EnergyModel) = print(io, typeof(m), " with ", length(axis(m, Bus)), " buses and ", length(m.devices), " devices")
Base.show(io::IO, c::Component) = print(io, typeof(c), "($(c.class))")
function Base.show(io::IO, c::Device{DF}) where DF <: DeviceFormulation
    print(io, typeof(c), "($(c.class))")
    demoted_DF = demote_formulation(modeltype(model(c)), DF)
    demoted_DF != DF && print(io, " (demoted to $demoted_DF)")
end

# function Base.show(io::IO, ::MIME"text/plain", d::Device)
#     println(io, d, " with ")
#     println(io, "* ", length(d.vars), " variables")
#     print(io, "* ", length(d.constrs), " constraints")
# end

model(m::EnergyModel) = m
model(c::Component) = c.model
modeltype(m::EnergyModel{MT,TF}) where {MT <: ModelType, TF <: PM.AbstractPowerFormulation} = MT
transmissionform(m::EnergyModel{MT,TF}) where {MT <: ModelType, TF <: PM.AbstractPowerFormulation} = TF

naming(c::Component) = c.class
naming(c::Component, args...) = Symbol(naming(c), flatten((:(::), a) for a=args)...)

Base.findall(pred::Base.Fix2{typeof(in), <:Axis}, d::Device) = intersect((findall(pred, get(d, attr, axis(d))) for attr = busattributes(d))...)
Base.findall(pred::Base.Fix2{typeof(in), <:Axis}, c::Bus) = findall(pred, axis(c).val)

function Base.getindex(m::EnergyModel, class::Symbol)
    for t = (:devices, :buses, :subnetworks)
        haskey(getfield(m, t), class) && return getfield(m, t)[class]
    end
    throw(KeyError(class))
end
Base.getindex(m::EnergyModel, ::Type{T}) where T <: Device = ContainerView(m, Dict{Symbol,T}(d.class=>d for d=devices(m, T)))
Base.getindex(m::EnergyModel, ::Type{Bus}) = ContainerView(m, Dict{Symbol,Bus}(d.class=>d for d=buses(m)))

Base.getindex(sn::SubNetwork, ::Type{T}) where T <: Device = SubContainerView(model(sn), Dict{Symbol,T}(d.class=>d for d=devices(sn, T)), sn.buses)
Base.getindex(sn::SubNetwork, ::Type{Bus}) = SubContainerView(model(sn), Dict{Symbol,Bus}(d.class=>d for d=buses(model(sn))), sn.buses)

"""
    Base.get(c::Component, attr::Symbol)
    Base.get(c::Component, attr::Symbol, axes...)

For the component `c` get the attribute `attr`, which is either the

- JuMP variable or constraint, or the
- parameter data from the `Data` object for the `class` of the component

If `axes` are provided they must be `Symbol`s or AxisArrays.Axis objects, to
specify which dimensions need to be added flexibly, ie by wrapping in a
WrappedArray.

If you want a JuMP object or a parameter specifically, use `getjump` or
`getparam` instead.

# Examples

```julia
c = m[:onwind]

get(c, :p)   # gets the JuMP variable

get(c, :p_max_pu) # gets the plant availability
                  # (which might be static or a timeseries)

get(c, :p_max_pu, axis(c), :snapshots) # plant availability as timeseries
```
"""
function Base.get(c::Component, attr::Symbol)
    ret = getjump(c, attr)
    !isnothing(ret) ? ret : getparam(c, attr)
end
Base.get(c::Component, attr::Symbol, axes...) = WrappedArray(get(c, attr), axes...)

Base.getindex(c::Component, attr::Symbol) = get(c, attr)

getjump(m::EnergyModel, c::Component, attr::Symbol) = getjump(m, c.class, attr)
function getjump(m::EnergyModel, class::Symbol, attr::Symbol)
    ret = get(get!(m.jumpobjects, class, Dict{Symbol}{Any}()), attr, nothing)
    if !isnothing(ret)
        AxisArray(ret)
    elseif !isnothing(m.parent)
        getjump(m.parent, class, attr)
    else
        nothing
    end
end
getjump(c::Component, attr::Symbol) = getjump(model(c), c, attr)
JuMP.getvalue(c::Component, attr::Symbol) = getvalue.(getjump(c, attr))
JuMP.getdual(c::Component, attr::Symbol) = getdual.(getjump(c, attr))
getparam(c::Component, attr::Symbol) = get(model(c).data, c, attr)

function axis(m::EnergyModel, name::Symbol)
    ax = get(m.axes, name, nothing)
    !isnothing(ax) ? ax : axis(m.data, name)
end
axis(m::EnergyModel, args...) = axis(m.data, args...)
axis(c::Component) = axis(model(c), c)
axis(c::Component, attr) = axis(model(c), attr)

axis(m::EnergyModel, T::Type{<:Component}) = axis(m[T])
axis(m::SubNetwork, T::Type{<:Component}) = axis(m[T])

# Could be specialized to not have to retrieve the whole axis (on the other
# hand, the axis should be cached, anyway)
Base.length(c::Component) = length(axis(c))

issimple(c::Component) = issimple(model(c).data, naming(c))
issimple(::SubNetwork) = false

withmodel(c::T, m::EnergyModel) where T <: Component = T(m, c.class)

add!(m::EnergyModel, ::Type{T}, class::Symbol; kwargs...) where T <: Component = add!(m, T, Axis{class}([class]); kwargs...)
add!(m::EnergyModel, ::Type{T}, class::Symbol, names; kwargs...) where T <: Component = add!(m, T, Axis{class}(names); kwargs...)
function add!(m::EnergyModel, ::Type{T}, ax::Axis{class}; suffix=nothing, parameters...) where {class, T <: Component}
    if !isnothing(suffix)
        ax = Axis{class}(Symbol.(ax.val, suffix))
    end

    push!(m.data, T, ax; parameters...)
    push!(m, T(m, class))
end


@testset "Evolution" begin

    struct MockCand <: AbstractCandidate
        val::Real
    end
    NaiveGAflux.fitness(c::MockCand) = c.val

    struct DummyCand <: AbstractCandidate end

    @testset "EliteSelection" begin
        pop = MockCand.([3, 7, 4, 5, 9, 0])
        @test fitness.(evolve(EliteSelection(3), pop)) == [9, 7, 5]
    end

    @testset "SusSelection" begin
        pop = MockCand.(0.1:0.1:0.8)
        @test fitness.(evolve(SusSelection(4, NoOpEvolution(), MockRng([0.05])), pop)) == [0.1, 0.4, 0.6, 0.7]
    end

    @testset "TournamentSelection" begin
        pop = MockCand.(1:10)
        Random.shuffle(::MockRng, a::AbstractVector) = reverse(a)
        Random.rand(r::MockRng, n::Int) = map(i -> rand(r), 1:n)

        @test fitness.(evolve(TournamentSelection(4, 2, 1, NoOpEvolution(), MockRng([0.05])), pop)) == [10, 8, 6, 4]
        @test fitness.(evolve(TournamentSelection(8, 3, 0.7, NoOpEvolution(), MockRng([0.01, 0.9, 0.01])), pop)) == [9, 9, 7, 4, 2, 8, 5, 2]

        @test fitness.(evolve(TournamentSelection(20, 2, 1, NoOpEvolution(), MockRng([0.01])), pop)) == repeat(10:-2:2, 4)

        # Edge case, nselect*k % popsize == 0
        @test fitness.(evolve(TournamentSelection(5, 2, 1, NoOpEvolution(), MockRng([0.01])), pop)) == 10:-2:2
        # Edge case, nselect*k % popsize == 1
        @test fitness.(evolve(TournamentSelection(2, 2, 1, NoOpEvolution(), MockRng([0.01])), MockCand.(1:3))) == [3, 2]
    end

    @testset "CombinedEvolution" begin
        pop = [DummyCand() for i in 1:5]
        @test evolve(CombinedEvolution(NoOpEvolution(), NoOpEvolution()), pop) == vcat(pop,pop)
    end

    @testset "EvolutionChain" begin
        dfitness = EvolveCandidates(mc -> MockCand(2 * fitness(mc)))
        sqrfitness = EvolveCandidates(mc -> MockCand(fitness(mc)^2))
        @test evolve(EvolutionChain(), MockCand.(1:5)) == MockCand.(1:5)
        @test evolve(EvolutionChain(dfitness, dfitness, sqrfitness), MockCand.(1:10)) == MockCand.((4:4:40).^2)
    end

    @testset "PairCandidates" begin
        struct PairConsumer <: AbstractEvolution end
        nseen = 0
        function NaiveGAflux._evolve(::PairConsumer, pop)
            nseen = length(pop)
            @test typeof(pop) == Vector{Tuple{MockCand, MockCand}}
            return pop
        end

        @test evolve(PairCandidates(NoOpEvolution()), []) == []
        pop = MockCand.(1:9)
        @test evolve(PairCandidates(PairConsumer()), pop) == pop
        @test nseen == 5
        pop = MockCand.(1:4)
        @test evolve(PairCandidates(PairConsumer()), pop) == pop
        @test nseen == 2
    end

    @testset "ShuffleCandidates" begin
        pop = MockCand.(1:10)
        @test evolve(ShuffleCandidates(MersenneTwister(1)), pop) != pop
    end

    @testset "EvolveCandidates" begin
        pop = MockCand.([2,3,5,9])
        @test fitness.(evolve(EvolveCandidates(mc -> MockCand(2mc.val)), pop)) == 2fitness.(pop)
    end

end