# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
"""
    struct ElectricTariff

- data for electric tariff in reopt model
- can be defined using custom rates or URDB rate
- very similar to the URDB struct but includes export rates and bins
"""
struct ElectricTariff
    energy_rates::AbstractArray{Float64, 2} # gets a second dim with tiers
    energy_tier_limits::AbstractArray{Float64,1}
    n_energy_tiers::Int

    monthly_demand_rates::AbstractArray{Float64, 2} # gets a second dim with tiers
    time_steps_monthly::AbstractArray{AbstractArray{Int64,1},1}  # length = 0 or 12
    monthly_demand_tier_limits::AbstractArray{Float64,1}
    n_monthly_demand_tiers::Int

    tou_demand_rates::AbstractArray{Float64, 2} # gets a second dim with tiers
    tou_demand_ratchet_time_steps::AbstractArray{AbstractArray{Int64,1},1}  # length = n_tou_demand_ratchets
    tou_demand_tier_limits::AbstractArray{Float64,1}
    n_tou_demand_tiers::Int

    demand_lookback_months::AbstractArray{Int,1}
    demand_lookback_percent::Float64
    demand_lookback_range::Int

    fixed_monthly_charge::Float64
    annual_min_charge::Float64
    min_monthly_charge::Float64

    export_rates::Dict{Symbol, AbstractArray}
    export_bins::AbstractArray{Symbol,1}

    coincident_peak_load_active_time_steps::AbstractVector{AbstractVector{Int64}}
    coincident_peak_load_charge_per_kw::AbstractVector{Float64}
    coincpeak_periods::AbstractVector{Int64}
end

"""
    ElectricTariff

ElectricTariff constructor
```julia
function ElectricTariff(;
    urdb_label::String="",
    urdb_response::Dict=Dict(),
    urdb_utility_name::String="",
    urdb_rate_name::String="",
    year::Int=2020,
    time_steps_per_hour::Int=1,
    NEM::Bool=false,
    wholesale_rate::T1=nothing,
    export_rate_beyond_net_metering_limit::T2=nothing,
    monthly_energy_rates::Array=[],
    monthly_demand_rates::Array=[],
    blended_annual_energy_rate::S=nothing,
    blended_annual_demand_rate::R=nothing,
    add_monthly_rates_to_urdb_rate::Bool=false,
    tou_energy_rates_per_kwh::Array=[],
    add_tou_energy_rates_to_urdb_rate::Bool=false,
    remove_tiers::Bool=false,
    demand_lookback_months::AbstractArray{Int64, 1}=Int64[],
    demand_lookback_percent::Float64=0.0,
    demand_lookback_range::Int=0,
    coincident_peak_load_active_time_steps::Vector{Vector{Int64}}=[Int64[]],
    coincident_peak_load_charge_per_kw::AbstractVector{<:Real}=Real[]
    ) where {
        T1 <: Union{Nothing, Int, Float64, Array},
        T2 <: Union{Nothing, Int, Float64, Array},
        S <: Union{Nothing, Int, Float64},
        R <: Union{Nothing, Int, Float64}
    }
```

!!! note
    The `NEM` boolean is determined by the ElectricUtility.net_metering_limit_kw. There is no need to pass in a `NEM`
    value.

"""
function ElectricTariff(;
    urdb_label::String="",
    urdb_response::Dict=Dict(),
    urdb_utility_name::String="",
    urdb_rate_name::String="",
    year::Int=2020,
    time_steps_per_hour::Int=1,
    NEM::Bool=false,
    wholesale_rate::T1=nothing,
    export_rate_beyond_net_metering_limit::T2=nothing,
    monthly_energy_rates::Array=[],
    monthly_demand_rates::Array=[],
    blended_annual_energy_rate::S=nothing,
    blended_annual_demand_rate::R=nothing,
    add_monthly_rates_to_urdb_rate::Bool=false,
    tou_energy_rates_per_kwh::Array=[],
    add_tou_energy_rates_to_urdb_rate::Bool=false,
    remove_tiers::Bool=false,
    demand_lookback_months::AbstractArray{Int64, 1}=Int64[],
    demand_lookback_percent::Float64=0.0,
    demand_lookback_range::Int=0,
    coincident_peak_load_active_time_steps::Vector{Vector{Int64}}=[Int64[]],
    coincident_peak_load_charge_per_kw::AbstractVector{<:Real}=Real[]
    ) where {
        T1 <: Union{Nothing, Int, Float64, Array},
        T2 <: Union{Nothing, Int, Float64, Array},
        S <: Union{Nothing, Int, Float64},
        R <: Union{Nothing, Int, Float64}
    }
    # TODO remove_tiers for multinode models
    nem_rate = Float64[]

    energy_tier_limits = Float64[]
    n_energy_tiers = 1
    monthly_demand_tier_limits = Float64[]
    n_monthly_demand_tiers = 1
    tou_demand_tier_limits = Float64[]
    n_tou_demand_tiers = 1
    time_steps_monthly = get_monthly_time_steps(year, time_steps_per_hour=time_steps_per_hour)

    u = nothing
    if !isempty(urdb_response)

        u = URDBrate(urdb_response, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(urdb_label)

        u = URDBrate(urdb_label, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(urdb_utility_name) && !isempty(urdb_rate_name)

        u = URDBrate(urdb_utility_name, urdb_rate_name, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(tou_energy_rates_per_kwh) && length(tou_energy_rates_per_kwh) == 8760*time_steps_per_hour

        tou_demand_rates = Float64[]
        tou_demand_ratchet_time_steps = []
        energy_rates = tou_energy_rates_per_kwh
        monthly_demand_rates = convert(Array{Float64}, monthly_demand_rates)

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end

    elseif !isempty(monthly_energy_rates)

        invalid_args = String[]
        if !(length(monthly_energy_rates) == 12)
            push!(invalid_args, "length(monthly_energy_rates) must equal 12, got length $(length(monthly_energy_rates))")
        end
        if !isempty(monthly_demand_rates) && !(length(monthly_demand_rates) == 12)
            push!(invalid_args, "length(monthly_demand_rates) must equal 12, got length $(length(monthly_demand_rates))")
        end
        if length(invalid_args) > 0
            error("Invalid argument values: $(invalid_args)")
        end

        if isempty(monthly_demand_rates)
            monthly_demand_rates = repeat([0.0], 12)
        end

        tou_demand_rates = Float64[]
        tou_demand_ratchet_time_steps = []
        energy_rates = Real[]
        for m in 1:12
            append!(energy_rates, [monthly_energy_rates[m] for ts in time_steps_monthly[m]])
        end

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end

    elseif !isnothing(blended_annual_energy_rate)

        tou_demand_rates = Float64[]
        tou_demand_ratchet_time_steps = []
        energy_rates = repeat(Real[blended_annual_energy_rate], 8760 * time_steps_per_hour)
        if !isnothing(blended_annual_demand_rate)
            monthly_demand_rates = repeat(Real[blended_annual_demand_rate], 12)
        end
        if isempty(monthly_demand_rates)
            monthly_demand_rates = repeat([0.0], 12)
        end

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end

    else
        error("Creating ElectricTariff requires at least urdb_label, urdb_response, monthly rates, annual rates, or tou_energy_rates_per_kwh.")
    end

    if !isnothing(u)  # use URDBrate

        if NEM
            t = get_tier_with_lowest_energy_rate(u)
            nem_rate = [-0.999 * x for x in u.energy_rates[t,:]]
        end

        energy_rates = u.energy_rates
        energy_tier_limits = u.energy_tier_limits
        n_energy_tiers = u.n_energy_tiers
        users_monthly_demand_rates = copy(monthly_demand_rates)
        monthly_demand_rates = u.monthly_demand_rates
        monthly_demand_tier_limits = u.monthly_demand_tier_limits
        n_monthly_demand_tiers = u.n_monthly_demand_tiers
        tou_demand_rates = u.tou_demand_rates
        tou_demand_tier_limits = u.tou_demand_tier_limits
        n_tou_demand_tiers = u.n_tou_demand_tiers

        if remove_tiers
            energy_rates, monthly_demand_rates, tou_demand_rates = remove_tiers_from_urdb_rate(u)
            energy_tier_limits, monthly_demand_tier_limits, tou_demand_tier_limits =
                Float64[], Float64[], Float64[]
            n_energy_tiers, n_monthly_demand_tiers, n_tou_demand_tiers = 1, 1, 1
        end

        tou_demand_ratchet_time_steps = u.tou_demand_ratchet_time_steps
        demand_lookback_months = u.demand_lookback_months
        demand_lookback_percent = u.demand_lookback_percent
        demand_lookback_range = u.demand_lookback_range
        fixed_monthly_charge = u.fixed_monthly_charge
        annual_min_charge = u.annual_min_charge
        min_monthly_charge = u.min_monthly_charge

        if add_monthly_rates_to_urdb_rate
            if length(monthly_energy_rates) == 12
                for tier in 1:size(energy_rates, 2), mth in 1:12, ts in time_steps_monthly[mth]
                    energy_rates[ts, tier] += monthly_energy_rates[mth]
                end
            end
            if length(users_monthly_demand_rates) == 12
                for tier in 1:size(monthly_demand_rates, 2), mth in 1:12
                    monthly_demand_rates[mth, tier] += users_monthly_demand_rates[mth]
                end
            end
        end

        if add_tou_energy_rates_to_urdb_rate && length(tou_energy_rates_per_kwh) == size(energy_rates, 1)
            for tier in 1:size(energy_rates, 2)
                energy_rates[1:end, tier] += tou_energy_rates_per_kwh
            end
        end
    else
        # need to reshape cost vectors to arrays (2nd dim is for tiers)
        energy_rates = reshape(energy_rates, :, 1)
        monthly_demand_rates = reshape(monthly_demand_rates, :, 1)
        tou_demand_rates = reshape(tou_demand_rates, :, 1)
    end

    #= export_rates
    There are three Export tiers and their associated export rates (negative values):
    1. NEM (Net Energy Metering)
    2. WHL (Wholesale)
    3. EXC (Excess, beyond NEM)

    Only one of NEM and Wholesale can be exported into due to the binary constraints.
    Excess can be exported into in the same time step as NEM.

    Excess is meant to be combined with NEM: NEM export is limited to the total grid purchased energy in a year and some
    utilities offer a compensation mechanism for export beyond the site load.
    The Excess tier is not available with the Wholesale tier.

    - if NEM then set ExportRate[:Nem, :] to energy_rate[tier_with_lowest_energy_rate, :]
    - user can provide either scalar wholesale rate or vector of time_steps,
    =#
    whl_rate = create_export_rate(wholesale_rate, length(energy_rates), time_steps_per_hour)
    if !isnothing(u) && sum(u.sell_rates) < 0
        whl_rate += u.sell_rates
    end
    exc_rate = create_export_rate(export_rate_beyond_net_metering_limit, length(energy_rates), time_steps_per_hour)

    if !NEM & (sum(whl_rate) >= 0)
        export_rates = Dict{Symbol, AbstractArray}()
        export_bins = Symbol[]
    elseif !NEM
        export_bins = [:WHL]
        export_rates = Dict(:WHL => whl_rate)
    elseif (sum(whl_rate) >= 0)
        export_bins = [:NEM]
        export_rates = Dict(:NEM => nem_rate)
        if sum(exc_rate) < 0
            push!(export_bins, :EXC)
            export_rates[:EXC] = exc_rate
        end
    else
        export_bins = [:NEM, :WHL]
        export_rates = Dict(:NEM => nem_rate, :WHL => whl_rate)
        if sum(exc_rate) < 0
            push!(export_bins, :EXC)
            export_rates[:EXC] = exc_rate
        end
    end

    coincpeak_periods = Int64[]
    if !isempty(coincident_peak_load_charge_per_kw)
        coincpeak_periods = collect(1:length(coincident_peak_load_charge_per_kw))
    end

    ElectricTariff(
        energy_rates,
        energy_tier_limits,
        n_energy_tiers,
        monthly_demand_rates,
        time_steps_monthly,
        monthly_demand_tier_limits,
        n_monthly_demand_tiers,
        tou_demand_rates,
        tou_demand_ratchet_time_steps,
        tou_demand_tier_limits,
        n_tou_demand_tiers,
        demand_lookback_months,
        demand_lookback_percent,
        demand_lookback_range,
        fixed_monthly_charge,
        annual_min_charge,
        min_monthly_charge,
        export_rates,
        export_bins,
        coincident_peak_load_active_time_steps,
        coincident_peak_load_charge_per_kw,
        coincpeak_periods
    )
end

function get_tier_with_lowest_energy_rate(u::URDBrate)
    """
    ExportRate should be lowest energy cost for tiered rates.
    Otherwise, ExportRate can be > FuelRate, which leads REopt to export all PV energy produced.
    """
    tier_with_lowest_energy_cost = 1
    if length(u.energy_tier_limits) > 1
        annual_energy_charge_sums = Float64[]
        for etier in u.energy_rates
            push!(annual_energy_charge_sums, sum(etier))
        end
        tier_with_lowest_energy_cost =
            findall(annual_energy_charge_sums .== minimum(annual_energy_charge_sums))[1]
    end
    return tier_with_lowest_energy_cost
end

"""
    function create_export_rate(e::Nothing, N::Int, ts_per_hour::Int=1)
No export rate provided by user: set to 0 dollars/kWh for all time
"""
function create_export_rate(e::Nothing, N::Int, ts_per_hour::Int=1)
    [0 for _ in range(1, stop=N) for ts in 1:ts_per_hour]
end

"""
    function create_export_rate(e::T, N::Int, ts_per_hour::Int=1) where T<:Real
Case for scaler export rate provided -> convert to array of time_steps
"""
function create_export_rate(e::T, N::Int, ts_per_hour::Int=1) where T<:Real
    repeat(float(-1*e), N * ts_per_hour)
end

"""
    function create_export_rate(e::AbstractArray{<:Real, 1}, N::Int, ts_per_hour::Int=1)

Check length of e and upsample if length(e) != N
"""
function create_export_rate(e::AbstractArray{<:Real, 1}, N::Int, ts_per_hour::Int=1)
    Ne = length(e)
    if Ne != Int(N/ts_per_hour) || Ne != N
        @error "Export rates do not have correct number of entries. Must be $(N) or $(Int(N/ts_per_hour))."
    end
    if Ne != N  # upsample
        export_rates = [-1*x for x in e for ts in 1:ts_per_hour]
    else
        export_rates = -1*e
    end
    return export_rates
end

"""
    get_monthly_time_steps(year::Int; time_steps_per_hour=1)

return Array{Array{Int64,1},1}, size = (12,)
"""
function get_monthly_time_steps(year::Int; time_steps_per_hour=1)
    a = Array[]
    i = 1
    for m in range(1, stop=12)
        n_days = daysinmonth(Date(string(year) * "-" * string(m)))
        stop = n_days * 24 * time_steps_per_hour + i - 1
        if m == 2 && isleapyear(year)
            stop -= 24 * time_steps_per_hour  # TODO support extra day in leap years?
        end
        steps = [step for step in range(i, stop=stop)]
        append!(a, [steps])
        i = stop + 1
    end
    return a
end

# TODO use this function only for URDBrate
function remove_tiers_from_urdb_rate(u::URDBrate)
    # tariff args: have to validate that there are no tiers
    if length(u.energy_tier_limits) > 1
        @warn "Energy rate contains tiers. Using the first tier!"
    end
    elec_rates = vec(u.energy_rates[1,:])

    if u.n_monthly_demand_tiers > 1
        @warn "Monthly demand rate contains tiers. Using the last tier!"
    end
    if u.n_monthly_demand_tiers > 0
        demand_rates_monthly = vec(u.monthly_demand_rates[:,u.n_monthly_demand_tiers])
    else
        demand_rates_monthly = vec(u.monthly_demand_rates)  # 0×0 Array{Float64,2}
    end

    if u.n_tou_demand_tiers > 1
        @warn "TOU demand rate contains tiers. Using the last tier!"
    end
    if u.n_tou_demand_tiers > 0
        demand_rates = vec(u.tou_demand_rates[:,u.n_tou_demand_tiers])
    else
        demand_rates = vec(u.tou_demand_rates)
    end

    return elec_rates, demand_rates_monthly, demand_rates
end


"""
    Controls(mitigate, remove, geoeng, adapt)

Create data structure for climate controls.

# Examples
```jldoctest
a = zeros(4);
C = Controls(a, a, a, a);
C.geoeng

# output
4-element Array{Float64,1}:
 0.0
 0.0
 0.0
 0.0
```

See also: [`ClimateModel`](@ref)
"""
mutable struct Controls
    mitigate::Array{Float64,1}
    remove::Array{Float64,1}
    geoeng::Array{Float64,1}
    adapt::Array{Float64,1}
end


#####
##### 'Tupled closure' implementation: 1-tuple, 2-tuple, and then n-tuple by induction
#####

#####
##### Stress divergences
#####

for stress_div in (:∂ⱼ_τ₁ⱼ, :∂ⱼ_τ₂ⱼ, :∂ⱼ_τ₃ⱼ)
    @eval begin
        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any}, clock, U, Ks, args...) =
                    $stress_div(i, j, k, grid, closures[1], clock, U, Ks[1], args...)

        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any}, clock, U, Ks, args...) = (
                    $stress_div(i, j, k, grid, closures[1], clock, U, Ks[1], args...)
                  + $stress_div(i, j, k, grid, closures[2], clock, U, Ks[2], args...))

        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any, <:Any}, clock, U, Ks, args...) = (
                    $stress_div(i, j, k, grid, closures[1], clock, U, Ks[1], args...)
                  + $stress_div(i, j, k, grid, closures[2], clock, U, Ks[2], args...)
                  + $stress_div(i, j, k, grid, closures[3], clock, U, Ks[3], args...))

        @inline $stress_div(i, j, k, grid::AbstractGrid, closures::Tuple, clock, U, Ks, args...) = (
                    $stress_div(i, j, k, grid, closures[1:2], clock, U, Ks[1:2], args...)
                  + $stress_div(i, j, k, grid, closures[3:end], clock, U, Ks[3:end], args...))
    end
end

#####
##### Tracer flux divergences
#####

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any}, c, iᶜ, clock, Ks, args...) =
        ∇_dot_qᶜ(i, j, k, grid, closures[1], c, iᶜ, clock, Ks[1], args...)

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any}, c, iᶜ, clock, Ks, args...) = (
        ∇_dot_qᶜ(i, j, k, grid, closures[1], c, iᶜ, clock, Ks[1], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[2], c, iᶜ, clock, Ks[2], args...))

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple{<:Any, <:Any, <:Any}, c, iᶜ, clock, Ks, args...) = (
        ∇_dot_qᶜ(i, j, k, grid, closures[1], c, iᶜ, clock, Ks[1], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[2], c, iᶜ, clock, Ks[2], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[3], c, iᶜ, clock, Ks[3], args...))

@inline ∇_dot_qᶜ(i, j, k, grid::AbstractGrid, closures::Tuple, c, iᶜ, clock, Ks, args...) = (
        ∇_dot_qᶜ(i, j, k, grid, closures[1:2], c, iᶜ, clock, Ks[1:2], args...)
      + ∇_dot_qᶜ(i, j, k, grid, closures[3:end], c, iᶜ, clock, Ks[3:end], args...))

#####
##### Utilities
#####

with_tracers(tracers, closure_tuple::Tuple) =
    Tuple(with_tracers(tracers, closure) for closure in closure_tuple)

function calculate_diffusivities!(diffusivity_fields_tuple, closure_tuple::Tuple, args...)
    for (α, closure) in enumerate(closure_tuple)
        @inbounds diffusivity_fields = diffusivity_fields_tuple[α]
        calculate_diffusivities!(diffusivity_fields, closure, args...)
    end
    return nothing
end

#####
##### Support for VerticallyImplicitTimeDiscretization
#####

const EC = AbstractTurbulenceClosure{<:ExplicitTimeDiscretization}
const VIC = AbstractTurbulenceClosure{<:VerticallyImplicitTimeDiscretization}

# Filter explicitly-discretized closures.
@inline z_diffusivity(clo::Tuple{<:EC},        iᶜ, Ks, args...) = tuple(0)
@inline z_diffusivity(clo::Tuple{<:VIC},       iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1], iᶜ, Ks[1], args...))
@inline z_diffusivity(clo::Tuple{<:VIC, <:EC}, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1], iᶜ, Ks[1], args...))
@inline z_diffusivity(clo::Tuple{<:EC, <:VIC}, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[2], iᶜ, Ks[2], args...))

@inline z_diffusivity(clo::Tuple{<:VIC, <:VIC}, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1], iᶜ, Ks[1], args...),
                                                                         z_diffusivity(clo[2], iᶜ, Ks[2], args...))

@inline z_diffusivity(clo::Tuple, iᶜ, Ks, args...) = tuple(z_diffusivity(clo[1:2],   iᶜ, Ks[1:2],   args...)...,
                                                           z_diffusivity(clo[3:end], iᶜ, Ks[3:end], args...)...)

@inline z_viscosity(clo::Tuple{<:EC},         Ks, args...) = tuple(0)
@inline z_viscosity(clo::Tuple{<:VIC},        Ks, args...) = tuple(z_viscosity(clo[1], Ks[1], args...))
@inline z_viscosity(clo::Tuple{<:VIC, <:EC},  Ks, args...) = tuple(z_viscosity(clo[1], Ks[1], args...))
@inline z_viscosity(clo::Tuple{<:EC, <:VIC},  Ks, args...) = tuple(z_viscosity(clo[2], Ks[2], args...))

@inline z_viscosity(clo::Tuple{<:VIC, <:VIC}, Ks, args...) = tuple(z_viscosity(clo[1], Ks[1], args...),
                                                                   z_viscosity(clo[2], Ks[2], args...))

@inline z_viscosity(clo::Tuple, Ks, args...) = tuple(z_viscosity(clo[1:2],   Ks[1:2], args...)...,
                                                     z_viscosity(clo[3:end], Ks[3:end], args...)...)

for coeff in (:νᶜᶜᶜ, :νᶠᶠᶜ, :νᶠᶜᶠ, :νᶜᶠᶠ, :κᶜᶜᶠ, :κᶜᶠᶜ, :κᶠᶜᶜ)
    @eval begin
        @inline $coeff(i, j, k, grid, clock, ν::Tuple{C1})     where C1       = $coeff(i, j, k, grid, clock, ν[1])
        @inline $coeff(i, j, k, grid, clock, ν::Tuple{C1, C2}) where {C1, C2} = $coeff(i, j, k, grid, clock, ν[1])   + $coeff(i, j, k, grid, clock, ν[2])
        @inline $coeff(i, j, k, grid, clock, ν::Tuple)                        = $coeff(i, j, k, grid, clock, ν[1:2]) + $coeff(i, j, k, grid, clock, ν[3:end])
    end
end

const ImplicitClosure = AbstractTurbulenceClosure{TD} where TD <: VerticallyImplicitTimeDiscretization
const ExplicitOrNothing = Union{ExplicitTimeDiscretization, Nothing}

@inline combine_time_discretizations(disc) = disc

@inline combine_time_discretizations(::ExplicitTimeDiscretization, ::VerticallyImplicitTimeDiscretization)           = VerticallyImplicitTimeDiscretization()
@inline combine_time_discretizations(::VerticallyImplicitTimeDiscretization, ::ExplicitTimeDiscretization)           = VerticallyImplicitTimeDiscretization()
@inline combine_time_discretizations(::VerticallyImplicitTimeDiscretization, ::VerticallyImplicitTimeDiscretization) = VerticallyImplicitTimeDiscretization()
@inline combine_time_discretizations(::ExplicitTimeDiscretization, ::ExplicitTimeDiscretization)                     = ExplicitTimeDiscretization()

@inline combine_time_discretizations(disc1, disc2, other_discs...) =
    combine_time_discretizations(combine_time_discretizations(disc1, disc2), other_discs...)

@inline time_discretization(closures::Tuple) = combine_time_discretizations(time_discretization.(closures)...)


"""
spin15.jl - T1 optimized pulses
"""

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin15"
const SAVE_PATH = abspath(joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME))

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
# state indices
const STATE1_IDX = 1:HDIM_ISO
const STATE2_IDX = STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO
const INTCONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
const CONTROLS_IDX = INTCONTROLS_IDX[end] + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT
const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const INTGAMMA_IDX = DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + 1
# control indices
const D2CONTROLS_IDX = 1:CONTROL_COUNT
const DT_IDX = D2CONTROLS_IDX[end] + 1:D2CONTROLS_IDX[end] + 1

# model
# DA = decay aware, TO = time optimal
struct Model{DA, TO} <: AbstractModel
    Model(DA::Bool=true, TO::Bool=true) = new{DA, TO}()
end
@inline RD.state_dim(::Model{false, TO}) where TO = ASTATE_SIZE_BASE
@inline RD.state_dim(::Model{true, TO}) where TO = ASTATE_SIZE_BASE + 1
@inline RD.control_dim(::Model{DA, false}) where DA = CONTROL_COUNT
@inline RD.control_dim(::Model{DA, true}) where DA = CONTROL_COUNT + 1

# dynamics
function RD.discrete_dynamics(::Type{RK3}, model::Model{DA, TO}, astate::StaticVector,
                              acontrols::StaticVector, time::Real, dt::Real) where {DA, TO}
    if TO
        dt = acontrols[DT_IDX[1]]^2
    end

    h_prop = exp(dt * (FQ_NEGI_H0_ISO + astate[CONTROLS_IDX[1]] * NEGI_H1_ISO))
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    intcontrols = astate[INTCONTROLS_IDX[1]] + dt * astate[CONTROLS_IDX[1]]
    controls = astate[CONTROLS_IDX[1]] + dt * astate[DCONTROLS_IDX[1]]
    dcontrols = astate[DCONTROLS_IDX[1]] + dt * acontrols[D2CONTROLS_IDX[1]]
    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
    ]

    if DA
        intgamma = (astate[INTGAMMA_IDX[1]] +
                    dt * amp_t1_reduced_spline(astate[CONTROLS_IDX[1]])^(-1))
        push!(astate_, intgamma)
    end

    return astate_
end

# main
function run_traj(;evolution_time=21.61, gate_type=zpiby2, time_optimal=false,
                  decay_aware=false, solver_type=altro, sqrtbp=false,
                  integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e2, 1e-1, 5e2],
                  smoke_test=false,
                  al_tol=1e-4, constraint_tol=1e-8, max_penalty=1e11,
                  pn_steps=2, save=true, verbose=true, dt_inv=Int64(1e1),
                  max_cost_value=1e8, max_iterations=Int64(2e5),
                  benchmark=false,)
    # model configuration
    dt = dt_inv^(-1)
    dt_max = (dt_inv / 2)^(-1)
    dt_min = (dt_inv * 2)^(-1)
    model = Model(decay_aware, time_optimal)
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.

    # initial state
    x0 = zeros(n)
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    x0 = SVector{n}(x0)

    # final state
    if gate_type == xpiby2
        target_state1 = Array(XPIBY2_ISO_1)
        target_state2 = Array(XPIBY2_ISO_2)
    elseif gate_type == ypiby2
        target_state1 = Array(YPIBY2_ISO_1)
        target_state2 = Array(YPIBY2_ISO_2)
    elseif gate_type == zpiby2
        target_state1 = Array(ZPIBY2_ISO_1)
        target_state2 = Array(ZPIBY2_ISO_2)
    end
    xf = zeros(n)
    xf[STATE1_IDX] = target_state1
    xf[STATE2_IDX] = target_state2
    xf = SVector{n}(xf)

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_max = SVector{n}(x_max)
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    x_min = SVector{n}(x_min)

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_max_boundary = SVector{n}(x_max_boundary)
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = SVector{n}(x_min_boundary)

    # dt bound
    u_max = SVector{m}([
        fill(Inf, CONTROL_COUNT);
        fill(sqrt(dt_max), eval(:($time_optimal ? 1 : 0))); #dt
    ])
    u_min = SVector{m}([
        fill(-Inf, CONTROL_COUNT);
        fill(sqrt(dt_min), eval(:($time_optimal ? 1 : 0))); #dt
    ])

    # initial trajectory
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}([
        fill(1e-4, CONTROL_COUNT);
        fill(sqrt(dt), eval(:($time_optimal ? 1 : 0)));
    ]) for k = 1:N - 1]
    X0 = [SVector{n}(
        fill(NaN, n)
    ) for k = 1:N]
    dt_ = time_optimal ? 1 : dt
    Z = Traj(X0, U0, dt_ * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], CONTROL_COUNT); # ∫a
        fill(qs[3], CONTROL_COUNT); # a
        fill(qs[4], CONTROL_COUNT); # ∂a
        fill(qs[5], eval(:($decay_aware ? 1 : 0))); # ∫γ1
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[6], CONTROL_COUNT); # ∂2a
        fill(qs[7], eval(:($time_optimal ? 1 : 0))); # Δt
    ]))
    objective = LQRObjective(Q, R, Qf, xf, N)

    # constraints
    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must satisfy dt bound
    dt_bnd = BoundConstraint(n, m, u_max=u_max, u_min=u_min)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx)
                        for idx in [STATE1_IDX, STATE2_IDX]]

    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N)
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end
    if time_optimal
        add_constraint!(constraints, dt_bnd, 1:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
                 projected_newton_tolerance=al_tol, n_steps=n_steps,
                 penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
                 projected_newton=projected_newton, iterations_inner=iterations_inner,
                 iterations_outer=iterations_outer, iterations=max_iterations,
                 max_cost_value=max_cost_value)
    if benchmark
        benchmark_result = Altro.benchmark_solve!(solver)
    else
        benchmark_result = nothing
        Altro.solve!(solver)
    end

    # Post-process.
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    dtidx_arr = Array(DT_IDX)
    # Square the dts.
    if time_optimal
        acontrols_arr[:, DT_IDX] = acontrols_arr[:, DT_IDX] .^2
        # acontrols_arr[:, DT_IDX] = map(abs, acontrols_arr[:, DT_IDX])
    end
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = iterations(solver)

    result = Dict(
            "acontrols" => acontrols_arr,
            "controls_idx" => cidx_arr,
            "d2controls_dt2_idx" => d2cidx_arr,
            "dt_idx" => dtidx_arr,
            "evolution_time" => evolution_time,
            "astates" => astates_arr,
            "Q" => Q_arr,
            "Qf" => Qf_arr,
            "R" => R_arr,
            "cmax" => cmax,
            "cmax_info" => cmax_info,
            "dt" => dt,
            "solver_type" => Integer(solver_type),
            "sqrtbp" => Integer(sqrtbp),
            "max_penalty" => max_penalty,
            "constraint_tol" => constraint_tol,
            "al_tol" => al_tol,
            "integrator_type" => Integer(integrator_type),
            "gate_type" => Integer(gate_type),
            "save_type" => Integer(jl),
            "iterations" => iterations_,
            "pn_steps" => pn_steps,
            "max_iterations" => max_iterations,
            "max_cost_value" => max_cost_value,
    )

    if time_optimal
        # sample the important metrics
        (controls_sample, d2controls_dt2_sample, evolution_time_sample
         ) = sample_controls(save_file_path)
        result["controls_sample"] = controls_sample
        result["d2controls_dt2_sample"] = d2controls_dt2_sample
        result["evolution_time_sample"] = evolution_time_sample
        result["save_type"] = Integer(samplejl)
    end

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    result = benchmark ? benchmark_result : result

    return result
end


module OneDTMazeUtils

using SparseArrays
using Distributions

import ..TMazeCumulantSchedules
import ..OneDTmazeConst
import ..OneDTMaze
import ..Learner
import ..check_goal
import ..range_check
import ..get_action_probs
import ..GVFHordes
import ..update
import ..Curiosity
import ..GVFSRHordes
import ..SRCreationUtils
import ..FeatureCreator
using ...StatsBase

const TMCS = TMazeCumulantSchedules
const ODTMC = OneDTmazeConst
const SRCU = Curiosity.SRCreationUtils

#####
# GVF Parameter Functions
####

struct GoalTermination <: GVFHordes.GVFParamFuncs.AbstractDiscount
    γ::Float64
end

function Base.get(gt::GoalTermination; state_tp1, kwargs...)
    any([check_goal(OneDTMaze, i, state_tp1) for i in 1:4]) ? 0.0 : gt.γ
end

struct GoalPolicy <: GVFHordes.GVFParamFuncs.AbstractPolicy
    goal::Int
end
(π::GoalPolicy)(s) = sample(Weights([get(π, state_t=s, action_t=a) for a ∈ 1:4]))

function Base.get(π::GoalPolicy; state_t, action_t, kwargs...)
    cur_x = state_t[1]
    cur_y = state_t[2]
    if π.goal == 1
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 1.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON)
            ODTMC.UP == action_t
        else
            ODTMC.LEFT == action_t
        end
    elseif π.goal == 2
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 1.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON)
            ODTMC.DOWN == action_t
        else
            ODTMC.LEFT == action_t
        end
    elseif π.goal == 3
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 0.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
            ODTMC.UP == action_t
        else
            ODTMC.RIGHT == action_t
        end
    elseif π.goal == 4
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 0.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
            ODTMC.DOWN == action_t
        else
            ODTMC.RIGHT == action_t
        end
    end
end

####
# Demon Creation
####

function create_demons(parsed, demon_projected_fc = nothing)
    action_space = 4
    demons = if parsed["demon_learner"] != "SR"
        GVFHordes.Horde(
            [GVFHordes.GVF(GVFHordes.GVFParamFuncs.FeatureCumulant(i+2),
                           GoalTermination(0.9),
                           GoalPolicy(i)) for i in 1:4])
    elseif parsed["demon_learner"] == "SR"
        @assert demon_projected_fc != nothing
        pred_horde =  GVFHordes.Horde(
                [GVFHordes.GVF(GVFHordes.GVFParamFuncs.FeatureCumulant(i+2),
                     GVFHordes.GVFParamFuncs.ConstantDiscount(0.0),
                     GoalPolicy(i)) for i in 1:4])

        SF_policies = [GoalPolicy(i) for i in 1:4]
        SF_discounts = [GoalTermination(0.9) for i in 1:4]
        num_SFs = length(SF_policies)
        # SF_horde = if parsed["demon_reward_feature_type"] == "state"
        #     SRCU.create_state_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)
        # else
        # SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)
        SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)

        # end
        GVFSRHordes.SRHorde(pred_horde, SF_horde, num_SFs, demon_projected_fc)
    else
        throw(ArgumentError("Cannot create demons"))
    end
    return demons
end

function make_behaviour_gvf(learner, γ, state_constructor, expl_strat) #discount, state_constructor_func, learner, exploration_strategy)
    function b_π(state_constructor, learner, exploration_strategy; kwargs...)
        s = state_constructor(kwargs[:state_t])
        preds = learner(s)
        return exploration_strategy(preds)[kwargs[:action_t]]
    end
    GVF_policy = GVFHordes.GVFParamFuncs.FunctionalPolicy((;kwargs...) -> b_π(state_constructor, learner, expl_strat; kwargs...))
    BehaviourGVF = GVFHordes.GVF(GVFHordes.GVFParamFuncs.RewardCumulant(),
                       GoalTermination(γ),
                       GVF_policy)
end

# ####
struct IdealStateActionDemonFeatures <: FeatureCreator
    num_actions::Int
end
function project_features(fc::IdealStateActionDemonFeatures, s_t, a_t, s_tp1)
    # new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, s_tp1) for i in 1:4]))
    goal_ind = findfirst([check_goal(OneDTMaze, i, s_tp1) for i in 1:4])
    reward_feature = zeros(Int(fc.num_actions*4))
    if !(goal_ind isa Nothing)
        reward_feature[Int((goal_ind-1)*fc.num_actions + a_t)] = 1
    end
    return sparsevec(reward_feature)
end
(FP::IdealStateActionDemonFeatures)(s_t,a_t,s_tp1) = project_features(FP, s_t,a_t,s_tp1)

Base.size(FP::IdealStateActionDemonFeatures) = 16
####
# Ideal Feature Creator
####
struct IdealDemonFeatures <: FeatureCreator
end

function project_features(fc::IdealDemonFeatures, state,action,state_tp1)
    new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state_tp1) for i in 1:4]))
    return new_state
end

(FP::IdealDemonFeatures)(state, action, next_state) = project_features(FP, state, action, next_state)
Base.size(FP::IdealDemonFeatures) = 4

struct MarthaIdealDemonFeatures <: FeatureCreator
    num_actions::Int
end

function project_features(fc::MarthaIdealDemonFeatures, state)
    new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state, ODTMC.ACTION_STEP + ODTMC.EPSILON + 0.00001) for i in 1:4]))
    # alt_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state_tp1) for i in 1:4]))
    return new_state
end

(FP::MarthaIdealDemonFeatures)(state) = project_features(FP, state)
Base.size(FP::MarthaIdealDemonFeatures) = 4

####
# GPI Feature Creation
####
struct TMazeEncoding <: FeatureCreator
    num_segments::Int
    num_partitions::Int
    num_actions::Int
    function TMazeEncoding()
        new(7,3,4)
    end
end
Base.size(FP::TMazeEncoding) = FP.num_segments * FP.num_partitions * FP.num_actions
function project_features(fc::FeatureCreator, obs, a_t, obs_tp1)
    segment = Inf
    partition = Inf
    x,y = obs

    #Segment 1: Top Right
    if (x == 0.0 && y >= 0.8)
        segment_length = (1.0 + ODTMC.ACTION_STEP) - (0.8)
        segment = 1
        partition = (y - 0.8) / segment_length
    #Segment 2: Bottom Right
    elseif (x == 0.0 && y < 0.8)
        segment_length = (0.8) - (0.6 - ODTMC.ACTION_STEP)
        segment = 2
        partition = (y - 0.8) / segment_length
    #Segment 3: Middle Left Branch
    elseif (y == 0.8 && range_check(x,0.0,0.5))
        segment_length = (0.5 - 0.0)
        segment = 3
        partition = (x) / segment_length
    # Segment 4: Middle Right Branch
    elseif (y == 0.8 && range_check(x,0.5,1.0))
        segment_length = (1.0-0.5)
        segment = 4
        partition = (x - 0.5)/segment_length
    #Segment 5: Bottom Branch
    elseif (x == 0.5)
        segment_length = (0.8 - 0.0)
        segment = 5
        partition = (y - 0.0)/segment_length
    # #Segment 6: Top Left Branch
    elseif (x == 1.0 && y >= 0.8)
        segment_length = (1.0 + ODTMC.ACTION_STEP) - (0.8)
        segment = 6
        partition = (y - 0.8) / segment_length
        #Segment 7: Bottom Left Branch
    elseif (x == 1.0 && y < 0.8)
        segment_length = (0.8) - (0.6 - ODTMC.ACTION_STEP)
        segment = 7
        partition = (y - 0.8) / segment_length
    else
        @warn "Not A Valid State (x,y)"
    end

    state_action = zeros(fc.num_segments * fc.num_partitions * fc.num_actions)
    # Need maximum since if agent is right on boundary, ceil (0.0) is 0.
    partition_offset = maximum([Int(ceil(partition * fc.num_partitions)),1])
    if partition_offset == 0
        @show x,y,segment
    end
    state_ind = (segment-1)*fc.num_partitions + partition_offset
    state_action_ind = state_ind * fc.num_actions + a_t

    state_action[state_action_ind] = 1
    return sparsevec(state_action)
end
(FP::TMazeEncoding)(state, action, state_tp1) = project_features(FP, state, action, state_tp1)

####
# Behaviour policies
####

Base.@kwdef struct RoundRobinPolicy <: Learner
    update = Nothing
end

Curiosity.update!(learner::RoundRobinPolicy, args...) = nothing

Base.get(π::RoundRobinPolicy; state_t, action_t, kwargs...) =
    get_action_probs(π, state_t, nothing)[action_t]

function Curiosity.get_action_probs(π::RoundRobinPolicy, features, state)
    cur_x = state[1]
    cur_y = state[2]
    ret = zeros(4)

    if cur_x == 0.5
        if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON) # Middle Junction
            ret[ODTMC.LEFT] = 0.5
            ret[ODTMC.RIGHT] = 0.5
        else # Middle Hallway
            ret[ODTMC.UP] = 1.0
        end
    elseif cur_y == 0.8 && range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON) # Left Junction
        ret[ODTMC.UP] = 0.5
        ret[ODTMC.DOWN] = 0.5
    elseif cur_y == 0.8 && range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
        ret[ODTMC.UP] = 0.5
        ret[ODTMC.DOWN] = 0.5
    elseif cur_x == 0.0
        if cur_y > 0.8
            ret[ODTMC.UP] = 1.0
        else
            ret[ODTMC.DOWN] = 1.0
        end
    elseif cur_x == 1.0
        if cur_y > 0.8
            ret[ODTMC.UP] = 1.0
        else
            ret[ODTMC.DOWN] = 1.0
        end
    elseif cur_x < 0.5
        ret[ODTMC.LEFT] = 1.0
    else
        ret[ODTMC.RIGHT] = 1.0
    end
    ret
end

####
# Cumulant Schedules
####
DrifterDistractor(parsed) = begin
    c_dist = Uniform(parsed["constant_target"][1],parsed["constant_target"][2])
    c1,c2 = rand(c_dist,2)
    if "drifter" ∈ keys(parsed)
        TMCS.DrifterDistractor(
            c1,
            c2,
            parsed["drifter"][1],
            parsed["drifter"][2],
            parsed["distractor"][1],
            parsed["distractor"][2])
    else
        TMCS.DrifterDistractor(
            c1,
            c2,
            parsed["drifter_init"],
            parsed["drifter_std"],
            parsed["distractor_mean"],
            parsed["distractor_std"])
    end
end

function get_cumulant_schedule(parsed)
    sched = parsed["cumulant_schedule"]
    if parsed["cumulant_schedule"] == "DrifterDistractor"
        DrifterDistractor(parsed)
    elseif parsed["cumulant_schedule"] == "Constant"
        if parsed["cumulant"] isa Number
            TMCS.Constant(parsed["cumulant"])
        else
            TMCS.Constant(parsed["cumulant"]...)
        end

    else
        throw("$(sched) Not Implemented")
    end
end

function get_true_values(env::Curiosity.OneDTMaze, eval_set)
    copy_eval_est = deepcopy(eval_set)
    num_gvfs = 4
    goal_cumulants = TMCS.get_cumulant_eval_values(env.cumulant_schedule)
    for i in 1:num_gvfs
        copy_eval_est[i, :] .*= goal_cumulants[i]
    end
    return copy_eval_est
end
function get_true_values(env::Curiosity.OneDTMaze, eval_set, gvf_idx)
    copy_eval_est = deepcopy(eval_set)
    goal_cumulants = TMCS.get_cumulant_eval_values(env.cumulant_schedule)
    copy_eval_est .*= goal_cumulants[gvf_idx]
    return copy_eval_est
end

end


"""
```
struct VarName{sym}
    indexing  ::    String
end
```

A variable identifier. Every variable has a symbol `sym` and `indices `indexing`.
The Julia variable in the model corresponding to `sym` can refer to a single value or
to a hierarchical array structure of univariate, multivariate or matrix variables. `indexing` stores the indices that can access the random variable from the Julia
variable.

Examples:

- `x[1] ~ Normal()` will generate a `VarName` with `sym == :x` and `indexing == "[1]"`.
- `x[:,1] ~ MvNormal(zeros(2))` will generate a `VarName` with `sym == :x` and
 `indexing == "[Colon(),1]"`.
- `x[:,1][2] ~ Normal()` will generate a `VarName` with `sym == :x` and
 `indexing == "[Colon(),1][2]"`.
"""
struct VarName{sym}
    indexing::String
end

"""
    @varname(var)

A macro that returns an instance of `VarName` given the symbol or expression of a Julia variable, e.g. `@varname x[1,2][1+5][45][3]` returns `VarName{:x}("[1,2][6][45][3]")`.
"""
macro varname(expr::Union{Expr, Symbol})
    expr |> varname |> esc
end
function varname(expr)
    ex = deepcopy(expr)
    (ex isa Symbol) && return quote
        DynamicPPL.VarName{$(QuoteNode(ex))}("")
    end
    (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!")
    inds = :(())
    while ex.head == :ref
        if length(ex.args) >= 2
            strs = map(x -> :($x === (:) ? "Colon()" : string($x)), ex.args[2:end])
            pushfirst!(inds.args, :("[" * join($(Expr(:vect, strs...)), ",") * "]"))
        end
        ex = ex.args[1]
        isa(ex, Symbol) && return quote
            DynamicPPL.VarName{$(QuoteNode(ex))}(foldl(*, $inds, init = ""))
        end
    end
    throw("VarName: Mis-formed variable name $(expr)!")
end

macro vsym(expr::Union{Expr, Symbol})
    expr |> vsym
end

"""
    vsym(expr::Union{Expr, Symbol})

Returns the variable symbol given the input variable expression `expr`. For example, if the input `expr = :(x[1])`, the output is `:x`.
"""
function vsym(expr::Union{Expr, Symbol})
    ex = deepcopy(expr)
    (ex isa Symbol) && return QuoteNode(ex)
    (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!")
    while ex.head == :ref
        ex = ex.args[1]
        isa(ex, Symbol) && return QuoteNode(ex)
    end
    throw("VarName: Mis-formed variable name $(expr)!")
end

"""
    @vinds(expr)

Returns a tuple of tuples of the indices in `expr`. For example, `@vinds x[1,:][2]` returns
`((1, Colon()), (2,))`.
"""
macro vinds(expr::Union{Expr, Symbol})
    expr |> vinds |> esc
end
function vinds(expr::Union{Expr, Symbol})
    ex = deepcopy(expr)
    inds = Expr(:tuple)
    (ex isa Symbol) && return inds
    (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!")
    while ex.head == :ref
        pushfirst!(inds.args, Expr(:tuple, ex.args[2:end]...))
        ex = ex.args[1]
        isa(ex, Symbol) && return inds
    end
    throw("VarName: Mis-formed variable name $(expr)!")
end

"""
    split_var_str(var_str, inds_as = Vector)

This function splits a variable string, e.g. `"x[1:3,1:2][3,2]"` to the variable's symbol `"x"` and the indexing `"[1:3,1:2][3,2]"`. If `inds_as = String`, the indices are returned as a string, e.g. `"[1:3,1:2][3,2]"`. If `inds_as = Vector`, the indices are returned as a vector of vectors of strings, e.g. `[["1:3", "1:2"], ["3", "2"]]`.
"""
function split_var_str(var_str, inds_as = Vector)
    ind = findfirst(c -> c == '[', var_str)
    if inds_as === String
        if ind === nothing
            return var_str, ""
        else
            return var_str[1:ind-1], var_str[ind:end]
        end
    end
    @assert inds_as === Vector
    inds = Vector{String}[]
    if ind === nothing
        return var_str, inds
    end
    sym = var_str[1:ind-1]
    ind = length(sym)
    while ind < length(var_str)
        ind += 1
        @assert var_str[ind] == '['
        push!(inds, String[])
        while var_str[ind] != ']'
            ind += 1
            if var_str[ind] == '['
                ind2 = findnext(c -> c == ']', var_str, ind)
                push!(inds[end], strip(var_str[ind:ind2]))
                ind = ind2+1
            else
                ind2 = findnext(c -> c == ',' || c == ']', var_str, ind)
                push!(inds[end], strip(var_str[ind:ind2-1]))
                ind = ind2
            end
        end
    end
    return sym, inds
end

@generated function inargnames(::VarName{s}, ::Model{_F, argnames}) where {s, argnames, _F}
    return s in argnames
end

@generated function inmissings(::VarName{s}, ::Model{_F, _a, _T, missings}) where {s, missings, _F, _a, _T}
    return s in missings
end


"""
    SMCModel(M!::F1, lG::F2, maxn::Int64, particle::Type, pScratch::Type) where
      {F1<:Function,F2<:Function}
- ```M!``` Mutation function
- ```lG``` Log potential function
- ```maxn``` Maximum n for which the model is well-defined
- ```particle``` Type of a particle
- ```pScratch``` Type of particle scratch space
"""
struct SMCModel{F1<:Function,F2<:Function}
  M!::F1
  lG::F2
  maxn::Int64
  particle::Type
  pScratch::Type
end

const SubVectorF64 = SubArray{Float64, 1, Array{Float64, 1},
  Tuple{UnitRange{Int64}}, true}
const SubVectorI64 = SubArray{Int64, 1, Array{Int64, 1},
  Tuple{UnitRange{Int64}}, true}
const SubVector{Particle} = SubArray{Particle,1,Array{Particle,1},Tuple{UnitRange{Int64}},true}

# the fields of this struct are used by the parallel SMC algorithm
# they are not intended for use by users of the package
mutable struct _SMCInternalParallel{Particle, ParticleScratch}
  Nperthread::Int64
  maxlws::Vector{Float64}
  partialSums::Vector{Float64}
  partialSumSqs::Vector{Float64}
  Ns::Vector{Int64}
  NsPartial::Vector{Int64}
  particleScratches::Vector{ParticleScratch}

  localAs::Vector{SubVectorI64}
  localEves::Vector{SubVectorI64}
  localOldEves::Vector{SubVectorI64}
  localZetas::Vector{SubVector{Particle}}
  localZetaAncs::Vector{SubVector{Particle}}
  localWs::Vector{SubVectorF64}
  localLogWs::Vector{SubVectorF64}
  localScratch1s::Vector{SubVectorF64}

  ## below are only populated if fullOutput = true
  localAllAs::Vector{Vector{SubVectorI64}}
  localAllZetas::Vector{Vector{SubVector{Particle}}}
  localAllEves::Vector{Vector{SubVectorI64}}
  localAllWs::Vector{Vector{SubVectorF64}}
end

function _SMCInternalParallel{Particle, ParticleScratch}(N::Int64, n::Int64,
  nthreads::Int64, fullOutput::Bool) where {Particle, ParticleScratch}
  @assert mod(N, nthreads) == 0 "N must be a multiple of nthreads"
  Nperthread = div(N, nthreads)

  if nthreads > 1
    maxlws::Vector{Float64} = Vector{Float64}(undef, nthreads)
    partialSums::Vector{Float64} = Vector{Float64}(undef, nthreads)
    partialSumSqs::Vector{Float64} = Vector{Float64}(undef, nthreads)
    Ns::Vector{Int64} = Vector{Int64}(undef, nthreads)
    NsPartial::Vector{Int64} = Vector{Int64}(undef, nthreads)

    particleScratches::Vector{ParticleScratch} =
      Vector{ParticleScratch}(undef, nthreads)
    # avoid false sharing
    Threads.@threads for i = 1:nthreads
      @inbounds particleScratches[i] = ParticleScratch()
    end

    localAs::Vector{SubVectorI64} =
      Vector{SubVectorI64}(undef, nthreads)
    localEves::Vector{SubVectorI64} =
      Vector{SubVectorI64}(undef, nthreads)
    localOldEves::Vector{SubVectorI64} =
      Vector{SubVectorI64}(undef, nthreads)
    localZetas::Vector{SubVector{Particle}} =
      Vector{SubVector{Particle}}(undef, nthreads)
    localZetaAncs::Vector{SubVector{Particle}} =
      Vector{SubVector{Particle}}(undef, nthreads)
    localWs::Vector{SubVectorF64} =
      Vector{SubVectorF64}(undef, nthreads)
    localLogWs::Vector{SubVectorF64} =
      Vector{SubVectorF64}(undef, nthreads)
    localScratch1s::Vector{SubVectorF64} =
      Vector{SubVectorF64}(undef, nthreads)
  else
    maxlws = Vector{Float64}(undef, 0)
    partialSums = Vector{Float64}(undef, 0)
    partialSumSqs = Vector{Float64}(undef, 0)
    Ns = Vector{Int64}(undef, 0)
    NsPartial = Vector{Int64}(undef, 0)
    particleScratches = Vector{ParticleScratch}(undef, 0)

    localAs = Vector{SubVectorI64}(undef, 0)
    localEves = Vector{SubVectorI64}(undef, 0)
    localOldEves = Vector{SubVectorI64}(undef, 0)
    localZetas = Vector{SubVector{Particle}}(undef, 0)
    localZetaAncs = Vector{SubVector{Particle}}(undef, 0)
    localWs = Vector{SubVectorF64}(undef, 0)
    localLogWs = Vector{SubVectorF64}(undef, 0)
    localScratch1s = Vector{SubVectorF64}(undef, 0)
  end

  if fullOutput && nthreads > 1
    localAllZetas::Vector{Vector{SubVector{Particle}}} =
      Vector{Vector{SubVector{Particle}}}(undef, n)
    localAllAs::Vector{Vector{SubVectorI64}} =
      Vector{Vector{SubVectorI64}}(undef, n-1)
    localAllEves::Vector{Vector{SubVectorI64}} =
      Vector{Vector{SubVectorI64}}(undef, n)
    localAllWs::Vector{Vector{SubVectorF64}} =
      Vector{Vector{SubVectorF64}}(undef, n)
    for p = 1:n
      localAllZetas[p] = Vector{SubVector{Particle}}(undef, nthreads)
      localAllEves[p] = Vector{SubVectorI64}(undef, nthreads)
      localAllWs[p] = Vector{SubVectorF64}(undef, nthreads)
      p < n && (localAllAs[p] = Vector{SubVectorI64}(undef, nthreads))
    end
  else
    localAllAs = Vector{Vector{SubVectorI64}}(undef, 0)
    localAllZetas = Vector{Vector{SubVector{Particle}}}(undef, 0)
    localAllEves = Vector{Vector{SubVectorI64}}(undef, 0)
    localAllWs = Vector{Vector{SubVectorF64}}(undef, 0)
  end

  return _SMCInternalParallel(Nperthread, maxlws, partialSums, partialSumSqs,
    Ns, NsPartial, particleScratches, localAs, localEves, localOldEves,
    localZetas, localZetaAncs, localWs, localLogWs, localScratch1s, localAllAs,
    localAllZetas, localAllEves, localAllWs)
end

# the fields of this struct are used by the SMC algorithm
# they are not intended for use by users of the package
mutable struct _SMCInternal{Particle, ParticleScratch}
  zetaAncs::Vector{Particle}
  oldEves::Vector{Int64}
  as::Vector{Int64}
  lws::Vector{Float64}
  scratch1::Vector{Float64}
  scratch2::Vector{Float64}
  vecOnes::Vector{Float64}
  nresamples::Int64

  sws::Float64
  mws::Float64
  maxlw::Float64
  particleScratch::ParticleScratch

  parallel::_SMCInternalParallel{Particle, ParticleScratch}
end

## constructor for _SMCInternal
function _SMCInternal{Particle, ParticleScratch}(N::Int64, n::Int64,
  nthreads::Int64, fullOutput::Bool) where {Particle, ParticleScratch}

  lws = Vector{Float64}(undef, N)
  as = Vector{Int64}(undef, N)
  fill!(as, 1)
  scratch1 = Vector{Float64}(undef, N)
  scratch2 = Vector{Float64}(undef, N)
  vecOnes = ones(N)
  nresamples::Int64 = 0
  zetaAncs = Vector{Particle}(undef, N)
  for i in 1:N
    zetaAncs[i] = Particle()
  end
  oldEves = Vector{Int64}(undef, N)

  ## assign user-defined particle scratch space
  @assert ParticleScratch == Nothing || !isbitstype(ParticleScratch)
  particleScratch = ParticleScratch()

  parallel::_SMCInternalParallel{Particle, ParticleScratch} =
    _SMCInternalParallel{Particle, ParticleScratch}(N, n, nthreads, fullOutput)

  return _SMCInternal(zetaAncs, oldEves, as, lws, scratch1, scratch2, vecOnes,
    nresamples, 0.0, 0.0, 0.0, particleScratch, parallel)
end

function _assignThreadViews(internal::_SMCInternal{Particle},
  zetas::Vector{Particle}, eves::Vector{Int64}, ws::Vector{Float64},
  allZetas::Vector{Vector{Particle}}, allWs::Vector{Vector{Float64}},
  allAs::Vector{Vector{Int64}}, allEves::Vector{Vector{Int64}}) where Particle

  zetaAncs::Vector{Particle} = internal.zetaAncs
  as::Vector{Int64} = internal.as
  oldEves::Vector{Int64} = internal.oldEves
  lws::Vector{Float64} = internal.lws
  scratch1::Vector{Float64} = internal.scratch1

  ip::_SMCInternalParallel = internal.parallel
  nthreads::Int64 = length(ip.localWs)
  Nperthread::Int64 = ip.Nperthread
  for i = 1:nthreads
    start = Nperthread * (i - 1) + 1
    finish = start + Nperthread - 1
    ip.localAs[i] = view(as, start:finish)
    ip.localEves[i] = view(eves, start:finish)
    ip.localOldEves[i] = view(oldEves, start:finish)
    ip.localZetas[i] = view(zetas, start:finish)
    ip.localZetaAncs[i] = view(zetaAncs, start:finish)
    ip.localWs[i] = view(ws, start:finish)
    ip.localLogWs[i] = view(lws, start:finish)
    ip.localScratch1s[i] = view(scratch1, start:finish)
  end

  fullOutput::Bool = length(ip.localAllWs) > 0
  if fullOutput && nthreads > 1
    n::Int64 = length(allWs)
    for p = 1:n
      for i = 1:nthreads
        start = Nperthread * (i - 1) + 1
        finish = start + Nperthread - 1
        ip.localAllZetas[p][i] = view(allZetas[p], start:finish)
        ip.localAllEves[p][i] = view(allEves[p], start:finish)
        ip.localAllWs[p][i] = view(allWs[p], start:finish)
        p < n && (ip.localAllAs[p][i] = view(allAs[p], start:finish))
      end
    end
  end
end

## SMC input / output struct. Also contains internal state for the SMC
## implementation
"""
    SMCIO{Particle, ParticleScratch}
Structs of this type should be constructed using the provided constructor.
Important fields:
- ```N::Int64``` Number of particles ``N``
- ```n::Int64``` Number of steps ``n``
- ```nthreads::Int64``` Number of threads
- ```fullOutput::Bool``` Whether particle system history should be recorded
- ```essThreshold::Float64``` Relative ESS Threshold ``\\tau``
- ```zetas::Vector{Particle}``` Time n particles ``\\zeta_n^1, \\ldots, \\zeta_n^N``
- ```eves::Vector{Int64}``` Time n Eve indices ``E_n^1, \\ldots, E_n^N``
- ```ws::Vector{Float64}``` Time n weights ``W_n^1, \\ldots, W_n^N``
- ```logZhats::Vector{Float64}``` ``\\log(\\hat{Z}^N_1), \\ldots, \\log(\\hat{Z}^N_n)``
- ```Vhat1s::Vector{Float64}``` ``\\hat{V}_1^N(1), \\ldots, \\hat{V}_n^N(1)``
- ```esses::Vector{Float64}``` Relative ESS values ``\\mathcal{E}_1^N, \\ldots, \\mathcal{E}_n^N``
- ```resample::Vector{Bool}``` Resampling indicators ``R_1, \\ldots, R_{n-1}``
Populated only if ```fullOutput == true```
- ```allZetas::Vector{Vector{Particle}}``` All the particles
- ```allWs::Vector{Vector{Float64}}``` All the weights
- ```allAs::Vector{Vector{Int64}}``` All the ancestor indices
- ```allEves::Vector{Vector{Int64}}``` All the Eve indices
"""
struct SMCIO{Particle, ParticleScratch}
  N::Int64
  n::Int64
  nthreads::Int64
  zetas::Vector{Particle}
  eves::Vector{Int64}
  ws::Vector{Float64}
  logZhats::Vector{Float64}
  Vhat1s::Vector{Float64}
  esses::Vector{Float64}
  resample::Vector{Bool}
  fullOutput::Bool
  essThreshold::Float64

  internal::_SMCInternal{Particle, ParticleScratch} # for internal use only

  # these are only populated if fullOutput = true
  allZetas::Vector{Vector{Particle}}
  allWs::Vector{Vector{Float64}}
  allAs::Vector{Vector{Int64}}
  allEves::Vector{Vector{Int64}}
end

"""
    SMCIO{Particle, ParticleScratch}(N::Int64, n::Int64, nthreads::Int64,
      fullOutput::Bool, essThreshold::Float64 = 2.0) where
      {Particle, ParticleScratch}
Constructor for ```SMCIO``` structs.
"""
function SMCIO{Particle, ParticleScratch}(N::Int64, n::Int64, nthreads::Int64,
  fullOutput::Bool, essThreshold::Float64 = 2.0) where {Particle,
  ParticleScratch}
  @assert hasmethod(Particle, ()) "Particle() must exist"
  @assert hasmethod(ParticleScratch, ()) "ParticleScratch() must exist"

  zetas = Vector{Particle}(undef, N)
  for i=1:N
    zetas[i] = Particle()
  end

  eves = Vector{Int64}(undef, N)
  ws = Vector{Float64}(undef, N)
  logZhats = Vector{Float64}(undef, n)
  Vhat1s = Vector{Float64}(undef, n)
  esses = Vector{Float64}(undef, n)
  resample = Vector{Bool}(undef, n-1)

  if fullOutput
    allZetas = Vector{Vector{Particle}}(undef, n)
    for i=1:n
      allZetas[i] = Vector{Particle}(undef, N)
      for j=1:N
        allZetas[i][j] = Particle()
      end
    end
    allWs = Vector{Vector{Float64}}(undef, n)
    for i=1:n
      allWs[i] = Vector{Float64}(undef, N)
    end
    allAs = Vector{Vector{Int64}}(undef, n-1)
    for i=1:n-1
      allAs[i] = Vector{Int64}(undef, N)
    end
    allEves = Vector{Vector{Int64}}(undef, n)
    for i=1:n
      allEves[i] = Vector{Int64}(undef, N)
    end
  else
    allZetas = Vector{Vector{Particle}}(undef, 0)
    allWs = Vector{Vector{Float64}}(undef, 0)
    allAs = Vector{Vector{Int64}}(undef, 0)
    allEves = Vector{Vector{Int64}}(undef, 0)
  end

  internal::_SMCInternal = _SMCInternal{Particle, ParticleScratch}(N, n,
    nthreads, fullOutput)

  _assignThreadViews(internal, zetas, eves, ws, allZetas, allWs, allAs, allEves)

  return SMCIO(N, n, nthreads, zetas, eves, ws, logZhats, Vhat1s, esses,
    resample, fullOutput, essThreshold, internal, allZetas, allWs, allAs,
    allEves)
end


function Base.get{S<:VehicleState,D,I,R}(::Feature_PosFt, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    FeatureValue(rec[pastframe][vehicle_index].state.posF.t)
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_PosFyaw, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    FeatureValue(rec[pastframe][vehicle_index].state.posF.ϕ)
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_Speed, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    FeatureValue(rec[pastframe][vehicle_index].state.v)
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_VelFs, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    FeatureValue(veh.state.v*cos(veh.state.posF.ϕ))
end
function Base.get{S<:VehicleState,D,I,R}(::Feature_VelFt, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    FeatureValue(veh.state.v*sin(veh.state.posF.ϕ))
end

generate_feature_functions("TurnRateG", :turnrateG, Float64, "rad/s")
function Base.get{S<:VehicleState,D,I,R}(::Feature_TurnRateG, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0; frames_back::Int=1)

    id = rec[pastframe][vehicle_index].id

    retval = FeatureValue(0.0, FeatureState.INSUF_HIST)
    pastframe2 = pastframe - frames_back
    if pastframe_inbounds(rec, pastframe2)

        veh_index_curr = vehicle_index
        veh_index_prev = findfirst(rec[pastframe2], id)

        if veh_index_prev != 0
            curr = rec[pastframe][veh_index_curr].state.posG.θ
            past = rec[pastframe2][veh_index_prev].state.posG.θ
            Δt = get_elapsed_time(rec, pastframe2, pastframe)
            retval = FeatureValue(deltaangle(past, curr) / Δt)
        end
    end

    retval
end
generate_feature_functions("TurnRateF", :turnrateF, Float64, "rad/s")
function Base.get{S<:VehicleState,D,I,R}(::Feature_TurnRateF, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    _get_feature_derivative_backwards(POSFYAW, rec, roadway, vehicle_index, pastframe)
end
generate_feature_functions("AngularRateG", :angrateG, Float64, "rad/s²")
function Base.get{S<:VehicleState,D,I,R}(::Feature_AngularRateG, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    _get_feature_derivative_backwards(TURNRATEG, rec, roadway, vehicle_index, pastframe)
end
generate_feature_functions("AngularRateF", :angrateF, Float64, "rad/s²")
function Base.get{S<:VehicleState,D,I,R}(::Feature_AngularRateF, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0)
    _get_feature_derivative_backwards(TURNRATEF, rec, roadway, vehicle_index, pastframe)
end
generate_feature_functions("DesiredAngle", :desang, Float64, "rad")
function Base.get{S<:VehicleState,D,I,R}(::Feature_DesiredAngle, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0;
    kp_desired_angle::Float64 = 1.0,
    )

    retval = FeatureValue(0.0, FeatureState.INSUF_HIST)
    if pastframe_inbounds(rec, pastframe) && pastframe_inbounds(rec, pastframe-1)

        id = rec[pastframe][vehicle_index].id

        pastϕ = get_state(rec, id, pastframe-1).posF.ϕ
        currϕ = get_state(rec, id, pastframe).posF.ϕ

        Δt = rec.timestep
        expconst = exp(-kp_desired_angle*Δt)
        retval = FeatureValue((currϕ - pastϕ*expconst) / (1.0 - expconst))
    end
    retval
end

generate_feature_functions("MarkerDist_Left", :d_ml, Float64, "m")
Base.get(::Feature_MarkerDist_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0) =
    FeatureValue(get_markerdist_left(rec[pastframe][vehicle_index], roadway))

generate_feature_functions("MarkerDist_Right", :d_mr, Float64, "m")
Base.get(::Feature_MarkerDist_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0) =
    FeatureValue(get_markerdist_right(rec[pastframe][vehicle_index], roadway))

generate_feature_functions("MarkerDist_Left_Left", :d_mll, Float64, "m", can_be_missing=true)
function Base.get(::Feature_MarkerDist_Left_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    #=
    Distance to the left lane marker one lane to the left
    =#

    veh = rec[pastframe][vehicle_index]
    lane = roadway[veh.state.posF.roadind.tag]
    if n_lanes_left(lane, roadway) > 0
        offset = veh.state.posF.t
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        FeatureValue(lane.width/2 - offset + lane_left.width)
    else
        FeatureValue(NaN, FeatureState.MISSING) # there is no left lane
    end
end
generate_feature_functions("MarkerDist_Right_Right", :d_mrr, Float64, "m", can_be_missing=true)
function Base.get(::Feature_MarkerDist_Right_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    #=
    Distance to the right lane marker one lane to the right
    =#

    veh = rec[pastframe][vehicle_index]
    lane = roadway[veh.state.posF.roadind.tag]
    if n_lanes_right(lane, roadway) > 0
        offset = veh.state.posF.t
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        FeatureValue(lane.width/2 + offset + lane_right.width)
    else
        FeatureValue(NaN, FeatureState.MISSING) # there is no right lane
    end
end
generate_feature_functions("RoadEdgeDist_Left", :d_edgel, Float64, "m")
function Base.get(::Feature_RoadEdgeDist_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    offset = veh.state.posF.t
    footpoint = get_footpoint(veh)
    seg = roadway[veh.state.posF.roadind.tag.segment]
    lane = seg.lanes[end]
    roadproj = proj(footpoint, lane, roadway)
    curvept = roadway[RoadIndex(roadproj)]
    lane = roadway[roadproj.tag]
    FeatureValue(lane.width/2 + abs(curvept.pos - footpoint) - offset)
end
generate_feature_functions("RoadEdgeDist_Right", :d_edger, Float64, "m")
function Base.get(::Feature_RoadEdgeDist_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    offset = veh.state.posF.t
    footpoint = get_footpoint(veh)
    seg = roadway[veh.state.posF.roadind.tag.segment]
    lane = seg.lanes[1]
    roadproj = proj(footpoint, lane, roadway)
    curvept = roadway[RoadIndex(roadproj)]
    lane = roadway[roadproj.tag]
    FeatureValue(lane.width/2 + abs(curvept.pos - footpoint) + offset)
end
generate_feature_functions("LaneOffsetLeft", :posFtL, Float64, "m", can_be_missing=true)
function Base.get(::Feature_LaneOffsetLeft, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh_ego = rec[pastframe][vehicle_index]
    t = veh_ego.state.posF.t
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_left(lane, roadway) > 0
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        lane_offset = t - lane.width/2 - lane_left.width/2
        FeatureValue(lane_offset)
    else
        FeatureValue(NaN, FeatureState.MISSING)
    end
end
generate_feature_functions("LaneOffsetRight", :posFtR, Float64, "m", can_be_missing=true)
function Base.get(::Feature_LaneOffsetRight, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh_ego = rec[pastframe][vehicle_index]
    t = veh_ego.state.posF.t
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_right(lane, roadway) > 0
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        lane_offset = t + lane.width/2 + lane_right.width/2
        FeatureValue(lane_offset)
    else
        FeatureValue(NaN, FeatureState.MISSING)
    end
end
generate_feature_functions("N_Lane_Right", :n_lane_right, Int, "-", lowerbound=0.0)
function Base.get(::Feature_N_Lane_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    nlr = rec[pastframe][vehicle_index].state.posF.roadind.tag.lane - 1
    FeatureValue(convert(Float64, nlr))
end
generate_feature_functions("N_Lane_Left", :n_lane_left, Int, "-", lowerbound=0.0)
function Base.get(::Feature_N_Lane_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    seg = roadway[veh.state.posF.roadind.tag.segment]
    nll = length(seg.lanes) - veh.state.posF.roadind.tag.lane
    FeatureValue(convert(Float64, nll))
end
generate_feature_functions("Has_Lane_Right", :has_lane_right, Bool, "-", lowerbound=0.0, upperbound=1.0)
function Base.get(::Feature_Has_Lane_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    val = get(N_LANE_RIGHT, rec, roadway, vehicle_index, pastframe).v > 0.0
    FeatureValue(convert(Float64, val))
end
generate_feature_functions("Has_Lane_Left", :has_lane_left, Bool, "-", lowerbound=0.0, upperbound=1.0)
function Base.get(::Feature_Has_Lane_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    val = get(N_LANE_LEFT, rec, roadway, vehicle_index, pastframe).v > 0.0
    FeatureValue(convert(Float64, val))
end
generate_feature_functions("LaneCurvature", :curvature, Float64, "1/m", can_be_missing=true)
function Base.get(::Feature_LaneCurvature, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    veh = rec[pastframe][vehicle_index]
    curvept = roadway[veh.state.posF.roadind]
    val = curvept.k
    if isnan(val)
        FeatureValue(0.0, FeatureState.MISSING)
    else
        FeatureValue(val)
    end
end

# Dist_Merge
# Dist_Split

generate_feature_functions("TimeToCrossing_Right", :ttcr_mr, Float64, "s", lowerbound=0.0, censor_hi=10.0)
function Base.get(::Feature_TimeToCrossing_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    d_mr = get(MARKERDIST_RIGHT, rec, roadway, vehicle_index, pastframe).v
    velFt = get(VELFT, rec, roadway, vehicle_index, pastframe).v

    if d_mr > 0.0 && velFt < 0.0
        FeatureValue(-d_mr / velFt)
    else
        FeatureValue(Inf)
    end
end
generate_feature_functions("TimeToCrossing_Left", :ttcr_ml, Float64, "s", lowerbound=0.0, censor_hi=10.0)
function Base.get(::Feature_TimeToCrossing_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    d_ml = get(MARKERDIST_RIGHT, rec, roadway, vehicle_index, pastframe).v
    velFt = get(VELFT, rec, roadway, vehicle_index, pastframe).v

    if d_ml > 0.0 && velFt < 0.0
        FeatureValue(-d_ml / velFs)
    else
        FeatureValue(Inf)
    end
end
generate_feature_functions("EstimatedTimeToLaneCrossing", :est_ttcr, Float64, "s", lowerbound=0.0, censor_hi=10.0)
function Base.get(::Feature_EstimatedTimeToLaneCrossing, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    ttcr_left = get(TIMETOCROSSING_LEFT, rec, roadway, vehicle_index, pastframe).v
    ttcr_right = get(TIMETOCROSSING_RIGHT, rec, roadway, vehicle_index, pastframe).v
    FeatureValue(min(ttcr_left, ttcr_right))
end
generate_feature_functions("A_REQ_StayInLane", :a_req_stayinlane, Float64, "m/s²", can_be_missing=true)
function Base.get(::Feature_A_REQ_StayInLane, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    velFt = get(VELFT, rec, roadway, vehicle_index, pastframe).v

    if velFt > 0.0
        d_mr = get(MARKERDIST_RIGHT, rec, roadway, vehicle_index, pastframe).v
        if d_mr > 0.0
            return FeatureValue(0.5velFt*velFt / d_mr)
        else
            return FeatureValue(NaN, FeatureState.MISSING)
        end
    else
        d_ml = get(MARKERDIST_LEFT, rec, roadway, vehicle_index, pastframe)
        if d_ml < 0.0
            return FeatureValue(-0.5velFt*velFt / d_ml)
        else
            return FeatureValue(NaN, FeatureState.MISSING)
        end
    end
end

generate_feature_functions("Time_Consecutive_Brake", :time_consec_brake, Float64, "s", lowerbound=0.0)
function Base.get(::Feature_Time_Consecutive_Brake, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)

    prev_accel = convert(Float64, get(ACC, rec, roadway, vehicle_index, pastframe))
    if prev_accel ≥ 0.0
        FeatureValue(0.0)
    else
        pastframe_orig = pastframe
        id = rec[pastframe][vehicle_index].id
        while pastframe_inbounds(rec, pastframe-1) &&
              get(ACC, rec, roadway, findfirst(rec[pastframe-1], id)) < 0.0

            pastframe -= 1
        end

        FeatureValue(get_elapsed_time(rec, pastframe, pastframe_orig))
    end
end
generate_feature_functions("Time_Consecutive_Accel", :time_consec_accel, Float64, "s", lowerbound=0.0)
function Base.get(::Feature_Time_Consecutive_Accel, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)

    prev_accel = convert(Float64, get(ACC, rec, roadway, vehicle_index, pastframe))
    if prev_accel ≤ 0.0
        FeatureValue(0.0)
    else

        pastframe_orig = pastframe
        id = rec[pastframe][vehicle_index].id
        while pastframe_inbounds(rec, pastframe-1) &&
              get(ACC, rec, roadway, findfirst(rec[pastframe-1], id)) > 0.0

            pastframe -= 1
        end

        FeatureValue(get_elapsed_time(rec, pastframe, pastframe_orig))
    end
end
generate_feature_functions("Time_Consecutive_Throttle", :time_consec_throttle, Float64, "s", lowerbound=0.0)
function Base.get(::Feature_Time_Consecutive_Throttle, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    tc_accel = get(TIME_CONSECUTIVE_ACCEL, rec, roadway, vehicle_index, pastframe).v
    tc_brake = get(TIME_CONSECUTIVE_BRAKE, rec, roadway, vehicle_index, pastframe).v
    FeatureValue(tc_accel ≥ tc_brake ? tc_accel : -tc_brake)
end

#############################################
#
# FRONT
#
#############################################

generate_feature_functions("Dist_Front", :d_front, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Front, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64=100.0,
    )

    if neighborfore.ind == 0
        FeatureValue(100.0, FeatureState.CENSORED_HI)
    else
        scene = rec[pastframe]
        len_ego = scene[vehicle_index].def.length
        len_oth = scene[neighborfore.ind].def.length
        FeatureValue(neighborfore.Δs - len_ego/2 - len_oth/2)
    end
end
generate_feature_functions("Speed_Front", :v_front, Float64, "m/s", can_be_missing=true)
function Base.get(::Feature_Speed_Front, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    )

    if neighborfore.ind == 0
        FeatureValue(0.0, FeatureState.MISSING)
    else
        FeatureValue(rec[pastframe][neighborfore.ind].state.v)
    end
end
generate_feature_functions("Timegap", :timegap, Float64, "s", can_be_missing=true)
function Base.get(::Feature_Timegap, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64 = 10.0,
    )

    v = rec[pastframe][vehicle_index].state.v

    if v ≤ 0.0 || neighborfore.ind == 0
        FeatureValue(censor_hi, FeatureState.CENSORED_HI)
    else
        scene = rec[pastframe]
        len_ego = scene[vehicle_index].def.length
        len_oth = scene[neighborfore.ind].def.length
        Δs = neighborfore.Δs - len_ego/2 - len_oth/2

        if Δs > 0.0
            FeatureValue(Δs / v)
        else
            FeatureValue(0.0) # collision!
        end
    end
end

generate_feature_functions("Inv_TTC", :inv_ttc, Float64, "1/s", can_be_missing=true)
function Base.get(::Feature_Inv_TTC, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64 = 10.0,
    )

    if neighborfore.ind == 0
        FeatureValue(0.0, FeatureState.MISSING)
    else
        scene = rec[pastframe]
        veh_fore = scene[neighborfore.ind]
        veh_rear = scene[vehicle_index]

        len_ego = veh_fore.def.length
        len_oth = veh_rear.def.length
        Δs = neighborfore.Δs - len_ego/2 - len_oth/2

        Δv = veh_fore.state.v - veh_rear.state.v

        if Δs < 0.0 # collision!
            FeatureValue(censor_hi, FeatureState.CENSORED_HI)
        elseif Δv > 0.0 # front car is pulling away
            FeatureValue(0.0)
        else
            f = -Δv/Δs
            if f > censor_hi
                FeatureValue(f, FeatureState.CENSORED_HI)
            else
                FeatureValue(f)
            end
        end
    end
end
generate_feature_functions("TTC", :ttc, Float64, "s", can_be_missing=true)
function Base.get(::Feature_TTC, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64 = 10.0,
    inv_ttc::FeatureValue = get(INV_TTC, rec, roadway, vehicle_index, pastframe, neighborfore=neighborfore, censor_hi=censor_hi),
    )

    if inv_ttc.i == FeatureState.MISSING
        # if the value is missing then front car not found and set TTC to censor_hi
        return FeatureValue(censor_hi, FeatureState.MISSING)
    else
        @assert is_feature_valid(inv_ttc) || inv_ttc.i == FeatureState.CENSORED_HI
        return FeatureValue(min(1.0 / inv_ttc.v, censor_hi))
    end
end

#############################################
#
# FRONT LEFT
#
#############################################
generate_feature_functions("Dist_Front_Left", :d_front_left, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Front_Left, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_left_lane(rec[pastframe], vehicle_index, roadway),
    )

    get(DIST_FRONT, rec, roadway, vehicle_index, pastframe, neighborfore=neighborfore)
end

#############################################
#
# FRONT RIGHT
#
#############################################

generate_feature_functions("Dist_Front_Right", :d_front_right, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Front_Right, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborfore::NeighborLongitudinalResult = get_neighbor_fore_along_right_lane(rec[pastframe], vehicle_index, roadway),
    )

    get(DIST_FRONT, rec, roadway, vehicle_index, pastframe, neighborfore=neighborfore)
end

#############################################
#
# REAR
#
#############################################

generate_feature_functions("Dist_Rear", :d_rear, Float64, "m", lowerbound=0.0, can_be_missing=true)
function Base.get(::Feature_Dist_Rear, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborrear::NeighborLongitudinalResult = get_neighbor_rear_along_lane(rec[pastframe], vehicle_index, roadway),
    censor_hi::Float64=100.0,
    )

    if neighborrear.ind == 0
        FeatureValue(100.0, FeatureState.CENSORED_HI)
    else
        scene = rec[pastframe]
        len_ego = scene[vehicle_index].def.length
        len_oth = scene[neighborrear.ind].def.length
        FeatureValue(neighborrear.Δs - len_ego/2 - len_oth/2)
    end
end
generate_feature_functions("Speed_Rear", :v_rear, Float64, "m/s", can_be_missing=true)
function Base.get(::Feature_Speed_Rear, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    neighborrear::NeighborLongitudinalResult = get_neighbor_rear_along_lane(rec[pastframe], vehicle_index, roadway),
    )

    if neighborrear.ind == 0
        FeatureValue(0.0, FeatureState.MISSING)
    else
        FeatureValue(rec[pastframe][neighborrear.ind].state.v)
    end
end

#############################################
#
# SCENE WISE
#
#############################################

generate_feature_functions("Is_Colliding", :is_colliding, Bool, "-", lowerbound=0.0, upperbound=1.0)
function Base.get{S,D,I,R}(::Feature_Is_Colliding, rec::EntityQueueRecord{S,D,I}, roadway::R, vehicle_index::Int, pastframe::Int=0;
    mem::CPAMemory=CPAMemory(),
    )

    scene = rec[pastframe]
    is_colliding = convert(Float64, get_first_collision(scene, vehicle_index, mem).is_colliding)
    FeatureValue(is_colliding)
end


"""
A `struct` for individuals that keeps individual-specific variables.
"""
mutable struct Ind{B<:AbstractFloat, C<:AbstractArray, D<:AbstractArray, E<:AbstractArray} <: AbstractAgent
  id::Int  # the individual ID
  pos::Tuple{Int, Int}  # the individuals position
  species::Int  # the species ID the individual belongs to
  biotic_phenotype::Vector{B}
  abiotic_phenotype::Vector{B}
  epistasisMat::C  # epistasis matrix
  pleiotropyMat::D  # pleiotropy matrix
  q::E  # expression array
  age::Int
  sex::Bool
  interaction_history::MArray{S, Int} where S # records the last interaction with all species
  energy::B  # determines whether need to feed (energy=0) or not.
  W::B  # survival probability
  isalive::Bool
end

struct Params{F<:AbstractFloat, I<:Int, N<:AbstractString}
  ngenes::Vector{I}
  nphenotypes::Vector{I}
  growthrates::Vector{F}
  selectionCoeffs::Vector{F}
  ploidy::Vector{I}
  optvals::Vector{Vector{Vector{Matrix{F}}}}
  optinds::Vector{Vector{I}}
  mutProbs::Vector{Vector{DiscreteNonParametric{Bool, Float64, Vector{Bool}, Vector{Float64}}}}
  mutMagnitudes::Vector{Vector{UnivariateDistribution{S} where S<:ValueSupport}}
  N::Vector{Vector{I}}
  E::Vector{Normal{F}}
  generations::I
  nspecies::I
  new_N::Vector{Vector{I}}
  migration_traits::Vector{I}
  vision_radius::Vector{I}
  check_fraction::Vector{F}
  migration_thresholds::Vector{F}
  step::MVector{1, Int64}
  nodes::Matrix{Tuple{I, I}}
  biotic_phenotypes::Vector{Vector{I}}
  abiotic_phenotypes::Vector{Vector{I}}
  max_ages::Vector{I}
  names::Dict{I, N}
  food_sources::Matrix{F}
  interactions::Matrix{F}
  resources::Matrix{I}
  resources_org::Matrix{I}
  recombination::Vector{Poisson{F}}
  initial_energy::Vector{F}
  bottlenecks::Vector{BitMatrix}
  repro_start::Vector{Int}
  repro_end::Vector{Int}
  seed::Int
end

const variance = 1.0

"""
    model_initiation(dd)

Innitializes the model.
"""
function model_initiation(dd)

  if !isnothing(dd[:model]["seed"])
    Random.seed!(dd[:model]["seed"])
  else
    rnd = Int(rand(1:1e6))
    dd[:model]["seed"] = rnd
    Random.seed!(rnd)
  end

  properties, species_arrays = create_properties(dd)
  epistasisMat, pleiotropyMat, expressionArrays = species_arrays

  space = dd[:model]["space"]

  if isnothing(space)
    fspace = GridSpace((1, 1))
  elseif typeof(space) <: Tuple
    fspace = GridSpace(space, periodic=dd[:model]["periodic"], metric=dd[:model]["metric"])
  end

  indtype = EvoDynamics.Ind{typeof(0.1), eltype(epistasisMat), eltype(pleiotropyMat), eltype(expressionArrays)}

  model = ABM(indtype, fspace, properties=properties, scheduler=Schedulers.randomly)

  # create and add agents
  for sp in 1:model.nspecies
    for (pos, n) in enumerate(model.N[sp])
      for ind in 1:n
        abiotic_ph = get_abiotic_phenotype(sp, epistasisMat[sp], pleiotropyMat[sp], expressionArrays[sp], model)
        biotic_ph = get_biotic_phenotype(sp, epistasisMat[sp], pleiotropyMat[sp], expressionArrays[sp], model)
        W = abiotic_fitness(abiotic_ph, sp, pos, model)
        W = adjust_fitness(W, sp, model)
        sex = false
        if model.ploidy[sp] == 2
          sex = rand((true, false))
        end
        interaction_history = MVector{model.nspecies, Int}(fill(0, model.nspecies))
        initial_energy = model.initial_energy[sp]
        add_agent!(model.nodes[pos], model, sp, biotic_ph, abiotic_ph, epistasisMat[sp], pleiotropyMat[sp], expressionArrays[sp], 0, sex, interaction_history, initial_energy, W, true)
      end
    end
  end

  return model
end

nnodes(x::GridSpace) = prod(size(x))
nnodes(x::ABM) = nnodes(x.space)

function create_properties(dd)
  nspecies = length(dd[:species])

  Ed = [Normal(0.0, dd[:species][i]["environmental noise"]) for i in 1:nspecies]
  Mdists = [[DiscreteNonParametric([true, false], [i, 1-i]) for i in dd[:species][arr]["mutation probabilities"]] for arr in 1:nspecies]  # μ (probability of change)
  Ddists = [[Normal(0, dd[:species][ar]["mutation magnitudes"][1]), DiscreteNonParametric([true, false], [dd[:species][ar]["mutation magnitudes"][2], 1-dd[:species][ar]["mutation magnitudes"][2]]), Normal(0, dd[:species][ar]["mutation magnitudes"][3])] for ar in 1:nspecies]  # amount of change in case of mutation

  # make single-element arrays 2D so that linAlg functions will work
  newA = Array{Array{Float64}}(undef, nspecies)
  newQ = Array{Array{Float64}}(undef, nspecies)
  for i in 1:nspecies
    if length(dd[:species][i]["epistasis matrix"]) == 1
      newA[i] = reshape(dd[:species][i]["epistasis matrix"], 1, 1)
      newQ[i] = reshape(dd[:species][i]["expression array"], 1, 1)
    else
      newA[i] = dd[:species][i]["epistasis matrix"]
      newQ[i] = dd[:species][i]["expression array"]
    end
  end

  epistasisMatS = [MArray{Tuple{size(newA[i])...}}(newA[i]) for i in eachindex(newA)]
  pleiotropyMatS = [MArray{Tuple{size(dd[:species][i]["pleiotropy matrix"])...}}(dd[:species][i]["pleiotropy matrix"]) for i in 1:nspecies]
  expressionArraysS = [MArray{Tuple{size(newQ[i])...}}(newQ[i]) for i in eachindex(newQ)]

  ngenes = [dd[:species][i]["number of genes"] for i in 1:nspecies]
  nphenotypes = [dd[:species][i]["number of phenotypes"] for i in 1:nspecies]
  growthrates = [dd[:species][i]["growth rate"] for i in 1:nspecies]
  selectionCoeffs = [dd[:species][i]["selection coefficient"] for i in 1:nspecies]
  ploidy = [dd[:species][i]["ploidy"] for i in 1:nspecies]
  optvals = [dd[:species][i]["optimal phenotype values"] for i in 1:nspecies]
  optinds = [dd[:species][i]["optimal phenotypes"] for i in 1:nspecies]
  Ns = [dd[:species][i]["N"] for i in 1:nspecies]
  migration_traits = [dd[:species][i]["migration phenotype"] for i in 1:nspecies]
  vision_radius = [dd[:species][i]["vision radius"] for i in 1:nspecies]
  check_fraction = [dd[:species][i]["check fraction"] for i in 1:nspecies]
  migration_thresholds = [dd[:species][i]["migration threshold"] for i in 1:nspecies]
  generations = dd[:model]["generations"]
  step = MVector{1, Int}(undef)
  step[1] = 0
  nnodes = Matrix{Tuple{Int, Int}}(undef, dd[:model]["space"]...)
  for col in 1:size(nnodes, 2)
    for row in 1:size(nnodes, 1)
      nnodes[row, col] = (row, col)
    end
  end
  biotic_phenotyps = [dd[:species][i]["biotic phenotypes"] for i in 1:nspecies]
  abiotic_phenotyps = [dd[:species][i]["abiotic phenotypes"] for i in 1:nspecies]
  max_ages = [dd[:species][i]["age"] for i in 1:nspecies]
  names = Dict(i => dd[:species][i]["name"] for i in 1:nspecies)
  recombination = [Poisson(dd[:species][i]["recombination"]) for i in 1:nspecies]
  initial_energy = [AbstractFloat(dd[:species][i]["initial energy"]) for i in 1:nspecies]
  bottlenecks = [dd[:species][i]["bottleneck times"] for i in 1:nspecies]
  repro_start = [dd[:species][i]["reproduction start age"] for i in 1:nspecies]
  repro_end = [dd[:species][i]["reproduction end age"] for i in 1:nspecies]

  properties = Params(ngenes, nphenotypes, growthrates, selectionCoeffs, ploidy, optvals, optinds, Mdists, Ddists, Ns, Ed, generations, nspecies, Ns, migration_traits, vision_radius, check_fraction, migration_thresholds, step, nnodes, biotic_phenotyps, abiotic_phenotyps, max_ages, names, dd[:model]["food sources"], dd[:model]["interactions"], dd[:model]["resources"], deepcopy(dd[:model]["resources"]), recombination, initial_energy, bottlenecks, repro_start, repro_end, dd[:model]["seed"])

  return properties, (epistasisMatS, pleiotropyMatS, expressionArraysS)
end

function return_opt_phenotype(species::Int, generation::Int, site::Int, model::ABM)
  nabiotic = length(model.abiotic_phenotypes[species])
  output = Array{typeof(0.1)}(undef, nabiotic)
  for trait in 1:nabiotic
    output[trait] = model.optvals[species][model.optinds[species][generation+1]][trait][site]
  end
  return output
end

function return_opt_phenotype(species::Int, generation::Int, site::Tuple{Int, Int}, model::ABM)
  nabiotic = length(model.abiotic_phenotypes[species])
  output = Array{typeof(0.1)}(undef, nabiotic)
  for trait in 1:nabiotic
    output[trait] = model.optvals[species][model.optinds[species][generation+1]][trait][site[1],site[2]]
  end
  return output
end

function model_step!(model::ABM)
  model.step[1] += 1
  model.resources .= model.resources_org
end

function agent_step!(agent::Ind, model::ABM)
  # update age
  agent.age += 1
  # abiotic survive
  abiotic_survive!(agent, model)  # the agent first survives then feeds, reproduces, and interacts.
  if !agent.isalive
    return
  end
  # use food
  burn_energy!(agent)
  # consume basic energy if agent can
  consume_food!(agent, model)
  # interact with other species
  interact!(agent, model)
  # Kill the agent if it doesn't have energy
  if agent.isalive && agent.energy < 0
    remove_agent!(agent, model)
    return
  end
  # survive
  survive!(agent, model)
  # migrate
  migrate!(agent, model)
  # reproduction for the haploid
  reproduce!(agent, model)
  # bottleneck
  if agent.isalive && model.bottlenecks[agent.species][coord2vertex(agent.pos, model), model.step[1]]
    if EvoDynamics.bottleneck(agent, model)
      remove_agent!(agent, model)
    end
  end
end

function bottleneck(agent, model)
  # @info "Package bottleneck function"
  return false
end

function remove_agent!(agent, model)
  agent.isalive = false
  kill_agent!(agent, model)
end

"""
    survive!(agent::Ind, model::ABM)

Kills the agent if it has no energy, or is too old.
"""
function survive!(agent::Ind, model::ABM)
  if !agent.isalive
    return
  elseif agent.energy < 0
    remove_agent!(agent, model)
  elseif agent.age ≥ model.max_ages[agent.species]
    remove_agent!(agent, model)
  # elseif rand() > agent.W
  #   remove_agent!(agent, model)
  end
end

"""
Kills the agent by chance given its fitness `W`
"""
function abiotic_survive!(agent::Ind, model::ABM)
  if agent.isalive && rand() > agent.W
    remove_agent!(agent, model)
  end
end

function adjust_fitness!(agent::Ind, model::ABM)
  W = agent.W < 0 ? 0.0 : agent.W
  newW = 1.0 - ( (1.0 - W) * model.selectionCoeffs[agent.species])
  agent.W = newW
end

function adjust_fitness(W, species, model::ABM)
  W2 = W < 0 ? 0.0 : W
  newW = 1.0 - ( (1.0 - W2) * model.selectionCoeffs[species])
  return newW
end

function burn_energy!(agent)
  agent.energy -= 1
end

function consume_food!(agent::Ind, model::ABM)
  environmental_consumption = model.food_sources[agent.species, agent.species]
  if environmental_consumption > 0 && model.resources[agent.pos...] >= environmental_consumption
    model.resources[agent.pos...] -= environmental_consumption
    agent.energy += environmental_consumption
  end
end

"Mutate an agent."
function mutate!(agent::Ind, model::ABM)
  mutated = false
  # mutate gene expression
  if rand(model.mutProbs[agent.species][1])
    agent.q .+= rand(model.mutMagnitudes[agent.species][1], model.ngenes[agent.species] * model.ploidy[agent.species])
    mutated = true
  end
  # mutate pleiotropy matrix
  if rand(model.mutProbs[agent.species][2])
    randnumbers = rand(model.mutMagnitudes[agent.species][2], size(agent.pleiotropyMat))
    agent.pleiotropyMat[randnumbers] .= .!agent.pleiotropyMat[randnumbers]
    mutated = true
  end
  # mutate epistasis matrix
  if rand(model.mutProbs[agent.species][3])
    agent.epistasisMat .+= rand(model.mutMagnitudes[agent.species][3], size(agent.epistasisMat))
    mutated = true
  end
  # update biotic and abiotic phenotypes and W
  if mutated
    update_fitness!(agent, model)
  end
end

coord2vertex(c, model) = c[1] + (size(model.space)[1] * (c[2]-1))

function update_fitness!(agent::Ind, model::ABM)
  abiotic_phenotype = get_abiotic_phenotype(agent, model)
  biotic_phenotype = get_biotic_phenotype(agent, model)
  W = abiotic_fitness(abiotic_phenotype, agent.species, agent.pos, model)
  W = adjust_fitness(W, agent.species, model)
  agent.biotic_phenotype .= biotic_phenotype
  agent.abiotic_phenotype .= abiotic_phenotype
  agent.W = W
end