"""
Opaque data structure holding the whole set of parameters used to compute a routing.

# Internals

These are implementation details that could evolve at any time. Many parameters here are already documented in the main
public constructor for `RoutingData`.

The first four fields are the most dependable of all of them (i.e. they are likely to stay unchanged for a long period of time):

* `g`: topology.
* `k`: demands.
* `solver`.
* `name`.

Several fields configure the solving process.

* Some parameters are always used:
  * `model_type`.
  * `model_simplifications`. TODO: reintroduce them.
* Some parameters are sometimes used:
  * `sub_model_type`.
* Some parameters are only useful for oblivious routing:
  * `model_all_traffic_matrices`: only for cutting-plane implementation.
  * `model_exact_opt_d`: only for cutting-plane implementation.
  * `model_robust_reformulation_traffic_matrices`: only for dual-reformulation implementation.
* Some parameters are only useful for problems without uncertainty.
  * `traffic_matrix`.

Several fields configure the output process.

* `logmessage`
* `plot_final_results`: TODO: reintroduce.
* `plot_each_iteration`: TODO: reintroduce.
* `export_lps`
* `export_lps_on_error`
* `output_folder`

From these graphs, several data structures are derived. Some of them are lazily computed when required (hence the
`mutable` aspect of this structure).

* For path-based formulations, paths are stored and uniquely identified:
  * `paths_edges`: all the paths, represented as list of edges (Edges). Path indices are always consistent
    with `paths`. Used as follows: `paths_edges[path_id]` is a list of edges.
  * `demand_to_path_ids`: maps a demand (arc from source to destination) to its path IDs.
  * `path_id_to_demand`: maps a path ID to the corresponding demand.
* In order to draw the graph, node positions are computed at most once:
  * `locs_x`: for each node, its x coordinate when drawing the graph (lazily created).
  * `locs_y`: for each node, its y coordinate when drawing the graph (lazily created).
  * `locs_f`: function that computes the two previous fields when requested. Takes a graph as input and outputs
    two vectors, `locs_x` and `locs_y`.

The time to create the `RoutingData` object is memorised in `time_precompute_ms`.
"""
mutable struct RoutingData
    g::MetaDiGraph
    k::AbstractGraph
    solver
    name::String
    time_precompute_ms::Float64

    model_type::ModelType
    sub_model_type::ModelType
    model_simplifications::Bool
    model_all_traffic_matrices::Bool
    model_exact_opt_d::Bool
    model_robust_reformulation_traffic_matrices::Bool

    traffic_matrix::Dict{Edge{Int}, Float64}

    paths_edges::Vector{Vector{Edge}}
    demand_to_path_ids::Dict{Edge{Int}, Vector{Int}}
    path_id_to_demand::Dict{Int, Edge}

    locs_x::Union{Vector{Float64}, Nothing}
    locs_y::Union{Vector{Float64}, Nothing}
    locs_f

    logmessage
    plot_final_results
    plot_each_iteration
    export_lps
    export_lps_on_error
    output_folder
end

"""
Creates an opaque `RoutingData` object based on the topology of the network `g` and the demands to route `k`.
Calling this function may take a while, as it precomputes many things used later on, based on the type of model to use.
Not all parameters are used in all cases.

Parameters that are always used (only the first three are required and positional):

* `g` is a directed graph with annotation (hence the type `MetaDiGraph`). Each edge is supposed to have a property
  `capacity` giving its capacity (units are not imposed by the package, but must be consistent). Each node must
  have a `name` property, which will be used for generating the output.
* `k` is another directed graph, without necessarily annotations (hence the type `AbstractSimpleGraph`).
* `solver` is the optimization solver to use for this problem (like `CPLEX.Optimizer`).
* `model_type` is the type of optimisation model used to compute the oblivious routing. Its type is `ModelType`.
* `name` is an optional name for the model.

The parameters almost always apply, and may have tremendous impact on solving performance:

* `model_simplifications` enables simplifications in the model. They do not imply any kind of approximation.
  Computing what parts of the model can be simplified may take more time than what the simplifications may improve
  the running times of the optimalisation part (especially for small networks and those with very high connexity).
* `remove_unreachable_nodes` determines whether the nodes in the topology `g` that are isolated, i.e. those that
  have no neighbours, should be removed
* `remove_unsatisfiable_demands` determines whether the demands that are not routable (i.e. there is no path in `g`
  from their source to their destination, the topology `g` provides no connectivity between these two nodes) should
  be removed.

Some models may use these parameters:

* `sub_model_type` is only used by algorithms using subproblems, like oblivious routing with cutting planes.
  All the following parameters are reused for the subproblems: `model_simplifications`, `npaths`.
* `npaths` is the number of shortest paths that should be computed beforehand for path-based formulations, solved with
  column generation or not. This number of paths is considered *for each demand*: there will be
  probably more than `npaths` that get generated, but never more than `npaths` times the number of demands).
  By default, 20 paths at most are generated for each demand.
  For column-generation-based formulations, this provides a starting set of columns: this set may be increased in size
  during the solving process. Otherwise, this is the total number of paths that will be used for solving the problem.

These parameters are only used for oblivious routing:

* `model_all_traffic_matrices` indicates whether all traffic matrices that violate the oblivious routing constraints
  should be added or not, at each iteration of the master problem. (The original algorithm only adds one matrix,
  but adding more at a time may improve running times by reducing the total number of iterations to perform.)
* `model_exact_opt_d` indicates whether, in the added constraints, the optimal congestion is computed for the generated
  traffic matrix. In theory, this value is always 1.0.
* `model_robust_reformulation_traffic_matrices` enables computing traffic matrices for the robust reformulations.
  Using this options increases the number of constraints in the problem to solve in order to get the right dual values.

These parametesr are only used for problems without uncertainty.

* `traffic_matrix` is the only traffic matrix that is considered.

These parameters only control the output of the solving process:

* `locs_f` function that computes the position of the nodes when plotting the graph in 2D. It takes a graph as input and
  outputs two vectors, first the x positions, then the y positions.
* `logmessage` is a function that will be called whenever the package produces any textual output. By default,
  text is redirected to the shell (`println`). This function should take a single string argument, the message to print.
* `verbose` indicates whether this function (not the complete package) should be explicit about what it is doing
  (for instance, if it removes demands). This option does not include regular progress logging for long operations.
* `plot_final_results`: whether plots should be produced when the solving process is done.
* `plot_each_iteration`: for iterative algorithms, whether plots should be produced at each iteration.
* `export_lps`: export the LP files (if possible) for each and every (sub)problem that is solved. If the master problem
  is iteratively built, each iteration will be exported. As the LP format can only represent MILP problems, not all
  combinations of problem types and solving algorithms can benefit from this parameter.
* `export_lps_on_error`: export the LP files (if possible) of the (sub)problem being solved when an error occurs.
  As the LP format can only represent MILP problems, not all combinations of problem types and solving algorithms
  can benefit from this parameter.
* `output_folder`: the folder where all file output is supposed to be taking place. This folder must exist beforehand
  or be an empty string to disable all output.
"""
function RoutingData(g::AbstractMetaGraph, # Not MetaDiGraph, to show a more interesting error message if the graph is not directed.
                     k::AbstractGraph,
                     solver,
                     model_type::ModelType;
                     name::String="",
                     sub_model_type::Union{ModelType, Nothing}=nothing,
                     model_simplifications::Bool=false,
                     model_all_traffic_matrices::Bool=false,
                     model_exact_opt_d::Bool=false,
                     model_robust_reformulation_traffic_matrices::Bool=false,
                     traffic_matrix::Dict{Edge{Int}, Float64}=Dict{Edge{Int}, Float64}(),
                     npaths::Int=20,
                     remove_unreachable_nodes::Bool=true,
                     remove_unsatisfiable_demands::Bool=true,
                     locs_f=spring_layout,
                     logmessage=println,
                     verbose::Bool=true,
                     plot_final_results::Bool=false,
                     plot_each_iteration::Bool=false,
                     export_lps::Bool=false,
                     export_lps_on_error::Bool=false,
                     output_folder::String="")
    start = time_ns()

    if sub_model_type === nothing
        sub_model_type = copy(model_type)
    end

    # Check whether the right properties have been defined in the graphs.
    for v in vertices(g)
        if ! has_prop(g, v, :name)
            error("Node $(v) has no name!")
        end
    end
    for e in edges(g)
        if ! has_prop(g, e, :capacity)
            error("Edge $(e) has no capacity!")
        end
    end

    if ! is_directed(g)
        error("The demand graph is not directed!")
    end
    if ! is_directed(k)
        error("The demand graph is not directed!")
    end

    # Check whether the other parameters have easily-checkable properties, before more cumbersome computations.
    if output_folder != "" && ! isdir(output_folder)
        error("The output path $(output_folder) does not exist!")
    end

    # Copy the graphs locally, so that the caller does not see when they are modified internally
    # (for instance, due to remove_unsatisfiable_demands).
    g = copy(g)
    k = copy(k)

    # If need be, remove some vertices because they are isolated.
    if remove_unreachable_nodes
        nremoved = remove_unreachable_nodes!(g)
        if verbose && nremoved > 0
            logmessage("Removed $(nremoved) vertices because they were not reachable.")
        end
    end

    # If need be, remove some demands due to lack of connectivity.
    if remove_unsatisfiable_demands
        nremoved = remove_unsatisfiable_demands!(g, k)
        if verbose && nremoved > 0
            logmessage("Removed $(nremoved) demands because they were not routable.")
        end
    end

    # If this is a path-based formulation, it may require some path generation.
    if (model_type.type == PathFormulation() || sub_model_type.type == PathFormulation()) && npaths > 0
        # Precompute some shortest paths and directly assign them numbers.
        # A path is a sequence of integers, each being the index of a node.
        paths = Vector{Int}[]
        demand_to_path_ids = Dict{Edge{Int}, Vector{Int}}()
        path_id_to_demand = Dict{Int, Edge}()

        for d in edges(k)
            demand_paths = yen_k_shortest_paths(g, src(d), dst(d), weights(g), npaths).paths
            if ! (model_type.cg || sub_model_type.cg) && length(demand_paths) == npaths
                @warn("The limit of paths to generate ($(npaths)) has been reached; the result will not be guaranteed to be optimum. " *
                      "Switch to a column-generation formulation for exact results.")
            end

            demand_to_path_ids[d] = Int[]
            for path in demand_paths
                push!(paths, path)
                p_id = length(paths)
                push!(demand_to_path_ids[d], p_id)
                path_id_to_demand[p_id] = d
            end
        end

        # Paths are returned as a list of nodes, transform it into a list of edges.
        paths_edges = [map(t -> Edge(t[1], t[2]), zip(p[1:end-1], p[2:end])) for p in paths]
    else
        # Not a path-based formulation: create empty collections of the right type.
        paths_edges = Vector{Edge}[]
        demand_to_path_ids = Dict{Edge{Int}, Vector{Int}}()
        path_id_to_demand = Dict{Int, Edge}()
    end

    stop = time_ns()

    return RoutingData(g, k, solver, name, (stop - start) / 1_000_000.,
                       # Solving parameters.
                       model_type, sub_model_type, model_simplifications,
                       model_all_traffic_matrices, model_exact_opt_d,
                       model_robust_reformulation_traffic_matrices,
                       traffic_matrix,
                       # Precomputed things.
                       paths_edges, demand_to_path_ids, path_id_to_demand,
                       # Output parameters (graph plotting and logging).
                       nothing, nothing, locs_f, logmessage, plot_final_results,
                       plot_each_iteration, export_lps, export_lps_on_error, output_folder)
end

function remove_unreachable_nodes!(g::AbstractGraph)
    # First, determine which nodes are unreachable; then, delete them.
    # If deletion is made in the loop over all vertices, as soon as a deletion
    # is performed, the vertex indices may change: vertices(g) is no more
    # consistent with inneighbors(g, v) and outneighbors(g, v).
    to_remove = Int[]
    for v in vertices(g)
        if length(inneighbors(g, v)) == 0 && length(outneighbors(g, v)) == 0
            push!(to_remove, v)
        end
    end
    for v in to_remove
        rem_vertex!(g, v)
    end
    return length(to_remove)
end

function remove_unsatisfiable_demands!(g::AbstractGraph, k::AbstractGraph)
    to_remove = Edge{Int}[]
    for e in edges(k)
        paths = dijkstra_shortest_paths(g, src(e), weights(g))
        if length(paths.parents) < dst(e) || paths.parents[dst(e)] == 0
            push!(to_remove, e)
        end
    end
    for e in to_remove
        rem_edge!(k, e)
    end
    return length(to_remove)
end

function add_path!(rd::RoutingData, demand::Edge, path::Vector{Edge})
    # Update RoutingData with the new path.
    # All demands already have a path, by assumption. Otherwise, the
    # master problem should not be feasible. Hence, computing the new
    # path ID is easy: just take the length of the vector.
    push!(rd.paths_edges, path)
    path_id = length(rd.paths_edges)

    rd.path_id_to_demand[path_id] = demand
    push!(rd.demand_to_path_ids[demand], path_id)

    return path_id
end

graph(rd::RoutingData) = rd.g
edges(rd::RoutingData) = edges(rd.g)
vertices(rd::RoutingData) = vertices(rd.g)
demands(rd::RoutingData) = edges(rd.k)

n_nodes(rd::RoutingData) = nv(rd.g)
n_edges(rd::RoutingData) = ne(rd.g)
n_demands(rd::RoutingData) = ne(rd.k)
n_paths(rd::RoutingData) = length(rd.paths_edges)

capacity(rd::RoutingData, e::Edge{Int}) = get_prop(rd.g, e, :capacity)

function find_path_ids_with_edge(rd::RoutingData, e::Edge)
    if n_paths(rd) > 0
        return filter(1:n_paths(rd)) do p
            e in rd.paths_edges[p]
        end
    else
        error("Precomputed paths are not available. Are you trying to use this function in a flow-based formulation?")
    end
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

include("vonmises.jl")

# Elasticity problems
abstract ElasticPlasticProblem            <: AbstractProblem
abstract PlaneStressElasticPlasticProblem <: ElasticPlasticProblem

function get_unknown_field_name{P<:ElasticPlasticProblem}(::Type{P})
    return "displacement"
end

function get_unknown_field_type{P<:ElasticPlasticProblem}(::Type{P})
    return Vector{Float64}
end

# 3D Elasticity problems
function ElasticPlasticProblem(dim::Int=3, elements=[])
    return Problem{ElasticPlasticProblem}("elasticplastic problem", dim, elements)
end

# 2D Plane stress elasticity problems
function PlaneStressElasticPlasticProblem(dim::Int=2, elements=[])
    return Problem{PlaneStressElasticPlasticProblem}("plane stress elasticplastic problem", dim, elements)
end

function get_residual_vector{P<:PlaneStressElasticPlasticProblem}(problem::Problem{P}, element::Element, ip::IntegrationPoint, time::Number; variation=nothing)
    r = zeros(Float64, problem.dim, length(element))

    J = get_jacobian(element, ip, time)

    # internal forces
    if haskey(element, "youngs modulus") && haskey(element, "poissons ratio")
        if !haskey(element, "integration points")
            if P == PlaneStressElasticPlasticProblem
                last_stress = zeros(2,2)
                last_strain = zeros(2,2)
            else
                last_stress = zeros(3,3)
                last_strain = zeros(3,3)
            end
        else
            for each_ip in element("integration points", time)
                if isapprox(each_ip.xi, ip.xi)
                    last_stress = ip("stress", time)
                    last_strain = ip("stress", time)
                    break
                end
            end
        end
        u = element("displacement", time, variation)
        grad = element(ip, time, Val{:grad})
        gradu = grad*u

        # deformation gradient
        F = I + gradu
        E = 1/2*(F'*F - I)

        #E = 1/2*(gradu + gradu') # finite strain (total)

        # material
        young = element("youngs modulus", ip, time)
        poisson = element("poissons ratio", ip, time)
        stress_y = element("yield stress", time).data
        dstrain = E - last_strain
        de_v = [dstrain[1,1], dstrain[2,2], dstrain[1,2]]
        material_model = element("material model", time)
        s = last_stress
        de = copy(ForwardDiff.get_value(dstrain))

        if P == PlaneStressElasticPlasticProblem
            C = stiffnessTensorPlaneStress(young, poisson)
            s_v = [s[1,1], s[2,2], s[1,2]]
            de_ = [de[1,1], de[2,2], de[1,2]]
            problem_stress_type = :PlaneStressElasticPlasticProblem
        else
            C = stiffnessTensor(young, poisson)
            s_v = [s[1,1], s[2,2], s[3,3], s[2,3], s[1,3], s[1,2]]
            de_ = [de[1,1], de[2,2], de[3,3], de[2,3], de[1,3], de[1,2]]
            problem_stress_type = :ElasticPlasticProblem
        end
        dep = zeros(3)
        stress_inc, dep = calculate_stress(de_,
                                           s_v,
                                           C,
                                           stress_y,
                                           Val{:vonMises},
                                           Val{problem_stress_type})
        info("%% ", dep)
        s_v += C * (de_v - dep)
        info("--: ", ForwardDiff.get_value(s_v))
        # stress
        if P == PlaneStressElasticPlasticProblem
            S = [s_v[1] s_v[3];
                 s_v[3] s_v[2]]
        else
        S = [s_v[1] s_v[6] s_v[5];
             s_v[6] s_v[2] s_v[4];
             s_v[5] s_v[4] s_v[3]]
        end
        r += F*S*grad*det(J)
    end

    # external forces - volume load
    if haskey(element, "displacement load")
        basis = element(ip, time)
        b = element("displacement load", ip, time)
        r -= b*basis*det(J)
    end

    # external forces - surface traction force
    if haskey(element, "displacement traction force")
        basis = element(ip, time)
        T = element("displacement traction force", ip, time)
        JT = transpose(J)
        s = size(JT, 2) == 1 ? JT : cross(JT[:,1], JT[:,2])
        r -= T*basis*norm(s)
    end

    return vec(r)
end

#=
function get_residual_vector{P<:ElasticPlasticProblem}(problem::Problem{P}, element::Element, ip::IntegrationPoint, time::Number; variation=nothing)
    r = zeros(Float64, problem.dim, length(element))

    J = get_jacobian(element, ip, time)

    info("_____________________")
    # internal forces
    if haskey(element, "youngs modulus") && haskey(element, "poissons ratio")

        if !haskey(element, "integration points")
            if P == PlaneStressElasticPlasticProblem
                last_stress = zeros(2,2)
                last_strain = zeros(2,2)
            else
                last_stress = zeros(3,3)
                last_strain = zeros(3,3)
            end
        else
            for each_ip in element("integration points", time)
                if isapprox(each_ip.xi, ip.xi)
                    last_stress = ip("stress", time)
                    last_strain = ip("stress", time)
                    break
                end
            end
        end
        u = element("displacement", time, variation)
        grad = element(ip, time, Val{:grad})
        gradu = grad*u

        # deformation gradient
        F = I + gradu

        # material
        young = element("youngs modulus", ip, time)
        poisson = element("poissons ratio", ip, time)
        mu = young/(2*(1+poisson))
        lambda = young*poisson/((1+poisson)*(1-2*poisson))
        if P == PlaneStressElasticityProblem
            lambda = 2*lambda*mu/(lambda + 2*mu)  # <- correction for 2d problems
        end

        # strain
        E = 1/2*(F'*F - I)
        #E = 1/2*(gradu + gradu') # finite strain (total)

        young = element("youngs modulus", ip, time)
        poisson = element("poissons ratio", ip, time)
        stress_y = element("yield stress", time).data
        dstrain = E - last_strain
        material_model = element("material model", time)
        s = last_stress
        de = ForwardDiff.get_value(dstrain)

        if P == PlaneStressElasticPlasticProblem
            C = stiffnessTensorPlaneStress(young, poisson)
            s_v = [s[1,1], s[2,2], s[1,2]]
            de_ = [de[1,1], de[2,2], de[1,2]]
            problem_stress_type = :PlaneStressElasticPlasticProblem
        else
            C = stiffnessTensor(young, poisson)
            s_v = [s[1,1], s[2,2], s[3,3], s[2,3], s[1,3], s[1,2]]
            de_ = [de[1,1], de[2,2], de[3,3], de[2,3], de[1,3], de[1,2]]
            problem_stress_type = :ElasticPlasticProblem
        end

        stress_inc, lambda = plastic_multiplier = calculate_stress(de_,
                                                           s_v,
                                                           C,
                                                           stress_y,
                                                           Val{:vonMises},
                                                           Val{problem_stress_type})

        # dep = lambda * dfds(s)
        # upate_material_parameters!(...)
        s_new = s_v + stress_inc
        #S = [s_v[1] s_v[6] s_v[5];
    #         s_v[6] s_v[2] s_v[4];
#             s_v[5] s_v[4] s_v[3]]
        S = [s_new[1] s_new[3];
             s_new[3] s_new[2]]
        # S = C * (E - dep)

        info("Stress: ", vec(ForwardDiff.get_value(S)))
        # stress
        #S = lambda*trace(E)*I + 2*mu*E

        r += F*S*grad*det(J)

    end

    # external forces - volume load
    if haskey(element, "displacement load")
        basis = element(ip, time)
        b = element("displacement load", ip, time)
        r -= b*basis*det(J)
    end

    # external forces - surface traction force
    if haskey(element, "displacement traction force")
        basis = element(ip, time)
        T = element("displacement traction force", ip, time)
        JT = transpose(J)
        s = size(JT, 2) == 1 ? JT : cross(JT[:,1], JT[:,2])
        r -= T*basis*norm(s)
    end

    return vec(r)
end
=# #fff


############################################################
# Split Layer
############################################################
@defstruct SplitLayer Layer (
  name :: AbstractString = "split",
  no_copy :: Bool = false,
  (tops :: Vector{Symbol} = Symbol[], length(tops) >= 1),
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == 1),
)
@characterize_layer(SplitLayer,
  can_do_bp => true
)

type SplitLayerState{N} <: LayerState
  layer      :: SplitLayer
  blobs      :: Vector{Blob}
  blobs_diff :: Vector{Blob}
end

function setup(backend::Backend, layer::SplitLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  N = length(layer.tops)

  if layer.no_copy
    # directly re-use the input blob
    blobs = Blob[inputs[1] for i = 1:N]
  else
    blobs = Array{Blob}(N)
    blobs[1] = inputs[1]
    for i = 2:N
      blobs[i] = make_blob(backend, eltype(inputs[1]), size(inputs[1]))
    end
  end
  blobs_diff = Array{Blob}(N)
  blobs_diff[1] = diffs[1] # re-use the first backward blob
  for i = 2:N
    if isa(diffs[1], NullBlob)
      blobs_diff[i] = NullBlob()
    else
      blobs_diff[i] = make_blob(backend, eltype(inputs[1]), size(inputs[1]))
    end
  end

  return SplitLayerState{N}(layer, blobs, blobs_diff)
end
function shutdown(backend::Backend, state::SplitLayerState)
  # some blobs are shared, but never mind, blob destroy function has
  # a guard that does not cause problems on double destroying
  map(destroy, state.blobs)
  map(destroy, state.blobs_diff)
end

function forward(backend::Backend, state::SplitLayerState, inputs::Vector{Blob})
  if !state.layer.no_copy
    for i = 2:length(state.blobs)
      copy!(state.blobs[i], inputs[1])
    end
  end
end

function backward{N}(backend::CPUBackend, state::SplitLayerState{N}, inputs::Vector{Blob}, diffs::Vector{Blob})
  if !isa(diffs[1], NullBlob)
    diff = diffs[1]
    len = length(diff)
    one = convert(eltype(diff), 1)
    for i = 2:N
      BLAS.axpy!(len, one, state.blobs_diff[i].data, 1, diff.data, 1)
    end
  end
end


using Test
using HTTP
using JSON2
using Carrot
using Mocking

Mocking.activate()

response = "{\"instrument\":\"EUR_USD\",\"granularity\":\"H1\",\"candles\":[{\"complete\":true,\"volume\":1013,\"time\":\"2020-06-26T04:00:00.000000000Z\",\"mid\":{\"o\":\"1.12227\",\"h\":\"1.12264\",\"l\":\"1.12211\",\"c\":\"1.12238\"}}]}"
data = JSON2.read(response)

@testset "fetch_data" begin
    @test create_time_series_metrix(data.candles) ==
          [1.12227 1.12264 1.12211 1.12238 1013.0]
    patch = @patch HTTP.request(a::String, b::String, c::Dict{String,String}; query="granularity=H1&count=1500", readtimeout=2) = HTTP.Response(200, body=Vector{UInt8}(response))
    apply(patch) do
        @test getCandles("EUR_GBP", "H1") == [1.12227 1.12264 1.12211 1.12238 1013.0]
    end

end


"""
```
OneAssetHANK{T} <: AbstractCTModel{T}
```

The `OneAssetHANK` type defines the structure of the one-asset HANK model
originally written in MATLAB by Ben Moll and SeHyoun Ahn at
https://sehyoun.com/EXAMPLE_one_asset_HANK_web.html

### Fields

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

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

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

#### Inputs to Measurement and Equilibrium Condition Equations

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

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

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

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

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

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

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

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

#### Model Specifications and Settings

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

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

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

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

#### Other Fields

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

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

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

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

    endogenous_states::OrderedDict{Symbol,Vector{Int}}      # these fields used to create matrices in the
    exogenous_shocks::OrderedDict{Symbol,Int}               # measurement & equilibrium condition equations
    expected_shocks::OrderedDict{Symbol,Vector{Int}}        #
    equilibrium_conditions::OrderedDict{Symbol,Vector{Int}} #
    endogenous_states_augmented::OrderedDict{Symbol,Int}    #
    observables::OrderedDict{Symbol,Int}                    #
    pseudo_observables::OrderedDict{Symbol,Int}             #

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

    observable_mappings::OrderedDict{Symbol, Observable}
    pseudo_observable_mappings::OrderedDict{Symbol, PseudoObservable}
end

description(m::OneAssetHANK) = "Julia implementation of the one-asset HANK model
originally written in MATLAB by Ben Moll and SeHyoun Ahn at
https://sehyoun.com/EXAMPLE_one_asset_HANK_web.html: OneAssetHANK, $(m.subspec)"

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

Arguments:
`m:: OneAssetHANK`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::OneAssetHANK)
    # Endogenous states
    endogenous_states = collect([:value_function, :inflation, :distribution,
                                 :monetary_policy, :w, :N, :C, :output, :B])

    # Exogenous shocks
    exogenous_shocks = collect([:mp_shock])

    # Expectations shocks
    expected_shocks = collect([:E_V, :E_π])

    # Equilibrium conditions
    equilibrium_conditions = collect([:value_function, :inflation, :distribution,
                                      :monetary_policy, :w, :N, :C, :Y, :B])

    # Additional states added after solving model
    # Lagged states and observables measurement error
    endogenous_states_augmented = []

    # Observables
    observables = keys(m.observable_mappings)

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

    # Assign dimensions
    m.endogenous_states[:value_function]  = collect(1:200)      # size of state space grid
    m.endogenous_states[:inflation]       = [201]               # inflation
    m.endogenous_states[:distribution]    = collect(202:400)    # size of grid minus 1 for ...
    m.endogenous_states[:monetary_policy] = [401]               # monetary policy shocks
    m.endogenous_states[:w]               = [402]               # static conditions from market-clearing
    m.endogenous_states[:N]               = [403]
    m.endogenous_states[:C]               = [404]
    m.endogenous_states[:output]          = [405]
    m.endogenous_states[:B]               = [406]
    m.expected_shocks[:E_V]               = m.endogenous_states[:value_function]
    m.expected_shocks[:E_π]               = m.endogenous_states[:inflation]

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

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

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

    # Initialize empty model
    m = OneAssetHANK{Float64}(
            # Model parameters and steady state values # before 2nd arg was Vector{Float64}()
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(),
            OrderedDict{Symbol,Int}(),

            # Model indices
            OrderedDict{Symbol,Vector{Int}}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Vector{Int}}(), OrderedDict{Symbol,Vector{Int}}(),
            OrderedDict{Symbol,Int}(),         OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Int}(),

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

    # Initialize parameters
    init_parameters!(m)
    init_model_indices!(m)

    # Set settings
    model_settings!(m)
    default_test_settings!(m)

    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Test to see if necessary values are defined, throw error otherwise
    try
        @assert typeof(get_setting(m, :reduce_state_vars)) == Bool
        @assert typeof(get_setting(m, :reduce_v)) == Bool
    catch
        error("Need to specify in settings whether to perform distribution
               and/or value function reduction")
    end

    if get_setting(m, :reduce_state_vars) || get_setting(m, :reduce_v)
        try
            @assert get_setting(m, :n_jump_vars) > 0
            @assert get_setting(m, :n_state_vars) > 0
            @assert get_setting(m, :n_state_vars_unreduce) >= 0
        catch
            error("Need to enter values in settings for n_jump_vars, n_state_vars,
                   and/or n_state_vars_unreduce in Settings. These indicate the
                   number of jump and aggregate state variables.")
        end
    end

    if get_setting(m, :reduce_state_vars)
        try
            @assert get_setting(m, :krylov_dim) > 0
        catch
            error("Need to specify in settings desired dimension of
                   Krylov subspace, krylov_dim.")
        end
    end

    if get_setting(m, :reduce_v)
        try
            @assert get_setting(m, :n_prior) > 0 && typeof(get_setting(m, :n_prior)) == Int
            @assert get_setting(m, :n_post) > 0 && typeof(get_setting(m, :n_post)) == Int
            @assert in(:knots_dict, keys(m.settings))
            @assert in(:spline_grid, keys(m.settings))
        catch
            error("Need to specify in settings n_prior, n_post, knots_dict
                   and/or spline_grid for value function reduction.
                   See reduction.jl for their uses.")
        end
    end

    # Compute steady state
    steadystate!(m)
    return m
end

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

Initializes the model's parameters, as well as empty values for the steady-state
parameters (in preparation for `steadystate!(m)` being called to initialize
those).
"""
function init_parameters!(m::OneAssetHANK)
    # Initialize parameters
    m <= parameter(:coefrra, 1.0, fixed=true,
                   description="Relative risk aversion",
                   tex_label="coefrra")

    m <= parameter(:frisch, 0.5, fixed=true,
                   description="Frisch elasticity",
                   tex_label="frisch")

    m <= parameter(:meanlabeff, 3.0, fixed=true,
                   description="meanlabeff: so that at h=1/3 output will be approximately = 1",
                   tex_label="meanlabeff")

    m <= parameter(:maxhours, 1.0, fixed=true,
                   description="maxhours:...",
                   tex_label="maxhours")

    # Production
    m <= parameter(:ceselast, 10., fixed=true,
                   description="CES elasticity",
                   tex_label="ceselast")

    m <= parameter(:priceadjust, 100., fixed=true,
                   description="priceadjust...",
                   tex_label="priceadjust")

    # Policy parameters
    m <= parameter(:taylor_inflation, 1.25, fixed=true,
                   description="Taylor rule coefficient on inflation",
                   tex_label="taylor_inflation")

    m <= parameter(:taylor_outputgap, 0., fixed=true,
                   description="Taylor rule coefficient on output",
                   tex_label="taylor_outputgap")

    m <= parameter(:labtax, 0.2, fixed=true,
                   description="Marginal tax rate on labor income",
                   tex_label="labtax")

    m <= parameter(:govbondtarget, 6., fixed=true,
                   description="govbondtarget: multiple of quarterly GDP",
                   tex_label="govbondtarget")

    m <= parameter(:lumptransferpc, 0.06, fixed=true,
                   description="lumptransferepc: 6% of quarterly GDP in steady state",
                   tex_label="lumptransferepc")

    m <= parameter(:govbcrule_fixnomB, 0., fixed=true,
                   description="govbcrule_fixnomB...",
                   tex_label="govbvrule_fixnomB")

    m <= parameter(:labdisutil, m[:meanlabeff] / ((0.75 ^(-m[:coefrra])) *
                                                  ((1. / 3.)^(1/m[:frisch]))), fixed=true,
                   description="Coefficient of labor disutility",
                   tex_label="labdisutil")

    # Aggregate shocks
    m <= parameter(:σ_MP, sqrt(0.05), (1e-20, 1e5), (1e-20, 1e5), ModelConstructors.Exponential(),
                   RootInverseGamma(4, .4), fixed=false,
                   description="Volatility of monetary policy shocks", tex_label="σ_MP")

    m <= parameter(:θ_MP, 0.25, fixed=true,
                   description="Rate of mean reversion in monetary policy shocks",
                   tex_label="θ_MP")

    m <= SteadyStateParameterArray(:V_ss, Vector{Float64}(),
                                   description = "Stacked steady-state value function")

    m <= SteadyStateParameter(:inflation_ss, NaN, description = "Steady-state rate of inflation")
    m <= SteadyStateParameterArray(:g_ss, Vector{Float64}(),
                                   description = "Stacked steady-state distribution")
    m <= SteadyStateParameter(:r_ss, NaN, description = "Steady-state real interest rate")
    m <= SteadyStateParameterArray(:u_ss, Vector{Float64}(), description = "Steady-state flow utility")
    m <= SteadyStateParameterArray(:c_ss, Vector{Float64}(),
                                   description = "Steady-state individual consumption rate")
    m <= SteadyStateParameterArray(:h_ss, Vector{Float64}(),
                                   description = "Steady-state individual labor hours")
    m <= SteadyStateParameterArray(:s_ss, Vector{Float64}(),
                                   description = "Steady-state individual savings rate")
    m <= SteadyStateParameter(:rnom_ss, NaN, description = "Steady-state nominal interest rate")
    m <= SteadyStateParameter(:B_ss, NaN, description = "Steady-state bond supply")
    m <= SteadyStateParameter(:N_ss, NaN, description = "Steady-state labor supply")
    m <= SteadyStateParameter(:Y_ss, NaN, description = "Steady-state output")
    m <= SteadyStateParameter(:labor_share_ss, NaN, description = "Frictionless labor share of income")
    m <= SteadyStateParameter(:w_ss, NaN, description = "Steady-state wage rate")
    m <= SteadyStateParameter(:profit_ss, NaN, description = "Steady-state profits")
    m <= SteadyStateParameter(:C_ss, NaN, description = "Steady-state aggregate consumption")
    m <= SteadyStateParameter(:T_ss, NaN, description = "Steady-state total taxes")
    m <= SteadyStateParameter(:G_ss, NaN, description = "Steady-state government spending")
    m <= SteadyStateParameter(:ρ_ss, NaN, description = "Steady-state discount rate")
end

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

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

    # Read in parameters
    coefrra        = m[:coefrra].value
    frisch         = m[:frisch].value
    meanlabeff     = m[:meanlabeff].value
    maxhours       = m[:maxhours].value
    ceselast       = m[:ceselast].value
    labtax         = m[:labtax].value
    govbondtarget  = m[:govbondtarget].value
    labdisutil     = m[:labdisutil].value
    lumptransferpc = m[:lumptransferpc].value

    # Read in grids
    I                = get_setting(m, :I)
    J                = get_setting(m, :J)
    a                = get_setting(m, :a)
    g_z              = get_setting(m, :g_z)
    zz               = get_setting(m, :zz)
    ymarkov_combined = get_setting(m, :ymarkov_combined)

    # Set necessary variables
    aa   = repeat(a, 1, J)
    amax = maximum(a)
    amin = minimum(a)

    # Read in initial rates
    iterate_r = get_setting(m, :iterate_r)
    r         = get_setting(m, :r0)
    rmin      = get_setting(m, :rmin)
    rmax      = get_setting(m, :rmax)
    iterate_ρ = get_setting(m, :iterate_ρ)
    ρ         = get_setting(m, :ρ0)
    ρmin      = get_setting(m, :ρmin)
    ρmax      = get_setting(m, :ρmax)

    # Read in approximation parameters
    Ir          = get_setting(m, :Ir)
    maxit_HJB   = get_setting(m, :maxit_HJB)
    tol_HJB     = get_setting(m, :tol_HJB)
    Δ_HJB       = get_setting(m, :Δ_HJB)
    maxit_kfe   = get_setting(m, :maxit_kfe)
    tol_kfe     = get_setting(m, :tol_kfe)
    Δ_kfe       = get_setting(m, :Δ_kfe)
    niter_hours = get_setting(m, :niter_hours)
    crit_S      = get_setting(m, :crit_S)

    # Initializing equilibrium objects
    labor_share_ss = (ceselast - 1) / ceselast
    w       = w_ss = labor_share_ss

    # compute initial guesses at steady state values given zz, labor_share_ss, etc.
    N_ss, Y_ss, B_ss, profit_ss, profshare, lumptransfer =
        calculate_ss_equil_vars(zz, labor_share_ss, meanlabeff, lumptransferpc, govbondtarget)

    # Initialize matrices for finite differences
    Vaf = Array{ComplexF64}(undef, I, J)
    Vab = Array{ComplexF64}(undef, I, J)

    cf  = Array{ComplexF64}(undef, I, J) # forward consumption difference
    hf  = Array{ComplexF64}(undef, I, J) # forward hours difference
    sf  = Array{ComplexF64}(undef, I, J) # forward saving difference
    cb  = Array{ComplexF64}(undef, I, J) # backward consumption difference
    hb  = Array{ComplexF64}(undef, I, J) # backward hours difference
    sb  = Array{ComplexF64}(undef, I, J) # backward saving difference
    c0  = Array{ComplexF64}(undef, I, J)
    A   = Array{ComplexF64}(undef, I*J, I*J)

    # Aswitch*v = \lambda_z(v(a,z') - v(a,z))
    # Captures expected change in value function due to jumps in z
    Aswitch = kron(ymarkov_combined, SparseMatrixCSC{ComplexF64}(LinearAlgebra.I, I, I))

    # Initialize steady state variables
    V  = Array{ComplexF64}(undef, I, J) # value function
    u  = Array{ComplexF64}(undef, I, J) # flow utility across state space
    s  = Array{ComplexF64}(undef, I, J) # savings across state space
    c  = Array{ComplexF64}(undef, I, J) # flow consumption
    h  = Array{ComplexF64}(undef, I, J) # flow hours of labor
    h0 = Array{ComplexF64}(undef, I, J) # guess of what h will be

    # Creates functions for computing flow utility, income earned, and labor done given
    # CRRA + frisch elasticity style labor disutility
    util, income, labor = construct_household_problem_functions(V, w, coefrra, frisch, labtax, labdisutil)

    # Setting up forward/backward difference grids for a. See helpers.jl
    daf, dab, azdelta = initialize_diff_grids(a, I, J)

    for ir = 1:Ir
        fill!(c,  complex(0.))
        fill!(h,  complex(1/3))
        fill!(h0, complex(1.))

        # Initial guess
        v = similar(h)

        for i in eachindex(h)
            inc  = income(h[i], zz[i], profshare[i], lumptransfer, r, aa[i]) # Get income
            v[i] = util(inc, h[i]) / ρ                                       # Value function guess
        end

        # Iterate HJB
        for ihjb = 1:maxit_HJB
            V = v
            Vaf, Vab, cf, hf, cb, hb = construct_initial_diff_matrices(V, Vaf, Vab,
                                                                       income, labor, h, h0, zz,
                                                                       profshare, lumptransfer,
                                                                       amax, amin,
                                                                       coefrra, r, daf,
                                                                       dab, maxhours)

            # Iterative method to find consistent forward/backward/neutral difference matrices for c and h
            cf, hf, cb, hb, c0, h0 = hours_iteration(income, labor, zz, profshare,
                                                     lumptransfer, aa,
                                                     coefrra, r, cf, hf, cb, hb, c0, h0,
                                                     maxhours, niter_hours)

            for i in eachindex(h0)
                c0[i] = income(h0[i], zz[i], profshare[i], lumptransfer, r, aa[i])
                sf[i] = income(hf[i], zz[i], profshare[i], lumptransfer, r, aa[i]) - cf[i]
                sb[i] = income(hb[i], zz[i], profshare[i], lumptransfer, r, aa[i]) - cb[i]
            end

            for j=1:J
                Vaf[I, j] = cf[I, j] ^ (-coefrra) # Forward difference for value function w.r.t. wealth
                Vab[1, j] = cb[1, j] ^ (-coefrra) # Backward difference for value function w.r.t. wealth
            end

            # See helpers.jl: this is a custom upwind function, not one from solve_hjb.jl
            V, A, u, h, c, s  = upwind(ρ, V, util, Aswitch, cf, cb, c0, hf, hb, h0, sf, sb, Vaf, Vab,
                                    daf, dab; Δ_HJB = Δ_HJB)

            # Check for convergence
            Vchange = V - v
            v = V

            err_HJB = maximum(abs.(Vchange))
            if err_HJB < tol_HJB
                break
            end
        end

        # Create initial guess for g0
        g0 = zeros(ComplexF64, I, J)
        # Assign stationary income distribution weight at a = 0, zero elsewhere
        g0[iszero.(a), :] = g_z
        # g_z is marginal distribution, so re-weight by some multiplier of Lebesgue measure
        g0 = g0 ./ reshape(azdelta, I, J)
        # Solve for distribution
        g = solve_kfe(A, g0, spdiagm(0 => azdelta),
                       maxit_kfe = maxit_kfe, tol_kfe = tol_kfe,
                       Δ_kfe = Δ_kfe)

        # Back out static conditions/steady state values given our value function and distribution
        N_ss, Y_ss, B_ss, profit_ss, profshare, lumptransfer, bond_err =
            calculate_ss_equil_vars(zz, h, g, azdelta, aa, labor_share_ss, meanlabeff,
                                lumptransferpc, govbondtarget)

        # Check bond market for market clearing
        r, rmin, rmax, ρ, ρmin, ρmax, clear_cond =
            check_bond_market_clearing(bond_err, crit_S, r, rmin, rmax, ρ, ρmin, ρmax,
                                       iterate_r, iterate_ρ)

        if clear_cond
            # Set steady state values
            m[:V_ss]           = real(vec(V))
            m[:inflation_ss]   = 0.
            m[:g_ss]           = real(g)
            m[:r_ss]           = r
            m[:u_ss]           = real(vec(u))
            m[:c_ss]           = real(vec(c))
            m[:h_ss]           = real(vec(h))
            m[:s_ss]           = real(vec(s))
            m[:rnom_ss]        = m[:r_ss].value + m[:inflation_ss].value
            m[:B_ss]           = sum(m[:g_ss].value .* vec(aa) .* azdelta)
            m[:N_ss]           = real(sum(vec(zz) .* m[:h_ss].value .* m[:g_ss].value .* azdelta))
            m[:Y_ss]           =  m[:N_ss]
            m[:labor_share_ss] = (ceselast - 1) / ceselast
            m[:w_ss]           =  m[:labor_share_ss].value
            m[:profit_ss]      = (1 - m[:labor_share_ss].value) .* m[:Y_ss].value
            m[:C_ss]           = sum(m[:c_ss].value .* m[:g_ss].value .* azdelta)
            m[:T_ss]           = real(lumptransfer)
            m[:ρ_ss]           = ρ
            m[:G_ss]           = labtax * m[:w_ss].value * m[:N_ss].value -
                                 m[:T_ss].value - m[:r_ss].value * m[:B_ss].value
            break
        end
    end
    return m
end

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

    # Steady state r iteration
    m <= Setting(:r0, 0.005, "Initial guess for real interest rate")
    m <= Setting(:rmax, 0.08, "Maximum guess for real interest rate")
    m <= Setting(:rmin, 0.001, "Minimum guess for real interest rate")

    # Steady state ρ iteration
    m <= Setting(:ρ0, 0.02, "Initial guess for discount rate")
    m <= Setting(:ρmax, 0.05, "Maximum guess for discount rate")
    m <= Setting(:ρmin, 0.005, "Minimum guess for discount rate")

    # State space grid
    m <= Setting(:I, 100, "Number of grid points")
    m <= Setting(:J, 2, "Number of income states")
    m <= Setting(:amin, 0., "Minimum asset grid value")
    m <= Setting(:amax, 40., "Maximum asset grid value")
    m <= Setting(:agridparam, 1, "Bending coefficient: 1 for linear")
    m <= Setting(:a, construct_asset_grid(get_setting(m,:I), get_setting(m, :agridparam),
                        get_setting(m, :amin), get_setting(m,:amax)), "Asset grid")
    m <= Setting(:ygrid_combined, [0.2, 1])
    m <= Setting(:ymarkov_combined, [-0.5 0.5; 0.0376 -0.0376], "Markov transition parameters")
    m <= Setting(:g_z, compute_stationary_income_distribution(get_setting(m, :ymarkov_combined),
                       get_setting(m, :J)), "Stationary income distribution")
    m <= Setting(:zz, construct_labor_income_grid(get_setting(m, :ygrid_combined),
                       get_setting(m, :g_z), m[:meanlabeff].value, get_setting(m, :I)),
                       "Labor income grid repeated across asset dimension")
    m <= Setting(:z, get_setting(m, :zz)[1, :], "Labor income grid")

    # Number of variables
    m <= Setting(:n_jump_vars, get_setting(m, :I) * get_setting(m, :J) + 1, "Number of jump variables")
    m <= Setting(:n_state_vars, get_setting(m, :I) * get_setting(m, :J), "Number of state variables")
    m <= Setting(:n_state_vars_unreduce, 0, "Number of state variables not being reduced")
    # R: Inserted from old code; not certain necessary
    m <= Setting(:n_static_relations, 5, "Number of static relations: bond-market clearing, labor market
                                          clearing, consumption, output, total assets")
    m <= Setting(:n_vars, get_setting(m, :n_jump_vars) + get_setting(m, :n_state_vars)
                 + get_setting(m, :n_static_relations),
                 "Number of variables, total")

    # Steady state approximation
    m <= Setting(:Ir, 100, "Max number of iterations for steady state")
    m <= Setting(:maxit_HJB, 500, "Max number of iterations for HJB")
    m <= Setting(:tol_HJB, 1e-8, "Tolerance for HJB error")
    m <= Setting(:Δ_HJB, 1e6, "Multiplier for implicit upwind scheme of HJB")
    m <= Setting(:maxit_kfe, 1000, "Max number of iterations for kfe")
    m <= Setting(:tol_kfe, 1e-12, "Tolerance for kfe error")
    m <= Setting(:Δ_kfe, 1e6, "Multiplier for implicit upwind scheme of kfe")
    m <= Setting(:niter_hours, 10, "Max number of iterations for finding hours worked")
    m <= Setting(:iterate_r, false, "Iterate on real interest rate or not")
    m <= Setting(:iterate_ρ, true, "Iterate on discount rate or not")
    m <= Setting(:crit_S, 1e-5, "Tolerance for error in bond markets")

    # Reduction
    m <= Setting(:n_knots, 12, "Number of knot points")
    m <= Setting(:c_power, 1, "Amount of bending of knot point locations to make them nonuniform")
    m <= Setting(:n_post, length(get_setting(m, :zz)[1,:]),
                 "Number of dimensions that need to be approximated by spline basis")
    m <= Setting(:n_prior, 1,
                 "Number of dimensions approximated by spline basis that
                  were not used to compute the basis matrix")

    knots = collect(range(get_setting(m, :amin), stop=get_setting(m, :amax), length=get_setting(m, :n_knots)-1))
    knots = (get_setting(m, :amax) - get_setting(m, :amin)) /
            (2^get_setting(m, :c_power)-1) * ((knots .- get_setting(m, :amin))
                                              / (get_setting(m, :amax) -
                                                 get_setting(m, :amin)) .+ 1) .^ get_setting(m, :c_power) .+
                                                 get_setting(m, :amin) .- (get_setting(m, :amax) -
                                                                          get_setting(m, :amin)) /
                                                                          (2 ^ get_setting(m, :c_power) - 1)
    m <= Setting(:knots_dict, Dict(1 => knots),
                 "Location of knot points for each dimension for value function reduction")
    m <= Setting(:krylov_dim, 20, "Krylov reduction dimension")
    m <= Setting(:reduce_state_vars, true, "Reduce state variables or not") #turned off reduction
    m <= Setting(:reduce_v, true, "Reduce value function or not") #turned off reduction
    m <= Setting(:spline_grid, get_setting(m, :a), "Grid of knot points for spline basis")

    # Sampling method
    m <= Setting(:sampling_method, :SMC, "Set sampling method to SMC")

    # SMC Settings
    m <= Setting(:n_particles, 3000)
    m <= Setting(:n_Φ, 200)
    m <= Setting(:λ, 2.0)
    m <= Setting(:n_smc_blocks, 1)
    m <= Setting(:use_parallel_workers, true)
    m <= Setting(:step_size_smc, 0.5)
    m <= Setting(:n_mh_steps_smc, 1)
    m <= Setting(:resampler_smc, :polyalgo)
    m <= Setting(:target_accept, 0.25)
    m <= Setting(:mixture_proportion, .9)
    m <= Setting(:tempering_target, 0.95)
    m <= Setting(:resampling_threshold, .5)
    m <= Setting(:use_fixed_schedule, true)
    m <= Setting(:smc_iteration, 0)

    # Track lag
    m <= Setting(:track_lag, false, "Add first lag when constructing measurement equation")

    # Forecast
    m <= Setting(:use_population_forecast, true,
                 "Whether to use population forecasts as data")
    m <= Setting(:forecast_zlb_value, 0.13,
        "Value of the zero lower bound in forecast periods, if we choose to enforce it")

    # Simulating states
    m <= Setting(:state_simulation_freq, 2,
                 "How many states you want to simulate between states + 1")
end


function get_init_values_for_comparison(sim::Simulation)
    bus_size = PSID.get_bus_count(sim.simulation_inputs)
    system = PSID.get_system(sim)
    V_R = Vector{Float64}(undef, bus_size)
    V_I = Vector{Float64}(undef, bus_size)
    Vm = Vector{Float64}(undef, bus_size)
    θ = Vector{Float64}(undef, bus_size)
    for bus in PSY.get_components(PSY.Bus, system)
        bus_n = PSY.get_number(bus)
        bus_ix = PSID.get_lookup(sim.simulation_inputs)[bus_n]
        V_R[bus_ix] = sim.x0_init[bus_ix]
        V_I[bus_ix] = sim.x0_init[bus_ix + bus_size]
        Vm[bus_ix] = sqrt(V_R[bus_ix]^2 + V_I[bus_ix]^2)
        θ[bus_ix] = angle(V_R[bus_ix] + V_I[bus_ix] * 1im)
    end
    results =
        Dict{String, Vector{Float64}}("V_R" => V_R, "V_I" => V_I, "Vm" => Vm, "θ" => θ)
    for device in PSY.get_components(PSY.DynamicInjection, system)
        states = PSY.get_states(device)
        name = PSY.get_name(device)
        global_index = PSID.get_global_index(sim.simulation_inputs)[name]
        x0_device = Vector{Float64}(undef, length(states))
        for (i, s) in enumerate(states)
            x0_device[i] = sim.x0_init[global_index[s]]
        end
        results[name] = x0_device
    end
    dyn_branches = PSY.get_components(PSY.DynamicBranch, system)
    if !isempty(dyn_branches)
        for br in dyn_branches
            states = PSY.get_states(br)
            name = PSY.get_name(br)
            global_index = PSID.get_global_index(sim.simulation_inputs)[name]
            x0_br = Vector{Float64}(undef, length(states))
            for (i, s) in enumerate(states)
                x0_br[i] = sim.x0_init[global_index[s]]
            end
            printed_name = "Line " * name
            results[printed_name] = x0_br
        end
    end
    return results
end

function clean_extra_timestep!(t::Vector{Float64}, δ::Vector{Float64})
    idx = unique(i -> t[i], 1:length(t))
    return t[idx], δ[idx]
end

function get_csv_delta(str::AbstractString)
    M = readdlm(str, ',')
    return clean_extra_timestep!(M[:, 1], M[:, 2])
end

function get_csv_data(str::AbstractString)
    M = readdlm(str, ',')
    return M
end


module TokyoVariantAnalysis4C2
using Logging
using DataFrames
using Dates
using Plots
using Plots.PlotMeasures

using ..TokyoVariantAnalysis4
using ..TokyoVariantAnalysis4A
using ..TokyoVariantAnalysis4B2

@info "--------"
@info "プロット・バックエンド・設定"
pyplot(
    titlefont=font("Meiryo",10),
    guidefont=font("Meiryo",10),
    tickfont=font("Meiryo",9)
)
const plot_date_start = Date("2021-05-25")
const plot_date_end   = plot_date_start + Day(100)
const date_ticks_dict = Dict(
    Date("2021-03-21") => "緊急事態宣言解除",
    Date("2021-04-01") => "2021\u201004\u201001",
    Date("2021-04-15") => "04\u201015",
    Date("2021-04-25") => "緊急事態宣言発令",
    Date("2021-05-01") => "2021\u201005\u201001",
    Date("2021-05-15") => "05\u201015",
    Date("2021-06-01") => "2021\u201006\u201001",
    Date("2021-06-15") => "06\u201015",
    Date("2021-06-20") => "緊急事態宣言解除",
    Date("2021-07-01") => "2021\u201007\u201001",
    Date("2021-07-04") => "都議会議員選挙",
    Date("2021-07-12") => "緊急事態宣言発令",
    Date("2021-07-15") => "07\u201015",
    Date("2021-07-23") => "オリンピック開会",
    Date("2021-08-01") => "2021\u201008\u201001",
    Date("2021-08-08") => "オリンピック閉会",
    Date("2021-08-15") => "08\u201015",
    Date("2021-08-22") => "宣言解除予定",
    Date("2021-09-01") => "2021\u201009\u201001",
)
const date_ticks_dates = sort(collect(keys(date_ticks_dict)))
const date_ticks_dict_en = Dict(
    Date("2021-03-21") => "End\\ of\\ the\\ SoE",
    Date("2021-04-01") => "2021\u201004\u201001",
    Date("2021-04-15") => "04\u201015",
    Date("2021-04-25") => "Decl.\\ of\\ the\\ SoE",
    Date("2021-05-01") => "2021\u201005\u201001",
    Date("2021-05-15") => "05\u201015",
    Date("2021-06-01") => "2021\u201006\u201001",
    Date("2021-06-15") => "06\u201015",
    Date("2021-06-20") => "End\\ of\\ the\\ SoE",
    Date("2021-07-01") => "2021\u201007\u201001",
    Date("2021-07-04") => "Assembly\\ election",
    Date("2021-07-12") => "Decl.\\ of\\ the\\ SoE",
    Date("2021-07-15") => "07\u201015",
    Date("2021-07-23") => "Olympics\\ starts",
    Date("2021-08-01") => "2021\u201008\u201001",
    Date("2021-08-08") => "Olympics\\ ends",
    Date("2021-08-15") => "08\u201015",
    Date("2021-08-22") => "SoE\\ ends\\ (sched)",
    Date("2021-09-01") => "2021\u201009\u201001",
)
const date_ticks_dates_en = sort(collect(keys(date_ticks_dict_en)))

plots = Dict{String,Plots.Plot}()
@info "--------"
@info "L452R 割合尤度プロット"
for once in 1:1
    pl = plot(
        size=(640,600),
        framestyle=:box,
        bottom_margin=8px, left_margin=16px, top_margin=16px, right_margin=-32px,
        title="東京（健安研調査・民間検査機関等調査）"
    )
    # x 軸設定
    sdate_value = date_to_value(plot_date_start)
    edate_value = date_to_value(plot_date_end)
    xlims!(pl, sdate_value, edate_value)
    nticks = length(date_ticks_dates)
    vdates = Float64[]
    labels = String[]
    for i in 1:nticks
        date = date_ticks_dates[i]
        if plot_date_start ≤ date ≤ plot_date_end
            push!(vdates, date_to_value(date))
            push!(labels, date_ticks_dict[date])
        end
    end
    xticks!(pl, vdates, labels; xtickfontrotation=90.0)
    # y 軸設定
    ylims!(pl, 0.0, 1.0)
    ylabel!(pl, "L452R 変異のある変異株（主に δ 株）の割合")
    yticks!(pl,
        [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
        ["0\\%", "20\\%", "40\\%", "60\\%", "80\\%", "100\\%"]
    )
    # 尤度ヒートマップ描画
    hm_xrange = range(hm_x_min, hm_x_max; length=hm_x_size)
    hm_yrange = range(hm_y_min, hm_y_max; length=hm_y_size)
    heatmap!(pl, hm_xrange, hm_yrange, hm_f;
        c=:Wistia, colorbar_title="相対尤度（最大尤度を 1 とする）", colorbar_titlefontfamily="Meiryo")
    plot!(pl, (t)->1.0/(1.0+1.0/exp(likelihood_max.α+likelihood_max.β*t)), lc=:red3, ls=:dot)
    # データ点
    pname = "L452R_prop"
    lmtl = -100.0
    lmtu = +100.0
    lmt(r) = min(max(r, lmtl), lmtu)
    ds  = variants_df.date_value
    #des = df.date_error
    rs  = variants_df[!,Symbol("$(pname)")]
    ls  = variants_df[!,Symbol("$(pname)_cil")]
    us  = variants_df[!,Symbol("$(pname)_ciu")]
    n_rows = nrow(variants_df)
    pm_ds  = Vector{Float64}(); sizehint!(pm_ds,  n_rows)
    #pm_des = Vector{Float64}(); sizehint!(pm_des, n_rows)
    pm_rs  = Vector{Float64}(); sizehint!(pm_rs,  n_rows)
    pm_els = Vector{Float64}(); sizehint!(pm_els, n_rows)
    pm_eus = Vector{Float64}(); sizehint!(pm_eus, n_rows)
    ps_ds  = Vector{Float64}(); sizehint!(ps_ds,  n_rows)
    #ps_des = Vector{Float64}(); sizehint!(ps_des, n_rows)
    ps_rs  = Vector{Float64}(); sizehint!(ps_rs,  n_rows)
    ps_els = Vector{Float64}(); sizehint!(ps_els, n_rows)
    ps_eus = Vector{Float64}(); sizehint!(ps_eus, n_rows)
    for i in 1:n_rows
        if !ismissing(rs[i])
            r  = lmt(rs[i])
            el = r - lmt(ls[i])
            eu = lmt(us[i]) - r
            if TokyoVariantAnalysis4A.lh_sdate_value ≤ ds[i] ≤ TokyoVariantAnalysis4A.lh_edate_value
                push!(pm_ds,  ds[i])
                #push!(pm_des, des[i])
                push!(pm_rs,  r)
                push!(pm_els, el)
                push!(pm_eus, eu)
            else
                push!(ps_ds,  ds[i])
                #push!(ps_des, des[i])
                push!(ps_rs,  r)
                push!(ps_els, el)
                push!(ps_eus, eu)
            end
        end
    end
    scatter!(pl, pm_ds, pm_rs; yerrors=(pm_els, pm_eus),
        color=RGB(0.0,0.5,1.0), msc=RGB(0.4,0.4,0.4), legend=:none)
    scatter!(pl, ps_ds, ps_rs; yerrors=(ps_els, ps_eus),
        color=RGB(0.0,0.5,1.0), msc=RGB(0.7,0.7,0.7), legend=:none)
    # グリッド
    vline!(pl, vdates; lc=RGBA(0.0,0.0,0.0,0.1), ls=:dot)
    hline!(pl, [0.2, 0.4, 0.6, 0.8]; lc=RGBA(0.0,0.0,0.0,0.1), ls=:dot)
    # アノテーション
    #coef = ratio_fit_coefs[pname][2]
    annotate!(pl, date_to_value(Date("2021-05-28")), 0.925, text("""
        「東京都新型コロナウイルス感染症対策本部報」より
        民間検査機関等における L452R の検査が報告されるようになった 2021-06-01 以降のデータ
        ただし、デルタ株クラスターが含まれているとみられる場合は、クラスターあたり 1 人分のみ計上
        赤点線は指数関数モデルでの最尤線. 背景の色は最大尤度に相対的な尤度を表す""",
        font("Meiryo",7), RGB(0.3,0.3,0.3), :left)
    )
    plname = "L452R"
    plots[plname] = pl
    @info plname
end
@info "--------"
@info "in English"
for once in 1:1
    pl = plot(
        size=(640,600),
        framestyle=:box,
        bottom_margin=8px, left_margin=16px, top_margin=16px, right_margin=-32px,
        title="Tokyo (Inst. of Public Health and Private Labs.)"
    )
    # x 軸設定
    sdate_value = date_to_value(plot_date_start)
    edate_value = date_to_value(plot_date_end)
    xlims!(pl, sdate_value, edate_value)
    nticks = length(date_ticks_dates_en)
    vdates = Float64[]
    labels = String[]
    for i in 1:nticks
        date = date_ticks_dates[i]
        if plot_date_start ≤ date ≤ plot_date_end
            push!(vdates, date_to_value(date))
            push!(labels, date_ticks_dict_en[date])
        end
    end
    xticks!(pl, vdates, labels; xtickfontrotation=90.0)
    # y 軸設定
    ylims!(pl, 0.0, 1.0)
    ylabel!(pl, "Proportion of the L452R mutation (mainly the δ variant)")
    yticks!(pl,
        [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
        ["0\\%", "20\\%", "40\\%", "60\\%", "80\\%", "100\\%"]
    )
    # 尤度ヒートマップ描画
    hm_xrange = range(hm_x_min, hm_x_max; length=hm_x_size)
    hm_yrange = range(hm_y_min, hm_y_max; length=hm_y_size)
    heatmap!(pl, hm_xrange, hm_yrange, hm_f;
        c=:Wistia, colorbar_title="relative likelihood (the maximum likelihood is 1)", colorbar_titlefontfamily="Meiryo")
    plot!(pl, (t)->1.0/(1.0+1.0/exp(likelihood_max.α+likelihood_max.β*t)), lc=:red3, ls=:dot)
    # データ点
    pname = "L452R_prop"
    lmtl = -100.0
    lmtu = +100.0
    lmt(r) = min(max(r, lmtl), lmtu)
    ds  = variants_df.date_value
    #des = df.date_error
    rs  = variants_df[!,Symbol("$(pname)")]
    ls  = variants_df[!,Symbol("$(pname)_cil")]
    us  = variants_df[!,Symbol("$(pname)_ciu")]
    n_rows = nrow(variants_df)
    pm_ds  = Vector{Float64}(); sizehint!(pm_ds,  n_rows)
    #pm_des = Vector{Float64}(); sizehint!(pm_des, n_rows)
    pm_rs  = Vector{Float64}(); sizehint!(pm_rs,  n_rows)
    pm_els = Vector{Float64}(); sizehint!(pm_els, n_rows)
    pm_eus = Vector{Float64}(); sizehint!(pm_eus, n_rows)
    ps_ds  = Vector{Float64}(); sizehint!(ps_ds,  n_rows)
    #ps_des = Vector{Float64}(); sizehint!(ps_des, n_rows)
    ps_rs  = Vector{Float64}(); sizehint!(ps_rs,  n_rows)
    ps_els = Vector{Float64}(); sizehint!(ps_els, n_rows)
    ps_eus = Vector{Float64}(); sizehint!(ps_eus, n_rows)
    for i in 1:n_rows
        if !ismissing(rs[i])
            r  = lmt(rs[i])
            el = r - lmt(ls[i])
            eu = lmt(us[i]) - r
            if TokyoVariantAnalysis4A.lh_sdate_value ≤ ds[i] ≤ TokyoVariantAnalysis4A.lh_edate_value
                push!(pm_ds,  ds[i])
                #push!(pm_des, des[i])
                push!(pm_rs,  r)
                push!(pm_els, el)
                push!(pm_eus, eu)
            else
                push!(ps_ds,  ds[i])
                #push!(ps_des, des[i])
                push!(ps_rs,  r)
                push!(ps_els, el)
                push!(ps_eus, eu)
            end
        end
    end
    scatter!(pl, pm_ds, pm_rs; yerrors=(pm_els, pm_eus),
        color=RGB(0.0,0.5,1.0), msc=RGB(0.4,0.4,0.4), legend=:none)
    scatter!(pl, ps_ds, ps_rs; yerrors=(ps_els, ps_eus),
        color=RGB(0.0,0.5,1.0), msc=RGB(0.7,0.7,0.7), legend=:none)
    # グリッド
    vline!(pl, vdates; lc=RGBA(0.0,0.0,0.0,0.1), ls=:dot)
    hline!(pl, [0.2, 0.4, 0.6, 0.8]; lc=RGBA(0.0,0.0,0.0,0.1), ls=:dot)
    # アノテーション
    #coef = ratio_fit_coefs[pname][2]
    annotate!(pl, date_to_value(Date("2021-05-28")), 0.925, text("""
        Data source: Daily reports of “Tōkyō-to Shingata Corona Virus Kansenshō Taisaku Hombu”
        Using data since 2021-06-01 when private laboratories began reporting L452R tests.
        Only one case included per cluster if case numbers contain Delta clusters.
        The red dotted line is the maximum likelihood line for the exponential model.
        The background color indicates the likelihood relative to the maximum.""",
        font("Meiryo",7), RGB(0.3,0.3,0.3), :left)
    )
    plname = "L452R_en"
    plots[plname] = pl
    @info plname
end

end # module


# let
#     W = 2
#     for T ∈ [Float32, Float64, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64]
#         Wmax = max_vector_width(T)
#         while W ≤ Wmax

#             W += W
#         end
#     end
# end
# function broadcast_str(W::Int, typ::String)
#     vtyp = "<$W x $typ>"
#     """
#         %ie = insertelement $vtyp undef, $typ %0, i32 0
#         %v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer
#         ret $vtyp %v
#     """
# end
@inline vzero(::Val{1}, ::Type{T}) where {T<:NativeTypes} = zero(T)
@inline vzero(::StaticInt{1}, ::Type{T}) where {T<:NativeTypes} = zero(T)
@generated function vzero(::Union{Val{W},StaticInt{W}}, ::Type{T}) where {W,T<:NativeTypes}
    # isone(W) && return Expr(:block, Expr(:meta,:inline), Expr(:call, :zero, T))
    typ = LLVM_TYPES[T]
    instrs = "ret <$W x $typ> zeroinitializer"
    quote
        $(Expr(:meta,:inline))
        Vec(llvmcall($instrs, _Vec{$W,$T}, Tuple{}))
    end
end
@generated function vbroadcast(::Union{Val{W},StaticInt{W}}, s::_T) where {W,_T<:NativeTypes}
    isone(W) && return :s
    if _T <: Integer && sizeof(_T) * W > REGISTER_SIZE
        T = pick_integer(W, sizeof(_T))
        if _T <: Unsigned
            T = unsigned(T)
        end
        # ssym = :(s % $T)
        ssym = :(convert($T, s))
    else
        T = _T
        ssym = :s
    end
    typ = LLVM_TYPES[T]
    vtyp = vtype(W, typ)
    instrs = """
        %ie = insertelement $vtyp undef, $typ %0, i32 0
        %v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer
        ret $vtyp %v
    """
    quote
        $(Expr(:meta,:pure,:inline))
        Vec(llvmcall($instrs, _Vec{$W,$T}, Tuple{$T}, $ssym))
    end
end

# for T ∈ [Float32,Float64,Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64]#, Float16]
#     maxW = pick_vector_width(T)
#     typ = LLVM_TYPES[T]
#     W = 2
#     while W ≤ maxW
#         instrs = "ret <$W x $typ> zeroinitializer"
#         @eval @inline vzero(::Val{$W}, ::Type{$T}) = Vec(llvmcall($instrs, Vec{$W,$T}, Tuple{}, ))
#         instrs = broadcast_str(W, typ)
#         # vtyp = "<$W x $typ>"
#         # instrs = """
#         # %ie = insertelement $vtyp undef, $typ %0, i32 0
#         # %v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer
#         # ret $vtyp %v
#         # """
#         @eval Base.@pure @inline vbroadcast(::Val{$W}, s::$T) = Vec(llvmcall($instrs, Vec{$W,$T}, Tuple{$T}, s))
#         W += W
#     end
# end

# @generated function vbroadcast(::Val{W}, s::Ptr{T}) where {W, T}
#     typ = JULIAPOINTERTYPE
#     instrs = broadcast_str(W, typ)
#     quote
#         $(Expr(:meta,:inline))
#         Vec(llvmcall( $instrs, _Vec{$W,Ptr{$T}}, Tuple{Ptr{$T}}, s ))
#     end
# end
# @generated function vbroadcast(::Val{W}, s::T) where {W, T <: NativeTypes}
#     typ = LLVM_TYPES[T]
#     instrs = broadcast_str(W, typ)
#     quote
#         $(Expr(:meta,:inline))
#         Vec(llvmcall( $instrs, _Vec{$W,$T}, Tuple{$T}, s))
#     end
# end
@generated function vbroadcast(::Union{Val{W},StaticInt{W}}, ptr::Ptr{T}) where {W, T}
    isone(W) && return Expr(:block, Expr(:meta, :inline), :(vload(ptr)))
    typ = LLVM_TYPES[T]
    ptyp = JULIAPOINTERTYPE
    vtyp = "<$W x $typ>"
    alignment = Base.datatype_alignment(T)
    instrs = """
        %ptr = inttoptr $ptyp %0 to $typ*
        %res = load $typ, $typ* %ptr, align $alignment
        %ie = insertelement $vtyp undef, $typ %res, i32 0
        %v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer
        ret $vtyp %v
    """
    quote
        $(Expr(:meta,:inline))
        Vec(llvmcall( $instrs, _Vec{$W,$T}, Tuple{Ptr{$T}}, ptr ))
    end
end
# @generated function Base.zero(::Type{Vec{W,T}}) where {W,T}
#     typ = LLVM_TYPES[T]
#     instrs = "ret <$W x $typ> zeroinitializer"
#     quote
#         $(Expr(:meta,:inline))
#         Vec(llvmcall($instrs, _Vec{$W,$T}, Tuple{}, ))
#     end
# end
@inline vbroadcast(::Union{Val{W},StaticInt{W}}, v::AbstractSIMDVector{W}) where {W} = v
@inline Vec{W,T}(v::Vec{W,T}) where {W,T} = v
# @inline vbroadcast(::Val{1}, s::T) where {T <: NativeTypes} = s
# @inline vbroadcast(::Val{1}, s::Ptr{T}) where {T <: NativeTypes} = s
@inline vzero(::Union{Val{W},StaticInt{W}}, ::Type{T}) where {W,T} = zero(Vec{W,T})
@inline Base.zero(::Type{Vec{W,T}}) where {W,T} = vzero(Val{W}(), T)
@inline Base.zero(::Vec{W,T}) where {W,T} = zero(Vec{W,T})
@inline Base.one(::Vec{W,T}) where {W,T} = vbroadcast(Val{W}(), one(T))

@inline Base.one(::Type{Vec{W,T}}) where {W,T} = vbroadcast(Val{W}(), one(T))
@inline Base.oneunit(::Type{Vec{W,T}}) where {W,T} = vbroadcast(Val{W}(), one(T))
@inline vzero(::Type{T}) where {T<:Number} = zero(T)
@inline vzero() = vzero(pick_vector_width_val(Float64), Float64)
# @inline sveczero(::Type{T}) where {T} = Svec(vzero(pick_vector_width_val(T)))
# @inline sveczero() = Svec(vzero(pick_vector_width_val(Float64)))

@inline Vec{W,T}(s::Integer) where {W,T<:Integer} = vbroadcast(Val{W}(), s % T)
@inline Vec{W,T}(s::Real) where {W,T} = vbroadcast(Val{W}(), T(s))
@inline Vec{W}(s::T) where {W,T<:NativeTypes} = vbroadcast(Val{W}(), s)
@inline Vec(s::T) where {T<:NativeTypes} = vbroadcast(pick_vector_width_val(T), s)

@generated function Base.zero(::Type{VecUnroll{N,W,T,V}}) where {N,W,T,V}
    t = Expr(:tuple); foreach(_ -> push!(t.args, :(zero(Vec{$W,$T}))), 0:N)
    Expr(:block, Expr(:meta, :inline), :(VecUnroll($t)))
end
@inline Base.zero(::VecUnroll{N,W,T,V}) where {N,W,T,V} = zero(VecUnroll{N,W,T,V})

@generated function VecUnroll{N,W,T,V}(x::S) where {N,W,T,V<:AbstractSIMDVector{W,T},S<:Real}
    t = Expr(:tuple)
    for n ∈ 0:N
        push!(t.args, :(convert($V, x)))
    end
    Expr(:block, Expr(:meta,:inline), :(VecUnroll($t)))
end
@inline VecUnroll{N,W,T}(x::NativeTypesV) where {N,W,T} = VecUnroll{N,W,T,Vec{W,T}}(x)

# @inline vbroadcast(::Union{Val{W},StaticInt{W}}, ::Type{T}, s::T) where {W,T} = vbroadcast(Val{W}(), s)
# @generated function vbroadcast(::Union{Val{W},StaticInt{W}}, ::Type{T}, s::S) where {W,T,S}
#     ex = if sizeof(T) < sizeof(S)
#         vbroadcast(Val{W}(), promote_type(T, S)(s))
#     else

#     end
#     Expr(:block, Expr(:meta, :inline), ex)
# end


"""
     word2vec(train, output; size=100, window=5, sample=1e-3, hs=0,  negative=5, threads=12, iter=5, min_count=5, alpha=0.025, debug=2, binary=1, cbow=1, save_vocal=Nothing(), read_vocab=Nothing(), verbose=false,)

    Parameters for training:
        train <file>
            Use text data from <file> to train the model
        output <file>
            Use <file> to save the resulting word vectors / word clusters
        size <Int>
            Set size of word vectors; default is 100
        window <Int>
            Set max skip length between words; default is 5
        sample <AbstractFloat>
            Set threshold for occurrence of words. Those that appear with
            higher frequency in the training data will be randomly
            down-sampled; default is 1e-5.
        hs <Int>
            Use Hierarchical Softmax; default is 1 (0 = not used)
        negative <Int>
            Number of negative examples; default is 0, common values are
            5 - 10 (0 = not used)
        threads <Int>
            Use <Int> threads (default 12)
        iter <Int>
            Run more training iterations (default 5)
        min_count <Int>
            This will discard words that appear less than <Int> times; default
            is 5
        alpha <AbstractFloat>
            Set the starting learning rate; default is 0.025
        debug <Int>
            Set the debug mode (default = 2 = more info during training)
        binary <Int>
            Save the resulting vectors in binary moded; default is 0 (off)
        cbow <Int>
            Use the continuous back of words model; default is 1 (skip-gram
            model)
        save_vocab <file>
            The vocabulary will be saved to <file>
        read_vocab <file>
            The vocabulary will be read from <file>, not constructed from the
            training data
        verbose <Bool>
            Print output from training
"""
function word2vec(train::AbstractString, output::AbstractString;
                  size::Int=100, window::Int=5, sample::AbstractFloat=1e-3,
                  hs::Int=0, negative::Int=5, threads::Int=12, iter::Int=5,
                  min_count::Int=5, alpha::AbstractFloat=0.025,
                  debug::Int=2, binary::Int=0, cbow::Int=1,
                  save_vocab=Nothing(), read_vocab=Nothing(),
                  verbose::Bool=false)

    parameters = AbstractString[]
    args = ["-train", "-output", "-size", "-window", "-sample", "-hs",
            "-negative", "-threads", "-iter", "-min-count", "-alpha",
            "-debug", "-binary", "-cbow"]
    values = [train, output, size, window, sample, hs, negative, threads,
              iter, min_count, alpha, debug, binary, cbow]

    for (arg, value) in zip(args, values)
        push!(parameters, arg)
        push!(parameters, string(value))
    end
    if save_vocab != Nothing()
        push!(parameters, "-save-vocab")
        push!(parameters, string(save_vocab))
    end
    if read_vocab != Nothing()
        push!(parameters, "-read-vocab")
        push!(parameters, string(read_vocab))
    end
    Word2Vec_jll.word2vec() do command
        run(`$(command) $(parameters)`)
    end
end

"""
```
     word2clusters(train, output, classes; size=100, window=5, sample=1e-3, hs=0,  negative=5, threads=1, iter=5, min_count=5, alpha=0.025, debug=2, binary=1, cbow=1, save_vocal=Nothing(), read_vocab=Nothing(), verbose=false)
```
    Parameters for training:
        train <file>
            Use text data from <file> to train the model
        output <file>
            Use <file> to save the resulting word vectors / word clusters
        size <Int>
            Set size of word vectors; default is 100
        window <Int>
            Set max skip length between words; default is 5
        sample <AbstractFloat>
            Set threshold for occurrence of words. Those that appear with
            higher frequency in the training data will be randomly
            down-sampled; default is 0 (off), useful value is 1e-5
        hs <Int>
            Use Hierarchical Softmax; default is 1 (0 = not used)
        negative <Int>
            Number of negative examples; default is 0, common values are 5 - 10
            (0 = not used)
        threads <Int>
            Use <Int> threads (default 1)
        iter <Int>
            Run more training iterations (default 5)
        min_count <Int>
            This will discard words that appear less than <Int> times
            (default 5)
        alpha <AbstractFloat>
            Set the starting learning rate; default is 0.025
        classes <Int>
            Number of word classes; if 0, output word classes rather than
            word vectors (default 0)
        debug <Int>
            Set the debug mode (default = 2 = more info during training)
        binary <Int>
            Save the resulting vectors in binary moded; default is 0 (off)
        cbow <Int>
            Use the continuous back of words model; default is 1 (0 for skip-gram
            model)
        save_vocab <file>
            The vocabulary will be saved to <file>
        read_vocab <file>
            The vocabulary will be read from <file>, not constructed from the
            training data
        verbose <Bool>
            Print output from training
"""
function word2clusters(train::AbstractString, output::AbstractString,
                       classes::Int; size::Int=100, window::Int=5,
                       sample::AbstractFloat=0., hs::Int=0,
                       negative::Int=5, threads::Int=1, iter::Int=5,
                       min_count::Int=5, alpha::AbstractFloat=0.025,
                       debug::Int=2, binary::Int=0, cbow::Int=1,
                       save_vocab=Nothing(), read_vocab=Nothing(),
                       verbose::Bool=false)

    parameters = AbstractString[]
    args = ["-train", "-output", "-size", "-window", "-sample", "-hs",
            "-negative", "-threads", "-iter", "-min-count", "-alpha",
            "-debug", "-binary", "-cbow", "-classes"]
    values = [train, output, size, window, sample, hs, negative, threads,
              iter, min_count, alpha, debug, binary, cbow, classes]
    for (arg, value) in zip(args, values)
        push!(parameters, arg)
        push!(parameters, string(value))
    end
    if save_vocab != Nothing()
        push!(parameters, "-save-vocab")
        push!(parameters, string(save_vocab))
    end
    if read_vocab != Nothing()
        push!(parameters, "-read-vocab")
        push!(parameters, string(read_vocab))
    end
    Word2Vec_jll.word2vec() do command
        run(`$(command) $(parameters)`)
    end
end

"""
     word2phrase(train, output; min_count=5, threshold=100, debug=2)

    Parameters for training:
	train <file>
	      Use text data from <file> to train the model
	output <file>
              Use <file> to save the resulting word vectors /
              word clusters / phrases
	min_count <Int>
              This will discard words that appear less than <int> times;
              default is 5
	threshold <AbstractFloat>
      	      The <AbstractFloat> value represents threshold for
              forming the phrases (higher means less phrases); default 100
	debug <Int>
	      Set the debug mode (default = 2 = more info during training)
"""
function word2phrase(train::AbstractString, output::AbstractString;
                     min_count::Int=5, threshold::Int=100, debug::Int=2)
    parameters = AbstractString[]
    args = ["-train", "-output", "-min-count", "-threshold", "-debug"]
    values = [train, output, min_count, threshold, debug]
    for (arg, value) in zip(args, values)
        push!(parameters, arg)
        push!(parameters, string(value))
    end
    Word2Vec_jll.word2phrase() do command
        run(`$(command) $(parameters)`)
    end
end


"""
    Simulation

Type used to run various simulations.

# Fields
- `t0::Float64`: initial time
- `tf::Float64`: final time
- `dt::Float64`: time-step
- `ts`: t0:dt:tf
"""
struct Simulation
    t0::Float64
    tf::Float64
    dt::Float64
    ts
end

"""
    Simulation(tf::Float64)
    Simulation(tf::Float64, dt::Float64)
    Simulation(t0::Float64, tf::Float64, dt::Float64)

Construct a simulation object from an initial time, final time, and time-step.
"""
Simulation(tf::Float64) = Simulation(0.0, tf, 0.01, 0.0:0.01:tf)
Simulation(tf::Float64, dt::Float64) = Simulation(0.0, tf, dt, 0.0:dt:tf)
Simulation(t0::Float64, tf::Float64, dt::Float64) = Simulation(t0, tf, dt, t0:dt:tf)

"""
    Base.length(sim::Simulation)

Get number of timesteps in simulation.
"""
Base.length(sim::Simulation) = length(sim.ts)

"""
    (sim::Simulation)(Σ::ControlAffineSystem)
    (sim::Simulation)(Σ::ControlAffineSystem, k::FeedbackPolicy)
    (sim::Simulation)(Σ::ControlAffineSystem, k::CLFQP)
    (sim::Simulation)((Σ::ControlAffineSystem, k::CBFQP)

Simulate a trajectory of a control affine system.
"""
function (sim::Simulation)(Σ::ControlAffineSystem)
    initialize!(Σ)
    for i in 1:(length(sim) - 1)
        u = Σ.m == 1 ? 0.0 : zeros(Σ.m)
        step!(Σ, u, sim.dt)
    end
    vec2mat!(Σ)

    return Σ
end

function (sim::Simulation)(Σ::ControlAffineSystem, k::FeedbackPolicy)
    initialize!(Σ)
    for i in 1:(length(sim) - 1)
        u = k(Σ.x)
        step!(Σ, u, sim.dt)
    end
    vec2mat!(Σ)

    return Σ
end

function (sim::Simulation)(Σ::ControlAffineSystem, k::CLFQP)
    initialize!(Σ)
    for i in 1:(length(sim) - 1)
        u = k(Σ.x)
        step!(Σ, u, sim.dt)
    end
    vec2mat!(Σ)

    return Σ
end

function (sim::Simulation)(Σ::ControlAffineSystem, k::CBFQP)
    initialize!(Σ)
    for i in 1:(length(sim) - 1)
        u = k(Σ.x)
        step!(Σ, u, sim.dt)
    end
    vec2mat!(Σ)

    return Σ
end

function (sim::Simulation)(Σ::ControlAffineSystem, k::TimeVaryingCLFQP)
    initialize!(Σ)
    for i in 1:(length(sim) - 1)
        t = sim.ts[i]
        u = k(Σ.x, t)
        step!(Σ, u, sim.dt)
    end
    vec2mat!(Σ)

    return Σ
end

function (sim::Simulation)(Σ::ControlAffineSystem, k::TimeVaryingCBFQP)
    initialize!(Σ)
    for i in 1:(length(sim) - 1)
        t = sim.ts[i]
        u = k(Σ.x, t)
        step!(Σ, u, sim.dt)
    end
    vec2mat!(Σ)

    return Σ
end


using ImageCore, Colors, FixedPointNumbers, ColorVectorSpace
using Base.Test

@testset "Deprecated" begin
    @testset "constructors" begin
        # grayim
        for a in (0x00:0x18, collect(0x00:0x18))
            B = reshape(a, 5, 5)
            img = grayim(B)
            @test colorspace(img) == "Gray"
            @test colordim(img) == 0
            @test img[5,2] === Gray{U8}(9/255)
        end
        for a in (0x00:0x7c, collect(0x00:0x7c))
            B = reshape(a, 5, 5, 5)
            img = grayim(B)
            @test colorspace(img) == "Gray"
            @test colordim(img) == 0
            @test img[5,2,3] === Gray{U8}(59/255)
        end
        for a in (0x0000:0x0018, collect(0x0000:0x0018))
            B = reshape(a, 5, 5)
            img = grayim(B)
            @test colorspace(img) == "Gray"
            @test colordim(img) == 0
            @test img[5,2] === Gray{U16}(9/65535)
        end
        for a in (0x0000:0x007c, collect(0x0000:0x007c))
            B = reshape(a, 5, 5, 5)
            img = grayim(B)
            @test colorspace(img) == "Gray"
            @test colordim(img) == 0
            @test img[5,2,3] === Gray{U16}(59/65535)
        end
        # colorim
        for a in (0x00:0x4a, collect(0x00:0x4a))
            C = reshape(a, 3, 5, 5)
            img = colorim(C)
            @test colorspace(img) == "RGB"
            @test colordim(img) == 0
            @test img[1,1] === RGB{U8}(0,1/255,2/255)
            C = reshape(a, 5, 5, 3)
            img = colorim(C)
            @test colorspace(img) == "RGB"
            @test colordim(img) == 0
            @test img[1,1] === RGB{U8}(0,25/255,50/255)
            @test_throws ErrorException colorim(reshape(a, 5, 3, 5))
        end
        for a in (0x0000:0x004a, collect(0x0000:0x004a))
            C = reshape(a, 3, 5, 5)
            img = colorim(C)
            @test colorspace(img) == "RGB"
            @test colordim(img) == 0
            @test img[1,1] === RGB{U16}(0,1/65535,2/65535)
            C = reshape(a, 5, 5, 3)
            img = colorim(C)
            @test colorspace(img) == "RGB"
            @test colordim(img) == 0
            @test img[1,1] === RGB{U16}(0,25/65535,50/65535)
            @test_throws ErrorException colorim(reshape(a, 5, 3, 5))
        end
        for a in (0x00:0x63, collect(0x00:0x63))
            for (S,T) in (("RGBA", RGBA), ("ARGB",ARGB))
                C = reshape(a, 4, 5, 5)
                img = colorim(C, S)
                @test colorspace(img) == S
                @test colordim(img) == 0
                @test img[1,1] === T{U8}(0,1/255,2/255,3/255)
                C = reshape(a, 5, 5, 4)
                img = colorim(C, S)
                @test colorspace(img) == S
                @test colordim(img) == 0
                @test img[1,1] === T{U8}(0,25/255,50/255,75/255)
            end
            @test_throws ErrorException colorim(reshape(a, 4, 5, 5), "RRRR")
        end
        # ambiguous sizes
        for a in (0x00:0x2c, collect(0x00:0x2c))
            C = reshape(a, 3, 5, 3)
            @test_throws ErrorException colorim(C)
        end
        for a in (0x00:0x4f, collect(0x00:0x4f))
            C = reshape(a, 4, 5, 4)
            @test_throws ErrorException colorim(C)
        end
    end
    @testset "traits" begin
        for (B,S) in ((rand(UInt16(1):UInt16(20), 3, 5),"Gray"),
                      (rand(Gray{Float32}, 3, 5),"Gray"),
                      (rand(RGB{Float16}, 3, 5),"RGB"),
                      (bitrand(3, 5),"Binary"),
                      (rand(UInt32, 3, 5),"RGB24"))
            T = eltype(B)
            @test data(B) === B
            @test colorspace(B) == S
            @test colordim(B) == 0
            @test timedim(B) == 0
            @test spatialorder(B) == [:y,:x]
            @test isdirect(B)
            @test limits(B) == (T==Bool ? (0,1) : (zero(T),one(T)))
            @test storageorder(B) == [:y,:x]
            @test ncolorelem(B) == (T <: Colorant ? length(T) : 1)
            assert2d(B)
            assert_scalar_color(B)
            @test isyfirst(B)
            @test !isxfirst(B)
            assert_yfirst(B)
            @test_throws ErrorException assert_xfirst(B)
            @test spatialproperties(B) == String[]
            @test_throws ErrorException permutedims(B, ("x", "y"))
            @test_throws ErrorException permutedims(B, (:x, :y))
        end
        for (B,S) in ((rand(UInt16(1):UInt16(20), 5),"Gray"),
                      (rand(Gray{Float32}, 5),"Gray"),
                      (rand(RGB{Float16}, 5),"RGB"),
                      (bitrand(5),"Binary"))
            @test colorspace(B) == S
        end
        for (B,S) in ((rand(Gray{Float32}, 5, 5, 5),"Gray"),
                      (rand(RGB{Float16}, 5, 5, 5),"RGB"),
                      (bitrand(5, 5, 5),"Binary"))
            @test colorspace(B) == S
        end
        @test colorspace(rand(Float32, 5, 5, 3)) == "RGB"
        @test_throws ErrorException colorspace(rand(Float32, 5, 5, 5))
    end
    @testset "Reinterpret, convert, separate, raw" begin
        # some of these are redundant with convert_reinterpret.jl, but
        # these are "classic" tests so it's good to make sure they
        # pass
        a = RGB{Float64}[RGB(1,1,0)]
        af = reinterpret(Float64, a)
        @test vec(af) == [1.0,1.0,0.0]
        @test size(af) == (3,1)
        @test_throws DimensionMismatch reinterpret(Float32, a)
        anew = reinterpret(RGB, af)
        @test anew == a
        anew = reinterpret(RGB, vec(af))
        @test anew[1] == a[1]
        @test ndims(anew) == 0
        anew = reinterpret(RGB{Float64}, af)
        @test anew == a
        @test_throws DimensionMismatch reinterpret(RGB{Float32}, af)
        Au8 = rand(0x00:0xff, 3, 5, 4)
        A8 = reinterpret(UFixed8, Au8)
        rawrgb8 = reinterpret(RGB, A8)
        @test eltype(rawrgb8) == RGB{UFixed8}
        @test reinterpret(UFixed8, rawrgb8) == A8
        @test reinterpret(UInt8, rawrgb8) == Au8
        rawrgb32 = convert(Array{RGB{Float32}}, rawrgb8)
        @test eltype(rawrgb32) == RGB{Float32}
        @test ufixed8(rawrgb32) == rawrgb8
        @test reinterpret(UFixed8, rawrgb8) == A8
        # imrgb8 = convert(Image, rawrgb8)
        # @test spatialorder(imrgb8) == ImageCore.yx
        # cvt = convert(Image, imrgb8)
        # @test cvt == imrgb8 && typeof(cvt) == typeof(imrgb8)
        # cvt = convert(Image{RGB{UFixed8}}, imrgb8)
        # @test cvt == imrgb8 && typeof(cvt) == typeof(imrgb8)
        imrgb8 = rawrgb8
        im8 = reinterpret(UFixed8, imrgb8)
        @test data(im8) == A8
        @test reinterpret(UInt8, imrgb8) == Au8
        @test reinterpret(RGB, im8) == imrgb8
        ims8 = separate(imrgb8)
        @test colorspace(ims8) == "RGB"
        # cvt = convert(Image, ims8)
        # @test cvt == ims8 && typeof(cvt) == typeof(ims8)
        # @test (cvt = convert(Image{UFixed8}, ims8)) == ims8 && typeof(cvt) == typeof(ims8)
        # @test (cvt = separate(ims8)) == ims8 && typeof(cvt) == typeof(ims8)
        # imrgb8_2 = convert(Image{RGB}, ims8)
        # @test isa(imrgb8_2, Image{RGB{UFixed8}})
        A = reinterpret(UFixed8, UInt8[1 2; 3 4])
        # imgray = convert(Image{Gray{UFixed8}}, A)
        imgray = convert(Array{Gray{U8}}, A)
        @test spatialorder(imgray) == ImageCore.yx
        @test data(imgray) == reinterpret(Gray{UFixed8}, [0x01 0x02; 0x03 0x04])
        # @test eltype(convert(Image{HSV{Float32}}, imrgb8)) == HSV{Float32}
        # @test eltype(convert(Image{HSV}, float32(imrgb8))) == HSV{Float32}
        # local img = Image(reinterpret(Gray{UFixed16}, rand(UInt16, 5, 5)))
        # imgs = subim(img, :, :)
        img = reinterpret(Gray{UFixed16}, rand(UInt16, 5, 5))
        imgs = view(img, :, :)
        @test isa(minimum(imgs), Gray{UFixed16})  # change in behavior from Images issue#232
        imgdata = rand(UInt16, 5, 5)
        # img = Image(reinterpret(Gray{UFixed16}, imgdata))
        img = reinterpret(Gray{UFixed16}, imgdata)
        @test all(raw(img) .== imgdata)
        # @test size(raw(Image(rawrgb8))) == (3,5,4)
        @test all(raw(imgdata) .== imgdata)
    end
end

nothing


module TopicModelsVB

using DelimitedFiles
using SpecialFunctions
using LinearAlgebra
using Random
using Distributions
using OpenCL
using Crayons

export Document, Corpus
export TopicModel
export check_doc, check_corp, check_model
export LDA, fLDA, CTM, fCTM, CTPF, gpuLDA, gpuCTM, gpuCTPF
export readcorp, writecorp, abridge_corp!, alphabetize_corp!, compact_corp!, condense_corp!, pad_corp!, remove_empty_docs!, remove_redundant!, stop_corp!, trim_corp!, trim_docs!, fixcorp!, showdocs, showtitles, getvocab, getusers
export train!
export @gpu
export gendoc, gencorp, showtopics, showlibs, showdrecs, showurecs, predict, topicdist

include("macros.jl")
include("utils.jl")
include("Corpus.jl")
include("TopicModel.jl")
include("LDA.jl")
include("fLDA.jl")
include("CTM.jl")
include("fCTM.jl")
include("CTPF.jl")
include("gpuLDA.jl")
include("gpuCTM.jl")
include("gpuCTPF.jl")
include("modelutils.jl")

end


# One action, zero observation, one timestep long, +1 reward every timestep: This isolates
# the value network. If my agent can't learn that the value of the only observation it ever
# sees it 1, there's a problem with the value loss calculation or the optimizer.
import AlphaZero.GI

struct GameSpec <: GI.AbstractGameSpec end

mutable struct GameEnv <: GI.AbstractGameEnv
  terminated :: Bool
end

const INITSTATE = (; terminated=false)

GI.init(::GameSpec, state=INITSTATE) = GameEnv(state.terminated)

GI.spec(::GameEnv) = GameSpec()

GI.two_players(::GameSpec) = false

const ACTIONS = collect(1:1)

GI.actions(::GameSpec) = ACTIONS

GI.actions_mask(g::GameEnv) = [true]

GI.current_state(g::GameEnv) = (; terminated=g.terminated)

GI.white_playing(g::GameEnv) = true

GI.game_terminated(g::GameEnv) = g.terminated

GI.white_reward(g::GameEnv) = 1.

function GI.play!(g::GameEnv, action)
  g.terminated = true
end

GI.vectorize_state(::GameSpec, state) = Float32[!state.terminated]


export resample

"""
	resample(A::AbstractGeoArray, resolution::Number;
			 crs::GeoFormat=crs(A), method::String="near")
	resample(A::AbstractGeoArray, snap::AbstractGeoArray; method::String="near")

`resample` uses `ArchGDAL.gdalwarp` to resample an `AbstractGeoArray`.

## Arguments
- `A`: The `AbstractGeoArray` to resample.
- `resolution`: A `Number` specifying the resolution for the output.
  If the keyword argument `crs` (described below) is specified, `resolution` must be in units of the `crs`.
- `snap`: an `AbstractGeoArray` whos resolution, crs and bounds will be snapped to.
  For best results it should roughly cover the same extent, or a subset of `A`.

## Keyword Arguments
- `crs`: A `GeoFormatTypes.GeoFormat` specifying an output crs
  (`A` with be reprojected to `crs` in addition to being resampled). Defaults to `crs(A)`
- `method`: A `String` specifying the method to use for resampling. Defaults to `"near"`
  (nearest neighbor resampling). See [resampling method](https://gdal.org/programs/gdalwarp.html#cmdoption-gdalwarp-r)
  in the gdalwarp docs for a complete list of possible values.

"""
function resample end
function resample(A::AbstractGeoArray, resolution::Number;
				  crs::GeoFormat=crs(A),
                  method::String="near")
    wkt = convert(String, convert(WellKnownText, crs))
    flags = ["-t_srs", "$(wkt)",
             "-tr", "$(resolution)", "$(resolution)",
             "-r", "$(method)"]
    AG.Dataset(A) do dataset
        AG.gdalwarp([dataset], flags) do warped
            GeoArray(warped)
        end
    end
end
function resample(A::AbstractGeoArray, snap::AbstractGeoArray; method::String="near")
    wkt = convert(String, convert(WellKnownText, crs(snap)))
    latres, lonres = map(abs ∘ step, span(snap, (Lat(), Lon())))
    (latmin, latmax), (lonmin, lonmax) = bounds(snap, (Lat(), Lon()))
    flags = ["-t_srs", "$(wkt)",
             "-tr", "$(latres)", "$(lonres)",
             "-te", "$(lonmin)", "$(latmin)", "$(lonmax)", "$(latmax)",
             "-r", "$(method)"]
    AG.Dataset(A) do dataset
        AG.gdalwarp([dataset], flags) do warped
            GeoArray(warped)
        end
    end
end