@testset "populate_lp_model!" begin
    form = _build_dmip_formulation()
    state = Cerberus.CurrentState()
    let node = Cerberus.Node()
        node_result = Cerberus.NodeResult(node)
        @inferred Cerberus.populate_lp_model!(
            state,
            form,
            node,
            node_result,
            CONFIG,
        )
        model = state.gurobi_model

        @test MOI.get(model, MOI.NumberOfVariables()) == 3
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1
        c1 = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_ET}())[1]
        f1 = MOI.get(model, MOI.ConstraintFunction(), c1)
        _is_equal(f1, 1.0 * _SV(_VI(1)) + 2.1 * _SV(_VI(2)) + 3.0 * _SV(_VI(3)))
        s1 = MOI.get(model, MOI.ConstraintSet(), c1)
        @test s1 == _ET(3.0)

        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 1
        c2 = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_LT}())[1]
        f2 = MOI.get(model, MOI.ConstraintFunction(), c2)
        _is_equal(f2, -3.5 * _SV(_VI(1)) + 1.2 * _SV(_VI(2)))
        s2 = MOI.get(model, MOI.ConstraintSet(), c2)
        @test s2 == _LT(4.0)
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 0

        @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(1), state),
        ) == (0.5, 1.0)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(2), state),
        ) == (-1.3, 2.3)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(3), state),
        ) == (0.0, 1.0)

        @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MIN_SENSE
        obj = MOI.get(model, MOI.ObjectiveFunction{_SAF}())
        _is_equal(obj, 1.0 * _SV(_VI(1)) - 1.0 * _SV(_VI(2)))
    end

    # Safe to do this now that we've populated the model once (undefined reference otherwise).
    model = state.gurobi_model
    # TODO: Test that model does NOT get created if the config/state says not to.
    state.rebuild_model = false
    let node = Cerberus.Node(
            [Cerberus.BoundUpdate(_CVI(1), _LT(0.0))],
            [Cerberus.BoundUpdate(_CVI(3), _GT(1.0))],
            2,
        )
        node_result = Cerberus.NodeResult(node)
        @inferred Cerberus.populate_lp_model!(
            state,
            form,
            node,
            node_result,
            CONFIG,
        )
        @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(1), state),
        ) == (0.5, 0.0)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(2), state),
        ) == (-1.3, 2.3)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(3), state),
        ) == (1.0, 1.0)
    end

    # Lie to Cerberus so that it doesn't reset the bounds
    state.on_a_dive = true
    let node = Cerberus.Node()
        node_result = Cerberus.NodeResult(node)
        @inferred Cerberus.populate_lp_model!(
            state,
            form,
            node,
            node_result,
            CONFIG,
        )
        @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(1), state),
        ) == (0.5, 0.0)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(2), state),
        ) == (-1.3, 2.3)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(3), state),
        ) == (1.0, 1.0)
    end

    state.on_a_dive = false
    let node = Cerberus.Node()
        node_result = Cerberus.NodeResult(node)
        @inferred Cerberus.populate_lp_model!(
            state,
            form,
            node,
            node_result,
            CONFIG,
        )
        @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(1), state),
        ) == (0.5, 1.0)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(2), state),
        ) == (-1.3, 2.3)
        @test MOI.Utilities.get_bounds(
            model,
            Float64,
            Cerberus.instantiate(_CVI(3), state),
        ) == (0.0, 1.0)
    end
end

@testset "update_node_bounds!" begin
    form = _build_dmip_formulation()
    state = _CurrentState()
    node = Cerberus.Node()
    node_result = Cerberus.NodeResult(node)
    @inferred Cerberus.populate_lp_model!(
        state,
        form,
        node,
        node_result,
        CONFIG,
    )
    model = state.gurobi_model
    @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 1.0)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 1.0)

    node = Cerberus.Node(
        [Cerberus.BoundUpdate(_CVI(1), _LT(0.0))],
        [Cerberus.BoundUpdate(_CVI(3), _GT(1.0))],
        Cerberus.AffineConstraint{_LT}[],
        Cerberus.AffineConstraint{_GT}[],
        2,
    )
    @inferred Cerberus.apply_branchings!(state, node)
    @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 0.0)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (1.0, 1.0)

    let f = _CSAF([1.2, 3.4], [_CVI(1), _CVI(2)], 5.6), s = _LT(7.8)
        bd = Cerberus.AffineConstraint(f, s)
        Cerberus.apply_branching!(node, bd)
        @inferred Cerberus.apply_branchings!(state, node)
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 2
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 0
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1
        cis = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_LT}())
        # NOTE: cis[1] should be the constraint from the "base" formulation.
        ci = cis[2]
        f_rt = MOI.get(model, MOI.ConstraintFunction(), ci)
        s_rt = MOI.get(model, MOI.ConstraintSet(), ci)
        # Constraint added with MOIU.normalize_and_add_constraint, which shifts
        # constant over to set.
        @test _is_equal(f_rt, Cerberus.instantiate(f, state) - 5.6)
        @test s_rt == _LT(7.8 - 5.6)
    end

    let f = _CSAF([2.4, 6.4], [_CVI(3), _CVI(1)], 0.0), s = _GT(3.5)
        bd = Cerberus.AffineConstraint(f, s)
        Cerberus.apply_branching!(node, bd)
        @inferred Cerberus.apply_branchings!(state, node)
        # NOTE: We are testing here that the _LT general branching constraint
        # only gets added to the model once.
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 2
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 1
        @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1
        cis = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_GT}())
        ci = cis[1]
        f_rt = MOI.get(model, MOI.ConstraintFunction(), ci)
        s_rt = MOI.get(model, MOI.ConstraintSet(), ci)
        @test _is_equal(f_rt, Cerberus.instantiate(f, state))
        @test s_rt == s
    end
end

@testset "remove_all_branchings!" begin
    form = _build_dmip_formulation()
    state = Cerberus.CurrentState()
    node = Cerberus.Node()
    Cerberus.create_base_model!(state, form, node, CONFIG)
    f_lt = _CSAF([1.2, 3.4], [_CVI(1), _CVI(2)], 0.0)
    s_lt = _LT(7.8)
    f_gt = _CSAF([2.4, 6.4], [_CVI(3), _CVI(1)], 0.0)
    s_gt = _GT(3.5)
    node_1 = Cerberus.Node(
        [Cerberus.BoundUpdate(_CVI(1), _LT(0.0))],
        [Cerberus.BoundUpdate(_CVI(3), _GT(1.0))],
        [Cerberus.AffineConstraint{_LT}(f_lt, s_lt)],
        [Cerberus.AffineConstraint{_GT}(f_gt, s_gt)],
        4,
    )
    Cerberus.apply_branchings!(state, node_1)

    model = state.gurobi_model
    @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 0.0)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (1.0, 1.0)
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 2
    let lt_cis = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_LT}())
        lt_ci = lt_cis[2]
        f_lt_rt = MOI.get(model, MOI.ConstraintFunction(), lt_ci)
        s_lt_rt = MOI.get(model, MOI.ConstraintSet(), lt_ci)
        @test _is_equal(f_lt_rt, Cerberus.instantiate(f_lt, state))
        @test s_lt_rt == s_lt
    end
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 1
    let gt_cis = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_GT}())
        gt_ci = gt_cis[1]
        f_gt_rt = MOI.get(model, MOI.ConstraintFunction(), gt_ci)
        s_gt_rt = MOI.get(model, MOI.ConstraintSet(), gt_ci)
        @test _is_equal(f_gt_rt, Cerberus.instantiate(f_gt, state))
        @test s_gt_rt == s_gt
    end
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1

    @inferred Cerberus.remove_all_branchings!(state, form)
    @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 1.0)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 1.0)
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 1
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 0
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1

    node_2 = Cerberus.Node(
        [Cerberus.BoundUpdate(_CVI(1), _LT(0.0))],
        Cerberus.BoundUpdate{_GT}[],
        Cerberus.AffineConstraint{_LT}[],
        [Cerberus.AffineConstraint{_GT}(f_gt, s_gt)],
        2,
    )
    Cerberus.apply_branchings!(state, node_2)
    @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 0.0)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 1.0)
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 1
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 1
    let gt_cis = MOI.get(model, MOI.ListOfConstraintIndices{_SAF,_GT}())
        gt_ci = gt_cis[1]
        f_gt_rt = MOI.get(model, MOI.ConstraintFunction(), gt_ci)
        s_gt_rt = MOI.get(model, MOI.ConstraintSet(), gt_ci)
        @test _is_equal(f_gt_rt, Cerberus.instantiate(f_gt, state))
        @test s_gt_rt == s_gt
    end
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1

    @inferred Cerberus.remove_all_branchings!(state, form)
    @test MOI.get(model, MOI.NumberOfConstraints{_SV,_IN}()) == 3
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 1.0)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOI.Utilities.get_bounds(
        model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 1.0)
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_LT}()) == 1
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_GT}()) == 0
    @test MOI.get(model, MOI.NumberOfConstraints{_SAF,_ET}()) == 1
end

@testset "apply_branchings!" begin
    form = _build_dmip_formulation()
    state = Cerberus.CurrentState()
    cs = state.constraint_state

    node = Cerberus.Node()
    Cerberus.create_base_model!(state, form, node, CONFIG)
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    @inferred Cerberus.apply_branchings!(state, node)
    @test cs.branch_state.num_lt_branches == 0
    @test cs.branch_state.num_gt_branches == 0
    @test isempty(cs.branch_state.lt_general_constrs)
    @test isempty(cs.branch_state.gt_general_constrs)

    node_result = Cerberus.NodeResult(node)
    Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (0.5, 1.0)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 1.0)

    bd_1 = Cerberus.BoundUpdate(_CVI(1), _GT(1.0))
    Cerberus.apply_branching!(node, bd_1)
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test Cerberus._unattached_bounds(cs, node, _GT) == [bd_1]
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    Cerberus.apply_branchings!(state, node)
    @test cs.branch_state.num_lt_branches == 0
    @test cs.branch_state.num_gt_branches == 1
    @test isempty(cs.branch_state.lt_general_constrs)
    @test isempty(cs.branch_state.gt_general_constrs)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (1.0, 1.0)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 1.0)
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    bd_2 = Cerberus.BoundUpdate(_CVI(3), _LT(0.0))
    Cerberus.apply_branching!(node, bd_2)
    @test Cerberus._unattached_bounds(cs, node, _LT) == [bd_2]
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    Cerberus.apply_branchings!(state, node)
    @test cs.branch_state.num_lt_branches == 1
    @test cs.branch_state.num_gt_branches == 1
    @test isempty(cs.branch_state.lt_general_constrs)
    @test isempty(cs.branch_state.gt_general_constrs)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (1.0, 1.0)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 0.0)
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    # Now we jump across the tree to a different node at the same depth.
    # Notably, we do not update the constraint_state. In particular, we still
    # believe that there has only been two variable branches. Therefore, the
    # model will still have the same bounds as above, and will NOT correspond
    # to the bounds in node below.
    node = Cerberus.Node()
    Cerberus.apply_branching!(node, Cerberus.BoundUpdate(_CVI(3), _GT(1.0)))
    Cerberus.apply_branching!(node, Cerberus.BoundUpdate(_CVI(1), _LT(0.0)))
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    Cerberus.apply_branchings!(state, node)
    @test cs.branch_state.num_lt_branches == 1
    @test cs.branch_state.num_gt_branches == 1
    @test isempty(cs.branch_state.lt_general_constrs)
    @test isempty(cs.branch_state.gt_general_constrs)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(1), state),
    ) == (1.0, 1.0)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(2), state),
    ) == (-1.3, 2.3)
    @test MOIU.get_bounds(
        state.gurobi_model,
        Float64,
        Cerberus.instantiate(_CVI(3), state),
    ) == (0.0, 0.0)
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _LT))
    @test isempty(Cerberus._unattached_constraints(cs, node, _GT))

    lt_gen_1 = Cerberus.AffineConstraint(2.3 * _SV(_VI(1)), _LT(2.3))
    lt_gen_2 = Cerberus.AffineConstraint(2.5 * _SV(_VI(3)), _LT(2.5))
    gt_gen = Cerberus.AffineConstraint(2.4 * _SV(_VI(2)), _GT(2.4))
    Cerberus.apply_branching!(node, lt_gen_1)
    Cerberus.apply_branching!(node, lt_gen_2)
    Cerberus.apply_branching!(node, gt_gen)
    @test isempty(Cerberus._unattached_bounds(cs, node, _LT))
    @test isempty(Cerberus._unattached_bounds(cs, node, _GT))
    @test Cerberus._unattached_constraints(cs, node, _LT) ==
          [lt_gen_1, lt_gen_2]
    @test Cerberus._unattached_constraints(cs, node, _GT) == [gt_gen]

    Cerberus.apply_branchings!(state, node)
    @test cs.branch_state.num_lt_branches == 1
    @test cs.branch_state.num_gt_branches == 1
    @test length(cs.branch_state.lt_general_constrs) == 2
    @test length(cs.branch_state.gt_general_constrs) == 1
    @test MOI.get(state.gurobi_model, MOI.NumberOfConstraints{_SAF,_LT}()) == 3
    @test MOI.get(state.gurobi_model, MOI.NumberOfConstraints{_SAF,_GT}()) == 1
end

@testset "formulate_disjunctions!" begin
    # min  y + 1.2
    # s.t. (-2 <= x <= -1 & y = -x - 1) or
    #          (-1 <= x <= +1 & y = 0) or
    #          (+1 <= x <= +2 & y = x - 1)
    #      x + y >= -0.5
    #      x in [-1.5, 3.0]
    #      y in [0.0, 0.5]
    form = _build_formulation_with_single_disjunction(
        DisjunctiveConstraints.NaiveBigM(
            DisjunctiveConstraints.IntervalArithmetic(),
        ),
    )
    let node = Cerberus.Node()
        # min  y + 1.2
        # s.t. z_1 + z_2 + z_3 = 1
        #      x + y >= -0.5
        #      x >= -1.5 - 0.5 z_1
        #      x >= -1.5 + 0.5 z_2
        #      x >= -1.5 + 2.5 z_3
        #      y >= -x - 1.5 + 0.5 z_1
        #      y >= 0
        # .    y >= x - 3 + 2.0 z_3
        #      x <=  3.0 - 4.0 z_1
        #      x <=  3.0 - 2.0 z_2
        #      x <=  3.0 - 1.0 z_3
        #      y <= -x + 3.5 - 4.5 z_1
        #      y <= 0.5 - 0.5 z_2
        #      y <= x + 2.0 - 3.0 z_3
        #      x in [-1.5, 3.0]
        #      y in [0.0, 0.5]
        #      z in [0,1]^3
        state = Cerberus.CurrentState()
        node_result = Cerberus.NodeResult(node)
        Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
        Cerberus.apply_branchings!(state, node)
        x = [_SV(Cerberus.instantiate(_CVI(i), state)) for i in 1:5]
        @assert [v.variable.value for v in x] == collect(1:5)

        expected_bounds =
            [(-1.5, 3.0), (0.0, 0.5), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)]
        expected_et_acs = [(1.0 * x[3] + 1.0 * x[4] + 1.0 * x[5], _ET(1.0))]
        expected_gt_acs = [
            (1.0 * x[1] + 1.0 * x[2], _GT(-0.5)),
            (1.0 * x[1] + 0.5 * x[3], _GT(-1.5)),
            (1.0 * x[1] - 0.5 * x[4], _GT(-1.5)),
            (1.0 * x[1] - 2.5 * x[5], _GT(-1.5)),
            (1.0 * x[1] + 1.0 * x[2] - 0.5 * x[3], _GT(-1.5)),
            (1.0 * x[2], _GT(0.0)),
            (-1.0 * x[1] + 1.0 * x[2] - 2.0 * x[5], _GT(-3.0)),
        ]
        expected_lt_acs = [
            (1.0 * x[1] + 4.0 * x[3], _LT(+3.0)),
            (1.0 * x[1] + 2.0 * x[4], _LT(+3.0)),
            (1.0 * x[1] + 1.0 * x[5], _LT(+3.0)),
            (1.0 * x[1] + 1.0 * x[2] + 4.5 * x[3], _LT(+3.5)),
            (1.0 * x[2] + 0.5 * x[4], _LT(0.5)),
            (-1.0 * x[1] + 1.0 * x[2] + 3.0 * x[5], _LT(+2.0)),
        ]

        _test_roundtrip_model(
            state.gurobi_model,
            expected_bounds,
            expected_lt_acs,
            expected_gt_acs,
            expected_et_acs,
        )
    end

    let node = Cerberus.Node(
            [Cerberus.BoundUpdate(_CVI(3), _LT(0.0))],
            Cerberus.BoundUpdate{_GT}[],
            1,
        )
        # min  y + 1.2
        # s.t. z_1 + z_2 + z_3 = 1
        #      x + y >= -0.5
        #      x >= -1.5 + 0.5 z_2
        #      x >= -1.5 + 2.5 z_3
        #      y >= 0
        # .    y >= x - 3 + 2.0 z_3
        #      x <=  3.0 - 2.0 z_2
        #      x <=  3.0 - 1.0 z_3
        #      y <= 0.5 - 0.5 z_2
        #      y <= x + 2.0 - 3.0 z_3
        #      x in [-1.5, 3.0]
        #      y in [0.0, 0.5]
        #      z in [0,1]^3
        #      z_1 = 0
        state = Cerberus.CurrentState()
        node_result = Cerberus.NodeResult(node)
        Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
        Cerberus.apply_branchings!(state, node)
        x = [_SV(Cerberus.instantiate(_CVI(i), state)) for i in 1:5]
        @assert [v.variable.value for v in x] == collect(1:5)

        expected_bounds =
            [(-1.5, 3.0), (0.0, 0.5), (0.0, 0.0), (0.0, 1.0), (0.0, 1.0)]
        expected_et_acs = [(1.0 * x[3] + 1.0 * x[4] + 1.0 * x[5], _ET(1.0))]
        expected_gt_acs = [
            (1.0 * x[1] + 1.0 * x[2], _GT(-0.5)),
            (1.0 * x[1] - 0.5 * x[4], _GT(-1.5)),
            (1.0 * x[1] - 2.5 * x[5], _GT(-1.5)),
            (1.0 * x[2], _GT(0.0)),
            (-1.0 * x[1] + 1.0 * x[2] - 2.0 * x[5], _GT(-3.0)),
        ]
        expected_lt_acs = [
            (1.0 * x[1] + 2.0 * x[4], _LT(+3.0)),
            (1.0 * x[1] + 1.0 * x[5], _LT(+3.0)),
            (1.0 * x[2] + 0.5 * x[4], _LT(0.5)),
            (-1.0 * x[1] + 1.0 * x[2] + 3.0 * x[5], _LT(+2.0)),
        ]
        _test_roundtrip_model(
            state.gurobi_model,
            expected_bounds,
            expected_lt_acs,
            expected_gt_acs,
            expected_et_acs,
        )
    end
end

@testset "MOI.optimize!" begin
    form = _build_dmip_formulation()
    state = _CurrentState()
    node = Cerberus.Node()
    node_result = Cerberus.NodeResult(node)
    Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
    model = state.gurobi_model
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.VariablePrimal(), _VI(1)) ≈ 0.5
    @test MOI.get(model, MOI.VariablePrimal(), _VI(2)) ≈ 2.5 / 2.1
    @test MOI.get(model, MOI.VariablePrimal(), _VI(3)) ≈ 0.0
end

@testset "_fill_solution!" begin
    form = _build_dmip_formulation()
    state = _CurrentState()
    node = Cerberus.Node()
    node_result = Cerberus.NodeResult(node)
    Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
    model = state.gurobi_model
    MOI.optimize!(model)
    @assert MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @inferred Cerberus._update_lp_solution!(state, form)
    x = state.current_solution
    @test length(x) == 3
    @test x ≈ [1 / 2, 2.5 / 2.1, 0.0]
end

@testset "get_basis" begin
    form = _build_dmip_formulation()
    state = _CurrentState()
    node = Cerberus.Node()
    node_result = Cerberus.NodeResult(node)
    Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
    model = state.gurobi_model
    MOI.optimize!(model)
    @assert MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    true_basis = @inferred Cerberus.get_basis(state)
    expected_basis = Cerberus.Basis(
        [MOI.NONBASIC_AT_LOWER, MOI.BASIC, MOI.NONBASIC_AT_LOWER],
        [MOI.BASIC],
        MOI.BasisStatusCode[],
        [MOI.NONBASIC],
        MOI.BasisStatusCode[],
        MOI.BasisStatusCode[],
    )
    @test true_basis.base_var_constrs == expected_basis.base_var_constrs
    @test true_basis.base_lt_constrs == expected_basis.base_lt_constrs
    @test true_basis.base_gt_constrs == expected_basis.base_gt_constrs
    @test true_basis.base_et_constrs == expected_basis.base_et_constrs
    @test true_basis.branch_lt_constrs == expected_basis.branch_lt_constrs
    @test true_basis.branch_gt_constrs == expected_basis.branch_gt_constrs
end

function _set_basis_model(basis::Cerberus.Basis)
    form = _build_dmip_formulation()
    state = _CurrentState()
    node = Cerberus.Node(
        Cerberus.BoundUpdate{_LT}[],
        Cerberus.BoundUpdate{_GT}[],
        Cerberus.AffineConstraint{_LT}[],
        Cerberus.AffineConstraint{_GT}[],
        0,
        -Inf,
    )
    node_result = Cerberus.NodeResult(node)
    Cerberus.populate_lp_model!(state, form, node, node_result, CONFIG)
    state.warm_starts[node] = basis
    Cerberus.set_basis_if_available!(state, node)
    return state.gurobi_model
end

@testset "set_basis_if_available!" begin
    # First, seed a suboptimal basis. This will disable presolve. It is only one pivot away from the optimal basis.
    let
        subopt_basis = Cerberus.Basis(
            [MOI.BASIC, MOI.NONBASIC_AT_LOWER, MOI.NONBASIC_AT_LOWER],
            [MOI.BASIC],
            MOI.BasisStatusCode[],
            [MOI.NONBASIC],
            MOI.BasisStatusCode[],
            MOI.BasisStatusCode[],
        )
        model = _set_basis_model(subopt_basis)
        MOI.optimize!(model)
        @test MOI.get(model, MOI.SimplexIterations()) == 1
    end

    # Now, seed the optimal basis. This will solve the problem without any simplex iterations.
    let
        opt_basis = Cerberus.Basis(
            [MOI.NONBASIC_AT_LOWER, MOI.BASIC, MOI.NONBASIC_AT_LOWER],
            [MOI.BASIC],
            MOI.BasisStatusCode[],
            [MOI.NONBASIC],
            MOI.BasisStatusCode[],
            MOI.BasisStatusCode[],
        )
        model = _set_basis_model(opt_basis)
        MOI.optimize!(model)
        @test MOI.get(model, MOI.SimplexIterations()) == 0
    end
end


module SSPS

using Gen
using GLMNet
using ArgParse
using JSON

export julia_main

include("dbn_preprocess.jl")
include("dbn_models.jl")
include("dbn_proposals.jl")
include("mcmc_inference.jl")
include("state_updates.jl")

@load_generated_functions()

"""
Given the results of a previous run, recover the
state of the markov chain and store it in the given
choicemap object.
"""
function recover_state!(cmap, samples)

    # Recover the lambda values
    for (j, svec) in enumerate(samples["lambda"])
        cmap[:lambda_vec => j => :lambda] = svec[end][2]
    end

    # Recover the parent sets
    for (j, pdict) in enumerate(samples["parent_sets"])
        pset = Int64[]
        for (k,v) in pdict
            if v[end][2] == 1
                append!(pset, parse(Int64, k))
            end
        end
        sort!(pset)
        cmap[:parent_sets => j => :parents] = pset
    end

end

"""
parse the script's arguments using ArgParse.
"""
function parse_script_arguments()

    s = ArgParseSettings()
    @add_arg_table! s begin
        "timeseries_filename"
            help = "name of timeseries data file"
            required = true
        "ref_graph_filename"
            help = "name of reference graph data file"
            required = true
        "output_path"
            help = "name of output JSON file"
            required = true
        "timeout"
            help = "execution timeout (in seconds). When reached, terminate and output results as they are."
            arg_type = Float64
        "--thinning"
            help = "number of proposals to make for each sample taken"
            arg_type = Int64
            default=1
        "--large-indeg"
            help = "an approximate upper bound on the number of parents for any node"
            arg_type = Float64
            default = 20.0
        "--lambda-max"
            help = "hyperparameter for lambda prior"
            default = 15.0
        "--regression-deg"
            help = "regression polynomial degree in Gaussian DBN (default: full)"
            arg_type = Int64
            default = -1
        "--lambda-prop-std"
            help = "standard deviation of lambda's proposal distribution"
            arg_type = Float64
            default = 0.25
        "--n-steps"
            help = "Terminate a markov chain after it runs this many steps."
            arg_type = Int64
            default = -1
        "--proposal"
            help = "(DEPRECATED. Future versions of SSPS will only use 'sparse'.) The graph proposal distribution to use in MCMC: 'sparse' or 'uniform'."
            arg_type = String
            default = "sparse"
            range_tester = x -> x in ("sparse", "uniform")
        "--resume-sampling"
            help = "Continue sampling the markov chain stored in the given JSON file"
            required = false
            arg_type = String
            default = ""
        "--multithread"
            help = "Parallelize each iteration with multiple threads. The number of threads is specified by the `JULIA_NUM_THREADS` environment variable. This flag overrides the `--proposal` argument, setting it to 'sparse'."
            action = :store_true
    end

    args = parse_args(s)
    arg_vec = transform_arguments(args)

    return arg_vec
end

"""
Load arguments into a vector
"""
function transform_arguments(parsed_arg_dict)
    arg_vec = []
    push!(arg_vec, parsed_arg_dict["timeseries_filename"])
    push!(arg_vec, parsed_arg_dict["ref_graph_filename"])
    push!(arg_vec, parsed_arg_dict["output_path"])
    push!(arg_vec, parsed_arg_dict["timeout"])
    push!(arg_vec, parsed_arg_dict["n-steps"])
    push!(arg_vec, parsed_arg_dict["thinning"])
    push!(arg_vec, parsed_arg_dict["large-indeg"])
    push!(arg_vec, parsed_arg_dict["lambda-max"])
    push!(arg_vec, parsed_arg_dict["regression-deg"])
    push!(arg_vec, parsed_arg_dict["lambda-prop-std"])
    push!(arg_vec, parsed_arg_dict["proposal"])
    push!(arg_vec, parsed_arg_dict["resume-sampling"])
    push!(arg_vec, parsed_arg_dict["multithread"])

    return arg_vec
end

"""
A wrapper around the MCMC inference procedure
"""
function perform_inference(timeseries_filename::String,
                           ref_graph_filename::String,
                           output_path::String,
                           timeout::Float64,
                           n_steps::Int64,
                           thinning::Int64,
                           large_indeg::Float64,
                           lambda_max::Float64,
                           regression_deg::Int64,
                           lambda_prop_std::Float64,
                           proposal::String,
			   resume_sampling::String,
			   multithread::Bool)

    clear_caches()

    ts_vec, ref_ps = load_formatted_data(timeseries_filename,
                                         ref_graph_filename)

    # Some data preprocessing
    Xminus, Xplus  = combine_X(ts_vec)
    Xplus = standardize_X(Xplus)[1]
    Xminus, Xplus = vectorize_X(Xminus, Xplus)
    V = length(Xplus)

    println("Invoking SSPS on input files:\n\t",
            timeseries_filename, "\n\t", ref_graph_filename)

    # prepare parameters for proposal distributions
    ref_parent_counts = [max(sum(values(ps)), 2.0) for ps in ref_ps]
    update_loop_args = 1.0 ./ log2.(V ./ ref_parent_counts)
    push!(update_loop_args, lambda_prop_std)

    # Choose a lower bound for the lambda variable
    lambda_min = log(max(V/large_indeg - 1.0, exp(0.5)))

    # Choose the right update loop
    update_loop_fn = vertex_lambda_update_loop
    if proposal == "uniform"
        update_loop_fn = uniform_update_loop
    end
    if multithread
        update_loop_fn = multithread_update_loop
    end

    # Load observations into a choice map
    observations = Gen.choicemap()
    for (i, Xp) in enumerate(Xplus)
        observations[:Xplus => i => :Xp] = Xp
    end

    # If necessary, reload state from a previous session
    if resume_sampling != ""
        println("Resuming Markov Chain stored at ", resume_sampling)

        f = open(resume_sampling, "r")
        str = read(f, String)
        close(f)
        samples = JSON.parse(str)

        recover_state!(observations, samples)
    end

    # Prepare the Gen model's arguments
    model_args = (ref_ps,
                  Xminus,
                  lambda_min,
                  lambda_max,
                  regression_deg)

    t_start = time()

    update_results_args = [V, true]

    chain_id = string("pid_",getpid())

    results = mcmc_inference(vertex_lambda_dbn_model,
                             model_args, observations,
                             update_loop_fn,
                             update_loop_args,
                             update_results_storediff,
                             update_results_args;
                             timeout=timeout,
                             n_steps=n_steps,
                             thinning=thinning,
			     chain_id=chain_id)

    # Some last updates to the `results` object
    results["t_elapsed"] = time() - t_start
    delete!(results, "prev_parents")

    println("Saving results to JSON file:")

    f = open(output_path, "w")
    JSON.print(f, results)
    close(f)
    println("\t", output_path)

end

function julia_main()::Cint

    arg_vec = parse_script_arguments()
    perform_inference(arg_vec...)
    return 0
end

# END MODULE
end


# Module for The King is Dead bot
# This houses all of the types

using DataFrames
using Random
using YAML
Random.seed!(100)

include("src/Types.jl")
include("src/Funcs.jl")
include("src/Display.jl")

inputs = YAML.load_file("inputs.yaml", dicttype=Dict{String, Union{Int64, Any}})
# inputs = YAML.load_file("Simple inputs.yaml", dicttype=Dict{String, Union{Int64, Any}})

court_size = 2
region_size = 4
faction_names = inputs["faction_names"]
region_names = inputs["region_names"]
player_names = inputs["player_names"]
card_names = inputs["card_names"]

# Set up the regions
regions = Array{Region}(undef, 0)
for (r, r_adj) in region_names
    followers = Dict(k=>0 for k in keys(faction_names))
    size = 0
    push!(regions, Region(r, followers, size, r_adj))
end

# Initialise the players
players = Array{Player}(undef, 0)
for name in player_names
    push!(players, Player(name))
end

# Initialise cards
cards = Array{Card}(undef, 0)
for (name, card) in card_names
    push!(cards, Card(name, card))
end

# Initialise the board
board = Board(regions, players)
# initialise_player_courts!(players, board)
initialise_board!(board)

println(board)
# println(board.supply)
# resolve_region!(board, "Essex")
# println(board.supply)

# Work out ALL possible board states after using a card
card = cards[2]
num_reg_to_check = 3
possible_boards = Board[]
num_reg_to_check = min(num_reg_to_check, length(board.unresolved_regions))
if card.type == "followers"
    # First do Assemble
    if (card.adjacent == false) & isnothing(card.swap)
        allocations = Dict[]
        let iteration_counter = 1
            for r in board.unresolved_regions[1:num_reg_to_check]
                for r2 in board.unresolved_regions[1:num_reg_to_check]
                    for r3 in board.unresolved_regions[1:num_reg_to_check]
                        push!(possible_boards, deepcopy(board))
                        allocation = Dict("English" => r.name,
                             "Welsh" => r2.name,
                             "Scottish" => r3.name)
                        push!(allocations, allocation)
                        for p in allocation
                            place_follower!(possible_boards[iteration_counter], p)
                        end
                        iteration_counter += 1
                    end
                end
            end
        end
    elseif (card.adjacent == false) & !isnothing(card.swap)
        # Now do Manoeuvre
    elseif isnothing(card.swap)
        # Now do Faction support
        # Determine all possible regions if an adjacent constraint is present
        # and followers are being placed
        # First work out which regions are controlled by the faction specified
        controlled_regions = [r.name for r in board.regions if r.controller == card.adjacent]
        possible_regions = []
        for r in board.unresolved_regions
            for adj in r.adjacent
                if adj in controlled_regions
                    push!(possible_regions, r.name)
                end
            end
        end
        for r_name in possible_regions
            new_board = deepcopy(board)
            for (f_name, f_num) in card.place
                println(f_name)
                for n in 1:f_num
                    place_follower!(new_board, r_name, f_name)
                end
            end
            push!(possible_boards, new_board)
        end
    else
        # Lastly do Outmanoeuvre
        # Determine all possible regions if an adjacent constraint is present
        # and followers are being swapped
        possible_regions = []
    end
end

# For every single board state created by playing a card,
# work out every single follower removal combination
player = players[1]
possible_boards_after_rem = Board[]
# possible_board = possible_boards[1]
for possible_board in possible_boards
    possible_removals = Pair[]
    for r in possible_board.unresolved_regions
        for (f_name,f_num) in r.followers
            if f_num > 0
                take_pair = r.name=>f_name
                push!(possible_removals, r.name=>f_name)
                board_after_rem = deepcopy(possible_board)
                take_follower!(board_after_rem, r.name, f_name)
                # println(evaluate_resolved_board(board_after_rem))
                push!(possible_boards_after_rem, board_after_rem)
            end
        end
    end
end


function maybe_track(arg, has_argument_grad::Bool, tape)
    has_argument_grad ? track(arg, tape) : arg
end

@noinline function ReverseDiff.special_reverse_exec!(
        instruction::ReverseDiff.SpecialInstruction{D}) where {D <: Distribution}
    dist::D = instruction.func
    args_maybe_tracked = instruction.input
    score_tracked = instruction.output
    arg_grads = logpdf_grad(dist, map(value, args_maybe_tracked)...)
    value_tracked = args_maybe_tracked[1]
    value_grad = arg_grads[1]
    if istracked(value_tracked)
        if has_output_grad(dist)
            increment_deriv!(value_tracked, value_grad * deriv(score_tracked))
        else
            error("Gradient required but not available for return value of distribution $dist")
        end
    end
    for (i, (arg_maybe_tracked, grad, has_grad)) in enumerate(
            zip(args_maybe_tracked[2:end], arg_grads[2:end], has_argument_grads(dist)))
        if istracked(arg_maybe_tracked)
            if has_grad
                increment_deriv!(arg_maybe_tracked, grad * deriv(score_tracked))
            else
                error("Gradient required but not available for argument $i of $dist")
            end
        end
    end
    nothing
end

###################
# accumulate_param_gradients! #
###################

mutable struct GFBackpropParamsState
    trace::DynamicDSLTrace
    score::TrackedReal
    tape::InstructionTape
    visitor::AddressVisitor
    scaler::Float64

    # only those tracked parameters that are in scope (not including splice calls)
    tracked_params::Dict{Symbol,Any}

    # tracked parameters for all (nested) splice calls
    splice_tracked_params::Dict{Tuple{GenerativeFunction,Symbol},Any}
end

function track_params(tape, params)
    tracked_params = Dict{Symbol,Any}()
    for (name, value) in params
        tracked_params[name] = track(value, tape)
    end
    tracked_params
end

function GFBackpropParamsState(trace::DynamicDSLTrace, tape, params, scaler)
    tracked_params = track_params(tape, params)
    splice_tracked_params = Dict{Tuple{GenerativeFunction,Symbol},Any}()
    score = track(0., tape)
    GFBackpropParamsState(trace, score, tape, AddressVisitor(), scaler,
        tracked_params, splice_tracked_params)
end

function read_param(state::GFBackpropParamsState, name::Symbol)
    value = state.tracked_params[name]
    value
end

function traceat(state::GFBackpropParamsState, dist::Distribution{T},
              args_maybe_tracked, key) where {T}
    local retval::T
    visit!(state.visitor, key)
    retval = get_choice(state.trace, key).retval
    args = map(value, args_maybe_tracked)
    score_tracked = track(logpdf(dist, retval, args...), state.tape)
    record!(state.tape, ReverseDiff.SpecialInstruction, dist,
        (retval, args_maybe_tracked...,), score_tracked)
    state.score += score_tracked
    retval
end

struct BackpropParamsRecord
    gen_fn::GenerativeFunction
    subtrace::Any
    scaler::Float64
end

function traceat(state::GFBackpropParamsState, gen_fn::GenerativeFunction{T},
              args_maybe_tracked, key) where {T}
    local retval::T
    visit!(state.visitor, key)
    subtrace = get_call(state.trace, key).subtrace
    retval = get_retval(subtrace)
    if accepts_output_grad(gen_fn)
        retval_maybe_tracked = track(retval, state.tape)
        @assert istracked(retval_maybe_tracked)
    else
        retval_maybe_tracked = retval
        @assert !istracked(retval_maybe_tracked)
    end
    record!(state.tape, ReverseDiff.SpecialInstruction,
        BackpropParamsRecord(gen_fn, subtrace, state.scaler),
        (args_maybe_tracked...,), retval_maybe_tracked)
    retval_maybe_tracked
end

function splice(state::GFBackpropParamsState, gen_fn::DynamicDSLFunction,
                args_maybe_tracked::Tuple)

    # save previous tracked parameter scope
    prev_tracked_params = state.tracked_params

    # construct new tracked parameter scope
    state.tracked_params = Dict{Symbol,Any}()
    for name in keys(gen_fn.params)
        if haskey(state.splice_tracked_params, (gen_fn, name))
            # parameter was already tracked in another splice call
            state.tracked_params[name] = state.splice_tracked_params[(gen_fn, name)]
        else
            # parameter was not already tracked
            tracked = track(get_param(gen_fn, name), state.tape)
            state.tracked_params[name] = tracked
            state.splice_tracked_params[(gen_fn, name)] = tracked
        end
    end

    retval_maybe_tracked = exec(gen_fn, state, args_maybe_tracked)

    # restore previous tracked parameter scope
    state.tracked_params = prev_tracked_params

    retval_maybe_tracked
end

@noinline function ReverseDiff.special_reverse_exec!(
        instruction::ReverseDiff.SpecialInstruction{BackpropParamsRecord})
    record = instruction.func
    gen_fn = record.gen_fn
    args_maybe_tracked = instruction.input
    retval_maybe_tracked = instruction.output
    if accepts_output_grad(gen_fn)
        @assert istracked(retval_maybe_tracked)
        retval_grad = deriv(retval_maybe_tracked)
    else
        @assert !istracked(retval_maybe_tracked)
        retval_grad = nothing
    end
    arg_grads = accumulate_param_gradients!(record.subtrace, retval_grad, record.scaler)

    # if has_argument_grads(gen_fn) is true for a given argument, then that
    # argument may or may not be tracked. if has_argument_grads(gen_fn) is
    # false for a given argument, then that argument must not be tracked.
    # otherwise it is an error.
    # note: code duplication with accumulate_param_gradients!
    for (i, (arg, grad, has_grad)) in enumerate(
            zip(args_maybe_tracked, arg_grads, has_argument_grads(gen_fn)))
        if has_grad && istracked(arg)
            increment_deriv!(arg, grad)
        elseif !has_grad && istracked(arg)
            error("Gradient required but not available for argument $i of $gen_fn")
        end
    end
    nothing
end

function accumulate_param_gradients!(trace::DynamicDSLTrace, retval_grad, scaler=1.)
    gen_fn = trace.gen_fn
    tape = new_tape()
    state = GFBackpropParamsState(trace, tape, gen_fn.params, scaler)
    args = get_args(trace)
    args_maybe_tracked = (map(maybe_track,
        args, gen_fn.has_argument_grads, fill(tape, length(args)))...,)
    retval_maybe_tracked = exec(gen_fn, state, args_maybe_tracked)

    # note: code duplication with choice_gradients
    if accepts_output_grad(gen_fn)
        if retval_grad == nothing
            error("Return value gradient required but not provided")
        end
        if istracked(retval_maybe_tracked)
            deriv!(retval_maybe_tracked, retval_grad)
        end
        # note: if accepts_output_grad(gen_fn) and
        # !istracked(retval_maybe_tracked), this means the return value did not
        # depend on the gradient source elements for this trace. that is okay.
    else
        if retval_grad != nothing
            error("Return value gradient not supported, but got $retval_grad != nothing")
        end
    end
    seed!(state.score)
    reverse_pass!(tape)

    # increment the gradient accumulators for trainable parameters in scope
    for (name, tracked) in state.tracked_params
        gen_fn.params_grad[name] += deriv(tracked) * state.scaler
    end

    # increment the gradient accumulators for trainable parameters in splice calls
    for ((spliced_gen_fn, name), tracked) in state.splice_tracked_params
        spliced_gen_fn.params_grad[name] += deriv(tracked) * state.scaler
    end

    # return gradients with respect to arguments with gradients, or nothing
    input_grads::Tuple = (map((arg, has_grad) -> has_grad ? deriv(arg) : nothing,
                             args_maybe_tracked, gen_fn.has_argument_grads)...,)
    input_grads
end

##################
# choice_gradients #
##################

mutable struct GFBackpropTraceState
    trace::DynamicDSLTrace
    score::TrackedReal
    tape::InstructionTape
    visitor::AddressVisitor
    params::Dict{Symbol,Any}
    selection::Selection
    tracked_choices::Trie{Any,TrackedReal}
    value_choices::DynamicChoiceMap
    gradient_choices::DynamicChoiceMap
end

function GFBackpropTraceState(trace, selection, params, tape)
    score = track(0., tape)
    visitor = AddressVisitor()
    tracked_choices = Trie{Any,TrackedReal}()
    value_choices = choicemap()
    gradient_choices = choicemap()
    GFBackpropTraceState(trace, score, tape, visitor, params,
        selection, tracked_choices, value_choices, gradient_choices)
end

function fill_gradient_map!(gradient_choices::DynamicChoiceMap,
                             tracked_trie::Trie{Any,TrackedReal})
    for (key, tracked) in get_leaf_nodes(tracked_trie)
        set_value!(gradient_choices, key, deriv(tracked))
    end
    # NOTE: there should be no address collision between these primitive
    # choices and the gen_fn invocations, as enforced by the visitor
    for (key, subtrie) in get_internal_nodes(tracked_trie)
        @assert !has_value(gradient_choices, key) && isempty(get_submap(gradient_choices, key))
        gradient_submap = choicemap()
        fill_gradient_map!(gradient_submap, subtrie)
        set_submap!(gradient_choices, key, gradient_submap)
    end
end

function fill_value_map!(value_choices::DynamicChoiceMap,
                          tracked_trie::Trie{Any,TrackedReal})
    for (key, tracked) in get_leaf_nodes(tracked_trie)
        set_value!(value_choices, key, value(tracked))
    end
    # NOTE: there should be no address collision between these primitive
    # choices and the gen_fn invocations, as enforced by the visitor
    for (key, subtrie) in get_internal_nodes(tracked_trie)
        @assert !has_value(value_choices, key) && !isempty(get_submap(value_choices, key))
        value_choices_submap = choicemap()
        fill_value_map!(value_choices_submap, subtrie)
        set_submap!(value_choices, key, value_choices_submap)
    end
end

function traceat(state::GFBackpropTraceState, dist::Distribution{T},
              args_maybe_tracked, key) where {T}
    local retval::T
    visit!(state.visitor, key)
    retval = get_choice(state.trace, key).retval
    args = map(value, args_maybe_tracked)
    score_tracked = track(logpdf(dist, retval, args...), state.tape)
    if key in state.selection
        tracked_retval = track(retval, state.tape)
        set_leaf_node!(state.tracked_choices, key, tracked_retval)
        record!(state.tape, ReverseDiff.SpecialInstruction, dist,
            (tracked_retval, args_maybe_tracked...,), score_tracked)
        state.score += score_tracked
        return tracked_retval
    else
        record!(state.tape, ReverseDiff.SpecialInstruction, dist,
            (retval, args_maybe_tracked...,), score_tracked)
        state.score += score_tracked
        return retval
    end
end

struct BackpropTraceRecord
    gen_fn::GenerativeFunction
    subtrace::Any
    selection::Selection
    value_choices::DynamicChoiceMap
    gradient_choices::DynamicChoiceMap
    key::Any
end

function traceat(state::GFBackpropTraceState, gen_fn::GenerativeFunction{T,U},
              args_maybe_tracked, key) where {T,U}
    local retval::T
    local subtrace::U
    visit!(state.visitor, key)
    subtrace = get_call(state.trace, key).subtrace
    get_gen_fn(subtrace) === gen_fn || gen_fn_changed_error(key)
    retval = get_retval(subtrace)
    if accepts_output_grad(gen_fn)
        retval_maybe_tracked = track(retval, state.tape)
        @assert istracked(retval_maybe_tracked)
    else
        retval_maybe_tracked = retval
        @assert !istracked(retval_maybe_tracked)
    end
    selection = state.selection[key]
    record = BackpropTraceRecord(gen_fn, subtrace, selection, state.value_choices,
        state.gradient_choices, key)
    record!(state.tape, ReverseDiff.SpecialInstruction, record, (args_maybe_tracked...,), retval_maybe_tracked)
    retval_maybe_tracked
end

function splice(state::GFBackpropTraceState, gen_fn::DynamicDSLFunction,
                args_maybe_tracked::Tuple)
    prev_params = state.params
    state.params = gen_fn.params
    retval = exec(gen_fn, state, args_maybe_tracked)
    state.params = prev_params
    retval
end

@noinline function ReverseDiff.special_reverse_exec!(
        instruction::ReverseDiff.SpecialInstruction{BackpropTraceRecord})
    record = instruction.func
    gen_fn = record.gen_fn
    args_maybe_tracked = instruction.input
    retval_maybe_tracked = instruction.output
    if accepts_output_grad(gen_fn)
        @assert istracked(retval_maybe_tracked)
        retval_grad = deriv(retval_maybe_tracked)
    else
        @assert !istracked(retval_maybe_tracked)
        retval_grad = nothing
    end
    (arg_grads, value_choices, gradient_choices) = choice_gradients(
        record.subtrace, record.selection, retval_grad)
    @assert isempty(get_submap(record.gradient_choices, record.key))
    @assert !has_value(record.gradient_choices, record.key)
    set_submap!(record.gradient_choices, record.key, gradient_choices)
    set_submap!(record.value_choices, record.key, value_choices)

    # if has_argument_grads(gen_fn) is true for a given argument, then that
    # argument may or may not be tracked. if has_argument_grads(gen_fn) is
    # false for a given argument, then that argument must not be tracked.
    # otherwise it is an error.
    # note: code duplication with accumulate_param_gradients!
    for (i, (arg, grad, has_grad)) in enumerate(
            zip(args_maybe_tracked, arg_grads, has_argument_grads(gen_fn)))
        if has_grad && istracked(arg)
            increment_deriv!(arg, grad)
        elseif !has_grad && istracked(arg)
            error("Gradient required but not available for argument $i of $gen_fn")
        end
    end
    nothing
end

function choice_gradients(trace::DynamicDSLTrace, selection::Selection, retval_grad)
    gen_fn = trace.gen_fn
    tape = new_tape()
    state = GFBackpropTraceState(trace, selection, gen_fn.params, tape)
    args = get_args(trace)
    args_maybe_tracked = (map(maybe_track,
        args, gen_fn.has_argument_grads, fill(tape, length(args)))...,)
    retval_maybe_tracked = exec(gen_fn, state, args_maybe_tracked)

    # note: code duplication with accumulate_param_gradients!
    if accepts_output_grad(gen_fn)
        if retval_grad == nothing
            error("Return value gradient required but not provided")
        end
        if istracked(retval_maybe_tracked)
            deriv!(retval_maybe_tracked, retval_grad)
        end
        # note: if accepts_output_grad(gen_fn) and
        # !istracked(retval_maybe_tracked), this means the return value did not
        # depend on the gradient source elements for this trace. that is okay.
    else
        if retval_grad != nothing
            error("Return value gradient not supported, but got $retval_grad != nothing")
        end
    end

    seed!(state.score)
    reverse_pass!(tape)

    # fill trace gradient with gradients with respect to primitive random choices
    fill_gradient_map!(state.gradient_choices, state.tracked_choices)
    fill_value_map!(state.value_choices, state.tracked_choices)

    # return gradients with respect to arguments with gradients, or nothing
    # NOTE: if a value isn't tracked the gradient is nothing
    input_grads::Tuple = (map((arg, has_grad) -> has_grad ? deriv(arg) : nothing,
                             args_maybe_tracked, gen_fn.has_argument_grads)...,)

    (input_grads, state.value_choices, state.gradient_choices)
end


using OmegaCore
using OmegaTest
using OmegaCore.Traits

using Test
using Distributions
using Random

function testsimple1()
  ω = defΩ()()
  ω = OmegaCore.tagrng(ω, Random.GLOBAL_RNG)
  ω2 = defΩ()()
  ω2 = OmegaCore.tagrng(ω2, Random.GLOBAL_RNG)
  x = 1 ~ StdNormal{Float64}()
  map(x, [ω, ω2])
end

function testsimple2()
  x = 1 ~ StdNormal{Float64}()
  rng = Random.GLOBAL_RNG
  OmegaCore.randsample(rng, x, 1; alg = RejectionSample)
end

function testsimple3a()
  x = 1 ~ StdNormal{Float64}()
  ω = defω()
  map(x, typeof(ω)[ω, ω])
end

function testsimple3()
  x = 1 ~ StdNormal{Float64}()
  ω = defω()
  x.([ω, ω])
end

@testset "infer types" begin
  @test isinferred(testsimple1)
  @test isinferred(testsimple2)
  @test isinferred(testsimple3)
  @test isinferred(testsimple3a)
end