import DataStructures: OrderedDict

"""
```
HetDSGELag{T} <: AbstractModel{T}
```

### Fields

#### Parameters and Steady-States
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model
  parameters.

* `steady_state::Vector{AbstractParameter}`: Model steady-state values, computed
  as a function of elements of `parameters`.

* `keys::OrderedDict{Symbol,Int}`: Maps human-readable names for all model
  parameters and steady-states to their indices in `parameters` and
  `steady_state`.

#### Inputs to Measurement and Equilibrium Condition Equations

The following fields are dictionaries that map human-readable names to row and
column indices in the matrix representations of of the measurement equation and
equilibrium conditions.

* `endogenous_states::OrderedDict{Symbol,Int}`: Maps each state to a column in
  the measurement and equilibrium condition matrices.

* `exogenous_shocks::OrderedDict{Symbol,Int}`: Maps each shock to a column in
  the measurement and equilibrium condition matrices.

* `expected_shocks::OrderedDict{Symbol,Int}`: Maps each expected shock to a
  column in the measurement and equilibrium condition matrices.

* `equilibrium_conditions::OrderedDict{Symbol,Int}`: Maps each equlibrium
  condition to a row in the model's equilibrium condition matrices.

* `observables::OrderedDict{Symbol,Int}`: Maps each observable to a row in the
  model's measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"an_schorfheide\", cached
  here for filepath computation.

* `subspec::String`: The model subspecification number, indicating that some
  parameters from the original model spec (\"ss0\") are initialized
  differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect computation
  without changing the economic or mathematical setup of the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If `true`,
  settings from `m.test_settings` are used in place of those in `m.settings`.

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

"""
mutable struct HetDSGELag{T} <: AbstractModel{T}
    parameters::ParameterVector{T}                         # vector of all time-invariant model parameters
    steady_state::ParameterVector{T}                       # model steady-state values

    # Temporary to get it to work. Need to
    # figure out a more flexible way to define
    # "grids" that are not necessarily quadrature
    # grids within the model
    grids::OrderedDict{Symbol,Union{Grid, Array}}

    keys::OrderedDict{Symbol,Int}                          # human-readable names for all the model
                                                           # parameters and steady-states

    state_variables::Vector{Symbol}                        # Vector of symbols of the state variables
    jump_variables::Vector{Symbol}                         # Vector of symbols of the jump variables
    normalized_model_states::Vector{Symbol}                # All of the distributional model
                                                           # state variables that need to be normalized
    endogenous_states_unnormalized::OrderedDict{Symbol,UnitRange} # Vector of unnormalized
                                                           # ranges of indices
    endogenous_states::OrderedDict{Symbol,UnitRange}       # Vector of ranges corresponding
                                                           # to normalized (post Klein solution) indices
    exogenous_shocks::OrderedDict{Symbol,Int}              #
    expected_shocks::OrderedDict{Symbol,Int}               #
    equilibrium_conditions::OrderedDict{Symbol,UnitRange}  #
    observables::OrderedDict{Symbol,Int}                   #

    spec::String                                           # Model specification number (eg "m990")
    subspec::String                                        # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                         # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}                    # Settings/flags for testing mode
    rng::MersenneTwister                                   # Random number generator
    testing::Bool                                          # Whether we are in testing mode or not

    observable_mappings::OrderedDict{Symbol, Observable}
end

description(m::HetDSGELag) = "HetDSGELag, $(m.subspec)"

"""
`init_model_indices!(m::HetDSGELag)`

Arguments:
`m:: HetDSGELag`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::HetDSGELag)
    # Endogenous states
    endogenous_states = collect([:l′_t1, :μ′_t1, :k′_t, :R′_t1, :i′_t1,
                                 :π′_t1,:π′_t2,:π′_t3,:y′_t1,:y′_t2,:y′_t3, :y′_t4,
                                 :z′_t1,:z′_t2,:z′_t3,:w′_t1,:I′_t1,:B′,:G′,:z′_t,:MU′,:LAMW′,
                                 :LAMF′,:MON′,:l′_t,:μ′_t,:R′_t,:i′_t,:t′_t,:w′_t,:L′_t,:π′_t,
                                 :wageinflation′_t,:mu′_t,:y′_t,:I′_t,:mc′_t,:Q′_t,:capreturn′_t])
                                 #=:l_t1,:μ_t1,:k_t,:R_t1,:i_t1,:π_t1,:π_t2,:π_t3,:y_t1,:y_t2,
                                 :y_t3,:y_t4,:z_t1,:z_t2,:z_t3,:w_t1,:I_t1,:B,:G,:Z,:MU,:LAMW,
                                 :LAMF,:MON,:ELL,:μ_t,:R_t,:i_t,
                                 :t_t,:w_t,:L_t,:π_t,:wageinflation_t,:mu_t,:y_t,:I_t,:mc_t,:Q_t,:capreturn_t]) =#

    # endogenous scalar states: k_t, R_t1, i_t1, π_t1, π_t2, π_t3, y_t1, y_t2, y_t3, y_t4, z_1, z_2, z_3, w_t1, x_t1

    # Exogenous shocks
    exogenous_shocks = collect([:b_sh,:g_sh,:z_sh,:μ_sh,:λ_w_sh, :λ_f_sh,:MON_sh])

    # Equilibrium conditions
    equilibrium_conditions = collect([:eq_euler,:eq_kolmogorov_fwd,:eq_lag_ell,:eq_lag_wealth,
                                      :eq_market_clearing,:eq_lambda, :eq_transfers,
                                      :eq_investment,:eq_tobin_q,:eq_capital_accumulation,
                                      :eq_wage_phillips,:eq_price_phillips,:eq_marginal_cost,
                                      :eq_gdp,:eq_optimal_kl,:eq_taylor, :eq_fisher,
                                      :eq_nominal_wage_inflation,:LR,:LI,:LPI,:L2PI,:L3PI,:LY,:L2Y,
                                      :L3Y,:L4Y,:LZ,:L2Z,:L3Z,:LW,:LX,:F33,:F34,:F35,:F36,:F37,
                                      :F38, :F39])

    # Observables
    observables = keys(m.observable_mappings)

    ########################################################################################
    # Setting indices of endogenous_states and equilibrium conditions manually for now

    setup_indices!(m)
    endo = m.endogenous_states_unnormalized
    eqcond = equilibrium_conditions

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

    m.normalized_model_states = [:μ′_t1, :l′_t1]

    for (i,k) in enumerate(exogenous_shocks);            m.exogenous_shocks[k]            = i end
    for (i,k) in enumerate(observables);                 m.observables[k]                 = i end
end

function HetDSGELag(subspec::String="ss0";
                   custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                   testing = false)

    # Model-specific specifications
    spec               = "het_dsge_lag"
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # initialize empty model
    m = HetDSGELag{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(),
            # grids and keys
            OrderedDict{Symbol,Union{Grid, Array}}(), OrderedDict{Symbol,Int}(),

            # normalized_model_states, state_inds, jump_inds
            Vector{Symbol}(), Vector{Symbol}(), Vector{Symbol}(),

            # model indices
            # endogenous states unnormalized, endogenous states normalized
            OrderedDict{Symbol,UnitRange}(), OrderedDict{Symbol,UnitRange}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,UnitRange}(), OrderedDict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}())

    m <= Setting(:nx, 300)
    m <= Setting(:ns, 2)

    default_settings!(m)

    # # Set observable transformations
    init_observable_mappings!(m)

    # Initialize model indices
    init_model_indices!(m)

    # Set settings
    model_settings!(m)
    # default_test_settings!(m)
    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Initialize parameters
    init_parameters!(m)

    # Initialize aggregate steady state parameters (necessary for grid construction)
    aggregate_steadystate!(m)

    # Initialize grids
    init_grids!(m)

    # # Solve for the steady state
    steadystate!(m)

    # # So that the indices of m.endogenous_states reflect the normalization
    normalize_model_state_indices!(m)

    init_subspec!(m)

    return m
end

"""
```
init_parameters!(m::HetDSGELag)
```

Initializes the model's parameters, as well as empty values for the steady-state
parameters (in preparation for `steadystate!(m)` being called to initialize
those).
"""
function init_parameters!(m::HetDSGELag)

    # Initialize parameters
    m <= parameter(:r, 0.01, fixed = true,
                   description = "r: Steady-state real interest rate.", tex_label = "r")
    m <= parameter(:α, 0.3, fixed = true,
                   description = "α: Capital share", tex_label = "\\alpha")
    m <= parameter(:H, 1.0, fixed = true,
                   description = "H: Aggregate hours worked", tex_label = "H")
    m <= parameter(:δ, 0.03, fixed = true,
                   description = "δ: Depreciation rate", tex_label = "\\delta")
    m <= parameter(:μ_sp, 0.0, fixed = true, description = "μ_sp: The trend in the skill process",
                   tex_label = "\\mu_{sp}")
    m <= parameter(:ρ_sp, 0.7, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(), BetaAlt(0.5, 0.2), fixed=true,
                   description="ρ_sp: AR(1) coefficient in the skill process.",
                   tex_label="\\rho_{sp}")
    m <= parameter(:σ_sp, 0.01, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.10), fixed=true,
                   description="σ_sp: The standard deviation of the skill process.",
                   tex_label="\\sigma_{sp}")
    m <= parameter(:γ, 0.0, description = "γ: TFP growth")
    #GoverY = 0.01
    m <= parameter(:g, 1/(1-0.01), description = "g: Steady-state government spending/gdp")
    m <= parameter(:η, 0.1, description = "η: Borrowing constraint (normalized by TFP)")

    m <= parameter(:ρ_B, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = "ρ_B: Persistence of discount factor shock", tex_label = "\\rho_B")
    m <= parameter(:ρ_G, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = "ρ_G: Persistence of govt spending shock", tex_label = "\\rho_G")
    m <= parameter(:ρ_z, 0.5,(1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = "ρ_z: Persistence of tfp growth shock", tex_label = "\\rho_z")
    m <= parameter(:ρ_μ, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = "ρ_μ: Persistence of investment shock", tex_label = "\\rho_\\mu")
    m <= parameter(:ρ_lamw, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = "ρ_λ_w: Persistence of wage mkup shock", tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_lamf, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = " # persistence of price mkup shock", tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_mon, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                     fixed=true, description = " # persistence of mon policy shock", tex_label = "\\rho_{mon}")

    m <= parameter(:σ_b, 0.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=true, description = "# standard deviation of discount factor shock", tex_label = "\\sigma_B")
    m <= parameter(:σ_g, 0.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=true, description = "# standard deviation of govt spending shock", tex_label = "\\sigma_G")
    m <= parameter(:σ_z, 0.5,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=false, description = "# standard deviation of tfp growth shock", tex_label = "\\sigma_z")
    m <= parameter(:σ_μ, 0.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=true, description = " # standard deviation of investment shock", tex_label = "\\sigma_\\mu")
    m <= parameter(:σ_λ_w, 0.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=true, description = "# standard deviation of wage mkup shock", tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_λ_f, 0.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=true, description = " # standard deviation of price mkup shock", tex_label = "\\sigma_{\\lambda_f}")
    m <= parameter(:σ_MON, 0.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.1),
                   fixed=true, description = " # standard deviation of mon policy shock", tex_label = "\\sigma_{mon}")

 m <= parameter(:spp, 4., description = "# second derivative of investment adjustment cost")
 m <= parameter(:lamw, 1.5, description = "# wage markup")
 m <= parameter(:ϕh , 2., description = "# inverse frisch elasticity")
 m <= parameter(:Φw , 10., description = "# rotemberg cost for wages")
 m <= parameter(:lamf, 1.5 , description = "# price markup")
 m <= parameter(:Φp , 1., description = "# rotemberg cost for prices")
 m <= parameter(:ρR , 0.75, description = "# persistence in taylor rule")
 m <= parameter(:ψπ , 1.5, description = "# weight on inflation in taylor rule")
 m <= parameter(:ψy , 0.5, description = "# weight on output growth in taylor rule")

    # Setting steady-state parameters
    nx = get_setting(m, :nx)
    ns = get_setting(m, :ns)

    # Steady state grids for functional/distributional variables and market-clearing discount rate
    m <= SteadyStateParameterGrid(:lstar, fill(NaN, nx*ns), description = "Steady-state expected discounted
                                  marginal utility of consumption", tex_label = "l_*")
    m <= SteadyStateParameterGrid(:cstar, fill(NaN, nx*ns), description = "Steady-state consumption",
                                  tex_label = "c_*")
    m <= SteadyStateParameterGrid(:μstar, fill(NaN, nx*ns), description = "Steady-state cross-sectional
                                  density of cash on hand", tex_label = "\\mu_*")
    m <= SteadyStateParameter(:βstar, NaN, description = "Steady-state discount factor",
                              tex_label = "\\beta_*")
end

"""
```
aggregate_steadystate!(m::HetDSGELag)
```
Steady state of aggregate scalar variables (Break out the "analytic" steady-state solution here instead of in steadystate!() since some of these parameters are required to construct the grids that are used in the functional/distributional steady state variable calculation (where init_grids! is called before steadystate! in model construction).
"""
function aggregate_steadystate!(m::HetDSGELag)
    m <= SteadyStateParameter(:Rkstar, m[:r] + m[:δ], description = "Rental rate on capital", tex_label = "Rk_*")
    m <= SteadyStateParameter(:ωstar, (m[:α]^(m[:α]/(1-m[:α])))*(1-m[:α])*m[:Rkstar]^(-m[:α]/(1-m[:α])), description = "Real wage", tex_label = "\\omega_*")
    m <= SteadyStateParameter(:klstar, (m[:α]/(1-m[:α]))*(m[:ωstar]/m[:Rkstar])*exp(m[:γ]), description = "Capital/Labor ratio", tex_label = "kl_*")
    m <= SteadyStateParameter(:kstar, m[:klstar]*m[:H], description = "Capital", tex_label = "k_*")
    m <= SteadyStateParameter(:xstar, (1-(1-m[:δ])*exp(-m[:γ]))*m[:kstar], description = "Investment", tex_label = "x_*")
    m <= SteadyStateParameter(:ystar, (exp(-m[:α]*m[:γ])*(m[:kstar]^m[:α])*m[:H]^(1-m[:α])), description = "GDP", tex_label = "y_*")
    m <= SteadyStateParameter(:Tstar, m[:Rkstar]*m[:kstar]*exp(-m[:γ]) - m[:xstar] - (1-(1/m[:g]))*m[:ystar], description = "Net transfer to households", tex_label = "T_*")

    return m
end

"""
```
init_grids!(m::HetDSGELag)
```
"""
function init_grids!(m::HetDSGELag)

    nx      = get_setting(m, :nx)
    ns      = get_setting(m, :ns)
    λ       = get_setting(m, :λ)

    grids = OrderedDict()

    # Skill grid
    #lsgrid, sprob, sscale = tauchen86(m[:μ_sp].value, m[:ρ_sp].value, m[:σ_sp].value, ns, λ)
    sprob = [[0.9 0.1];[0.1 0.9]]
    sgrid = [0.8;1.2]
    (λs, vs) = eigen(Matrix{Float64}(sprob'))
    order_λs = sortperm(λs, rev = true)
    vs = vs[:,order_λs]
    ss_skill_distr = vs[:,1]/sum(vs[:,1])
    means = ss_skill_distr'*sgrid
    meanz = (get_setting(m, :zhi)+get_setting(m, :zlo))/2.
    sgrid = sgrid/(meanz*means) # so that skills integrate to 1
    sscale = sgrid[2] - sgrid[1]
    swts = (sscale/ns)*ones(ns)
    #sgrid = exp.(lsgrid)
    grids[:sgrid] = Grid(sgrid, swts, sscale)

 # Calculate endogenous grid bounds (bounds are analytic functions of model parameters,
    # which ensure there is no mass in the distributions across x where there shouldn't be)
    # Calculate the lowest possible cash on hand in steady state
    zlo = get_setting(m, :zlo)
    smin = minimum(sgrid)*zlo #exp(m[:μ_sp]/(1-m[:ρ_sp]) - get_setting(m, :λ)*sqrt(m[:σ_sp]^2/(1-m[:ρ_sp])^2))*get_setting(m, :zlo)
    m <= Setting(:xlo, m[:ωstar]*smin*m[:H] - (1+m[:r])*m[:η]*exp(-m[:γ]) + m[:Tstar] + sgrid[1]*m[:ωstar]*m[:H]*0.05, "Lower bound on cash on hand")
    m <= Setting(:xhi, max(get_setting(m, :xlo)*2, get_setting(m, :xlo)+5.), "Upper Bound on cash on hand")
    m <= Setting(:xscale, get_setting(m, :xhi) - get_setting(m, :xlo), "Size of the xgrid")
    xlo     = get_setting(m, :xlo)
    xhi     = get_setting(m, :xhi)
    xscale  = get_setting(m, :xscale)

    # Cash on hand grid
    grids[:xgrid] = Grid(uniform_quadrature(xscale), xlo, xhi, nx, scale = xscale)

    # Markov transition matrix for skill
    grids[:fgrid] = sprob./swts

    # Total grid vectorized across both dimensions
    grids[:sgrid_total] = kron(sgrid, ones(nx))
    grids[:xgrid_total] = kron(ones(ns), grids[:xgrid].points)
    grids[:weights_total] = kron(swts, grids[:xgrid].weights)

    m.grids = grids
end

function model_settings!(m::HetDSGELag)
    default_settings!(m)

    # Defaults
    # Data settings for released and conditional data. Default behavior is to set vintage
    # of data to today's date.
    vint = Dates.format(now(), DSGE_DATE_FORMAT)
    m <= Setting(:data_vintage, vint, true, "vint", "Data vintage")

    saveroot = normpath(joinpath(dirname(@__FILE__), "../../../","save"))
    datapath = normpath(joinpath(dirname(@__FILE__), "../../../","save","input_data"))

    m <= Setting(:saveroot, saveroot, "Root of data directory structure")
    m <= Setting(:dataroot, datapath, "Input data directory path")

    # Solution method
    m <= Setting(:solution_method, :klein)

    m <= Setting(:krylov_reduce, false)

    # Likelihood method
    m <= Setting(:use_chand_recursion, true)

    # Anticipated shocks
    m <= Setting(:n_anticipated_shocks, 0,
                 "Number of anticipated policy shocks")
    # m <= Setting(:n_anticipated_shocks_padding, 20,
                 # "Padding for anticipated policy shocks")

    # Whether one wishes to re-compute βstar
    m <= Setting(:use_last_βstar, false, "Flag to avoid recomputing steady-state.")

    # Number of states and jumps
    m <= Setting(:normalize_distr_variables, true, "Whether or not to perform the
                 normalization of the distributional states in the Klein solution step")
    m <= Setting(:n_predetermined_variables, 0, "Number of predetermined variables after
                 removing the densities. Calculated with the Jacobian normalization.
                 This set to 0 is just the default setting, since it will always be
                 overwritten once the Jacobian is calculated.")

    endo = m.endogenous_states_unnormalized
    state_indices = [endo[:l′_t1]; endo[:μ′_t1];  endo[:k′_t]; endo[:R′_t1];endo[:i′_t1];endo[:π′_t1];
                     endo[:π′_t2];endo[:π′_t3];endo[:y′_t1];endo[:y′_t2];endo[:y′_t3];endo[:y′_t4];
                     endo[:z′_t1];endo[:z′_t2];endo[:z′_t3];endo[:w′_t1];endo[:I′_t1];endo[:B′];
                     endo[:G′];endo[:z′_t];endo[:MU′];endo[:LAMW′]; endo[:LAMF′];endo[:MON′]]
                     #=endo[:l_t1];
                     endo[:μ_t1];endo[:k_t];endo[:R_t1];endo[:i_t1]; endo[:π_t1];endo[:π_t2];
                     endo[:π_t3]; endo[:y_t1];endo[:y_t2];endo[:y_t3]; endo[:y_t4];endo[:z_t1];
                     endo[:z_t2]; endo[:z_t3];endo[:w_t1];endo[:I_t1]; endo[:B];endo[:G];endo[:Z];
                     endo[:MU]; endo[:LAMW];endo[:LAMF];endo[:MON]]=#

    jump_indices = [endo[:l′_t]; endo[:μ′_t];endo[:R′_t];endo[:i′_t];endo[:t′_t];endo[:w′_t];
                    endo[:L′_t];endo[:π′_t];endo[:wageinflation′_t];endo[:mu′_t];endo[:y′_t];
                    endo[:I′_t];endo[:mc′_t]; endo[:Q′_t];endo[:capreturn′_t]]
                    #=endo[:ELL];
                    endo[:μ_t];endo[:R_t];endo[:i_t];endo[:t_t];endo[:w_t]; endo[:L_t];endo[:π_t];
                    endo[:wageinflation_t];endo[:mu_t];endo[:y_t];endo[:I_t];endo[:mc_t];endo[:Q_t];
                    endo[:capreturn_t]]=#

    m <= Setting(:state_indices, state_indices, "Which indices of m.endogenous_states correspond to
                 backward looking state variables")
    m <= Setting(:jump_indices, jump_indices, "Which indices of m.endogenous_states correspond to jump
                 variables")
    #m.state_variables = m.endogenous_states.keys[get_setting(m, :state_indices)]
    #m.jump_variables = m.endogenous_states.keys[get_setting(m, :jump_indices)]

    # Mollifier setting parameters
    m <= Setting(:In, 0.443993816237631, "Normalizing constant for the mollifier")
    m <= Setting(:zlo, 1.0, "Lower bound on second income shock to mollify actual income")
    m <= Setting(:zhi, 5.0, "Upper bound on second income shock to mollify actual income")

    # s: Skill Distribution/ "Units of effective labor" Grid Setup
    m <= Setting(:ns, 2, "Skill distribution grid points")
    m <= Setting(:λ, 2.0, "The λ parameter in the Tauchen distribution calculation")

    # x: Cash on Hand Grid Setup
    m <= Setting(:nx, 300, "Cash on hand distribution grid points")

    m <= Setting(:trunc_distr, true)
    m <= Setting(:rescale_weights, true)
    m <= Setting(:mindens, 1e-8)

    # Function-valued variables include distributional variables
    m <= Setting(:n_function_valued_backward_looking_states, 2, "Number of function-valued
                 backward looking state variables")
    m <= Setting(:n_backward_looking_distributional_vars, 2, "Number of state variables that are
                 distributional variables.")
    m <= Setting(:n_function_valued_jumps, 2, "Number of function-valued jump variables")
    m <= Setting(:n_jump_distributional_vars, 2, "Number of jump variables that are distributional
                 variables.")

    # Note, these settings assume normalization.
    # The n degrees of freedom removed depends on the distributions/dimensions
    # of heterogeneity that we have discretized over, in this case,
    # cash on hand and the skill distribution. In general the rule of
    # thumb is, remove one degree of freedom for the first endogenous distribution (cash on
    # hand), then one additional degree of freedom for each exogenous distribution (skill
    # distribution). Multiple endogenous distributions only permit removing a single degree
    # of freedom since it is then non-trivial to obtain the marginal distributions.
    m <= Setting(:n_degrees_of_freedom_removed, 1, "Number of degrees of freedom from the
                 distributional variables to remove.")
    n_dof_removed = get_setting(m, :n_degrees_of_freedom_removed)

    ####################################################
    # Calculating the number of backward-looking states
    ####################################################
    n_backward_looking_distr_vars = get_setting(m, :n_backward_looking_distributional_vars)
    m <= Setting(:backward_looking_states_normalization_factor,
                 n_dof_removed*n_backward_looking_distr_vars, "The number of dimensions removed from the
                 backward looking state variables for the normalization.")

    n = get_setting(m, :nx)*get_setting(m, :ns)
    n_backward_looking_vars = length(get_setting(m, :state_indices))
    n_backward_looking_function_valued_vars = get_setting(m, :n_function_valued_backward_looking_states)
    n_backward_looking_scalar_vars = n_backward_looking_vars - n_backward_looking_function_valued_vars

    m <= Setting(:n_backward_looking_states, n*n_backward_looking_distr_vars +
                 n_backward_looking_scalar_vars - get_setting(m, :backward_looking_states_normalization_factor),
                 "Number of state variables, in the true sense (fully
                  backward looking) accounting for the discretization across the grid
                  of function-valued variables and the normalization of
                  distributional variables.")

    ##################################
    # Calculating the number of jumps
    ##################################
    n_jump_distr_vars = get_setting(m, :n_jump_distributional_vars)
    m <= Setting(:jumps_normalization_factor,
                 n_dof_removed*n_jump_distr_vars, "The number of dimensions removed from the
                 jump variables for the normalization.")

    n_jump_vars = length(get_setting(m, :jump_indices))
    n_jump_function_valued_vars = get_setting(m, :n_function_valued_jumps)
    n_jump_scalar_vars = n_jump_vars - n_jump_function_valued_vars

    m <= Setting(:n_jumps, n*n_jump_function_valued_vars +
                 n_jump_scalar_vars - get_setting(m, :jumps_normalization_factor),
                 "Number of jump variables (forward looking) accounting for
                  the discretization across the grid of function-valued variables and the
                  normalization of distributional variables.")

    m <= Setting(:n_model_states, get_setting(m, :n_backward_looking_states) + get_setting(m, :n_jumps),
                 "Number of 'states' in the state space model. Because backward and forward
                 looking variables need to be explicitly tracked for the Klein solution
                 method, we have n_states and n_jumps")
end

function setup_indices!(m::HetDSGELag)
    nx = get_setting(m, :nx)
    ns = get_setting(m, :ns)
    endo = m.endogenous_states_unnormalized
    eqconds = m.equilibrium_conditions
    nxns = nx*ns

    # Endogenous function-valued states
    endo[:l′_t1]    = 1:nxns            # lagged ell function
    endo[:μ′_t1]    = nxns+1:2*nxns     # lagged wealth distribution
    #endogenous scalar-valued states
    endo[:k′_t]  = 2*nxns+1:2*nxns+1             # capital –dont get confused with steadystate object K
    endo[:R′_t1] = 2*nxns+2:2*nxns+2              # lagged real interest rate
    endo[:i′_t1] = 2*nxns+3:2*nxns+3             # lagged nominal interest rate
    endo[:π′_t1] = 2*nxns+4:2*nxns+4             # lagged inflation
    endo[:π′_t2] = 2*nxns+5:2*nxns+5             # double-lagged inflation
    endo[:π′_t3] = 2*nxns+6:2*nxns+6             # triple-lagged inflation
    endo[:y′_t1] = 2*nxns+7:2*nxns+7             # lagged gdp
    endo[:y′_t2] = 2*nxns+8:2*nxns+8             # double-lagged gdp
    endo[:y′_t3] = 2*nxns+9:2*nxns+9             # triple-lagged gdp
    endo[:y′_t4] = 2*nxns+10:2*nxns+10           # quad-lagged gdp
    endo[:z′_t1] = 2*nxns+11:2*nxns+11           # lag tfp growth
    endo[:z′_t2] = 2*nxns+12:2*nxns+12           # double-lag tfp growth
    endo[:z′_t3] = 2*nxns+13:2*nxns+13           # triple-lag tfp growth
    endo[:w′_t1] = 2*nxns+14:2*nxns+14           # lag real wages
    endo[:I′_t1] = 2*nxns+15:2*nxns+15           # lag investment–don't get this confused with i, the nominal interest rate
    # exogenous scalar-valued states:
    endo[:B′]    = 2*nxns+16:2*nxns+16        # discount factor shock
    endo[:G′]    = 2*nxns+17:2*nxns+17        # govt spending
    endo[:z′_t]    = 2*nxns+18:2*nxns+18        # tfp growth
    endo[:MU′]   = 2*nxns+19:2*nxns+19        # investment shock
    endo[:LAMW′] = 2*nxns+20:2*nxns+20        # wage markup
    endo[:LAMF′] = 2*nxns+21:2*nxns+21        # price markup
    endo[:MON′]  = 2*nxns+22:2*nxns+22        # monetary policy shock

 # function-valued jumps
    endo[:l′_t]  = 2*nxns+23:3*nxns+22 # ell function
    endo[:μ′_t]  = 3*nxns+23:4*nxns+22 # wealth distribution
    #scalar-valued jumps
    endo[:R′_t]  = 4*nxns+23:4*nxns+23        # real interest rate
    endo[:i′_t]  = 4*nxns+24:4*nxns+24        # nominal interest rate
    endo[:t′_t]  = 4*nxns+25:4*nxns+25        # transfers + dividends
    endo[:w′_t]  = 4*nxns+26:4*nxns+26        # real wage
    endo[:L′_t]  = 4*nxns+27:4*nxns+27        # hours worked
    endo[:π′_t]  = 4*nxns+28:4*nxns+28        # inflation
    endo[:wageinflation′_t] = 4*nxns+29:4*nxns+29        # nominal wage inflation
    endo[:mu′_t] =  4*nxns+30:4*nxns+30        # average marginal utility
    endo[:y′_t]  = 4*nxns+31:4*nxns+31       # gdp
    endo[:I′_t]  = 4*nxns+32:4*nxns+32        # investment
    endo[:mc′_t] = 4*nxns+33:4*nxns+33        # marginal cost - this is ζ in HetDSGELagₖd.pdf
    endo[:Q′_t]  = 4*nxns+34:4*nxns+34        # Tobin's qfunction
    endo[:capreturn′_t] = 4*nxns+35:4*nxns+35        # return on capital

    nvars = 4*nxns+35
    nscalars = 35 # num eqs which output scalars
    nyscalars = 13 # num scalar jumps
    nxscalars = nscalars - nyscalars # num scalar states

#=endo[:l_t1] = endo[:l′_t1] .+ nvars #LELL
endo[:μ_t1] = endo[:μ′_t1] .+ nvars #LM
endo[:k_t] = endo[:k′_t] .+ nvars #KK
endo[:R_t1] = endo[:R′_t1] .+ nvars #LRR
endo[:i_t1] = endo[:i′_t1] .+ nvars #LII
endo[:π_t1] = endo[:π′_t1]  .+ nvars #LPI
endo[:π_t2] = endo[:π′_t2]  .+ nvars #L2PI
endo[:π_t3] = endo[:π′_t3]  .+ nvars #L3PI
endo[:y_t1] = endo[:y′_t1] .+ nvars #LY
endo[:y_t2] = endo[:y′_t2] .+ nvars #L2Y
endo[:y_t3] = endo[:y′_t3] .+ nvars #l3Y
endo[:y_t4] = endo[:y′_t4] .+ nvars #L4Y
endo[:z_t1] = endo[:z′_t1] .+ nvars #LZ
endo[:z_t2] = endo[:z′_t2] .+ nvars #L2Z
endo[:z_t3] = endo[:z′_t3] .+ nvars #L3Z
endo[:w_t1] = endo[:w′_t1] .+ nvars #LW
endo[:I_t1] = endo[:I′_t1] .+ nvars #LX
endo[:B] = endo[:BP] .+ nvars #B
endo[:G] = endo[:GP] .+ nvars #G
endo[:Z] = endo[:ZP] .+ nvars #Z
endo[:MU] = endo[:MUP] .+ nvars #MU
endo[:LAMW] = endo[:LAMWP] .+ nvars #LAMW
endo[:LAMF] = endo[:LAMFP] .+ nvars #LAMF
endo[:MON] = endo[:MONP] .+ nvars #MON
endo[:ELL] = endo[:l′_t] .+ nvars #ELL
endo[:μ_t] = endo[:μ′_t] .+ nvars #M
endo[:R_t] = endo[:R′_t] .+ nvars #RR
endo[:i_t] = endo[:i′_t] .+ nvars #II
endo[:t_t] = endo[:t′_t]  .+ nvars #TT
endo[:w_t] = endo[:w′_t] .+ nvars #W
endo[:L_t] = endo[:L′_t] .+ nvars #HH
endo[:π_t] = endo[:π′_t] .+ nvars #PI
endo[:wageinflation_t] = endo[:wageinflation′_t] .+ nvars #PIW
endo[:mu_t] = endo[:mu′_t] .+ nvars #LAM
endo[:y_t] = endo[:y′_t] .+ nvars #Y
endo[:I_t] = endo[:I′_t]  .+ nvars #X
endo[:mc_t] = endo[:mc′_t]  .+ nvars #MC
endo[:Q_t] = endo[:Q′_t]  .+ nvars #Q
endo[:capreturn_t] = endo[:capreturn′_t] .+ nvars #RK =#

# create objects needed for solve.jl
# we will order function blocks as follows:
# 1. all function blocks which output a function (first real eqs, then lags)
# 2. all function blocks which map functions to scalars
# 3. all scalar blocks involving endogenous vbls (first real eqs, then lags)
# 4. shock processes
funops = 1:4 # which operators output a function

    # function blocks which output a function
    eqconds[:eq_euler]              = 1:nxns
    eqconds[:eq_kolmogorov_fwd]     = nx*ns+1:2*nx*ns
    eqconds[:eq_lag_ell]            = 2*nxns+1:3*nxns
    eqconds[:eq_lag_wealth]         = 3*nxns+1:4*nxns
    # function blocks which map functions to scalars
    eqconds[:eq_market_clearing]    = 4*nxns+1:4*nxns+1
    eqconds[:eq_lambda]             = 4*nxns+2:4*nxns+2
    #scalar blocks involving endogenous variables
    eqconds[:eq_transfers]            = 4*nxns+3:4*nxns+3 # transfers
    eqconds[:eq_investment]          = 4*nxns+4:4*nxns+4 # investment
    eqconds[:eq_tobin_q]              = 4*nxns+5:4*nxns+5 # tobin's q
    eqconds[:eq_capital_accumulation] = 4*nxns+6:4*nxns+6 # capital accumulation
    eqconds[:eq_wage_phillips] = 4*nxns+7:4*nxns+7 # wage phillips curve
    eqconds[:eq_price_phillips] = 4*nxns+8:4*nxns+8 # price phillips curve
    eqconds[:eq_marginal_cost]  = 4*nxns+9:4*nxns+9 # marginal cost
    eqconds[:eq_gdp]  = 4*nxns+10:4*nxns+10 # gdp
    eqconds[:eq_optimal_kl] = 4*nxns+11:4*nxns+11 # optimal K/L ratio
    eqconds[:eq_taylor] = 4*nxns+12:4*nxns+12 # taylor rule
    eqconds[:eq_fisher] = 4*nxns+13:4*nxns+13 # fisher eqn
    eqconds[:eq_nominal_wage_inflation] = 4*nxns+14:4*nxns+14 # nominal wage inflation
    # lagged variables
    eqconds[:LR] = 4*nxns+15:4*nxns+15 # LR
    eqconds[:LI] = 4*nxns+16:4*nxns+16 # LI
    eqconds[:LPI] = 4*nxns+17:4*nxns+17 # LPI
    eqconds[:L2PI] = 4*nxns+18:4*nxns+18 # L2PI
    eqconds[:L3PI] = 4*nxns+19:4*nxns+19 # L3PI
    eqconds[:LY] = 4*nxns+20:4*nxns+20 # LY
    eqconds[:L2Y] = 4*nxns+21:4*nxns+21 # L2Y
    eqconds[:L3Y] = 4*nxns+22:4*nxns+22 # L3Y
    eqconds[:L4Y] = 4*nxns+23:4*nxns+23 # L4Y
    eqconds[:LZ]  = 4*nxns+24:4*nxns+24 # LZ
    eqconds[:L2Z] = 4*nxns+25:4*nxns+25 # L2Z
    eqconds[:L3Z] = 4*nxns+26:4*nxns+26 # L3Z
    eqconds[:LW]  = 4*nxns+27:4*nxns+27 # LW
    eqconds[:LX] = 4*nxns+28:4*nxns+28 # LX
    # shocks
    eqconds[:F33] = 4*nxns+29:4*nxns+29 # discount factor B
    eqconds[:F34] = 4*nxns+30:4*nxns+30 # govt spending G
    eqconds[:F35] = 4*nxns+31:4*nxns+31 # tfp growth Z
    eqconds[:F36] = 4*nxns+32:4*nxns+32 # investment MU
    eqconds[:F37] = 4*nxns+33:4*nxns+33 # wage mkup LAMW
    eqconds[:F38] = 4*nxns+34:4*nxns+34 # price mkup LAMF
    eqconds[:F39] = 4*nxns+35:4*nxns+35 # monetary policy MON

# Total grid x*s
m <= Setting(:n, get_setting(m, :nx) * get_setting(m, :ns), "Total grid size, multiplying
                     across grid dimensions.")

m <= Setting(:nvars, 4*get_setting(m, :n) +35, "num variables")
m <= Setting(:nscalars, 35, " # num eqs which output scalars")
m <= Setting(:nyscalars, 13, "num scalar jumps")
m <= Setting(:nxscalars, 35 - 13, "num scalar states")

m.endogenous_states = deepcopy(endo)

end


function NewKrest(r)
    for side in 0:3
        GoToTheWallAndMark(r, HorizonSide(side))
    end
end

function GoToTheWallAndMark(r::Robot, side::HorizonSide)
    x = 0
    while !mainborder(r,side)
        putmarker!(r)
        trymove!(r, side)
        x = x + 1
    end
    putmarker!(r)
    while x != 0
        trymove!(r, HorizonSide((Int(side) + 2) % 4))
        x = x - 1
    end
end

function trymove!(r::Robot, side::HorizonSide)
    x = 0
    if isborder(r,side)
        while isborder(r, side)
            move!(r,HorizonSide((Int(side) + 1) % 4))
            x = x + 1
        end
        move!(r,side)
        while isborder(r,HorizonSide((Int(side) + 3) % 4))
            move!(r,side)
        end
        while x != 0
            move!(r,HorizonSide((Int(side) + 3) % 4))
            x = x - 1
        end
    else
        move!(r,side)
    end
end

function mainborder(r,side)
    x = 0
    b = true
    while isborder(r,side) && !isborder(r,HorizonSide((Int(side) + 3) % 4))
        move!(r,HorizonSide((Int(side) + 3) % 4))
        x = x + 1
    end
    if isborder(r,side)
        b = true
    else
        b = false
    end
    while x != 0
        move!(r,HorizonSide((Int(side) + 1) % 4))
        x = x - 1
    end
    return b
end


using IRTracker
using DynamicPPL
using Distributions
using DataStructures: SortedDict

export strip_dependencies, makegraph, showgraph

try_getindex(expr::TapeReference) = expr[]
try_getindex(expr::TapeConstant) = getvalue(expr)

const CallingNode{F} = Union{PrimitiveCallNode{<:Any, F}, NestedCallNode{<:Any, F}}

istilde(::CallingNode{typeof(DynamicPPL.tilde_assume)}) = true
istilde(::CallingNode{typeof(DynamicPPL.tilde_observe)}) = true
istilde(::CallingNode{typeof(DynamicPPL.dot_tilde_assume)}) = true
istilde(::CallingNode{typeof(DynamicPPL.dot_tilde_observe)}) = true
istilde(::AbstractNode) = false

"""
    tilde_parameters(node)

Extract a tuple of variablen name, distribution, and value argument of any tilde call.

The variable name can be `nothing` for constant observations.  The whole value is `nothing` for
nodes that are not tilde calls.
"""
tilde_parameters(node::AbstractNode) = nothing

function tilde_parameters(node::CallingNode{typeof(DynamicPPL.tilde_assume)})
    args = getarguments(node) # rng, ctx, sampler, right, vn, inds, vi
    vn, dist, value, vi = args[5], args[4], getvalue(node), args[7]
    return vn, dist, TapeConstant(value), vi
end

function tilde_parameters(node::CallingNode{typeof(DynamicPPL.tilde_observe)})
    args = getarguments(node)
    if length(args) == 7
        # ctx, sampler, right, left, vname, vinds, vi
        vn, dist, value, vi = args[5], args[3], args[4], args[7]
        return vn, dist, value, vi
    elseif length(args) == 5
        # ctx, sampler, right, left, vi
        dist, value, vi = args[3], args[4], args[5]
        return nothing, dist, value, vi
    else
        throw(ArgumentError("$node has unknown argument structure, you should not have reached this point!"))
    end
end

function tilde_parameters(node::CallingNode{typeof(DynamicPPL.dot_tilde_assume)})
    args = getarguments(node) # rng, ctx, sampler, right, left, vn, inds, vi
    vn, dist, value, vi = args[6], args[4], args[5], args[8]
    return vn, dist, value, vi
end

function tilde_parameters(node::CallingNode{typeof(DynamicPPL.dot_tilde_observe)})
    args = getarguments(node)
    if length(args) == 7
        # ctx, sampler, right, left, vn, inds, vi
        vn, dist, value, vi = args[5], args[3], args[4], args[7]
        return vn, dist, value, vi
    elseif length(args) == 5
        # ctx, sampler, right, left, vi
        dist, value, vi = args[3], args[4], args[5]
        return nothing, dist, value, vi
    else
        throw(ArgumentError("$node has unknown argument structure, you should not have reached this point!"))
    end
end

"""
    model_argument_nodes(root)

Extract from a stripped model trace all model argument `getindex` calls.
"""
function model_argument_nodes(root)
    # From the beginning of the trace,
    # ```
    # @1: [Arg:§1:%1] @11#1 → ##evaluator#465
    # @2: [Arg:§1:%2] @11#2 → Random._GLOBAL_RNG()
    # @3: [Arg:§1:%3] @11#3 → Model{...}
    # @4: [Arg:§1:%4] @11#4 → VarInfo (0 variables, dimension 0; logp: 0.0)
    # @5: [Arg:§1:%5] @11#5 → SampleFromPrior()
    # @6: [Arg:§1:%6] @11#6 → DefaultContext()
    # @7: [§1:%7] ⟨getproperty⟩(@3, ⟨:args⟩) → (x = [0.1, 0.05, 1.0],)
    # @8: [§1:%8] ⟨getproperty⟩(@7, ⟨:x⟩) → [0.1, 0.05, 1.0]
    # ```
    # extract only the `getproperty(@n, ⟨:varname⟩)` line for each argument.

    argument_nodes = Vector{AbstractNode}()
    model_node = getchildren(root)[3]
    modelargs_nodes = Vector{AbstractNode}()

    for child in getchildren(root)
        if child isa CallingNode{typeof(getproperty)}
            value, property_name = getarguments(child)
            if getvalue(property_name) == :args && try_getindex(value) == model_node
                push!(modelargs_nodes, child)
            elseif try_getindex(value) ∈ modelargs_nodes
                # a `getproperty(args, symbol)` whose parent is a `args = getproperty(model, :args)`
                # is a model argument, for sure
                push!(argument_nodes, child)
            end
        end
    end

    return argument_nodes
end

"""
    add_candidates!(dependencies, candidates, mutants, node)

Traverse refernces of `node` backwards and turn the reached `candidates` into `dependencies`.
"""
function add_candidates!(dependencies, candidates, mutants, node)
    current_refs = Vector{AbstractNode}(referenced(node))
    push!(dependencies, node)

    while !isempty(current_refs)
        node = pop!(current_refs)
        node ∈ dependencies && continue

        if node ∈ candidates
            push!(dependencies, node)
            union!(current_refs, referenced(node))
        end

        if haskey(mutants, node)
            push!(dependencies, node)
            union!(current_refs, mutants[node])
        end
    end
end

"""
    strip_dependencies(root)

Slice out only those nodes that are between tildes or model arguments, taking some care to get
mutated things (vector + `setindex!` etc.) right.
"""
function strip_dependencies(root)
    # "dependencies" are nodes that are in the dependency graph for sure (i.e., tildes).
    # "candidates" are nodes that follow a dependency node; are included as dependencies, if the
    # chain ends up in a dependency node again (e.g., x ~ D, y = f(x), z ~ D2(y)) -- this is
    # what `add_candidates!` does by working backwards.
    # "mutants" are backedges (from mutated to mutators) in the dependency graph, resulting from
    # later operations mutating earlier values(e.g, z = zeros(), z[1] ~ D, x ~ D2(z[1])).

    dependencies = model_argument_nodes(root)
    candidates = Set{AbstractNode}()
    mutants = Dict{AbstractNode, Vector{AbstractNode}}()

    for child in getchildren(root)
        params = tilde_parameters(child)

        if !isnothing(params)
            # tilde node: is a dependency for sure
            push!(dependencies, child)
        else
            # non-tilde node: make candidate, if any parent is a candidate or dependency

            # this `t` is critical -- exploding type inference!
            t = any(r ∈ candidates || r ∈ dependencies for r in referenced(child))
            if t
                push!(candidates, child)

                if child isa CallingNode{typeof(setindex!)}
                    # if the candidate is mutating an array, also make the mutated object a candidate
                    # and record the backedge from the mutated one to:
                    # `n = setindex!(x, v, i)` makes `x` a candidate, and a backedge from `x` to `n`.

                    mutated = getargument(child, 1)
                    if mutated isa TapeReference
                        push!(candidates, mutated[])
                        push!(get!(Vector{AbstractNode}, mutants, mutated[]), child)
                    end
                end
            end
        end
    end

    # Go over nodes once more and turn candidates between dependencies into dependencies, too.
    for child in getchildren(root)
        params = tilde_parameters(child)
        if !isnothing(params)
            vn, dist, value = params
            vn isa TapeReference && add_candidates!(dependencies, candidates, mutants, vn[])
            dist isa TapeReference && add_candidates!(dependencies, candidates, mutants, dist[])
            value isa TapeReference && add_candidates!(dependencies, candidates, mutants, value[])
        end
    end

    return sort(dependencies; by=getposition) |> unique
end

struct Reference
    number::Int
end

Base.show(io::IO, r::Reference) = print(io, "⟨", r.number, "⟩")
Base.isless(q::Reference, r::Reference) = isless(q.number, r.number)
Base.hash(r::Reference, h::UInt) = hash(r.number, h)
Base.:(==)(q::Reference, r::Reference) = q.number == r.number

abstract type Statement end

struct Assumption{dotted, TN<:VarName, TVal} <: Statement
    vn::TN
    dist_ref::Reference
    value::TVal
    logp::Float64
end

Assumption{dotted}(vn, dist_ref, value, logp) where {dotted} =
    Assumption{dotted, typeof(vn), typeof(value)}(vn, dist_ref, value, logp)

struct Observation{dotted, TN<:Union{Nothing, VarName}, TVal} <: Statement
    vn::TN
    dist_ref::Reference
    value::TVal
    logp::Float64
end

Observation{dotted}(vn, dist_ref, value, logp) where {dotted} =
    Observation{dotted, typeof(vn), typeof(value)}(vn, dist_ref, value, logp)

struct Call{TD<:Union{Nothing, Tuple{VarName, Reference}}, TF, TArgs<:Tuple, TVal} <: Statement
    definition::TD
    f::TF
    args::TArgs
    value::TVal
end

Call(f, args, value) = Call(nothing, f, args, value)

struct Constant{TVal} <: Statement
    value::TVal
end

const Tilde{dotted} = Union{Assumption{dotted}, Observation{dotted}}

_shortname(d::Type{<:Distribution}) = string(nameof(d))
_shortname(other) = string(other)

function Base.show(io::IO, stmt::Assumption)
    print(io, stmt.vn, (isdotted(stmt) ? " .~ " : " ~ "), stmt.dist_ref, " → ", stmt.value)
end
function Base.show(io::IO, stmt::Observation)
    print(io, (isnothing(stmt.vn) ? stmt.value : stmt.vn))
    print(io, (isdotted(stmt) ? " .⩪ " : " ⩪ "), stmt.dist_ref,  " ← ", stmt.value)
end
function Base.show(io::IO, stmt::Call)
    if !isnothing(stmt.definition)
        vn = stmt.definition[1]
        print(io, vn, " = ")
    end
    print(io, stmt.f)
    print(io, "(")
    join(io, stmt.args, ", ")
    print(io, ") → ", stmt.value)
end
Base.show(io::IO, stmt::Constant) = print(io, stmt.value)

IRTracker.getvalue(stmt::Statement) = stmt.value

isdotted(::Type{<:Assumption{dotted}}) where {dotted} = dotted
isdotted(::Type{<:Observation{dotted}}) where {dotted} = dotted
isdotted(tilde::Tilde) = isdotted(typeof(tilde))

getvn(stmt::Tilde) = stmt.vn
getvn(stmt::Call) = isnothing(stmt.definition) ? nothing : stmt.definition[1]
getvn(::Constant) = nothing

"""
    parent_variables(graph, stmt)

Return all `Assumption`s (potentially with indexing) that the tilde `stmt` depends on directly.
"""
function parent_variables(graph, stmt::Tilde)
    result = Set{Pair{VarName, Assumption}}()
    parent_variables!(result, graph, graph[stmt.dist_ref])
    stmt.value isa Reference && parent_variables!(result, graph, stmt.value)
    return result
end

parent_variables!(result, graph, stmt) = result
function parent_variables!(result, graph, ref::Reference)
    # follow back implicit dependencies on unmarked locations without indexing, e.g.
    # ⟨2⟩ = [0.1, -0.2]
    # ⟨3⟩ = Array{Float64,1}(array initializer with undefined values, 2) → ...
    # ⟨6⟩ = m[1] ~ Normal() → ...
    # ⟨7⟩ = m[1] = getindex(⟨3⟩, 1) → ...
    # ⟨10⟩ = m[2] ~ Normal() → ...
    # ⟨11⟩ = m[2] = getindex(⟨3⟩, 2) → ...
    # ⟨13⟩ = x ⩪ MvNormal(⟨3⟩) ← ⟨2⟩
    # `x` depends on the whole of `m`, stored in `⟨3⟩`, and thus implicitly
    # on `m[1]` and `m[2]` (cf. reverse_deps test case).

    if haskey(graph.rv_locations, ref)
        for (ix, loc) in graph.rv_locations[ref]
            loc == ref && continue # skip one-line array tildes, e.g., z ~ filldist(...)
            parent_variables!(result, graph, loc)
        end
    end

    return parent_variables!(result, graph, graph[ref])
end
function parent_variables!(result, graph, stmt::Assumption)
    # direct dependency:
    # ⟨31⟩ = w ~ Dirichlet()
    # ⟨42⟩ = z[1] ~ DiscreteNonParametric(⟨31⟩)
    return push!(result, stmt.vn => stmt)
end
function parent_variables!(result, graph, stmt::Observation)
    @warn "The parent statement is an observation ($(stmt)), something weird has happened..."
    return result
end
function parent_variables!(result, graph, stmt::Call{<:Nothing})
    # go through all argumens of an unmarked location:
    # ⟨85⟩ = getindex(⟨2⟩, ⟨72⟩)
    # ⟨86⟩ = x[3] ⩪ Normal(⟨80⟩, 1.0) ← ⟨85⟩

    for arg in stmt.args
        parent_variables!(result, graph, arg)
    end

    return result
end
function parent_variables!(result, graph, stmt::Call{<:Tuple, typeof(getindex)})
    # special case: don't add whole array (`z`) for indexing calls
    # ⟨12⟩ = z ~ filldist(⟨10⟩, ⟨5⟩)
    # ⟨30⟩ = z[2] = getindex(⟨12⟩, ⟨28⟩)
    # ⟨31⟩ = μ[2] = getindex(⟨7⟩, ⟨30⟩)
    for arg in Base.tail(stmt.args)
        parent_variables!(result, graph, arg)
    end
    vn, location = stmt.definition
    push!(result, vn => graph[location])
    return result
end

struct Graph
    """Entries in graph: mapping from references to statements."""
    statements::SortedDict{Reference, Statement}

    """Association which nodes in the trace are associated with which references in the graph"""
    reference_mapping::Dict{AbstractNode, Reference}

    """
    Remembers in which references random variables are actually stored; this is either the location
    of the tilde statement (`beta ~ Normal()`, `z ~ filldist(...)`), or the array position it is
    assigned to (`z[i] ~ Normal()`).

    The container is first ordered by reference of storage, then by index for `subsumes` testing.
    """
    rv_locations::Dict{Reference, Vector{Pair{Tuple, Reference}}}
end

Graph() = Graph(SortedDict{Reference, Statement}(),
                Dict{AbstractNode, Reference}(),
                Dict{Reference, Vector{Pair{Tuple, Reference}}}())

Base.IteratorSize(::Type{Graph}) = Base.HasLength()
Base.length(graph::Graph) = length(graph.statements)
Base.IteratorEltype(::Type{Graph}) = Base.HasEltype()
Base.eltype(graph::Graph) = eltype(graph.statements)
Base.getindex(graph::Graph, ref) = graph.statements[ref]
Base.setindex!(graph::Graph, stmt, ref) = graph.statements[ref] = stmt
Base.firstindex(graph::Graph) = firstindex(graph.statements)
Base.lastindex(graph::Graph) = lastindex(graph.statements)
Base.get(graph::Graph, ref, default) = get(graph.statements, ref, default)
Base.haskey(graph::Graph, ref) = haskey(graph.statements, ref)
Base.delete!(graph::Graph, ref) = delete!(graph.statements, ref)
Base.keys(graph::Graph) = keys(graph.statements)
Base.values(graph::Graph) = values(graph.statements)
Base.iterate(graph::Graph) = iterate(graph.statements)
Base.iterate(graph::Graph, state) = iterate(graph.statements, state)

# just intended for debugging
Base.getindex(graph::Graph, i::Int) = graph[Reference(i)]

# to convert an atom to a runtime value: either the value of the referenced node, or a literal
tovalue(graph, ref::Reference) = getvalue(graph[ref])
tovalue(graph, other) = other

function Base.mapreduce(f, op, graph::Graph; init)
    for kv in graph
        init = op(init, f(kv))
    end

    return init
end

function Base.filter(f, graph::Graph; init=Vector{eltype(graph)}())
    for kv in graph
        f(kv) && push!(init, kv)
    end

    return init
end

Base.map(f, graph::Graph; init=Vector{eltype(graph)}()) = mapreduce(f, push!; init=init)

getstatements(graph::Graph) = graph.statements

hasmapping(graph::Graph, node) = haskey(graph.reference_mapping, node)
function getmapping(graph::Graph, node)
    ref = graph.reference_mapping[node]
    return ref
end
function getmapping(graph::Graph, node, default)
    if hasmapping(graph, node)
        getmapping(graph, node)
    else
        return default
    end
end
setmapping!(graph::Graph, (node, ref)::Pair) = graph.reference_mapping[node] = ref

function deletemapping!(graph::Graph, node)
    if hasmapping(graph, node)
        delete!(graph, getmapping(graph, node))
    end
    return graph
end

"""Look up the matching location of an RV in graph.rv_locations[ref] for given index tuple."""
function _find_indexing(indexings, indexing)
    for (ix, r) in indexings
        if DynamicPPL.subsumes(ix, indexing)
            return r
        end
    end
    return nothing
end

function getrv(graph::Graph, ref::Reference, indexing)
    _find_indexing(graph.rv_locations[ref], indexing)
end

function setrv!(graph, ((ref, indexing), vn))
    push!(get!(graph.rv_locations, ref, valtype(graph.rv_locations)()),
          indexing => vn)
end

function hasrv(graph, ref, indexing)
    indexings = get(graph.rv_locations, ref, nothing)
    return !isnothing(indexings) && !isnothing(_find_indexing(indexings, indexing))
end

function makereference!(graph, node)
    newref = Reference(length(graph.reference_mapping) + 1)
    setmapping!(graph, node => newref)
    return newref
end

function Base.show(io::IO, graph::Graph)
    for (ref, stmt) in graph
        println(io, ref, " = ", stmt)
    end
end

function convertcall(graph, node::CallingNode)
    f = convertvalue(graph, getfunction(node))
    args = convertvalue.(Ref(graph), getarguments(node))
    return Call(f, args, getvalue(node))
end

convertvalue(graph, expr::TapeReference) = getmapping(graph, expr[], getvalue(expr))
convertvalue(graph, expr::TapeConstant) = getvalue(expr)

convertdist!(graph, dist_expr::TapeConstant) = Constant(getvalue(dist))
function convertdist!(graph, dist_expr::TapeReference)
    ref = getmapping(graph, dist_expr[])
    return ref
end

convertvn!(graph, vn_expr::Nothing) = nothing
convertvn!(graph, vn_expr::TapeConstant) = getvalue(vn_expr)
function convertvn!(graph, vn_expr::TapeReference)
    # @108: ⟨tuple⟩(@82) = (2,)
    # @109: ⟨tuple⟩(@108) = ((2,),)
    # @110: ⟨VarName⟩(⟨:x⟩, @109) = x[2]

    # extract the nodes that compose the indices of a varname
    vn_node = vn_expr[] # @110[]
    vn_name = getvalue(getargument(vn_node, 1))
    indices_node = try_getindex(getargument(vn_node, 2)) # @109[]
    index_element_nodes = getindex.(getarguments(indices_node)) # (@108[],)
    index_nodes = Tuple(try_getindex.(getarguments(index)) for index in index_element_nodes) # ((@82[],),)
    index_refs = map(ix -> getmapping.(Ref(graph), ix, ix), index_nodes)  # ((getmapping(graph, @82[]),),)

    # delete statements of all nodes that were involved in the construction (ie., @108, @109, @110)
    deletemapping!(graph, vn_node)
    deletemapping!(graph, indices_node)
    deletemapping!.(Ref(graph), index_element_nodes)

    return VarName(vn_name, map.(ix -> tovalue(graph, ix), index_refs))
end

function pushtilde!(graph, callingnode, tilde_constructor)
    vn_expr, dist_expr, value_expr, vi_expr = tilde_parameters(callingnode)
    vn = convertvn!(graph, vn_expr)
    dist_ref = convertdist!(graph, dist_expr)
    value = convertvalue(graph, value_expr)

    ref = makereference!(graph, callingnode)
    graph[ref] = tilde_constructor(vn, dist_ref, value, -Inf)

    if !isnothing(vn)
        indexing_values = tovalue.(Ref(graph), DynamicPPL.getindexing(vn))
        if isdotted(tilde_constructor)
            @warn "Broadcasted tildes ($(graph[ref])) are not fully supported!"
            setrv!(graph, (value, indexing_values) => ref)
        elseif indexing_values == ()
            # this is the case where an unindexed variable is sampled -- the value is equivalent to
            # the current reference; otherwise, a `setindex!` should follow in the next line.
            setrv!(graph, (ref, indexing_values) => ref)
        end
    end

    return graph
end

pushnode!(graph, node::CallingNode{typeof(DynamicPPL.tilde_assume)}) =
    pushtilde!(graph, node, Assumption{false})
pushnode!(graph, node::CallingNode{typeof(DynamicPPL.dot_tilde_assume)}) =
    pushtilde!(graph, node, Assumption{true})
pushnode!(graph, node::CallingNode{typeof(DynamicPPL.tilde_observe)}) =
    pushtilde!(graph, node, Observation{false})
pushnode!(graph, node::CallingNode{typeof(DynamicPPL.dot_tilde_observe)}) =
    pushtilde!(graph, node, Observation{true})

function pushnode!(graph, node::CallingNode)
    ref = makereference!(graph, node)
    graph[ref] = convertcall(graph, node)
    return graph
end

function pushnode!(graph, node::CallingNode{typeof(DynamicPPL.matchingvalue)})
    # special handling for model arguments
    # @7: ⟨getproperty⟩(@6, ⟨:x⟩) = [0.5, 1.1]
    # @8: ⟨DynamicPPL.matchingvalue⟩(@4, @3, @7) = [0.5, 1.1]

    value_expr = getargument(node, 3) # @7 above
    if value_expr isa TapeReference
        getproperty_node = value_expr[]
        delete!(graph, getmapping(graph, getproperty_node))
    else # value vas a literal, not a function argument
        @warn "this probably shouldn't have happened..."
    end

    ref = makereference!(graph, node)
    graph[ref] = Constant(getvalue(node))
    return graph
end

function pushnode!(graph, node::CallingNode{typeof(setindex!)})
    # @32: [§5:%34] ⟨DynamicPPL.tilde_assume⟩(@5, @4, @20, @29, @31, @3) = -0.031672055938280076
    # @33: [§5:%35] ⟨setindex!⟩(@19, @32, ⟨1⟩) = [...]
    argument_exprs = try_getindex.(getarguments(node))
    arguments = getmapping.(Ref(graph), argument_exprs, argument_exprs)
    mutated, value, indexing = arguments[1], arguments[2], arguments[3:end]
    indexing_values = (tovalue.(Ref(graph), indexing),)

    if value isa Reference && graph[value] isa Tilde
        # this is the case where in the previous line, an indexed RV was sampled; we need
        # to associate the value with the array position it is assigned to here.
        setrv!(graph, (mutated, indexing_values) => value)

        # # we can also replace the `setindex!` call that followed the tilde node
        # # by a `getindex` on the same value, to preserve the information of the dependency of
        # # between the variable, the array, and the index
        # # PROBABLY UNNEEDED? see reverse_deps test case in test_conditionals
        definition = (graph[value].vn, value)
        ref = makereference!(graph, node)
        graph[ref] = Call(definition, getindex, (mutated, indexing...), tovalue(graph, value))
    else
        invoke(pushnode!, Tuple{Graph, CallingNode}, graph, node)
    end

    return graph
end

function pushnode!(graph, node::CallingNode{typeof(getindex)})
    # @180: ⟨getindex⟩(@9, @135) → -0.05
    argument_exprs = try_getindex.(getarguments(node))
    arguments = getmapping.(Ref(graph), argument_exprs, argument_exprs)
    array, indexing = arguments[1], arguments[2:end]
    indexing_values = (tovalue.(Ref(graph), indexing),)

    if hasrv(graph, array, indexing_values)
        rv = getrv(graph, array, indexing_values)
        ref = makereference!(graph, node)
        definition = (VarName(graph[rv].vn, indexing_values), rv)
        graph[ref] = Call(definition, getindex, (array, indexing...), getvalue(node))
    else
        invoke(pushnode!, Tuple{Graph, CallingNode}, graph, node)
    end

    return graph
end

function pushnode!(graph, node::ConstantNode)
    ref = makereference!(graph, node)
    graph[ref] = Constant(getvalue(node))
    return graph
end

function pushnode!(graph, node::ArgumentNode)
    if !isnothing(node.branch_node)
        # don't add, but remember branch arguments
        original_ref = getmapping(graph, referenced(node)[1])
        setmapping!(graph, node => original_ref)
    else
        # handle function argument nodes like constants
        ref = makereference!(graph, node)
        graph[ref] = Constant(getvalue(node))
    end

    return graph
end

function makegraph(slice::Vector{<:AbstractNode})
    return foldl(slice, init=Graph()) do graph, node
        pushnode!(graph, node)
    end
end


module ProgramConstantsModule

const MAX_DEGREE = 2
const CONST_TYPE = Float32
const BATCH_DIM = 2
const FEATURE_DIM = 1
const RecordType = Dict{String,Any}

"""Enum for concurrency type (to get function specialization)"""
abstract type SRConcurrency end
struct SRSerial <: SRConcurrency end
struct SRThreaded <: SRConcurrency end
struct SRDistributed <: SRConcurrency end

end


n_actions(p::LaserTagPOMDP) = p.diag_actions ? 9 : 5
POMDPs.actions(p::LaserTagPOMDP) = 1:n_actions(p)
POMDPs.actionindex(p::LaserTagPOMDP, a::Int) = a

const ACTION_NAMES = SVector("north",
                             "east",
                             "south",
                             "west",
                             "tag",
                             "northeast",
                             "southeast",
                             "southwest",
                             "northwest")

const TAG_ACTION = 5

const ACTION_DIRS = SVector(Coord(0,1),
                            Coord(1,0),
                            Coord(0,-1),
                            Coord(-1,0),
                            Coord(0,0),
                            Coord(1,1),
                            Coord(1,-1),
                            Coord(-1,-1),
                            Coord(-1,1))

const CARDINALS = SVector(Coord(0,1),
                            Coord(1,0),
                            Coord(0,-1),
                            Coord(-1,0))

const DIAGONALS = SVector(Coord(1,1),
                         Coord(1,-1),
                         Coord(-1,-1),
                         Coord(-1,1))

const DIR_TO_ACTION = Dict{Coord, Int}(c=>i for (i,c) in enumerate(ACTION_DIRS))


module Borg

include("BorgC.jl")

export createProblem, solve

function eval_wrapper(vars_raw::Ptr{Float64}, objs_raw::Ptr{Float64}, consts_raw::Ptr{Float64}, userParams::Ptr{Void})
    problem = unsafe_pointer_to_objref(userParams)::BorgProblem

    vars = pointer_to_array(vars_raw, problem.numberOfVariables)
    objs = pointer_to_array(objs_raw, problem.numberOfObjectives)
    consts = pointer_to_array(consts_raw, problem.numberOfConstraints)

    problem.eval_fn(vars, objs, consts)

    return
end

type BorgProblem
	ref::BorgC.BORG_Problem

	numberOfVariables::Int64
	numberOfObjectives::Int64
	numberOfConstraints::Int64

	eval_fn::Function

	function BorgProblem(
            numberOfVariables::Int,
            numberOfObjectives::Int,
            numberOfConstraints::Int,
            eval_fn::Function)

		problem = new()

        problem.numberOfVariables = numberOfVariables
        problem.numberOfObjectives = numberOfObjectives
        problem.numberOfConstraints = numberOfConstraints
        problem.eval_fn = eval_fn
		problem.ref = BorgC.BORG_Problem_create(numberOfVariables, numberOfObjectives, numberOfConstraints, eval_wrapper, problem)

        return problem
	end
end

function createProblem(
    varLowerBounds::Vector{Float64},
    varUpperBounds::Vector{Float64},
    objEpsilons::Vector{Float64},
    numberOfConstraints::Int,
    eval_fn::Function)

    if length(varLowerBounds) != length(varUpperBounds)
        error("varLowerBounds and varUpperBounds must be of same length.")
    end

    numberOfVariables = length(varLowerBounds)
    numberOfObjectives = length(objEpsilons)

    problem = BorgProblem(numberOfVariables, numberOfObjectives, numberOfConstraints, eval_fn)

    for i in 1:numberOfVariables
        BorgC.BORG_Problem_set_bounds(problem.ref, i, varLowerBounds[i], varUpperBounds[i])
    end

    for i in 1:numberOfObjectives
        BorgC.BORG_Problem_set_epsilon(problem.ref, i, objEpsilons[i])
    end

    return problem
end

function solve(problem::BorgProblem, maxIt::Int)
    archive = BorgC.BORG_Algorithm_run(problem.ref, maxIt)
    return archive
end

end


#=
@doc """
Facilities to help map GUI input / expression into a function that aggregates columns.
We would store that function and use it whenever we need to aggregate a column.

Three ways to instantiate the DataFrameAggr
* Symbol. e.g. `:mean`. It's lowered into `mean(:_)`. See below for the meaning of `:_`
* Expr. e.g. `:( mean(:_,:wcol) )`, `:( quantile( :_, 0.75 ) )`
   symbols (quoted) are interpreted as the named column.
   The underscore `:_` column is special -- it's interpreted on-the-fly as the column that
   needs to be aggregated. So multiple columns can use the same expression `mean(_,:wcol)` and
   they will all use the same weight (from `wcol`), but produces target specific mean.
* Function: (most likely lambda) e.g. `df -> quantile( df[:col], 0.5 )`. It is expected
   to generate either a value (typically a scalar consistent to the column type), or
   a 1-row, 1-col dataframe.
""" ->
=#

DataFrameAggrCache = Dict{Any,Function}()

defaultAggr( ::Type{} ) = :uniqvalue
defaultAggr{T<:Real}( ::Type{T} ) = :sum
defaultAggr{T}( ::Type{Array{T,1}} ) = :unionall

function liftAggrSpecToFunc( c::Symbol, dfa::UTF8String )
    if haskey( DataFrameAggrCache, (c, dfa ) )
        return DataFrameAggrCache[ (c, dfa ) ]
    end
    ret = liftAggrSpecToFunc( c, parse( dfa ) )
    DataFrameAggrCache[ (c, dfa) ] = ret
end
liftAggrSpecToFunc( c::Symbol, dfa::ASCIIString ) = liftAggrSpecToFunc( c, utf8( dfa ) )

function liftAggrSpecToFunc( c::Symbol, dfa::Union{ Function, Symbol, Expr } )
    if typeof( dfa ) == Function
        return dfa
    end
    if haskey( DataFrameAggrCache, (c, dfa) )
        return DataFrameAggrCache[ (c, dfa ) ]
    end
    # "mean" or "Module.aggrfunc"
    if typeof( dfa ) == Symbol || typeof( dfa ) == Expr && dfa.head == :(.) && all( x->typeof(x)==Symbol, dfa.args )
        funnameouter = gensym( "DFAggr" )
        code = :(
            function $funnameouter( _df_::AbstractDataFrame )
            end )
        push!( code.args[2].args, Expr( :call, dfa, Expr( :call, :getindex, :_df_, QuoteNode(c) ) ) )
        ret = eval( code )
    else # expr
        if !Base.Meta.isexpr( dfa, :call )
            error( string( dfa ) * " does not look like an aggregator function")
        end
        # disallow mutating function
        fname = dfa.args[1]
        if Base.Meta.isexpr( dfa, :curly )
            error( "DataFrameAggr: curly not supported")
        elseif !( typeof( fname ) <: Symbol ) && !Base.Meta.isexpr( fname, :(.) )
            error( "DataFrameAggr: only simple function name please")
        end
        if contains( string(fname), "!" )
            error( string(fname) * " seems to have side effects" )
        end

        # replace _ with _df_[$c], and then leverage @with

        # before we do that, note that
        # (A) in DataFramesMeta, macro converts :x to Expr( :quote, :x ) but
        #     :( :x ) or parse( ":x" ) it is actually QuoteNode( :x ),
        #     so we need to do a little conversion in order to leverage that package
        cdfa = deepcopy( dfa )
        convertExpression!( cdfa, c )

        membernames = Dict{Union{Symbol,Expr}, Symbol}()
        cdfa = DataFramesMeta.replace_syms(cdfa, membernames)
        funargs = map(x -> :( getindex( _df_, $(x)) ), collect(keys(membernames)))
        funnameouter = gensym("DFAggr")
        funname = gensym()
        code = quote
            function $funnameouter( _df_::AbstractDataFrame )
                function $funname($(collect(values(membernames))...))
                    $cdfa
                end
                $funname($(funargs...))
            end
        end
        ret = eval( code )
    end
    DataFrameAggrCache[ (c, dfa) ] = ret
end

# used by aggregation lifting and calcpivot lifting
function convertExpression!( ex::Expr, column_ctx::Symbol = Symbol("") )
    for i in 1:length( ex.args )
        a = ex.args[i]
        if typeof( a ) == QuoteNode
            if a.value == :_ && column_ctx != Symbol("")
                ex.args[i] = Expr( :quote, column_ctx )
            else
                ex.args[i] = Expr( :quote, a.value )
            end
        elseif typeof( a ) == Expr
            convertExpression!( a )
        end
    end
end

function unionall( x::AbstractDataArray )
    l = dropna( x )
    t = eltype( eltype( x ) )
    s = Set{t}()
    for el in l
        push!( s, el... )
    end
    collect( s )
end

function unionall( x::Array )
    t = eltype( eltype( x ) )
    s = Set{t}()
    for el in x
        push!( s, el... )
    end
    collect( s )
end

function uniqvalue( x::AbstractDataArray; skipna::Bool=true )
    lvls = DataArrays.levels(x)
    if skipna
        l = dropna( lvls )
        if length(l) == 1
            return l[1]
        end
        return NA
    end
    if length(lvls) == 1
        return lvls[1]
    end
    return NA
end

function uniqvalue{T<:AbstractString}( x::Union{ Array{T}, DataArray{T}, PooledDataArray{T} }; skipna::Bool=true, skipempty::Bool=true )
    lvls = DataArrays.levels(x)
    if skipna
        l = dropna( lvls )
        if skipempty
            emptyidx = findfirst( l, "" )
            if length( l ) == 1 && emptyidx == 0
                return l[1]
            elseif length( l ) == 2 && emptyidx != 0
                if emptyidx == 1
                    return l[2]
                else
                    return l[1]
                end
            end
        elseif length( l ) == 1
            return l[1]
        end
        return NA
    end
    if skipempty
        emptyidx = findfirst( lvls, "" )
        if length( lvls ) == 1 && emptyidx == 0
            return lvls[1] # could be NA
        elseif length( lvls ) == 2 && emptyidx != 0
            if emptyidx == 1
                return lvls[2]
            else
                return lvls[1]
            end
        end
    elseif length( lvls ) == 1
        return lvls[1]
    end
    return NA
end

immutable CalcPivot
    spec::Expr
    by::Array{Symbol,1}
    CalcPivot( x::UTF8String, by::Array{Symbol,1}=Symbol[] ) = CalcPivot( parse(x), by )
    CalcPivot( x::ASCIIString, by::Array{Symbol,1}=Symbol[] ) = CalcPivot( parse(utf8(x)), by )
    CalcPivot( x::UTF8String, by::Symbol ) = CalcPivot( parse(x), Symbol[ by ] )
    CalcPivot( x::ASCIIString, by::Symbol ) = CalcPivot( parse(utf8(x)), Symbol[ by ] )
    function CalcPivot( x::Expr, by::Symbol )
        CalcPivot( x, Symbol[ by ] )
    end
    function CalcPivot( ex::Expr, by::Array{Symbol,1}=Symbol[] )
        if !Base.Meta.isexpr( ex, :call )
            error( string( ex ) * " does not look like an aggregator function")
        end
        # disallow mutating function
        fname = ex.args[1]
        if Base.Meta.isexpr( ex, :curly )
            error( "CalcPivot: curly not supported")
        elseif !( typeof( fname ) <: Symbol ) && !Base.Meta.isexpr( fname, :(.) )
            error( "CalcPivot: only simple function name please")
        end
        if contains( string(fname), "!" )
            error( string(fname) * " seems to have side effects" )
        end

        if fname == :topnames # ensure we have the name in "by"
            # expect the first argument is a QuoteNode, or Expr( :quote, symbol )
            if typeof( ex.args[2] ) == QuoteNode
                name_col = ex.args[2].value
            elseif Base.Meta.isexpr( ex.args[2], :quote )
                name_col = ex.args[2].args[1]
            else
                throw( "topnames: 1st argument expects a symbol (name column)")
            end
            if !in( name_col, by )
                new( ex, Symbol[ by..., name_col ] )
            else
                new( ex, by )
            end
        else
            new( ex, by )
        end
    end
end

CalcPivotFuncCache = Dict{ Any, Function }()
CalcPivotAggrDepCache = Dict{ Any, Array{Symbol,1} }()

# lifting a CalcPivot is a more complicated beast compared to aggregation.
# It is basically the fancier "apply" component of the # split-apply-combine strategy,
# with the "apply" also doing its nested split-apply-combine steps if necessary.

function liftCalcPivotToFunc( ex::Expr, by::Array{Symbol,1} )
    if haskey( CalcPivotFuncCache, (ex, by ) )
        return CalcPivotFuncCache[ (ex, by ) ]
    end

    cex = deepcopy( ex )
    convertExpression!( cex )
    funnameouter = gensym("calcpvt")
    funname = gensym()

    membernames = Dict{Union{Symbol,Expr}, Symbol}()
    cex = DataFramesMeta.replace_syms(cex, membernames)
    # keys are the columns. values are the unique gensyms

    if !isempty( by ) # micro split-apply-combine
        funargs = map(x -> :( getindex( byf, $(Meta.quot(x))) ), collect(keys(membernames)))
        aggregates = setdiff( collect( keys( membernames ) ), by )
        CalcPivotAggrDepCache[ (ex, by ) ] = aggregates
        # basically we want to do
        # byf = by(_df_, by, lambdadf -> DataFrame(b = aggrfuncs[:b](lambdadf), c = aggrfuncs[:c](lambdadf), ... )
        aggr_args = Any[]
        lambdasym = gensym( "lambdadf" )
        for a in aggregates
            push!( aggr_args, Expr( :kw, a, Expr( :call, Expr( :ref, :aggrfuncs, QuoteNode( a ) ), lambdasym ) ) )
        end

        bycolsexpr=Expr( :vcat, map( _->QuoteNode(_), by )... ) # [ :a, :b, :c ... ]
        aggrcode = Expr( :->, lambdasym, Expr( :call, DataFrame, aggr_args... ) )
        code = :(
            function $funnameouter( _df_::AbstractDataFrame, c::Symbol; kwargs... ) # we need kwargs here for aggregate specs
                aggrfuncs = Dict{Symbol,Function}()
                for (aggrc,spec) in kwargs
                    # this would throw if spec is not compliant, as usual
                    aggrfuncs[aggrc] = TermWin.liftAggrSpecToFunc(aggrc, spec )
                end
                # if kwargs doesn't have the required columns, it'll throw in this line
                byf = DataFrames.by( _df_, $bycolsexpr, $aggrcode )
                function $funname($(collect(values(membernames))...))
                    $cex
                end
                vals = $funname($(funargs...))
                ret = byf[ $bycolsexpr ]
                ret[ c ] = vals
                ret # This has all the "by" columns and the result, named by c
                # combine is done by the caller
            end
        )
    else # much simplier, we are doing line-by-line apply.
        # Common use case: a simple bucketing, or a line by line ranking.
        # a "pure" reading on empty "by" array may mean we aggregate everything in the
        # dataframe into one row, and then attach this row's value (a constant) into every row.
        # However, I have not known any legitimate use case for this.
        CalcPivotAggrDepCache[ (ex, by )] = Symbol[]
        funargs = map(x -> :( getindex( _df_, $(x)) ), collect(keys(membernames)))
        code = :(
            function $funnameouter( _df_::AbstractDataFrame ) # we need kwargs here for aggregate specs
                function $funname($(collect(values(membernames))...))
                    $cex
                end
                $funname($(funargs...))
                # the creation of column is done by the caller
            end
        )
    end
    ret = eval( code )
    CalcPivotFuncCache[ (ex, by ) ] = ret
end

function cut_categories{S<:Real, T<:Real}( ::Type{S}, breaks::Vector{T};
    boundedness = :unbounded,
    leftequal=true, # t1 <= x < t2 or t1 < x <= t2?
    absolute=false, # t1 <= |x| < t2?
    rank=true, # add a rank to the string output for easier sorting?
    ranksep = ". ", # "1. t1 <= x < t2"?
    label = "", # if not compact, what label do we use for x?
    compact=(label==""), # <t1, [t1,t2), t2+. Further shortened for integer intervals with length=1
    reverse=false, # reverse the rank from the largest first?
    # the following format the boundary numbers
    # see Formatting.jl
    prefix="", suffix="", scale=1, precision=-1,
    commas=false,stripzeros=(precision==-1),parens=false,
    mixedfraction=false,autoscale=:none,conversion=""
    )
    n = length(breaks)
    breakstrs = UTF8String[]
    function formatter(_)
        prefix * format( _*scale,
            precision=precision,
            commas=commas,
            stripzeros=stripzeros,
            parens=parens,
            mixedfraction=mixedfraction,
            autoscale=autoscale,
            conversion=conversion ) * suffix
    end
    for b in breaks
        push!( breakstrs, formatter( b ) )
    end
    if boundedness == :unbounded
        ncategories = n + 1
    elseif boundedness == :bounded
        ncategories = n-1
    else
        ncategories = n
    end
    pool = Array(UTF8String, ncategories )
    if rank
        rankwidth = length(string(ncategories))
    end
    if !rank
        rankprefixfunc = _->""
    elseif reverse
        rankprefixfunc = j -> format( n+2-j, width=rankwidth ) * ranksep
    else
        rankprefixfunc = j -> format( j, width=rankwidth ) * ranksep
    end
    if compact
        if S <: Integer && T <: Integer && scale == 1
            # we use 1...5, 6, 7...10, 11+etc.
            if leftequal
                breakminus1strs = UTF8String[]
                for b in breaks
                    push!( breakminus1strs, formatter( b-1 ) )
                end
                poolindexshift = -1
                if boundedness in [ :unbounded, :boundedabove ]
                    pool[1] = rankprefixfunc(1) * "≤" * breakminus1strs[1]
                    poolindexshift = 0
                end
                for i in 2:n
                    if breaks[i-1] == breaks[i]-1
                        pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift)*breakstrs[i-1]
                    else
                        pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift)*breakstrs[i-1]*"…"*breakminus1strs[i]
                    end
                end
                if boundedness in [ :unbounded, :boundedbelow ]
                    pool[n+1+poolindexshift] = rankprefixfunc(n+1+poolindexshift)*breakstrs[n]*"+"
                end
            else
                breakplus1strs = UTF8String[]
                for b in breaks
                    push!( breakminus1strs, formatter( b+1 ) )
                end
                poolindexshift = -1
                if boundedness in [ :unbounded, :boundedabove ]
                    pool[1] = rankprefixfunc(1) *"≤ "* breakstrs[1]
                    poolindexshift = 0
                end
                for i in 2:n
                    if breaks[i-1]+1 == breaks[i]
                        pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift)*breakstrs[i]
                    else
                        pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift)*breakplus1strs[i-1]*"…"*breakstrs[i]
                    end
                end
                if boundedness in [ :unbounded, :boundedbelow ]
                    pool[n+1+poolindexshift] = rankprefixfunc(n+1+poolindexshift)*breakplus1strs[n]*"+"
                end
            end
        else # by the way, we don't show absolute in compact
            if leftequal
                brackL = "["
                brackR = ")"
                compareL = utf8( "<" )
                compareR = "≥"
            else
                brackL = "("
                brackR = "]"
                compareL = "≤"
                compareR = utf8(">")
            end
            poolindexshift = -1
            if boundedness in [ :unbounded, :boundedabove ]
                pool[1] = rankprefixfunc(1) * compareL * breakstrs[1]
                poolindexshift = 0
            end
            for i in 2:n
                if i == 2 && boundedness in [ :boundedbelow, :bounded ]
                    pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift) * "[" * breakstrs[i-1]* "," *breakstrs[i] * brackR
                elseif i == n && boundedness in [ :boundedabove, :bounded ]
                    pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift) * brackL * breakstrs[i-1]* "," *breakstrs[i] * "]"
                else
                    pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift) * brackL * breakstrs[i-1]* "," *breakstrs[i] * brackR
                end
            end
            if boundedness in [ :unbounded, :boundedbelow ]
                pool[n+1+poolindexshift] = rankprefixfunc(n+1+poolindexshift) * compareR * breakstrs[n]
            end
        end
    else
        if absolute
            label2 = "|"*label*"|"
        else
            label2 = label
        end
        if leftequal
            compareL = " ≤ "
            compareR = utf8(" < ")
        else
            compareL = utf8(" < ")
            compareR = " ≤ "
        end
        poolindexshift = -1
        if boundedness in [ :unbounded, :boundedabove ]
            pool[1] = rankprefixfunc( 1 ) * label2 * compareR * breakstrs[1]
            poolindexshift = 0
        end
        for i in 2:n
            if i == 2 && boundedness in [ :boundedbelow, :bounded ]
                pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift) * breakstrs[i-1] * " ≤ " * label2 * compareR * breakstrs[i]
            elseif i == n && boundedness in [ :boundedabove, :bounded ]
                pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift) * breakstrs[i-1] * compareL * label2 * " ≤ " * breakstrs[i]
            else
                pool[i+poolindexshift] = rankprefixfunc(i+poolindexshift) * breakstrs[i-1] * compareL * label2 * compareR * breakstrs[i]
            end
        end
        if boundedness in [ :unbounded, :boundedbelow ]
            pool[n+1+poolindexshift] = rankprefixfunc(n+1+poolindexshift) * breakstrs[n] * compareL * label2
        end
    end
    return pool
end

# boundedness:
#    unbounded    gives n+1 categories for n breaks.
#    boundedbelow gives n   categories for n breaks. Values below min will be NA
#    boundedabove gives n   categories for n breaks. Values above max will be NA
#    bounded      gives n-1 categories for n breaks. Values below min or above max will be NA
function discretize{S<:Real, T<:Real}(x::AbstractArray{S,1}, breaks::Vector{T};
    boundedness = :unbounded,
    bucketstrs = UTF8String[], # if provided, all of below will be ignored. length must be length(breaks)+1
    leftequal=true, # t1 <= x < t2 or t1 < x <= t2?
    absolute=false, # t1 <= |x| < t2?
    rank=true, # add a rank to the string output for easier sorting?
    ranksep = ". ", # "1. t1 <= x < t2"?
    label = "", # if not compact, what label do we use for x?
    compact=(label==""), # <t1, [t1,t2), t2+. Further shortened for integer intervals with length=1
    reverse=false, # reverse the rank from the largest first?
    # the following format the boundary numbers
    # see Formatting.jl
    prefix="", suffix="", scale=1, precision=-1,
    commas=false,stripzeros=(precision==-1),parens=false,
    mixedfraction=false,autoscale=:none,conversion=""
    )
    if !issorted(breaks)
        sort!(breaks)
    end
    refs = fill(zero(DataArrays.DEFAULT_POOLED_REF_TYPE), length(x))
    n = length(breaks)
    if absolute
        x2 = abs(x)
    else
        x2 = x
    end

    if boundedness == :unbounded
        below_min_mult = 1
        above_max_mult = 1
        ref_shift = 1
        ncategories = length( breaks ) + 1
    elseif boundedness == :boundedbelow
        below_min_mult = 0
        above_max_mult = 1
        ref_shift = 0
        ncategories = length( breaks )
    elseif boundedness == :boundedabove
        below_min_mult = 1
        above_max_mult = 0
        ref_shift = 1
        ncategories = length( breaks )
    elseif boundedness == :bounded
        below_min_mult = 0
        above_max_mult = 0
        ref_shift = 0
        ncategories = length( breaks ) - 1
    end

    if ncategories < 1
        error( "Too few categories. Change boundedness or add breaks")
    end

    if leftequal
        for i in 1:length(x)
            if isna( x, i )
                refs[i] = 0
            elseif x2[i] <  breaks[1]
                refs[i] = below_min_mult
            elseif x2[i] > breaks[end]
                refs[i] = (n+ref_shift) * above_max_mult
            elseif x2[i] == breaks[end]
                if boundedness in [ :bounded, :boundedabove ]
                    refs[i] = ncategories
                else
                    refs[i] = n+ref_shift
                end
            else
                refs[i] = searchsortedlast(breaks, x2[i]) + ref_shift
            end
        end
    else
        for i in 1:length(x)
            if isna( x, i )
                refs[i] = 0
            elseif x2[i] < breaks[1]
                refs[i] = below_min_mult
            elseif x2[i] > breaks[end]
                refs[i] = (n+ref_shift) * above_max_mult
            else
                refs[i] = searchsortedfirst(breaks, x2[i])
            end
        end
    end

    if length( bucketstrs ) != 0
        if length( bucketstrs ) != ncategories
            error( "bucketstrs expected to have size " * string( ncategories ) *
                 ". Got " * string( length( bucketstrs ) ) )
        end
        if maximum( refs ) > ncategories
            maxref = maximum( refs )
            s = "ncategories < max refs \n maxref="  * string( maxref ) * "\n ncategories=" * string( ncategories )
            s *= "\n buckets = " * string( bucketstrs )
            idx = findfirst( refs, maxref )
            s *= "\n Example x = " * string(x2[idx] )
            s *= "\n breaks" * string( breaks )
            error( s )
        end
        return DataArrays.PooledDataArray(DataArrays.RefArray(refs), bucketstrs )
    end
    pool = cut_categories( S, breaks,
        boundedness = boundedness,
        leftequal   = leftequal,
        absolute    = absolute,
        rank        = rank,
        ranksep     = ranksep,
        label       = label,
        compact     = compact,
        reverse     = reverse,
        prefix=prefix, suffix=suffix, scale=scale, precision=precision,
        commas=commas,stripzeros=stripzeros,parens=parens,
        mixedfraction=mixedfraction,autoscale=autoscale,conversion=conversion
        )

    DataArrays.PooledDataArray(DataArrays.RefArray(refs), pool)
end

# quantile-based auto-breaks
# weighted quantile is not implemented
# use scale=100.0, suffix="%", to express the quantiles in percentages
function discretize{S<:Real}(x::AbstractArray{S,1}; quantiles = Float64[], ngroups::Int = 4, kwargs ... )
    if length( quantiles ) != 0
        if any( _ -> _ < 0.0 || _ > 1.0 , quantiles )
            error( "illegal quantile numbers outside [0,1]")
        end
        if !issorted(quantiles)
            sort!(quantiles)
        end
        if quantiles[1] != 0.0
            insert!( quantiles, 1, 0.0 )
        end

        if quantiles[end] != 1.0
            push!( quantiles, 1.0 )
        end
        bucketstrs = cut_categories( Float64, quantiles; boundedness = :bounded, kwargs... )
        discretize( x, quantile( x, quantiles ); bucketstrs = bucketstrs, boundedness = :bounded, kwargs... )
    else
        qs = [0:ngroups]/ngroups
        bucketstrs = cut_categories( Float64, qs; boundedness = :bounded, kwargs... )
        discretize( x, quantile( x, qs); bucketstrs = bucketstrs, boundedness = :bounded, kwargs... )
    end
end

# names are expected to be unique
# n is the maximum rank number to report. Actual outcome may depend on existence of a tie, and dense option
function topnames{S<:AbstractString,T<:Real}( name::AbstractArray{S,1}, measure::AbstractArray{T,1}, n::Int;
    absolute=false,
    ranksep = ". ",
    dense = true, # if there is a tie in the 2nd place, do we do "1,2,2,4", or "1,2,2,3"
    tol = 0,  # if absolute, what is the smallest contribution that we would consider
    others = "Others",
    parens = false # put parentheses around names with negative measure?
    )

    if absolute
        df = DataFrame( name = name, measure = measure, absmeasure = abs(measure) )
        if tol > 0 # filter out too small names
            dfsorted = sort( df[ df[ :absmeasure ] >= tol ], cols = [:absmeasure, :measure ], rev=[true,true] )
        else
            dfsorted = sort( df, cols = [ :absmeasure, :measure ], rev = [ true, true ] )
        end
    else
        df = DataFrame( name = name, measure = measure )
        dfsorted = sort( df, cols = [ :measure ], rev = [ true ] )
    end

    rankcount = 1
    rankwidth = length( string( n ) )
    nr = nrow( dfsorted )

    if !absolute
        pool = UTF8String[]
        refs = fill(zero(DataArrays.DEFAULT_POOLED_REF_TYPE), nr )
        lastval  = zero( T )
        lastrank = 0
        for r in 1:nr
            if isna( dfsorted[:measure], r )
                continue
            else
                val = dfsorted[ r, :measure ]
                if lastrank != 0 && lastval == val # tie
                    push!( pool, format( lastrank, width=rankwidth ) * ranksep * dfsorted[ r, :name ] )
                    refs[r] = length( pool )
                    if !dense
                        rankcount += 1
                    end
                elseif rankcount > n
                    break
                else
                    push!( pool, format( rankcount, width=rankwidth ) * ranksep * dfsorted[ r, :name ] )
                    lastrank = rankcount
                    lastval = val
                    refs[r] = length( pool )
                    rankcount += 1
                end
            end
        end
        dfsorted[ :rankstr ] = DataArrays.PooledDataArray(DataArrays.RefArray(refs), pool)
        rdf = dfsorted[ [:name, :rankstr ] ]
        jdf = join( df, rdf, on = :name, kind = :left )
    else
        rankedflag = fill( zero( Bool ), nr )
        lastval  = zero( T )
        lastrank = 0
        for r in 1:nr
            if isna( dfsorted[:measure], r )
                continue
            else
                val = dfsorted[ r, :measure ]
                if lastrank != 0 && lastval == val # tie
                    rankedflag[r] = true
                    if !dense
                        rankcount += 1
                    end
                elseif rankcount > n
                    break
                else
                    rankedflag[r] = true
                    lastrank = rankcount
                    lastval = val
                    rankcount += 1
                end
            end
        end
        dfsorted[ :rankedflag ] = rankedflag
        dfsorted2 = sort( dfsorted, cols = [ :measure ], rev = [ true ] )
        rankstr = DataArray(UTF8String,nr)
        rankcount = 1
        lastval  = zero( T )
        lastrank = 0
        for r in 1:nr
            if isna( dfsorted2[ :measure ], r )
                continue
            elseif dfsorted2[ r, :rankedflag ]
                val = dfsorted2[ r, :measure ]
                if lastrank != 0 && lastval == val # tie
                    if parens && val < 0
                        rankstr[r] = format( lastrank, width=rankwidth ) * ranksep * "("*dfsorted2[r,:name]*")"
                    else
                        rankstr[r] = format( lastrank, width=rankwidth ) * ranksep * dfsorted2[r,:name]
                    end
                    if !dense
                        rankcount += 1
                    end
                elseif rankcount > n
                    break
                else
                    if parens && val < 0
                        rankstr[r] = format( rankcount, width=rankwidth ) * ranksep * "("*dfsorted2[r,:name]*")"
                    else
                        rankstr[r] = format( rankcount, width=rankwidth ) * ranksep * dfsorted2[r,:name]
                    end
                    lastrank = rankcount
                    lastval = val
                    rankcount += 1
                end
            end
        end
        dfsorted2[ :rankstr ] = rankstr
        jdf = join( df, dfsorted2[ [:name, :rankstr] ], on = :name, kind=:left )
    end

    # replace NA with "others"
    ret = DataArrays.PooledDataArray( jdf[ :rankstr ] )
    push!( ret.pool, others )
    poollen = length( ret.pool )
    for i = 1:length( ret.refs )
        if ret.refs[i] == 0
            ret.refs[i] = poollen
        end
    end
    ret
end

import DataFrames.describe
export describe

function describe{T}( io, dv::Array{T,1} )
    describe( io, DataArray( dv ) )
end


#Provide utility functions for interoperability with Compose.jl
#Define conversions from Shapefile.jl shape primitives to their Compose.jl counterparts

#Convert Shapefile rectangle to Compose rectangle
import Compose: rectangle
rectangle{T<:Real}(R::Shapefile.Rect{T}) = rectangle(R.left,R.top,R.right-R.left,R.bottom-R.top)

#Compose polygons cannot be disjoint but Shapefile.Polygons can
#Need to convert Shapefile.Polygon to list of Compose polygons
function convert(::Type{Vector{Compose.Form{Compose.PolygonPrimitive}}},
        shape::Shapefile.Polygon)
    points = {}
    polygons={}
    currentpart=2
    for (i,p) in enumerate(shape.points)
        push!(points, p)
        if i==length(shape.points) || (currentpart≤length(shape.parts) && i==shape.parts[currentpart])
            push!(polygons, polygon([(p.x,p.y) for p in points]))
            currentpart += 1
            points = {}
        end
    end
    polygons
end

polygons(shape::Shapefile.Polygon) = convert(Vector{Compose.Form{Compose.PolygonPrimitive}},
    shape)


abstract type BinaryModel end
abstract type BitModel <: BinaryModel end
abstract type SpinModel <: BinaryModel end

abstract type AbstractVariable end
abstract type AbstractBitVariable <: AbstractVariable end
abstract type AbstractSpinVariable <: AbstractVariable end


abstract type AbstractBufferPair{T} end

"""
    struct BufferPair{T}  <: AbstractBufferPair{T}
A helper struct used to store pair of `BufferWord` used for rewriting.
`BufferPair`s are used in conjunction with `kbWork` struct.
"""
struct BufferPair{T} <: AbstractBufferPair{T}
    _vWord::BufferWord{T}
    _wWord::BufferWord{T}
end

BufferPair{T}() where {T} = BufferPair(one(BufferWord{T}), one(BufferWord{T}))

"""
    function rewrite_from_left!(bp::BufferPair, u::AbstractWord, rewriting)
Rewrites a word from left using buffer words from `BufferPair` and `rewriting` object.

Note: this implementation returns an instance of `BufferWord`!
"""
function rewrite_from_left!(bp::BufferPair, u::AbstractWord, rewriting)
    if isempty(rewriting)
        store!(bp._vWord, u)
        return bp._vWord
    end
    empty!(bp._vWord)
    store!(bp._wWord, u)
    v = rewrite_from_left!(bp._vWord, bp._wWord, rewriting)
    return v
end

mutable struct kbWork{T}
    lhsPair::BufferPair{T}
    rhsPair::BufferPair{T}
    tmpPair::BufferPair{T}
end

kbWork{T}() where {T} = kbWork(BufferPair{T}(), BufferPair{T}(), BufferPair{T}())


using DiffEqBase, DiffEqJump
using Test
const DJ = DiffEqJump

# test data
minpriority = 2.0^exponent(1e-12)
maxpriority = 2.0^exponent(1e12)
priorities = [1e-13, .99*minpriority, minpriority,1.01e-4, 1e-4, 5., 0., 1e10]

mingid = exponent(minpriority)   # = -40
ptog = priority -> DJ.priortogid(priority, mingid)
pt = DJ.PriorityTable(ptog, priorities, minpriority, maxpriority)

#display(priorities)
#display(pt)

# test insert
grpcnt = DJ.numgroups(pt)
push!(priorities, maxpriority*.99)
DJ.insert!(pt, length(priorities), priorities[end])
@test grpcnt == DJ.numgroups(pt)
@test pt.groups[end].pids[1] == length(priorities)

push!(priorities, maxpriority*.99999)
DJ.insert!(pt, length(priorities), priorities[end])
@test grpcnt == DJ.numgroups(pt)
@test pt.groups[end].pids[2] == length(priorities)

numsmall = length(pt.groups[2].pids)
push!(priorities, minpriority*.6)
DJ.insert!(pt, length(priorities), priorities[end])
@test grpcnt == DJ.numgroups(pt)
@test pt.groups[2].pids[end] == length(priorities)

push!(priorities, maxpriority)
DJ.insert!(pt, length(priorities), priorities[end])
@test grpcnt == DJ.numgroups(pt)-1
@test pt.groups[end].pids[1] == length(priorities)

# test updating
DJ.update!(pt, 5, priorities[5], 2*priorities[5])   # group 29
priorities[5] *= 2
@test pt.groups[29].numpids == 1
@test pt.groups[30].numpids == 1

DJ.update!(pt, 9, priorities[9], maxpriority*1.01)
priorities[9] = maxpriority*1.01
@test pt.groups[end].numpids == 2
@test pt.groups[end-1].numpids == 1

DJ.update!(pt, 10, priorities[10], 0.)
priorities[10] = 0.
@test pt.groups[1].numpids == 2

# test sampling
cnt = 0
Nsamps = Int(1e7)
for i = 1:Nsamps
    global cnt; pid = DJ.sample(pt, priorities)
    (pid == 8) && (cnt += 1)
end
@test abs(cnt // Nsamps -  0.008968535978248484) / 0.008968535978248484 < .05