using Revise
module QFactors
using Revise

using  CSV, DataFrames, Serialization, Dates, Statistics,
  Distributed, Finometrics, Measures, StatsBase,
  Formatting, CodecLz4, CodecZlib,  Distributions, LinearAlgebra

if  pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

import Base.eltype, Base.getindex, Base.iterate, Base.length

#necessary due to exporting conflict
describe(df::DataFrame) = DataFrames.describe(df)

############################CONSTANTS

const OVERLEAF_PATH = "C:\\Users\\Clinton\\Dropbox\\Apps\\Overleaf\\QFactors"
const TABLE_PATH = "$OVERLEAF_PATH\\sub\\tab"

const OUT_PATH = pwd() * "\\output"
const DATA_PATH = pwd() * "\\data"

const CRSP_TYPE = :monthly
const CRSP_SUFFIX = CRSP_TYPE == :monthly ? 'M' : 'D'
const CRSP_NAME = "CRSP-$CRSP_SUFFIX"
const CRSP_DATE_FORMAT = dateformat"yyyymmdd"

const COMP_NAME = "COMP"
const COMP_A_NAME = "COMP-A"
const COMP_Q_NAME = "COMP-Q"
const COMP_DATE_FORMAT = dateformat"yyyymmdd"

const CCM_NAME = "CCM"
const CCM_DATE_FORMAT = dateformat"yyyymmdd"

const LAST_DATE_STRING = "20191231"
const LAST_DATE = Date(LAST_DATE_STRING, CCM_DATE_FORMAT)

const UNIV_NAME = "$(CRSP_NAME)_$(COMP_NAME)_UNIV"

const FACT_NAME = "$(CRSP_NAME)_$(COMP_NAME)_FACT"

const TEST_OUTPUT = true

const TIME_2_STALE = Month(6)
const MIN_DATE = Date(1972,1,1)
const MAX_DATE = Date(2018,12,31)

const VWCRSP_NAME = "VWCRSP-D"
const RFCRSP_NAME = "RFCRSP-M"
const IDXCRSP_PATH = "$DATA_PATH\\crspidx"
const IDXCRSP_DATE_FORMAT = CRSP_DATE_FORMAT
const IDX_DATE_COL = :caldt

const FF_PORT_NAME = "FF-25-SIZE-BM-D"
const FF_PATH = "$DATA_PATH\\french"
const FF_PORT_RENAME =   Dict(:smalllobm=>:me1bm1, :smallhibm=>:me1bm5,
  :biglobm=>:me5bm1, :bighibm=>:me5bm5)
const FF_PORT_VALS = [Symbol("me$(i)bm$(j)") for i ∈ 1:5 for j ∈ 1:5]
const FF_DATE_FORMAT = CRSP_DATE_FORMAT

#############################Multi-threading

#available lock if needed- use lock(spinlock) unlock(spinlock) etc
const spinlock = Threads.SpinLock()

const PARALLEL = Ref{Bool}(true)

#NOTE: This macro allows for the parallelization routines to be shut down
#by flicking the above PARALLEL switch
macro mpar(expr)
  if PARALLEL[]
    quote
      :($(Threads.@threads($expr)))
    end
  else
    quote
      :($($expr))
    end
  end
end

###########################IO Compression/Decompression routines
IN_CSV_STREAM(p::String) = LZ4DecompressorStream(open(p))
IN_CSV_STREAM(F::Function, p::String) = open(F, LZ4DecompressorStream, p)

IN_JLS_STREAM(p::String) = LZ4DecompressorStream(open(p))
IN_JLS_STREAM(F::Function, p::String) = open(F, LZ4DecompressorStream, p)

OUT_JLS_STREAM(p::String) = LZ4CompressorStream(open(p, "w"))
OUT_JLS_STREAM(F::Function, p::String) = open(F, LZ4CompressorStream, p, "w")

##############################EXPORTS

############################FILES
include("QUtilities.jl")
include("FormFactors\\PreprocessCRSP.jl")
include("FormFactors\\PreprocessCOMP.jl")
include("FormFactors\\MergeCOMPCCM.jl")
include("FormFactors\\MergeCRSPCOMP.jl")
include("FormFactors\\PortfolioConstruction.jl")
include("FormFactors\\FormFactors.jl")
include("FormFactors\\MergeFACTOthers.jl")
include("AnalyzeFactors.jl")
include("AnalyzeQFactors.jl")

end


"""
    ArrayManifold{M <: Manifold} <: Manifold

A manifold to encapsulate manifolds working on array representations of
`MPoints` and `TVectors` in a transparent way, such that for these manifolds its
not necessary to introduce explicit types for the points and tangent vectors,
but they are encapusalted/stripped automatically when needed.
"""
struct ArrayManifold{M <: Manifold} <: Manifold
    manifold::M
end
convert(::Type{M},m::ArrayManifold{M}) where M <: Manifold = m.manifold
convert(::Type{ArrayManifold{M}},m::M) where M <: Manifold = ArrayManifold(M)

manifold_dimension(M::ArrayManifold) = manifold_dimension(M.manifold)

@traitimpl IsDecoratorManifold{ArrayManifold}

struct ArrayMPoint{V <: AbstractArray{<:Number}} <: MPoint
    value::V
end
convert(::Type{V},x::ArrayMPoint{V}) where V <: AbstractArray{<:Number} = x.value
convert(::Type{ArrayMPoint{V}},x::V) where V <: AbstractArray{<:Number} = ArrayPoint{V}(x)
eltype(::Type{ArrayMPoint{V}}) where V = eltype(V)
similar(x::ArrayMPoint) = ArrayMPoint(similar(x.value))
similar(x::ArrayMPoint, ::Type{T}) where T = ArrayMPoint(similar(x.value, T))

struct ArrayTVector{V <: AbstractArray{<:Number}} <: TVector
    value::V
end
convert(::Type{V},v::ArrayTVector{V}) where V <: AbstractArray{<:Number} = v.value
convert(::Type{ArrayTVector{V}},v::V) where V <: AbstractArray{<:Number} = ArrayTVector{V}(v)
eltype(::Type{ArrayTVector{V}}) where V = eltype(V)
similar(x::ArrayTVector) = ArrayTVector(similar(x.value))
similar(x::ArrayTVector, ::Type{T}) where T = ArrayTVector(similar(x.value, T))

array_value(x::AbstractArray) = x
array_value(x::ArrayMPoint) = x.value
array_value(v::ArrayTVector) = v.value

(+)(v1::ArrayTVector, v2::ArrayTVector) = ArrayTVector(v1.value + v2.value)
(-)(v1::ArrayTVector, v2::ArrayTVector) = ArrayTVector(v1.value - v2.value)
(-)(v::ArrayTVector) = ArrayTVector(-v.value)
(*)(a::Number, v::ArrayTVector) = ArrayTVector(a*v.value)

function isapprox(M::ArrayManifold, x, y; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_manifold_point(M, y; kwargs...)
    return isapprox(M.manifold, array_value(x), array_value(y); kwargs...)
end

function isapprox(M::ArrayManifold, x, v, w; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_tangent_vector(M, x, v; kwargs...)
    is_tangent_vector(M, x, w; kwargs...)
    return isapprox(M.manifold, array_value(x), array_value(v), array_value(w); kwargs...)
end

function project_tangent!(M::ArrayManifold, w, x, v; kwargs...)
    is_manifold_point(M, x; kwargs...)
    project_tangent!(M.manifold, w.value, array_value(x), array_value(v))
    is_tangent_vector(M, x, w; kwargs...)
    return w
end

function distance(M::ArrayManifold, x, y; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_manifold_point(M, y; kwargs...)
    return distance(M.manifold, array_value(x), array_value(y))
end

function inner(M::ArrayManifold, x, v, w; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_tangent_vector(M, x, v; kwargs...)
    is_tangent_vector(M, x, w; kwargs...)
    return inner(M.manifold, array_value(x), array_value(v), array_value(w))
end

function exp!(M::ArrayManifold, y, x, v; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_tangent_vector(M, x, v; kwargs...)
    exp!(M.manifold, array_value(y), array_value(x), array_value(v))
    is_manifold_point(M, y; kwargs...)
    return y
end

function log(M::ArrayManifold, x, y; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_manifold_point(M, y; kwargs...)
    v = ArrayTVector(log(M.manifold, array_value(x), array_value(y)))
    is_tangent_vector(M, x, v; kwargs...)
    return v
end

function log!(M::ArrayManifold, v, x, y; kwargs...)
    is_manifold_point(M, x; kwargs...)
    is_manifold_point(M, y; kwargs...)
    log!(M.manifold, array_value(v), array_value(x), array_value(y))
    is_tangent_vector(M, x, v; kwargs...)
    return v
end

function zero_tangent_vector!(M::ArrayManifold, v, x; kwargs...)
    is_manifold_point(M, x; kwargs...)
    zero_tangent_vector!(M.manifold, array_value(v), array_value(x); kwargs...)
    is_tangent_vector(M, x, v; kwargs...)
    return v
end

function zero_tangent_vector(M::ArrayManifold, x; kwargs...)
    is_manifold_point(M, x; kwargs...)
    w = zero_tangent_vector(M.manifold, array_value(x))
    is_tangent_vector(M, x, w; kwargs...)
    return w
end

function vector_transport(M::ArrayManifold, x, v, y)
    return vector_transport(M.manifold,
                            array_value(x),
                            array_value(v),
                            array_value(y))
end

export ArrayManifold,
    ArrayMPoint,
    ArrayTVector


push!(LOAD_PATH, joinpath(dirname(@__FILE__), "../../.."))
push!(LOAD_PATH, dirname(@__FILE__))
ENV["USE_CUDA"] = true
workspace()
using Merlin

function bench()
  a =  Array{Float32,2}[]
  push!(a, zeros(Float32,100,100))
  push!(a, zeros(Float32,100,100))
  for i = 1:1000
    Merlin.concat(a, 1)
    #cat(1,a...)
  end
end

@time bench()

include("token.jl")
include("model.jl")
include("train.jl")
path = "C:/Users/hshindo/Dropbox/tagging"

@time train(path)


using .CuArrays
using Serialization
import Serialization: serialize

const CuFlatS0{P,T,M<:CuArray} = FlatS0{P,T,M}

### broadcasting
preprocess(dest::F, bc::Broadcasted) where {F<:CuFlatS0} =
    Broadcasted{Nothing}(CuArrays.cufunc(bc.f), preprocess_args(dest, bc.args), map(OneTo,size_2d(F)))
preprocess(dest::F, arg) where {M,F<:CuFlatS0{<:Any,<:Any,M}} =
    adapt(M,broadcastable(F, arg))
function copyto!(dest::F, bc::Broadcasted{Nothing}) where {F<:CuFlatS0}
    bc′ = preprocess(dest, bc)
    copyto!(firstfield(dest), bc′)
    return dest
end
BroadcastStyle(::FlatS0Style{F,Array}, ::FlatS0Style{F,CuArray}) where {P,F<:FlatS0{P}} =
    FlatS0Style{basetype(F){P},CuArray}()

# always adapt to Array storage when serializing since we may deserialize in an
# environment that does not have CuArrays loaded
serialize(s::AbstractSerializer, f::CuFlatS0) = serialize(s, adapt(Array,f))

### misc
# the generic versions of these trigger scalar indexing of CuArrays, so provide
# specialized versions:

function copyto!(dst::F, src::F) where {F<:CuFlatS0}
    copyto!(firstfield(dst),firstfield(src))
    dst
end
pinv(D::Diagonal{T,<:CuFlatS0}) where {T} = Diagonal(@. ifelse(isfinite(inv(D.diag)), inv(D.diag), $zero(T)))
inv(D::Diagonal{T,<:CuFlatS0}) where {T} = any(Array((D.diag.==0)[:])) ? throw(SingularException(-1)) : Diagonal(inv.(D.diag))
fill!(f::CuFlatS0, x) = (fill!(firstfield(f),x); f)
dot(a::CuFlatS0, b::CuFlatS0) = sum_kbn(Array(Map(a).Ix .* Map(b).Ix)) * fieldinfo(a).Δx^2
≈(a::CuFlatS0, b::CuFlatS0) = (firstfield(a) ≈ firstfield(b))

# some pretty low-level hacks to get a few thing broadcasting correctly for
# Complex arguments that don't currently work in CuArrays
CuArrays.CUDAnative.isfinite(x::Complex) = Base.isfinite(x)
CuArrays.CUDAnative.sqrt(x::Complex) = CuArrays.CUDAnative.sqrt(CuArrays.CUDAnative.abs(x)) * CuArrays.CUDAnative.exp(im*CuArrays.CUDAnative.angle(x)/2)
CuArrays.culiteral_pow(::typeof(^), x::Complex, ::Val{2}) = x * x


using Distributed
addprocs(4)

@everywhere begin
    using Pkg
    Pkg.activate(".")
end

sleep(0.1)

@everywhere begin
    using CSV
    using POMDPs
    using POMDPPolicies
    using POMDPSimulators
    using MultiAgentPOMDPs
    using FactoredValueMCTS
    using MultiUAVDelivery
    using Random
    using Statistics
    using Distributions
    include("pipeline.jl")
end

const PSET = [(nagents=8, XY_AXIS_RES=0.2, XYDOT_LIM=0.2, XYDOT_STEP=0.2, NOISESTD=0.1),
               (nagents=16, XY_AXIS_RES=0.1, XYDOT_LIM=0.1, XYDOT_STEP=0.1, NOISESTD=0.05),
               (nagents=32, XY_AXIS_RES=0.08, XYDOT_LIM=0.08, XYDOT_STEP=0.08, NOISESTD=0.05),
               (nagents=48, XY_AXIS_RES=0.05, XYDOT_LIM=0.05, XYDOT_STEP=0.05, NOISESTD=0.02)]

const REWSET = (goal_bonus=1000.0, prox_scaling = 1.0, repul_pen=10.0, dynamics_scaling=10.0)

const MAX_STEPS = 200

function get_maxplus_solver(d, n, c, k, agent_utils, node_exp, edge_exp)
    solver = FVMCTSSolver(
        depth=d,
        n_iterations=n,
        exploration_constant=c,
        rng=MersenneTwister(8),
        coordination_strategy = MaxPlus(
            message_iters=k,
            message_norm=true,
            use_agent_utils=agent_utils,
            node_exploration=node_exp,
            edge_exploration=edge_exp,
        )
    )
    return solver
end

function get_varel_solver(d, n, c)
    solver = FVMCTSSolver(depth=d,
                             n_iterations=n,
                             exploration_constant=c,
                             rng=MersenneTwister(8),
                             coordination_strategy=VarEl())
    return solver
end

function run_experiment(prob, solver_name, nevals, logdir)
    mdp = FirstOrderMultiUAVDelivery(pset=PSET[prob], rewset=REWSET)

    psetvals = join(("$p1=$p2" for (p1, p2) in pairs(PSET[prob])), ",")
    rsetvals = join(("$p1=$p2" for (p1, p2) in pairs(REWSET)), ",")
    exp_name = "UAV_delivery_$(psetvals)_$(rsetvals)_$(solver_name)"
    if solver_name == :random
        solver = RandomSolver(rng=MersenneTwister(1999))
        @info exp_name
        df = online_evaluate(Dict(solver_name=>solver), mdp, nevals, MAX_STEPS)
        CSV.write(joinpath(logdir, "$(exp_name).csv"), df)
        for k in unique(df.policy)
            fdf = filter(r->r[:policy] == k, df)
            vals = fdf.discret
            uvals = fdf.undiscret
            @info k mean=mean(vals) std=std(vals)
            @info k mean=mean(uvals) std=std(uvals)
        end
        return nothing
    end

    d = 10
    n = 5*PSET[prob].nagents

    if prob == 1
        c = 5
        k = 10
    elseif prob == 2
        c = 10
        k = 20
    elseif prob == 3
        c = 20
        k = 20
    else
        c = 30
        k = 20
    end

    if solver_name == :maxplus
        exp_name *= "_d=$(d),n=$(n),c=$(c),k=$(k)"
        solver = get_maxplus_solver(d, n, c, k, true, true, false)
    elseif solver_name == :varel
        exp_name *= "_d=$(d),n=$(n),c=$(c)"
        solver = get_varel_solver(d, n, c)
    end
    @info exp_name
    df = online_evaluate(Dict(solver_name=>solver), mdp, nevals, MAX_STEPS)
    CSV.write(joinpath(logdir, "$(exp_name).csv"), df)
    for k in unique(df.policy)
        fdf = filter(r->r[:policy] == k, df)
        vals = fdf.discret
        uvals = fdf.undiscret
        @info "discounted" k mean=mean(vals) std=std(vals)
        @info "undiscounted" k mean=mean(uvals) std=std(uvals)
    end
end

function main(args)
    prob = parse(Int, args[1])
    solver = Symbol(args[2])
    nevals = parse(Int, args[3])

    if length(args) < 4
        logdir = "."
    else
        logdir = args[4]
    end
    !isdir(logdir) && mkpath(logdir)
    run_experiment(prob, solver, nevals, logdir)
end

(abspath(PROGRAM_FILE) == @__FILE__) && main(ARGS)


module StaticNets

using Distributions
using StatsBase
using Optim
using LineSearches
using StatsFuns
using Roots
using MLBase
using Statistics
using LinearAlgebra
using Random
using Distributed

using ..AReg
using ..Utilities
using ..Scalings
using ..ErgmRcall

## STATIC NETWORK MODEL
abstract type Ergm end
abstract type ErgmDirBin <: Ergm end

#constants
targetErrValStaticNets = 1e-2
export targetErrValStaticNets

targetErrValStaticNetsW = 1e-5
export targetErrValStaticNetsW

bigConstVal = 10^6
export bigConstVal

maxLargeVal =  1e40# 1e-10 *sqrt(prevfloat(Inf))
export maxLargeVal

minSmallVal = 1e2*eps()
export minSmallVal

# region StaticNet interface (work in progress)

name(x::T where T <: Ergm) = error("to be defined")

type_of_obs(model::T where T <: Ergm) = error("to be defined")

obj_fun(model::T where T <: Ergm, obs_t, N, par) = error("to be defined")

grad_obj_fun(model::T where T <: Ergm, obs_t, N, par) = error("to be defined")

hessian_obj_fun(model::T where T <: Ergm, obs_t, N, par) = error("to be defined")

fisher_info_obj_fun(model::T where T <: Ergm, obs_t, N, par) = error("to be defined")

estimate(model::T where T <: Ergm, A) = error("to be defined")

estimate_sequence(model::T where T <: Ergm, obsT) = error("to be defined")

sample_ergm(model::T where T <: Ergm, N, parVec, nSample) = error("to be defined")  # return a 3D array N x N x nSample

"""
return a 4D array N x N x T x nSample
"""
function sample_ergm_sequence(model::T where T <: Ergm, N, parVecSeq_T::Matrix, nSample)

    T = size(parVecSeq_T)[2]
    A_T_dgp = zeros(Int8, N, N, T, nSample)
    for t=1:T
        A_T_dgp[:, :, t, :] = sample_ergm(model, get_N_t(N,t), parVecSeq_T[:, t], nSample)
    end

    return A_T_dgp
end

#endregion

include("./StaticNets_models/StaticNets_DirBin1.jl")
include("./StaticNets_models/StaticNets_DirBin0Rec0.jl")
include("./StaticNets_models/StaticNets_ErgmPML.jl")
include("./StaticNets_models/StaticNets_conf_intervals.jl")

end


mutable struct SortingDirectJumpAggregation{T,S,F1,F2,RNG,DEPGR} <: AbstractSSAJumpAggregator
    next_jump::Int
    next_jump_time::T
    end_time::T
    cur_rates::Vector{T}
    sum_rate::T
    ma_jumps::S
    rates::F1
    affects!::F2
    save_positions::Tuple{Bool,Bool}
    rng::RNG
    dep_gr::DEPGR
    jump_search_order::Vector{Int}
    jump_search_idx::Int
    SortingDirectJumpAggregation{T,S,F1,F2,RNG}(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S,
                                                rs::F1, affs!::F2, sps::Tuple{Bool,Bool}, rng::RNG,
                                                dep_gr::DEPGR, jtoidx::Vector{Int}) where {T,S,F1,F2,RNG,DEPGR} =
      new{T,S,F1,F2,RNG,DEPGR}(nj, njt, et, crs, sr, maj, rs, affs!, sps, rng, dep_gr, jtoidx, zero(Int))
  end

function SortingDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T,
                                      maj::S, rs::F1, affs!::F2, sps::Tuple{Bool,Bool},
                                      rng::RNG; dep_graph=nothing, kwargs...) where {T,S,F1,F2,RNG}

    # a dependency graph is needed and must be provided if there are constant rate jumps
    if dep_graph == nothing
        if (get_num_majumps(maj) == 0) || !isempty(rs)
            error("To use ConstantRateJumps with the SortingDirect algorithm a dependency graph must be supplied.")
        else
            dg = make_dependency_graph(get_num_majumps(maj), maj)
        end
    else
        dg = dep_graph
    end

    # map jump idx to idx in cur_rates
    jtoidx = collect(1:length(crs))
    SortingDirectJumpAggregation{T,S,F1,F2,RNG}(nj, njt, et, crs, sr, maj, rs, affs!, sps, rng, dg, jtoidx)
end

  ########### The following routines should be templates for all SSAs ###########

# condition for jump to occur
@inline function (p::SortingDirectJumpAggregation)(u, t, integrator)
    p.next_jump_time == t
end

# executing jump at the next jump time
function (p::SortingDirectJumpAggregation)(integrator)
    execute_jumps!(p, integrator, integrator.u, integrator.p, integrator.t)
    generate_jumps!(p, integrator, integrator.u, integrator.p, integrator.t)
    register_next_jump_time!(integrator, p, integrator.t)
    nothing
end

# setting up a new simulation
function (p::SortingDirectJumpAggregation)(dj, u, t, integrator) # initialize
    initialize!(p, integrator, u, integrator.p, t)
    register_next_jump_time!(integrator, p, t)
    nothing
end

  # creating the JumpAggregation structure (function wrapper-based constant jumps)
function aggregate(aggregator::SortingDirect, u, p, t, end_time, constant_jumps,
                   ma_jumps, save_positions, rng; kwargs...)

    # handle constant jumps using function wrappers
    rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps)

    build_jump_aggregation(SortingDirectJumpAggregation, u, p, t, end_time, ma_jumps,
                           rates, affects!, save_positions, rng; kwargs...)
end

# set up a new simulation and calculate the first jump / jump time
function initialize!(p::SortingDirectJumpAggregation, integrator, u, params, t)
    fill_rates_and_sum!(p, u, params, t)
    generate_jumps!(p, integrator, u, params, t)
    nothing
end

# execute one jump, changing the system state
function execute_jumps!(p::SortingDirectJumpAggregation, integrator, u, params, t)
    # execute jump
    num_ma_rates = get_num_majumps(p.ma_jumps)
    if p.next_jump <= num_ma_rates
        @inbounds executerx!(u, p.next_jump, p.ma_jumps)
    else
        idx = p.next_jump - num_ma_rates
        @inbounds p.affects![idx](integrator)
    end

    # update search order
    jso   = p.jump_search_order
    jsidx = p.jump_search_idx
    if jsidx != 1
        @inbounds tmp          = jso[jsidx]
        @inbounds jso[jsidx]   = jso[jsidx-1]
        @inbounds jso[jsidx-1] = tmp
    end

    # update current jump rates
    update_dependent_rates!(p, u, params, t)
    nothing
end

# calculate the next jump / jump time
function generate_jumps!(p::SortingDirectJumpAggregation, integrator, u, params, t)
    @fastmath p.next_jump_time = t + calc_next_jump!(p, u, params, t)
    nothing
end

######################## SSA specific helper routines ########################

# recalculate jump rates for jumps that depend on the just executed jump (p.next_jump)
function update_dependent_rates!(p, u, params, t)
    @inbounds dep_rxs = p.dep_gr[p.next_jump]
    num_majumps = get_num_majumps(p.ma_jumps)
    cur_rates   = p.cur_rates
    sum_rate    = p.sum_rate
    majumps     = p.ma_jumps
    @inbounds for rx in dep_rxs
        sum_rate -= cur_rates[rx]
        if rx <= num_majumps
            @inbounds cur_rates[rx] = evalrxrate(u, rx, majumps)
        else
            @inbounds cur_rates[rx] = p.rates[rx-num_majumps](u, params, t)
        end
        sum_rate += cur_rates[rx]
    end

    p.sum_rate = sum_rate
end

function fill_rates_and_sum!(p, u, params, t)
    sum_rate = zero(typeof(p.sum_rate))

    # mass action jumps
    majumps   = p.ma_jumps
    cur_rates = p.cur_rates
    @inbounds for i in 1:get_num_majumps(majumps)
        cur_rates[i] = evalrxrate(u, i, majumps)
        sum_rate    += cur_rates[i]
    end

    # constant rates (IGNORED)
    rates = p.rates
    idx   = get_num_majumps(majumps) + 1
    @inbounds for i in eachindex(rates)
        cur_rates[idx] = rates[i](u, params, t)
        sum_rate += cur_rates[idx]
        idx += 1
    end

    p.sum_rate = sum_rate
end

# searches down the rate list for the next reaction
@fastmath function calc_next_jump!(p, u, params, t)

    # next jump type
    cur_rates = p.cur_rates
    numjumps  = length(cur_rates)
    jso       = p.jump_search_order
    rn        = p.sum_rate * rand(p.rng)
    @inbounds for idx = 1:numjumps
        rn -= cur_rates[jso[idx]]
        if rn < zero(rn)
            p.jump_search_idx = idx
            break
        end
    end
    @inbounds p.next_jump = jso[p.jump_search_idx]

    # return time to next jump
    randexp(p.rng) / p.sum_rate

end


module SBMLToolkit

using Catalyst
using SBML
using SymbolicUtils

include("reactionsystem.jl")

export ReactionSystem, ODESystem
export readSBML, set_level_and_version, convert_simplify_math

end


#=
Loads a single image from a filename and converts to the expected 256x256 size
=#
function load_single_image(filename::String,imgs,ii,x_off,y_off)

    this_img=load(filename)

    x_off[ii] = 256 / size(this_img,2)
    y_off[ii] = 256 / size(this_img,1)

    new_img=imresize(this_img,(256,256))
    for i=1:size(imgs,2)
        for j=1:size(imgs,1)
            imgs[i,j,1,ii]=convert(Float32,new_img[j,i].r)
            imgs[i,j,2,ii]=convert(Float32,new_img[j,i].g)
            imgs[i,j,3,ii]=convert(Float32,new_img[j,i].b)
        end
    end

    nothing
end

function make_all_labels(ll,xo,yo,j_a)
    for i=1:size(ll,4)
        add_single_label(ll,i,xo,yo,j_a);
    end
end

function add_single_label(ll,ii,x_off,y_off,j_a,ds=4)

    for i=1:size(ll,3)
        xo=round(Int64,x_off[ii] * j_a[1,i,ii]/ds)
        yo=round(Int64,y_off[ii] * j_a[2,i,ii]/ds)

        ll[:,:,i,ii] = gaussian_2d(1.0:1.0:size(ll,1),1.0:1.0:size(ll,2),xo,yo)

        if maximum(ll[:,:,i,ii])>0
            ll[:,:,i,ii] = ll[:,:,i,ii] ./ maximum(ll[:,:,i,ii])
        end

    end
end

#=
Saving Hourglass Model
=#
function save_hourglass(name,x)
    count=[1];
    file=matopen(name,"w")
        save_nn(x,file,count)
    close(file)
end

function save_nn(x,file,count)

    for f in fieldnames(typeof(x))
        f1=getfield(x,f)
        if typeof(f1) <: NN
            save_nn(f1,file,count)
        elseif eltype(f1) <: NN
            for i=1:length(f1)
                save_nn(f1[i],file,count)
            end
        else
            if f == :w
                write(file,string("w_",count[1]),convert(Array,f1.value))
                count[1]+=1
            elseif f == :b
                write(file,string("b_",count[1]),convert(Array,f1.value))
                count[1]+=1
            elseif f == :ms
                write(file,string("ms_mo_",count[1]),f1.momentum)
                write(file,string("ms_mean_",count[1]),convert(Array,f1.mean))
                write(file,string("ms_var_",count[1]),convert(Array,f1.var))
                count[1]+=1
            elseif f == :bn_p
                write(file,string("bn_p_",count[1]),convert(Array,f1.value))
                count[1]+=1
            end
        end
    end
end

#=
Loading Hourglass Model
=#
function load_hourglass(name,x,atype=KnetArray)
    count=[1];
    file=matopen(name,"r")
        load_nn(x,file,count,atype)
    close(file)
end

function load_nn(x,file,count,atype=KnetArray)

    for f in fieldnames(typeof(x))
        f1=getfield(x,f)
        if typeof(f1) <: NN
            load_nn(f1,file,count,atype)
        elseif eltype(f1) <: NN
            for i=1:length(f1)
                load_nn(f1[i],file,count,atype)
            end
        else
            if f == :w
                setfield!(x,f,Param(convert(atype,read(file,string("w_",count[1])))))
                count[1]+=1
            elseif f == :b
                xx = read(file,string("b_",count[1]))
                if typeof(xx) == Float32
                    xxx = zeros(Float32,1,1,1,1)
                    xxx[1] = xx
                    setfield!(x,f,Param(convert(atype,xxx)))
                else
                    setfield!(x,f,Param(convert(atype,xx)))
                end

                count[1]+=1
            elseif f == :ms
                x.ms.momentum=read(file,string("ms_mo_",count[1]))

                xx=read(file,string("ms_mean_",count[1]))
                if typeof(xx) == Float32
                    x.ms.mean = convert(atype,[xx])
                else
                    x.ms.mean = convert(atype,xx)
                end

                xx=read(file,string("ms_var_",count[1]))
                if typeof(xx) == Float32
                    x.ms.var = convert(atype,[xx])
                else
                    x.ms.var = convert(atype,xx)
                end

                count[1]+=1
            elseif f == :bn_p
                xx = read(file,string("bn_p_",count[1]))
                if typeof(xx) == Float32
                    x.bn_p = Param(convert(atype,[xx]))
                else
                    x.bn_p = Param(convert(atype,xx))
                end
                count[1]+=1
            end
        end
    end
end

#=
Change the number of input or output features or dimensions for hourglass model
=#
function change_hourglass(hg::HG2,feature_num::Int,input_dim::Int,output_dim::Int,atype=KnetArray)

    change_hourglass_input(hg,feature_num,input_dim,atype)
    change_hourglass_output(hg,feature_num,output_dim,atype)

    nothing
end

function change_hourglass_input(hg::HG2,feature_num::Int,input_dim::Int,atype=KnetArray)

    #Input transform
    hg.fb.c1.w = Param(convert(atype,xavier_normal(Float32,7,7,input_dim,64)))
    hg.fb.c1.b = Param(convert(atype,xavier_normal(Float32,1,1,64,1)))
    hg.fb.c1.bn_p = Param(convert(atype{Float32,1},bnparams(input_dim)))
    hg.fb.c1.ms = bnmoments()

    nothing
end

function change_hourglass_output(hg::HG2,feature_num::Int,output_dim::Int,atype=KnetArray)

    for i=1:length(hg.c1)
        hg.c1[i].w = Param(convert(atype,xavier_normal(Float32,1,1,feature_num,output_dim)))
        hg.c1[i].b = Param(convert(atype,xavier_normal(Float32,1,1,output_dim,1)))
    end

    for i=1:length(hg.merge_preds)
        hg.merge_preds[i].w = Param(convert(atype,xavier_normal(Float32,1,1,output_dim,feature_num)))
        hg.merge_preds[i].b = Param(convert(atype,xavier_normal(Float32,1,1,feature_num,1)))
    end

    nothing
end


function removeBubble(bubble, lc, rc)
    Lver, Rver = bubble.lver, bubble.rver
    para = bubble.parent.para
    chan = bubble.channel

    if NoBubble in para.filter
        if Lver.para.innerLoopNum == 0 && Rver.para.innerLoopNum == 0
            if chan == T || chan == U
                if lc == DI && rc == DI
                    return true
                end
            end
        end
    end

    return false
end

function notProper(para, K)
    if Proper in para.filter
        transferLoop = para.transferLoop
        @assert isempty(transferLoop) == false "Please initialize para.transferLoop to check proper diagrams."
        if transferLoop ≈ K
            return true
        end
    end
    return false
end

# check if G exist without creating objects in the pool
function isValidG(filter, innerLoopNum::Int)
    if (NoFock in filter) && (innerLoopNum == 1)
        return false
    end

    if (Girreducible in filter) && (innerLoopNum > 0)
        return false
    end

    return true
end

function isValidG(para::GenericPara)
    @assert para.diagType == GreenDiag
    return isValidG(para.filter, para.innerLoopNum)
end

function isValidSigma(filter, innerLoopNum::Int, subdiagram::Bool)
    @assert innerLoopNum >= 0
    if innerLoopNum == 0
        return false
    end
    if subdiagram && (Girreducible in filter)
        return false
    end

    if subdiagram && (NoFock in filter) && innerLoopNum == 1
        return false
    end

    return true
end

function isValidPolarization(filter, innerLoopNum::Int, subdiagram::Bool)
    @assert innerLoopNum >= 0
    if innerLoopNum == 0
        return false
    end
    if subdiagram && (Wirreducible in filter)
        return false
    end
    if subdiagram && (NoBubble in filer) && innerLoopNum == 1
        return false
    end
    return true
end


using DrWatson
@quickactivate
using HierarchicalAD
using FileIO, BSON
using StatsBase
using ProgressMeter
using Colors

N_samples = 100000

font_size = 24
img_size = (28, 28)
fonts = ["Georgia", "", "Arial", "Ubuntu Mono", "Free Mono", "Eufm10", "Lobster Two", "Purisa", "Manjari Thin", "Chilanka", "Uroob", "MathJax_Typewriter"]
accents = ["", " Bold"]
digs = 0:9
cols = collect(keys(Colors.color_names))
rots = [-20, -10, -5, -3, 0, 3, 5, 10, 20]
shearxs = -20:20
shearys = -20:20
zooms =0.9:0.01:1.2

sample_input() = map(x->sample(x,1)[1], (fonts, accents, digs, cols, rots, shearxs, shearys, zooms))

inputs = map(i->sample_input(), 1:N_samples)
inputf = datadir("digits_rgb/labels.bson")
save(inputf, Dict(:labels => inputs))

@info "Generating digits..."
outputs = @showprogress map(i->HierarchicalAD.create_digit_rgb(font_size, img_size, i...), inputs)
outputf = datadir("digits_rgb/digits.bson")
save(outputf, Dict(:digits=> outputs))


using DrWatson, PyCall
@quickactivate "GCN_HM_GRN-Integration"

pushfirst!(PyVector(pyimport("sys")."path"), "");
run_sim = pyimport("large_batch").run_sim

allparams = Dict(
    :layer => ["arma", "sage", "tag"],
    :batch => [100, 200],
    :cl => ["E116"]
)

dicts = dict_list(allparams)

function makesim(d::Dict)
    @unpack layer, batch, cl = d
    best_auc, losses_test, losses_train, auc_l = run_sim(cl, batch, layer)
    fulld = copy(d)
    fulld[:auc] = best_auc
    fulld[:losses_test] = losses_test
    fulld[:losses_train] = losses_train
    fulld[:auc_list] = auc_l
    return fulld
end

for (i, d) in enumerate(dicts)
    f = makesim(d)
    print("Finished $(i)")
    @tagsave(datadir("large_batch", savename(d, "bson")), f; safe = true)
end


mutable struct DynamicGenerator{
    M <: Machine,
    S <: Shaft,
    A <: AVR,
    TG <: TurbineGov,
    P <: PSS,
} <: DynamicInjection
    number::Int64
    name::String
    bus::Bus
    ω_ref::Float64
    V_ref::Float64
    P_ref::Float64
    Q_ref::Float64
    machine::M
    shaft::S
    avr::A
    tg::TG
    pss::P
    n_states::Int64
    states::Vector{Symbol}
    ext::Dict{String, Any}
    function DynamicGenerator(
        number::Int64,
        name::String,
        bus::Bus,
        ω_ref::Float64,
        V_ref::Float64,
        P_ref::Float64,
        Q_ref::Float64,
        machine::M,
        shaft::S,
        avr::A,
        tg::TG,
        pss::P,
    ) where {M <: Machine, S <: Shaft, A <: AVR, TG <: TurbineGov, P <: PSS}

        n_states =
            (machine.n_states + shaft.n_states + avr.n_states + tg.n_states + pss.n_states)

        states = vcat(machine.states, shaft.states, avr.states, tg.states, pss.states)

        new{M, S, A, TG, P}(
            number,
            name,
            bus,
            ω_ref,
            V_ref,
            P_ref,
            Q_ref,
            machine,
            shaft,
            avr,
            tg,
            pss,
            n_states,
            states,
            Dict{String, Any}(),
        )
    end
end

get_Sbase(device::DynamicGenerator) = device.machine.s_rated
get_bus(device::DynamicGenerator) = device.bus
get_states(device::DynamicGenerator) = device.states
get_n_states(device::DynamicGenerator) = device.n_states
get_number(device::DynamicGenerator) = device.number
get_name(device::DynamicGenerator) = device.name
get_ω_ref(device::DynamicGenerator) = device.ω_ref
get_V_ref(device::DynamicGenerator) = device.V_ref
get_P_ref(device::DynamicGenerator) = device.P_ref
get_Q_ref(device::DynamicGenerator) = device.Q_ref
get_machine(device::DynamicGenerator) = device.machine
get_shaft(device::DynamicGenerator) = device.shaft
get_avr(device::DynamicGenerator) = device.avr
get_tg(device::DynamicGenerator) = device.tg
get_pss(device::DynamicGenerator) = device.pss
get_ext(device::DynamicGenerator) = device.ext


"""
```
shock_decompositions(m, system, histshocks)

shock_decompositions(system, forecast_horizons, histshocks, start_index,
    end_index)

shock_decompositions(m, system, histshocks, start_date, end_date)

shock_decompositions(system, forecast_horizons, histshocks, start_index,
    end_index, regime_inds)
```

### Inputs

- `system::System{S}` or `RegimeSwitchingSystem{S}`: state-space system matrices
- `forecast_horizons::Int`: number of periods ahead to forecast
- `histshocks::Matrix{S}`: matrix of size `nshocks` x `hist_periods` of
  historical smoothed shocks
- `start_index::Int`: first index from which to return computed shock
  decompositions
- `end_index::Int`: last index for which to return computed shock decompositions
- `start_date::Date`: initial date for deterministic trends
- `end_date::Date`: final date for deterministic trends
- `regime_inds::Vector{UnitRange}`: indices of the data corresponding to each regime.

where `S<:AbstractFloat`.

### Outputs

- `states::Array{S, 3}`: matrix of size `nstates` x `nperiods` x `nshocks` of
  state shock decompositions
- `obs::Array{S, 3}`: matrix of size `nobs` x `nperiods` x `nshocks` of
  observable shock decompositions
- `pseudo::Array{S, 3}`: matrix of size `npseudo` x `nperiods` x `nshocks` of
  pseudo-observable shock decompositions

where `nperiods = `end_index - start_index + 1`.
"""
function shock_decompositions(m::AbstractDSGEModel,
    system::System{S}, histshocks::Matrix{S}) where {S<:AbstractFloat}

    horizon   = forecast_horizons(m)
    start_ind = index_shockdec_start(m)
    end_ind   = index_shockdec_end(m)

    shock_decompositions(system, horizon, histshocks, start_ind, end_ind)
end

function shock_decompositions(system::System{S},
    forecast_horizons::Int, histshocks::Matrix{S},
    start_index::Int, end_index::Int) where {S<:AbstractFloat}

    # Setup
    nshocks      = size(system[:RRR], 2)
    nstates      = size(system[:TTT], 2)
    nobs         = size(system[:ZZ], 1)
    npseudo      = size(system[:ZZ_pseudo], 1)
    histperiods  = size(histshocks, 2)
    allperiods   = histperiods + forecast_horizons

    states = zeros(S, nstates, allperiods, nshocks)
    obs    = zeros(S, nobs, allperiods, nshocks)
    pseudo = zeros(S, npseudo, allperiods, nshocks)

    # Check dates
    if forecast_horizons <= 0 || start_index < 1 || end_index > allperiods
        throw(DomainError("At least one of forecast_horizons, the index for the start date of the shock decomposition, and the index for the end date of the shock decomposition is outside the permissible range."))
    end

    # Set constant system matrices to 0
    system = zero_system_constants(system)

    z0 = zeros(S, nstates)
    for i = 1:nshocks
        # Isolate single shock
        shocks = zeros(S, nshocks, allperiods)
        shocks[i, 1:histperiods] = histshocks[i, :]

        # Iterate state space forward
        states[:, :, i], obs[:, :, i], pseudo[:, :, i], _ = forecast(system, z0, shocks)
    end

    # Return shock decompositions in appropriate range
    if start_index == 1 && end_index == allperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range, :], obs[:, range, :], pseudo[:, range, :]
    end
end

function shock_decompositions(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, histshocks::Matrix{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m))) where {S<:AbstractFloat}
    # Set up dates
    horizon     = forecast_horizons(m)
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Set up regime indices. Note that we do not need to account
    # for ZLB split b/c `histshocks` should have zeros for anticipated shocks
    # in the pre-ZLB periods.
    regime_inds = regime_indices(m, start_date, end_date)
    if regime_inds[1][1] < 1 # remove periods occuring before desired start date
        regime_inds[1] = 1:regime_inds[1][end]
    end

    shock_decompositions(system, horizon, histshocks, start_index, end_index, regime_inds)
end

function shock_decompositions(system::RegimeSwitchingSystem{S},
                              forecast_horizons::Int, histshocks::Matrix{S},
                              start_index::Int, end_index::Int,
                              regime_inds::Vector{UnitRange{Int}}) where {S<:AbstractFloat}

    # Setup
    nshocks     = size(system[1, :RRR], 2)
    nstates     = size(system[1, :TTT], 2)
    nobs        = size(system[1, :ZZ], 1)
    npseudo     = size(system[1, :ZZ_pseudo], 1)
    histperiods = size(histshocks, 2)
    allperiods  = histperiods + forecast_horizons

    states = zeros(S, nstates, allperiods, nshocks)
    obs    = zeros(S, nobs,    allperiods, nshocks)
    pseudo = zeros(S, npseudo, allperiods, nshocks)
    shocks = zeros(S, nshocks, histperiods)

    # Check dates
    if forecast_horizons <= 0 || start_index < 1 || end_index > allperiods
        throw(DomainError("At least one of forecast_horizons, the index for the start date of the shock decomposition, and the index for the end date of the shock decomposition is outside the permissible range."))
    end

    # Set constant system matrices to 0 and start at stationary steady state
    system = zero_system_constants(system)
    z0     = zeros(S, nstates)

    fcast_shocks = zeros(S, nshocks, forecast_horizons) # these are always zero
    for i = 1:nshocks
        # Isolate single shock
        shocks[i, :] = histshocks[i, :]
        # Use forecast to iterate state-space system forward without shocks or the ZLB procedure
        for (reg_num, reg_ind) in enumerate(regime_inds)
            init_state = (reg_num == 1) ? zeros(S, nstates) : states[:, reg_ind[1] - 1, i] # update initial state

            states[:, reg_ind, i], obs[:, reg_ind, i], pseudo[:, reg_ind, i], _ = forecast(system[reg_num], init_state, shocks[:, reg_ind])
        end

        fcast_inds = (histperiods + 1):allperiods
        states[:, fcast_inds, i], obs[:, fcast_inds, i], pseudo[:, fcast_inds, i], _ =
            forecast(system[length(regime_inds)], states[:, regime_inds[end][end], i], fcast_shocks)

        # zero out shocks
        shocks[i, :] .= 0.
    end

    # Return shock decompositions in appropriate range
    if start_index == 1 && end_index == allperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range, :], obs[:, range, :], pseudo[:, range, :]
    end
end

"""
```
deterministic_trends(m, system, z0)

deterministic_trends(system, z0, nperiods, start_index, end_index)

deterministic_trends(m, system, z0, start_date, end_date)

deterministic_trends(system, z0, nperiods, regime_inds, regimes)
```

Compute deterministic trend values of states, observables, and
pseudo-observables, given a model object and system matrices. The deterministic
trend for a single draw is simply the series that would be obtained by iterating
the state-space system forward, beginning from a state vector `z0` in the last
presample period.

### Inputs

- `m::AbstractDSGEModel`: model object
- `system::System{S}` or `RegimeSwitchingSystem`: state-space system matrices
- `z0`::Vector{S}: initial state vector
- `start_date::Date`: initial date for deterministic trends
- `end_date::Date`: final date for deterministic trends
- `regime_inds::Vector{UnitRange}`: indices of the data corresponding to each regime.
- `regimes::UnitRange`: which regimes are involved in the date range for which
    we want to compute the deterministic trends

where `S<:AbstractFloat`.

### Outputs

- `states::Matrix{S}`: matrix of size `nstates` x `nperiods` of state
  steady-state values
- `obs::Matrix{S}`: matrix of size `nobs` x `nperiods` of observable
  steady-state values
- `pseudo::Matrix{S}`: matrix of size `npseudo` x `nperiods` of
  pseudo-observable steady-state values

where `nperiods` is the number of quarters between
`date_shockdec_start(m)` and `date_shockdec_end(m)`, inclusive.
"""
function deterministic_trends(m::AbstractDSGEModel{S},
                              system::System{S}, z0::Vector{S}) where {S<:AbstractFloat}

    # Dates: We compute the deterministic trend starting from the
    # first historical period.  However, since it is only used to
    # compute shock decompositions, we truncate and only store
    # results for periods corresponding to the shockdec period.
    nperiods    = subtract_quarters(date_forecast_end(m), date_mainsample_start(m)) + 1
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    deterministic_trends(system, z0, nperiods, start_index, end_index)
end

function deterministic_trends(system::System{S}, z0::Vector{S}, nperiods::Int,
                              start_index::Int, end_index::Int) where {S<:AbstractFloat}

    # Set constant system matrices to 0
    system = zero_system_constants(system)

    # Construct matrix of 0 shocks for entire history and forecast horizon
    nshocks  = size(system[:RRR], 2)
    shocks   = zeros(S, nshocks, nperiods)

    # Use forecast to iterate state-space system forward without shocks or the ZLB procedure
    states, obs, pseudo, _ = forecast(system, z0, shocks)

    # Return deterministic trends in appropriate range
    if start_index == 1 && end_index == nperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range], obs[:, range], pseudo[:, range]
    end
end

function deterministic_trends(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, z0::Vector{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m))) where {S<:AbstractFloat}

    # Dates: We compute the deterministic trend starting from the
    # first historical period.  However, since it is only used to
    # compute shock decompositions, we truncate and only store
    # results for periods corresponding to the shockdec period.
    nperiods    = subtract_quarters(date_forecast_end(m), date_mainsample_start(m)) + 1
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Now intersect these periods with regime indices
    regime_inds = regime_indices(m, start_date, end_date) # do not need to account for ZLB split b/c shocks don't matter for trends
    if regime_inds[1][1] < 1
        regime_inds[1] = 1:regime_inds[1][end]
    end
    regime_inds[end] = regime_inds[end][1]:end_index      # if the end index is in the middle of a regime or is past the regime's end

    return deterministic_trends(system, z0, nperiods, start_index, end_index, regime_inds)
end

function deterministic_trends(system::RegimeSwitchingSystem{S}, z0::Vector{S}, nperiods::Int,
                              start_index::Int, end_index::Int,
                              regime_inds::Vector{UnitRange{Int}}) where {S<:AbstractFloat}

    # Set constant system matrices to 0
    system = zero_system_constants(system)

    # Initialize storage matrices
    nshocks = size(system[1, :RRR], 2) # this is used in the forecasting loops to create shock matrices
    nstates = size(system[1, :TTT], 1)
    nobs    = size(system[1, :ZZ], 1)
    npseudo = size(system[1, :ZZ_pseudo], 1)
    states = Matrix{S}(undef, nstates, nperiods)
    obs    = Matrix{S}(undef, nobs,    nperiods)
    pseudo = Matrix{S}(undef, npseudo, nperiods)

    # Use forecast to iterate state-space system forward without shocks or the ZLB procedure
    for (reg_num, reg_ind) in enumerate(regime_inds)
        init_state = (reg_num == 1) ? z0 : states[:, reg_ind[1] - 1] # update initial state

        # Construct matrix of 0 shocks for this regime
        shocks = zeros(S, nshocks, length(reg_ind))
        states[:, reg_ind], obs[:, reg_ind], pseudo[:, reg_ind], _ = forecast(system[reg_num], init_state, shocks)
    end
    if regime_inds[end][end] < nperiods
        shocks   = zeros(S, nshocks, nperiods - regime_inds[end][end])
        last_ind = (1 + regime_inds[end][end]):nperiods
        states[:, last_ind], obs[:, last_ind], pseudo[:, last_ind], _ = forecast(system[length(regime_inds)], states[:, regime_inds[end][end]], shocks)
    end

    if start_index == 1 && end_index == nperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range], obs[:, range], pseudo[:, range]
    end
end

"""
```
trends(system::System{S}) where {S<:AbstractFloat}
trends(system::RegimeSwitchingSystem{S}) where {S<:AbstractFloat}
```

Compute trend (steady-state) states, observables, and pseudo-observables. The
trend is used for plotting shock decompositions.
"""
function trends(system::System{S}) where {S<:AbstractFloat}

    state_trend  = system[:CCC]
    obs_trend    = system[:ZZ]*system[:CCC] + system[:DD]
    pseudo_trend = system[:ZZ_pseudo]*system[:CCC] + system[:DD_pseudo]

    return state_trend, obs_trend, pseudo_trend
end

function trends(system::RegimeSwitchingSystem{S}) where {S<:AbstractFloat}

    nreg  = n_regimes(system)
    input = 1:nreg

    state_trends  = Matrix{S}(undef, length(system[1, :CCC]), nreg)
    obs_trends    = Matrix{S}(undef, length(system[1, :DD]), nreg)
    pseudo_trends = Matrix{S}(undef, length(system[1, :DD_pseudo]), nreg)
    for i in input
        state_trends[:, i]  = system[i, :CCC]
        obs_trends[:, i]    = system[i, :ZZ] * system[i, :CCC] + system[i, :DD]
        pseudo_trends[:, i] = system[i, :ZZ_pseudo] * system[i, :CCC] + system[i, :DD_pseudo]
    end

    return state_trends, obs_trends, pseudo_trends
end


using Random
using Distributions
using Statistics
using SparseArrays
using MinimalRLCore

mutable struct TabularTMaze <: MinimalRLCore.AbstractEnvironment
    world::Array{Array{String}}
    start_state::Array{Int64}
    current_state::Array{Int64}
    goal_states::Array{String}
    feature_size::Int64
    step::Int
    exploring_starts::Bool
    cumulant_schedule::CumulantSchedule
    extrinsic_reward::Float64

    function TabularTMaze(exploring_starts, cumulant_schedule; extrinsic_reward = -0.01)
        feature_size = 21
        num_actions = 4
        step = 1

        world = [["G1", "0", "0", "0", "0", "0", "G3"],
                 ["1", "0", "0", "0", "0", "0", "1"],
                 ["1", "1", "1", "1", "1", "1", "1"],
                 ["1", "0", "0", "1", "0", "0", "1"],
                 ["G2", "0", "0", "1", "0", "0", "G4"],
                 ["0", "0", "0", "1", "0", "0", "0"],
                 ["0", "0", "0", "1", "0", "0", "0"],
                 ["0", "0", "0", "1", "0", "0", "0"],
                 ["0", "0", "0", "1", "0", "0", "0"]]
        goal_states = ["G1", "G2", "G3", "G4"]
        start_state = [9,4]
        new(world, start_state, [9,4], goal_states, feature_size,step, exploring_starts, cumulant_schedule, extrinsic_reward)
    end
end

Base.size(e::TabularTMaze) = e.feature_size

function generate_obs(state::Array{Int64})
    obs = zeros(Int64, 9, 7)
    obs[state[1],state[2]] = 1
    toReturn = findfirst(!iszero, obs[valid_state_mask()])
    return toReturn
end

function valid_state_mask()
    """
    return a mask of valid states that is 9x7
    """
    world = [["G1", "0", "0", "0", "0", "0", "G3"],
             ["1", "0", "0", "0", "0", "0", "1"],
             ["1", "1", "1", "1", "1", "1", "1"],
             ["1", "0", "0", "1", "0", "0", "1"],
             ["G2", "0", "0", "1", "0", "0", "G4"],
             ["0", "0", "0", "1", "0", "0", "0"],
             ["0", "0", "0", "1", "0", "0", "0"],
             ["0", "0", "0", "1", "0", "0", "0"],
             ["0", "0", "0", "1", "0", "0", "0"]]

    world = permutedims(hcat(world...))
    valid_states = findall(x-> x!="0", world)
    return valid_states
end

function valid_state_action_mask()
"""
returns a mask for the state-action mask (63,4)
"""
    num_actions = 4
    world = [["G1", "0", "0", "0", "0", "0", "G3"],
                      ["1", "0", "0", "0", "0", "0", "1"],
                      ["1", "1", "1", "1", "1", "1", "1"],
                      ["1", "0", "0", "1", "0", "0", "1"],
                      ["G2", "0", "0", "1", "0", "0", "G4"],
                      ["0", "0", "0", "1", "0", "0", "0"],
                      ["0", "0", "0", "1", "0", "0", "0"],
                      ["0", "0", "0", "1", "0", "0", "0"],
                      ["0", "0", "0", "1", "0", "0", "0"]]

    flattened_world = reshape(permutedims(hcat(world...)),(63))
    valid_states = findall(x-> x!="0", flattened_world)
    return valid_states
end

MinimalRLCore.get_reward(env::TabularTMaze) = env.extrinsic_reward
is_terminal(env::TabularTMaze, pos) = env.world[pos[1]][pos[2]] in env.goal_states
MinimalRLCore.is_terminal(env::TabularTMaze) = is_terminal(env, env.current_state)
MinimalRLCore.get_actions(env::TabularTMaze) = 1:4

function MinimalRLCore.get_state(env::TabularTMaze)
    obs = generate_obs(env.current_state)
    cumulants = get_cumulants(env, env.cumulant_schedule, env.current_state)

    return vcat(obs,cumulants)
end

function get_cumulants(env::TabularTMaze, cs::CumulantSchedule, pos)
    num_cumulants = 4
    cumulants = zeros(num_cumulants)
    if env.world[pos[1]][pos[2]] == "G1"
        cumulants[1] = get_cumulant(cs, "G1")
    elseif env.world[pos[1]][pos[2]] == "G2"
        cumulants[2] = get_cumulant(cs, "G2")
    elseif env.world[pos[1]][pos[2]] == "G3"
        cumulants[3] = get_cumulant(cs, "G3")
    elseif env.world[pos[1]][pos[2]] == "G4"
        cumulants[4] = get_cumulant(cs, "G4")
    end
    return cumulants
end

function MinimalRLCore.reset!(environment::TabularTMaze, rng::AbstractRNG=Random.GLOBAL_RNG)
    if environment.exploring_starts == false
        environment.current_state = environment.start_state
    elseif environment.exploring_starts == true
        possible_start_states = findall(x -> x == "1", hcat(environment.world...))
        start_state = possible_start_states[rand(rng, 1:length(possible_start_states))]
        # The start states are flipped since hcat "transposes" a list of list when converting it to a matrix
        environment.current_state = [start_state[2],start_state[1]]
    end
end

function MinimalRLCore.reset!(environment::TabularTMaze,start_state::Int)
    s = valid_state_mask()[start_state]
    environment.current_state = [s[1],s[2]]
end

function MinimalRLCore.reset!(environment::TabularTMaze, start_state::CartesianIndex)
    # throw("Implement env_start with a start_state")
    # environment.current_state = [start_state[2], start_state[1]]
    environment.current_state = [start_state[1], start_state[2]]
end

function MinimalRLCore.environment_step!(environment::TabularTMaze, action, rng::AbstractRNG=Random.GLOBAL_RNG)
    actions = [(-1, 0), (0, 1), (1, 0), (0, -1)] # up, right, down, left
    reward = 0.0
    terminal = false

    potential_state = [environment.current_state[1] + actions[action][1],
                       environment.current_state[2] + actions[action][2]]

    if potential_state[1] in [0, 10] || potential_state[2] in [0, 8] # No nothing, agent doesn't move
    elseif environment.world[potential_state[1]][potential_state[2]] == "1"
        environment.current_state = potential_state
    elseif is_terminal(environment, potential_state)
        environment.current_state = potential_state
        terminal = true
    end
    # update!(environment, environment.cumulant_schedule, environment.current_state)
    update!(environment.cumulant_schedule, environment.current_state)
end

using RecipesBase, Colors
@recipe function f(env::TabularTMaze)
    ticks := nothing
    foreground_color_border := nothing
    grid := false
    legend := false
    aspect_ratio := 1
    xaxis := false
    yaxis := false
    yflip := false

    SIZE=20
    BG = Colors.RGB(1.0, 1.0, 1.0)
    BORDER = Colors.RGB(0.0, 0.0, 0.0)
    WALL = Colors.RGB(0.3, 0.3, 0.3)
    AC = Colors.RGB(0.69921875, 0.10546875, 0.10546875)
    GOAL = Colors.RGB(0.796875, 0.984375, 0.76953125)

    cell = fill(BG, SIZE, SIZE)
    cell[1, :] .= BORDER
    cell[end, :] .= BORDER
    cell[:, 1] .= BORDER
    cell[:, end] .= BORDER

    s_y = length(env.world)
    s_x = length(env.world[1])

    screen = fill(BG, (s_y + 2)*SIZE, (s_x + 2)*SIZE)

    screen[:, 1:SIZE] .= WALL
    screen[1:SIZE, :] .= WALL
    screen[end-(SIZE-1):end, :] .= WALL
    screen[:, end-(SIZE-1):end] .= WALL

    for j ∈ 1:s_x
        for i ∈ 1:s_y
            sqr_i = ((i)*SIZE + 1):((i+1)*SIZE)
            sqr_j = ((j)*SIZE + 1):((j+1)*SIZE)
            if env.current_state[1] == i && env.current_state[2] == j
                v = @view screen[sqr_i, sqr_j]
                v .= cell
                v[Int(SIZE/2)-4:Int(SIZE/2)+5, Int(SIZE/2)-4:Int(SIZE/2)+5] .= AC
            elseif env.world[i][j] == "0"
                screen[sqr_i, sqr_j] .= WALL
            elseif env.world[i][j] == "1"
                screen[sqr_i, sqr_j] .= cell
            elseif env.world[i][j][1] == 'G'
                v = @view screen[sqr_i, sqr_j]
                v .= GOAL
                v[1, :] .= BORDER
                v[:, 1] .= BORDER
                v[end, :] .= BORDER
                v[:, end] .= BORDER
            end
        end
    end
    screen[end:-1:1,:]
end


## ------------  ParamInfo

export ParamInfo
# export set_value!

# function ParamInfo(pName :: Symbol, startIdx :: Integer,
#     valueV :: T1, lbV :: T2, ubV :: T2) where {T1 <: AbstractArray, T2 <: Real}

#     sz = size(valueV);
#     return ParamInfo{T1}(pName, startIdx, valueV, fill(lbV, sz), fill(ubV, sz));
# end

function Base.isapprox(p1 :: ParamInfo{F1}, p2 :: ParamInfo{F1};
    atol :: Real = 1e-8) where F1
    return isequal(p1.pName, p2.pName)  &&
        # isapprox(p1.valueV, p2.valueV; atol)  &&
        isapprox(p1.lbV, p2.lbV; atol)  &&
        isapprox(p1.ubV, p2.ubV; atol);
end

Base.isapprox(p1 :: ParamInfo{F1}, p2 :: ParamInfo{F2}) where {F1, F2} = false;

lb(pInfo :: ParamInfo{F1}) where F1 = pInfo.lbV;
ub(pInfo :: ParamInfo{F1}) where F1 = pInfo.ubV;
n_values(pInfo :: ParamInfo{F1}) where F1 = length(pInfo.lbV);
indices(pInfo :: ParamInfo{F1}) where F1 <: Real = pInfo.startIdx;
indices(pInfo :: ParamInfo{F1}) where F1 =
    pInfo.startIdx .+ (0 : (n_values(pInfo) - 1));

random_values(pInfo :: ParamInfo{F1}, rng :: AbstractRNG) where F1 =
    pInfo.lbV .+ (pInfo.ubV .- pInfo.lbV) .* rand(rng, size(pInfo.lbV));
random_values(pInfo :: ParamInfo{F1}, rng :: AbstractRNG) where F1 <: Real =
    pInfo.lbV .+ (pInfo.ubV .- pInfo.lbV) .* rand(rng);

# Reshape an input to match dimensions of ParamInfo.
reshape_vector(pInfo :: ParamInfo{F1}, v :: F1) where F1 <: Real = v;
reshape_vector(pInfo :: ParamInfo{F1}, v) where F1 <: AbstractArray{<:Real} =
    reshape(v, size(pInfo.lbV));

# function set_value!(pInfo :: ParamInfo{F1}, v :: Real) where F1 <: Real
#     pInfo.valueV = v;
# end

# function set_value!(pInfo :: ParamInfo{F1}, v :: AbstractArray{F2}) where {F1 <: Real, F2 <: Real}
#     pInfo.valueV = only(v);
# end

# function set_value!(pInfo :: ParamInfo{<: AbstractArray{F1}}, v :: AbstractArray{F2}) where {F1 <: Real, F2 <: Real}
#     pInfo.valueV = reshape(v, size(pInfo.lbV));
# end

# function value_from_vector!(pInfo :: ParamInfo{F1},
#     valueV :: AbstractVector{F1}) where F1 <: Real
#     pInfo.valueV = valueV[pInfo.startIdx];
# end

# function value_from_vector!(pInfo :: ParamInfo{<:AbstractArray{F1}},
#     valueV :: AbstractVector{F1}) where F1 <: Real
#     pInfo.valueV = reshape(valueV[indices(pInfo)], size(pInfo.lbV));
# end

function make_test_param_info(pName, sz; startIdx = 10)
    if isempty(sz)
        lbV = 0.3;
        ubV = 5.8;
        # valueV = 3.5;
    else
        lbV = fill(0.3, sz);
        ubV = fill(5.8, sz);
        # valueV = 0.3 .* lbV .+ 0.7 .* ubV;
    end
    ParamInfo(pName, startIdx, lbV, ubV);
end

## ------------  ValueVector

# ValueVector
# export ValueVector, set_values!, values, lb, ub

ValueVector{F1}() where F1 = ValueVector{F1}(OrderedDict{Symbol, ParamInfo}());

function make_test_value_vector(n)
    vv = ValueVector{ValueType}();
    startIdx = 1;
    for j = 1 : n
        pName = Symbol("p$j");
        if j == 1
            sz = ();
        elseif j == 2
            sz = (j, );
        else
            sz = (j, 2);
        end
        vv.d[pName] = make_test_param_info(pName, sz; startIdx);
        nElem = prod(sz);
        startIdx += nElem;
    end
    return vv
end

function validate_param_match(pvec :: ParamVector, vv :: ValueVector{F1}) where F1
    isValid = true;
    nCal, _ = n_calibrated_params(pvec);
    if length(vv) != nCal
        isValid = false;
        @warn """
            No of calibrated params differs for $pvec
            pvec: $(nCal)
            vv:   $(length(vv))
            """;
    end

    if isValid  &&  (length(vv) > 0)
        pList = calibrated_params(pvec);
        for p in pList
            pName = name(p);
            if haskey(vv.d, pName)
                pInfo = vv.d[pName];
                if size(calibrated_value(p)) != size(lb(pInfo))
                    isValid = false;
                    @warn "Size mismatch for $pName";
                end
            else
                isValid = false;
                @warn "Param $pName missing from ValueVector";
            end
        end
    end
    return isValid
end

function validate_vv(vv :: ValueVector{F1}) where F1
    isValid = true;
    isValid = isValid  &&  validate_indices(vv);
    return isValid
end

function validate_indices(vv :: ValueVector{F1}) where F1
    isValid = true;
    nValues = n_values(vv);
    if nValues > 0
        takenV = zeros(Int, nValues);
        for (_, pInfo) in vv.d
            # This can be a Range of scalar.
            idxV = indices(pInfo) .- start_index(vv) .+ 1;
            if idxV[1] < 1
                @warn "Indices below start index";
                isValid = false;
            end
            if last(idxV) > nValues
                @warn "Indices past last index";
                isValid = false;
            end
            # Must make this a range for broadcasting to work
            isValid  &&  (takenV[first(idxV) : last(idxV)] .+= 1);
        end
        isValid = all(takenV .== 1);
    end
    return isValid
end

Base.isempty(vv :: ValueVector{F1}) where F1 =
    n_values(vv) == 0;
Base.length(vv :: ValueVector{F1}) where F1 = length(vv.d);

"""
	$(SIGNATURES)

Total number of scalar values across all parameters.
"""
function n_values(vv :: ValueVector{F1}) where F1
    nValues = 0;
    for (_, pInfo) in vv.d
        nValues += n_values(pInfo);
    end
    return nValues
end

function start_index(vv :: ValueVector{F1}) where F1
    (_, pInfo) = first(vv.d);
    return pInfo.startIdx;
end

Base.last(vv :: ValueVector{F1}) where F1 = vv.d[last(collect(keys(vv.d)))];

function last_index(vv :: ValueVector{F1}) where F1
    pInfo = last(vv);
    return last(indices(pInfo))
end

# values(vv :: ValueVector{F1}) where F1 = vv.valueV;
# lb(vv :: ValueVector{F1}) where F1 = vv.lbV;
# ub(vv :: ValueVector{F1}) where F1 = vv.ubV;
# pnames(vv :: ValueVector{F1}) where F1 = vv.pNameV;
# Base.length(vv :: ValueVector{F1}) where F1 = Base.length(vv.valueV);

function Base.isapprox(vv1 :: ValueVector{F1}, vv2 :: ValueVector{F1};
    atol :: Real = 1e-8) where F1

    if length(vv1.d) != length(vv2.d)
        return false;
    end
    for (pName, pInfo1) in vv1.d
        if haskey(vv2.d, pName)
            pInfo2 = vv2.d[pName]
            if !isapprox(pInfo1, pInfo2; atol)
                return false
            end
        else
            return false
        end
    end
    return true
end

# """
# 	$(SIGNATURES)

# Set values from a Vector{Real}. These are usually the transformed values for the entire `ModelObject`.
# """
# function set_values!(vv :: ValueVector{F1}, valueV :: AbstractVector{F1}) where F1
#     @assert length(valueV) >= last_index(vv)  "Vector too short";
#     for (_, pInfo) in vv.d
#         set_value!(pInfo, valueV[indices(pInfo)]);
#     end
#     return nothing
# end

"""
	$(SIGNATURES)

Get *transformed* values for all parameters in the `ValueVector`.
"""
function get_values(pvec :: ParamVector, vv :: ValueVector{F1}) where F1
    # @assert isequal(get_object_id(pvec), get_object_id(vv))  """
    #     ObjectId mismatch:
    #     $(get_object_id(pvec))  vs  $(get_object_id(vv))
    #     """;
    valueV = Vector{F1}();
    for (pName, pInfo) in vv.d
        p = retrieve(pvec, pName);
        @assert size(lb(pInfo)) == size(calibrated_value(p));
        append!(valueV, transform_param(pvec.pTransform, p));
    end

    return valueV
end

function random_guess(vv :: ValueVector{F1}, rng :: AbstractRNG) where F1
    guessV = Vector{F1}();
    for (_, pInfo) in vv.d
        rv = random_values(pInfo, rng);
        if rv isa Real
            push!(guessV, rv);
        else
            append!(guessV, vec(rv));
        end
    end
    return guessV
end

# ----------------


mutable struct BOMCPTree{B,A}
    ### for each belief node
    n_b::Vector{Int} # Number of visits, N(b)
    b_children::Vector{Vector{Int}} #Children, Ch(b)
    b_labels::Vector{B} # Labels, b
    v::Vector{Float64} # Value, V(b) (Unused)

    m_b::Vector{Int} # Number of visits from gen, B(bab')

    ### for each state-action node
    n_a::Vector{Int} #Number of times this action node has been visited
    a_children::Vector{Vector{Int}} #spnode
    p_nodes::Vector{Int} # The parent of this node (state not state ID)
    p_labels::Vector{B} # The parent of this node (state not state ID)
    a_labels::Vector{A} # The action associated with this node
    q::Vector{Float64} #Current Q estimate
    transitions::Dict{Tuple{Int64, B}, Tuple{Int, Int}} # Belief:(Bnode, n)
    r::Vector{Float64} # Mean reward

    function BOMCPTree{B,A}(sz::Int=1000) where {B,A}
        sz = min(sz, 100_000)
        return new(sizehint!(Int[], sz),
                   sizehint!(Vector{Int}[], sz),
                   sizehint!(B[], sz),
                   sizehint!(Float64[], sz),

                   sizehint!(Int[], sz),

                   sizehint!(Int[], sz),
                   sizehint!(Vector{Int}[], sz),
                   sizehint!(Int[], sz),
                   sizehint!(B[], sz),
                   sizehint!(A[], sz),
                   sizehint!(Float64[], sz),
                   Dict{Tuple{Int64, B}, Tuple{Int64,Int64,Float64}}(),
                   sizehint!(Float64[], sz)
                  )
    end
end

function insert_belief_node!(tree::BOMCPTree{B,A}, b::B) where {B,A}
    push!(tree.n_b, 0)
    push!(tree.b_children, Int[])
    push!(tree.b_labels, b)
    push!(tree.v, 0.0)

    push!(tree.m_b, 1)

    bnode = length(tree.n_b)
    return bnode
end

function insert_action_node!(tree::BOMCPTree{B,A}, bnode::Int, a::A, n0::Int, q0::Float64) where {B,A}
    push!(tree.n_a, n0)
    push!(tree.a_children, Int[])
    push!(tree.p_labels, tree.b_labels[bnode])
    push!(tree.p_nodes, bnode)
    push!(tree.a_labels, a)
    push!(tree.q, q0)
    push!(tree.r, 0.0)
    banode = length(tree.n_a)
    push!(tree.b_children[bnode], banode)
    return banode
end

Base.isempty(tree::BOMCPTree) = isempty(tree.n) && isempty(tree.q)

struct BOMCPBeliefNode{B,A}
    tree::BOMCPTree{B,A}
    index::Int
end

children(n::BOMCPBeliefNode) = n.tree.b_children[n.index]
n_children(n::BOMCPBeliefNode) = length(children(n))
isroot(n::BOMCPBeliefNode) = n.index == 1


# This type represents an AND-node (Q-node) of the belief tree, branching
# on observations. It maps the set of observations seen during simulations to
# the VNodes that the simulations reach at the next level

  # Fields:
  # obs_to_particles: a mapping from observations seen to the particles (after
  # the transition) that generated the observations. The particles for each
  # observation become the representative set for the corresponding v-node
  # at the next level.
  # depth: depth of the V-node *above* this node.
  # action: action that led to this q-node.
  # first_step_reward: The average first step reward of particles when they
  # took action 'action'.
  # history: history up to the V-node *above* this node.
  # debug: Flag controlling debugging output.

mutable struct _QNode{S,A,O,B,T} #Workaround to achieve a circular type definition
    obs_to_particles::Dict{O, Vector{DESPOTParticle{S}}}
    depth::Int64
    action::A
    first_step_reward::Float64
    history::History
    weight_sum::Float64
    obs_and_nodes::Vector{Tuple{O,T}}
    n_visits::Int64                # Needed for large problems
    bounds::B

      # default constructor
      function _QNode{S,A,O,B,T}(
                    pomdp::POMDP{S,A,O},
                    bounds::B,
                    obs_to_particles::Dict{O,Vector{DESPOTParticle{S}}},
                    depth::Int64,
                    action::A,
                    first_step_reward::Float64,
                    history::History{A,O},
                    config::DESPOTConfig) where {S,A,O,B,T}

            this = new()
            this.obs_to_particles = obs_to_particles
            this.depth = depth
            this.action = action
            this.first_step_reward = first_step_reward
            this.history = history
            this.weight_sum = 0
            this.obs_and_nodes = Tuple{O,T}[]
            this.n_visits = 0
            this.bounds = bounds

            for (obs, particles) in this.obs_to_particles
                obs_weight_sum = 0.0
                for p in particles
                    obs_weight_sum += p.weight
                end
                this.weight_sum += obs_weight_sum
                add(this.history, action, obs)
                remove_last(this.history)
                push!(this.obs_and_nodes, (obs,
                                           VNode{S,A,O,B}(pomdp,
                                                          particles,
                                                          bounds,
                                                          this.depth+1,  # TODO: check depth
                                                          obs_weight_sum,
                                                          false,
                                                          config)))
            end
            return this
        end
end

# This type ("Value Node") encapsulates a belief node (and recursively, a
# belief tree). It stores the set of particles associated with the node, an
# AND-node for each action, and some bookkeeping information.

mutable struct VNode{S,A,O,B}
  particles::Array{DESPOTParticle{S},1}
  lbound::Float64
  ubound::Float64
  depth::Int64
  default_value::Float64            # Value of the default policy (equals to lbound value
                                    # before any backups are performed)
  pruned_action::A                  # Best action at the node after pruning
  weight::Float64                   # Sum of particle weights at this belief
  best_ub_action::A                 # Action that gives the highest upper bound
  best_lb_action::A                 # Action that gives the highest lower bound
  in_tree::Bool                     # True if the node is visited by trial().
                                    # In order to determine if a node is a fringe node
                                    # of the belief tree, we need to expand it one level.
                                    # The nodes added during this expansion of a fringe
                                    # node are not considered to be within the tree unless
                                    # explicitly visited by trial(), so we use
                                    # this indicator variable.
  n_tree_nodes::Int64               # Number of nodes with in_tree == true in the subtree
                                    # rooted at this node
  q_nodes::Dict{A,_QNode{S,A,O,B,
                VNode{S,A,O,B}}}    # Dict of children q-nodes
  n_visits::Int64                   # Needed for large domains
  n_actions_allowed::Int64          # current number of action branches allowed in the node, needed for large domains
  q_star::Float64                   # best current Q-value, needed for large domains

  # default constructor
  function VNode{S,A,O,B}(
               pomdp::POMDP{S,A,O},
               particles::Vector{DESPOTParticle{S}},
               b::B,
               depth::Int64,
               weight::Float64,
               in_tree::Bool,
               config::DESPOTConfig) where {S,A,O,B}

        this = new()
        this.particles          = particles
        this.lbound,
        this.ubound             = bounds(b, pomdp, particles, config)
        this.depth              = depth
        this.default_value      = this.lbound
        this.weight             = weight
        this.in_tree            = in_tree
        this.n_tree_nodes       = in_tree ? 1:0
        this.q_nodes            = Dict{A,QNode{S,A,O,B}}()
        this.n_visits           = 0
        this.n_actions_allowed  = 0
        this.q_star             = -Inf

        # note these are left undefined for now. Are they needed?
        # this.pruned_action      = #undef
        # this.best_ub_action     = #undef
        # this.best_lb_action     = #undef

        validate_bounds(this.lbound, this.ubound, config)
        return this
  end
end

const QNode{S,A,O,B} = _QNode{S,A,O,B,VNode{S,A,O,B}}

function get_upper_bound{S,A,O,B}(qnode::QNode{S,A,O,B})
  ubound::Float64 = 0.0
  for (obs, node) in qnode.obs_and_nodes
      ubound += node.ubound * node.weight
  end
  return ubound/qnode.weight_sum
end

function get_lower_bound{S,A,O,B}(qnode::QNode{S,A,O,B})
  lbound::Float64 = 0.0
  for (obs, node) in qnode.obs_and_nodes
      lbound += node.lbound * node.weight
  end
  return lbound/qnode.weight_sum
end

#TODO: Fix this

# function prune{S,A,O,L,U}(qnode::QNode{S,A,O,L,U}, total_pruned::Int64, config::DESPOTConfig)
#   cost::Float64 = 0.0
#   total_pruned::Int64 = 0
#
#   for (obs,node) in qnode.obs_to_node
#     cost, total_pruned += prune(node, total_pruned, config)
#   end
#   return cost, total_pruned
# end

function get_lb_action{S,A,O,B}(node::VNode{S,A,O,B}, config::DESPOTConfig, discount::Float64)
  local a_star
  q_star::Float64 = -Inf
  remaining_reward::Float64 = 0.0

  for (a,q_node) in node.q_nodes
    remaining_reward = get_lower_bound(q_node)
    if q_node.first_step_reward + discount * remaining_reward > q_star + config.tiny
      q_star = q_node.first_step_reward + discount * remaining_reward
      a_star = a
    end
  end
  return a_star
end

#TODO: fix pruning

# function prune{S,A,O,L,U}(node::VNode{S,A,O,L,U}, total_pruned::Int64, config::DESPOTConfig)
#   # Cost if the node was pruned
#   cost = (config.discount^node.depth) * node.weight * node.default_value
#                 - config.pruning_constant
#   if !node.inTree # Leaf
#     @assert(nodes.n_tree_nodes == 0)
#     return cost, total_pruned
#   end
#
#   for q_node in node.q_nodes
#     first_step_reward = config.discount^depth * weight * q_node.first_step_reward
#     best_child_value, total_pruned = prune(q_node, total_pruned, config)
#
#     new_cost = first_step_reward + best_child_value - config.pruning_constant
#     if new_cost > cost
#       cost = new_cost
#       pruned_action = q_node.action
#     end
#   end
#   if pruned_action == -1
#     total_pruned +=1
#   end
#   return cost, total_pruned
# end


four(::Type{Float64}) = 4.0
four(::Type{Float32}) = 4.0f0
four(::Type{Float16}) = Float16(4.0f0)

two(::Type{Float64}) = 2.0
two(::Type{Float32}) = 2.0f0
two(::Type{Float16}) = Float16(2.0f0)

negone(::Type{Float64}) = -1.0
negone(::Type{Float32}) = -1.0f0
negone(::Type{Float16}) = Float16(-1.0f0)

onehalf(::Type{Float64}) = 0.5
onehalf(::Type{Float32}) = 0.5f0
onehalf(::Type{Float16}) = Float16(0.5f0)

onequarter(::Type{Float64}) = 0.25
onequarter(::Type{Float32}) = 0.25f0
onequarter(::Type{Float16}) = Float16(0.25f0)


"""
    AbstractKKTVector{T, VT}

Supertype for KKT's right-hand-side vectors ``(x, s, y, z, ν, w)``.
"""
abstract type AbstractKKTVector{T, VT} end

"""
    full(X::AbstractKKTVector)

Return the all the values stored inside the KKT vector `X`.
"""
function full end

Base.length(rhs::AbstractKKTVector) = length(full(rhs))

"""
    number_primal(X::AbstractKKTVector)

Get total number of primal values ``(x, s)`` in KKT vector `X`.
"""
number_primal(rhs::AbstractKKTVector) = length(primal(rhs))

"""
    number_dual(X::AbstractKKTVector)

Get total number of dual values ``(y, z)`` in KKT vector `X`.
"""
number_dual(rhs::AbstractKKTVector) = length(dual(rhs))

"""
    primal(X::AbstractKKTVector)

Return the primal values ``(x, s)`` stored in the KKT vector `X`.
"""
function primal end

"""
    dual(X::AbstractKKTVector)

Return the dual values ``(y, z)`` stored in the KKT vector `X`.
"""
function dual end

"""
    primal_dual(X::AbstractKKTVector)

Return both the primal and the dual values ``(x, s, y, z)`` stored in the KKT vector `X`.
"""
function primal_dual end

"""
    dual_lb(X::AbstractKKTVector)

Return the dual values ``ν`` associated to the lower-bound
stored in the KKT vector `X`.
"""
function dual_lb end

"""
    dual_ub(X::AbstractKKTVector)

Return the dual values ``w`` associated to the upper-bound
stored in the KKT vector `X`.
"""
function dual_ub end

function Base.fill!(rhs::AbstractKKTVector{T}, val::T) where T
    fill!(full(rhs), val)
end

# Overload basic BLAS operations.
norm(X::AbstractKKTVector, p::Real) = norm(full(X), p)
dot(X::AbstractKKTVector, Y::AbstractKKTVector) = dot(full(X), full(Y))
function axpy!(a::Number, X::AbstractKKTVector, Y::AbstractKKTVector)
    axpy!(a, full(X), full(Y))
end

#=
    ReducedKKTVector
=#

"""
    ReducedKKTVector{T, VT<:AbstractVector{T}} <: AbstractKKTVector{T, VT}

KKT vector ``(x, s, y, z)``, associated to a [`AbstractReducedKKTSystem`](@ref).

Compared to [`UnreducedKKTVector`](@ref), it does not store
the dual values associated to the primal's lower and upper bounds.
"""
struct ReducedKKTVector{T, VT<:AbstractVector{T}} <: AbstractKKTVector{T, VT}
    values::VT
    xp::VT # unsafe view
    xl::VT # unsafe view
end

ReducedKKTVector(n::Int, m::Int, nlb::Int, nub::Int) = ReducedKKTVector(n, m)
function ReducedKKTVector(n::Int, m::Int)
    x = Vector{Float64}(undef, n + m)
    fill!(x, 0.0)
    # Wrap directly array x to avoid dealing with views
    pp = pointer(x)
    xp = unsafe_wrap(Vector{Float64}, pp, n)
    pd = pointer(x, n + 1)
    xl = unsafe_wrap(Vector{Float64}, pd, m)
    return ReducedKKTVector{Float64, Vector{Float64}}(x, xp, xl)
end
function ReducedKKTVector(rhs::AbstractKKTVector)
    return ReducedKKTVector(number_primal(rhs), number_dual(rhs))
end

full(rhs::ReducedKKTVector) = rhs.values
primal(rhs::ReducedKKTVector) = rhs.xp
dual(rhs::ReducedKKTVector) = rhs.xl
primal_dual(rhs::ReducedKKTVector) = rhs.values

#=
    UnreducedKKTVector
=#

"""
    UnreducedKKTVector{T, VT<:AbstractVector{T}} <: AbstractKKTVector{T, VT}

Full KKT vector ``(x, s, y, z, ν, w)``, associated to a [`AbstractUnreducedKKTSystem`](@ref).

"""
struct UnreducedKKTVector{T, VT<:AbstractVector{T}} <: AbstractKKTVector{T, VT}
    values::VT
    x::VT  # unsafe view
    xp::VT # unsafe view
    xl::VT # unsafe view
    xzl::VT # unsafe view
    xzu::VT # unsafe view
end

function UnreducedKKTVector(n::Int, m::Int, nlb::Int, nub::Int)
    values = Vector{Float64}(undef, n + m + nlb + nub)
    fill!(values, 0.0)
    # Wrap directly array x to avoid dealing with views
    pp = pointer(values)
    x = unsafe_wrap(Vector{Float64}, pp, n + m) # Primal-Dual
    xp = unsafe_wrap(Vector{Float64}, pp, n) # Primal
    pd = pointer(values, n + 1)
    xl = unsafe_wrap(Vector{Float64}, pd, m) # Dual
    pzl = pointer(values, n + m + 1)
    xzl = unsafe_wrap(Vector{Float64}, pzl, nlb) # Lower bound
    pzu = pointer(values, n + m + nlb + 1)
    xzu = unsafe_wrap(Vector{Float64}, pzu, nub) # Upper bound
    return UnreducedKKTVector{Float64, Vector{Float64}}(values, x, xp, xl, xzl, xzu)
end

full(rhs::UnreducedKKTVector) = rhs.values
primal(rhs::UnreducedKKTVector) = rhs.xp
dual(rhs::UnreducedKKTVector) = rhs.xl
primal_dual(rhs::UnreducedKKTVector) = rhs.x
dual_lb(rhs::UnreducedKKTVector) = rhs.xzl
dual_ub(rhs::UnreducedKKTVector) = rhs.xzu


using Turing, Random, Test

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@stage_testset "stan" "stan-interface.jl" begin
    using CmdStan

    chn = sample(gdemo_default, 2000, 1000, false, 1,
        CmdStan.Adapt(), CmdStan.Hmc(stepsize_jitter=0))
    check_gdemo(chn)
end


using StingerGraphs
using BenchmarkTools
using JLD
using FileIO

function bench(filename)
    s = Stinger()
    fields = [
        max_nv, max_neblock, max_netypes, max_nvtypes, num_insertions,
        num_deletions, num_insertions_last_batch, num_deletions_last_batch
    ]

    getrng = srand(1)

    getbench = @benchmarkable $s[f] samples=10000000  seconds=10 setup=(f=rand($getrng, $fields))

    setrng = srand(2)
    setbench = @benchmarkable $s[f]=val samples=10000000 setup=f,val=rand($setrng, $fields),rand($setrng,0:1000) seconds=10#For a reasonable run time

    info("Running getbench")
    gettrial = run(getbench)
    info("Running setbench")
    settrial = run(setbench)

    jldopen(filename, "w") do f
        write(f, "gettrial", gettrial)
        write(f, "settrial", settrial)
    end

    return gettrial, settrial
end


module MixedModelsExtras

export icc

using MixedModels

include("icc.jl")

end


"""
A `SolutionObject` is typically a struct with multiple Arrays being iterated.

Example:

```
struct VFIObject{T} <: SolutionObject
     vfun::Array{Float64, T}
     pol::Array{Int64, T}
end
```

Importantly, you have to define an `unpack` method for the subtypes of
`SolutionObject`. That should return a Tuple with the objects that are used to
determine convergence.

In the example above, if one is doing conventional VFI, they would do

```
unpack(vfiobj::VFIObject)=(vfiobj.vfun,)
```

"""
abstract type SolutionObject end

errmsg = "Please define an `unpack` method for "
unpack(x::SolutionObject)=error(errmsg * "$(typeof(x))")

" Interface for governing whether and how many of the iterations are displayed. "
abstract type IterationDisplay end

" Display whenever current iteration is a multiple of `mod` "
struct DisplayMod <: IterationDisplay
    mod :: Int64
end

display_iter(i, dm::DisplayMod) =
    (mod(i, dm.mod) == 0) && println("Iteration $i")

struct NoDisplay <: IterationDisplay end

display_iter(i, ::NoDisplay)=nothing

@with_kw struct IterationSpec{T <: IterationDisplay}
    maxiter::Int64 = 3000
    tol::Float64 = 1e-8
    show_convergence::Bool = true
    iterdisp::T = NoDisplay()
    lastiter_showdist::Bool = false
end

" Conservative specification for value function iteration "
const _vfi_conserv = IterationSpec(maxiter=3000,
                             tol = 1e-9,
                             iterdisp = DisplayMod(100),
                             lastiter_showdist=false)

function basic_message(niter, loop_cont, iterspec)
    if iterspec.show_convergence
        if loop_cont
            println("Stopped due to maximum iterations (N = $(iterspec.maxiter))")
        else
            println("Converged after $niter iterations")
        end
    end
end

"""
    contract!(container, xiter, other_args, F!, vfi::IterationSpec)

Simple interface for contraction mappings when all iteration objects
relevant to a given problem are stored in a container object, but such
container object might also include other fields.

A common use case of this is a regular VFI procedure. `contract!` allows
you to store both policy and value functions in the same object (`container`)
but to only iterate on the value function. In regular VFI, the policy function
is obtained as a "residual" of the VFI step, and it might be efficient to save it.

`container` will typically be a 'container type' that stores objects in a value
function iteration procedure. For example, container might contain a (array
representation of) a value function and a policy function.

`xiter` is typically a subcollection of the objects in `container` that will be
actually iterated on.

`F!` should be the function that takes iteration objects and executes the VFI
step.

**IMPORTANT: Has side-effects on both `container` and `xiter`**

"""
function contract!(container, xiter, other_args, F!, vfi::IterationSpec=_vfi_conserv)
    niter = 0
    loop_cont = true

    while loop_cont && (niter <= vfi.maxiter)
        niter += 1

        # Execute iteration
        F!(container, xiter, other_args...)

        # Evaluate convergence
        loop_cont = ! within_tolerance(container,
                                       xiter, vfi.tol)

        # Diagnostics
        display_iter(niter, vfi.iterdisp)

        # Replace `xiter` with container's updated version
        transfer!(xiter, container)
    end

    if vfi.lastiter_showdist
        niter += 1
        F!(container, xiter, other_args...)
        dd = maxdist_unroll(xiter, container)
        println(" Maximum distance: $dd ")
    end

    basic_message(niter, loop_cont, vfi)
end

unpack(x::Tuple)=x

transfer!(target::SolutionObject,
          source::SolutionObject)=transfer!(unpack(target), unpack(source))

@unroll function transfer!(tup_target::Tuple, tup_source::Tuple)
    i = 0
    @unroll for mu in tup_target
        i+=1
        copyto!(mu, tup_source[i])
    end
end

transfer!(xiter::AbstractArray, container)=copyto!(xiter, container)

within_tolerance(a,b,tol)=error("`within_tolerance` not defined for $(typeof(a))")
within_tolerance(a::T,b::T,tol) where {T <: SolutionObject} =
    within_tolerance(unpack(a), unpack(b), tol)

#
# Implements efficient tolerance check for l^∞
function within_tolerance(a::AbstractArray,b,tol)
    for i in eachindex(b)
        if !within_tolerance(a[i],b[i],tol)
            return false
        end
    end
    return true
end

@inline within_tolerance(a::Number, b, tol)=abs(b-a) <= tol
@inline within_tolerance(a::Int64, b::Int64, tol)=(a == b)

@unroll function within_tolerance_unroll(tup1::Tuple, tup2::Tuple, tol)
    i = 0

    @unroll for X in tup1
        i += 1
        Xpr = tup2[i]

        if !within_tolerance(X, Xpr, tol)
            return false
        end
    end

    return true
end

within_tolerance(tup1::Tuple, tup2::Tuple, tol)=within_tolerance_unroll(tup1, tup2, tol)

maxdist_unroll(x1::T, x2::T) where {T <: SolutionObject} =
    maxdist_unroll(unpack(x1), unpack(x2))

@unroll function maxdist_unroll(tup1, tup2)
    i = 0
    m = -Inf
    @unroll for X in tup1
        i += 1
        s = maximum(abs.(X .- tup2[i]))
        m = max(s, m)
    end

    m
end


using Metrics
using Test
using Random

Random.seed!(0)

@testset "Metrics.jl" begin
    include("regression.jl")
    include("classification.jl")
    include("rank.jl")
    include("nlp.jl")
    include("cv.jl")
end