# Test indexes
@testset "Indexes" begin
    # Abstract types
    @test AbstractInfOptIndex isa DataType
    @test ObjectIndex isa DataType
    @test ObjectIndex <: AbstractInfOptIndex
    # IndependentParameterIndex
    @test IndependentParameterIndex <: ObjectIndex
    @test IndependentParameterIndex(1).value == 1
    # DependentParametersIndex
    @test DependentParametersIndex <: ObjectIndex
    @test DependentParametersIndex(1).value == 1
    # DependentParameterIndex
    @test DependentParameterIndex <: AbstractInfOptIndex
    idx = DependentParametersIndex(1)
    @test DependentParameterIndex(idx, 1).object_index == idx
    @test DependentParameterIndex(idx, 1).param_index == 1
    # InfiniteParameterIndex
    @test InfiniteParameterIndex <: AbstractInfOptIndex
    # FiniteParameterIndex
    @test FiniteParameterIndex <: ObjectIndex
    @test FiniteParameterIndex(1).value == 1
    # ParameterFunctionIndex
    @test ParameterFunctionIndex <: ObjectIndex
    @test ParameterFunctionIndex(1).value == 1
    # InfiniteVariableIndex
    @test InfiniteVariableIndex <: ObjectIndex
    @test InfiniteVariableIndex(1).value == 1
    # SemiInfiniteVariableIndex
    @test SemiInfiniteVariableIndex <: ObjectIndex
    @test SemiInfiniteVariableIndex(1).value == 1
    # PointVariableIndex
    @test PointVariableIndex <: ObjectIndex
    @test PointVariableIndex(1).value == 1
    # FiniteVariableIndex
    @test FiniteVariableIndex <: ObjectIndex
    @test FiniteVariableIndex(1).value == 1
    # DerivativeIndex
    @test DerivativeIndex <: ObjectIndex
    @test DerivativeIndex(1).value == 1
    # MeasureIndex
    @test MeasureIndex <: ObjectIndex
    @test MeasureIndex(1).value == 1
    # InfOptConstraintIndex
    @test InfOptConstraintIndex <: ObjectIndex
    @test InfOptConstraintIndex(1).value == 1
    # test CleverDict extensions
    @testset "CleverDict Extensions" begin
        # index_to_key
        @test MOIUC.index_to_key(IndependentParameterIndex, Int64(42)) == IndependentParameterIndex(42)
        @test MOIUC.index_to_key(DependentParametersIndex, Int64(42)) == DependentParametersIndex(42)
        @test MOIUC.index_to_key(FiniteParameterIndex, Int64(42)) == FiniteParameterIndex(42)
        @test MOIUC.index_to_key(ParameterFunctionIndex, Int64(42)) == ParameterFunctionIndex(42)
        @test MOIUC.index_to_key(InfiniteVariableIndex, Int64(42)) == InfiniteVariableIndex(42)
        @test MOIUC.index_to_key(SemiInfiniteVariableIndex, Int64(42)) == SemiInfiniteVariableIndex(42)
        @test MOIUC.index_to_key(PointVariableIndex, Int64(42)) == PointVariableIndex(42)
        @test MOIUC.index_to_key(FiniteVariableIndex, Int64(42)) == FiniteVariableIndex(42)
        @test MOIUC.index_to_key(DerivativeIndex, Int64(42)) == DerivativeIndex(42)
        @test MOIUC.index_to_key(MeasureIndex, Int64(42)) == MeasureIndex(42)
        @test MOIUC.index_to_key(InfOptConstraintIndex, Int64(42)) == InfOptConstraintIndex(42)
        # index_to_key
        @test MOIUC.key_to_index(MeasureIndex(42)) == 42
    end
    # test Base.length
    @testset "Base.length" begin
        @test length(MeasureIndex(1)) == 1
    end
    # test Base.broadcastable
    @testset "Base.broadcastable" begin
        @test Base.broadcastable(MeasureIndex(1)) isa Ref
    end
    # test Base.iterate
    @testset "Base.iterate" begin
        @test iterate(MeasureIndex(1)) == (MeasureIndex(1), true)
        @test iterate(MeasureIndex(1), true) isa Nothing
    end
    # test Base.hash
    @testset "Base.hash" begin
        @test Base.hash(MeasureIndex(1), UInt(1)) isa UInt
    end
end

# Test Infinite Domains
@testset "Infinite Domains" begin
    # Abstract types
    @test AbstractInfiniteDomain isa DataType
    @test InfiniteScalarDomain <: AbstractInfiniteDomain
    @test InfiniteArrayDomain <: AbstractInfiniteDomain
    # IntervalDomain
    @test IntervalDomain <: InfiniteScalarDomain
    @test IntervalDomain(0, 1) == IntervalDomain(0.0, 1.0)
    @test_throws ErrorException IntervalDomain(1, 0)
    # UniDistributionDomain
    @test UniDistributionDomain <: InfiniteScalarDomain
    @test UniDistributionDomain(Normal()) isa UniDistributionDomain{<:Normal}
    # NonUnivariateDistribution
    @test InfiniteOpt.NonUnivariateDistribution <: Distribution
    # MultiDistributionDomain
    dist = Dirichlet([0.5, 0.5])
    @test MultiDistributionDomain <: InfiniteArrayDomain
    @test MultiDistributionDomain(dist) isa MultiDistributionDomain{<:Dirichlet}
    @test MultiDistributionDomain(dist) == MultiDistributionDomain(dist)
    @test MultiDistributionDomain(dist) != MultiDistributionDomain(MvNormal([1, 1], [1 0; 0 1]))
    # CollectionDomain
    @test CollectionDomain <: InfiniteArrayDomain
    @test CollectionDomain([IntervalDomain(0, 1), IntervalDomain(2, 3)]) isa CollectionDomain{IntervalDomain}
end

# Test Generative Support Info
@testset "Generative Support Info" begin
    # Abstract types
    @test AbstractGenerativeInfo isa DataType
    # NoGenerativeSupports
    @test NoGenerativeSupports <: AbstractGenerativeInfo
    @test NoGenerativeSupports() isa NoGenerativeSupports
    @test NoGenerativeSupports() == NoGenerativeSupports()
    # UniformGenerativeInfo
    @test UniformGenerativeInfo <: AbstractGenerativeInfo
    @test UniformGenerativeInfo([0.0, 1.0], InternalLabel) isa UniformGenerativeInfo
    @test UniformGenerativeInfo([-2, 0, 1], InternalLabel, -2, 1).support_basis == [0, 2.0/3.0, 1]
    @test UniformGenerativeInfo([0.0, 1.0], InternalLabel) == UniformGenerativeInfo([2, 6], InternalLabel, 2, 6)
    @test UniformGenerativeInfo([0.0, 1.0], InternalLabel) != UniformGenerativeInfo([2, 5], InternalLabel, 2, 6)
    @test_throws ErrorException UniformGenerativeInfo([2, 6], InternalLabel)
end

# Test Derivative Evaluation Methods
@testset "Derivative Methods" begin
    # abstract types
    @test AbstractDerivativeMethod isa DataType
    @test GenerativeDerivativeMethod <: AbstractDerivativeMethod
    @test NonGenerativeDerivativeMethod <: AbstractDerivativeMethod
    # test orthogonal collocation (defined in derivative_evaluations.jl)
    @test OrthogonalCollocation <: GenerativeDerivativeMethod
    @test OrthogonalCollocation{GaussLobatto}(3, GaussLobatto()) isa OrthogonalCollocation
    @test OrthogonalCollocation(3, GaussLobatto()) == OrthogonalCollocation{GaussLobatto}(3, GaussLobatto())
    @test OrthogonalCollocation(3) == OrthogonalCollocation(3, GaussLobatto())
    @test_throws ErrorException OrthogonalCollocation(1)
    # test finite difference
    @test FDTechnique isa DataType
    @test Forward <: FDTechnique
    @test Central <: FDTechnique
    @test Backward <: FDTechnique
    @test FiniteDifference <: NonGenerativeDerivativeMethod
    @test FiniteDifference(Forward()) isa FiniteDifference
    @test FiniteDifference() == FiniteDifference(Backward())
    @test FiniteDifference(Forward(), false) isa FiniteDifference
end

# Test parameter datatypes
@testset "Parameters" begin
    # abstract types
    @test InfOptParameter <: AbstractVariable
    @test ScalarParameter <: InfOptParameter
    @test AbstractDataObject isa DataType
    # test IndependentParameter
    @test IndependentParameter <: ScalarParameter
    dict = SortedDict{Float64, Set{DataType}}(2 => Set([All]))
    method = FiniteDifference()
    info = NoGenerativeSupports()
    @test IndependentParameter(IntervalDomain(0, 1), dict, 6, method, info).domain isa IntervalDomain
    # test FiniteParameter
    @test FiniteParameter <: ScalarParameter
    @test FiniteParameter(1) == FiniteParameter(1)
    # test DependentParameters
    @test DependentParameters <: InfOptParameter
    @test DependentParameters(CollectionDomain([IntervalDomain(0, 1)]),
                              Dict(zeros(1) => Set([All])), 6, [method]).domain isa CollectionDomain
    # test ScalarParameterData
    @test ScalarParameterData <: AbstractDataObject
    @test ScalarParameterData(FiniteParameter(42), 1, 1, "bob").name == "bob"
    # test MultiParameterData
    @test MultiParameterData <: AbstractDataObject
    params = DependentParameters(CollectionDomain([IntervalDomain(0, 1)]),
                                 Dict(zeros(1) => Set([All])), 6, [method])
    @test MultiParameterData(params, 1, 1:1, ["par[1]"]) isa MultiParameterData
end

# Test the InfiniteModel datatype
@testset "InfiniteModel" begin
    # test basic
    @test InfiniteModel <: JuMP.AbstractModel
    @test InfiniteModel().ext isa Dict
    # prepare optimizer constructor
    mockoptimizer = () -> MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()),
                                             eval_objective_value=false)
    mockattributes = MOI.OptimizerWithAttributes(mockoptimizer, MOI.Silent() => true)
    # test optimizer constructors
    @test InfiniteModel(mockoptimizer).optimizer_constructor == mockoptimizer
    @test InfiniteModel(mockattributes).optimizer_constructor == mockoptimizer
    m = InfiniteModel();
    @test isa(Base.broadcastable(m), Base.RefValue{InfiniteModel})
    @test length(JuMP.object_dictionary(m)) == 0
    @test InfiniteModel(caching_mode = MOIU.MANUAL) isa JuMP.AbstractModel
    @test InfiniteModel(mockoptimizer,
                        caching_mode = MOIU.MANUAL) isa JuMP.AbstractModel
    # test accessors
    @test InfiniteOpt._last_param_num(m) == 0
    @test InfiniteOpt._param_object_indices(m) isa Vector{Union{IndependentParameterIndex, DependentParametersIndex}}
    # test other methods
    @test empty!(InfiniteModel(mockoptimizer)).optimizer_constructor == mockoptimizer
end

# Test reference variable datatypes
@testset "References" begin
    m = InfiniteModel();
    # Abstract types
    @test DispatchVariableRef <: AbstractVariableRef
    @test FiniteRef <: DispatchVariableRef
    # test GeneralVariableRef
    @test GeneralVariableRef <: AbstractVariableRef
    @test GeneralVariableRef(m, 1, MeasureIndex) isa GeneralVariableRef
    # test IndependentParameterRef
    @test IndependentParameterRef <: DispatchVariableRef
    idx = IndependentParameterIndex(1)
    @test IndependentParameterRef(m, idx) isa IndependentParameterRef
    # test DependentParameterRef
    @test DependentParameterRef <: DispatchVariableRef
    idx = DependentParameterIndex(DependentParametersIndex(1), 1)
    @test DependentParameterRef(m, idx) isa DependentParameterRef
    # test FiniteParameterRef
    @test FiniteParameterRef <: FiniteRef
    idx = FiniteParameterIndex(1)
    @test FiniteParameterRef(m, idx) isa FiniteParameterRef
    # InfiniteVariableRef
    @test InfiniteVariableRef <: DispatchVariableRef
    idx = InfiniteVariableIndex(1)
    @test InfiniteVariableRef(m, idx) isa InfiniteVariableRef
    # SemiInfiniteVariableRef
    @test SemiInfiniteVariableRef <: DispatchVariableRef
    idx = SemiInfiniteVariableIndex(1)
    @test SemiInfiniteVariableRef(m, idx) isa SemiInfiniteVariableRef
    # ParameterFunctionRef
    @test ParameterFunctionRef <: DispatchVariableRef
    idx = ParameterFunctionIndex(1)
    @test ParameterFunctionRef(m, idx) isa ParameterFunctionRef
    # PointVariableRef
    @test PointVariableRef <: FiniteRef
    idx = PointVariableIndex(1)
    @test PointVariableRef(m, idx) isa PointVariableRef
    # FiniteVariableRef
    @test FiniteVariableRef <: FiniteRef
    idx = FiniteVariableIndex(1)
    @test FiniteVariableRef(m, idx) isa FiniteVariableRef
    # DerivativeRef
    @test DerivativeRef <: DispatchVariableRef
    idx = DerivativeIndex(1)
    @test DerivativeRef(m, idx) isa DerivativeRef
    # MeasureRef
    @test MeasureRef <: DispatchVariableRef
    idx = MeasureIndex(1)
    @test MeasureRef(m, idx) isa MeasureRef
    # InfOptConstraintRef
    @test InfOptConstraintRef isa DataType
    idx = InfOptConstraintIndex(1)
    @test InfOptConstraintRef(m, idx) isa InfOptConstraintRef
end

# Test DomainRestrictions
@testset "Domain Restrictions" begin
    # prepare the data
    m = InfiniteModel()
    idx = DependentParametersIndex(1)
    idx1 = DependentParameterIndex(idx, 1)
    idx2 = DependentParameterIndex(idx, 2)
    idx3 = IndependentParameterIndex(1)
    par1 = GeneralVariableRef(m, 1, DependentParameterIndex, 1)
    par2 = GeneralVariableRef(m, 1, DependentParameterIndex, 2)
    par3 = GeneralVariableRef(m, 1, IndependentParameterIndex, -1)
    pars = [par1, par2]
    # test datatype
    @testset "DataType" begin
        @test DomainRestrictions isa UnionAll
        d = Dict(par3 => IntervalDomain(0, 1))
        @test DomainRestrictions(d) isa DomainRestrictions{GeneralVariableRef}
        @test DomainRestrictions() isa DomainRestrictions{GeneralVariableRef}
    end
    # test _expand_parameter_tuple
    @testset "_expand_parameter_tuple" begin
        # test good to go
        d = (par3 => IntervalDomain(0, 1),)
        @test InfiniteOpt._expand_parameter_tuple(d) == Dict(d...)
        # test with array of parameters
        d = (pars => IntervalDomain(0, 1), par3 => IntervalDomain(0, 1))
        @test InfiniteOpt._expand_parameter_tuple(d) isa Dict
        # test symbolic domains
        d = (pars => [0, 1], par3 => 0)
        @test InfiniteOpt._expand_parameter_tuple(d) isa Dict
        # test bad vector domain
        d = (pars => [0, 1, 3],)
        @test_throws ErrorException InfiniteOpt._expand_parameter_tuple(d)
        # test fallback
        d = (pars => :a,)
        @test_throws ErrorException InfiniteOpt._expand_parameter_tuple(d)
    end
    # test expansion definition
    @testset "DomainRestrictions Expansion" begin
        d = (par3 => IntervalDomain(0, 1),)
        @test DomainRestrictions(d).intervals == Dict(d...)
        d = (pars => IntervalDomain(0, 1), par3 => IntervalDomain(0, 1))
        @test DomainRestrictions(d).intervals isa Dict{GeneralVariableRef, IntervalDomain}
        @test DomainRestrictions(par3 => 0).intervals isa Dict
    end
    dr = DomainRestrictions(par3 => [0, 1])
    # test intervals
    @testset "intervals" begin
        @test intervals(dr) == dr.intervals
    end
    # test Base.length
    @testset "Base.length" begin
        @test length(dr) == 1
    end
    # test Base.isempty
    @testset "Base.isempty" begin
        @test !isempty(dr)
        @test isempty(DomainRestrictions())
    end
    # test Base.:(==)
    @testset "Base.:(==)" begin
        @test dr == DomainRestrictions((par3 => IntervalDomain(0, 1),))
        @test dr != DomainRestrictions()
    end
    # test Base.copy
    @testset "Base.copy" begin
        @test copy(dr) == dr
    end
    # test Base.getindex
    @testset "Base.getindex" begin
        @test dr[par3] == IntervalDomain(0, 1)
    end
    # test Base.setindex!
    @testset "Base.setindex!" begin
        @test (dr[par3] = IntervalDomain(0, 2)) == IntervalDomain(0, 2)
    end
    # test Base.haskey
    @testset "Base.haskey" begin
        @test haskey(dr, par3)
        @test !haskey(dr, par2)
    end
    # test Base.keys
    @testset "Base.keys" begin
        @test keys(dr) == keys(intervals(dr))
    end
    # test Base.iterate
    @testset "Base.iterate" begin
        @test [r for r in dr] == [par3 => IntervalDomain(0, 2)]
    end
    # test Base.merge
    @testset "Base.merge" begin
        new_dr = DomainRestrictions()
        @test merge(new_dr, dr) == dr
    end
    # test Base.merge!
    @testset "Base.merge!" begin
        new_dr = DomainRestrictions()
        @test merge!(new_dr, dr) isa DomainRestrictions
        @test new_dr == dr
    end
    # test Base.filter
    @testset "Base.filter" begin
        @test filter(e -> e[1] != par3, dr) == DomainRestrictions()
    end
    # test Base.delete!
    @testset "Base.delete!" begin
        @test delete!(dr, par3) isa DomainRestrictions
        @test length(dr) == 0
    end
end

# Test parameter functions
@testset "Parameter Functions" begin
    # useful data
    m = InfiniteModel()
    pref = GeneralVariableRef(m, 1, IndependentParameterIndex, -1)
    vt = IC.VectorTuple(pref)
    # test ParameterFunction
    @test ParameterFunction(sin, vt, [1], [1]) isa ParameterFunction
    # test ParameterFunctionData
    @test ParameterFunctionData(ParameterFunction(sin, vt, [1], [1])) isa ParameterFunctionData
end

# Test variable datatypes
@testset "Variables" begin
    # useful data
    m = InfiniteModel()
    num = Float64(0)
    sample_info = VariableInfo(true, num, true, num, true, num, true, num, true, true)
    func = (x) -> NaN
    inf_info = VariableInfo(true, num, true, num, true, num, false, func, true, true)
    pref = GeneralVariableRef(m, 1, IndependentParameterIndex, -1)
    vref = GeneralVariableRef(m, 1, InfiniteVariableIndex)
    # Infinite variable
    @test InfiniteVariable <: JuMP.AbstractVariable
    @test InfiniteVariable(inf_info, IC.VectorTuple(pref), [1], [1], true) isa InfiniteVariable
    # Semi-Infinite variable
    @test SemiInfiniteVariable <: JuMP.AbstractVariable
    @test SemiInfiniteVariable(vref, Dict(1 => Float64(2)), [1], [1]) isa SemiInfiniteVariable
    # Point variable
    @test PointVariable <: JuMP.AbstractVariable
    @test PointVariable(sample_info, vref, Float64[1]) isa PointVariable
    # VariableData
    @test VariableData <: AbstractDataObject
    @test VariableData(ScalarVariable(sample_info)) isa VariableData
end

# Test derivative datatypes
@testset "Derivatives" begin
    # useful data
    m = InfiniteModel()
    num = Float64(0)
    sample_info = VariableInfo(true, num, true, num, true, num, true, num, true, true)
    func = (x) -> NaN
    inf_info = VariableInfo(true, num, true, num, true, num, false, func, true, true)
    pref = GeneralVariableRef(m, 1, IndependentParameterIndex, -1)
    vref = GeneralVariableRef(m, 1, InfiniteVariableIndex)
    dref = GeneralVariableRef(m, 1, DerivativeIndex)
    # derivative
    @test Derivative <: JuMP.AbstractVariable
    @test Derivative(inf_info, true, vref, pref) isa Derivative
    # Semi-infinite derivative
    @test SemiInfiniteVariable(dref, Dict(1 => Float64(2)), [1], [1]) isa SemiInfiniteVariable
    # Point derivative
    @test PointVariable(sample_info, dref, Float64[1]) isa PointVariable
    # VariableData
    @test VariableData(Derivative(inf_info, true, vref, pref)) isa VariableData
end

# Test the measure datatypes
@testset "Measures" begin
    # Helper data
    m = InfiniteModel()
    pref = GeneralVariableRef(m, 2, IndependentParameterIndex)
    w(t) = 1
    info = NoGenerativeSupports()
    # abstract types
    @test AbstractMeasureData isa DataType
    # DiscreteMeasureData
    @test DiscreteMeasureData <: AbstractMeasureData
    @test DiscreteMeasureData(pref, ones(2), ones(2), All, w, NaN, NaN, false) isa DiscreteMeasureData
    @test DiscreteMeasureData([pref], ones(2), ones(1, 2), All, w, [NaN], [NaN], false) isa DiscreteMeasureData
    # FunctionalDistributionDomain
    @test FunctionalDiscreteMeasureData <: AbstractMeasureData
    @test FunctionalDiscreteMeasureData(pref, w, 2, All, info, w, NaN, NaN, false) isa FunctionalDiscreteMeasureData
    @test FunctionalDiscreteMeasureData([pref], w, 2, All, w, [NaN], [NaN], false) isa FunctionalDiscreteMeasureData
    @test FunctionalDiscreteMeasureData([pref], w, 2, All, info, w, [NaN], [NaN], false) isa FunctionalDiscreteMeasureData
    # Measure
    @test Measure isa UnionAll
    @test Measure(zero(AffExpr),
                  DiscreteMeasureData(pref, ones(2), ones(2), All, w, NaN, NaN, false),
                  [1], [1], false) isa Measure
    # MeasureData
    @test MeasureData <: AbstractDataObject
    @test MeasureData(Measure(zero(AffExpr), DiscreteMeasureData(pref,
                      ones(2), ones(2), All, w, NaN, NaN, false), [1], [1], true)) isa MeasureData
end

# Test the constraint datatypes
@testset "Constraints" begin
    # Setup
    m = InfiniteModel()
    pref = GeneralVariableRef(m, 2, IndependentParameterIndex, -1)
    rs = DomainRestrictions(pref => [0, 1])
    con = ScalarConstraint(zero(AffExpr), MOI.Integer())
    # DomainRestrictedConstraint
    @test DomainRestrictedConstraint <: JuMP.AbstractConstraint
    @test DomainRestrictedConstraint(con, rs).restrictions == rs
    # ConstraintData
    @test ConstraintData <: AbstractDataObject
    @test ConstraintData(con, [1], "", MeasureIndex[], false) isa ConstraintData
end


cd(@__DIR__)
lines = readlines("../monitor.txt")
lines = split.(lines, ":")

filterlines(name) = [eval(Meta.parse(l[2])) for l in filter(x->x[1]==name, lines)]

rws, wws = filterlines("readweight"), filterlines("writeweight")

rws = hcat(rws...)
wws = hcat(wws...)


macro run_exps(db)
    #addprocs(1)
    #n = procs()[end]
    #db = fetch(@spawnat n $(quote
    _results = gensym()
    all_implementations = gensym()
    functions = gensym()
    parameters = gensym()
    results = gensym()
    current_dir = gensym()
    x = gensym()
    esc(quote
        using JuliaDB
        using Query

        # Extract the implementations
        $all_implementations = Int[]
        $x = @from i in db.runs_table begin
            @from j in db.experiments_table
            @where j.Execute && i.Experiment_ID == j.ID && !i.Executed && !i.Executing
            @select i.Implementation_ID
            @collect
        end
        $all_implementations = unique(x)

        # Extract the functions, parameters and results
        $functions = Dict{Int, Symbol}()
        $parameters = Dict{Int, Vector{Tuple{Symbol, Int}}}()
        $results = Dict{Int, Vector{Int}}()
        for implementation in $all_implementations
            $parameters[implementation] = Tuple{Symbol,Int}[]
            $results[implementation] = Int[]

            i = db.implementations_table[implementation]
            $functions[implementation] = i.Function
            if i.Repository != "" && i.Commit != "" && i.Commit != ComputExp.HEAD(i.Repository)
                ComputExp.STASH(i.Repository)
                ComputExp.CHECKOUT(i.Repository, i.Commit)
            end
            if i.Repository != ""
                eval(parse("using $(i.Repository)"))
            end
            include(i.Script)

            params = @from j in db.implementation_parameters_table begin
                @where j.Implementation_ID == implementation
                @select {j.Parameter_symbol, j.ID}
                @collect
            end
            for j in params
                push!($parameters[implementation], (j.Parameter_symbol, j.ID))
            end

            # Retrieves the results of the implementation, order is assumed the same order returned by the function
            $results[implementation] = @from j in db.implementation_results_table begin
                @where j.Implementation_ID == implementation
                @select j.ID
                @collect
            end
        end

        $current_dir = pwd()

        # Run the experiments
        for exper in db.experiments_table
            if exper.Execute
                # Get the run ids of the experiment
                exper_runs = @from i in db.runs_table begin
                    @where i.Experiment_ID == exper.ID
                    @select i.ID
                    @collect
                end

                # Change the directory
                cd(exper.Directory)

                # Run the run
                for rn in exper_runs
                    r = db.runs_table[rn]
                    if !r.Executed && !r.Executing
                        column(db.runs_table, :Executing)[rn] = true
                        problem_features = Dict(@from i in db.problem_features_table begin
                            @where i.Problem_ID == r.Problem_ID
                            @select {i.Feature_symbol, i.Feature_value}
                            @collect
                        end)

                        run_parameters = Dict(@from i in db.run_implementation_parameter_table begin
                            @from j in db.implementation_parameters_table
                            @where i.Run_ID == r.ID && i.Implementation_parameter_ID == j.ID
                            @select {j.Parameter_symbol, i.Value}
                            @collect
                        end)

                        func = db.implementations_table[r.Implementation_ID].Function
                        $_results = eval(func)(; problem = problem_features, parameters = run_parameters)

                        column(db.runs_table, :Executed)[rn] = true
                        column(db.runs_table, :Executing)[rn] = false

                        for (rid, rv) in zip($results[r.Implementation_ID], $_results)
                            resultid = ComputExp.newid(db.run_implementation_result_table)
                            ComputExp.add_row!(db.run_implementation_result_table, @NT(ID = resultid, Implementation_result_ID = rid, Run_ID = r.ID, Value = rv))
                        end
                    end
                end
                #column(db.experiments_table, :Execute)[exper.ID] = false
            end
        end
        cd($current_dir)
    #end))
    #rmprocs(n)
    end)
end


function set_bdx(Qbase, rho, u, v, p, ch, nch, icell, j)
    Qbase[icell, j, 1] = rho
    Qbase[icell, j, 2] = u
    Qbase[icell, j, 3] = v
    Qbase[icell, j, 4] = p
    for ns in 1:nch
        Qbase[icell, j, 4+ns] = ch[ns]
    end
    return Qbase
end

function set_bdy(Qbase, rho, u, v, p, ch, nch, jcell, i)
    Qbase[i, jcell, 1] = rho
    Qbase[i, jcell, 2] = u
    Qbase[i, jcell, 3] = v
    Qbase[i, jcell, 4] = p
    for ns in 1:nch
        Qbase[i, jcell, 4+ns] = ch[ns]
    end
    return Qbase
end

function set_boundary(Qbase, cellxmax, cellymax, vecAx, vecAy, bdcon, Rhat, nval, nch)
    """bdcon[i][j]
        i : 境界番号(x-,x+ ,y-,y+)
        j=1-6 : "bd1_con":"2",
                "bd1_rho":"1.0",
                "bd1_u"  :"300.0",
                "bd1_v"  :"0.0",
                "bd1_p"  :"1.0",
                "bd1_N2" :"1.0",
                "bd1_N"  :"0.0",
                "bd1_T"  :"300.0",
    """
    rho = 0.0
    u   = 0.0
    v   = 0.0
    p   = 0.0
    ch  = zeros(nch)

    # bd1 = x-
    # bd2 = x+
    bdx         = [1, cellxmax]
    bdx_nei     = [2, cellxmax-1]
    bdx_nei_rev = [cellxmax-1, 2]
    for dim in 1:2
        if Int(bdcon[dim][1]) == 00
            for j in 1:cellymax
                rho = bdcon[dim][2]
                u   = bdcon[dim][3]
                v   = bdcon[dim][4]
                p   = bdcon[dim][5]
                for ns in 1:nch
                    ch[ns] =  bdcon[dim][5+ns]
                end
                set_bdx(Qbase, rho, u, v, p, ch, nch, bdx[dim], j)
            end
        elseif Int(bdcon[dim][1]) == 01
            for j in 1:cellymax
                rho = bdcon[dim][2]
                u   = bdcon[dim][3]
                v   = bdcon[dim][4]
                T   = bdcon[dim][6+nch]
                p   = rho*Rhat[bdx_nei[dim],j]*T
                for ns in 1:nch
                    ch[ns] =  bdcon[dim][5+ns]
                end
                set_bdx(Qbase, rho, u, v, p, ch, nch, bdx[dim], j)
            end
        elseif Int(bdcon[dim][1]) == 20
            for j in 1:cellymax
                rho = Qbase[bdx_nei[dim],j,1]
                u   = Qbase[bdx_nei[dim],j,2]
                v   = Qbase[bdx_nei[dim],j,3]
                p   = Qbase[bdx_nei[dim],j,4]
                for ns in 1:nch
                    ch[ns] =  Qbase[bdx_nei[dim],j,4+ns]
                end
                set_bdx(Qbase, rho, u, v, p, ch, nch, bdx[dim], j)
            end
        elseif Int(bdcon[dim][1]) == 30
            for j in 1:cellymax
                rho = Qbase[bdx_nei[dim],j,1]
                u   = -Qbase[bdx_nei[dim],j,2]
                v   = -Qbase[bdx_nei[dim],j,3]
                p   = Qbase[bdx_nei[dim],j,4]
                for ns in 1:nch
                    ch[ns] =  Qbase[bdx_nei[dim],j,4+ns]
                end
                set_bdx(Qbase, rho, u, v, p, ch, nch, bdx[dim], j)
            end
        elseif Int(bdcon[dim][1]) == 31
            for j in 1:cellymax
                u   = -Qbase[bdx_nei[dim],j,2]
                v   = -Qbase[bdx_nei[dim],j,3]
                p   = Qbase[bdx_nei[dim],j,4]
                T   = bdcon[dim][6+nch]
                rho = p/(Rhat[bdx_nei[dim],j]*T)
                for ns in 1:nch
                    ch[ns] =  Qbase[bdx_nei[dim],j,4+ns]
                end
                set_bdx(Qbase, rho, u, v, p, ch, nch, bdx[dim], j)
            end
        elseif Int(bdcon[dim][1]) == 40
            for j in 1:cellymax
                rho = Qbase[bdx_nei_rev[dim],j,1]
                u   = Qbase[bdx_nei_rev[dim],j,2]
                v   = Qbase[bdx_nei_rev[dim],j,3]
                p   = Qbase[bdx_nei_rev[dim],j,4]
                for ns in 1:nch
                    ch[ns] =  Qbase[bdx_nei_rev[dim],j,4+ns]
                end
                set_bdx(Qbase, rho, u, v, p, ch, nch, bdx[dim], j)
            end
        end
    end

    # bd3 = y-
    # bd4 = y+
    bdy     = [1, cellymax]
    bdy_nei     = [2, cellymax-1]
    bdy_nei_rev = [cellymax-1, 2]
    for dim in 3:4
        if Int(bdcon[dim][1]) == 00
            for i in 1:cellxmax
                rho = bdcon[dim][2]
                u   = bdcon[dim][3]
                v   = bdcon[dim][4]
                p   = bdcon[dim][5]
                for ns in 1:nch
                    ch[ns] =  bdcon[dim][5+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 01
            for i in 1:cellxmax
                rho = bdcon[dim][2]
                u   = bdcon[dim][3]
                v   = bdcon[dim][4]
                T   = bdcon[dim][6+nch]
                p   = rho*Rhat[i,bdy_nei[dim-2]]*T
                for ns in 1:nch
                    ch[ns] =  bdcon[dim][5+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 02
            temp1 = Int64(round(cellxmax/4)+1) # 1/4流出
            temp2 = Int64((temp1-1)*3+1)       # 1/4~3/4流入
            for i in 1:temp1
                rho = Qbase[i,bdy_nei[dim-2],1]
                u   = Qbase[i,bdy_nei[dim-2],2]
                v   = Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
            for i in temp1+1:temp2
                rho = bdcon[dim][2]
                u   = bdcon[dim][3]
                v   = bdcon[dim][4]
                p   = bdcon[dim][5]
                for ns in 1:nch
                    ch[ns] =  bdcon[dim][5+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
            for i in temp2+1:cellxmax
                rho = Qbase[i,bdy_nei[dim-2],1]
                u   = Qbase[i,bdy_nei[dim-2],2]
                v   = Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 03
            temp1 = Int64(round(cellxmax/4)+1) # 1/4流出
            temp2 = Int64((temp1-1)*3+1)       # 1/4~3/4流入
            for i in 1:temp1
                rho = Qbase[i,bdy_nei[dim-2],1]
                u   = Qbase[i,bdy_nei[dim-2],2]
                v   = Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
            for i in temp1+1:temp2
                rho = bdcon[dim][2]
                u   = bdcon[dim][3]
                v   = bdcon[dim][4]
                T   = bdcon[dim][6+nch]
                p   = rho*Rhat[i,bdy_nei[dim-2]]*T
                for ns in 1:nch
                    ch[ns] =  bdcon[dim][5+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
            for i in temp2+1:cellxmax
                rho = Qbase[i,bdy_nei[dim-2],1]
                u   = Qbase[i,bdy_nei[dim-2],2]
                v   = Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 20
            for i in 1:cellxmax
                rho = Qbase[i,bdy_nei[dim-2],1]
                u   = Qbase[i,bdy_nei[dim-2],2]
                v   = Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 30
            for i in 1:cellxmax
                rho = Qbase[i,bdy_nei[dim-2],1]
                u   = -Qbase[i,bdy_nei[dim-2],2]
                v   = -Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 31
            for i in 1:cellxmax
                u   = -Qbase[i,bdy_nei[dim-2],2]
                v   = -Qbase[i,bdy_nei[dim-2],3]
                p   = Qbase[i,bdy_nei[dim-2],4]
                T   = bdcon[dim][6+nch]
                rho = p/(Rhat[i,bdy_nei[dim-2]]*T)
                for ns in 1:nch
                    ch[ns] =  Qbase[i,bdy_nei[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        elseif Int(bdcon[dim][1]) == 40
            for i in 1:cellxmax
                rho = Qbase[i,bdy_nei_rev[dim-2],1]
                u   = Qbase[i,bdy_nei_rev[dim-2],2]
                v   = Qbase[i,bdy_nei_rev[dim-2],3]
                p   = Qbase[i,bdy_nei_rev[dim-2],4]
                for ns in 1:nch
                    ch[ns] = Qbase[i,bdy_nei_rev[dim-2],4+ns]
                end
                set_bdy(Qbase, rho, u, v, p, ch, nch, bdy[dim-2],i)
            end
        end
    end

    return Qbase
end

function check_bd(bdcon)
    bd = -1
    for l in 1:4
        if     Int(bdcon[l][1]) == 00
            bd = 00
        elseif Int(bdcon[l][1]) == 01
            bd = 01
        elseif Int(bdcon[l][1]) == 02
            bd = 02
        elseif Int(bdcon[l][1]) == 03
            bd = 03

        elseif Int(bdcon[l][1]) == 20
            bd = 20

        elseif Int(bdcon[l][1]) == 30
            bd = 30
        elseif Int(bdcon[l][1]) == 31
            bd = 31

        elseif Int(bdcon[l][1]) == 40
            bd = 40
        else
            println("\n check boundary condition ! \n")
            throw(UndefVarError(:x))
        end
        println("bd condition")
        println(bd)
    end
end


# Implementation is taken from [Flux.jl](https://github.com/FluxML/Flux.jl)
"""
    Chain(layers...)
Chain multiple layers / functions together, so that they are called in sequence
on a given input.
```julia
m = Chain(x -> x^2, x -> x+1)
m(5) == 26
m = Chain(Dense(input=10, output=5), Dense(input=5, output=2))
x = rand(10)
m(x) == m[2](m[1](x))
```
`Chain` also supports indexing and slicing, e.g. `m[2]` or `m[1:end-1]`.
`m[1:3](x)` will calculate the output of the first three layers.
"""
struct Chain{T<:Tuple}
  layers::T
   Chain(xs...) = new{typeof(xs)}(xs)
end

children(c::Chain) = c.layers
mapchildren(f, c::Chain) = Chain(f.(c.layers)...)

applychain(::Tuple{}, x) = x
applychain(fs::Tuple, x) = applychain(Base.tail(fs), first(fs)(x))

(c::Chain)(x) = applychain(c.layers, x)

Base.getindex(c::Chain, i::AbstractArray) = Chain(c.layers[i]...)
Base.getindex(c::Chain, i::Integer) = c.layers[i]
Base.getindex(c::Chain, ::Colon) = c


using JSON, Compat, Mmap, HDF5, Dates, DataStructures, YAML

"""
    load_data(filename)

Load data file
"""
# add functionality to read new Auspex data files (format on develop branch)

function load_data(filename::AbstractString; load_var = false)
    # maintain original data format when variance is not avilable for compatibility with plotters and other data analysis. Use variance for tomography
    if load_var
        datasets    = Dict{String, Dict{String, Array{Any}}}()
    else
        datasets    = Dict{String, Array{Any}}()
    end
    descriptors = Dict{String,Any}()
    # per qubit
    for group in readdir(filename)
        if load_var
            datasets[group] = Dict{String, Any}()
        end
        datafile = filter(x-> occursin("data.dat", x), readdir(joinpath(filename, group)))
        varfile = filter(x-> occursin("variance.dat", x), readdir(joinpath(filename, group)))
        metafile = filter(x-> occursin("meta", x), readdir(joinpath(filename, group)))
        metadata = JSON.parsefile(joinpath(filename, group, metafile[1]), dicttype=DataStructures.OrderedDict)
        open(joinpath(filename, group, datafile[1]), "r") do op
            dims = length(metadata["shape"])>1 ? tuple(metadata["shape"]...) : tuple(metadata["shape"][1],1)
            data = Mmap.mmap(op, Array{Complex{Float64},1}, prod(dims))
            data = reshape(data, reverse(dims))
            if load_var
                datasets[group]["data"] = data
            else
                datasets[group] = data
            end
        end
        if load_var
            if ~isempty(varfile)
                open(joinpath(filename, group, varfile[1]), "r") do op
                    dims = length(metadata["shape"])>1 ? tuple(metadata["shape"]...) : tuple(metadata["shape"][1],1)
                    data = Mmap.mmap(op, Array{Complex{Float64},1}, prod(dims))
                    data = reshape(data, reverse(dims))
                    datasets[group]["variance"] = data
                end
            else
              @warn("Variance not available for dataset: $(group)")
            end
        end
        descriptors[group] = metadata
    end
    return datasets, descriptors
end

"""
    load_data(datapath, filenum; subdir, auspex)

Search file number filenum in folder datapath/subdir. Subdir default is current date in "yymmdd"
auspex: boolean for file format (default = true)
"""
function load_data(datapath::AbstractString, filenum::Int, subdir=Dates.format(Dates.today(),"yymmdd"), h5=false; load_var=false)
  # optionally, search for filenum instead of filename
  # search in a subdirectory with today's date, if not specified
  #regexpr = auspex? r"(\d{4})(.h5)" : r"(\d+)(_\w+)(.h5)"
  ext = h5 ? ".h5" : ".auspex"
  regexpr = Regex("(.+)($(lpad(filenum, 4, "0")))(\\$ext)")
  files = filter(x -> occursin(regexpr, x), readdir(joinpath(datapath, subdir)))
  if length(files) == 0
    println("No matching files found!")
    return
  elseif length(files) > 1
    println("More than one matching file found!")
    return
  end
  filename = joinpath(datapath, subdir, files[1])
  h5 ? load_h5_data(filename) : load_data(filename, load_var=load_var)
end

## TO BE UPDATED
# """
# 	load_latest_data()
#     load_latest_data(ind_from_last)
#
# Load latest data file, or ind_from_last files before the latest.
# """
# function load_latest_data(logpath::AbstractString, ind_from_last::Int = 0)
# 	df = CSV.read(logpath, delim = "\t"; use_mmap = false)
# 	@assert ind_from_last < length(df.columns[1]) "Not enough data files in the history."
# 	filename = df.columns[1][end-ind_from_last]
# 	load_auspex_data(filename)
# end

"""
    load_h5_data(filename)

Load hdf5 data and descriptor in legacy auspex format

FILE STRUCTURE:
file.hdf5:
   header/ (with yml settings, if saved)
   group1/
      data/
         axis_1 (contains all tuples visited on this axis)
         axis_2 (contains all tuples visited on this axis)
         quantity_1 (the measured values)
         quantity_2 (the measured values)
      descriptor (contains references to the axes datasets below)
      axis_1 (contains the values of the axis  see note #1)
      axis_2 (contains the values of the axis - see note #1)
   group2/
      ... same structure as above ...

Note #1: Any axes referenced by the descriptor can themselves contain
         references to multiple subaxes. This indicates that the axis
         is "unstructured", i.e. that multiple coordinates where changed
         simultaneously.
"""
function load_h5_data(filename::AbstractString)

    datasets    = Dict{String, Dict{String, Vector{Any}}}()
    descriptors = Dict{String, Vector{Dict{String,Any}}}()
	header =   Dict{String, Any}()
	#info on file, such as name, ...?

	header["filename"] = filename
    h5open(filename, "r") do f
        # Find all of the group names, which will correspond to qubits when
        # using exp_factory inside of auspex. The default group name is simply
        # "main".
        group_names = names(f)
		#load measurement/instrument settings
		if "header" in group_names
			header["settings"] = YAML.load(read(attrs(f["header"])["settings"]))
		end

        for group_name in filter!(name -> name != "header", group_names)
            g = f[group_name]

            # Read in the descriptor
            desc_refs = read(g["descriptor"])

            desc = Vector{Dict{String,Any}}()
            for desc_ref in desc_refs
                axis = Dict{String,Any}()
                if reinterpret(Bool, read(attrs(g[desc_ref])["unstructured"]))
                    axis_refs      = read(g[desc_ref])
                    axis["points"] = [read(g[axis_ref]) for axis_ref in axis_refs]
                    axis["name"]   = [read(attrs(g[axis_ref])["name"]) for axis_ref in axis_refs]
                    axis["unit"]   = [read(attrs(g[axis_ref])["unit"]) for axis_ref in axis_refs]
                else
                    axis["points"] = read(g[desc_ref])
                    axis["name"]   = read(attrs(g[desc_ref])["name"])
                    axis["unit"]   = read(attrs(g[desc_ref])["unit"])
                end
                # See if we need to get metadata:
                if "metadata" in names(attrs(g[desc_ref]))
                    meta_ref = read(attrs(g[desc_ref])["metadata"])
                    meta_enum_ref = read(attrs(g[desc_ref])["metadata_enum"])
                    axis["metadata"] = [read(g[meta_enum_ref])[m+1] for m in read(g[meta_ref])]
                end

                push!(desc, axis)
            end

            descriptors[group_name] = desc

            # Read in the actual data
            datacol_names = names(g["data"])
            datasets[group_name] = Dict{String, Vector{Any}}()
            for datacol_name in datacol_names
              dataset = read(g["data"][datacol_name])
              # heuristic to read complex numbers stored as HDF5Compounds
              datasets[group_name][datacol_name] = convert_dataset(dataset)
            end
        end
    end
    datasets, descriptors, header
end

convert_dataset(ds) = ds
# heuristic to read complex numbers stored as HDF5Compounds
function convert_dataset(ds::Vector{NamedTuple{(:r, :i), Tuple{Float32, Float32}}})
  return [complex(d.data[1], d.data[2]) for d in ds]
end


import CLIMAParameters

@testset "Override defaults" begin
    struct EarthParameterSet <: AbstractEarthParameterSet end

    ps = EarthParameterSet()
    CLIMAParameters.Planet.grav(::EarthParameterSet) = 2.0
    @test grav(ps) ≈ 2.0

end


using OpenQASM
using OpenQASM.Types
using OpenQASM.Tools
using MLStyle
using RBNF: Token
using Test

@testset "cmp_exp" begin
    @test cmp_exp(Neg(qasm_f64(0.2)), qasm_f64(-0.2))
    @test cmp_exp(qasm_f64(-0.2), Neg(qasm_f64(0.2)))
end

@testset "conversion" begin
    @test convert(String, Token{:str}("\"abc\"")) == "abc"
    @test convert(Symbol, Token{:id}("abc")) === :abc
    @test convert(Float64, Token{:float64}("0.23")) == 0.23
    @test convert(Int, Token{:int}("1")) == 1
end

@testset "qasm token helper" begin
    @test qasm_id("abc") ≈ Token{:id}("abc")
    @test qasm_id(:abc) ≈ Token{:id}("abc")
    @test qasm_int(2) ≈ Token{:int}("2")
    @test qasm_f64(2.3) ≈ Token{:float64}("2.3")
    @test qasm_str("abc") ≈ Token{:str}("\"abc\"")
end

@testset "Bit(::String[, addrs])" begin
    bit = Bit("qreg")
    @test bit.name == qasm_id("qreg")
    @test bit.address === nothing

    bit = Bit("qreg", 2)
    @test bit.name == qasm_id("qreg")
    @test bit.address == qasm_int(2)
end

@testset "qasm parser" begin

    qasm = """OPENQASM 2.0;
    include "qelib1.inc";
    gate custom(lambda) a {
        u1(sin(lambda) + 1) a;
    }
    // comment
    gate g a
    {
        U(0,0,0) a;
    }

    qreg q[4];
    creg c1[1];
    creg c2[1];
    U(-1.0, pi/2+3, 3.0) q[2];
    CX q[1], q[2];
    custom(0.3) q[3];
    barrier q;
    h q[0];
    measure q[0] -> c0[0];
    if(c0==1) z q[2];
    u3(0.1 + 0.2, 0.2, 0.3) q[0];
    """

    ast = OpenQASM.parse(qasm)
    println(ast)

    @testset "mainprogram" begin
        @test ast isa MainProgram
        @test ast.version == v"2.0.0"
    end

    @test @match ast begin
        MainProgram(v"2.0.0", progs) => true
        _ => false
    end

    @testset "include" begin
        @test ast.prog[1] isa Include
        @test ast.prog[1].file isa Token{:str}
        @test ast.prog[1].file.str == "\"qelib1.inc\""

        token = @match ast.prog[1] begin
            Include(file) => file
            _ => nothing
        end
        @test token.str == "\"qelib1.inc\""
    end

    @testset "gate custom(lambda) a" begin
        @test ast.prog[2] isa Gate
        custom = ast.prog[2]

        @test @match custom begin
            Gate(GateDecl(name, cargs, qargs), body) => true
            _ => false
        end

        @test custom.decl.name isa Token{:id}
        @test custom.decl.name.str == "custom"
        @test length(custom.decl.cargs) == 1
        @test custom.decl.cargs[1] isa Token{:id}
        @test custom.decl.cargs[1].str == "lambda"
        @test length(custom.decl.qargs) == 1
        @test custom.decl.qargs[1] isa Token{:id}
        @test custom.decl.qargs[1].str == "a"

        @test length(custom.body) == 1
        @test custom.body[1] isa Instruction
        inst = custom.body[1]
        @test inst.name == "u1"
        @test length(inst.cargs) == 1
        @test length(inst.qargs) == 1
        @test inst.cargs[1] isa Tuple
        @test inst.qargs[1] isa Bit
        carg = inst.cargs[1]
        qarg = inst.qargs[1]
        @test length(carg) == 3
        @test carg[1] isa Call
        @test carg[2] isa Token{:reserved}
        @test carg[3] isa Token{:int}
        @test carg[1].name === :sin
        @test carg[1].args.str == "lambda"
        @test carg[2].str == "+"
        @test carg[3].str == "1"
        @test qarg.name.str == "a"
        @test qarg.address === nothing
    end

    @testset "gate g" begin
        @test ast.prog[3] isa Gate
        g = ast.prog[3]
        @test isempty(g.decl.cargs)
        @test length(g.decl.qargs) == 1
        @test g.decl.qargs[1] isa Token{:id}
        @test g.decl.qargs[1].str == "a"
        @test g.decl.name isa Token{:id}
        @test g.decl.name.str == "g"

        @test length(g.body) == 1
        @test g.body[1] isa UGate
        ugate = g.body[1]
        @test ugate.qarg isa Bit
        @test ugate.qarg.name.str == "a"
        @test ugate.qarg.address === nothing
        @test ugate.y isa Token{:int}
        @test ugate.y.str == "0"
        @test ugate.z1.str == "0"
        @test ugate.z2.str == "0"
    end

    @testset "qreg" begin
        @test ast.prog[4] isa RegDecl
        reg = ast.prog[4]
        @test reg.name isa Token{:id}
        @test reg.name.str == "q"
        @test reg.size isa Token{:int}
        @test reg.size.str == "4"
        @test reg.type isa Token{:reserved}
        @test reg.type.str == "qreg"
    end

    @testset "creg" begin
        @test ast.prog[5] isa RegDecl
        reg = ast.prog[5]
        @test reg.name isa Token{:id}
        @test reg.name.str == "c1"
        @test reg.size isa Token{:int}
        @test reg.size.str == "1"
        @test reg.type isa Token{:reserved}
        @test reg.type.str == "creg"
    end

    @testset "U" begin
        @test ast.prog[7] isa UGate
        U = ast.prog[7]
        @test U.qarg isa Bit
        @test U.y isa Tuple
        @test U.z1 isa Neg
        @test U.z2 isa Token{:float64}
        @test U.qarg.name.str == "q"
        @test U.qarg.address.str == "2"
        @test U.y[2] isa Token{:reserved}
        @test U.y[2].str == "+"
        @test U.y[1] isa Tuple
        @test U.y[1][1] isa Token{:reserved}
        @test U.y[1][2] isa Token{:reserved}
        @test U.y[1][1].str == "pi"
        @test U.y[1][2].str == "/"
    end

    @testset "CX" begin
        @test ast.prog[8] isa CXGate
        cx = ast.prog[8]
        @test cx.ctrl isa Bit
        @test cx.qarg isa Bit
        @test cx.ctrl.name.str == "q"
        @test cx.qarg.name.str == "q"
        @test cx.ctrl.address.str == "1"
        @test cx.qarg.address.str == "2"
    end

    @testset "inst" begin
        @test ast.prog[9] isa Instruction
        inst = ast.prog[9]
        @test inst.name == "custom"
        @test inst.cargs[1].str == "0.3"
        @test inst.qargs[1].name.str == "q"
        @test inst.qargs[1].address.str == "3"
    end

    @testset "barrier" begin
        @test ast.prog[10] isa Barrier
        barrier = ast.prog[10]
        @test length(barrier.qargs) == 1
        @test barrier.qargs[1] isa Bit
        @test barrier.qargs[1].name.str == "q"
        @test barrier.qargs[1].address === nothing
    end

    @testset "h q[0]" begin
        @test ast.prog[11] isa Instruction
        @test ast.prog[11].name == "h"
        @test length(ast.prog[11].qargs) == 1
        @test length(ast.prog[11].cargs) == 0
    end

    @testset "measure" begin
        m = ast.prog[12]
        @test m isa Measure
        @test m.qarg isa Bit
        @test m.carg isa Bit
        @test m.qarg.name isa Token{:id}
        @test m.qarg.name.str == "q"
        @test m.carg.name isa Token{:id}
        @test m.carg.name.str == "c0"
        @test m.qarg.address isa Token{:int}
        @test m.carg.address isa Token{:int}
        @test m.qarg.address.str == "0"
        @test m.carg.address.str == "0"
    end

    @testset "ifstmt" begin
        ifstmt = ast.prog[13]
        @test ifstmt isa IfStmt
        @test ifstmt.left isa Token{:id}
        @test ifstmt.left.str == "c0"
        @test ifstmt.right isa Token{:int}
        @test ifstmt.right.str == "1"
        @test ifstmt.body isa Instruction
        @test ifstmt.body.name == "z"
        @test length(ifstmt.body.qargs) == 1
        @test length(ifstmt.body.cargs) == 0
        @test ifstmt.body.qargs[1].name.str == "q"
        @test ifstmt.body.qargs[1].address.str == "2"
    end

    @testset "u3" begin
        inst = ast.prog[14]
        @test inst isa Instruction
        @test inst.name == "u3"
        @test length(inst.cargs) == 3
        @test inst.cargs[1][1].str == "0.1"
        @test inst.cargs[1][2].str == "+"
        @test inst.cargs[1][3].str == "0.2"
        @test inst.cargs[2].str == "0.2"
        @test inst.cargs[3].str == "0.3"
    end

end

@testset "cmp_ast" begin
    qasm1 = """OPENQASM 2.0;
    include "qelib1.inc";
    gate custom(lambda) a {
        u1(sin(lambda) + 1) a;
    }
    // comment
    gate g a
    {
        U(0,0,0) a;
    }

    qreg q[4];
    creg c1[1];
    creg c2[1];
    U(-1.0, pi/2+3, 3.0) q[2];
    CX q[1], q[2];
    custom(0.3) q[3];
    barrier q;
    h q[0];
    measure q[0] -> c0[0];
    if(c0==1) z q[2];
    u3(0.1 + 0.2, 0.2, 0.3) q[0];
    """

    qasm2 = """OPENQASM 2.0;
    include "qelib1.inc";
    gate custom(lambda) a {
        u1(sin(lambda) + 1) a;
    }
    // comment
    gate g a
    {
        U(0,0,0) a;
    }

    qreg q[4];
    creg c1[1];
    creg c2[1];
    U(-1.0, pi/2+3, 3.0) q[2];
    CX q[1], q[2];
    custom(0.3) q[3];
    barrier q;
    h q[0];
    measure q[0] -> c0[0];
    if(c0==1) z q[2];
    """

    ast1 = OpenQASM.parse(qasm1)
    ast2 = OpenQASM.parse(qasm2)
    println(ast1)
    println(ast2)

    @test !(ast1 ≈ ast2)
    @test ast1 ≈ ast1
    @test ast2 ≈ ast2

    s = """
    gate test_gate(theta, phi) qreg_2, qreg_3, qreg_1 {
    x qreg_1;
    z qreg_2;
    rx(-(sin(theta))+2.0) qreg_3;
    rx(sin(theta)+tan(phi)) qreg_3;
    ry(cos(phi)-ln(phi)) qreg_3;
    ry(cos(phi)*sqrt(phi)) qreg_3;
    CX qreg_1, qreg_3;
    }
    """
    ast3 = OpenQASM.parse_gate(s)
    println(ast3)
    @test ast3 ≈ ast3
end


## Deprecated

# deprecated at 0.13
@deprecate kmpp(X, k) initseeds(:kmpp, X, k)
@deprecate kmpp_by_costs(costs, k) initseeds_by_costs(:kmpp, costs, k)

# deprecated at 0.13.1
@deprecate copyseeds(X, iseeds) copyseeds!(Matrix{eltype(X)}(undef, size(X, 1), length(iseeds)), X, iseeds)

# deprecated as of 0.13.2
@deprecate varinfo(k1::Int, a1::AbstractVector{Int},
                   k2::Int, a2::AbstractVector{Int}) varinfo(a1, a2)
@deprecate varinfo(R::ClusteringResult, k0::Int, a0::AbstractVector{Int}) varinfo(R, a0)

# FIXME remove after deprecation period for merge/labels/height/method
Base.propertynames(hclu::Hclust, private::Bool = false) =
    (:merges, :heights, :order, :linkage,
     #= deprecated as of 0.12 =# :height, :labels, :merge, :method)

# FIXME remove after deprecation period for merge/labels/height/method
@inline function Base.getproperty(hclu::Hclust, prop::Symbol)
    if prop === :height
        Base.depwarn("Hclust::height is deprecated, use Hclust::heights", Symbol("Hclust::height"))
        return getfield(hclu, :heights)
    elseif prop === :labels
        Base.depwarn("Hclust::labels is deprecated and will be removed in future versions", Symbol("Hclust::labels"))
        return 1:nnodes(hclu)
    elseif prop === :merge
        Base.depwarn("Hclust::merge is deprecated, use Hclust::merges", Symbol("Hclust::merge"))
        return getfield(hclu, :merges)
    elseif prop === :method
        Base.depwarn("Hclust::method is deprecated, use Hclust::linkage", Symbol("Hclust::method"))
        return getfield(hclu, :linkage)
    else
        return getfield(hclu, prop)
    end
end

# FIXME remove after deprecation period for cweights
Base.propertynames(clu::KmeansResult, private::Bool = false) =
    (:centers, :assignments, :costs, :counts, :wcounts,
     :totalcost, :iterations, :converged,
     #= deprecated as of 0.13.2 =# :cweights)

# FIXME remove after deprecation period for cweights
@inline function Base.getproperty(clu::KmeansResult, prop::Symbol)
    if prop === :cweights
        Base.depwarn("KmeansResult::cweights is deprecated, use wcounts(clu::KmeansResult)",
                     Symbol("KmeansResult::cweights"))
        return clu.wcounts
    else
        return getfield(clu, prop)
    end
end


module TypeEmulator

export
    emulate, MainModule, CoreModule, NewEmpty,
    NewModule, NewDataType, NewAny, NewBottom, NewFunction,
    NewType, NewTypeName, NewTypeVar, NewUnionAll, NewUnion,
    NewFunction, NewMethod, NewMethodList, NewMethodTable, NewTypeMapEntry,
    supertypes, isnewsubtype, linearize, register!

import Base:    show, subtypes, supertype

include("metadata.jl")
include("types.jl")
include("typeconversions.jl")
include("linear.jl")
include("relations.jl")
include("variables.jl")

end # module


""" Syntax for a generalized algebraic theory (GAT).

Unlike instances of a theory, syntactic expressions don't necessarily satisfy
the equations of the theory. For example, the default syntax operations for the
`Category` theory don't form a category because they don't satisfy the category
laws, e.g.,
```
compose(f, id(A)) != compose(f)
```
Whether dependent types are enforced at runtime and whether expressions are
automatically brought to normal form depends on the particular syntax. In
general, a single theory may have many different syntaxes. The purpose of this
module to make the construction of syntax simple but flexible.
"""
module Syntax
export @syntax, GATExpr, SyntaxDomainError, head, args, type_args, first, last,
  invoke_term, functor, to_json_sexpr, parse_json_sexpr, show_sexpr,
  show_unicode, show_latex

import Base: first, last
import Base.Meta: ParseError, show_sexpr
using Compat
using Match

using ..GAT: Context, Signature, TypeConstructor, TermConstructor, Typeclass
import ..GAT
import ..GAT: invoke_term
using ..Meta

# Data types
############

""" Base type for expression in the syntax of a GAT.

We define Julia types for each *type constructor* in the theory, e.g., object,
morphism, and 2-morphism in the theory of 2-categories. Of course, Julia's
type system does not support dependent types, so the type parameters are
incorporated in the Julia types. (They are stored as extra data in the
expression instances.)

The concrete types are structurally similar to the core type `Expr` in Julia.
However, the *term constructor* is represented as a type parameter, rather than
as a `head` field. This makes dispatch using Julia's type system more
convenient.
"""
abstract type GATExpr{T} end

head(::GATExpr{T}) where T = T
args(expr::GATExpr) = expr.args
first(expr::GATExpr) = first(args(expr))
last(expr::GATExpr) = last(args(expr))
type_args(expr::GATExpr) = expr.type_args

function Base.:(==)(e1::GATExpr, e2::GATExpr)
  head(e1) == head(e2) && args(e1) == args(e2) && type_args(e1) == type_args(e2)
end
function Base.hash(e::GATExpr, h::UInt)
  hash(args(e), hash(head(e), h))
end

function Base.show(io::IO, expr::GATExpr)
  print(io, head(expr))
  print(io, "(")
  join(io, args(expr), ",")
  print(io, ")")
end
function Base.show(io::IO, expr::GATExpr{:generator})
  value = first(expr)
  if isnothing(value)
    show(io, value) # Value `nothing` cannot be printed
  else
    print(io, value)
  end
end

struct SyntaxDomainError <: Exception
  constructor::Symbol
  args::Vector
end

function Base.showerror(io::IO, exc::SyntaxDomainError)
  print(io, "Domain error in term constructor $(exc.constructor)(")
  join(io, exc.args, ",")
  print(io, ")")
end

# Syntax
########

""" Define a *syntax* system for a generalized algebraic theory (GAT).

A syntax system consists of Julia types (with top type `GATExpr`) for each type
constructor in the signature, plus Julia functions for

1. *Generators*: creating new generator terms, e.g., objects or morphisms
2. *Accessors*: accessing type parameters, e.g., domains and codomains
3. *Term constructors*: applying term constructors, e.g., composition and
   monoidal products

Julia code for all this is generated by the macro. Any of the methods can be
overriden with custom simplification logic.
"""
macro syntax(syntax_head, mod_name, body=Expr(:block))
  @assert body.head == :block
  syntax_name, base_types = @match syntax_head begin
    Expr(:call, [name::Symbol, args...]) => (name, args)
    name::Symbol => (name, [])
    _ => throw(ParseError("Ill-formed syntax signature $syntax_head"))
  end
  functions = map(parse_function, strip_lines(body).args)

  expr = Expr(:call, :syntax_code, Expr(:quote, syntax_name),
              esc(Expr(:ref, :Type, base_types...)),
              esc(mod_name), esc(nameof(__module__)), functions)
  Expr(:block,
    Expr(:call, esc(:eval), expr),
    :(Core.@__doc__ $(esc(syntax_name))))
end
function syntax_code(name::Symbol, base_types::Vector{Type},
                     signature_module::Module, outer_module::Module,
                     functions::Vector)
  class = signature_module.class()
  signature = class.signature
  signature_ref = GlobalRef(parentmodule(signature_module),
                            nameof(signature_module))

  # Generate module with syntax types and type/term generators.
  mod = Expr(:module, true, name,
    Expr(:block, [
      # Prevents error about export not being at toplevel.
      # https://github.com/JuliaLang/julia/issues/28991
      LineNumberNode(0);
      Expr(:export, [cons.name for cons in signature.types]...);
      Expr(:using, Expr(:., :., :., nameof(outer_module)));
      :(signature() = $signature_ref);
      gen_types(signature, base_types);
      gen_type_accessors(signature);
      gen_term_generators(signature, outer_module);
      gen_term_constructors(signature, outer_module);
    ]...))

  # Generate toplevel functions.
  toplevel = []
  bindings = Dict{Symbol,Any}(
    c.name => Expr(:(.), name, QuoteNode(c.name)) for c in signature.types)
  bindings[:Super] = name
  syntax_fns = Dict(parse_function_sig(f) => f for f in functions)
  for f in interface(class)
    sig = parse_function_sig(f)
    if haskey(syntax_fns, sig)
      # Case 1: The method is overriden in the syntax body.
      expr = generate_function(replace_symbols(bindings, syntax_fns[sig]))
    elseif !isnothing(f.impl)
      # Case 2: The method has a default implementation in the signature.
      expr = generate_function(replace_symbols(bindings, f))
    else
      # Case 3: Call the default syntax method.
      params = [ gensym("x$i") for i in eachindex(sig.types) ]
      call_expr = Expr(:call, sig.name,
        [ Expr(:(::), p, t) for (p,t) in zip(params, sig.types) ]...)
      body = Expr(:call, Expr(:(.), name, QuoteNode(sig.name)), params...)
      f_impl = JuliaFunction(call_expr, f.return_type, body)
      expr = generate_function(replace_symbols(bindings, f_impl))
    end
    push!(toplevel, expr)
  end
  Expr(:toplevel, mod, toplevel...)
end

""" Complete set of Julia functions for a syntax system.
"""
function interface(class::Typeclass)::Vector{JuliaFunction}
  sig = class.signature
  [ GAT.interface(class);
    [ GAT.constructor(constructor_for_generator(cons), sig)
      for cons in sig.types ]; ]
end

""" Generate syntax type definitions.
"""
function gen_type(cons::TypeConstructor, base_type::Type=Any)::Expr
  base_expr = GlobalRef(Syntax, :GATExpr)
  base_name = if base_type == Any
    base_expr
  else
    GlobalRef(parentmodule(base_type), nameof(base_type))
  end
  expr = :(struct $(cons.name){T} <: $base_name{T}
    args::Vector
    type_args::Vector{$base_expr}
  end)
  generate_docstring(strip_lines(expr, recurse=true), cons.doc)
end
function gen_types(sig::Signature, base_types::Vector{Type})::Vector{Expr}
  if isempty(base_types)
    map(gen_type, sig.types)
  else
    map(gen_type, sig.types, base_types)
  end
end

""" Generate accessor methods for type parameters.
"""
function gen_type_accessors(cons::TypeConstructor)::Vector{Expr}
  fns = []
  sym = gensym(:x)
  for (i, param) in enumerate(cons.params)
    call_expr = Expr(:call, param, Expr(:(::), sym, cons.name))
    return_type = GAT.strip_type(cons.context[param])
    body = Expr(:ref, Expr(:(.), sym, QuoteNode(:type_args)), i)
    push!(fns, generate_function(JuliaFunction(call_expr, return_type, body)))
  end
  fns
end
function gen_type_accessors(sig::Signature)::Vector{Expr}
  vcat(map(gen_type_accessors, sig.types)...)
end

""" Generate methods for syntax term constructors.
"""
function gen_term_constructor(cons::TermConstructor, sig::Signature,
                              mod::Module; dispatch_type::Symbol=Symbol())::Expr
  head = GAT.constructor(cons, sig)
  call_expr, return_type = head.call_expr, head.return_type
  if dispatch_type == Symbol()
    dispatch_type = cons.name
  end
  body = Expr(:block)

  # Create expression to check constructor domain.
  eqs = GAT.equations(cons, sig)
  if !isempty(eqs)
    clauses = [ Expr(:call,:(==),lhs,rhs) for (lhs,rhs) in eqs ]
    conj = foldr((x,y) -> Expr(:(&&),x,y), clauses)
    insert!(call_expr.args, 2,
      Expr(:parameters, Expr(:kw, :strict, false)))
    push!(body.args,
      Expr(:if,
        Expr(:(&&), :strict, Expr(:call, :(!), conj)),
        Expr(:call, :throw,
          Expr(:call, GlobalRef(Syntax, :SyntaxDomainError),
            Expr(:quote, cons.name),
            Expr(:vect, cons.params...)))))
  end

  # Create call to expression constructor.
  type_params = gen_term_constructor_params(cons, sig, mod)
  push!(body.args,
    Expr(:call,
      Expr(:curly, return_type, Expr(:quote, dispatch_type)),
      Expr(:vect, cons.params...),
      Expr(:vect, type_params...)))

  generate_function(JuliaFunction(call_expr, return_type, body))
end
function gen_term_constructors(sig::Signature, mod::Module)::Vector{Expr}
  [ gen_term_constructor(cons, sig, mod) for cons in sig.terms ]
end

""" Generate expressions for type parameters of term constructor.

Besides expanding the implicit variables, we must handle two annoying issues:

1. Add types for method dispatch where necessary (see `GAT.add_type_dispatch`)
   FIXME: We are currently only handling the nullary case (e.g., `munit()`).
   To handle the general case, we need to do basic type inference.

2. Rebind the term constructors to ensure that user overrides are preferred over
   the default term constructors.
"""
function gen_term_constructor_params(cons, sig, mod)::Vector
  expr = GAT.expand_term_type(cons, sig)
  raw_params = @match expr begin
    Expr(:call, [name::Symbol, args...]) => args
    _::Symbol => []
  end

  bindings = Dict(c.name => GlobalRef(mod, c.name) for c in sig.terms)
  params = []
  for expr in raw_params
    expr = replace_nullary_constructors(expr, sig)
    expr = replace_symbols(bindings, expr)
    push!(params, expr)
  end
  params
end
function replace_nullary_constructors(expr, sig)
  @match expr begin
    Expr(:call, [name::Symbol]) => begin
      terms = sig.terms[findall(cons -> cons.name == name, sig.terms)]
      @assert length(terms) == 1
      Expr(:call, name, terms[1].typ)
    end
    Expr(:call, [name::Symbol, args...]) =>
      Expr(:call, name, [replace_nullary_constructors(a,sig) for a in args]...)
    _ => expr
  end
end

""" Generate methods for term generators.

Generators are extra term constructors created automatically for the syntax.
"""
function gen_term_generator(cons::TypeConstructor, sig::Signature, mod::Module)::Expr
  gen_term_constructor(constructor_for_generator(cons), sig, mod;
                       dispatch_type = :generator)
end
function gen_term_generators(sig::Signature, mod::Module)::Vector{Expr}
  [ gen_term_generator(cons, sig, mod) for cons in sig.types ]
end
function constructor_for_generator(cons::TypeConstructor)::TermConstructor
  value_param = :__value__
  params = [ value_param; cons.params ]
  typ = Expr(:call, cons.name, cons.params...)
  context = merge(Context(value_param => :Any), cons.context)
  TermConstructor(cons.name, params, typ, context)
end

# Reflection
############

""" Invoke a term constructor by name in a syntax system.

This method provides reflection for syntax systems. In everyday use the generic
method for the constructor should be called directly, not through this function.
"""
function invoke_term(syntax_module::Module, constructor_name::Symbol, args...)
  signature_module = syntax_module.signature()
  signature = signature_module.class().signature
  syntax_types = Tuple(getfield(syntax_module, cons.name) for cons in signature.types)
  invoke_term(signature_module, syntax_types, constructor_name, args...)
end

""" Name of constructor that created expression.
"""
function constructor_name(expr::GATExpr)::Symbol
  if head(expr) == :generator
    nameof(typeof(expr))
  else
    head(expr)
  end
end

""" Create generator of the same type as the given expression.
"""
function generator_like(expr::GATExpr, value)::GATExpr
  invoke_term(
    parentmodule(typeof(expr)),
    nameof(typeof(expr)),
    value,
    type_args(expr)...
  )
end

# Functors
##########

""" Functor from GAT expression to GAT instance.

Strictly speaking, we should call these "structure-preserving functors" or,
better, "model homomorphisms of GATs". But this is a category theory library,
so we'll go with the simpler "functor".

A functor is completely determined by its action on the generators. There are
several ways to specify this mapping:

  1. Simply specify a Julia instance type for each doctrine type, using the
     required `types` tuple. For this to work, the generator constructors
     must be defined for the instance types.

  2. Explicitly map each generator term to an instance value, using the
     `generators` dictionary.

  3. For each doctrine type (e.g., object and morphism), specify a function
     mapping generator terms of that type to an instance value, using the
     `terms` dictionary.

The `terms` dictionary can also be used for special handling of non-generator
expressions. One use case for this capability is defining forgetful functors,
which map non-generators to generators.
"""
function functor(types::Tuple, expr::GATExpr;
                 generators::AbstractDict=Dict(), terms::AbstractDict=Dict())
  # Special case: look up a specific generator.
  if head(expr) == :generator && haskey(generators, expr)
    return generators[expr]
  end

  # Special case: look up by type of term (usually a generator).
  name = constructor_name(expr)
  if haskey(terms, name)
    return terms[name](expr)
  end

  # Otherwise, we need to call a term constructor (possibly for a generator).
  # Recursively evalute the arguments.
  term_args = []
  for arg in args(expr)
    if isa(arg, GATExpr)
      arg = functor(types, arg; generators=generators, terms=terms)
    end
    push!(term_args, arg)
  end

  # Invoke the constructor in the codomain category!
  syntax_module = parentmodule(typeof(expr))
  signature_module = syntax_module.signature()
  invoke_term(signature_module, types, name, term_args...)
end

# Serialization
###############

""" Serialize expression as JSON-able S-expression.

The format is an S-expression encoded as JSON, e.g., "compose(f,g)" is
represented as ["compose", f, g].
"""
function to_json_sexpr(expr::GATExpr;
    by_reference::Function = default_to_json_by_reference
  )
  if head(expr) == :generator && by_reference(first(expr))
    to_json_sexpr(first(expr))
  else
    [ string(constructor_name(expr)); map(to_json_sexpr, args(expr)) ]
  end
end
to_json_sexpr(::Nothing) = nothing
to_json_sexpr(x::String) = x
to_json_sexpr(x::Real) = x
to_json_sexpr(x::Bool) = x
to_json_sexpr(x) = string(x)

default_to_json_by_reference(x) = false

""" Deserialize expression from JSON-able S-expression.

If `symbols` is true (the default), strings are converted to symbols.
"""
function parse_json_sexpr(syntax_module::Module, sexpr;
    parse_head::Function = identity,
    parse_reference::Function = disable_parse_json_reference,
    parse_value::Function = identity,
    symbols::Bool = true,
  )
  signature_module = syntax_module.signature()
  signature = signature_module.class().signature
  type_lens = Dict(cons.name => length(cons.params) for cons in signature.types)

  # The "value" type is
  # - false if the sexpr should be parsed as a term (GAT expression)
  # - true if the sexpr should be parsed as a value (Julia object)
  function parse_impl(sexpr::Vector, ::Type{Val{false}})
    name = Symbol(parse_head(symbols ? Symbol(sexpr[1]) : sexpr[1]))
    nargs = length(sexpr) - 1
    args = [
      parse_impl(arg, Val{
        (i == 1 && get(type_lens, name, nothing) == nargs-1) ||
        isa(arg, Bool) || isa(arg, Number) || isa(arg, Nothing)
      })
      for (i, arg) in enumerate(sexpr[2:end])
    ]
    invoke_term(syntax_module, name, args...)
  end
  parse_impl(x, ::Type{Val{true}}) = parse_value(x)
  parse_impl(x::String, ::Type{Val{false}}) = parse_reference(symbols ? Symbol(x) : x)
  parse_impl(x::String, ::Type{Val{true}}) = parse_value(symbols ? Symbol(x) : x)

  parse_impl(sexpr, Val{false})
end

disable_parse_json_reference(x) = error("Loading terms by name is not enabled")

# Pretty-print
##############

""" Show the syntax expression as an S-expression.

Cf. the standard library function `Meta.show_sexpr`.
"""
show_sexpr(expr::GATExpr) = show_sexpr(stdout, expr)

function show_sexpr(io::IO, expr::GATExpr)
  if head(expr) == :generator
    print(io, repr(first(expr)))
  else
    print(io, "(")
    join(io, [string(head(expr));
              [sprint(show_sexpr, arg) for arg in args(expr)]], " ")
    print(io, ")")
  end
end

""" Show the expression in infix notation using Unicode symbols.
"""
show_unicode(expr::GATExpr) = show_unicode(stdout, expr)
show_unicode(io::IO, x::Any; kw...) = show(io, x)

function show_unicode(io::IO, expr::GATExpr; kw...)
  # By default, show in prefix notation.
  print(io, head(expr))
  print(io, "[")
  join(io, [sprint(show_unicode, arg) for arg in args(expr)], ",")
  print(io, "]")
end

function show_unicode(io::IO, expr::GATExpr{:generator}; kw...)
  print(io, first(expr))
end

function show_unicode_infix(io::IO, expr::GATExpr, op::String;
                            paren::Bool=false)
  show_unicode_paren(io, expr) = show_unicode(io, expr; paren=true)
  if (paren) print(io, "(") end
  join(io, [sprint(show_unicode_paren, arg) for arg in args(expr)], op)
  if (paren) print(io, ")") end
end

""" Show the expression in infix notation using LaTeX math.

Does *not* include `\$` or `\\[begin|end]{equation}` delimiters.
"""
show_latex(expr::GATExpr) = show_latex(stdout, expr)
show_latex(io::IO, sym::Symbol; kw...) = print(io, sym)
show_latex(io::IO, x::Any; kw...) = show(io, x)

function show_latex(io::IO, expr::GATExpr; kw...)
  # By default, show in prefix notation.
  print(io, "\\mathop{\\mathrm{$(head(expr))}}")
  print(io, "\\left[")
  join(io, [sprint(show_latex, arg) for arg in args(expr)], ",")
  print(io, "\\right]")
end

function show_latex(io::IO, expr::GATExpr{:generator}; kw...)
  # Try to be smart about using text or math mode.
  content = string(first(expr))
  if all(isletter, content) && length(content) > 1
    print(io, "\\mathrm{$content}")
  else
    print(io, content)
  end
end

function show_latex_infix(io::IO, expr::GATExpr, op::String;
                          paren::Bool=false, kw...)
  show_latex_paren(io, expr) = show_latex(io, expr; paren=true, kw...)
  sep = op == " " ? op : " $op "
  if (paren) print(io, "\\left(") end
  join(io, [sprint(show_latex_paren, arg) for arg in args(expr)], sep)
  if (paren) print(io, "\\right)") end
end

function show_latex_script(io::IO, expr::GATExpr, head::String;
                           super::Bool=false, kw...)
  print(io, head, super ? "^" : "_", "{")
  join(io, [sprint(show_latex, arg) for arg in args(expr)], ",")
  print(io, "}")
end

end


# Variables #
# ========================== #
function JuMP.build_variable(_error::Function, variable::JuMP.ScalarVariable, set::DecisionSet)
    return VariableConstrainedOnCreation(variable, SingleDecisionSet(Decision(variable.info, Float64), set.constraint))
end

function JuMP.build_variable(_error::Function, variables::Vector{<:JuMP.ScalarVariable}, set::DecisionSet)
    return VariablesConstrainedOnCreation(variables, MultipleDecisionSet([Decision(variable.info, Float64) for variable in variables], set.constraint))
end

function JuMP.build_variable(_error::Function, variable::JuMP.ScalarVariable, ::KnownSet)
    return VariableConstrainedOnCreation(variable, SingleKnownSet(KnownDecision(variable.info, Float64)))
end

function JuMP.build_variable(_error::Function, variables::Vector{<:JuMP.ScalarVariable}, ::KnownSet)
    return VariablesConstrainedOnCreation(variables, MultipleKnownSet([KnownDecision(variable.info, Float64) for variable in variables]))
end

function JuMP.add_variable(model::Model, variable::VariableConstrainedOnCreation{<:Union{SingleDecisionSet, SingleKnownSet}}, name::String)
    decisions = get_decisions(model)
    if decisions isa IgnoreDecisions
        # Create a regular JuMP variable if decisions are not handled
        return JuMP.add_variable(model, variable.scalar_variable, name)
    end
    var_index, con_index = MOI.add_constrained_variable(backend(model), variable.set)
    # Map to model decisions after indices are known
    if !has_decision(decisions, var_index)
        # Store decision if is seen for the first time
        set_decision!(decisions, var_index, variable.set)
    else
        # Reuse if decision has been created already
        MOI.set(backend(model), MOI.ConstraintSet(), con_index, reuse(variable.set, decision(decisions, var_index)))
    end
    # Add any given decision constraints
    _moi_constrain_decision(backend(model), var_index, variable.scalar_variable.info, variable.set)
    # Finally, set any given name
    if !isempty(name)
        MOI.set(backend(model), MOI.VariableName(), var_index, name)
    end
    # Return created decision as DecisionRef/KnownRef
    return VariableRef(model, var_index, variable.set)
end

function JuMP.add_variable(model::Model, variable::VariablesConstrainedOnCreation{<:Union{MultipleDecisionSet, MultipleKnownSet}}, names)
    decisions = get_decisions(model)
    if decisions isa IgnoreDecisions
        # Create regular JuMP variables if decisions are not handled
        var_refs = map(zip(variable.scalar_variables, names)) do (scalar_variable, name)
            JuMP.add_variable(model, scalar_variable, name)
        end
        return reshape_vector(var_refs, variable.shape)
    end
    var_indices, con_index = MOI.add_constrained_variables(backend(model), variable.set)
    # Map to model decisions after indices are known
    seen_decisions = Vector{Decision{Float64}}()
    for (i,var_index) in enumerate(var_indices)
        if !has_decision(decisions, var_index)
            # Store decision if is seen for the first time
            set_decision!(decisions, var_index, i, variable.set)
        else
            # Reuse if decision has been created already
            push!(seen_decisions, decision(decisions, var_index))
        end
    end
    if !isempty(seen_decisions)
        # Sanity check
        length(seen_decisions) == length(variable.scalar_variables) || error("Inconsistency in number of seen decisions and created variables.")
        # Update decision set for reuse
        MOI.set(backend(model), MOI.ConstraintSet(), con_index, reuse(variable.set, seen_decisions))
    end
    # Add any given decision constraints
    for (index, scalar_variable) in zip(var_indices, variable.scalar_variables)
        _moi_constrain_decision(backend(model), index, scalar_variable.info, variable.set)
    end
    # Finally, set any given names
    for (var_index, name) in zip(var_indices, JuMP.vectorize(names, variable.shape))
        if !isempty(name)
            MOI.set(backend(model), MOI.VariableName(), var_index, name)
        end
    end
    # Return created decisions as DecisionRefs/KnownRefs
    refs = map(var_indices) do index
        VariableRef(model, index, variable.set)
    end
    return reshape_vector(refs, variable.shape)
end

# Containers #
# ========================== #
const DenseAxisArray = JuMP.Containers.DenseAxisArray
const SparseAxisArray = JuMP.Containers.SparseAxisArray

struct DecisionDenseAxisArray{V <: Union{DecisionRef, KnownRef}, A <: DenseAxisArray, S <: MOI.AbstractSet}
    array::A
    constraint::S

    function DecisionDenseAxisArray{V}(array::DenseAxisArray, constraint::MOI.AbstractSet) where V <: Union{DecisionRef, KnownRef}
        A = typeof(array)
        S = typeof(constraint)
        return new{V, A, S}(array, constraint)
    end
end

struct DecisionSparseAxisArray{V <: Union{DecisionRef, KnownRef}, A <: SparseAxisArray, S <: MOI.AbstractSet}
    array::A
    constraint::S

    function DecisionSparseAxisArray{V}(array::SparseAxisArray, constraint::MOI.AbstractSet) where V <: Union{DecisionRef, KnownRef}
        A = typeof(array)
        S = typeof(constraint)
        return new{V, A, S}(array, constraint)
    end
end

function JuMP.build_variable(_error::Function, variables::DenseAxisArray{<:JuMP.ScalarVariable}, set::DecisionSet)
    return DecisionDenseAxisArray{DecisionRef}(variables, set.constraint)
end

function JuMP.build_variable(_error::Function, variables::DenseAxisArray{<:JuMP.ScalarVariable}, ::KnownSet)
    return DecisionDenseAxisArray{KnownRef}(variables, NoSpecifiedConstraint())
end

function JuMP.build_variable(_error::Function, variables::SparseAxisArray{<:JuMP.ScalarVariable}, set::DecisionSet)
    return DecisionSparseAxisArray{DecisionRef}(variables, set.constraint)
end

function JuMP.build_variable(_error::Function, variables::SparseAxisArray{<:JuMP.ScalarVariable}, ::KnownSet)
    return DecisionSparseAxisArray{KnownRef}(variables, NoSpecifiedConstraint())
end

function JuMP.add_variable(model::Model, variable::DecisionDenseAxisArray{V}, names::DenseAxisArray{String}) where V <: Union{DecisionRef, KnownRef}
    array = variable.array
    refs = DenseAxisArray{V}(undef, array.axes...)
    for idx in eachindex(array)
        var = array[idx]
        refs[idx] = JuMP.add_variable(model,
                                      VariableConstrainedOnCreation(var, set(var, V, variable.constraint)),
                                      names[idx])
    end
    return refs
end

function JuMP.add_variable(model::Model,
                           variable::DecisionSparseAxisArray{V, SparseAxisArray{T,N,K}},
                           names::SparseAxisArray{String}) where {V <: Union{DecisionRef, KnownRef}, T, N, K}
    refs = SparseAxisArray(Dict{K,V}())
    for (idx, var) in variable.array.data
        refs[idx] = JuMP.add_variable(model,
                                      VariableConstrainedOnCreation(var, set(var, V, variable.constraint)),
                                      names[idx])
    end
    return refs
end

# Constraints #
# ========================== #
function JuMP.build_constraint(_error::Function, aff::Union{DecisionAffExpr, DecisionQuadExpr}, set::MOI.AbstractScalarSet)
    offset = constant(aff.variables)
    JuMP.add_to_expression!(aff.variables, -offset)
    shifted_set = MOIU.shift_constant(set, -offset)
    return JuMP.ScalarConstraint(aff, shifted_set)
end

function JuMP.build_constraint(_error::Function, aff::Union{DecisionAffExpr, DecisionQuadExpr}, lb, ub)
    JuMP.build_constraint(_error, aff, MOI.Interval(lb, ub))
end

# Helper function #
# ========================== #
function _moi_constrain_decision(backend::MOI.ModelLike, index, info, set::Union{SingleDecisionSet, MultipleDecisionSet})
    # We don't call the _moi* versions (e.g., _moi_set_lower_bound) because they
    # have extra checks that are not necessary for newly created variables.
    nothing_added = true
    if info.has_lb
        MOI.add_constraint(backend, SingleDecision(index),
                           MOI.GreaterThan{Float64}(info.lower_bound))
        nothing_added &= false
    end
    if info.has_ub
        MOI.add_constraint(backend, SingleDecision(index),
                           MOI.LessThan{Float64}(info.upper_bound))
        nothing_added &= false
    end
    if info.has_fix
        MOI.add_constraint(backend, SingleDecision(index),
                           MOI.EqualTo{Float64}(info.fixed_value))
        nothing_added &= false
    end
    if info.binary
        MOI.add_constraint(backend, SingleDecision(index),
                           MOI.ZeroOne())
        nothing_added &= false
    end
    if info.integer
        MOI.add_constraint(backend, SingleDecision(index), MOI.Integer())
        nothing_added &= false
    end
    if info.has_start
        MOI.set(backend, MOI.VariablePrimalStart(), index,
                Float64(info.start))
        nothing_added &= false
    end
    # Mark decision as free if not bounded or otherwise constrained
    if nothing_added && set.constraint isa NoSpecifiedConstraint
        MOI.add_constraint(backend, SingleDecision(index), FreeDecision())
    end
end

function _moi_constrain_decision(backend::MOI.ModelLike, index, info, ::Union{SingleKnownSet, MultipleKnownSet})
    # Known decisions are not constrained, just return
    return nothing
end


# This file was generated, do not modify it. # hide
DecisionTreeRegressor = @load DecisionTreeRegressor pkg=DecisionTree
forest = EnsembleModel(model=DecisionTreeRegressor())


### Methods ###

function ThalamusModel(config::Dict)
  numTR = get(config, "numTR", 8)
  numtRN = get(config, "numtRN", 1)
  szTR = get(config, "szTR", 64)
  sztRN = get(config, "sztRN", 64)

  agent = Agent()

  agent_TR = Agent()
  for i = 1:numTR
    addnode!(agent_TR, GenericNeurons(szTR; b=0.25, d=0.05); name="$i")
  end
  agent_tRN = Agent()
  for i = 1:numtRN
    addnode!(agent_tRN, GenericNeurons(sztRN; a=0.1, d=2); name="$i")
  end
  agent["TR"] = agent_TR
  agent["tRN"] = agent_tRN

  # Reciprocally connect TR and tRN
  for i = 1:numTR
    for j = 1:numtRN
      # TODO: Delays
      # TODO: Masking
      addnode!(agent, GenericSynapses(szTR, sztRN; outputMask=-1); name="TR_$(i)_tRN_$(j)")
      addedge!(agent, agent["TR_$(i)_tRN_$(j)"], (
        (agent["TR"]["$i"], :input),
        (agent["tRN"]["$j"], :output)
      ))
      addnode!(agent, GenericSynapses(sztRN, szTR); name="tRN_$(j)_TR_$(i)")
      addedge!(agent, agent["tRN_$(j)_TR_$(i)"], (
        (agent["tRN"]["$j"], :input),
        (agent["TR"]["$i"], :output)
      ))
    end
  end

  return agent
end


# ============================================================
# Structs of Particle
# ============================================================

"""
    Particle(M, X, V, E, IDX, FLAG = zeros(eltype(IDX), axes(IDX, 1)), T = zero(M))

Struct of arrays for particle simulation

@vars: m, x, v, e, idx, flag, tc

"""
mutable struct Particle{T1,T2,T3,T4,T5} <: AbstractParticle

    m::T1
    x::T2
    v::T3
    e::T1
    idx::T4
    ref::T4
    flag::T5
    tc::T1

    function Particle(
        M,
        X,
        V,
        E,
        IDX,
        REF = zeros(eltype(IDX), axes(IDX, 1)),
        FLAG = zeros(eltype(IDX), axes(IDX, 1)),
        T = zero(M),
    )
        m = deepcopy(M)
        x = deepcopy(X)
        v = deepcopy(V)
        e = deepcopy(E)
        idx = deepcopy(IDX)
        ref = deepcopy(REF)
        flag = deepcopy(FLAG)
        tc = deepcopy(T)

        new{typeof(m),typeof(x),typeof(v),typeof(idx),typeof(flag)}(
            m,
            x,
            v,
            e,
            idx,
            ref,
            flag,
            tc,
        )
    end

end

"""
    Particle1D(M::AbstractFloat, X::Real, V::AbstractArray, IDX::Integer, T::Real)

1D particle

@vars: m, x, v, e, idx, tc

"""
mutable struct Particle1D{T1,T2,T3} <: AbstractParticle1D

    m::T1
    x::T1
    v::T2
    e::T1
    idx::T3
    flag::T3
    tc::T1

    function Particle1D(M, X, V, E, IDX::Integer, FLAG = zero(IDX), T = zero(M))
        m = deepcopy(M)
        x = deepcopy(X)
        v = deepcopy(V)
        e = deepcopy(E)
        idx = deepcopy(IDX)
        flag = deepcopy(FLAG)
        tc = deepcopy(T)

        new{typeof(m),typeof(v),typeof(idx)}(m, x, v, e, idx, flag, tc)
    end

end

"""
    Particle2D(M::AbstractFloat, X::Real, V::AbstractArray, IDX::Integer, T::Real)

2D particle

@vars: m, x, v, idx, tb

"""
mutable struct Particle2D{T1,T2,T3} <: AbstractParticle2D

    m::T1
    x::T1
    y::T1
    v::T2
    e::T1
    idx::T3
    idy::T3
    flag::T3
    tc::T1

    function Particle2D(
        M::AbstractFloat,
        X::Real,
        Y::Real,
        V::AbstractArray,
        E,
        IDX::Integer,
        IDY::Integer,
        FLAG = zero(IDX),
        T = zero(M)::Real,
    )
        m = deepcopy(M)
        x = deepcopy(X)
        y = deepcopy(Y)
        v = deepcopy(V)
        e = deepcopy(E)
        idx = deepcopy(IDX)
        idy = deepcopy(IDY)
        flag = deepcopy(FLAG)
        tc = deepcopy(T)

        new{typeof(m),typeof(v),typeof(idx)}(m, x, y, v, e, idx, idy, flag, tc)
    end

end

"""
    ControlVolume1D(X::Real, DX::Real, W::AbstractArray, PRIM::AbstractArray)

1D control volume in correspondence with particle simulation

@vars: x, dx, w, prim, sw, wg, τ, np, vrmax, remainder

"""
mutable struct ControlVolumeParticle1D{F,A,I<:Integer} <: AbstractControlVolume1D

    x::F
    dx::F

    w::A
    prim::A
    sw::A

    wf::A
    wp::A
    τ::F

    np::I
    ip::I
    vrmax::F
    remainder::F

    function ControlVolumeParticle1D(
        X::Real,
        DX::Real,
        W::AbstractArray,
        PRIM::AbstractArray,
        TAU = 0.0::Real,
        NP = 0::Integer,
        IP = 0::Integer,
        VR = 3.0 * sqrt(PRIM[end])::Real,
        RE = 0::Integer,
    )
        x = deepcopy(X)
        dx = deepcopy(DX)

        w = deepcopy(W)
        prim = deepcopy(PRIM)
        sw = zeros(eltype(w), axes(w))

        wf = deepcopy(W)
        wp = zero(W)
        τ = deepcopy(TAU)

        np = deepcopy(NP)
        ip = deepcopy(IP)
        vrmax = deepcopy(VR)
        remainder = deepcopy(RE)

        new{typeof(x),typeof(w),typeof(np)}(
            x,
            dx,
            w,
            prim,
            sw,
            wf,
            wp,
            τ,
            np,
            ip,
            vrmax,
            remainder,
        )
    end

end

"""
    ControlVolume2D(X::Real, DX::Real, Y::Real, DY::Real, W::AbstractArray, PRIM::AbstractArray)

2D control volume in correspondence with particle simulation

@vars: x, y, dx, dy, w, prim, sw, wg, τ, np, vrmax, remainder

"""
mutable struct ControlVolumeParticle2D{F,A,B,I} <: AbstractControlVolume2D

    x::F
    y::F
    dx::F
    dy::F

    w::A
    prim::A
    sw::B

    wf::A
    wp::A
    τ::F

    np::I
    ip::I
    vrmax::F
    remainder::F

    function ControlVolumeParticle2D(
        X::Real,
        DX::Real,
        Y::Real,
        DY::Real,
        W::AbstractArray,
        PRIM::AbstractArray,
        TAU = 0.0::Real,
        NP = 0::Integer,
        IP = 0::Integer,
        VR = 3.0 * sqrt(PRIM[end])::Real,
        RE = 0::Integer,
    )

        x = deepcopy(X)
        dx = deepcopy(DX)
        y = deepcopy(Y)
        dy = deepcopy(DY)

        w = deepcopy(W)
        prim = deepcopy(PRIM)
        sw = zeros(eltype(W), (axes(W)..., Base.OneTo(2)))

        wf = deepcopy(W)
        wp = zero(W)
        τ = deepcopy(TAU)

        np = deepcopy(NP)
        ip = deepcopy(IP)
        vrmax = deepcopy(VR)
        remainder = deepcopy(RE)

        new{typeof(x),typeof(w),typeof(sw),typeof(np)}(
            x,
            dx,
            y,
            dy,
            w,
            prim,
            sw,
            wf,
            wp,
            τ,
            np,
            ip,
            vrmax,
            remainder,
        )

    end

end


export K2KrylovParams

"""
Type to use the K2 formulation with a Krylov method, using the package
[`Krylov.jl`](https://github.com/JuliaSmoothOptimizers/Krylov.jl).
The outer constructor

    K2KrylovParams(; uplo = :L, kmethod = :minres, preconditioner = :Identity,
                   rhs_scale = true, form_mat = false,
                   atol0 = 1.0e-4, rtol0 = 1.0e-4,
                   atol_min = 1.0e-10, rtol_min = 1.0e-10,
                   ρ0 = sqrt(eps()) * 1e5, δ0 = sqrt(eps()) * 1e5,
                   ρ_min = 1e2 * sqrt(eps()), δ_min = 1e2 * sqrt(eps()),
                   memory = 20)

creates a [`RipQP.SolverParams`](@ref) that should be used to create a [`RipQP.InputConfig`](@ref).
The available methods are:
- `:minres`
- `:minres_qlp`
- `:symmlq`

The list of available preconditioners for this solver is displayed here: [`RipQP.PreconditionerDataK2`](@ref).
"""
mutable struct K2KrylovParams <: AugmentedParams
  uplo::Symbol
  kmethod::Symbol
  preconditioner::Symbol
  rhs_scale::Bool
  form_mat::Bool
  atol0::Float64
  rtol0::Float64
  atol_min::Float64
  rtol_min::Float64
  ρ0::Float64
  δ0::Float64
  ρ_min::Float64
  δ_min::Float64
  mem::Int
end

function K2KrylovParams(;
  uplo::Symbol = :L,
  kmethod::Symbol = :minres,
  preconditioner::Symbol = :Identity,
  rhs_scale::Bool = true,
  form_mat::Bool = false,
  atol0::T = 1.0e-4,
  rtol0::T = 1.0e-4,
  atol_min::T = 1.0e-10,
  rtol_min::T = 1.0e-10,
  ρ0::T = sqrt(eps()) * 1e5,
  δ0::T = sqrt(eps()) * 1e5,
  ρ_min::T = 1e2 * sqrt(eps()),
  δ_min::T = 1e2 * sqrt(eps()),
  mem::Int = 20,
) where {T <: Real}
  if uplo == :L && form_mat
    error("matrix can only be created if uplo == :U")
  end
  return K2KrylovParams(
    uplo,
    kmethod,
    preconditioner,
    rhs_scale,
    form_mat,
    atol0,
    rtol0,
    atol_min,
    rtol_min,
    ρ0,
    δ0,
    ρ_min,
    δ_min,
    mem,
  )
end

mutable struct MatrixTools{T}
  diag_Q::SparseVector{T, Int} # Q diag
  diagind_K::Vector{Int}
end

mutable struct PreallocatedDataK2Krylov{
  T <: Real,
  S,
  M <: Union{LinearOperator{T}, AbstractMatrix{T}},
  MT <: Union{MatrixTools{T}, Int},
  Pr <: PreconditionerDataK2,
  Ksol <: KrylovSolver,
} <: PreallocatedDataAugmentedKrylov{T, S}
  pdat::Pr
  D::S                                  # temporary top-left diagonal
  rhs::S
  rhs_scale::Bool
  regu::Regularization{T}
  δv::Vector{T}
  K::M # augmented matrix
  mt::MT
  KS::Ksol
  atol::T
  rtol::T
  atol_min::T
  rtol_min::T
end

function opK2prod!(
  res::AbstractVector{T},
  nvar::Int,
  Q::Union{AbstractMatrix{T}, AbstractLinearOperator{T}},
  D::AbstractVector{T},
  A::Union{AbstractMatrix{T}, AbstractLinearOperator{T}},
  δv::AbstractVector{T},
  v::AbstractVector{T},
  α::T,
  β::T,
  uplo::Symbol,
) where {T}
  @views mul!(res[1:nvar], Q, v[1:nvar], -α, β)
  res[1:nvar] .+= α .* D .* v[1:nvar]
  if uplo == :U
    @views mul!(res[1:nvar], A, v[(nvar + 1):end], α, one(T))
    @views mul!(res[(nvar + 1):end], A', v[1:nvar], α, β)
  else
    @views mul!(res[1:nvar], A', v[(nvar + 1):end], α, one(T))
    @views mul!(res[(nvar + 1):end], A, v[1:nvar], α, β)
  end
  res[(nvar + 1):end] .+= @views (α * δv[1]) .* v[(nvar + 1):end]
end

function PreallocatedData(
  sp::K2KrylovParams,
  fd::QM_FloatData{T},
  id::QM_IntData,
  itd::IterData{T},
  pt::Point{T},
  iconf::InputConfig{Tconf},
) where {T <: Real, Tconf <: Real}

  # init Regularization values
  D = similar(fd.c, id.nvar)
  if iconf.mode == :mono
    regu = Regularization(T(sp.ρ0), T(sp.δ0), T(sp.ρ_min), T(sp.δ_min), :classic)
    D .= -T(1.0e0) / 2
  else
    regu =
      Regularization(T(sp.ρ0), T(sp.δ0), T(sqrt(eps(T)) * 1e0), T(sqrt(eps(T)) * 1e0), :classic)
    D .= -T(1.0e-2)
  end
  δv = [regu.δ] # put it in a Vector so that we can modify it without modifying opK2prod!
  if sp.form_mat
    diag_Q = get_diag_Q(fd.Q.data.colptr, fd.Q.data.rowval, fd.Q.data.nzval, id.nvar)
    K = Symmetric(create_K2(id, D, fd.Q.data, fd.A, diag_Q, regu), fd.uplo)
    diagind_K = get_diag_sparseCSC(K.data.colptr, id.ncon + id.nvar)
    mt = MatrixTools(diag_Q, diagind_K)
  else
    K = LinearOperator(
      T,
      id.nvar + id.ncon,
      id.nvar + id.ncon,
      true,
      true,
      (res, v, α, β) -> opK2prod!(res, id.nvar, fd.Q, D, fd.A, δv, v, α, β, fd.uplo),
    )
    mt = 0
  end

  rhs = similar(fd.c, id.nvar + id.ncon)
  KS = init_Ksolver(K, rhs, sp)
  pdat = eval(sp.preconditioner)(id, fd, regu, D, K)

  return PreallocatedDataK2Krylov(
    pdat,
    D,
    rhs,
    sp.rhs_scale,
    regu,
    δv,
    K, #K
    mt,
    KS,
    T(sp.atol0),
    T(sp.rtol0),
    T(sp.atol_min),
    T(sp.rtol_min),
  )
end

function solver!(
  dd::AbstractVector{T},
  pad::PreallocatedDataK2Krylov{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
  step::Symbol,
) where {T <: Real}
  pad.rhs .= dd
  if pad.rhs_scale
    rhsNorm = kscale!(pad.rhs)
  end
  ksolve!(pad.KS, pad.K, pad.rhs, pad.pdat.P, verbose = 0, atol = pad.atol, rtol = pad.rtol)
  update_kresiduals_history!(res, pad.K, pad.KS.x, pad.rhs)
  if pad.rhs_scale
    kunscale!(pad.KS.x, rhsNorm)
  end

  dd .= pad.KS.x

  return 0
end

function update_pad!(
  pad::PreallocatedDataK2Krylov{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
) where {T <: Real}
  if cnts.k != 0
    update_regu!(pad.regu)
  end

  if pad.atol > pad.atol_min
    pad.atol /= 10
  end
  if pad.rtol > pad.rtol_min
    pad.rtol /= 10
  end

  pad.D .= -pad.regu.ρ
  pad.D[id.ilow] .-= pt.s_l ./ itd.x_m_lvar
  pad.D[id.iupp] .-= pt.s_u ./ itd.uvar_m_x
  pad.δv[1] = pad.regu.δ
  if typeof(pad.K) <: Symmetric{T, SparseMatrixCSC{T, Int}}
    pad.D[pad.mt.diag_Q.nzind] .-= pad.mt.diag_Q.nzval
    pad.K.data.nzval[view(pad.mt.diagind_K, 1:(id.nvar))] = pad.D
    pad.K.data.nzval[view(pad.mt.diagind_K, (id.nvar + 1):(id.ncon + id.nvar))] .= pad.regu.δ
  end

  update_preconditioner!(pad.pdat, pad, itd, pt, id, fd, cnts)

  return 0
end


using Revise
using Printf
using Random
using POMDPs
using POMDPPolicies
using POMDPSimulators
using BeliefUpdaters
using AutomotiveDrivingModels
using AutomotivePOMDPs
using AutoViz
using RLInterface
using DeepQLearning

rng = MersenneTwister(1)

const K = 4

pomdp = OCPOMDP(ΔT = 0.5, p_birth = 0.3, max_peds = 1)

env = KMarkovEnvironment(pomdp, k=K)

up = KMarkovUpdater(K)

policy = RandomPolicy(pomdp, rng=rng)

hr = HistoryRecorder(rng=rng, max_steps = 100)
s0 = rand(rng, initialstate(pomdp))
initial_observation = rand(rng, initialobs(pomdp, s0))
initial_obs_vec = fill(initial_observation, 4)
hist = POMDPs.simulate(hr, pomdp, policy, up, initial_obs_vec, s0)

# Visualize
using Reel
frames = Frames(MIME("image/png"), fps=4)
AutoViz.render(s0, pomdp.env, cam = StaticCamera(VecE2(25.0,0.0), 15.0))
for step in eachstep(hist, "s,a,r,sp")
    s, a, r, sp = step
    push!(frames, AutoViz.render(sp, pomdp.env, cam=StaticCamera(VecE2(25.0,0.0), 15.0)))
end
write("out.gif", frames)

pomdp.env.obstacles


using CuArrays
using Distributions: pdf
using Random

RLBase.get_prob(p::AbstractPolicy, obs, ::RLBase.AbstractActionStyle, a) =
    pdf(get_prob(p, obs), a)

Random.rand(s::MultiContinuousSpace{<:CuArray}) = rand(CuArrays.CURAND.generator(), s)