"""
```
COTHANK{T} <: AbstractRepModel{T}
```

The `THANK` type defines the structure of THANK, an implementation
of the DSGE model in Bilbiie, Primiceri, Tambalotti (2019)

### 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.

* `endogenous_states_augmented::OrderedDict{Symbol,Int}`: Maps lagged states to
  their columns in the measurement and equilibrium condition equations. These
  are added after `gensys` solves the model.

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

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

#### Model Specifications and Settings

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

* `subspec::String`: The model subspecification number, indicating that
  some parameters from the original model spec (\"ss10\") 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.

* `pseudo_observable_mappings::OrderedDict{Symbol,PseudoObservable}`: A
  dictionary that stores names and transformations to/from model units. See
  `PseudoObservable` for further details.
"""
mutable struct COTHANK{T} <: AbstractRepModel{T}
    parameters::ParameterVector{T}                         # vector of all time-invariant model parameters
    steady_state::ParameterVector{T}                       # model steady-state values
    keys::OrderedDict{Symbol,Int}                          # human-readable names for all the model
                                                           # parameters and steady-states

    endogenous_states::OrderedDict{Symbol,Int}             # these fields used to create matrices in the
    exogenous_shocks::OrderedDict{Symbol,Int}              # measurement and equilibrium condition equations.
    expected_shocks::OrderedDict{Symbol,Int}               #
    equilibrium_conditions::OrderedDict{Symbol,Int}        #
    endogenous_states_augmented::OrderedDict{Symbol,Int}   #
    observables::OrderedDict{Symbol,Int}                   #
    pseudo_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}
    pseudo_observable_mappings::OrderedDict{Symbol, PseudoObservable}
end

description(m::COTHANK) = "Bilbiie, Primiceri, Tambalotti (2019), $(m.subspec)."

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

Arguments:
`m:: COTHANK`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::COTHANK)
    # Endogenous states
    endogenous_states = [:y_t, :y1_t, :y2_t, :k_t, :k_S_t, :L1_t,  :L2_t, :klR1_t, :klR2_t,   # check if missing ι terms
                         :w1_t, :w2_t, :π_t, :π1_t, :π2_t, :mc1_t, :mc2_t, :λ_p_t, :λ_w_t,      # sticky prices
                         :λ_S_t, :λ_H1_t, :λ_H2_t, :λ_SH1_t, :λ_SH2_t, :c_t,  :c_H1_t, :c_H2_t, :c_S_t, :i_S_t,
                         :R_t, :u_t, :ϕ_t, :z_t, :ρ_t, :a_t, :d_t, :x_t, :b_t, :t_t, #:φ_t,
                         :g_t, :g_w1_t, :g_w2_t, :μ_t, :bᴿ_t, :η_mp_t, :t_S_t, :t_H1_t, :t_H2_t,
                         :τ_t, :τ_S_t, :τ_H1_t, :τ_H2_t, :π_S_t, :price1_t, :price2_t, :Eπ1_t, :Eπ2_t, :Eλ_S_t,
                         :Eλ_H1_t, :Eλ_H2_t, :Ex_t, :Eϕ_t, :Eρ_t, :Ei_S_t, :Ew1_t, :Eπ_t, :Ew2_t]
                                                                                # CHECK BACK IF ANY CAN BE AUGMENTED VARIABLES
#=                         :y_f_t, :k_f_t, :L_f_t, :Rk_f_t, :w_f_t, :mc_f_t,       # flexible prices
                         :λ_f_t, :c_f_t, :R_f_t, :u_f_t, :ϕ_f_t,
                         :i_f_t, :kbar_f_t, :wgap_f_t, :gdp_f_t, :Ec_f_t, :Eλ_f_t,
                         :Eϕ_f_t, :ERk_f_t, :Ei_f_t, :λ_h_f_t, :λ_s_f_t, :c_h_f_t, :c_s_f_t,
                         :t_h_f_t, :Ey_f_t, :Eλ_h_f_t, :Eλ_s_f_t]
=#
    # Exogenous shocks
    exogenous_shocks = [:a_sh, :d_sh, :λ_p_sh, :b_sh, :z_sh, :τ_H1_sh, :τ_H2_sh, :μ_sh, :λ_w_sh,
                        :τ_S_sh, :g_sh, :η_mp_sh]

    # Expectations shocks
    expected_shocks = [:Eπ1_sh, :Eπ2_sh, :Eλ_S_sh, :Eλ_H1_sh, :Eλ_H2_sh,
                       :Ex_sh, :Eϕ_sh, :Eρ_sh, :Ei_S_sh, :Ew1_sh, :Eπ_sh, :Ew2_sh] # sticky prices

    # Equilibrium conditions
    equilibrium_conditions = [:eq_KL_ratio_good1, :eq_mc_good1, :eq_KL_ratio_good2, :eq_mc_good2, :eq_production_good1, :eq_production_good2,
                              :eq_production_finalgood, :eq_price_index, :eq_demand_good1, :eq_π1, :eq_π2, :eq_price_phillips_good1,
                              :eq_price_phillips_good2, :eq_marg_utility_S, :eq_marg_utility_H1, :eq_marg_utility_H2, :eq_average_c,
                              :eq_euler, :eq_c_H1, :eq_c_H2, :eq_cap_util, :eq_ϕ, :eq_i, :eq_cap_input, :eq_cap_accum, :eq_wage_phillips_good1,
                              :eq_marg_sub_good1, :eq_avg_marg_util_SH1, :eq_wage_phillips_good2, :eq_marg_sub_good2, :eq_avg_marg_util_SH2,
                              :eq_mp, :eq_revenue, :eq_transfer, :eq_realdebt, :eq_fiscal_rule, :eq_redistribute_H1, :eq_redistribute_H2,
                              :eq_real_profits,
                              :eq_cap_market_clear, :eq_goods_market_clear, :eq_gdp,
                              :eq_a, :eq_d, :eq_λ_p, :eq_b, :eq_z, :eq_τ_H1, :eq_τ_H2, :eq_μ, :eq_λ_w, :eq_τ_S, :eq_g, :eq_η_mp, #eq_φ
                               :eq_Eπ1, :eq_Eπ2, :eq_Eλ_S, :eq_Eλ_H1, :eq_Eλ_H2, :eq_Ex, :eq_Eϕ,
                              :eq_Eρ, :eq_Ei_S, :eq_Ew1, :eq_Eπ, :eq_Ew2]
#=
                              :eq_y_f, :eq_k_f, :eq_L_f, :eq_Rk_f, :eq_w_f, :eq_mc_f,       # flexible prices
                              :eq_λ_f, :eq_c_f, :eq_R_f, :eq_u_f, :eq_ϕ_f,
                              :eq_i_f, :eq_kbar_f, :eq_wgap_f, :eq_gdp_f, :eq_Ec_f, :eq_Eλ_f,
                              :eq_Eϕ_f, :eq_ERk_f, :eq_Ei_f, :eq_λ_h_f, :eq_λ_s_f, :eq_c_h_f,
                              :eq_c_s_f, :eq_t_h_f, :eq_Ey_f, :eq_Eλ_h_f, :eq_Eλ_s_f]
=#
    # Additional states added after solving model
    # Lagged states and observables measurement error
    endogenous_states_augmented = []

    # Observables
    observables = keys(m.observable_mappings)

    # Pseudo-observables
    pseudo_observables = keys(m.pseudo_observable_mappings)

    for (i,k) in enumerate(endogenous_states);           m.endogenous_states[k]           = i end
    for (i,k) in enumerate(exogenous_shocks);            m.exogenous_shocks[k]            = i end
    for (i,k) in enumerate(expected_shocks);             m.expected_shocks[k]             = i end
    for (i,k) in enumerate(equilibrium_conditions);      m.equilibrium_conditions[k]      = i end
    for (i,k) in enumerate(endogenous_states);           m.endogenous_states[k]           = i end
    for (i,k) in enumerate(endogenous_states_augmented); m.endogenous_states_augmented[k] = i+length(endogenous_states) end
    for (i,k) in enumerate(observables);                 m.observables[k]                 = i end
    for (i,k) in enumerate(pseudo_observables);          m.pseudo_observables[k]          = i end
end

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

    # Model-specific specifications
    spec               = split(basename(@__FILE__),'.')[1]
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # Initialize empty model
    m = COTHANK{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(), OrderedDict{Symbol,Int}(),

            # model indices
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),

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

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

    # Set observable and pseudo-observable transformations
    init_observable_mappings!(m)
    #init_pseudo_observable_mappings!(m)

    # Initialize parameters
    init_parameters!(m)

    init_model_indices!(m)
    init_subspec!(m)
    steadystate!(m)

    return m
end

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

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::COTHANK)
    # Calibrated parameters#
    m <= parameter(:δ, 0.025, fixed=true,
                   description="δ: The capital depreciation rate.",
                   tex_label = "\\delta" ) # from Giorgio's matlab script

    # Non standard devation parameters
    m <= parameter(:α, 1-0.6/(1.15), (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Normal(0.30, 0.05), fixed = false,
                   description="α: Capital elasticity in the intermediate goods sector's production function (also known as the capital share).",
                   tex_label = "\\alpha") # from Giorgio's matlab script

    m <= parameter(:ι_p, 0.21, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.5, 0.15), fixed = false,
                   description = "ι_p: The weight attributed to last period's inflation in price indexation. " *
                   "(1-ι_p) is the weight attributed to steady-state inflation.",
                   tex_label = "\\iota_p") # from THANK

    m <= parameter(:ι_w, 0.15, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.5, 0.15), fixed = false,
                   description="ι_w: The weight attributed to last period's wage in wage indexation. "
                    * "(1-ι_w) is the weight attributed to steady-state wages.",
                   tex_label = "\\iota_w") # from THANK

    m <= parameter(:h, 0.5, (-5.0, 5.0), (-5., 5.), ModelConstructors.Untransformed(), Beta(0.5, 0.1), fixed = false,
                   description="h: habit formation parameter.",
                   tex_label = "h") # from Giorgio's matlab script

    m <= parameter(:λ_p_ss, 0.15,(-1000.,1000.),(-1000., 1000.), ModelConstructors.Untransformed(), Normal(0.15, 0.05),  fixed = false,
                   description="λ_p_ss: The steady state net price markup.",
                   tex_label = "\\lambda_p") # from Giorgio's matlab script

    m <= parameter(:λ_w_ss, 0.12,(-1000.,1000.),(-1000., 1000.), ModelConstructors.Untransformed(), Normal(0.15, 0.05), fixed = false,
                   description = "λ_w_ss: The steady state net wage markup, which affects the elasticity of substitution between differentiated labor services.",
                   tex_label = "\\lambda_w") # from Giorgio's matlab script

    m <= parameter(:L_ss, 0.5, (-1000., 1000.), (-1e3, 1e3), ModelConstructors.Untransformed(), Normal(0.0,0.5), fixed=false,
                   description="L_ss: The steady state for log hours.",
                   tex_label = "\\log(L)^{ss}") # from Giogio's matlab script

    m <= parameter(:Fβ, 0.1, (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(), Gamma(0.25, 0.1), fixed=false,
                   description = "Fβ: Discount rate transformed.",

                   tex_label = "100(\\beta^{-1} - 1)") # from THANK

    m <= parameter(:ν, 4., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(), Gamma(2, 0.75), fixed = false,
                   description="ν_l: The inverse Frisch elasticity.",
                   tex_label = "\\nu") # from THANK

    m <= parameter(:ξ_p, 0.85, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.66, 0.1), fixed=false,
                   description="ξ_p: The Calvo parameter. In every period, intermediate goods producers optimize prices with probability (1-ξ_p). "
                   * "With probability ξ_p, prices are adjusted according to a weighted average of the previous period's inflation "
                   * "(π_t1) and steady-state inflation (π_ss).",
                   tex_label = "\\xi_p") # from THANK

    m <= parameter(:ξ_w, 0.75, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.66, 0.1), fixed=false,
                   description="ξ_w: (1-ξ_w) is the probability with which households can freely choose wages in each period. "
                   * "With probability ξ_w, wages increase at a geometrically weighted average of the steady state rate of "
                   * "wage increases and last period's productivity times last period's inflation.",
                   tex_label = "\\xi_w") # from THANK

    m <= parameter(:χ, 5., (0., 10.), (1e-5, 0.), ModelConstructors.Exponential(), Gamma(5., 1.), fixed = false,
                   description="χ: The elasticity of the capital utilization cost function.",
                   tex_label = "\\chi") # from THANK

    m <= parameter(:S′′, 2.5, (-15., 15.), (-15., 15.), ModelConstructors.Untransformed(), Gamma(4., 1), fixed = false,
                   description="S′′: The investment adjust cost.",
                   tex_label = "S^{\\prime\\prime}") # from THANK

    m <= parameter(:ϕ_1, 2., (1e-5, 10.), (1e-5, 10.00), ModelConstructors.Exponential(), Normal(1.7, 0.3), fixed = false,
                   description="ψ₁: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1") # from THANK

    m <= parameter(:ϕ_2, 0.1, (-0.5, 0.5), (-0.5, 0.5), ModelConstructors.Untransformed(), Normal(0.125, 0.05), fixed = false,
                   description="ψ₂: Weight on output gap in monetary policy rule.",
                   tex_label = "\\psi_2") # from THANK

    m <= parameter(:ϕ_3, 0.25, (-0.5, 0.5), (-0.5, 0.5), ModelConstructors.Untransformed(), Normal(0.125, 0.05), fixed = false,

                   description="ψ₃: Weight on ouptut growth in the monetary policy rule.",
                   tex_label = "\\psi_3") # from THANK

    # the following new parameters (from θ to v) are inputs into the steady state calculations given by Giorgio's matlab script
    m <= parameter(:θ, 0.25, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:τ_x, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:t_x, 0.2, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script

    m <= parameter(:f_H1, 1/4, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) #from Giorgio's MATLAB script
    m <= parameter(:f_S1, 3/4, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:H1, 1.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:H2, 1.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:A2, 0.7, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:s, 0.995, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:g_x, 0.18, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # from Giorgio's MATLAB script
    m <= parameter(:v, 0.7, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true)
    m <= parameter(:vars, 1.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true)
    m <= parameter(:ψ_H1, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true)
    m <= parameter(:ψ_H2, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true)

    # the following parameters (from ζ to sx) are new but had no given value, so current value is random placeholder
    m <= parameter(:ζ, 1.5, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # ???
         # the above parameter appears as \varzeta in the final good producers equilibrium conditions
         # but no \varzeta unicode characters exists in Julia so we use \zeta instead
    m <= parameter(:ζ_X, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # ???
    m <= parameter(:ζ_G, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # ???
    m <= parameter(:ζ_B, 1.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # ???
    m <= parameter(:s_x, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Untransformed(), RootInverseGamma(0.5, 1), fixed = true) # ???

    # exogenous processes - autocorrelation
    m <= parameter(:ρ_R, 0.8, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_R: persistence in the monetary policy rule.",
                   tex_label = "\\rho_{R}") # from THANK

    m <= parameter(:ρ_η_mp, 0.8, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_R: persistence in the monetary policy rule.",
                   tex_label = "\\rho_{R}") # from THANK

    m <= parameter(:ρ_z, 0.25, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_{z}") # from THANK

    m <= parameter(:ρ_a, 0.15, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.4, 0.2), fixed = false,
                   description="ρ_a: AR(1) coefficient", # fix description
                   tex_label = "\\rho_{a}") # ???

    m <= parameter(:ρ_d, 0.15, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.4, 0.2), fixed = false,
                   description="ρ_d: AR(1) coefficient", # fix description
                   tex_label = "\\rho_{d}") # ???

    m <= parameter(:ρ_φ, 0.15, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.4, 0.2), fixed = false,
                   description="ρ_φ: AR(1) coefficient.", # fix description
                   tex_label = "\\rho_{\\varphi}") # ???

    m <= parameter(:ρ_μ, 0.7, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}") # from THANK

    m <= parameter(:ρ_b, 0.7, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_b") # from THANK

    m <= parameter(:ρ_λ_p, 0.95, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_λ_p: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_p}") # from THANK

    m <= parameter(:ρ_λ_w, 0.98, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}") # from THANK

    m <= parameter(:ρ_g, 0.98, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g") # from THANK

    m <= parameter(:ρ_τ_S, 0.98, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_τ_S: AR(1) coefficient.", # fix description
                   tex_label = "\\rho_\\tau_S") # ???

    m <= parameter(:ρ_τ_H1, 0.98, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_τ_H1: AR(1) coefficient.", # fix description
                   tex_label = "\\rho_\\tau_H1") # ???

    m <= parameter(:ρ_τ_H2, 0.98, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_τ_H2: AR(1) coefficient.", # fix description
                   tex_label = "\\rho_g") # ???

    m <= parameter(:ρ_T, 0.8, (0., 1.), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.6, 0.2), fixed = false,
                   description="ρ_T: parameter governing fiscal rule.", # Update this description later
                   tex_label = "\\rho_{R}") # ???

    m <= parameter(:η_λ_p, 0.0, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.50, 0.20), fixed = false,
                   description="η_λ_p: Moving average component in the price markup shock.",
                   tex_label = "\\eta_{\\lambda_p}") # from THANK

    m <= parameter(:η_λ_w, 0.0, (1e-5, 0.999), (1e-5, 0.999), ModelConstructors.SquareRoot(), Beta(0.50, 0.20), fixed = false,
                   description="η_λ_w: Moving average component in the wage markup shock.",
                   tex_label = "\\eta_{\\lambda_w}") # from THANK

    # exogenous processes - standard deviation
    m <= parameter(:σ_z, 0.9, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_z: The standard deviation of the technology process.",
                   tex_label = "\\sigma_{z}") # from THANK

    m <= parameter(:σ_a, 0.9, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_a: The standard deviation of the COVID-19 shock",
                   tex_label = "\\sigma_{a}") # ???

    m <= parameter(:σ_d, 0.9, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_d: ", # fix description
                   tex_label = "\\sigma_{d}") # ???

    m <= parameter(:σ_φ, 4., (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_φ: The standard deviation of the shock to disutility from labor supply", # fix description
                   tex_label = "\\sigma_{\\mu}") # ???, set to 4, following some values used in literature (e.g. Del Negro, Smets, and Schorfheide (2007) (DSGE-VAR paper))

    m <= parameter(:σ_μ, 5., (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}") # from THANK

    m <= parameter(:σ_b, 0.05, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.1, 1), fixed = false,
                   description="σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}") # from THANK

    m <= parameter(:σ_λ_p, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.1, 1), fixed = false,
                   description="σ_λ_p: The standard deviation of the price markup shock",
                   tex_label = "\\sigma_{\\lambda_p}") # from THANK

    m <= parameter(:σ_λ_w, 0.2, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.1, 1), fixed = false,
                   description="σ_λ_w: The standard deviation of the wage markup shock",
                   tex_label = "\\sigma_{\\lambda_w}") # from THANK

    m <= parameter(:σ_η_mp, 0.238, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.10), fixed = false,
                   description="σ_η_mp: The standard deviation of the monetary policy shock.", # fix description
                   tex_label = "\\sigma_{\\eta_mp}") # ???, set to a reasonable calibration from the FRBNY DSGE model

    m <= parameter(:σ_g, 0.35, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}") # from THANK

    m <= parameter(:σ_τ_S, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_τ_S.", # fix description
                   tex_label = "\\sigma_{\\tau_S}") # ???, zeroed out for now

    m <= parameter(:σ_τ_H1, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_τ_H1:.", # fix description
                   tex_label = "\\sigma_{\\tau_{H1}}") # ???, zeroed out for now

    m <= parameter(:σ_τ_H2, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(0.5, 1), fixed = false,
                   description="σ_τ_H2: .", # fix description
                   tex_label = "\\sigma_{\\tau_{H2}}") # ???, zeroed out for now

    # steady states
    m <= SteadyStateParameter(:γ, NaN, tex_label = "\\gamma")
    m <= SteadyStateParameter(:β, NaN, tex_label = "\\beta")
    m <= SteadyStateParameter(:r_ss, NaN, tex_label = "r_ss")
    m <= SteadyStateParameter(:π_ss, NaN, tex_label = "\\pi_ss")
    m <= SteadyStateParameter(:π_S_ss, NaN, tex_label = "\\pi_ss")
    m <= SteadyStateParameter(:expγ, NaN, tex_label = "exp^\\gamma")
#=
    m <= SteadyStateParameter(:ψ_H1, NaN, tex_label = "\\psi_H1")
    m <= SteadyStateParameter(:ψ_H2, NaN, tex_label = "\\psi_H2")
=#
    m <= SteadyStateParameter(:L1, NaN, tex_label = "L1")
    m <= SteadyStateParameter(:L2, NaN, tex_label = "L2")
    m <= SteadyStateParameter(:ρ_ss, NaN, tex_label = "\\rho")
    m <= SteadyStateParameter(:w1_ss, NaN, tex_label = "w1")
    m <= SteadyStateParameter(:w2_ss, NaN, tex_label = "w2")
    m <= SteadyStateParameter(:klR1_ss, NaN, tex_label = "klr1_{ss}")
    m <= SteadyStateParameter(:klR2_ss, NaN, tex_label = "klr2_{ss}")
    m <= SteadyStateParameter(:k_ss, NaN, tex_label = "k_{ss}")
    m <= SteadyStateParameter(:F1_ss, NaN, tex_label = "F1_{ss}")
    m <= SteadyStateParameter(:y1_ss, NaN, tex_label = "y1_ss")
    m <= SteadyStateParameter(:y_ss, NaN, tex_label = "y_{ss}")
    m <= SteadyStateParameter(:y2_ss, NaN, tex_label = "y2_{ss}")
#    m <= SteadyStateParameter(:F2_ss, NaN, tex_label = "F2_{ss}")
    m <= SteadyStateParameter(:x_ss, NaN, tex_label = "x_{ss}")
    m <= SteadyStateParameter(:i_S_ss, NaN, tex_label = "i_s_{ss}")
    m <= SteadyStateParameter(:g_ss, NaN, tex_label = "g_{ss}")
    m <= SteadyStateParameter(:c_ss, NaN, tex_label = "c_{ss}")
    m <= SteadyStateParameter(:τ_ss, NaN, tex_label = "\\tau_{ss}")
    m <= SteadyStateParameter(:τ_H1_ss, NaN, tex_label = "\\tau_h1_{ss}")
    m <= SteadyStateParameter(:τ_H2_ss, NaN, tex_label = "\\tau_H2_{ss}")
    m <= SteadyStateParameter(:τ_s_ss, NaN, tex_label = "\\tau_s_{ss}")

    m <= SteadyStateParameter(:t_ss, NaN, tex_label = "t_{ss}")
    m <= SteadyStateParameter(:t_H1_ss, NaN, tex_label = "t_H1_{ss}")
    m <= SteadyStateParameter(:t_H2_ss, NaN, tex_label = "t_H2_{ss}")
    m <= SteadyStateParameter(:t_S_ss, NaN, tex_label = "t_s_{ss}")

    m <= SteadyStateParameter(:c_H1_ss, NaN, tex_label = "c_h10_{ss}")
    m <= SteadyStateParameter(:c_H2_ss, NaN, tex_label = "c_h20_{ss}")
    m <= SteadyStateParameter(:c_S_ss, NaN, tex_label = "c_s0_{ss}")
    m <= SteadyStateParameter(:λ_H1_ss, NaN, tex_label = "\\lambda_h10_{ss}")
    m <= SteadyStateParameter(:λ_H2_ss, NaN, tex_label = "\\lambda_h20_{ss}")
    m <= SteadyStateParameter(:λ_S_ss, NaN, tex_label = "\\lambda_s0_{ss}")
    m <= SteadyStateParameter(:λ_SH1_ss, NaN, tex_label = "\\lambda_sh1_{ss}")
    m <= SteadyStateParameter(:λ_SH2_ss, NaN, tex_label = "\\lambda_sh1_{ss}")
    m <= SteadyStateParameter(:bᴿ_ss, NaN, tex_label = "bR0_{ss}")
    m <= SteadyStateParameter(:R_ss, NaN, tex_label = "R0_{ss}")
    m <= SteadyStateParameter(:mc1_ss, NaN, tex_label = "mc1_{ss}")
    m <= SteadyStateParameter(:mc2_ss, NaN, tex_label = "mc2_{ss}")
    m <= SteadyStateParameter(:price1_ss, NaN, tex_label = "price1_{ss}")
    m <= SteadyStateParameter(:price2_ss, NaN, tex_label = "price2_{ss}")

end

"""
```
steadystate!(m::COTHANK)
```

Calculates the model's steady-state values. `steadystate!(m)` must be called whenever
the parameters of `m` are updated.
"""
function steadystate!(m::COTHANK)

    # These steady state values are inputs given by Giorgio's matlab script
    m[:γ]        = log(1.02)/4
    m[:expγ]     = exp(m[:γ])
    m[:β]        = 0.99
    m[:π_ss]     = 1.02^(1/4)
    m[:π_S_ss]   = 1.02^(1/4)
    m[:r_ss]     = exp(m[:γ]) / m[:β] - 1.

    # Lines 553-595 are taken from Giorgio's MATLAB script for solving steady state, steady states that we
    # are unsure about are in lines 598-604
#    m[:ψ_H1]     = m[:θ] * m[:f_H1] / 1.1 # from Giorgio's MATLAB script for solving steady state
#    m[:ψ_H2]     = m[:θ] * (1 - m[:f_H1]) / 1.1 # from Giorgio's MATLAB script for solving steady state

    m[:L1]   = ((1-m[:θ]) * m[:f_S1] + m[:θ] * m[:f_H1]) * m[:H1] # from Giorgio's MATLAB script for solving steady state
    m[:L2]   = ((1-m[:θ]) * (1 - m[:f_S1]) + m[:θ] * (1 - m[:f_H1])) *m[:H2] # from Giorgio's MATLAB script for solving steady state

    price1ss_price2ss = (m[:v] * m[:L2] * m[:A2] / m[:L1] / (1 - m[:v])) ^ ((1 - m[:α]) / (m[:α] + m[:vars] - m[:α] * m[:vars]))
    m[:price1_ss]     = ( m[:v] + (1 - m[:v]) * price1ss_price2ss ^ (m[:vars] - 1)) ^ (-1 / (1 - m[:vars]))
    m[:price2_ss]     = m[:price1_ss] / price1ss_price2ss

    m[:ρ_ss]    = exp(m[:γ]) / m[:β] - (1 - m[:δ]) # from Giorgio's MATLAB script for solving steady state

    m[:w1_ss]   = (m[:price1_ss] * m[:α] ^ m[:α] * (1 - m[:α]) ^ (1 - m[:α]) / (1 + m[:λ_p_ss]) / (m[:ρ_ss] ^ m[:α])) ^ (1 / (1 - m[:α])) # from Giorgio's MATLAB script for solving steady state
    m[:w2_ss]   = m[:A2] * m[:w1_ss] * price1ss_price2ss ^ (1 / (m[:α] - 1)) # from Giorgio's MATLAB script for solving steady state

    m[:klR1_ss] = m[:w1_ss] / m[:ρ_ss] * m[:α] / (1 - m[:α]) # from Giorgio's MATLAB script for solving steady state
    m[:klR2_ss] = m[:A2] * m[:klR1_ss] * price1ss_price2ss ^ (1 / (m[:α] - 1)) # from Giorgio's MATLAB script for solving steady state
    m[:k_ss]    = m[:klR1_ss] * m[:L1] + m[:klR2_ss] * m[:L2] # from Giorgio's MATLAB script for solving steady state

    m[:F1_ss]   = m[:λ_p_ss] / (1 + m[:λ_p_ss]) * m[:L1] * m[:klR1_ss] ^ m[:α] # from Giorgio's MATLAB script for solving steady state
    m[:y1_ss]   = m[:L1] * m[:klR1_ss] ^ m[:α] - m[:F1_ss] # from Giorgio's MATLAB script for solving steady state
    m[:y_ss]    = m[:y1_ss] / m[:v] * m[:price1_ss] ^ m[:vars] # from Giorgio's MATLAB script for solving steady state
    m[:y2_ss]   = (1 - m[:v]) * m[:y_ss] * m[:price2_ss] ^ (-m[:vars]) # from Giorgio's MATLAB script for solving steady state
#    m[:F2_ss]   = m[:λ_p_ss] * m[:y2_ss] / m[:A2] # unneeded for now
    m[:x_ss]    = m[:y_ss] # from Giorgio's MATLAB script for solving steady state

    m[:i_S_ss]   = (1 - (1 - m[:δ]) * exp(-m[:γ])) * exp(m[:γ]) * m[:k_ss] / (1 - m[:θ]) # from Giorgio's MATLAB script for solving steady state
    m[:g_ss]    = m[:g_x] * m[:x_ss] # from Giorgio's MATLAB script for solving steady state
    m[:c_ss]    = m[:y_ss] - m[:g_ss] - (1 - m[:θ]) * m[:i_S_ss] # from Giorgio's MATLAB script for solving steady state

    m[:τ_ss]    = m[:τ_x] * m[:x_ss] # from Giorgio's MATLAB script for solving steady state
    m[:τ_H1_ss] = m[:τ_ss] * m[:ψ_H1] / m[:θ] / m[:f_H1] # from Giorgio's MATLAB script for solving steady state
    m[:τ_H2_ss] = m[:τ_ss] *m[:ψ_H2] / m[:θ] / (1 - m[:f_H1]) # from Giorgio's MATLAB script for solving steady state
    m[:τ_s_ss]  = (m[:τ_ss] - m[:θ] * (m[:f_H1] * m[:τ_H1_ss] + (1 - m[:f_H1]) * m[:τ_H2_ss])) / (1 - m[:θ])

    m[:t_ss]    = m[:t_x] * m[:x_ss] # from Giorgio's MATLAB script for solving steady state
    m[:t_H1_ss] = m[:t_ss] * m[:ψ_H1] / m[:θ] / m[:f_H1] # from Giorgio's MATLAB script for solving steady state
    m[:t_H2_ss] = m[:t_ss] * m[:ψ_H2] / m[:θ] / (1 - m[:f_H1]) # from Giorgio's MATLAB script for solving steady state
    m[:t_S_ss]  = (m[:t_ss] - m[:θ] * ( m[:f_H1] * m[:t_H1_ss] + (1 - m[:f_H1]) * m[:t_H2_ss])) / (1 - m[:θ]) # from Giorgio's MATLAB script for solving steady state

    # the values below will be used as initial values for solving the fixed point problem
    m[:c_H1_ss] = m[:w1_ss] * m[:H1] - m[:t_H1_ss] + m[:τ_H1_ss] # from Giorgio's MATLAB script for solving steady state
    m[:c_H2_ss] =  m[:w2_ss] * m[:H2] - m[:t_H2_ss] + m[:τ_H2_ss] # from Giorgio's MATLAB script for solving steady state
    m[:c_S_ss]  = (m[:c_ss] - m[:c_H1_ss] * m[:θ] * m[:f_H1] - m[:c_H2_ss] * m[:θ] * (1-m[:f_H1])) / (1 - m[:θ]) # from Giorgio's MATLAB script for solving steady state
    m[:λ_H1_ss] = exp(m[:γ]) / (exp(m[:γ]) - m[:h]) / m[:c_H1_ss] # from Giorgio's MATLAB script for solving steady state
    m[:λ_H2_ss] = exp(m[:γ]) / (exp(m[:γ]) - m[:h]) / m[:c_H2_ss] # from Giorgio's MATLAB script for solving steady state
    m[:λ_S_ss]  = exp(m[:γ]) / (exp(m[:γ]) - m[:h]) / m[:c_S_ss] # from Giorgio's MATLAB script for solving steady state
    m[:R_ss]    = 1.025 / 4.
    m[:bᴿ_ss]   = (m[:g_ss] - m[:t_ss] + m[:τ_ss]) / (1 - m[:R_ss] / exp(m[:γ]) / m[:π_ss]) # from Giorgio's MATLAB script for solving steady state

    # PROBLEM STEADY STATES
    m[:mc1_ss]  = m[:ρ_ss] ^ m[:α] * m[:w1_ss] ^ (1 - m[:α]) / (m[:α] ^ m[:α] * (1 - m[:α]) ^ (1 - m[:α])) / m[:price1_ss]
    m[:mc2_ss]  = m[:ρ_ss] ^ m[:α] * (m[:w2_ss] / m[:A2]) ^ (1 - m[:α]) / (m[:α] ^ m[:α] * (1 - m[:α]) ^ (1 - m[:α])) / m[:price2_ss]
    m[:π_S_ss]  = (m[:y_ss] - m[:price1_ss] * m[:mc1_ss] * m[:L1] * m[:klR1_ss] ^ m[:α] - m[:price2_ss] * m[:mc2_ss] * m[:L2] * m[:klR2_ss] ^ m[:α] * m[:A2] ^
                      (1 - m[:α])) / (1 - m[:θ])

    # function to compute residuals for solving system of nonlinear equations
    function distSS!(r::AbstractVector{S}, x::AbstractVector{S}) where {S <: Real}

        c_H1 = x[1]
        c_H2 = x[2]
        c_S  = x[3]
        λ_H1 = x[4]
        λ_H2 = x[5]
        λ_S  = x[6]
        bR   = x[7]
        R    = x[8]

        # Define residual in r
        r[1] = m[:w1_ss] * m[:H1] - m[:t_H1_ss] + m[:τ_H1_ss] + R * bR * m[:f_S1] * (1 - m[:s]) / m[:θ] / m[:f_H1] / exp(m[:γ]) / m[:π_ss] - c_H1
        r[2] = m[:w2_ss] * m[:H2] - m[:t_H2_ss] + m[:τ_H2_ss] + R * bR * (1 - m[:f_S1]) * (1 - m[:s]) / m[:θ] /(1-m[:f_H1]) / exp(m[:γ]) / m[:π_ss] - c_H2
        r[3] = m[:c_ss] / (1 - m[:θ]) - m[:θ] * (m[:f_H1] * c_H1 + (1 - m[:f_H1]) * c_H2) / (1 - m[:θ]) - c_S

        r[4] = exp(m[:γ]) / (exp(m[:γ]) - m[:h]) / c_H1 - λ_H1
        r[5] = exp(m[:γ]) / (exp(m[:γ]) - m[:h]) / c_H2 - λ_H2
        r[6] = exp(m[:γ]) / (exp(m[:γ]) - m[:h]) / c_S - λ_S

        r[7] = (m[:g_ss] - m[:t_ss] + m[:τ_ss]) / (1 - R / exp(m[:γ]) / m[:π_ss]) - bR

        r[8] = exp(m[:γ]) * m[:π_ss] / m[:β] * λ_S / (m[:s] * λ_S + (1 - m[:s]) * (m[:f_S1] * λ_H1 + (1 - m[:f_S1]) * λ_H2)) - R

        return r
    end

    init = [m[:c_H1_ss].value, m[:c_H2_ss].value, m[:c_S_ss].value, m[:λ_H1_ss].value, m[:λ_H2_ss].value, m[:λ_S_ss].value,
            m[:bᴿ_ss].value, m[:R_ss].value]
    ss = nlsolve(distSS!, init)

    @assert ss.f_converged "Steady state did not converge"

    solution = ss.zero

    m[:c_H1_ss], m[:c_H2_ss], m[:c_S_ss], m[:λ_H1_ss], m[:λ_H2_ss], m[:λ_S_ss], m[:bᴿ_ss], m[:R_ss] = solution

    m[:λ_SH1_ss] = (1 - m[:θ]) * m[:f_S1] / ((1 - m[:θ]) * m[:f_S1] + m[:θ] * m[:f_H1]) * m[:λ_S_ss] + m[:θ] *
                       m[:f_H1] / ((1 - m[:θ]) * m[:f_S1] + m[:θ] * m[:f_H1]) * m[:λ_H1_ss]
    m[:λ_SH2_ss] = (1 - m[:θ]) * (1 - m[:f_S1]) / ((1 - m[:θ]) * (1 - m[:f_S1]) + m[:θ] * (1 - m[:f_H1])) * m[:λ_S_ss] +
                       m[:θ] * (1 - m[:f_H1]) / ((1 - m[:θ]) * (1 - m[:f_S1]) + m[:θ] * (1 - m[:f_H1])) * m[:λ_H2_ss]
    return m
end

function model_settings!(m::COTHANK)

    DSGE.default_settings!(m)

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

    # Data
    m <= Setting(:data_id, 3, "Dataset identifier")
    if get_setting(m, :cond_id) in collect(1:5)
        m <= Setting(:cond_full_names, [:obs_gdp, :obs_corepce, :obs_spread, :obs_nominalrate, :obs_longrate],
                     "Observables used in conditional forecasts")
    elseif get_setting(m, :cond_id) == 6
        m <= Setting(:cond_full_names, [:obs_gdp, :obs_corepce, :obs_spread, :obs_nominalrate, :obs_longrate,
                                        :obs_gdpdeflator], "Observables used in conditional forecasts")
    end
    m <= Setting(:cond_semi_names, [:obs_spread, :obs_nominalrate, :obs_longrate],
                 "Observables used in semiconditional forecasts")

    # Forecast
    m <= Setting(:use_population_forecast, true,
                 "Whether to use population forecasts as data")
    m <= Setting(:shockdec_startdate, Nullable(quartertodate("2007-Q1")),
                 "Date of start of shock decomposition output period. If null, then shockdec starts at date_mainsample_start")
end


#=============================#
#== put needed modules here ==#
#=============================#
import PyCall
ptl = PyCall.pyimport("qcportal")

import JuliaChem

import JSON

#================================#
#== JuliaChem execution script ==#
#================================#
function script()
  #== get molecule information from QCArchive ==#
  client = ptl.FractalClient()
  s22_database = client.get_collection("ReactionDataset", "S22")

  mol_index = s22_database.get_index()[1]
  mol_object_full = s22_database.get_molecules(first) #doesn't seem to work
  mol_object = get(mol_object_full.molecule, 0)

  #== create input system ==#
  molecule = JSON.parse(mol_object.json())
  display(molecule)

  driver = "energy"

  model = Dict(
    "method" => "RHF",
    "basis" => "STO-3G"
  )

  keywords = Dict(
    "scf" => Dict(
      "niter" => 50,
      "ndiis" => 3,
      "dele" => 1E-8,
      "rmsd" => 1E-6,
      "prec" => "Float64",
      "direct" => false,
      "debug" => false
    )
  )

  #== generate basis set ==#
  basis = JuliaChem.JCBasis.run(molecule, model)

  shell = basis[2]
  display(shell)
end

JuliaChem.initialize()
script()
JuliaChem.finalize()


using Test
using ImagineHardware
using Unitful
import Unitful: V, s, Hz, m, μm, A, mA, N, g, μF

p = retrieve_positioner("NanoSX800")
@test name(p) == "NanoSX800"
@test vmin_in(p) == -20.0V
@test vmax_in(p) == 130.0V
@test vspan_in(p) == vmax_in(p) - vmin_in(p)
@test ImagineHardware.mon_slope(p, true) == 10.0V / 800.0μm
@test ImagineHardware.mon_slope(p, false) == 10.0V / 900.0μm
@test closed_loop_pad(p) == 50.0μm
@test vmin_out(p, true) == 0.0V
@test vmax_out(p, true) == 10.0V
@test vmin_out(p, true) == 0.0V
@test vmax_out(p, true) == 10.0V
@test vmin_out(p, false) == 0.0V
@test vmax_out(p, false) == 10.0V
@test vspan_out(p, true) == vmax_out(p, true) - vmin_out(p, true)
@test vspan_out(p, false) == vmax_out(p, false) - vmin_out(p, false)
@test pmin(p, true) == 0.0μm
@test pmax(p, true) == 800.0μm
@test pmin(p, false) == -50.0μm
@test pmax(p, false) == 850.0μm
@test pspan(p, true) == pmax(p, true) - pmin(p, true)
@test pspan(p, false) == pmax(p, false) - pmin(p, false)
@test ImagineHardware.blocked_pad(p) == 125μm
@test closed2open(0.0μm, p) == 0.0μm
@test closed2open(800.0μm, p) == 800.0μm
#@test_throws Exception closed2open(-10.0μm, p) #relaxed this requirement because the piezo can go this far in closed loop mode, though it's not supposed to
@test open2closed(0.0μm, p) == 0.0μm
@test open2closed(800.0μm, p) == 800.0μm
@test_throws Exception open2closed(-10.0μm, p)
@test max_displacement(p) == 1150μm
@test capacitance(p) == 7.0μF
@test motion_axis(p) == 1
@test stiffness(p) == stiffness(p, 1) == 0.2N/μm
@test stiffness_all(p) == (0.2N/μm, 2.5N/μm, 2.5N/μm)
@test effective_mass(p) == 77.1917g
@test resonance(p) == resonance(p, 1) == 210.0Hz
@test resonance_all(p) == (210.0Hz, 500.0Hz, 700.0Hz)

@test mon2pos(0.0V, p, true) == 0.0μm
@test mon2pos(0.0V, p, false) == -50.0μm
@test mon2pos(10.0V, p, true) == 800.0μm
@test mon2pos(10.0V, p, false) == 850.0μm
@test_throws Exception mon2pos(10.1V, p, false)
@test_throws Exception mon2pos(10.1V, p, true)
@test pos2mon(0.0μm, p, true) == 0.0V
@test pos2mon(-50.0μm, p, false) == 0.0V
@test pos2mon(800.0μm, p, true) == 10.0V
@test pos2mon(850.0μm, p, false) == 10.0V
@test_throws Exception pos2mon(1000μm, p, false)
#@test_throws Exception pos2mon(808μm, p, true) #relaxed this requirement becase the piezo can go this far in closed loop mode, though it's not supposed to

@test displacement2pos(0μm, p, false) == -50μm #blocked internally
@test displacement2pos(125μm, p, false) == -50μm
@test displacement2pos(125μm, p, false) == -50μm
@test displacement2pos(1025μm, p, false) == 850μm
@test displacement2pos(175μm, p, true) == 0μm
@test displacement2pos(975μm, p, true) == 800μm
@test_throws Exception displacement2pos(0μm, p, true)
@test_throws Exception displacement2pos(125μm, p, true)
@test_throws Exception displacement2pos(1025μm, p, true)
#including outside force
@test displacement2pos(125μm, p, false, 0.0N) == displacement2pos(125μm, p, false)
@test displacement2pos(125μm, p, false, 10.0N) == displacement2pos(125μm + uconvert(μm, 10.0N/stiffness(p)), p, false)

@test_throws Exception pos2displacement(-50μm, p, false)
@test_throws Exception pos2displacement(850μm, p, false)
@test pos2displacement(0μm, p, true) == 175μm
@test pos2displacement(800μm, p, true) == 975μm
@test pos2displacement(400μm, p, true) == 575μm
@test pos2displacement(400μm, p, false) == 575μm
#including outside force
@test pos2displacement(400μm, p, true, 0.0N) == 575μm
@test pos2displacement(400μm, p, false, 0.0N) == 575μm
@test pos2displacement(400μm, p, true, 10.0N) == 575μm - 10.0N/stiffness(p)
@test pos2displacement(400μm, p, true, 10.0N) == 575μm - 10.0N/stiffness(p)

@test input2displacement(-20.0V, p) == 0.0μm
@test input2displacement(130.0V, p) == max_displacement(p)
@test input2displacement(55.0V, p) == max_displacement(p) / 2
@test_throws Exception input2displacement(-30.0V, p)
@test_throws Exception input2displacement(140.0V, p)

@test displacement2input(0.0μm, p) == -20.0V
@test displacement2input(max_displacement(p) / 2, p) == 55.0V
@test displacement2input(max_displacement(p), p) == 130.0V
@test_throws Exception displacement2input(1200μm, p)
@test_throws Exception displacement2input(-10μm, p)

@test input2pos(-20.0V, p, false) == displacement2pos(0.0μm, p, false)
@test input2pos(-20.0V, p, false) == displacement2pos(0.0μm, p, false)
@test input2pos(130.0V, p, false) == displacement2pos(1150.0μm, p, false)
@test input2pos(130.0V, p, false) == displacement2pos(1150.0μm, p, false)
@test input2pos(55.0V, p, false) == displacement2pos(max_displacement(p)/2, p, false)
@test input2pos(55.0V, p, true) == displacement2pos(max_displacement(p)/2, p, true)
@test_throws Exception input2pos(-30.0V, p)
@test_throws Exception input2pos(140.0V, p)
#including outside force
@test input2pos(-20.0V, p, false, 0.0N) == input2pos(-20.0V, p, false)
@test input2pos(130.0V, p, false, 0.0N) == input2pos(130.0V, p, false)
@test input2pos(55.0V, p, false, 0.0N) == input2pos(55.0V, p, false)
@test input2pos(55.0V, p, true, 0.0N) == input2pos(55.0V, p, true)
@test input2pos(-20.0V, p, false, 10.0N) == displacement2pos(0.0μm, p, false, 10.0N)

#including outside force
@test pos2input(0μm, p, true, 0.0N) == pos2input(0μm, p, true)
@test pos2input(800μm, p, true, 0.0N) == pos2input(800μm, p, true)
@test pos2input(0μm, p, true, 10.0N) == displacement2input(pos2displacement(0μm, p, true, 10.0N), p)

#adjusted_displacement - compression = mon_displacement
#adjusted_displacement = mon_displacement + compression
#                  = mon_displacement + f / stiffness(p)
mon2displacement(0.0V, p, true) == pos2displacement(0.0μm, p, true)
@test_throws Exception mon2displacement(0.0V, p, false)
@test mon2displacement(5.0V, p, true) == max_displacement(p)/2
@test mon2displacement(0.0V, p, true, 0.0N) == pos2displacement(0.0μm, p, true, 0.0N)
@test_throws Exception mon2displacement(0.0V, p, false, 0.0N)
@test mon2displacement(0.0V, p, true, 1.0N) == pos2displacement(0.0μm, p, true, 1.0N)
@test_throws Exception mon2displacement(0.0V, p, false, 1.0N)

@test max_force(p, 0μm, 1) == stiffness(p) * max_displacement(p)
@test max_force(p, 800μm, -1) == stiffness(p) * -800μm # - pos2displacement(800.0μm, p, true))
@test max_force(p, 0μm, 800μm, true) == max_force(p, pos2displacement(0.0μm, p, true), 1)
@test max_force(p, 800μm, 0μm, true) == max_force(p, pos2displacement(800.0μm, p, true), -1) == -max_force(p, pos2displacement(0.0μm, p, true), 1)
@test max_force(p, 0μm, 800μm, false) == max_force(p, pos2displacement(0.0μm, p, false), 1)
@test max_force(p, 800μm, 0μm, false) == max_force(p, pos2displacement(800.0μm, p, false), -1) == -max_force(p, pos2displacement(0.0μm, p, false), 1)

@test loaded_resonance(100.0Hz, 70.0g, 100.0g) == uconvert(Unitful.Hz, 100.0Hz * sqrt(70.0g/(70.0g+100.0g)))
@test loaded_resonance(p, 500g) == loaded_resonance(resonance(p), effective_mass(p), 500g)


__precompile__()
module NetABM
using LightGraphs, Random, DelimitedFiles, StatsBase, SparseArrays, Distributions

export lectura_uw, lfr_network, Agent, Params, get_coop, set_adapt_agents!, update_coop!, set_coop_effect!,
    set_adapt_sign_agents!

export init_demographics!,
    get_coop!, next_state!, update_coop_infections!, update_coop_distance!
    #  , update_effect_given_distance_coop!, update_effect_given_distance_coop_free_dir!
    #  update_single_effect_distance_coop!, update_single_effect_distance_coop_free_dir!

export initialize_demographics!, set_fixed_coop_agents!, set_coop_agents!, SI_attitude!

export assign_contacts!,
    get_next_state!,
    update_state!,
    update_single_given_distance!,
    update_coop_given_distance!,
    rand_sign

export update_all_agents!,
    get_populations,
    export_parameters,
    update_single_effect_distance!,
    update_effect_given_distance!,
    update_single_effect_distance_coop!,
    update_single_effect_distance_coop_free_dir!,
    update_effect_given_distance_coop!,
    update_effect_given_distance_coop_free_dir!,
    update_effect_given_distance_coop_prior!,
    update_single_effect_distance_coop_prior!,
    update_effect_given_distance_coop_stochastic!,
    update_single_effect_distance_coop_stochastic!,
    update_effect_given_distance_coop_alpha!,
    update_single_effect_distance_alpha!,
    risk_assessment!, update_risk_assessment!,
    risk_up!, update_past!, update_single_neighbors!,
    update_all_neighbors!, update_effect_given_distance_coop_alpha_risk!

"""
    Agent(id, neighs, p_cop, state)
Agent definition for the simulations with parameters:
`id          ::Int64 ` -> ID
`state       ::String ` -> Infection state (S, I, R)
`days        ::Int64 ` -> Days the agent will remain infected
`new_state   ::String ` -> Agent's state after meetings
`num_meets   ::Int64 ` -> Number of meetings in a time step
`recovery_t  ::Int64 ` -> Recovery time
`infection_t ::Int64 ` -> Time of infection
`contacts_t  ::Vector{Int64}` -> Contacts at time t
`age_group   ::Int64 ` -> Agent's age group (demographic)
`degree_t    ::Int64` -> Degree (Number of contacts) at time t
`ϵ           ::Float64` -> Bound of confidence
`p_cop       ::Float64 ` -> Cooperation probability
`attitude    ::String ` -> Risk aversion vs risk tolerant behavior: RA, RT
`coop_effect ::Float64 ` -> The mitigation efficacy of the complaince rules take by the agent
`at_home     ::Bool ` -> Flag to represent Agent is at home
`adapter     ::Bool ` -> Flag to represent Agent is willing to change behavior
`prior_bel   ::Float64 ` -> The mean value of the direction change of adaptive behavior
"""
mutable struct Agent
    id          ::Int64
    state       ::String
    new_state   ::String
    previous    ::Array{String}
    num_meets   ::Int64
    recovery_t  ::Int64
    infection_t ::Int64
    contacts_t  ::Vector{Int64}
    coopf       ::Vector{Int64}
    non_coopf   ::Vector{Int64}
    counter     ::Int64
    age_group   ::Int64
    degree_t    ::Int64
    ϵ           ::Float64
    p_cop       ::Float64
    attitude    ::String
    coop_effect ::Float64
    new_coop_effect ::Float64
    at_home     ::Bool
    adapter     ::Bool
    prior_bel   ::Float64
    # DEFAULT CONSTRUCTOR
    Agent(id) = new(
        id,
        "S",
        "S",
        Vector{String}(),
        1,
        5,
        0,
        Vector{Int64}(),
        Vector{Int64}(),
        Vector{Int64}(),
        0,
        0,
        0,
        0.0,
        1.0,
        "rt",
        0.0,
        0.0,
        false,
        true,
        5.0
    )
end

##=================####==============##

"""
    Params(N, p_link, attack)
Global parameters for the simulation
`N::Int64 `  -> Number of agents in the population
`now_t::Int64 `  -> Current time (iteration)
`p_link::Float64` -> Link probability (Erdos-Renyi)
`p_infected_t0::Float64` -> Initial probability (fraction) of infected agents
`p_coop_agents::Float64` -> Initial probability (fraction) of lockdown agents
`attack_rate::Float64` -> Virus' attack rate
"""
mutable struct Params
    num_agents      ::Int64
    now_t           ::Int64
    μ_recovT_agents ::Int64
    p_link          ::Float64
    p_infected_t0   ::Float64
    μ_cop_agents    ::Float64
    σ_cop_agents    ::Float64
    attack_rate     ::Float64
    repetition      ::Int
    # DEFAULT CONSTRUCTOR
    Params() = new(0, 0, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0)
end

#  include("agents.jl")
include("initialize.jl")
include("interactions.jl")
include("environmental.jl")
include("update_schemes.jl")
include("opinion_dynamics.jl")
include("data_out.jl")
end # module


############################ Accessing Solution Values #####################

observables(sr::SimResult) = observables(sr.sim)
observables(sim::Simulation) = collect(keys(sim.vals.u[1]))
observables(sim::MCResult) = observables(sim[1])
constants(sim::SimResult) = sim.scenario.prob.p.constants

@inline Base.getindex(sim::Simulation, I...) = sim.vals[I...]
@inline Base.getindex(sim::Simulation, i::Symbol,::Colon) = [sim.vals[j][i] for j in 1:length(sim.vals)]
@inline Base.getindex(sr::SimResult, I...) = sr.sim[I...]
@inline Base.getindex(mc::MCResult, I...) = mc.sim[I...]

solat(sr::SimResult, args...) = solat(sr.sim, args...)

solat(sim::Simulation, t, idx, scope) = solat(sim, t, scope)[idx]
solat(sim::Simulation, t, idx::Colon, scope) = solat(sim, t, scope)

function solat(sim::Simulation, t, scope)
  if scope == :ode_ || scope == :start_
      _id = findfirst(x->x==t, sim.vals.t) # change to searchsortedfirst ?
  else
      t_id = findall(x->x==t, sim.vals.t)
      _id = t_id[findfirst(x->x==scope, @view(sim.scope[t_id]))]
  end
  return sim[_id]
end

(s::Simulation)(t, idx=:, scope=:ode_) = solat(s, t, idx, scope)
(sr::SimResult)(t, idx=:, scope=:ode_) = solat(sr.sim, t, idx, scope)

############################ DataFrames ########################################

function DataFrame(s::Simulation; vars=observables(s))
  df = DataFrame(t=s.vals.t)

  [df[!, v] = s[v,:] for v in vars[in.(vars, Ref(observables(s)))]]
  !isnothing(s.scope) && (df[!,:scope]=s.scope)

  return df
end

DataFrame(sr::SimResult; kwargs...) = DataFrame(sr.sim; kwargs...)

function DataFrame(sr::Pair{Symbol,S}; kwargs...) where S<:SimResult
  df = DataFrame(last(sr); kwargs...)
  df[!, :scenario] .= first(sr) # add new column

  return df
end

function DataFrame(res::Vector{Pair{Symbol,S}}; kwargs...) where S<:SimResult
  df_vectors = DataFrame.(res; kwargs...)

  return vcat(df_vectors...; cols=:union)
end

function DataFrame(mcr::MCResult; kwargs...)
  # df performance
  df = DataFrame()

  for (i,s) in enumerate(mcr.sim)
      dfs = DataFrame(s; kwargs...)
      insertcols!(dfs, 1, :iter => fill(i, length(s)))
      df = vcat(df,dfs)
  end

  return df
end

function DataFrame(mcr::Pair{Symbol,S}; kwargs...) where S<:MCResult
  df = DataFrame(last(mcr); kwargs...)
  df[!, :scenario] .= first(mcr) # add new column

  return df
end

function DataFrame(res::Vector{Pair{Symbol,S}}; kwargs...) where S<:MCResult
  df_vectors = DataFrame.(res; kwargs...)

  return vcat(df_vectors...; cols=:union)
end

############################ Save Result ########################################
"""
    save_results(filepath::String, sim::AbstractResult)

Save results as csv file

Arguments:

- `filepath`: path and name of the file to write to
- `sim`: simulation results of `AbstractResult` type
"""
save_results(filepath::String, sim::AbstractResult) = save_results(filepath, DataFrame(sim))

save_results(filepath::String, df::DataFrame) = CSV.write(filepath, df, delim=";")

#=FIXME
function save_results(path::String, mcsim::MCResult; groupby::Symbol=:observables)
  if groupby == :simulations
    for i in 1:length(mcsim)
      save_results("$path/$i.csv", mcsim[i])
    end
  elseif groupby == :observables
    obs = observables(mcsim[1])
    lobs = length(obs)
    for ob in obs
      df = DataFrame(t = mcsim[1].t)
      [df[!,string(i)] = mcsim[i][ob,:] for i in 1:length(mcsim)]
      CSV.write("$path/$ob.csv", df, delim=";")
    end
  end
end
=#


using MarketData

facts("diagnostic tests produce accurate results") do

  context("jaques-berra test") do
      @fact length(cl.values) => 500
  end
end


using JuliaDB
using MXNet
using PyPlot
using ProgressMeter
using JLD
using LaTeXStrings
include("self_confidence.jl")

mutable struct SQ_model
    input_sch::ML.Schema
    output_sch::ML.Schema
    X3_1net::MXNet.mx.FeedForward #neural network
    X3_2net::MXNet.mx.FeedForward #neural network
    range::Array{Float64}
end

function data_source(batchsize::Int64,train_in,train_out,valid_in,valid_out)
  train = mx.ArrayDataProvider(
    :data => train_in,
    :label => train_out,
    batch_size = batchsize,
    shuffle = true,
    )
  valid = mx.ArrayDataProvider(
    :data => valid_in,
    :label => valid_out,
    batch_size = batchsize,
    shuffle = true,
    )
  println("DONE")
  return train, valid
end

function make_nn_SQ_model(train_fname::String,valid_fname::String,log_fname::String; input_dict::Dict=Dict(),
                          output_dict::Dict=Dict(),nn_epoc::Int64=150,nn_batch_size::Int64=-1,
                          train_subsample::Int64=1,valid_subsample::Int64=1)

    training_in, training_out, training_table, input_sch, output_sch = return_data(train_fname,inputs=input_dict,outputs=output_dict,subsample=train_subsample)
    valid_in, valid_out, valid_table, _notused, _notused = return_data(valid_fname,inputs=input_dict,outputs=output_dict,subsample=valid_subsample)

    output_list = collect(keys(output_dict))
    input_list = collect(keys(input_dict))
    info("*****outputs: $output_list\n")
    info("*****inputs: $input_list")

    # create a DNN
    data = mx.Variable(:data)
    label = mx.Variable(:label)
    net = @mx.chain     mx.Variable(:data) =>
                        mx.FullyConnected(num_hidden=10) =>
                        mx.Activation(act_type=:relu) =>
                        mx.FullyConnected(num_hidden=10) =>
                        mx.Activation(act_type=:relu) =>
                        mx.FullyConnected(num_hidden= 1) =>
                        mx.LinearRegressionOutput(mx.Variable(:label))

    # final network definition, don't change, except if using gpu
    X3_1net = mx.FeedForward(net, context=mx.cpu())
    X3_2net = mx.FeedForward(net, context=mx.cpu())

    # set up the optimizer: select one, explore parameters, if desired
    #  optimizer = mx.SGD(η=0.01, μ=0.9, λ=0.00001)
    optimizer = mx.ADAM()

    # get data providers
    # Array(training_out[1,:]') -- this is some crazy conversion so the output is the right type....
    if nn_batch_size < 0
        # no value was provided
        nn_batch_size = round(Int,0.25*length(training_in))
    end

    X3_1trainprovider, X3_1evalprovider = data_source(#= batchsize =# nn_batch_size,training_in,Array(training_out[1,:]'),valid_in,Array(valid_out[1,:]'))
    X3_2trainprovider, X3_2evalprovider = data_source(#= batchsize =# nn_batch_size,training_in,Array(training_out[2,:]'),valid_in,Array(valid_out[2,:]'))

    # train, reporting loss for training and evaluation sets
    info("Training X3_1")
    mx.train(X3_1net, optimizer, X3_1trainprovider,
           initializer = mx.NormalInitializer(0.0, 0.1),
           eval_metric = mx.MSE(),
           eval_data = X3_1evalprovider,
           n_epoch = nn_epoc,
           callbacks = [mx.speedometer(),mx.do_checkpoint(string(log_fname,"_net1"),frequency=nn_epoc)])
    info("Training X3_2")
    mx.train(X3_2net, optimizer, X3_2trainprovider,
           initializer = mx.NormalInitializer(0.0, 0.1),
           eval_metric = mx.MSE(),
           eval_data = X3_2evalprovider,
           n_epoch = nn_epoc,
           callbacks = [mx.speedometer(),mx.do_checkpoint(string(log_fname,"_net2"),frequency=nn_epoc)])

    #  println(size(training_out))
    training_range = [minimum(training_out,2) maximum(training_out,2)]

    # put into a SQ_model for return
    SQ = SQ_model(input_sch,output_sch,X3_1net,X3_2net,training_range)
    return SQ
end

function SQ_predict(SQ::SQ_model,inputs::Array;use_eng_units::Bool=false)
    println("inputs:")
    display(inputs)
    X3_1plotprovider = mx.ArrayDataProvider(:data => inputs, shuffle=false)
    X3_2plotprovider = mx.ArrayDataProvider(:data => inputs, shuffle=false)
    X3_1fit = mx.predict(SQ.X3_1net, X3_1plotprovider)
    X3_2fit = mx.predict(SQ.X3_2net, X3_2plotprovider)
    println("Predicted Output:")
    display(X3_1fit)
    display(X3_2fit)
    if use_eng_units
        return restore_eng_units(inputs,SQ.input_sch), restore_eng_units([X3_1fit; X3_2fit],SQ.output_sch)
    else
        return inputs, [X3_1fit; X3_2fit]
    end
end

# convert data back to original units
function restore_eng_units(ary::Array,sch::ML.Schema)
    info("Restoring Engineering Units")
    println(size(ary))
    arrays = Dict()
    i = 1
    for entry in sch
        if entry.second != nothing
            println("key: $(entry.first), value: $(entry.second)")
            entry_mean = mean(entry.second)
            entry_std = std(entry.second)
            arrays[entry.first] = ary[i,:].*entry_std + entry_mean
            i += 1
        end
    end
    return arrays
    #  return [x.second for x in arrays]
end

function return_data(fname::String;inputs::Dict=Dict(),outputs::Dict=Dict(),
                     schema::Array=[],subsample::Int64=1)

    function my_splitschema(xs::ML.Schema,ks::Vector{Symbol})
        return filter((k,v) -> k ∉ ks, xs), filter((k,v) -> k ∈ ks, xs)
    end

    println("Processing $fname")
    table = loadtable(fname,escapechar='"')
    input_sch = nothing
    output_sch = nothing
    schema_dict = Dict()

    if isempty(schema)
        #  p = Progress(length(colnames(table)),dt=0.1,desc="Processing",output=STDOUT)
        for itm_name in colnames(table)
            if itm_name ∈ keys(inputs)
                #  ProgressMeter.next!(p,showvalues=[itm_name])
                key_type = nothing
                if contains(inputs[itm_name],"Continuous")
                    key_type = ML.Continuous
                elseif contains(inputs[itm_name],"Categorical")
                    key_type = ML.Categorical
                end
                schema_dict[itm_name] = key_type
            elseif itm_name ∈ keys(outputs)
                key_type = nothing
                if contains(outputs[itm_name],"Continuous")
                    key_type = ML.Continuous
                elseif contains(outputs[itm_name],"Categorical")
                    key_type = ML.Categorical
                end
                #  ProgressMeter.next!(p,showvalues=[itm_name])
                schema_dict[itm_name] = key_type
            else
                #  ProgressMeter.next!(p,showvalues=[itm_name])
                # ignore variables that aren't inputs or outputs
                schema_dict[itm_name] = nothing
            end
        end

        output_list = collect(keys(outputs))
        input_list = collect(keys(inputs))

        sch = ML.schema(table, hints=schema_dict)
        input_sch, output_sch = my_splitschema(sch,output_list)
    else
        input_sch = schema[1]
        output_sch = schema[2]
    end

    # allow subsampling data, subsample=1 means every data point is used
    # subsample=2 means every other point is used, etc...
    in_data = ML.featuremat(input_sch, table)[:,1:subsample:end]
    out_data = ML.featuremat(output_sch, table)[:,1:subsample:end]

    if any(isnan.(in_data)) || any(isnan.(out_data))
        println(input_sch)
        display(in_data)
        display(output_sch)
        display(out_data)
        # this can happen from mis-spelled keys, OR data like a zero in the standard deviation
        error("NaN in data!!!")
    end

    return in_data, out_data, table, input_sch, output_sch
end

function add_sq_scatter3d_annotation(ax1::PyCall.PyObject,ax2::PyCall.PyObject,in::Array,out_eng_ary::Dict,i1::Symbol,i2::Symbol,poi::Array,SQmodel::SQ_model;yval::Symbol=:X3_1,yval_std::Symbol=:X3_2,marker::AbstractString="*",s::Int64=100,color=:black,fontsize::Int64=12)

    in_eng, pred_outputs = SQ_predict(SQmodel,in,use_eng_units=true)
    display(in_eng)

    x_idx = nearest_to_x(in,poi,val_and_idx=false)
    println("#####################")
    display(in_eng)
    println(size(in_eng[i1]))
    println(x_idx)
    println(poi)
    println("[$(in_eng[i1][x_idx]) $(in_eng[i2][x_idx])]")
    println("#####################")
    #  error()

    limits = restore_eng_units(SQmodel.range,SQmodel.output_sch)
    #  println(limits)
    rwd_rng = [limits[yval][1],limits[yval][2]]

    #  println(in_eng)
    #  println(pred_outputs)
    sq = X3(Normal(out_eng_ary[yval][x_idx],out_eng_ary[yval_std][x_idx]),Normal(pred_outputs[yval][x_idx],pred_outputs[yval_std][x_idx]),global_rwd_range=rwd_rng)
    sq_txt = @sprintf("%0.3f",sq)

    println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
    println("vals: $([out_eng_ary[yval][x_idx], out_eng_ary[yval_std][x_idx]])")
    println("preds: $([pred_outputs[yval][x_idx], pred_outputs[yval_std][x_idx]])")
    println("pred_outputs: $pred_outputs")
    println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
    ax2[:errorbar](0.1,pred_outputs[yval][x_idx],pred_outputs[yval_std][x_idx],color=:blue,label="Trusted",fmt="o",capsize=3,alpha=0.6)
    ax2[:errorbar](-0.1,out_eng_ary[yval][x_idx],out_eng_ary[yval_std][x_idx],color=:red,fmt="o",capsize=3,alpha=0.6,label="Candidate")
    ax2[:set_xlim]([-0.2,0.2])
    ax2[:set_xticklabels]("")
    #  ax2[:set_ylim](rwd_rng)
    ax2[:axhline](limits[:X3_1][2])
    ax2[:axhline](limits[:X3_1][1])
    ax2[:text](0.8*0.2,limits[:X3_1][2],L"r_H",fontsize=fontsize,va="top",ha="right")
    ax2[:text](0.8*0.2,limits[:X3_1][1],L"r_L",fontsize=fontsize,va="bottom",ha="right")
    ax2[:set_title](string("SQ at ",marker,": $sq_txt"),fontsize=fontsize)
    ax2[:set_ylabel]("Reward")
    ax2[:set_adjustable]("box-forced")
    #  ax2[:set_aspect](1/(5*(limits[:X3_1][2]-limits[:X3_2][1])))
end
function norm2(ary::Array)
    n = NaN*ones(1,size(ary,2))
    for i = 1:size(ary,2)
        n[i] = norm(ary[:,i])
    end
    return n
end
function nearest_to_x(ary::Array,x::Float64;val_and_idx::Bool=false)
    resid = abs.(ary-x)
    @assert size(resid,2) == size(ary,2)
    idx_min = indmin(resid)
    if val_and_idx
        return idx_min,ary[idx_min]
    else
        return idx_min
    end
end
function nearest_to_x(ary::Array,x::Array;val_and_idx::Bool=false)
    resid = norm2(ary.-x)
    idx_min = indmin(resid)
    println(idx_min)
    if val_and_idx
        return idx_min,ary[:,idx_min]
    else
        return idx_min
    end
end

function add_sq_annotation(ax::PyCall.PyObject,x_ary::Dict,y_ary::Dict,preds::Dict,idx::Int64,xvar::Symbol,yvar_mean::Symbol,yvar_std::Symbol,SQmodel::SQ_model;fontsize::Int64=12,txt_loc::Array=[])
    #  println("########## Debug Annotation ##########")
    #  idx, idx_val = nearest_to_x(tst_in_eng[xvar],target_val,val_and_idx=true)
    limits = restore_eng_units(SQmodel.range,SQmodel.output_sch)

    note_x = x_ary[xvar][idx]

    ex_c_m = y_ary[yvar_mean][idx]
    ex_c_s = y_ary[yvar_std][idx]
    ex_c = Normal(ex_c_m,ex_c_s)

    ex_t_m = preds[yvar_mean][idx]
    ex_t_s = preds[yvar_std][idx]
    ex_t = Normal(ex_t_m,ex_t_s)

    ex_X3 = X3(ex_c,ex_t,global_rwd_range=[limits[yvar_mean][1],limits[yvar_mean][2]])

    if isempty(txt_loc)
        # auto calculate the text location
        x_lims = (minimum(x_ary[xvar]),maximum(x_ary[xvar]))
        y_lims = (limits[yvar_mean][1],limits[yvar_mean][2])

        rng_xlims = x_lims[2]-x_lims[1]
        rng_ylims = y_lims[2]-y_lims[1]
        note_pct_x = (note_x-x_lims[1])/rng_xlims*1.1
        note_pct_y = (ex_c_m-y_lims[1])/rng_ylims*1.1
        txt_loc = [note_pct_x*rng_xlims+x_lims[1],note_pct_y*rng_ylims+y_lims[1]]
        #  println("##########")
        #  println("limits: $limits")
        #  println("nx: $note_x, xlims: $x_lims")
        #  println("ny: $ex_c_m, ylims: $y_lims")
        #  #  println("pct_coords: $")
        #  println("arrow coords: $([note_x,ex_c_m])")
        #  println("##########")
    end

    ax[:annotate](@sprintf("SQ:%0.2f",ex_X3),xy=[note_x,ex_c_m],xytext=txt_loc,arrowprops=Dict(:arrowstyle=>"->"),fontsize=fontsize)
    #  ax[:annotate](".",xy=[note_x,ex_c_m],xytext=[0.8,0.],arrowprops=Dict(:arrowstyle=>"->"),fontsize=fontsize)
    #  ax[:errorbar](note_x, ex_t_m,yerr=ex_t_s/2,fmt="o",capsize=3,color=:black,alpha=0.6,label="")
    #  println("########## END Debug Annotation ##########")
end

function return_cat_stats(y::Dict,yval::Symbol,x::Dict,xval::Symbol,yval_std::Symbol;return_std::Bool=false)
    unq_x = unique(x[xval])
    num_xs = length(unq_x)
    ymat = zeros(num_xs)
    t_ucl = zeros(num_xs)
    t_lcl = zeros(num_xs)
    std_val = NaN
    for u_x in unq_x
        u_x_ind = convert(Int,round(u_x))
        vals = y[yval][x[xval].==u_x]
        if length(unique(vals)) == 1
            #using model to predict
            ymat[u_x_ind] = mean(vals)
            std_val = y[yval_std][u_x_ind]/2
            t_ucl[u_x_ind] = ymat[u_x_ind]+std_val
            t_lcl[u_x_ind] = ymat[u_x_ind]-std_val
        else
            #  println(vals)
            ymat[u_x_ind] = mean(vals)
            #  t_ucl[u_x_ind] = ymat[u_x_ind]+y[yval_std][u_x_ind]/2
            #  t_lcl[u_x_ind] = ymat[u_x_ind]-y[yval_std][u_x_ind]/2
            std_val = std(vals)
            t_ucl[u_x_ind] = ymat[u_x_ind]+std_val
            t_lcl[u_x_ind] = ymat[u_x_ind]-std_val
        end
    end
    if return_std
        return ymat,unq_x,std_val
    else
        return ymat,unq_x,t_ucl,t_lcl
    end
end

function scatter_with_conf_bnds(ax::PyCall.PyObject,x::Dict,y::Dict,xval::Symbol,yval::Symbol,yval_std::Symbol,color::Symbol;subsample::Array{Int64}=collect(1:length(x[xval])),label::String="",bar::Bool=false)
    unq_x = unique(x[xval])
    num_xs = length(unq_x)
    if bar
        x_srt2 = sortperm(x[xval])[subsample]
        ax[:errorbar](x[xval][subsample], y[yval][subsample],yerr=y[yval_std][subsample],fmt="o",capsize=3,color=color,alpha=0.6,label=label)
    else
        x_srt = sortperm(x[xval])[subsample]

        ax[:plot](x[xval][x_srt],y[yval][x_srt],label=label,color=color)
        t_ucl = y[yval]+y[yval_std]
        t_lcl = y[yval]-y[yval_std]
        #  println(t_ucl)
        #  println(x_srt)

        ax[:scatter](x[xval][x_srt],t_ucl[x_srt],label="",s=0)
        ax[:scatter](x[xval][x_srt],t_lcl[x_srt],label="",s=0)

        ax[:fill_between](x[xval][x_srt],y[yval][x_srt],t_ucl[x_srt],alpha=0.2,color=color)
        ax[:fill_between](x[xval][x_srt],y[yval][x_srt],t_lcl[x_srt],alpha=0.2,color=color)
    end

end
function make_label_from_keys(d::Dict)
    z = ""
    for x in collect(keys(d))
        z = string(z,x)
    end
    return z
end
function make_label_from_keys(a::Array)
    d = Dict()
    for z in a
        d[z] = ""
    end
    return make_label_from_keys(d)
end

function make_sq_model(net_type::String,inpts::Array{Symbol};num_epoc::Int64=250,subsample::Int64=1,
                       trusted_fname::String="",net_folder::String="nn_logs")
    train_fname = trusted_fname
    test_fname = trusted_fname #see if this works... I don't this is important for training

    #  inputs = Dict(:tprob=>"ML.Continuous",:e_mcts=>"ML.Continuous")
    inputs = Dict()
    for i in inpts
        inputs[i] = "ML.Continuous"
    end
    outputs = Dict(:X3_1=>"ML.Continuous",:X3_2=>"ML.Continuous")

    log_fname = joinpath(net_folder,"$(net_type)_$(make_label_from_keys(inputs))")

    training_subsample = subsample
    # make model
    SQmodel = make_nn_SQ_model(train_fname,test_fname,log_fname,input_dict=inputs,output_dict=outputs,nn_epoc=num_epoc,train_subsample=training_subsample)

    info("Writing variable to file:")
    println(log_fname)
    jldopen(log_fname*"_SQmodel.jld","w") do file
        JLD.addrequire(file,MXNet)
        JLD.addrequire(file,JuliaDB)
        write(file,"train_fname",train_fname)
        write(file,"test_fname",test_fname)
        write(file,"input_sch",SQmodel.input_sch)
        write(file,"output_sch",SQmodel.output_sch)
        write(file,"inputs",inputs)
        write(file,"outputs",outputs)
        write(file,"range",SQmodel.range)
    end
end


include("./ParticleSwarmAlignment.jl")
using .ParticleSwarmAlignment
using Combinatorics
using Printf

function load_fasta_file(fname)
    sequences = String[]
    line_counter = 0
    open(fname) do f
        seq = ""
        for line in eachline(f)
            if startswith(line, ">")
                if length(seq) > 0
                    push!(sequences, seq)
                end
                seq = ""
            else
                seq = string(seq, line)
            end
        end
        if length(seq) > 0
            push!(sequences, seq)
        end
    end
    return sequences
end

__precompile__()
function bigmain()
    # Load single-sequence FASTA files for different species
    files = [
        "california-sea-lion.fasta",
        "cattle.fasta",
        "chicken.fasta",
        "chimpanzee.fasta",
        #"fission-yeast.fasta",
        "human.fasta",
        "long-finned-pilot-whale.fasta",
        #"norway-rat.fasta",
        #"pacific-white-sided-dolphin.fasta",
        #"rhesus-monkey.fasta",
        #"tropical-clawed-frog.fasta",
        #"water-buffalo.fasta",
        #"white-tufted-ear-marmoset.fasta"
    ]
    for i in 1:length(files)
        files[i] = string("../fasta/HBA1/", files[i])
    end
    sequences = String[]

    for fname in files
        append!(sequences, load_fasta_file(fname))
    end

    iterations = 20
    num_particles = 10
    @time pso_score, pso_aligned = PSO_MSA(sequences, iterations, num_particles)

    println("Performing exhaustive search")
    @time exhaustive_score, exhaustive_aligned = exhaustive_search(sequences)
    @printf("BEST EXHAUSTIVE SEQUENCES (score=%d)\n", exhaustive_score)
    for seq in exhaustive_aligned
        println(seq)
    end

    @printf("PSO best score: %d, exhaustive best score: %d\n", pso_score, exhaustive_score)
end

function smallmain()
    filename = "../fasta/ENA-ebi.fasta"
    sequences = load_fasta_file(filename)
    iterations = 50
    score, aligned_sequences = PSO_MSA(sequences, iterations)

    println("Performing exhaustive search")
    best_score_exhaustive, best_aligned_exhaustive = exhaustive_search(sequences)
    @printf("BEST EXHAUSTIVE SEQUENCES (score=%d)\n", best_score_exhaustive)
    for seq in best_aligned_exhaustive
        println(seq)
    end

    @printf("PSO best score: %d, exhaustive best score: %d\n", score, best_score_exhaustive)

end

function exhaustive_search(sequences, weighted_progressive::Bool=false)
    perms = permutations(sequences)
    total_perms = length(perms)
    best_score = 0
    best_perm = String[]
    best_aligned = String[]
    i = 1
    for perm in perms
        @printf("Exhaustive Search step %d/%d\n", i, total_perms)
        p_edges = Array{Tuple{String,String,Int},1}(undef,length(perm)-1)
        for sidx in 2:length(perm)
            seqA_orig = perm[sidx-1]
            seqB_orig = perm[sidx]
            if weighted_progressive
                align_score, seqA_aligned, seqB_aligned = global_align(seqA_orig, seqB_orig)
                scoreAB = score_sequences([seqA_aligned, seqB_aligned])
            else
                # setting the edge weight to will just do the alignment inorder
                scoreAB = 0
            end
            p_edges[sidx-1] = (seqA_orig, seqB_orig, scoreAB)
        end
        aligned_sequences = progressive_alignment_inorder(perm, p_edges)
        score = score_sequences(aligned_sequences)
        if score > best_score
            best_score = score
            best_perm = copy(perm)
            best_aligned = copy(aligned_sequences)
        end
        i += 1
    end
    return best_score, best_aligned
end

function scoremain()
    sequences = [
        "-------AT-T----TAA-CT--CAGA---ACTGTC-GG---T--A-GACC----G-TG-A--AG--G-CGT--TCCC",
        "--------T-T----T---CC--CAGAG--AG-GTA-GGCTTT--ATGAGC----G--G-ATTAG--G--GTGAT---",
        "---T-ACAC-CTTAGTG--CT--CTCA---AT-GT--GG--TTG-A-GATCT---G-TT-A--A-----C-T------",
        "-AGTGA-AC-C---GTG--CG--C--A---AT--T------TTGTA-CATCT---G-TA-AG-AG--G-C-TG-T---",
        "-GGAGG-AG-C----AGA-CT--C------AC--TA-GG---T--G-TAT-----G----AG-AG--G-C-TG-T---",
        "---C-G-GTGT----TG-TCCGTC-GGC--AC-CTG-AG---TG-A---CTT---GATGCAT-------C-T--T---",
        "TGAC-G-GC-T----TGATCC--C-G----A---TG-GG--CTG-A--ACC----GATACAT-AG--G-C-C------",
        "-----A-AC-CC---AG---T--C-AA---AC-ATACGGA-ACC-A-CGCCTACAG-C--A--AG--G-C----C---",
        "-------AC------AGG-GT--CTAA---G--GTAC----TCC-T-TGTCCAC-G-C--A--TGACGTC-TA-C---",
        "-----ATGC-C-----GC-TT--C-GGTTTA--GTAC----TTG-A--GTCC-CAG-CG-A--AG-CG-C--G-----"
    ]
    score = score_sequences(sequences)
    println(score)
end

#scoremain()
bigmain()
#smallmain()


function Vector3D(arg0::Vector{jdouble})
    return Vector3D((Vector{jdouble},), arg0)
end

function Vector3D(arg0::jdouble, arg1::Vector3D)
    return Vector3D((jdouble, Vector3D), arg0, arg1)
end

function Vector3D(arg0::jdouble, arg1::Vector3D, arg2::jdouble, arg3::Vector3D)
    return Vector3D((jdouble, Vector3D, jdouble, Vector3D), arg0, arg1, arg2, arg3)
end

function Vector3D(arg0::jdouble, arg1::Vector3D, arg2::jdouble, arg3::Vector3D, arg4::jdouble, arg5::Vector3D)
    return Vector3D((jdouble, Vector3D, jdouble, Vector3D, jdouble, Vector3D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function Vector3D(arg0::jdouble, arg1::Vector3D, arg2::jdouble, arg3::Vector3D, arg4::jdouble, arg5::Vector3D, arg6::jdouble, arg7::Vector3D)
    return Vector3D((jdouble, Vector3D, jdouble, Vector3D, jdouble, Vector3D, jdouble, Vector3D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function Vector3D(arg0::jdouble, arg1::jdouble)
    return Vector3D((jdouble, jdouble), arg0, arg1)
end

function Vector3D(arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return Vector3D((jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function add(obj::Vector3D, arg0::JVector)
    return jcall(obj, "add", Vector3D, (JVector,), arg0)
end

function add(obj::Vector3D, arg0::jdouble, arg1::JVector)
    return jcall(obj, "add", Vector3D, (jdouble, JVector), arg0, arg1)
end

function angle(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "angle", jdouble, (Vector3D, Vector3D), arg0, arg1)
end

function cross_product(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "crossProduct", Vector3D, (Vector3D, Vector3D), arg0, arg1)
end

function cross_product(obj::Vector3D, arg0::JVector)
    return jcall(obj, "crossProduct", Vector3D, (JVector,), arg0)
end

function distance(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "distance", jdouble, (Vector3D, Vector3D), arg0, arg1)
end

function distance(obj::Vector3D, arg0::Point)
    return jcall(obj, "distance", jdouble, (Point,), arg0)
end

function distance1(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "distance1", jdouble, (Vector3D, Vector3D), arg0, arg1)
end

function distance1(obj::Vector3D, arg0::JVector)
    return jcall(obj, "distance1", jdouble, (JVector,), arg0)
end

function distance_inf(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "distanceInf", jdouble, (Vector3D, Vector3D), arg0, arg1)
end

function distance_inf(obj::Vector3D, arg0::JVector)
    return jcall(obj, "distanceInf", jdouble, (JVector,), arg0)
end

function distance_sq(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "distanceSq", jdouble, (Vector3D, Vector3D), arg0, arg1)
end

function distance_sq(obj::Vector3D, arg0::JVector)
    return jcall(obj, "distanceSq", jdouble, (JVector,), arg0)
end

function dot_product(::Type{Vector3D}, arg0::Vector3D, arg1::Vector3D)
    return jcall(Vector3D, "dotProduct", jdouble, (Vector3D, Vector3D), arg0, arg1)
end

function dot_product(obj::Vector3D, arg0::JVector)
    return jcall(obj, "dotProduct", jdouble, (JVector,), arg0)
end

function equals(obj::Vector3D, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_alpha(obj::Vector3D)
    return jcall(obj, "getAlpha", jdouble, ())
end

function get_delta(obj::Vector3D)
    return jcall(obj, "getDelta", jdouble, ())
end

function get_norm(obj::Vector3D)
    return jcall(obj, "getNorm", jdouble, ())
end

function get_norm1(obj::Vector3D)
    return jcall(obj, "getNorm1", jdouble, ())
end

function get_norm_inf(obj::Vector3D)
    return jcall(obj, "getNormInf", jdouble, ())
end

function get_norm_sq(obj::Vector3D)
    return jcall(obj, "getNormSq", jdouble, ())
end

function get_space(obj::Vector3D)
    return jcall(obj, "getSpace", Space, ())
end

function get_x(obj::Vector3D)
    return jcall(obj, "getX", jdouble, ())
end

function get_y(obj::Vector3D)
    return jcall(obj, "getY", jdouble, ())
end

function get_z(obj::Vector3D)
    return jcall(obj, "getZ", jdouble, ())
end

function get_zero(obj::Vector3D)
    return jcall(obj, "getZero", Vector3D, ())
end

function hash_code(obj::Vector3D)
    return jcall(obj, "hashCode", jint, ())
end

function is_infinite(obj::Vector3D)
    return jcall(obj, "isInfinite", jboolean, ())
end

function is_na_n(obj::Vector3D)
    return jcall(obj, "isNaN", jboolean, ())
end

function negate(obj::Vector3D)
    return jcall(obj, "negate", Vector3D, ())
end

function normalize(obj::Vector3D)
    return jcall(obj, "normalize", Vector3D, ())
end

function orthogonal(obj::Vector3D)
    return jcall(obj, "orthogonal", Vector3D, ())
end

function scalar_multiply(obj::Vector3D, arg0::jdouble)
    return jcall(obj, "scalarMultiply", Vector3D, (jdouble,), arg0)
end

function subtract(obj::Vector3D, arg0::JVector)
    return jcall(obj, "subtract", Vector3D, (JVector,), arg0)
end

function subtract(obj::Vector3D, arg0::jdouble, arg1::JVector)
    return jcall(obj, "subtract", Vector3D, (jdouble, JVector), arg0, arg1)
end

function to_array(obj::Vector3D)
    return jcall(obj, "toArray", Vector{jdouble}, ())
end

function to_string(obj::Vector3D)
    return jcall(obj, "toString", JString, ())
end

function to_string(obj::Vector3D, arg0::NumberFormat)
    return jcall(obj, "toString", JString, (NumberFormat,), arg0)
end


using SHA

@testset "Client" begin

    host = "127.0.0.1"
    server_port = 8080
    server_url = "http://" * host * ":" * string(server_port)

    num_comm_rounds = 1
    fraction_clients = 0.1f0
    num_total_clients = 10
    chunksize = 256

    newconfig() = Fed.Config.QuantizedConfig{UInt8}(
        Fed.Config.BaseConfig(
            server_url,
            num_comm_rounds,
            fraction_clients,
            num_total_clients
        )
    )

    function start_testing_server()
        # dummy values for testing purpose
        weights = rand(Float32, 199210)
        strategy = Fed.Server.federated_averaging
        eval_hook = (weights::Vector{Float32}) -> (0.0f0, 0.0f0)
        config = newconfig()

        central_node = Fed.Server.CentralNode(
            host,
            server_port,
            weights,
            strategy,
            eval_hook,
            config
        )

        @async Fed.Server.start(central_node)
    end

    @testset "Start client" begin
        server = start_testing_server()

        train_hook = (weights::Vector{Float32}) -> weights
        config = newconfig()

        node = Fed.Client.Node(
            host,
            8081,
            train_hook,
            config
        )

        @async Fed.Client.start(node)
        sleep(2)
        @test true
    end

end


export visualize!

"""
    visualize!(simulation; save_as, fps, frame_nums, frame_size, xlim, ylim, particle_colors, multithreaded)

# Arguments
- `simulation::Simulation`: simulation to visualize
- `save_as::Tuple{Symbol, String}` = (:gif, "temp.gif"):
- `fps::Int64 = 10`: frames per second if creating an animation
- `frame_nums::Union{Array{Int64, 1}, Symbol} = :all`: frame numbers to use.
- `frame_size::Tuple{Int64, Int64} = (600, 600)`: size of frame in pixels
- `xlim::Union{Array{Float64, 1}, Symbol} = :all`: frame limits along x-axis
- `ylim::Union{Array{Float64, 1}, Symbol} = :all`: frame limits along y-axis
- `particle_colors::Union{DefaultDict{Symbol, String, String}, Dict{Symbol, String}} = DefaultDict{Symbol, String}("black")`: particle colors
- `multithreaded::Bool = false`: use threads to update frames
"""
function visualize!(simulation::Simulation; save_as::Tuple{Symbol, String} = (:gif, "temp.gif"),
                                            fps::Int64 = 10,
                                            frame_nums::Union{Array{Int64, 1}, Symbol} = :all,
                                            frame_size::NTuple{2, Int64} = (600, 600),
                                            xlim::Union{Array{Float64, 1}, Symbol} = :all, ylim::Union{Array{Float64, 1}, Symbol} = :all,
                                            particle_colors::Union{DefaultDict{Symbol, String, String}, Dict{Symbol, String}} = DefaultDict{Symbol, String}("black"),
                                            multithreaded::Bool = false)
    save_type, save_str = save_as
    if !isdir(dirname(save_str))
        mkpath(dirname(save_str))
    end

    if frame_nums == :all
        frame_nums = [f for f = 1 : length(simulation.history)]
    elseif frame_nums == :start
        frame_nums = [1]
    elseif frame_nums == :end
        frame_nums = [length(simulation.history)]
    elseif frame_nums == :startend
        frame_nums = [1, length(simulation.history)]
    end

    if xlim == :all
        xlim = [0.0, simulation.L_x]
    end
    if ylim == :all
        ylim = [0.0, simulation.L_y]
    end

    if save_type == :gif
        print_message("GENERATING GIF")
        animation = Animation()
    elseif save_type == :images
        print_message("GENERATING IMAGES")
    end

    N_particles = length(simulation.things_to_save.particles)
    N_bonds = length(simulation.things_to_save.bonds)
    scene = plot(size = frame_size, legend = false, axis = false, grid = false, aspect_ratio = 1, xlim = xlim, ylim = ylim)
    for (f, frame_num) = enumerate(frame_nums)
        if f == 1
            θ = LinRange(0.0, 2 * pi, 20)
            unit_circle = (xs = cos.(θ), ys = sin.(θ))
            for particle in simulation.history[frame_num].particles
                cxs, cys = particle.R .* unit_circle.xs .+ particle.x, particle.R .* unit_circle.ys .+ particle.y
                color = particle_colors[particle.ptype]
                plot!(cxs, cys, seriestype = [:shape,], color = color, linecolor = color, fillalpha = 0.3)
            end

            for (particle_1, particle_2) in simulation.history[frame_num].bonds
                x_1, y_1 = particle_1.x, particle_1.y
                x_2, y_2 = particle_2.x, particle_2.y
                Δx = wrap_displacement(x_1 - x_2; period = simulation.periodic_in_x ? simulation.L_x : -1.0)
                Δy = wrap_displacement(y_1 - y_2; period = simulation.periodic_in_y ? simulation.L_y : -1.0)
                plot!([x_2, x_2 + Δx], [y_2, y_2 + Δy], color = "black", linealpha = 0.3)
                plot!([x_1, x_1 - Δx], [y_1, y_1 - Δy], color = "black", linealpha = 0.3)
            end
        else
            prev_frame_num = frame_nums[f - 1]
            for n = 1 : N_particles
                prev_particle, particle = simulation.history[prev_frame_num].particles[n], simulation.history[frame_num].particles[n]
                Δx, Δy = particle.x - prev_particle.x, particle.y - prev_particle.y
                scene.series_list[n][:x] .+= Δx
                scene.series_list[n][:y] .+= Δy
            end
            for n = N_particles + 1 : 2 : N_particles + 2 * N_bonds
                particle_1, particle_2 = simulation.history[frame_num].bonds[div(n - N_particles - 1, 2) + 1]
                x_1, y_1 = particle_1.x, particle_1.y
                x_2, y_2 = particle_2.x, particle_2.y
                Δx = wrap_displacement(x_1 - x_2; period = simulation.periodic_in_x ? simulation.L_x : -1.0)
                Δy = wrap_displacement(y_1 - y_2; period = simulation.periodic_in_y ? simulation.L_y : -1.0)
                scene.series_list[n][:x], scene.series_list[n][:y] = [x_2, x_2 + Δx], [y_2, y_2 + Δy]
                scene.series_list[n + 1][:x], scene.series_list[n + 1][:y] = [x_1, x_1 - Δx], [y_1, y_1 - Δy]
            end
        end

        if save_type == :gif
            frame(animation)
        elseif save_type == :images
            savefig(string(save_str, "Frame ", frame_num, ".png"))
        end
        println("Frame ", frame_num, " done")
    end

    if save_type == :gif
        gif(animation, save_str, fps = fps)
        print_message("GIF GENERATED")
    elseif save_type == :images
        print_message("IMAGES GENERATED")
    end
end


export Conn
import Base.show

abstract type Obj end

mutable struct Conn <: Obj
    input::Float64
    output::Float64
    state::Array{Float64}
    params::Array{Float64}
    source::Ref
    dest::Ref
end

function Conn(cfg::Dict, source::Obj, dest::Obj, params::Array{Float64})
    Conn(0.0, 0.0, zeros(cfg["n_conn_state"]), params, Ref(source), Ref(dest))
end

function get_params(c::Conn, n::Int64)
    vcat(c.params, zeros(n - length(c.params)))
end

function show(io::IO, x::Conn)
    print(io, string("Conn(", x.input, ", ", x.output, ", ", x.state, ", ",
                     x.params, ", ", x.source[].params, ", ", x.dest[].params, ")"));
end


"""
Abstract MCMC sampler type
"""
abstract type MCMCSampler end

"""
MCMC sampler struct for AdvancedHMC NUTS

* `model`: model function that accepts data
* `config`: sampler configution settings
* `name`: a unique identifer given to each sampler
"""
mutable struct AHMCNUTS{T1} <: MCMCSampler
    model::T1
    config::Turing.Inference.NUTS
    name::Symbol
    Nsamples::Int64
    autodiff::Symbol
end

AHMCNUTS(model, config)=AHMCNUTS(model, config, :AHMCNUTS, 0, :forward)

"""
MCMC sampler struct for CmdStan NUTS

* `model`: model configuration
* `dir`: probject directory
* `id`: a unique identifier for each instance of CmdStan in parallel applications
* `name`: a unique identifer given to each sampler
"""
mutable struct CmdStanNUTS{T1} <: MCMCSampler
    model::T1
    dir::String
    id::String
    name::Symbol
    autidiff::Symbol
end

CmdStanNUTS(model, dir) = CmdStanNUTS(model, dir, model.name, :CmdStanNUTS, :reverse)

"""
MCMC sampler struct for DynamicHMC NUTS

* `model`: model function that accepts data
* `config`: sampler configution settings
* `name`: a unique identifer given to each sampler
"""
mutable struct DHMCNUTS{T1,T2} <: MCMCSampler
    model::T1
    Nsamples::T2
    name::Symbol
    autodiff::Symbol
end

DHMCNUTS(model) = DHMCNUTS(model,0,:DHMCNUTS,:forward)

"""
Primary function that performs mcmc benchmark repeatedly on a set of samplers
and records the results.
* 'sampler': tuple of sampler objects
* `results`: DataFrame containing benchmark results
* `csr̂`: cross sampler r̂
* `simulate`: data generating function
* `Nreps`: number of repetitions for a given set of simulation parameters. Default = 100
* `kwargs`: optional keyword arguments that are passed to modifyConfig!, updateResults! and
runSampler, providing flexibility in benchmark simulations.
"""
function benchmark!(samplers, results, csr̂, simulate, Nreps, chains;kwargs...)
    for rep in 1:Nreps
      data = simulate(;kwargs...)
      schains = Chains[]
      for s in samplers
          modifyConfig!(s;kwargs...)
          println("\nSampler: $(typeof(s))")
          println("Simulation: $simulate")
          println("Repetition: $rep of $Nreps\n")
          foreach(x->println(x),kwargs)
          performance = @timed runSampler(s, data;kwargs...)
          push!(schains,performance[1])
          allowmissing!(results)
          results = updateResults!(s, performance, results;kwargs...)
          savechain!(s ,chains, performance;kwargs...)
      end
      csr̂=cross_samplerRhat!(schains, csr̂;kwargs...)
    end
    return results,csr̂
end

function benchmark!(sampler::T, results, csr̂, simulate, Nreps, chains;kwargs...) where {T<:MCMCSampler}
    return benchmark!((sampler, ), results, simulate, Nreps, chains;kwargs...)
end

"""
Runs the benchmarking procedure and returns the results

* `samplers`: a tuple of samplers or a single sampler object
* `simulate`: model simulation function with keyword Nd
* `Nreps`: number of times the benchmark is repeated for each factor combination
"""
 function benchmark(samplers, simulate, Nreps, chains=();kwargs...)
     results = DataFrame()
     csr̂ = DataFrame()
     for p in Permutation(kwargs)
         results,csr̂=benchmark!(samplers, results, csr̂, simulate, Nreps, chains;p...)
     end
     return [results csr̂]
 end

 """
 Runs the benchmarking procedure defined in benchmark in parallel and returns the results

 * `samplers`: a tuple of samplers or a single sampler object
 * `simulate`: model simulation function with keyword Nd
 * `Nreps`: number of times the benchmark is repeated for each factor combination
 """
 function pbenchmark(samplers, simulate, Nreps;kwargs...)
     compile(samplers, simulate;kwargs...)
     pfun(rep) = benchmark(samplers, simulate, rep, chains;kwargs...)
     reps = setreps(Nreps)
     presults = pmap(rep->pfun(rep),reps)
     return vcat(presults..., cols=:union)
 end

"""
Extracts model and configuration from sampler object and performs
parameter estimation
* `s`: sampler object
* `data`: data for benchmarking
* `Nchains`: number of chains ran serially. Default =  1
"""
function runSampler(s::AHMCNUTS, data;Nchains=1, kwargs...)
    f() = sample(s.model(data...), s.config, s.Nsamples;discard_adapt=false)
    return reduce(chainscat, map(x->f(),1:Nchains))
end

function runSampler(s::CmdStanNUTS, data;Nchains=1, kwargs...)
    f() = stan(s.model, toDict(data), summary=false, s.dir)[2]
    return reduce(chainscat, map(x->f(), 1:Nchains))
end

function runSampler(s::DHMCNUTS, data;Nchains=1, kwargs...)
    f() = s.model(data...,s.Nsamples,s.autodiff)
    return reduce(chainscat, map(x->f(),1:Nchains))
end

"""
Update the results DataFrame on each iteration
* `s`: MCMC sampler object
* `performance`: includes MCMC Chain, execution time, and memory measurements
* `results`: DataFrame containing benchmark results
"""
function updateResults!(s::AHMCNUTS, performance, results;kwargs...)
    chain = performance[1]
    newDF = DataFrame()
    chain=removeBurnin(chain;kwargs...)
    df = MCMCChains.describe(chain)[1].df
    addChainSummary!(newDF, chain,df,:ess)
    addChainSummary!(newDF, chain,df,:r_hat)
    addESStime!(newDF, chain, df, performance)
    addHPD!(newDF, chain)
    addMeans!(newDF, df)
    permutecols!(newDF, sort!(names(newDF)))#ensure correct order
    dfi=MCMCChains.describe(chain, sections=[:internals])[1]
    newDF[!,:epsilon] = [dfi[:step_size,:mean][1]]
    newDF[!,:tree_depth] = [dfi[:tree_depth, :mean][1]]
    addPerformance!(newDF,performance)
    newDF[!,:sampler] = [s.name]
    addKW!(newDF;kwargs...)
    return vcat(results, newDF, cols=:union)
end

function updateResults!(s::CmdStanNUTS, performance, results;kwargs...)
    chain = performance[1]
    newDF = DataFrame()
    chain=removeBurnin(chain;kwargs...)
    df = MCMCChains.describe(chain)[1].df
    addChainSummary!(newDF, chain, df, :ess)
    addChainSummary!(newDF,chain,df,:r_hat)
    addESStime!(newDF, chain, df, performance)
    addHPD!(newDF, chain)
    addMeans!(newDF, df)
    permutecols!(newDF, sort!(names(newDF)))#ensure correct order
    dfi=MCMCChains.describe(chain, sections=[:internals])[1]
    newDF[!,:epsilon] = [dfi[:stepsize__, :mean][1]]
    newDF[!,:tree_depth] = [dfi[:treedepth__, :mean][1]]
    addPerformance!(newDF, performance)
    newDF[!,:sampler] = [s.name]
    addKW!(newDF;kwargs...)
    return vcat(results, newDF, cols=:union)
end

function updateResults!(s::DHMCNUTS, performance, results;kwargs...)
    chain = performance[1]
    newDF = DataFrame()
    chain = removeBurnin(chain;kwargs...)
    df = MCMCChains.describe(chain)[1].df
    addChainSummary!(newDF, chain, df, :ess)
    addChainSummary!(newDF, chain, df, :r_hat)
    addESStime!(newDF, chain, df, performance)
    addHPD!(newDF, chain)
    addMeans!(newDF, df)
    permutecols!(newDF, sort!(names(newDF)))#ensure correct order
    dfi=MCMCChains.describe(chain, sections=[:internals])[1]
    newDF[!,:epsilon] = [dfi[:lf_eps, :mean][1]]
    newDF[!,:tree_depth] = [dfi[:tree_depth, :mean][1]]
    addPerformance!(newDF,performance)
    newDF[!,:sampler] = [s.name]
    addKW!(newDF;kwargs...)
    return vcat(results, newDF, cols=:union)
end

"""
Modifies MCMC sampler configuration, including number of samples, target
acceptance rate and others depending on the specific sampler.
* `s`: sampler object
* `Nsamples`: total number of MCMC samples
* `Nadapt`: number of adaption samples during warm up
* `delta`: target acceptance rate.
"""
function modifyConfig!(s::AHMCNUTS;Nsamples, Nadapt, delta, autodiff=:forward, kwargs...)
    s.config = Turing.NUTS(Nadapt, delta)
    s.Nsamples = Nsamples
    s.autodiff = autodiff == :reverse ? Turing.setadbackend(:reverse_diff) : Turing.setadbackend(:forward_diff)
end

function modifyConfig!(s::CmdStanNUTS;Nsamples, Nadapt, delta, kwargs...)
    s.model.num_samples = Nsamples-Nadapt
    s.model.num_warmup = Nadapt
    s.model.method.adapt.delta = delta
    s.model.method.num_samples = Nsamples-Nadapt
    s.model.method.num_warmup = Nadapt
    id = myid()
    if id != 1
        s.model.name = string(s.id,id)
    end
end

function modifyConfig!(s::DHMCNUTS;Nsamples, autodiff=:forward, kwargs...)
    s.Nsamples = Nsamples
    s.autodiff = autodiff == :reverse ? :ReverseDiff : :ForwardDiff
end