using BSON: @save, @load
using GeometryBasics
using CairoMakie  # GLMakie
using SLAM
using ProgressMeter
using FileIO

include("kitty.jl")

function main(n_frames)
    base_dir = "/mnt/zy_data/kitty-dataset/"
    sequence = "00"
    stereo = true

    dataset = KittyDataset(base_dir, sequence; stereo)
    println(dataset)

    save_dir = joinpath("/mnt/zy_data/kitty-dataset/projects", "kitty-$sequence")
    isdir(save_dir) || mkdir(save_dir)
    @info "Save directory: $save_dir"

    fx, fy = dataset.K[1, 1], dataset.K[2, 2]
    cx, cy = dataset.K[1:2, 3]
    height, width = 376, 1241
    # height, width = 370, 1226

    camera = SLAM.Camera(fx, fy, cx, cy, 0, 0, 0, 0, height, width)
    right_camera = SLAM.Camera(fx, fy, cx, cy, 0, 0, 0, 0, height, width; Ti0=dataset.Ti0)

    params = Params(;
        stereo,
        window_size=9, max_distance=35, pyramid_levels=3,
        max_nb_keypoints=1000, max_reprojection_error=3.0,
        do_local_bundle_adjustment=false, map_filtering=true)

    saver = ReplaySaver()
    visualizer = nothing
    # visualizer = Visualizer((900, 600))
    # display(visualizer)

    slam_manager = SlamManager(params, camera; right_camera, visualizer=saver)
    slam_manager_thread = Threads.@spawn run!(slam_manager)

    t1 = time()
    @showprogress for i in 1:n_frames
        timestamp = dataset.timestamps[i]
        left_frame, right_frame = dataset[i]
        left_frame = Gray{Float64}.(left_frame)

        if params.stereo
            right_frame = Gray{Float64}.(right_frame)
            add_stereo_image!(slam_manager, left_frame, right_frame, timestamp)
        else
            add_image!(slam_manager, left_frame, timestamp)
        end

        if visualizer ≢ nothing
            SLAM.set_image!(visualizer, rotr90(left_frame))
            process_frame_wc!(visualizer)
        end

        q_size = get_queue_size(slam_manager)
        f_size = length(slam_manager.mapper.estimator.frame_queue)
        m_size = length(slam_manager.mapper.keyframe_queue)
        while q_size > 0 || f_size > 0 || m_size > 0
            sleep(1e-2)
            q_size = get_queue_size(slam_manager)
            f_size = length(slam_manager.mapper.estimator.frame_queue)
            m_size = length(slam_manager.mapper.keyframe_queue)
        end

        sleep(1e-2)
    end

    slam_manager.exit_required = true
    wait(slam_manager_thread)

    t2 = time()
    @info "SLAM took: $(t2 - t1) seconds."

    SLAM.save(saver, save_dir)
    slam_manager, visualizer
end

function replay(n_frames)
    base_dir = "/mnt/zy_data/kitty-dataset/"
    sequence = "00"
    stereo = true

    dataset = KittyDataset(base_dir, sequence; stereo)
    println(dataset)

    save_dir = joinpath("/mnt/zy_data/kitty-dataset/projects", "kitty-$sequence")
    isdir(save_dir) || mkdir(save_dir)
    @info "Save directory: $save_dir"

    saver = ReplaySaver()
    SLAM.load!(saver, save_dir)
    @assert length(saver.positions) == n_frames - 1

    resolution = (900, 600)
    image_resolution = (1241, 376)
    # image_resolution = (1226, 370)
    visualizer = Visualizer(;resolution, image_resolution)
    display(visualizer)

    t1 = time()
    for i in 2:n_frames
        left_frame, right_frame = dataset[i]
        left_frame = Gray{Float64}.(left_frame)

        position = saver.positions[i - 1]
        set_image!(visualizer, rotr90(left_frame))
        set_position!(visualizer, position)

        sleep(1 / 60)
    end
    t2 = time()
    @info "Visualization took: $(t2 - t1) seconds."

    visualizer
    save(joinpath(save_dir, "out_1.jpg"), visualizer.image)  # image  figure
    display(visualizer)
end

# main(4541)
replay(4541)

# 可正常运行, 就是plot和save_plot时没成功.
# http://www.cvlibs.net/datasets/kitti/eval_odometry.php   vslam, kitti


using Pkg;
Pkg.activate("");
using CUTEst, Dates, NLPModels, SolverBenchmark

using ARCTR, JSOSolvers, NLPModelsIpopt

nmax = 10000
problems = readlines("list_problems_$nmax.dat")
cutest_problems = (CUTEstModel(p) for p in problems)

max_time = 300.0 # 3600.0
max_ev = typemax(Int)
max_iter = typemax(Int)
atol = 1e-5
rtol = 1e-6

solvers = Dict(
    :ARCqKOp10201z025 =>
        nlp -> ARCqKOp(
            nlp,
            verbose = false,
            atol = atol,
            rtol = rtol,
            max_time = max_time,
            max_iter = max_iter,
            shifts = 10.0 .^ (collect(-10.0:1.0:20.0)),
            ζ = 0.25,
        ),
    :ARCqKOp10201z05 =>
        nlp -> ARCqKOp(
            nlp,
            verbose = false,
            atol = atol,
            rtol = rtol,
            max_time = max_time,
            max_iter = max_iter,
            shifts = 10.0 .^ (collect(-10.0:1.0:20.0)),
            ζ = 0.5,
        ),
    :ARCqKOp10201z075 =>
        nlp -> ARCqKOp(
            nlp,
            verbose = false,
            atol = atol,
            rtol = rtol,
            max_time = max_time,
            max_iter = max_iter,
            shifts = 10.0 .^ (collect(-10.0:1.0:20.0)),
            ζ = 0.75,
        ),
    :ARCqKOp10201z1 =>
        nlp -> ARCqKOp(
            nlp,
            verbose = false,
            atol = atol,
            rtol = rtol,
            max_time = max_time,
            max_iter = max_iter,
            shifts = 10.0 .^ (collect(-10.0:1.0:20.0)),
            ζ = 1.0,
        ),
)
rmv_problems = [
    "SBRYBND",
    "GENHUMPS",
    "NELSONLS",
    "NONMSQRT",
    "FMINSRF2",
    "JIMACK",
    "EIGENCLS",
    "INDEF",
    "FLETCBV3",
    "FMINSURF",
    "FLETCHBV",
    "PARKCH",
    "EIGENBLS",
]
stats = bmark_solvers(
    solvers,
    cutest_problems,
    skipif = prob -> (prob.meta.name in rmv_problems),
)

using JLD2
@save "$(today())_$(prod(String.(keys(solvers)) .* "_"))cutest_$(string(length(problems)))_$(nmax).jld2" stats


export PrerecordedDriver

type PrerecordedDriver <: DriverModel{NextState, IntegratedContinuous}
    trajdata::Trajdata # log we pull from
    time_start::Float64   # starting time [s]
    t_current::Float64
    timestep::Float64
    nextstate::VehicleState
end
function PrerecordedDriver(trajdata::Trajdata, frame_start::Int, egoid::Int, timestep::Float64=get_mean_timestep(trajdata))
    time_start = get_time(trajdata, frame_start)
    t_current = NaN
    nextstate = get_vehiclestate(trajdata, egoid, frame_start)
    PrerecordedDriver(trajdata, time_start, t_current, timestep, nextstate)
end

get_name(::PrerecordedDriver) = "PrerecordedDriver"
action_context(model::PrerecordedDriver) = IntegratedContinuous(model.timestep, 1)
function reset_hidden_state!(model::PrerecordedDriver)
    model.t_current = model.time_start
    model
end
function observe!(model::PrerecordedDriver, scene::Scene, roadway::Roadway, egoid::Int)
    model.t_current += model.timestep

    i = searchsortedfirst(model.trajdata.frames, model.t_current, by=frame_or_time->isa(frame_or_time, TrajdataFrame) ? frame_or_time.t : frame_or_time) # index of first frame ≥ t_current

    if i == 1
        lo, hi = 1, 2
    elseif i ≤ nframes(model.trajdata)
        lo, hi = i-1, i
    else # t_current > all times in model.trajdata.frames
        lo, hi = i-2, i-1
    end

    t = model.t_current
    t_lo = get_time(model.trajdata, lo)
    t_hi = get_time(model.trajdata, hi)

    γ = (t - t_lo) / (t_hi - t_lo)

    if (γ == 0.0 && !iscarinframe(model.trajdata, egoid, lo)) ||
       (γ == 1.0 && !iscarinframe(model.trajdata, egoid, hi)) ||
       (!iscarinframe(model.trajdata, egoid, lo) || !iscarinframe(model.trajdata, egoid, hi))

        model.nextstate = VehicleState()
    else
        state_lo = get_vehiclestate(model.trajdata, egoid, lo)
        state_hi = get_vehiclestate(model.trajdata, egoid, hi)
        model.nextstate = lerp(state_lo, state_hi, γ, model.trajdata.roadway)
    end

    model
end

Base.rand(model::PrerecordedDriver) = NextState(model.nextstate)
Distributions.pdf(model::PrerecordedDriver, a::NextState) = a.s == model.nextstate ? Inf : 0.0
Distributions.logpdf(model::PrerecordedDriver, a::NextState) = a.s == model.nextstate ? Inf : -Inf


export RandomPolicy

using Random

"""
    RandomPolicy(action_space, rng)

Randomly return a valid action.
"""
struct RandomPolicy{S<:AbstractSpace,R<:AbstractRNG} <: AbstractPolicy
    action_space::S
    rng::R
end

Base.show(io::IO, p::RandomPolicy) = print(io, "RandomPolicy($(p.action_space))")

Random.seed!(p::RandomPolicy, seed) = Random.seed!(p.rng, seed)

"""
    RandomPolicy(action_space; seed=nothing)
"""
RandomPolicy(s; seed = nothing) = RandomPolicy(s, MersenneTwister(seed))

"""
    RandomPolicy(env::AbstractEnv; seed=nothing)
"""
RandomPolicy(env::AbstractEnv; seed = nothing) =
    RandomPolicy(get_action_space(env), MersenneTwister(seed))

function (p::RandomPolicy)(obs, ::FullActionSet)
    legal_actions = get_legal_actions(obs)
    length(legal_actions) == 0 ? get_invalid_action(obs) : rand(p.rng, legal_actions)
end

(p::RandomPolicy)(obs, ::MinimalActionSet) = rand(p.rng, p.action_space)
(p::RandomPolicy)(obs::BatchObs, ::MinimalActionSet) =
    [rand(p.rng, p.action_space) for _ in 1:length(obs)]

RLBase.update!(p::RandomPolicy, experience) = nothing

RLBase.get_prob(p::RandomPolicy, s) =
    fill(1 / length(p.action_space), length(p.action_space))
RLBase.get_prob(p::RandomPolicy, s, a) = 1 / length(p.action_space)


using Test

using MathOptInterface
const MOI = MathOptInterface

include("dummy.jl")

"""
    MutLessThan{T<:Real} <: MOI.AbstractScalarSet

A mutable `LessThan`-like set to test `copy` of indicator set
"""
mutable struct MutLessThan{T<:Real} <: MOI.AbstractScalarSet
    upper::T
    MutLessThan(v::T) where {T<:Real} = new{T}(v)
end

Base.copy(mlt::MutLessThan) = MutLessThan(Base.copy(mlt.upper))

@testset "Sets" begin
    @testset "==" begin
        # By default, `==` redirects to `===`, it works for bits type
        # but not for `BigInt`s. We define functions creating different
        # instances so that `a() !== a()`.
        a() = big(1)
        b() = big(2)
        @test a() !== a()
        @test b() !== b()
        for S in [MOI.LessThan, MOI.GreaterThan, MOI.EqualTo, MOI.PowerCone, MOI.DualPowerCone]
            @test S(a()) == S(a())
            @test S(a()) != S(b())
            @test S(b()) == S(b())
            @test S(b()) != S(a())
        end
        for S in [MOI.Interval, MOI.Semicontinuous, MOI.Semiinteger]
            @test S(a(), b()) == S(a(), b())
            @test S(a(), b()) != S(b(), a())
            @test S(a(), b()) != S(b(), b())
            @test S(a(), b()) != S(a(), a())
            @test S(a(), a()) != S(b(), b())
            @test S(a(), a()) == S(a(), a())
        end
        S = MOI.IndicatorSet
        A() = MOI.LessThan(a())
        B() = MOI.LessThan(b())
        @test S{MOI.ACTIVATE_ON_ZERO}(A()) == S{MOI.ACTIVATE_ON_ZERO}(A())
        @test S{MOI.ACTIVATE_ON_ZERO}(A()) != S{MOI.ACTIVATE_ON_ONE}(A())
        @test S{MOI.ACTIVATE_ON_ZERO}(A()) != S{MOI.ACTIVATE_ON_ZERO}(B())
        @test S{MOI.ACTIVATE_ON_ONE}(A()) != S{MOI.ACTIVATE_ON_ONE}(B())
    end
    @testset "Copy" begin
        @testset "for $S" for S in [MOI.SOS1, MOI.SOS2]
            s = S([1.0])
            s_copy = copy(s)
            s_copy.weights[1] = 2.0
            @test s.weights[1] == 1.0
        end
        @testset "IndicatorSet" begin
            s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(4.0))
            s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.GreaterThan(4.0))
            s1_copy = copy(s1)
            s2_copy = copy(s2)
            @test s1_copy isa MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}
            @test s1 == s1_copy
            @test s2_copy isa MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}
            @test s2 == s2_copy
            s3 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MutLessThan(4.0))
            s3_copy = copy(s3)
            @test s3.set.upper ≈ 4.0
            s3_copy.set.upper = 5.0
            @test s3.set.upper ≈ 4.0
            @test s3_copy.set.upper ≈ 5.0
        end
    end
    @testset "Broadcast" begin
        model = DummyModelWithAdd()
        x = MOI.add_variables(model, 3)
        cis = MOI.add_constraint.(model, x, MOI.EqualTo(0.0))
        @test cis isa Vector{MOI.ConstraintIndex{MOI.SingleVariable,
                                                 MOI.EqualTo{Float64}}}
        @test length(cis) == 3
    end
end


# This is a collection of C functions provided by the
# plotting library MathGL

# Constants and Types
#-------------------------------------------------------------------------#

const lib     = Libdl.find_library(["libmgl", "libmgl2"])
# For debugging purposes
#=const lib     = Libdl.find_library(["/home/thomas/downloads/mathgl-2.3.5.1/src/libmgl2.so"])=#
const libqt   = Libdl.find_library(["libmgl-qt"])
const libfltk = Libdl.find_library(["libmgl-fltk"])
const libglut = Libdl.find_library(["libmgl-glut"])

typealias HMGL Ptr{Void}
typealias HCDT Ptr{Void}
typealias HMDT Ptr{Void}
typealias HADT Ptr{Void}
typealias HMPR Ptr{Void}
typealias HMEX Ptr{Void}
typealias HMAX Ptr{Void}

typealias DrawCallback Ptr{Void}
typealias LoadCallback Ptr{Void}

typealias Float Float64
typealias Dual  Complex{Float}

include("constants.jl")

# Functions
#-------------------------------------------------------------------------#

# The "raw" access to the C interface provided by the mgl libs is handled
# by the getf macro, defined in getfunction.jl
include("getfunction.jl")
# The functions are called as they are in the C api, though this notation
# does not meet the usual julia styling conventions (the C interface is not
# intended to be exposed to users who do not explicitly want everything to
# be like in the C api).

include("functions.jl")

# For the functions requireing qt/glut/fltk no test were/are conducted
include("qtfunctions.jl")
include("glutfunctions.jl")
include("fltkfunctions.jl")


module EvolutionaryStrategies

using Cambrian
using Random
using LinearAlgebra

abstract type ESState end

include("individual.jl")
include("exponential_nes.jl")
include("separable_nes.jl")

end


"""
Spikes

Includes types and methods for handling spike trains.
"""
module Spikes

using JLD
import Base: show, isequal, hash
include("metadata.jl")

################################################################################
#### Include appropriate files
################################################################################
include("SpikeTrains.jl")
include("poissonproc.jl")
include("fileio.jl")
# include("CRCNS/CRCNS.jl")

################################################################################
#### Miscellaneous functions and constants
################################################################################

export show, isequal, hash, savespikes, loadspikes,
       SpikeTrains, histogram, raster, n_cells,
       inhomogeneous_poisson_process

end #module


using YAML
using ArgParse

function add_args(s::ArgParseSettings)
    @add_arg_table(
        s,
        "--cfg", arg_type = String, default = "cfg/default.yaml",
        "--vinput", arg_type = Int64, default = 16,
        "--noise", arg_type = Float64, default = 5.0,
        "--theta_plus", arg_type = Float64, default = 0.4,
        "--refrac_e", arg_type = Float64, default = 15.0,
        "--refrac_i", arg_type = Float64, default = 4.0,
        "--ttheta", arg_type = Float64, default = 4.7,
        "--tdop", arg_type = Float64, default = 6.2,
        "--lr", arg_type = Float64, default = 6.0,
        "--target", arg_type = Float64, default = 0.4,
        "--wrandom", arg_type = Bool, default = true,
        "--wstart", arg_type = Float64, default = 0.9,
        "--winh", arg_type = Float64, default = 0.4,
        "--wmax", arg_type = Float64, default = 3.5,
        "--mu", arg_type = Float64, default = 0.1
    )
    s
end

function get_config(argtable::ArgParseSettings)
    argtable = add_args(argtable)
    args = parse_args(argtable)
    cfg = YAML.load_file(args["cfg"])

    for k in ["noise", "vinput", "theta_plus", "refrac_e", "refrac_i", "lr",
              "target", "wrandom", "wstart", "winh", "wmax", "mu"]
        cfg[k] = args[k]
    end

    for k in ["ttheta", "tdop"]
        cfg[k] = args[k] * 100
    end
    args, cfg
end


using Oulu
using Test

@testset "Oulu.jl" begin
    # Write your tests here.
    @test Oulu.addNumbers(1, 2) == 3
    @test Oulu.addNumbers(1, 2) != 4
    @code_warntype Oulu.addNumbers(1, 2)
end


function hcat_by_sum(xs::AbstractVector)
    T = eltype(xs[1])
    dim = size(xs[1], 1)
    N = length(xs)

    ns = map(x->size(x,2), xs)
    pushfirst!(ns, 1)
    cumsum!(ns, ns)

    A = similar(xs[1], T, dim, ns[end]-1)
    for i in 1:N
        A[:, ns[i]:(ns[i+1]-1)] .= xs[i]
    end
    return A
end

function ChainRulesCore.rrule(::typeof(hcat_by_sum), xs::AbstractVector)
    N = length(xs)

    ns = map(x->size(x,2), xs)
    pushfirst!(ns, 1)
    cumsum!(ns, ns)

    hcat_by_sum_pullback(Δ) = (NoTangent(), ntuple(i->view(Δ,:,ns[i]:(ns[i+1]-1)), N))
    hcat_by_sum(xs), hcat_by_sum_pullback
end

function ChainRulesCore.rrule(::typeof(parent), A::Base.SubArray)
    parent_pullback(Δ) = (NoTangent(), view(Δ, A.indices...))
    parent(A), parent_pullback
end


using Oceananigans.BoundaryConditions: ChannelSolutionBCs

"""
    ChannelModel(; kwargs...)

Construct a `Model` with walls in the y-direction. This is done by imposing
`FreeSlip` boundary conditions in the y-direction instead of `Periodic`.

kwargs are passed to the regular `Model` constructor.
"""
ChannelModel(; boundary_conditions=ChannelSolutionBCs(), kwargs...) =
    Model(; boundary_conditions=boundary_conditions, kwargs...)


"Compilation to Julia (and other targets, if you want)"
module Compile

using ..AExpressions, ..CompileUtils
import MacroTools: striplines

export compiletojulia, runprogram

"compile `aexpr` into Expr"
function compiletojulia(aexpr::AExpr)::Expr

  # dictionary containing types/definitions of global variables, for use in constructing init func.,
  # next func., etcetera; the three categories of global variable are external, initnext, and lifted
  historydata = Dict([("external" => []), # :typedecl aexprs for all external variables
               ("initnext" => []), # :assign aexprs for all initnext variables
               ("lifted" => []), # :assign aexprs for all lifted variables
               ("types" => Dict())]) # map of global variable names (symbols) to types

  if (aexpr.head == :program)
    # handle AExpression lines
    lines = filter(x -> x !== :(), map(arg -> compile(arg, historydata, aexpr), aexpr.args))

    # construct STATE struct and initialize state::STATE
    stateStruct = compilestatestruct(historydata)
    initStateStruct = compileinitstate(historydata)

    # handle init, next, prev, and built-in functions
    initnextFunctions = compileinitnext(historydata)
    prevFunctions = compileprevfuncs(historydata)
    builtinFunctions = compilebuiltin()

    # remove empty lines
    lines = filter(x -> x != :(),
            vcat(builtinFunctions, lines, stateStruct, initStateStruct, prevFunctions, initnextFunctions))

    # construct module
    expr = quote
      module CompiledProgram
        export init, next
        import Base.min
        using Distributions
        using MLStyle: @match
        $(lines...)
      end
    end
    expr.head = :toplevel
    striplines(expr)
  else
    throw(AutumnCompileError("AExpr Head != :program"))
  end
end

"Run `prog` for finite number of time steps"
function runprogram(prog::Expr, n::Int)
  mod = eval(prog)
  state = mod.init(mod.Click(5, 5))

  for i in 1:n
    externals = [nothing, mod.Click(rand([1:10;]), rand([1:10;]))]
    state = mod.next(mod.next(state, externals[rand(Categorical([0.7, 0.3]))]))
  end
  state
end

end


@testset "testimports" begin
    # Check that everything is imported without an error. This is quite a weak criterion,
    # but I'm not sure what else is easily doable.
    for op in DLA.ops
        @eval $(import_expr(op))
    end
end


# main.jl
# Code to generate various practice flashcards
# for Julia's || and && operators.

# Change this if you want to generate more or fewer flashcards. There are 127 possibilities in total.
NUM_FLASHCARDS = 30

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

operators = (:(||), :(&&))
bools = (true, false)

# It's good practice to preallocate the arrays.

A = Array{Bool}(undef, NUM_FLASHCARDS, 4)
B = Array{Symbol}(undef, NUM_FLASHCARDS, 3)

# Now we fill them with randomly-selected members from the
# tuples above.

Random.rand!(A, bools)
Random.rand!(B, operators)

# Finally, let's make the flashcards themselves. S will have
# the question slide, and E will have the answer slide.

S = Array{String}(undef, NUM_FLASHCARDS)
E = Array{Bool}(undef, NUM_FLASHCARDS)

for i in 1:size(S)[1]
    S[i] = "$(A[i,1]) $(B[i,1]) $(A[i,2]) $(B[i,2]) $(A[i,3]) $(B[i,3]) $(A[i,4])"
end

@. E = eval(Meta.parse(S))

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

open("..\\Julia-Short-Circuit-Flashcard-Generators\\generated_flashcards.txt", "w") do io
    for i in 1:size(S)[1]
        print(io, "$(S[i])\n$(E[i])\n\n")
    end
end

println("Flashcards generated.~")
# Good work, my friends.

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


"""
    _InitialValue

A singleton type for representing "universal" initial value (identity element).

The idea is that, given `op` for `mapfoldl`, virtually, we define an "extended"
version of it by

    op′(::_InitialValue, x) = x
    op′(acc, x) = op(acc, x)

This is just a conceptually useful model to have in mind and we don't actually
define `op′` here  (yet?).  But see `Base.BottomRF` for how it might work in
action.

(It is related to that you can always turn a semigroup without an identity into
a monoid by "adjoining" an element that acts as the identity.)
"""
struct _InitialValue end

@inline _first(a1, as...) = a1

################
## map / map! ##
################

# In 0.6 the three methods below could be replaced with
# `map(f, as::Union{<:StaticArray,AbstractArray}...)` which included at least one `StaticArray`
# this is not the case on 0.7 and we instead hope to find a StaticArray in the first two arguments.
@inline function Base.map(f, a1::TupleVector, as::AbstractArray...)
    _map(f, a1, as...)
end
@inline function Base.map(f, a1::AbstractArray, a2::TupleVector, as::AbstractArray...)
    _map(f, a1, a2, as...)
end
@inline function Base.map(f, a1::TupleVector, a2::TupleVector, as::AbstractArray...)
    _map(f, a1, a2, as...)
end

@generated function _map(f, a::AbstractArray...)
    first_tuplevector = findfirst(ai -> ai <: TupleVector, a)
    if first_tuplevector === nothing
        return :(throw(ArgumentError("No TupleVector found in argument list")))
    end
    # Passing the Val as an argument to _map leads to inference issues when
    # recursively mapping over nested TupleVectors (see issue #593). Calling
    # Val in the generator here is valid because a[first_staticarray] is known to be a
    # TupleVector for which the default Val method is correct. If wrapped
    # TupleVector (with a custom Val method) are to be supported, this will
    # no longer be valid.
    S = length(a[first_tuplevector])

    if S == 0
        # In the empty case only, use inference to try figuring out a sensible
        # eltype, as is done in Base.collect and broadcast.
        # See https://github.com/JuliaArrays/StaticArrays.jl/issues/528
        eltys = [:(eltype(a[$i])) for i ∈ 1:length(a)]
        return quote
            @_inline_meta
            S = same_size(a...)
            T = Core.Compiler.return_type(f, Tuple{$(eltys...)})
            @inbounds return similar_type(a[$first_staticarray], T, S)()
        end
    end

    exprs = Vector{Expr}(undef, S)
    for i ∈ 1:S
        tmp = [:(a[$j][$i]) for j ∈ 1:length(a)]
        exprs[i] = :(f($(tmp...)))
    end

    return quote
        @_inline_meta
        S = same_size(a...)
        @inbounds elements = tuple($(exprs...))
        @inbounds return similar_type(typeof(_first(a...)), eltype(elements), S)(elements)
    end
end

@inline function Base.map!(f, dest::TupleVector{n}, a::TupleVector{n}...) where n
    _map!(f, dest, Val(n), a...)
end

# Ambiguities with Base:
@inline function map!(f, dest::TupleVector{n}, a::TupleVector{n}) where n
    _map!(f, dest, Val(n), a)
end
@inline function map!(f,dest::TupleVector{n},a::TupleVector{n},b::TupleVector{n}) where n
    _map!(f, dest, Val(n), a, b)
end

@generated function _map!(f, dest, ::Val{S}, a::TupleVector...) where {S}
    exprs = Vector{Expr}(undef, S)
    for i ∈ 1:S
        tmp = [:(a[$j][$i]) for j ∈ 1:length(a)]
        exprs[i] = :(dest[$i] = f($(tmp...)))
    end
    return quote
        @_inline_meta
        @inbounds $(Expr(:block, exprs...))
    end
end

###############
## mapreduce ##
###############

@inline function Base.mapreduce(f, op, a::TupleVector, b::TupleVector...; dims=:, init = _InitialValue())
    _mapreduce(f, op, dims, init, same_size(a, b...), a, b...)
end

@inline _mapreduce(args::Vararg{Any,N}) where N = _mapfoldl(args...)

@generated function _mapfoldl(f, op, dims::Colon, init, ::Val{S}, a::TupleVector...) where {S}
    if S == 0
        if init === _InitialValue
            if length(a) == 1
                return :(Base.mapreduce_empty(f, op, $(eltype(a[1]))))
            else
                return :(throw(ArgumentError("reducing over an empty collection is not allowed")))
            end
        else
            return :init
        end
    end
    tmp = [:(a[$j][1]) for j ∈ 1:length(a)]
    expr = :(f($(tmp...)))
    if init === _InitialValue
        expr = :(Base.reduce_first(op, $expr))
    else
        expr = :(op(init, $expr))
    end
    for i ∈ 2:S
        tmp = [:(a[$j][$i]) for j ∈ 1:length(a)]
        expr = :(op($expr, f($(tmp...))))
    end
    return quote
        @_inline_meta
        @inbounds return $expr
    end
end

@inline function _mapreduce(f, op, D::Int, init, sz::Val{S}, a::TupleVector) where {S}
    # Body of this function is split because constant propagation (at least
    # as of Julia 1.2) can't always correctly propagate here and
    # as a result the function is not type stable and very slow.
    # This makes it at least fast for three dimensions but people should use
    # for example any(a; dims=Val(1)) instead of any(a; dims=1) anyway.
    if D == 1
        return _mapreduce(f, op, Val(1), init, sz, a)
    else
        return _mapreduce(f, op, Val(D), init, sz, a)
    end
end

@generated function _mapfoldl(f, op, dims::Val{1}, init,
                               ::Val{S}, a::TupleVector) where {S,D}

    exprs = Array{Expr}(undef, 1)
    for i ∈ Base.product(1:1)
        expr = :(f(a[$(i...)]))
        if init === _InitialValue
            expr = :(Base.reduce_first(op, $expr))
        else
            expr = :(op(init, $expr))
        end

        exprs[i...] = expr
    end

    return quote
        @_inline_meta
        @inbounds elements = tuple($(exprs...))
        @inbounds return similar_type(a, eltype(elements), Val(1))(elements)
    end
end

############
## reduce ##
############

@inline Base.reduce(op, a::TupleVector; dims = :, init = _InitialValue()) =
    _reduce(op, a, dims, init)

# disambiguation
Base.reduce(::typeof(vcat), A::TupleVector{n,<:AbstractVecOrMat}) where n =
    Base._typed_vcat(mapreduce(eltype, promote_type, A), A)
Base.reduce(::typeof(vcat), A::TupleVector{n,<:TupleVectorLike}) where n =
    _reduce(vcat, A, :, _InitialValue())

Base.reduce(::typeof(hcat), A::TupleVector{n,<:AbstractVecOrMat}) where n =
    Base._typed_hcat(mapreduce(eltype, promote_type, A), A)
Base.reduce(::typeof(hcat), A::TupleVector{n,<:TupleVectorLike}) where n =
    _reduce(hcat, A, :, _InitialValue())

@inline _reduce(op, a::TupleVector, dims, init = _InitialValue()) =
    _mapreduce(identity, op, dims, init, length_val(a), a)

################
## (map)foldl ##
################

# Using `where {R}` to force specialization. See:
# https://docs.julialang.org/en/v1.5-dev/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing-1
# https://github.com/JuliaLang/julia/pull/33917

@inline Base.mapfoldl(f::F, op::R, a::TupleVector; init = _InitialValue()) where {F,R} =
    _mapfoldl(f, op, :, init, length_val(a), a)
@inline Base.foldl(op::R, a::TupleVector; init = _InitialValue()) where {R} =
    _foldl(op, a, :, init)
@inline _foldl(op::R, a, dims, init = _InitialValue()) where {R} =
    _mapfoldl(identity, op, dims, init, length_val(a), a)

#######################
## related functions ##
#######################

# These are all similar in Base but not @inline'd
#
# Implementation notes:
#
# 1. mapreduce and mapreducedim usually do not take initial value, because we don't
# always know the return type of an arbitrary mapping function f.  (We usually want to use
# some initial value such as one(T) or zero(T), where T is the return type of f, but
# if users provide type-unstable f, its return type cannot be known.)  Therefore, mapped
# versions of the functions implemented below usually require the collection to have at
# least two entries.
#
# 2. Exceptions are the ones that require Boolean mapping functions.  For example, f in
# all and any must return Bool, so we know the appropriate initial value is true and false,
# respectively.  Therefore, all(f, ...) and any(f, ...) are implemented by mapreduce(f, ...)
# with an initial value v0 = true and false.
#
# TODO: change to use Base.reduce_empty/Base.reduce_first
@inline Base.iszero(a::TupleVector{N,T}) where {N,T} = reduce((x,y) -> x && iszero(y), a, init=true)

@inline Base.sum(a::TupleVector{N,T}; dims=:) where {N,T} = _reduce(+, a, dims)
@inline Base.sum(f, a::TupleVector{N,T}; dims=:) where {N,T} = _mapreduce(f, +, dims, _InitialValue(), Val(N), a)
@inline Base.sum(f::Union{Function, Type}, a::TupleVector{N,T}; dims=:) where {N,T} = _mapreduce(f, +, dims, _InitialValue(), Val(N), a) # avoid ambiguity

@inline Base.prod(a::TupleVector{N,T}; dims=:) where {N,T} = _reduce(*, a, dims)
@inline Base.prod(f, a::TupleVector{N,T}; dims=:) where {N,T} = _mapreduce(f, *, dims, _InitialValue(), Val(N), a)
@inline Base.prod(f::Union{Function, Type}, a::TupleVector{N,T}; dims=:) where {N,T} = _mapreduce(f, *, dims, _InitialValue(), Val(N), a)

@inline Base.count(a::TupleVector{N,Bool}; dims=:) where N = _reduce(+, a, dims)
@inline Base.count(f, a::TupleVector{N}; dims=:) where N = _mapreduce(x->f(x)::Bool, +, dims, _InitialValue(), Val(N), a)

@inline Base.all(a::TupleVector{N,Bool}; dims=:) where N = _reduce(&, a, dims, true)  # non-branching versions
@inline Base.all(f::Function, a::TupleVector{N}; dims=:) where N = _mapreduce(x->f(x)::Bool, &, dims, true, Val(N), a)

@inline Base.any(a::TupleVector{N,Bool}; dims=:) where N = _reduce(|, a, dims, false) # (benchmarking needed)
@inline Base.any(f::Function, a::TupleVector{N}; dims=:) where N = _mapreduce(x->f(x)::Bool, |, dims, false, Val(N), a) # (benchmarking needed)

@inline Base.in(x, a::TupleVector{N}) where N = _mapreduce(==(x), |, :, false, Val(N), a)

#=_mean_denom(a, dims::Colon) = length(a)
_mean_denom(a, dims::Int) = size(a, dims)
_mean_denom(a, ::Val{D}) where {D} = size(a, D)
_mean_denom(a, ::Type{Val{D}}) where {D} = size(a, D)

@inline Statistics.mean(a::TupleVector; dims=:) = _reduce(+, a, dims) / _mean_denom(a, dims)
@inline Statistics.mean(f::Function, a::TupleVector{N}; dims=:) = _mapreduce(f, +, dims, _InitialValue(), Val(N), a) / _mean_denom(a, dims)=#

@inline Base.minimum(a::TupleVector; dims=:) = _reduce(min, a, dims) # base has mapreduce(idenity, scalarmin, a)
@inline Base.minimum(f::Function, a::TupleVector{N}; dims=:) where N = _mapreduce(f, min, dims, _InitialValue(), Val{N}, a)

@inline Base.maximum(a::TupleVector; dims=:) = _reduce(max, a, dims) # base has mapreduce(idenity, scalarmax, a)
@inline Base.maximum(f::Function, a::TupleVector; dims=:) = _mapreduce(f, max, dims, _InitialValue(), Val(N), a)

# Diff is slightly different
@inline LinearAlgebra.diff(a::TupleVector{N};dims=Val(1)) where N = _diff(Val(N),a,dims)

@inline function _diff(sz::Val, a::TupleVector, D::Int)
    _diff(sz,a,Val(D))
end
@generated function _diff(::Val{N}, a::TupleVector, ::Val{1}) where N
    Snew = N-1
    exprs = Array{Expr}(undef, Snew)
    for i1 = Base.product(1:Snew)
        i2 = copy([i1...])
        i2[1] = i1[1] + 1
        exprs[i1...] = :(a[$(i2...)] - a[$(i1...)])
    end
    return quote
        @_inline_meta
        elements = tuple($(exprs...))
        @inbounds return similar_type(a, eltype(elements), Val($Snew))(elements)
    end
end

_maybe_val(dims::Integer) = Val(Int(dims))
_maybe_val(dims) = dims
_valof(::Val{D}) where D = D

@inline Base.accumulate(op::F, a::TupleVector; dims = :, init = _InitialValue()) where {F} =
    _accumulate(op, a, _maybe_val(dims), init)

@inline function _accumulate(op::F, a::TupleVector, dims::Union{Val,Colon}, init) where {F}
    # Adjoin the initial value to `op` (one-line version of `Base.BottomRF`):
    rf(x, y) = x isa _InitialValue ? Base.reduce_first(op, y) : op(x, y)

    if isempty(a)
        T = return_type(rf, Tuple{typeof(init), eltype(a)})
        return similar_type(a, T)()
    end

    results = _foldl(
        a,
        dims,
        (similar_type(a, Union{}, Val(0))(), init),
    ) do (ys, acc), x
        y = rf(acc, x)
        # Not using `push(ys, y)` here since we need to widen element type as
        # we iterate.
        (vcat(ys, TV[y]), y)
    end
    dims === (:) && return first(results)

    ys = map(first, results)
    # Now map over all indices of `a`.  Since `_map` needs at least
    # one `TupleVector` to be passed, we pass `a` here, even though
    # the values of `a` are not used.
    data = _map(a, CartesianIndices(a)) do _, CI
        D = _valof(dims)
        I = Tuple(CI)
        J = Base.setindex(I, 1, D)
        ys[J...][I[D]]
    end
    return similar_type(a, eltype(data))(data)
end

@inline Base.cumsum(a::TupleVector; kw...) = accumulate(Base.add_sum, a; kw...)
@inline Base.cumprod(a::TupleVector; kw...) = accumulate(Base.mul_prod, a; kw...)


using Test
using QuantumLattices.Prerequisites.TypeTraits

struct WithTrait{F1,F2,F3}
    f1::F1
    f2::F2
    f3::F3
end
Base.:(==)(fc1::WithTrait,fc2::WithTrait) = ==(efficientoperations,fc1,fc2)
Base.isequal(fc1::WithTrait,fc2::WithTrait) = isequal(efficientoperations,fc1,fc2)
Base.:<(fc1::WithTrait,fc2::WithTrait) = <(efficientoperations,fc1,fc2)
Base.isless(fc1::WithTrait,fc2::WithTrait) = isless(efficientoperations,fc1,fc2)
Base.isapprox(fc1::WithTrait,fc2::WithTrait;atol::Real=10^-5,rtol::Real=10^-5) = isapprox(efficientoperations,Val((:f1,:f2)),fc1,fc2;atol=atol,rtol=rtol)
Base.replace(fc::WithTrait;kwargs...) = replace(efficientoperations,fc;kwargs...)

@testset "efficientoperations" begin
    @test ==(efficientoperations,(),())
    @test isequal(efficientoperations,(),())
    @test ==(efficientoperations,(1,2),(1,2,3))==false
    @test isequal(efficientoperations,(1,2),(1,2,3))==false

    fc1,fc2=WithTrait(1.0,2,3),WithTrait(1,2,3)
    @test fc1==fc2
    @test isequal(fc1,fc2)
    fc1,fc2=WithTrait(1.0,2,3),WithTrait(1,2,4.0)
    @test fc1<fc2
    @test isless(fc1,fc2)

    fc1,fc2=WithTrait(1.0+10^-6,2,3),WithTrait(1,2,3)
    @test fc1≈fc2
    fc1,fc2=WithTrait(1.0,2-10^-6,3),WithTrait(1,2,3)
    @test fc1≈fc2
    fc1,fc2=WithTrait(1.0,2,3-10^-6),WithTrait(1,2,3)
    @test fc1≉ fc2

    @test replace(WithTrait(1,2,3),f1='c')==WithTrait('c',2,3)
end

@testset "forder/corder" begin
    dims=(2,2,2)
    finds=[(1,1,1),(2,1,1),(1,2,1),(2,2,1),(1,1,2),(2,1,2),(1,2,2),(2,2,2)]
    cinds=[(1,1,1),(1,1,2),(1,2,1),(1,2,2),(2,1,1),(2,1,2),(2,2,1),(2,2,2)]
    @test collect(indtosub(dims,i,forder) for i=1:prod(dims))==finds
    @test collect(indtosub(dims,i,corder) for i=1:prod(dims))==cinds
    @test collect(subtoind(dims,inds,forder) for inds in finds)==collect(1:prod(dims))
    @test collect(subtoind(dims,inds,corder) for inds in cinds)==collect(1:prod(dims))
end


module FSimBase

using Reexport
import DiffEqBase
using DiffEqBase: DECallback, DEIntegrator, ODEProblem, DiscreteProblem, CallbackSet, isinplace, init, solve
using DiffEqCallbacks: SavedValues, SavingCallback
@reexport using SimulationLogger
using Transducers: Map
using DataFrames: DataFrame
using NamedTupleTools: namedtuple
using DataFrames

export AbstractEnv, State, Params, Dynamics, Dynamics!
export AbstractController, Command
export apply_inputs
export Simulator
export solve, reinit!, log!, step!, step_until!

include("APIs/APIs.jl")

end


module Alpaca

    ##############################################################################
    ##
    ## Dependencies
    ##
    ##############################################################################

    using Base
    using Statistics
    using RCall
    using DataFrames
    using Distributions
    using Printf
    import StatsBase: coef, nobs, coeftable, vcov, predict, residuals, var,
        RegressionModel, model_response, stderror, confint, fit, CoefTable,
        dof_residual, r2, adjr2, deviance, mss, rss, islinear, response
    using Reexport
    using CategoricalArrays
    @reexport using StatsBase
    @reexport using StatsModels

    ##############################################################################
    ##
    ## Exported methods and types
    ##
    ##############################################################################

    export feglm

    ##############################################################################
    ##
    ## Load files
    ##
    ##############################################################################

    include("RegressionResult.jl")
    include("feglm.jl")

    # instantiate R's alpaca package
    # R"if(\"alpaca\" %in% rownames(installed.packages()) == FALSE) {install.packages(\"alpaca\")}"
    # R"if(\"formula.tools\" %in% rownames(installed.packages()) == FALSE) {install.packages(\"formula.tools\")}"

end


function lazyCheck(mdp::MDP, vi::ValueIteration)
  checkArgumentOrder(mdp)
  checkTransition(mdp)
  checkDiscretize(mdp, vi)
  info(string("mdp and value iteration solver passed basic checks"))
end

# Check that all |mdp| RangeVar variables have valid discretization schemes in |vi|
function checkDiscretize(mdp::MDP, vi::ValueIteration)

  for lazyvar in values(mdp.statemap)
    if isa(lazyvar, RangeVar)
      if !haskey(vi.statemap, lazyvar.varname)
        error(string(
          "state variable ", lazyvar.varname,
          " does not have a discretization scheme"))
      elseif lazyvar.maxval - lazyvar.minval < vi.statemap[lazyvar.varname].step
        error(string(
          "state variable ", lazyvar.varname,
          " has a discretization step size larger than than its range"))
      end
    end
  end

  for lazyvar in values(mdp.actionmap)
    if isa(lazyvar, RangeVar)
      if !haskey(vi.actionmap, lazyvar.varname)
        error(string(
          "action variable ", lazyvar.varname,
          " does not have a discretization scheme"))
      elseif lazyvar.maxval - lazyvar.minval < vi.actionmap[lazyvar.varname].step
        error(string(
          "action variable ", lazyvar.varname,
          " has a discretization step size larger than than its range"))
      end
    end
  end

end


abstract type AbstractIterState{T<:Real} end
####Fixed point methods
#contains various state variables for running the solver
mutable struct IterFPState{TΘ⁺<:VolumePartsIter,T} <: AbstractIterState{T}
  Θ⁺::TΘ⁺
  λ⁺::T
  ftol::T

  griddelta::T
  gridstart::T
  gridend::T
  remainingpregriddelta::Vector{T}
  currentgriddelta::T

  starttime::Float64
  besttime::Float64
  maxelapsedsincebest::Float64
  maxelapsed::Float64

  iter::Int
  bestiter::Int
  maxitersincebest::Int
  maxiter::Int

  lastreporttime::Float64
  stop::Union{Nothing,Symbol}
end

#updates the problem state given the solution for an iteration
function updatefpstate!(Ψ::IterFPState{TΘ⁺,T}, Θ::VolumePartsIter, λ::T) where {TΘ⁺,T}
  Ψ.iter+=1
  t=time()
  if abs(Ψ.λ⁺ - λ) > Ψ.ftol
    Ψ.besttime=t
    Ψ.bestiter=Ψ.iter
  end

  if λ < Ψ.λ⁺ #keep track of the best solution, even if it doesn't meet ftol criteria
    Ψ.Θ⁺.Π .= Θ.Π
    Ψ.λ⁺ = λ
  end

  #adjusts the grid delta if a pre-conditioning glide path was provided
  if length(Ψ.remainingpregriddelta) > 0
    Ψ.currentgriddelta = popfirst!(Ψ.remainingpregriddelta)
    if length(Ψ.remainingpregriddelta) == 0
      Ψ.currentgriddelta = Ψ.griddelta #the last point should be the current grid
      @info "Grid delta pre-conditioning glide path complete"
    end
    Ψ.besttime=t
    Ψ.bestiter=Ψ.iter
    return nothing #this short-circuits the stopping criteria
  end

  #run through the stopping criteria
  elapsed::Int = round(time() - Ψ.starttime)
  elapsedsincebest::Int = round(time() - Ψ.besttime)
  itersincebest::Int = Ψ.iter - Ψ.bestiter
  if  elapsed > Ψ.maxelapsed
    @info "STOP: max time $elapsed/$(Ψ.maxelapsed) reached"
    Ψ.stop=:maxtime
  elseif elapsedsincebest>Ψ.maxelapsedsincebest
    @info "STOP: max time with minimal improvement $elapsedsincebest/$(Ψ.maxelapsedsincebest) reached"
    Ψ.stop=:maxtimesincebest
  elseif Ψ.iter ≥ Ψ.maxiter
    @info "STOP: max iter $(Ψ.iter) reached."
    Ψ.stop=:maxiter
  elseif itersincebest  ≥ Ψ.maxitersincebest
    @info "STOP: max iter since best $itersincebest reached."
    Ψ.stop=:maxitersincebest
  end

  return nothing
end

#contains some defaults for creating teh initial solver state
IterFPState(Θ::VolumePartsIter, ::Type{T};
  Θ⁺::VolumePartsIter = deepcopy(Θ),
  λ⁺::T=T(Inf),
  ftol::T=max(T(PARAM[:fpftol]), zero(T)),

  griddelta::T = PARAM[:fpgriddelta] |> T,
  gridstart::T = PARAM[:fpgridstart] |> T,
  gridend::T = PARAM[:fpgridend] |> T,
  remainingpregriddelta::Vector{T}= T.(PARAM[:fppregriddelta]),
  currentgriddelta::T= (
    length(remainingpregriddelta) > 0 ? popfirst!(remainingpregriddelta) : griddelta) |> T,

  starttime::Float64=time(),
  besttime::Float64=time(),
  maxelapsedsincebest::Float64=PARAM[:fpmaxelapsedsincebest],
  maxelapsed::Float64=PARAM[:fpmaxelapsed],

  iter::Int=0,
  bestiter::Int=0,
  maxitersincebest::Int=PARAM[:fpmaxitersincebest],
  maxiter::Int=PARAM[:fpmaxiter],

  lastreporttime::Float64=time(),
  stop::Union{Nothing,Symbol} = nothing,
  ) where T = IterFPState(Θ⁺, λ⁺, ftol,
    griddelta, gridstart, gridend, remainingpregriddelta, currentgriddelta,
    starttime, besttime, maxelapsedsincebest, maxelapsed,
    iter, bestiter, maxitersincebest, maxiter,
    lastreporttime, stop)

#reports the current status of the solver
function status(Ψ::IterFPState, λ; minreporttime::Float64 = PARAM[:iterminreporttime])
  (time()<minreporttime+Ψ.lastreporttime) && return nothing
  Ψ.lastreporttime=time()

  elapsed::Int = round(time() - Ψ.starttime)
  elapsedsincebest::Int = round(time() - Ψ.besttime)

  itersincebest::Int = Ψ.iter - Ψ.bestiter

  println("iter=$(Ψ.iter)/$(Ψ.maxiter), itersincebest=$itersincebest/$(Ψ.maxitersincebest)," *
    "loss=$λ, bestloss=$(Ψ.λ⁺), " *
    "time=$elapsed/$(Ψ.maxelapsed), " *
    "timesincebest = $elapsedsincebest/$(round(Ψ.maxelapsedsincebest)), " *
    "griddelta=$(Ψ.currentgriddelta)")

  return nothing
end

function solveA!(iterloss, Θ::VolumePartsIter,
    Xv::AbstractXYIter{TM,TV,T}, _Xv::XYIterAlloc{TM,TV,T}, RegressionType::Val, ::Val{:fixedpoint}
    ) where {TM,TV,T}

  Ψ::IterFPState = IterFPState(Θ, T)

  try
    while true #main solving loop

      #update the growths and values
      λ::T = updateA!(iterloss, Θ,Xv,_Xv,RegressionType, ΔG=Ψ.currentgriddelta)

      #update the problem state
      updatefpstate!(Ψ,Θ,λ)

      #check the stopping criteria and report status
      status(Ψ, λ)
      (!(Ψ.stop === nothing)) && break

    end
  catch err
    @warn "Aborted solver with error $err"
    Ψ.stop = :aborted
  finally #do this as a finallizer in case there is an error

    #ensure that we capture the best λ
    Θ.Π .= Ψ.Θ⁺.Π

    b = IOBuffer()
    elapsed::Int = round(time() - Ψ.starttime)
    elapsedsincebest::Int = round(time() - Ψ.besttime)
    itersincebest::Int = Ψ.iter - Ψ.bestiter
    write(b, "***Results info:\n" *
      "stop_reason: $(Ψ.stop)\n" *
      "iterations: $(Ψ.iter)/$(Ψ.maxiter)\n" *
      #"start_time: $(Ψ.starttime)\n" *
      "elapsed_at_completion: $(elapsed)/$(Ψ.maxelapsed)\n" *
      "elapsed since best: ~$(elapsedsincebest)/$(Ψ.maxelapsedsincebest)\n" *
      "iter since best: ~$(itersincebest)/$(Ψ.maxitersincebest)\n" *
      "Best resid: $(Ψ.λ⁺)\n")

    write(b, "\n\n***Other Simulation parameters:\n" *
      "regression method: $(RegressionType)\n" *
      "report interval: $(PARAM[:iterminreporttime])\n" *
      "Types (Matrix, Vector, T): $TM, $TV, $T\n" *
      "grid (excludes preconditioning) start:delta:end: $(Ψ.gridstart):$(Ψ.griddelta):$(Ψ.gridend)\n" *
      "currentgriddelta: $(Ψ.currentgriddelta)\n")

    if length(PARAM[:iterpregriddelta]) > 0
      write(b,
      "precondition max: $(maximum(PARAM[:iterpregriddelta]))\n" *
      "precondition length: $(length(PARAM[:iterpregriddelta]))\n"
      )
    end
    #compute the full jacobian with AD

    #∇Θver = gradient(()->loss(Θver,X,y),Flux.params(Θver))
    #rmsver = sqrt(mean(vec∇(Θver, ∇Θver).^2))

    #write(b, "\n\nlst squares gradient rms: $rmsver\n")

    iterinfo = String(take!(b))
    println(iterinfo)
    open("$(PARAM[:outputpath])\\$(iterresultsname(Ψ.λ⁺))_info.txt", "w+") do f
      write(f, iterinfo)
    end

    return Ψ.λ⁺
  end

  @assert false #should never hit this point due to the finalizer
end


"""
```
estimate(m, data; verbose=:low, proposal_covariance=Matrix())
```

Estimate the DSGE parameter posterior distribution.

### Arguments:
- `m::AbstractModel`: model object

### Optional Arguments:
- `data`: well-formed data as `Matrix` or `DataFrame`. If this is not provided, the `load_data` routine will be executed.

### Keyword Arguments:
- `verbose::Symbol`: The desired frequency of function progress messages printed to standard out.
   - `:none`: No status updates will be reported.
   - `:low`: Status updates will be provided in csminwel and at each block in
     Metropolis-Hastings.
   - `:high`: Status updates provided at each iteration in Metropolis-Hastings.
- `proposal_covariance::Matrix`: Used to test the metropolis_hastings algorithm with a precomputed
  covariance matrix for the proposal distribution. When the Hessian is singular,
  eigenvectors corresponding to zero eigenvectors are not well defined, so eigenvalue
  decomposition can cause problems. Passing a precomputed matrix allows us to ensure that
  the rest of the routine has not broken.
- `mle`: Set to true if parameters should be estimated by maximum likelihood directly.
    If this is set to true, this function will return after estimating parameters.
- `run_MH`: Set to false to disable sampling from the posterior.
"""
function estimate(m::AbstractModel, df::DataFrame;
                  verbose::Symbol=:low,
                  proposal_covariance::Matrix=Matrix(0,0),
                  mle::Bool = false,
                  run_MH::Bool = true)
    data = df_to_matrix(m, df)
    estimate(m, data; verbose=verbose, proposal_covariance=proposal_covariance,
             mle = mle, run_MH = run_MH)
end
function estimate(m::AbstractModel;
                  verbose::Symbol=:low,
                  proposal_covariance::Matrix=Matrix(0,0),
                  mle::Bool = false,
                  run_MH::Bool = true)
    # Load data
    df = load_data(m; verbose=verbose)
    estimate(m, df; verbose=verbose, proposal_covariance=proposal_covariance,
             mle = mle, run_MH = run_MH)
end
function estimate(m::AbstractModel, data::Matrix{Float64};
                  verbose::Symbol=:low,
                  proposal_covariance::Matrix=Matrix(0,0),
                  mle::Bool = false,
                  run_MH::Bool = true)

    ########################################################################################
    ### Step 1: Find posterior/likelihood mode (if reoptimizing, run optimization routine)
    ########################################################################################

    # Specify starting mode

    vint = get_setting(m, :data_vintage)
    if reoptimize(m)
        println("Reoptimizing...")

        # Inputs to optimization algorithm
        n_iterations       = get_setting(m, :optimization_iterations)
        ftol               = get_setting(m, :optimization_ftol)
        xtol               = get_setting(m, :optimization_xtol)
        gtol               = get_setting(m, :optimization_gtol)
        step_size          = get_setting(m, :optimization_step_size)
        converged          = false

        # If the algorithm stops only because we have exceeded the maximum number of
        # iterations, continue improving guess of modal parameters
        total_iterations = 0
        optimization_time = 0
        max_attempts = get_setting(m, :optimization_attempts)
        attempts = 1

        while !converged
            tic()
            out, H = optimize!(m, data;
                               method = get_setting(m, :optimization_method),
                               ftol=ftol, grtol = gtol, xtol = xtol,
                               iterations=n_iterations, show_trace=true, step_size=step_size,
                               verbose=verbose,
                               mle = mle)

            attempts += 1
            total_iterations += out.iterations
            converged = out.converged || attempts > max_attempts

            if VERBOSITY[verbose] >= VERBOSITY[:low]
                @printf "Total iterations completed: %d\n" total_iterations
                @printf "Optimization time elapsed: %5.2f\n" optimization_time += toq()
            end

            # Write params to file after every `n_iterations` iterations
            params = map(θ->θ.value, m.parameters)
            h5open(rawpath(m, "estimate", "paramsmode.h5"),"w") do file
                file["params"] = params
            end
        end

        # write parameters to file one last time so we have the final mode
        h5open(rawpath(m, "estimate", "paramsmode.h5"),"w") do file
            file["params"] = params
        end
    end

    params = map(θ->θ.value, m.parameters)

    # Sampling does not make sense if mle=true
    if mle || !run_MH
        return nothing
    end

    ########################################################################################
    ### Step 2: Compute proposal distribution
    ###
    ### In Metropolis-Hastings, we draw sample parameter vectors from
    ### the proposal distribution, which is a degenerate multivariate
    ### normal centered at the mode. Its variance is the inverse of
    ### the hessian. We find the inverse via eigenvalue decomposition.
    ########################################################################################

    # Calculate the Hessian at the posterior mode
    hessian = if calculate_hessian(m)
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            println("Recalculating Hessian...")
        end

        hessian, _ = hessian!(m, params, data; verbose=verbose)

        h5open(rawpath(m, "estimate","hessian.h5"),"w") do file
            file["hessian"] = hessian
        end

        hessian

    # Read in a pre-calculated Hessian
    else
        fn = hessian_path(m)
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            println("Using pre-calculated Hessian from $fn")
        end

        hessian = h5open(fn,"r") do file
            read(file, "hessian")
        end

        hessian
    end

    # Compute inverse hessian and create proposal distribution, or
    # just create it with the given cov matrix if we have it
    propdist = if isempty(proposal_covariance)
        # Make sure the mode and hessian have the same number of parameters
        n = length(params)
        @assert (n, n) == size(hessian)

        # Compute the inverse of the Hessian via eigenvalue decomposition
        S_diag, U = eig(hessian)
        big_eig_vals = find(x -> x > 1e-6, S_diag)
        rank = length(big_eig_vals)

        S_inv = zeros(n, n)
        for i = (n-rank+1):n
            S_inv[i, i] = 1/S_diag[i]
        end

        hessian_inv = U*sqrt.(S_inv) #this is the inverse of the hessian
        DSGE.DegenerateMvNormal(params, hessian_inv)
    else
        DSGE.DegenerateMvNormal(params, proposal_covariance)
    end

    if DSGE.rank(propdist) != n_parameters_free(m)
        println("problem –    shutting down dimensions")
    end

    ########################################################################################
    ### Step 3: Sample from posterior using Metropolis-Hastings algorithm
    ########################################################################################

    # Set the jump size for sampling
    cc0 = get_setting(m, :mh_cc0)
    cc = get_setting(m, :mh_cc)

    metropolis_hastings(propdist, m, data, cc0, cc; verbose=verbose)

    ########################################################################################
    ### Step 4: Calculate and save parameter covariance matrix
    ########################################################################################

    compute_parameter_covariance(m)

    return nothing
end

"""
```
metropolis_hastings{T<:AbstractFloat}(propdist::Distribution, m::AbstractModel,
    data::Matrix{T}, cc0::T, cc::T; verbose::Symbol = :low)
```

Implements the Metropolis-Hastings MCMC algorithm for sampling from the posterior
distribution of the parameters.

### Arguments

- `propdist`: The proposal distribution that Metropolis-Hastings begins sampling from.
- `m`: The model object
- `data`: Data matrix for observables
- `cc0`: Jump size for initializing Metropolis-Hastings.
- `cc`: Jump size for the rest of Metropolis-Hastings.

### Optional Arguments

- `verbose`: The desired frequency of function progress messages printed to
  standard out. One of:

```
   - `:none`: No status updates will be reported.
   - `:low`: Status updates provided at each block.
   - `:high`: Status updates provided at each draw.
```
"""
function metropolis_hastings{T<:AbstractFloat}(propdist::Distribution,
                                               m::AbstractModel,
                                               data::Matrix{T},
                                               cc0::T,
                                               cc::T;
                                               verbose::Symbol=:low)

    # If testing, set the random seeds at fixed numbers
    if m.testing
        srand(m.rng, 654)
    end

    # Set number of draws, how many we will save, and how many we will burn
    # (initialized here for scoping; will re-initialize in the while loop)
    n_blocks = 0
    n_sim    = 0
    mhthin   = 0
    n_burn   = 0
    n_params = n_parameters(m)

    # Initialize algorithm by drawing para_old from a normal distribution centered on the
    # posterior mode until the parameters are within bounds or the posterior value is
    # sufficiently large.
    para_old = rand(propdist, m; cc=cc0)
    post_old = -Inf

    n_blocks = n_mh_blocks(m)
    n_sim    = n_mh_simulations(m)
    n_burn   = n_mh_burn(m)
    mhthin   = mh_thin(m)

    initialized = false
    while !initialized
        post_old = posterior!(m, para_old, data; mh = true)
        if post_old > -Inf
            propdist.μ = para_old
            initialized = true
        else
            para_old = rand(propdist, m; cc=cc0)
        end
    end

    # Report number of blocks that will be used
    if VERBOSITY[verbose] >= VERBOSITY[:low]
        println("Blocks: $n_blocks")
        println("Draws per block: $n_sim")
    end

    # For n_sim*mhthin iterations within each block, generate a new parameter draw.
    # Decide to accept or reject, and save every (mhthin)th draw that is accepted.
    all_rejections = 0

    # Initialize matrices for parameter draws and transition matrices
    mhparams = zeros(n_sim, n_parameters(m))

    # Open HDF5 file for saving parameter draws
    simfile = h5open(rawpath(m,"estimate","mhsave.h5"),"w")
    n_saved_obs = n_sim * (n_blocks - n_burn)
    parasim = d_create(simfile, "mhparams", datatype(Float32),
                       dataspace(n_saved_obs,n_params),
                       "chunk", (n_sim,n_params))

    # Keep track of how long metropolis_hastings has been sampling
    total_sampling_time = 0.

    for block = 1:n_blocks

        tic()
        block_rejections = 0

        for j = 1:(n_sim*mhthin)

            # Draw para_new from the proposal distribution
            para_new = rand(propdist, m; cc = cc)

            # Solve the model (checking that parameters are within bounds and
            # gensys returns a meaningful system) and evaluate the posterior
            post_new = posterior!(m, para_new, data; mh = true)

            if VERBOSITY[verbose] >= VERBOSITY[:high]
                println("Block $block, Iteration $j: posterior = $post_new")
            end

            # Choose to accept or reject the new parameter by calculating the
            # ratio (r) of the new posterior value relative to the old one We
            # compare min(1, r) to a number drawn randomly from a uniform (0, 1)
            # distribution. This allows us to always accept the new draw if its
            # posterior value is greater than the previous draw's, but it gives
            # some probability to accepting a draw with a smaller posterior
            # value, so that we may explore tails and other local modes.
            r = exp(post_new - post_old)
            x = rand(m.rng)

            if x < min(1.0, r)
                # Accept proposed jump
                para_old = para_new
                post_old = post_new
                propdist.μ = para_new

                if VERBOSITY[verbose] >= VERBOSITY[:high]
                    println("Block $block, Iteration $j: accept proposed jump")
                end
            else
                # Reject proposed jump
                block_rejections += 1

                if VERBOSITY[verbose] >= VERBOSITY[:high]
                    println("Block $block, Iteration $j: reject proposed jump")
                end
            end

            # Save every (mhthin)th draw
            if j % mhthin == 0
                draw_index = convert(Int, j/mhthin)
                mhparams[draw_index, :]  = para_old'
            end
        end # of block

        all_rejections += block_rejections
        block_rejection_rate = block_rejections/(n_sim*mhthin)

        ## Once every iblock times, write parameters to a file

        # Calculate starting and ending indices for this block (corresponds to a new chunk in memory)
        block_start = n_sim*(block-n_burn-1)+1
        block_end   = block_start+n_sim-1

        # Write parameters to file if we're past n_burn blocks
        if block > n_burn
            parasim[block_start:block_end, :] = map(Float32, mhparams)
        end

        # Calculate time to complete this block, average block time, and
        # expected time to completion
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            block_time = toq()
            total_sampling_time += block_time
            total_sampling_time_minutes = total_sampling_time/60
            expected_time_remaining_sec     = (total_sampling_time/block)*(n_blocks - block)
            expected_time_remaining_minutes = expected_time_remaining_sec/60

            println("Completed $block of $n_blocks blocks.")
            println("Total time to compute $block blocks: $total_sampling_time_minutes minutes")
            println("Expected time remaining for Metropolis-Hastings: $expected_time_remaining_minutes minutes")
            println("Block $block rejection rate: $block_rejection_rate \n")
        end

    end # of loop over blocks

    close(simfile)

    rejection_rate = all_rejections / (n_blocks*n_sim*mhthin)
    if VERBOSITY[verbose] >= VERBOSITY[:low]
        println("Overall rejection rate: $rejection_rate")
    end
end

"""
```
compute_parameter_covariance(m::AbstractModel)
```

Calculates the parameter covariance matrix from saved parameter draws, and writes it to the
parameter_covariance.h5 file in the `workpath(m, "estimate")` directory.

### Arguments
* `m::AbstractModel`: the model object
"""
function compute_parameter_covariance(m::AbstractModel)

    # Read in saved parameter draws
    param_draws_path = rawpath(m,"estimate","mhsave.h5")
    if !isfile(param_draws_path)
        @printf STDERR "Saved parameter draws not found.\n"
        return
    end
    param_draws = h5open(param_draws_path, "r") do f
        read(f, "mhparams")
    end

    # Calculate covariance matrix
    param_covariance = cov(param_draws)

    # Write to file
    h5open(workpath(m, "estimate","parameter_covariance.h5"),"w") do f
        f["mhcov"] = param_covariance
    end

    return param_covariance
end

"""
```
get_estimation_output_files(m)
```

Returns a `Dict{Symbol, String}` with all files created by `estimate(m)`.
"""
function get_estimation_output_files(m::AbstractModel)
    output_files = Dict{Symbol, String}()

    for file in [:paramsmode, :hessian, :mhsave]
        output_files[file] = rawpath(m, "estimate", "$file.h5")
    end

    for file in [:paramsmean, :parameter_covariance]
        output_files[file] = workpath(m, "estimate", "$file.h5")
    end

    for file in [:priors, :prior_posterior_means, :moments]
        output_files[file] = tablespath(m, "estimate", "$file.tex")
    end

    return output_files
end


function event_occurred(obj::ODEEventHandler, arg0::ODEStateAndDerivative, arg1::jboolean)
    return jcall(obj, "eventOccurred", Action, (ODEStateAndDerivative, jboolean), arg0, arg1)
end

function g(obj::ODEEventHandler, arg0::ODEStateAndDerivative)
    return jcall(obj, "g", jdouble, (ODEStateAndDerivative,), arg0)
end

function init(obj::ODEEventHandler, arg0::ODEStateAndDerivative, arg1::jdouble)
    return jcall(obj, "init", void, (ODEStateAndDerivative, jdouble), arg0, arg1)
end

function reset_state(obj::ODEEventHandler, arg0::ODEStateAndDerivative)
    return jcall(obj, "resetState", ODEState, (ODEStateAndDerivative,), arg0)
end


module GeodesicBase

import Parameters: @with_kw
import SciMLBase

include("metric-params.jl")
include("physical-quantities.jl")
include("geodesic-solutions.jl")

export AbstractMetricParams, metric_params, metric

end # module


module PCA
    using Pluto

    run_notebook() = Pluto.run(notebook=joinpath(@__DIR__, "../notebook/pca.jl"))
end


"""
    ConvertEltype <: AugmentorBase.Operation

Description
--------------

Convert the element type of the given array/image into the given
`eltype`. This operation is especially useful for converting
color images to grayscale (or the other way around). That said,
the operation is not specific to color types and can also be used
for numeric arrays (e.g. with separated channels).

Note that this is an element-wise convert function. Thus it can
not be used to combine or separate color channels. Use
[`SplitChannels`](@ref) or [`CombineChannels`](@ref) for those
purposes.

Usage
--------------

    ConvertEltype(eltype)

Arguments
--------------

- **`eltype`** : The eltype of the resulting array/image.

See also
--------------

[`CombineChannels`](@ref), [`SplitChannels`](@ref), [`augment`](@ref)

Examples
--------------

```julia-repl
julia> using Augmentor, Colors

julia> A = rand(RGB, 10, 10) # three color channels
10×10 Array{RGB{Float64},2}:
[...]

julia> augment(A, ConvertEltype(Gray{Float32})) # convert to grayscale
10×10 Array{Gray{Float32},2}:
[...]
```
"""
struct ConvertEltype{T} <: Operation
    eltype::Type{T}
end

@inline supports_lazy(::Type{<:ConvertEltype}) = true

function applyeager(op::ConvertEltype{T}, input::AbstractArray, param) where T
    maybe_copy(convert(AbstractArray{T}, input))
end

function applylazy(op::ConvertEltype{T}, input::AbstractArray, param) where T
    mappedarray(c->convert(T,c), input)
end

function showconstruction(io::IO, op::ConvertEltype)
    print(io, typeof(op).name.name, '(')
    _showcolor(io, op.eltype)
    print(io, ')')
end

function Base.show(io::IO, op::ConvertEltype)
    if get(io, :compact, false)
        print(io, "Convert eltype to ")
        _showcolor(io, op.eltype)
    else
        print(io, "AugmentorBase.")
        showconstruction(io, op)
    end
end


import YAML
import Dolang
path = Pkg.dir("Dolo","src","linter.jl")
include(path)

filename = Pkg.dir("Dolo","linter/models/rbc_iid.yaml")
filename

#DL = DoloLinter
d = yaml_node_from_file(filename)
errswarrs = lint(filename)

model_equation_types = keys(d[:equations])
model_equations = []
model_equations_vec = []
for key in keys(d[:equations])
    n = length(d[:equations][key].value)
    for i=1:n
        # would be cooler to do directly for sym in syms[vg]
        eq = d[:equations][key].value[i]
        push!(model_equations_vec, (key, eq.value, eq))
        push!(model_equations, eq.value)
    end
end

model_symbol_types = keys(d[:symbols])
model_symbols = []
model_symbols_vec = []
for key in keys(d[:symbols])
    n = length(d[:symbols][key].value)
    for i=1:n
        # would be cooler to do directly for sym in syms[vg]
        sym = d[:symbols][key].value[i]
        push!(model_symbols_vec, (key, sym.value, sym))
        push!(model_symbols, sym.value)
    end
end

for eq in model_equations_vec
  if eq[1] == "arbitrage"

    equ , comp = split(eq[2], " | ")
    for sym in keys(Dolang.IncidenceTable(parse(equ)).by_var)

    end
  end
end

import Dolo: RECIPES

errors = LinterWarning[]
warnings = LinterWarning[]

eqg = "arbitrage"
spec = RECIPES[:dtcc][:specs][:arbitrage][:eqs]
Symbol_set = Dict(eqg => Set(([el[2] for el in spec if el[1]==eqg])) for eqg in Set(el[1] for el in spec))

for eq_ in d[:equations][eqg].value
      eqs = eq_.value
      equ , comp = split(eqs, " | ")
      inc = Dolang.IncidenceTable(parse(equ)).by_var
      for sym in keys(inc)
        sym = string(sym)
        if !(sym in model_symbols)  && !(string(sym) in keys(d[:definitions]))
          errvalue = sym
          errtype = "Unkown Symbol"
          loc = get_loc(eq_)
          msg = string(sym, "  in the '", eqg , "' equation is not defined in symbols or definitions sections")
          println(sym)
          push!(errors, LinterWarning(errvalue, errtype, msg, loc, filename))

        elseif  (string(sym) in model_symbols)
          floc = findfirst(isequal(sym), model_symbols)
          declared_sym_type = model_symbols_vec[floc][1]
          println(declared_sym_type)

          if !(inc[sym] in Symbol_set[declared_sym_type])
            errvalue = sym
          ##  errtype = "Timing convention error"
          ##  loc = get_loc(eq_)
          ##  msg = string(sym, "  in the '", eqg , "' equation ")
          ##  println(sym)
          ##  push!(errors, LinterWarning(errvalue, errtype, msg, loc, filename))
          end
        end

      end

end

function format_human(exc::Union{LinterException, LinterWarning})
    # Output should be a nice one line human-readable strings
    # colors, spacing, etc...
end

function format_human(errors::Vector{LinterException}, warnings::Vector{LinterWarning})
    # ouptputs all errors then all warnings, line-by-line
    # example from python's `dolo-lint --format=human examples/models/rbc_iid.yaml`
    # error: 12, 30: Symbol 'beta' already declared as 'parameters'. (pos (12, 16))
end

function format_dict(exc::Union{LinterException, LinterWarning})
    # Output should be a dictionary like:
    # {"type": "error", "source": "examples/models/rbc_iid.yaml", "range": [[11, 29], [11, 33]], "text": "Symbol 'beta' already declared as 'parameters'."}
end

function format_json(errors::Vector{LinterException}, warnings::Vector{LinterWarning})
    # Structure compatible with atom linter
    # see python's `dolo-lint --format=json examples/models/rbc_iid.yaml` for an example
    # [{"type": "error", "source": "examples/models/rbc_iid.yaml", "range": [[11, 29], [11, 33]], "text": "Symbol 'beta' already declared as 'parameters'. (pos (12, 16))"}]
end


module TestTableTraits
using Test, TableTraits, DataFrames

struct ColumnSource
end

TableTraits.supports_get_columns_copy_using_missing(::ColumnSource) = true

function TableTraits.get_columns_copy_using_missing(x::ColumnSource)
    return (a=[1,2,3], b=[4.,5.,6.], c=["A", "B", "C"])
end

@testset "TableTraits" begin

    df = DataFrame(ColumnSource())

    @test size(df)==(3,3)
    @test df.a==[1,2,3]
    @test df.b==[4.,5.,6.]
    @test df.c==["A", "B", "C"]

end

end


export get_state_timeseries, get_state_cumulative,
       get_vaccine_doses_administered_statewise, get_statewise_daily,
       get_active_cases_state

# 2. Statewise timeseries of Confirmed, Recovered and Deceased numbers.
function get_state_timeseries(state::String)
    df = read_CSV_url(states_csv_url)
    df.Active = df.Confirmed - df.Recovered - df.Deceased
    return df[df[!,:State].==state,:]
end

function get_state_timeseries()
    df = read_CSV_url(states_csv_url)
    df.Active = df.Confirmed - df.Recovered - df.Deceased
    return df
end

# 4. Statewise per day delta of Confirmed, Recovered and Deceased numbers.
function get_statewise_daily()
    df = read_CSV_url(state_wise_daily_csv_url)
    return df
end

# 5. Statewise cumulative numbers till date.
function get_state_cumulative(brief::Bool=true)
    df = read_CSV_url(state_wise_csv_url)
    if brief == true
        colindex = columnindex(df,:State_Notes)
        select!(df,Not(colindex))
    end
    return df
end

# Function to create state_code to state and vice versa dictionary objects.
function generate_dict_statecodes()
    df = get_state_cumulative()
    codeToStateDict = Dict{String,String}()
    stateToCodeDict = Dict{String,String}()
    nrows = size(df)[1]
    for i in 1:nrows
        state = df[i,:State]
        code = df[i,:State_code]
        codeToStateDict[code]=state
        stateToCodeDict[state]=code
    end

    return stateToCodeDict, codeToStateDict
end

state_to_code, code_to_state = generate_dict_statecodes()

# Function to calculate active cases on a particular date in all states.
function get_active_cases_state(date::Date)
    df = get_state_timeseries()
    df_today = df[df.Date.==date,:]
    sort!(df_today, :State)

    dateY = date-Dates.Day(1)# date yesterday
    df_yesterday = df[df.Date.==dateY,:]
    sort!(df_yesterday, :State)

    df_today.DeltaDeceased = df_today.Deceased - df_yesterday.Deceased
    df_today.DeltaActive = df_today.Active - df_yesterday.Active

    return df_today
end
# function calculate_active_cases_states(date::Date)
#     df = Covid19India.get_statewise_daily()
#     dfdate = df[df[!,:Date_YMD].==date,:]
#     dfout = DataFrame(State=String[], State_code=String[], Confirmed=Int[], Recovered=Int[], Deceased=Int[], Active=Int[], DeltaActive=Int[])
#     nstates = size(dfdate)[2]-3
#     colnames = names(dfdate)
#     for col_id in 1:nstates
#         state_code = colnames[col_id+3]
#         state=try
#             code_to_state[state_code]
#         catch e
#             continue
#         end
#         println(state)
#         confirmed = dfdate[dfdate.Status.=="Confirmed", state_code][1]
#         recovered = dfdate[dfdate.Status.=="Recovered", state_code][1]
#         deceased = dfdate[dfdate.Status.=="Deceased", state_code][1]

#         df_timeseries = Covid19India.get_state_timeseries(state)
#         #df_timeseries_today = df_timeseries[df_timeseries.Date.==date,:]
#         #df_timeseries_yesterday = df_timeseries[df_timeseries.Date=(date-Dates.Day(1))]
#         # active_today = df_timeseries_today[1,"Confirmed"]-df_timeseries_today[1,"Recovered"]-df_timeseries_today[1,"Deceased"]
#         # active_yesterday = df_timeseries_yesterday[1,"Confirmed"]-df_timeseries_yesterday[1,"Recovered"]-df_timeseries_yesterday[1,"Deceased"]
#         # active_delta = active_today - active_yesterday

#         # push!(dfout.State, code_to_state[state_code])
#         # push!(dfout.State_code, state_code)
#         # push!(dfout.Confirmed, confirmed)
#         # push!(dfout.Recovered, recovered)
#         # push!(dfout.Deceased, deceased)
#         # push!(dfout.Active, active_today)
#         # push!(dfout.DeltaActive, active_delta)
#     end

#     return df_timeseries
# end

# VACCINE DOSES ADMINISTERED.
# 12. Number of vaccine doses administered statewise [Reading]
function get_vaccine_doses_administered_timeseries_statewise()
    df = read_CSV_url(vaccine_doses_administered_statewise_csv_url)
    return df
end

function get_vaccine_doses_administered_timeseries_statewise(state::String)
    df = get_vaccine_doses_administered_timeseries_statewise()
    return df[df[!,:State].==state,:]
end

function get_vaccine_doses_administered_yesterday()
    df = get_vaccine_doses_administered_timeseries_statewise()
    #dttoday = Dates.format(Dates.now(), "dd/mm/yyyy")
    dtyesterday = Dates.format(Dates.now() - Dates.Day(1), "dd/mm/yyyy")
    dt_dby = Dates.format(Dates.now() - Dates.Day(2), "dd/mm/yyyy")
    dt_dbdby = Dates.format(Dates.now() - Dates.Day(3), "dd/mm/yyyy")
    df = df[!,["State",dt_dbdby, dt_dby,dtyesterday]]
    df.delta_today = df[!,4]-df[!,3]
    df.delta_yesterday = df[!,3]-df[!,2]
    df.delta_diff = df.delta_today - df.delta_yesterday
    sort!(df,:delta_today, rev=true)
    return df
end

# 13. Key data points from CoWin database at a state level
function get_cowin_vaccine_data_statewise(brief::Bool=false)
    df = read_CSV_url(cowin_vaccine_data_statewise_csv_url)
    if brief == true
        df = df[!,[]]
    end
    return df
end


struct composite{T<:Array{Float64}, S<:Float64}
    a::T
    b::S
end

struct composite2{T<:Array{Float64}, S<:Int64}
    a::T
    b::S
end

obj = composite([1.0;2;3], 5.6)
obj2 = composite2([1.0;3], 5)
function gg(x, y::Union{Int64, Float64, Nothing}=nothing)
    if y==nothing
        return x^2
    else
        return x
end
end


struct DurationResolver{T₁ <: Number} <: Function
    intersample_duration::T₁
    nested_interval::NestedInterval
end

function (resolve_duration::DurationResolver)(n::Nothing)
    # All times are in milliseconds
    nested_interval = resolve_duration.nested_interval
    tₛ = (get_stop(nested_interval) - get_start(nested_interval)) * resolve_duration.intersample_duration
    # Display timestamp.
    #timestamp = Dates.epochms2datetime(Dates.value(Millisecond(tₛ)))
    #timestr = Dates.format(timestamp , dateformat"HH:MM:SS")
    timestr = Dates.canonicalize(Dates.CompoundPeriod(Dates.Millisecond(tₛ)))
    string(timestr)
end


# plotting tools
include("generateVTK.jl")
include("vm2d_particle.jl")

# # set freestream
# function Uinf(x::Array{Float64,1})
#     Uinf = [0.0,0.0,0.0]
#     return Uinf
# end

# plot x-y view
function plot_xy(particles::Array{Particle,1}; velocity::Bool=false, xs=nothing, ys=nothing, zs=nothing, circulation::Bool=false, title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[1],particle.x[2],c=la.norm(particle.gamma))#la.norm(particle.gamma))
    end

    # velocity field
    if velocity
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        plt.quiver(gridxv, gridyv, vx, vy)
    end

    # circulation vectors
    if circulation
        gridxc, gridyc, gridzc, gx, gy, gz = gfield(particles)
        plt.quiver(gridxc, gridyc, gx, gy)
    end

    plt.title(title)
    plt.xlabel("x")
    plt.ylabel("y")
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

# plot y-z view
function plot_yz(particles::Array{Particle,1}; velocity::Bool=false, xs=nothing, ys=nothing, zs=nothing, circulation::Bool=false, title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[2],particle.x[3],c=la.norm(particle.gamma))
    end

    # velocity field
    if velocity
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        plt.quiver(gridyv, gridzv, vy, vz)
    end

    # circulation vectors
    if circulation
        gridxc, gridyc, gridzc, gx, gy, gz = gfield(particles)
        plt.quiver(gridyc, gridzc, gy, gz)
    end

    plt.title(title)
    plt.xlabel("y")
    plt.ylabel("z")
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

# plot y-z view
function plot_2d(particles::Array{Particle,1}, dims::String; velocity::Bool=false, xs=nothing, ys=nothing, zs=nothing, circulation::Bool=false, title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # set up axes
    xaxis = collect(dims)[1]
    yaxis = collect(dims)[2]
    if xaxis == 'x'
        xindex = 1
    elseif xaxis == 'y'
        xindex = 2
    elseif xaxis == 'z'
        xindex = 3
    end
    if yaxis == 'x'
        yindex = 1
    elseif yaxis == 'y'
        yindex = 2
    elseif yaxis == 'z'
        yindex = 3
    end

    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[xindex],particle.x[yindex],c=la.norm(particle.gamma))
    end

    # velocity field
    if velocity
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        gridv = [gridxv, gridyv, gridzv]
        v = [vx, vy, vz]
        plt.quiver(gridv[xindex], gridv[yindex], v[xindex], v[yindex])
    end

    # circulation vectors
    if circulation
        gridxc, gridyc, gridzc, gx, gy, gz = gfield(particles)
        gridc = [gridxc, gridyc, gridzc]
        g = [gx, gy, gz]
        plt.quiver(gridc[xindex], gridc[yindex], g[xindex], g[yindex])
    end

    plt.title(title)
    plt.xlabel(string(xaxis))
    plt.ylabel(string(yaxis))
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

# plot 3-D view
function plot_3d(particles::Array{Particle,1}; title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[2],particle.x[3],c=la.norm(particle.gamma))
    end
    plt.title(title)
    plt.xlabel("y")
    plt.ylabel("z")
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

Uinf(x) = [0.0,0.0,0.0] # default freestream velocity function
# animate
function animate(particles,timestart,timestep,numtimesteps; Uinf=Uinf)
    currenttime = timestart
    # set up plot
    plt.plot()
    plt.cool()
    plt.axis("equal")
    for particle in particles
        plt.scatter(particle.x[1],particle.x[2],c=particle.gamma[3])
    end
    plt.title("time: "*string(round(currenttime; digits=3)))
    # animate
    for i in range(1,length=numtimesteps)
        ## advance timestep
        advance(particles,timestep,Uinf)
        currenttime += timestep
        ## update plot
        plot_xy(particles; title="time: "*string(round(currenttime; digits=3)))
    end
end

"""
`packVTK(particles::Array{Particle,1},xs,ys,zs)`

Uses Eduardo Alvarez' `generateVTK()` function to generate VTK text files for visualizations with paraview.

Optional arguments `xs`, `ys`, and `zs` define the cartesian grid for visualizing the velocity field.
"""
function packVTK(filename, particles::Array{Particle,1}; xs=Nothing, ys=Nothing, zs=Nothing, num=nothing)
    # make array of points defining vortex particles
    points = Array{Float64,1}[]
    for particle in particles
        push!(points,particle.x)
    end

    # compile circulation and particle velocity vector data
    gammavecs = Array{Float64,1}[]
    particlevelocities = Array{Number,1}[]
    for particle in particles
        push!(gammavecs, particle.gamma)
        push!(particlevelocities, particle.v)
    end
    gammadata = Dict(
        "field_name" => "Circulation",
        "field_type" => "vector",
        "field_data" => gammavecs
    )
    particlevdata = Dict(
        "field_name" => "ParticleVelocity",
        "field_type" => "vector",
        "field_data" => particlevelocities
    )

    # compile all vtk data for particles
    point_data = [gammadata, particlevdata]

    # generate VTK file for particles
    generateVTK(filename, points; point_data = point_data, path = "vtk", num=num)

    if xs != Nothing
        # compile velocity field data
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        velocityvecs = Array{Number,1}[]
        velocitycoordinates = Array{Number,1}[]
        for i in range(1, length=length(gridxv))
            push!(velocityvecs, [vx[i], vy[i], vz[i]])
            push!(velocitycoordinates, [gridxv[i], gridyv[i], gridzv[i]])
        end
        velocitydata = Dict(
            "field_name" => "Velocity",
            "field_type" => "vector",
            "field_data" => velocityvecs
        )

        # compile all vtk data for velocity field
        velocity_data = [velocitydata]

        # generate VTK file for velocity field
        generateVTK(filename*"_velocityfield", velocitycoordinates; point_data = velocity_data, path = "vtk", num=num)
    end

end


mutable struct NNPolicy <: Policy
    pomdp::OCPOMDP
    nn::PyObject
    nstack::Int64
    input_shape::Tuple{Int64, Int64, Int64}
    action_map::Vector{OCAction}

    NNPolicy(pomdp::OCPOMDP, nn::PyObject) = new(pomdp, nn, 4, (84,84,1), ordered_actions(pomdp))
end

mutable struct NNBelief
    n::Int64 # number of observation to stack
    obs::Array{Float64}
end

mutable struct NNUpdater <: Updater
    problem::OCPOMDP
    n::Int64
    input_shape::Tuple{Int64, Int64, Int64}

    NNUpdater(problem::OCPOMDP) = new(problem, 4, (84,84,1))
end

function POMDPs.update(bu::NNUpdater, bold::NNBelief, a::Float64, o::OCObs)
    @assert bu.n == bold.n
    o_mat = convert_o(o, bu.problem)
    w, h = size(o_mat)
    new_obs = zeros(w, h, bu.n)
    # dequeue and append new obs
    for i=1:bu.n-1
        new_obs[:,:,i] = bold.obs[:,:,i+1]
    end
    new_obs[:,:,bu.n] = o_mat
    return NNBelief(bu.n, new_obs)
end

const OFF_KEY = -1

function POMDPs.update(bu::NNUpdater, bold::Dict{Int64,NNBelief}, a::Float64, o::Dict{Int64,OCObs})
    bnew = Dict{Int64, NNBelief}()
    for oid in keys(o)
        if haskey(bold, oid) && oid != OFF_KEY
            bnew[oid] = update(up, bold[oid], a, o[oid])
        elseif oid == OFF_KEY
            bnew[OFF_KEY] = update(up, bold[OFF_KEY], a, o[oid])
        else # cars appeared
            bnew[oid] = update(up, bold[OFF_KEY], a, o[oid])
        end
    end
    return bnew
end

function initialize_off_belief(policy::NNPolicy, ego::VehicleState)
    ped = get_off_the_grid(policy.pomdp)
    o0 = OCState(false, ego, ped)
    b0 = Dict{Int64, NNBelief}()
    o_mat = convert_o(o0, policy.pomdp)
    w, h = size(o_mat)
    o_stacked = reshape(repmat(o_mat, 1, policy.nstack), (w, h, policy.nstack))
    b0[OFF_KEY] = NNBelief(policy.nstack, o_stacked)
    return b0
end

# function normalize_o(s::OCState, pomdp::OCPOMDP)
#     o_vec = zeros(4)
#     o_vec[1], o_vec[2] = s.ego.posG.x/pomdp.x_goal, s.ego.v/pomdp.env.params.speed_limit
#     o_vec[3] = s.ped.posG.y/pomdp.y_goal
#     o_vec[4] = s.ped.v
#     if isinf(o_vec[4])
#         o_vec[4] = 0.
#     end
#     o_vec[4] /= pomdp.env.params.ped_max_speed
#     return o_vec
# end
#
# function normalize_o(x::Float64, v_ego::Float64, y::Float64, v_ped::Float64, pomdp::OCPOMDP)
#     o_vec = zeros(4)
#     o_vec[1], o_vec[2] = x/pomdp.x_goal, v_ego/pomdp.env.params.speed_limit
#     o_vec[3] = y/pomdp.y_goal
#     o_vec[4] = v_ped
#     if isinf(o_vec[4])
#         o_vec[4] = 0.
#     end
#     o_vec[4] /= pomdp.env.params.ped_max_speed
#     return o_vec
# end

function POMDPs.action(policy::NNPolicy, b::NNBelief)
    ai = policy.nn[:action](b.obs)
    return policy.action_map[ai]
end

function POMDPs.value(policy::NNPolicy, b::NNBelief)
    return policy.nn[:value](b.obs)
end

function action(policy::Policy, b::Dict{Int64, NNBelief})
    return POMDPs.action(policy, b).acc
end

const NNDriverModel = CrosswalkDriver{Float64, Dict{Int64,NNBelief},Dict{Int64, OCObs},NNPolicy,NNUpdater,POMDPSensor}

function reset_model!(model::NNDriverModel,
                      ego::Vehicle)
        model.tick = 0
        model.a = 0.
        model.b = initialize_off_belief(model.policy, ego.state)
        model.o = Dict{Int64, OCObs}()
        model.ego = ego
        reset_policy!(model.policy)
end

function reset_policy!(policy::NNPolicy)
    policy.nn[:reset]()
end


mutable struct GFExtendState
    prev_trace::GFTrace
    trace::GFTrace
    args_change::Any
    constraints::Any
    score::Float64
    weight::Float64
    visitor::AddressVisitor
    params::Dict{Symbol,Any}
    retchange::Union{Some{Any},Nothing}
    callee_output_changes::HomogenousTrie{Any,Any}
end

function GFExtendState(args_change, prev_trace, constraints, params)
    visitor = AddressVisitor()
    GFExtendState(prev_trace, GFTrace(), args_change, constraints, 0., 0.,
        visitor, params, nothing, HomogenousTrie{Any,Any}())
end

get_args_change(state::GFExtendState) = state.args_change

function set_ret_change!(state::GFExtendState, value)
    if state.retchange === nothing
        state.retchange = value
    else
        lightweight_retchange_already_set_err()
    end
end

function get_addr_change(state::GFExtendState, addr)
    get_leaf_node(state.callee_output_changes, addr)
end

function extend_no_change_err(addr)
    error("Attempted to change value of random choice at $addr during extend")
end

function addr(state::GFExtendState, dist::Distribution{T}, args, addr) where {T}
    visit!(state.visitor, addr)
    has_previous = has_primitive_call(state.prev_trace, addr)
    constrained = has_leaf_node(state.constraints, addr)
    if has_previous && constrained
        extend_no_change_err(addr)
    end
    if has_internal_node(state.constraints, addr)
        lightweight_got_internal_node_err(addr)
    end
    local retval::T
    local call::CallRecord
    if has_previous
        call = get_primitive_call(state.prev_trace, addr)
        if call.args != args
            extend_no_change_err(addr)
        end
        retval = call.retval
        score = call.score
    elseif constrained
        retval = get_leaf_node(state.constraints, addr)
        score = logpdf(dist, retval, args...)
        state.weight += score
        call = CallRecord(score, retval, args)
    else
        retval = random(dist, args...)
        score = logpdf(dist, retval, args...)
        call = CallRecord(score, retval, args)
    end
    state.trace = assoc_primitive_call(state.trace, addr, call)
    state.score += score
    if constrained && has_previous
        # there was a change and this was the previous value
        retchange = (true, prev_call.retval)
    elseif has_previous
        retchange = NoChange()
    else
        # retchange is null, because the address is new
        retchange = nothing
    end
    set_leaf_node!(state.callee_output_changes, addr, retchange)
    retval
end

function addr(state::GFExtendState, gen::Generator{T}, args, addr, args_change) where {T}
    visit!(state.visitor, addr)
    if has_internal_node(state.constraints, addr)
        constraints = get_internal_node(state.constraints, addr)
    elseif has_leaf_node(state.constraints, addr)
        lightweight_got_leaf_node_err(addr)
    else
        constraints = EmptyAssignment()
    end
    if has_subtrace(state.prev_trace, addr)
        prev_trace = get_subtrace(state.prev_trace, addr)
        (trace, weight, retchange) = extend(gen, args, args_change, prev_trace, constraints)
        set_leaf_node!(state.callee_output_changes, addr, retchange)
    else
        (trace, weight) = generate(gen, args, constraints)
    end
    call::CallRecord = get_call_record(trace)
    retval::T = call.retval
    state.trace = assoc_subtrace(state.trace, addr, trace)
    state.score += call.score
    state.weight += weight
    retval
end

function splice(state::GFExtendState, gen::GenFunction, args::Tuple)
    exec(gf, state, args)
end

function extend(gf::GenFunction, args, args_change, trace::GFTrace, constraints)
    state = GFExtendState(args_change, trace, constraints, gf.params)
    retval = exec(gf, state, args)
    call = CallRecord(state.score, retval, args)
    state.trace.call = call
    (state.trace, state.weight, state.retchange)
end


x = 3
y = 2

println(2x)

# same as y / x
println(x \ y)

# pipe
println([1,2,3] |> x -> x * 10 |> sum)

# like lisp
println(*(10, 2, 20))

# like natural formula
println(2(10 * 10)x)