text:
# Test internal info methods
@testset "Internal Info Get/Set" begin
    # initialize model and infinite variable
    m = InfiniteModel()
    info = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, false)
    new_info = VariableInfo(true, 0., true, 0., true, 0., true, 0., true, false)
    param = InfOptParameter(IntervalSet(0, 1), Number[], false)
    pref = add_parameter(m, param)
    var = InfiniteVariable(info, (pref,))
    m.vars[1] = var
    vref = InfiniteVariableRef(m, 1)
    bounds = ParameterBounds()
    # _variable_info
    @testset "_variable_info" begin
        @test InfiniteOpt._variable_info(vref) == info
    end
    # _update_variable_info (Infinite)
    @testset "_update_variable_info (Infinite)" begin
        @test isa(InfiniteOpt._update_variable_info(vref, new_info), Nothing)
        @test InfiniteOpt._variable_info(vref) == new_info
    end
    # _update_variable_info (Point)
    @testset "_update_variable_info (Point)" begin
        # initialize point variable
        var = PointVariable(info, vref, (0,))
        m.vars[2] = var
        vref = PointVariableRef(m, 2)
        # test function
        @test isa(InfiniteOpt._update_variable_info(vref, new_info), Nothing)
        @test InfiniteOpt._variable_info(vref) == new_info
    end
    # _update_variable_info (Hold)
    @testset "_update_variable_info (Hold)" begin
        # initialize hold variable
        var = HoldVariable(info, bounds)
        m.vars[3] = var
        vref = HoldVariableRef(m, 3)
        # test function
        @test isa(InfiniteOpt._update_variable_info(vref, new_info), Nothing)
        @test InfiniteOpt._variable_info(vref) == new_info
    end
end

# Test the lower bound methods
@testset "Lower Bound" begin
    # initialize model and 3 test variables
    m = InfiniteModel()
    info1 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, false)
    info2 = VariableInfo(true, 0., false, 0., false, 0., false, 0., false, false)
    info3 = VariableInfo(false, 0., false, 0., true, 0., false, 0., false, false)
    bounds = ParameterBounds()
    var1 = HoldVariable(info1, bounds)
    var2 = HoldVariable(info2, bounds)
    var3 = HoldVariable(info3, bounds)
    m.vars[1] = var1
    m.vars[2] = var2
    m.vars[3] = var3
    m.var_to_name[1] = "var1"
    m.var_to_name[2] = "var2"
    m.var_to_name[3] = "var3"
    m.var_to_lower_bound[2] = 1
    vref1 = HoldVariableRef(m, 1)
    vref2 = HoldVariableRef(m, 2)
    vref3 = HoldVariableRef(m, 3)
    # _update_var_constr_mapping
    @testset "_update_var_constr_mapping" begin
        # test normal with no previous constrs
        @test isa(InfiniteOpt._update_var_constr_mapping([vref2], 1), Nothing)
        @test m.var_to_constrs[JuMP.index(vref2)] == [1]
        # test normal with previous constrs
        @test isa(InfiniteOpt._update_var_constr_mapping([vref2], 2), Nothing)
        @test m.var_to_constrs[JuMP.index(vref2)] == [1, 2]
        # undo changes
        delete!(m.var_to_constrs, JuMP.index(vref2))
    end
    # set_name
    @testset "JuMP.set_name" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.GreaterThan(0.))
        cref = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test isa(set_name(cref, "bob"), Nothing)
        @test m.constr_to_name[1] == "bob"
        # undo changes
        delete!(m.constr_to_name, 1)
    end
    # index
    @testset "JuMP.index" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.GreaterThan(0.))
        cref = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test JuMP.index(cref) == 1
    end
    # add_constraint
    @testset "JuMP.add_constraint" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.GreaterThan(0.))
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test add_constraint(m, constr, "bob") == expected
        @test m.constrs[1] == constr
        @test m.constr_to_name[1] == "bob"
        @test !m.constr_in_var_info[1]
        m.constr_in_var_info[1] = true
    end
    # is_fixed
    @testset "JuMP.is_fixed" begin
        @test is_fixed(vref3)
        @test !is_fixed(vref2)
    end
    # has_lower_bound
    @testset "JuMP.has_lower_bound" begin
        @test !has_lower_bound(vref1)
        @test has_lower_bound(vref2)
    end
    # lower_bound
    @testset "JuMP.lower_bound" begin
        @test lower_bound(vref2) == 0
        @test_throws ErrorException lower_bound(vref1)
    end
    # _lower_bound_index
    @testset "JuMP._lower_bound_index" begin
        @test JuMP._lower_bound_index(vref2) == 1
        @test_throws ErrorException JuMP._lower_bound_index(vref1)
    end
    # set_lower_bound_index
    @testset "JuMP.set_lower_bound_index" begin
        # test function
        @test isa(InfiniteOpt._set_lower_bound_index(vref2, 2), Nothing)
        @test m.var_to_lower_bound[2] == 2
        # undo changes
        m.var_to_lower_bound[2] = 1
    end
    # set_lower_bound
    @testset "JuMP.set_lower_bound" begin
        # test adding new lower bound
        @test isa(set_lower_bound(vref1, 1), Nothing)
        @test has_lower_bound(vref1)
        @test lower_bound(vref1) == 1
        @test m.constr_in_var_info[2]
        @test !optimizer_model_ready(m)
        @test m.constrs[2] == ScalarConstraint(vref1, MOI.GreaterThan(1.))
        # test changing lower bound
        @test isa(set_lower_bound(vref2, 1.5), Nothing)
        @test has_lower_bound(vref2)
        @test lower_bound(vref2) == 1.5
        @test m.constr_in_var_info[1]
        @test !optimizer_model_ready(m)
        @test m.constrs[1] == ScalarConstraint(vref2, MOI.GreaterThan(1.5))
        # test fixed variable error
        @test_throws AssertionError set_lower_bound(vref3, 0)
    end
    # LowerBoundRef
    @testset "JuMP.LowerBoundRef" begin
        # prepare constraint reference
        m.constr_to_name[1] = ""
        constr = ScalarConstraint(vref2, MOI.GreaterThan(0.))
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test for finite variable
        @test LowerBoundRef(vref2) == expected
        # prepare infinite variable
        param = InfOptParameter(IntervalSet(0, 1), Number[], false)
        pref = add_parameter(m, param)
        var = InfiniteVariable(info1, (pref,))
        m.vars[4] = var
        m.var_to_name[4] = "var"
        vref = InfiniteVariableRef(m, 4)
        # prepare lower bound constraint
        set_lower_bound(vref, 0)
        constr = ScalarConstraint(vref, MOI.GreaterThan(0.))
        expected = InfiniteConstraintRef(m, 3, JuMP.shape(constr))
        # test for infinite variable
        @test LowerBoundRef(vref) == expected
    end
end

# Test the upper bound methods
@testset "Upper Bound" begin
    # initialize model and 3 test variables
    m = InfiniteModel()
    info1 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, false)
    info2 = VariableInfo(false, 0., true, 0., false, 0., false, 0., false, false)
    info3 = VariableInfo(false, 0., false, 0., true, 0., false, 0., false, false)
    bounds = ParameterBounds()
    var1 = HoldVariable(info1, bounds)
    var2 = HoldVariable(info2, bounds)
    var3 = HoldVariable(info3, bounds)
    m.vars[1] = var1
    m.vars[2] = var2
    m.vars[3] = var3
    m.var_to_name[1] = "var1"
    m.var_to_name[2] = "var2"
    m.var_to_name[3] = "var3"
    m.var_to_upper_bound[2] = 1
    vref1 = HoldVariableRef(m, 1)
    vref2 = HoldVariableRef(m, 2)
    vref3 = HoldVariableRef(m, 3)
    # add_constraint
    @testset "JuMP.add_constraint" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.LessThan(0.))
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test add_constraint(m, constr, "bob") == expected
        @test m.constrs[1] == constr
        @test m.constr_to_name[1] == "bob"
        @test !m.constr_in_var_info[1]
        m.constr_in_var_info[1] = true
    end
    # has_upper_bound
    @testset "JuMP.has_upper_bound" begin
        @test !has_upper_bound(vref1)
        @test has_upper_bound(vref2)
    end
    # upper_bound
    @testset "JuMP.upper_bound" begin
        @test upper_bound(vref2) == 0
        @test_throws ErrorException upper_bound(vref1)
    end
    # _upper_bound_index
    @testset "JuMP._upper_bound_index" begin
        @test JuMP._upper_bound_index(vref2) == 1
        @test_throws ErrorException JuMP._upper_bound_index(vref1)
    end
    # _set_upper_bound_index
    @testset "_set_upper_bound_index" begin
        # test function
        @test isa(InfiniteOpt._set_upper_bound_index(vref2, 2), Nothing)
        @test m.var_to_upper_bound[2] == 2
        # undo changes
        m.var_to_upper_bound[2] = 1
    end
    # set_upper_bound
    @testset "JuMP.set_upper_bound" begin
        # test adding new upper bound
        @test isa(set_upper_bound(vref1, 1), Nothing)
        @test has_upper_bound(vref1)
        @test upper_bound(vref1) == 1
        @test m.constr_in_var_info[2]
        @test !optimizer_model_ready(m)
        @test m.constrs[2] == ScalarConstraint(vref1, MOI.LessThan(1.))
        # test changing upper bound
        @test isa(set_upper_bound(vref2, 1.5), Nothing)
        @test has_upper_bound(vref2)
        @test upper_bound(vref2) == 1.5
        @test m.constr_in_var_info[1]
        @test !optimizer_model_ready(m)
        @test m.constrs[1] == ScalarConstraint(vref2, MOI.LessThan(1.5))
        # test fixed variable error
        @test_throws AssertionError set_upper_bound(vref3, 0)
    end
    # upperBoundRef
    @testset "JuMP.UpperBoundRef" begin
        # prepare constraint reference
        m.constr_to_name[1] = ""
        constr = ScalarConstraint(vref2, MOI.LessThan(0.))
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test for finite variable
        @test UpperBoundRef(vref2) == expected
        # prepare infinite variable
        param = InfOptParameter(IntervalSet(0, 1), Number[], false)
        pref = add_parameter(m, param)
        var = InfiniteVariable(info1, (pref,))
        m.vars[4] = var
        m.var_to_name[4] = "var"
        vref = InfiniteVariableRef(m, 4)
        # prepare Upper bound constraint
        set_upper_bound(vref, 0)
        constr = ScalarConstraint(vref, MOI.LessThan(0.))
        expected = InfiniteConstraintRef(m, 3, JuMP.shape(constr))
        # test for infinite variable
        @test UpperBoundRef(vref) == expected
    end
end

# Test the fix methods
@testset "Fix" begin
    # initialize model and 4 test variables
    m = InfiniteModel()
    info1 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, false)
    info2 = VariableInfo(false, 0., false, 0., true, 0., false, 0., false, false)
    bounds = ParameterBounds()
    var1 = HoldVariable(info1, bounds)
    var2 = HoldVariable(info2, bounds)
    var3 = HoldVariable(info1, bounds)
    var4 = HoldVariable(info1, bounds)
    m.vars[1] = var1
    m.vars[2] = var2
    m.vars[3] = var3
    m.vars[4] = var4
    m.var_to_name[1] = "var1"
    m.var_to_name[2] = "var2"
    m.var_to_name[3] = "var3"
    m.var_to_name[4] = "var4"
    m.var_to_fix[2] = 1
    vref1 = HoldVariableRef(m, 1)
    vref2 = HoldVariableRef(m, 2)
    vref3 = HoldVariableRef(m, 3)
    vref4 = HoldVariableRef(m, 4)
    # add_constraint
    @testset "JuMP.add_constraint" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.EqualTo(0.))
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test add_constraint(m, constr, "bob") == expected
        @test m.constrs[1] == constr
        @test m.constr_to_name[1] == "bob"
        @test !m.constr_in_var_info[1]
        m.constr_in_var_info[1] = true
    end
    # is_fixed
    @testset "is_fixed" begin
        @test !is_fixed(vref1)
        @test is_fixed(vref2)
    end
    # fix_value
    @testset "JuMP.fix_value" begin
        @test fix_value(vref2) == 0
        @test_throws ErrorException fix_value(vref1)
    end
    # _fix_index
    @testset "JuMP._fix_index" begin
        @test JuMP._fix_index(vref2) == 1
        @test_throws ErrorException JuMP._fix_index(vref1)
    end
    # _set_fix_index
    @testset "_set_fix_index" begin
        # test function
        @test isa(InfiniteOpt._set_fix_index(vref2, 2), Nothing)
        @test m.var_to_fix[2] == 2
        # undo changes
        m.var_to_fix[2] = 1
    end
    # fix
    @testset "JuMP.fix" begin
        # test adding new fix
        @test isa(fix(vref1, 1), Nothing)
        @test is_fixed(vref1)
        @test fix_value(vref1) == 1
        @test m.constr_in_var_info[2]
        @test !optimizer_model_ready(m)
        @test m.constrs[2] == ScalarConstraint(vref1, MOI.EqualTo(1.))
        # test changing fix
        @test isa(fix(vref2, 1.5), Nothing)
        @test is_fixed(vref2)
        @test fix_value(vref2) == 1.5
        @test m.constr_in_var_info[1]
        @test !optimizer_model_ready(m)
        @test m.constrs[1] == ScalarConstraint(vref2, MOI.EqualTo(1.5))
        # add lower and upper bounds to vars 3 and 4
        set_lower_bound(vref3, 0.0)
        set_upper_bound(vref4, 0.0)
        # test lower/upper bound  error
        @test_throws ErrorException fix(vref3, 0)
        @test_throws ErrorException fix(vref4, 0)
        # test forcing with lower
        @test isa(fix(vref3, 1, force = true), Nothing)
        @test is_fixed(vref3)
        @test fix_value(vref3) == 1
        @test m.constr_in_var_info[5]
        @test !optimizer_model_ready(m)
        @test m.constrs[5] == ScalarConstraint(vref3, MOI.EqualTo(1.))
        # test forcing with upper
        @test isa(fix(vref4, 1.5, force = true), Nothing)
        @test is_fixed(vref4)
        @test fix_value(vref4) == 1.5
        @test m.constr_in_var_info[6]
        @test !optimizer_model_ready(m)
        @test m.constrs[6] == ScalarConstraint(vref4, MOI.EqualTo(1.5))
    end
    # FixRef
    @testset "JuMP.FixRef" begin
        # prepare constraint reference
        m.constr_to_name[1] = ""
        constr = ScalarConstraint(vref2, MOI.EqualTo(0.))
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test for finite variable
        @test FixRef(vref2) == expected
        # prepare infinite variable
        param = InfOptParameter(IntervalSet(0, 1), Number[], false)
        pref = add_parameter(m, param)
        var = InfiniteVariable(info1, (pref,))
        m.vars[5] = var
        m.var_to_name[5] = "var"
        vref = InfiniteVariableRef(m, 5)
        # prepare Upper bound constraint
        fix(vref, 0)
        constr = ScalarConstraint(vref, MOI.EqualTo(0.))
        expected = InfiniteConstraintRef(m, 7, JuMP.shape(constr))
        # test for infinite variable
        @test FixRef(vref) == expected
    end
end

# Test the start value methods
@testset "Start Value" begin
    # initialize model and 4 test variables
    m = InfiniteModel()
    info1 = VariableInfo(false, 0., false, 0., false, 0., true, 0., false, false)
    bounds = ParameterBounds()
    var = HoldVariable(info1, bounds)
    m.vars[1] = var
    vref = HoldVariableRef(m, 1)
    # start_value
    @testset "JuMP.start_value" begin
        @test start_value(vref) == 0
    end
    # set_start_value
    @testset "JuMP.set_start_value" begin
        @test isa(set_start_value(vref, 1.5), Nothing)
        @test start_value(vref) == 1.5
        @test !optimizer_model_ready(m)
    end
end

# Test the binary methods
@testset "Binary" begin
    # initialize model and 3 test variables
    m = InfiniteModel()
    info1 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, false)
    info2 = VariableInfo(false, 0., false, 0., false, 0., false, 0., true, false)
    info3 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, true)
    bounds = ParameterBounds()
    var1 = HoldVariable(info1, bounds)
    var2 = HoldVariable(info2, bounds)
    var3 = HoldVariable(info3, bounds)
    m.vars[1] = var1
    m.vars[2] = var2
    m.vars[3] = var3
    m.var_to_name[1] = "var1"
    m.var_to_name[2] = "var2"
    m.var_to_name[3] = "var3"
    m.var_to_zero_one[2] = 1
    vref1 = HoldVariableRef(m, 1)
    vref2 = HoldVariableRef(m, 2)
    vref3 = HoldVariableRef(m, 3)
    # add_constraint
    @testset "JuMP.add_constraint" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.ZeroOne())
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test add_constraint(m, constr, "bob") == expected
        @test m.constrs[1] == constr
        @test m.constr_to_name[1] == "bob"
        @test !m.constr_in_var_info[1]
        m.constr_in_var_info[1] = true
    end
    # is_binary
    @testset "JuMP.is_binary" begin
        @test !is_binary(vref1)
        @test is_binary(vref2)
    end
    # _binary_index
    @testset "JuMP._binary_index" begin
        @test JuMP._binary_index(vref2) == 1
        @test_throws ErrorException JuMP._binary_index(vref1)
    end
    # _set_binary_index
    @testset "_set_binary_index" begin
        # test function
        @test isa(InfiniteOpt._set_binary_index(vref2, 2), Nothing)
        @test m.var_to_zero_one[2] == 2
        # undo changes
        m.var_to_zero_one[2] = 1
    end
    # set_binary
    @testset "JuMP.set_binary" begin
        # test adding binary
        @test isa(set_binary(vref1), Nothing)
        @test is_binary(vref1)
        @test m.constr_in_var_info[2]
        @test !optimizer_model_ready(m)
        @test m.constrs[2] == ScalarConstraint(vref1, MOI.ZeroOne())
        # test setting binary again
        @test isa(set_binary(vref2), Nothing)
        @test is_binary(vref2)
        @test m.constr_in_var_info[1]
        @test !optimizer_model_ready(m)
        @test m.constrs[1] == ScalarConstraint(vref2, MOI.ZeroOne())
        # test integer variable error
        @test_throws ErrorException set_binary(vref3)
    end
    # BinaryRef
    @testset "BinaryRef" begin
        # prepare constraint reference
        m.constr_to_name[1] = ""
        constr = ScalarConstraint(vref2, MOI.ZeroOne())
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test for finite variable
        @test BinaryRef(vref2) == expected
        # prepare infinite variable
        param = InfOptParameter(IntervalSet(0, 1), Number[], false)
        pref = add_parameter(m, param)
        var = InfiniteVariable(info1, (pref,))
        m.vars[4] = var
        m.var_to_name[4] = "var"
        vref = InfiniteVariableRef(m, 4)
        # prepare Upper bound constraint
        set_binary(vref)
        constr = ScalarConstraint(vref, MOI.ZeroOne())
        expected = InfiniteConstraintRef(m, 3, JuMP.shape(constr))
        # test for infinite variable
        @test BinaryRef(vref) == expected
    end
end

# Test the integer methods
@testset "Integer" begin
    # initialize model and 3 test variables
    m = InfiniteModel()
    info1 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, false)
    info2 = VariableInfo(false, 0., false, 0., false, 0., false, 0., false, true)
    info3 = VariableInfo(false, 0., false, 0., false, 0., false, 0., true, false)
    bounds = ParameterBounds()
    var1 = HoldVariable(info1, bounds)
    var2 = HoldVariable(info2, bounds)
    var3 = HoldVariable(info3, bounds)
    m.vars[1] = var1
    m.vars[2] = var2
    m.vars[3] = var3
    m.var_to_name[1] = "var1"
    m.var_to_name[2] = "var2"
    m.var_to_name[3] = "var3"
    m.var_to_integrality[2] = 1
    vref1 = HoldVariableRef(m, 1)
    vref2 = HoldVariableRef(m, 2)
    vref3 = HoldVariableRef(m, 3)
    # add_constraint
    @testset "JuMP.add_constraint" begin
        # prepare constraint reference
        constr = ScalarConstraint(vref2, MOI.Integer())
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test function
        @test add_constraint(m, constr, "bob") == expected
        @test m.constrs[1] == constr
        @test m.constr_to_name[1] == "bob"
        @test !m.constr_in_var_info[1]
        m.constr_in_var_info[1] = true
    end
    # is_integer
    @testset "JuMP.is_integer" begin
        @test !is_integer(vref1)
        @test is_integer(vref2)
    end
    # _integer_index
    @testset "JuMP._integer_index" begin
        @test JuMP._integer_index(vref2) == 1
        @test_throws ErrorException JuMP._integer_index(vref1)
    end
    # _set_integer_index
    @testset "_set_integer_index" begin
        # test function
        @test isa(InfiniteOpt._set_integer_index(vref2, 2), Nothing)
        @test m.var_to_integrality[2] == 2
        # undo changes
        m.var_to_integrality[2] = 1
    end
    # set_integer
    @testset "JuMP.set_integer" begin
        # test adding integer
        @test isa(set_integer(vref1), Nothing)
        @test is_integer(vref1)
        @test m.constr_in_var_info[2]
        @test !optimizer_model_ready(m)
        @test m.constrs[2] == ScalarConstraint(vref1, MOI.Integer())
        # test setting integer again
        @test isa(set_integer(vref2), Nothing)
        @test is_integer(vref2)
        @test m.constr_in_var_info[1]
        @test !optimizer_model_ready(m)
        @test m.constrs[1] == ScalarConstraint(vref2, MOI.Integer())
        # test integer variable error
        @test_throws ErrorException set_integer(vref3)
    end
    # IntegerRef
    @testset "IntegerRef" begin
        # prepare constraint reference
        m.constr_to_name[1] = ""
        constr = ScalarConstraint(vref2, MOI.Integer())
        expected = FiniteConstraintRef(m, 1, JuMP.shape(constr))
        # test for finite variable
        @test IntegerRef(vref2) == expected
        # prepare infinite variable
        param = InfOptParameter(IntervalSet(0, 1), Number[], false)
        pref = add_parameter(m, param)
        var = InfiniteVariable(info1, (pref,))
        m.vars[4] = var
        m.var_to_name[4] = "var"
        vref = InfiniteVariableRef(m, 4)
        # prepare Upper bound constraint
        set_integer(vref)
        constr = ScalarConstraint(vref, MOI.Integer())
        expected = InfiniteConstraintRef(m, 3, JuMP.shape(constr))
        # test for infinite variable
        @test IntegerRef(vref) == expected
    end
end


text:
using JLD
using DataFrames
using Dates
using CSV
using Printf
using RCall

function create_weekly_data(data_name,output_name)
    df = DataFrame(CSV.File(data_name))

    # """ Pre-vaccine analysis """
    tot_days = size(df)[1]
    time_interval = 7

    df_weekly = df[1:time_interval:tot_days,:]

    scale_SIR(SIR) = SIR./1e4

    transform!(df_weekly,[:S] => ByRow(scale_SIR) => :S)
    transform!(df_weekly,[:I] => ByRow(scale_SIR) => :I)
    transform!(df_weekly,[:R] => ByRow(scale_SIR) => :R)

    CSV.write(output_name,df_weekly)
end

function process_result(data_path,data_name)
    dat = load(data_path)
    key = [i for i in keys(dat)]
    dat = dat[key[1]]
    long = map(x-> dat[x,:,:],[1;2;3])
    R"saveRDS($long,$data_name)"
end

# save weekly data
covid_csv_files = readdir("data")
covid_csv_files = covid_csv_files[occursin.("covid19_bc_",covid_csv_files)]
output_names = "data/weekly/" .* covid_csv_files
covid_csv_files = "data/" .* covid_csv_files
create_weekly_data.(covid_csv_files,output_names)

# save results as rds files
results = readdir("results/covid/solution")
output_names = "results/covid/covid_rds/" .* replace.(results,".jld" => ".rds")
results = "results/covid/solution/" .* results
process_result.(results,output_names)

ext: jl


text:
##

using ACE, ACEbase, StaticArrays
using Printf, Test, LinearAlgebra, ACE.Testing, Random
using ACE: evaluate, evaluate_d, SymmetricBasis, PIBasis,
           grad_config, grad_params, O3
using ACEbase.Testing: fdtest, println_slim

randconfig(B1p, nX) = ACEConfig( rand(PositionState{Float64}, B1p.bases[1], nX) )

##

@info("Basic test of LinearACEModel construction and evaluation")

# construct the 1p-basis
maxdeg = 6
ord = 3
Bsel = SimpleSparseBasis(ord, maxdeg)

B1p = ACE.Utils.RnYlm_1pbasis(; maxdeg=maxdeg)

# generate a configuration
cfg = randconfig(B1p, 10)
Xs = cfg.Xs

φ = ACE.Invariant()
basis = SymmetricBasis(φ, B1p, O3(), Bsel)

BB = evaluate(basis, cfg)
c = rand(length(BB)) .- 0.5
naive = ACE.LinearACEModel(basis, c, evaluator = :naive)
standard = ACE.LinearACEModel(basis, c, evaluator = :standard)

## FIO

@info("Check FIO")
using ACEbase.Testing: test_fio
println_slim(@test(all(test_fio(naive; warntype = false))))
println_slim(@test(all(test_fio(standard; warntype = false))))

##

evaluate_ref(basis, cfg, c) = sum(evaluate(basis, cfg) .* c)
grad_config_ref(basis, cfg, c) = permutedims(evaluate_d(basis, cfg)) * c
grad_params_ref(basis, cfg, c) = evaluate(basis, cfg)
grad_params_config_ref(basis, cfg, c) = evaluate_d(basis, cfg)

evaluate(naive, cfg) ≈  evaluate(standard, cfg)
evaluate_ref(basis, cfg, c) ≈ evaluate(naive, cfg)
grad_params(naive, cfg) ≈ grad_params(standard, cfg)
grad_config(naive, cfg) ≈ grad_config(standard, cfg)

(fun, funref, str) = (ACE.grad_params_config, grad_params_config_ref, "grad_params_config")

for (fun, funref, str) in [
         (evaluate, evaluate_ref, "evaluate"),
         (ACE.grad_config, grad_config_ref, "grad_config"),
         (ACE.grad_params, grad_params_ref, "grad_params"),
         (ACE.grad_params_config, grad_params_config_ref, "grad_params_config"),
      ]
   @info("Testing `$str` for different model evaluators")
   for ntest = 1:30
      local c, cfg
      cfg = randconfig(B1p, 10)
      c = rand(length(basis)) .- 0.5
      ACE.set_params!(naive, c)
      ACE.set_params!(standard, c)
      val = funref(basis, cfg, c)
      val_naive = fun(naive, cfg)
      val_standard = fun(standard, cfg)
      print_tf(@test( val ≈ val_naive ≈ val_standard ))
   end
   println()
end

##

@info("Evaluate LinearACEModel with vector")
for _ = 1:20
   cfg = randconfig(B1p, rand(8:15))
   print_tf(@test evaluate(standard, cfg) ≈ evaluate(standard, cfg.Xs))
   print_tf(@test grad_config(standard, cfg) ≈ grad_config(standard, cfg.Xs))
   print_tf(@test grad_params(standard, cfg) ≈ grad_params(standard, cfg.Xs))
end
println()

##

@info("Test a Linear Model with EuclideanVector output")
maxdeg = 6; ord = 3
Bsel = SimpleSparseBasis(ord, maxdeg)
B1p = ACE.Utils.RnYlm_1pbasis(; maxdeg=maxdeg)
φ = ACE.EuclideanVector{Float64}()
basis = SymmetricBasis(φ, B1p, O3(), Bsel; isreal=true)

##
@info(" test evaluation of basis vs model ")

for ntest = 1:30
   local cfg, c
   cfg = randconfig(B1p, 10)
   BB = evaluate(basis, cfg)
   c = randn(length(BB)) ./ (1:length(BB)).^2
   model = ACE.LinearACEModel(basis, c, evaluator = :standard)
   val1 = sum(c .* BB)
   val2 = evaluate(model, cfg)
   print_tf(@test( val1 ≈ val2 ))
end
println()

##
@info("test gradients of LinearACEModel of equivariant vector")

BB = evaluate(basis, cfg)
c = randn(length(BB)) ./ (1:length(BB)).^2
model = ACE.LinearACEModel(basis, c, evaluator = :standard)

cfg = randconfig(B1p, 10)
evaluate(model, cfg)

dBB = evaluate_d(basis, cfg)
val1 = sum(c[i] * dBB[i,:] for i = 1:length(basis))
val2 = grad_config(model, cfg)
println_slim( @test( val1 ≈ val2 ))

##

@info("And a finite-difference test for good measure")

for ntest = 1:30
   w = randn(3)
   Us = randn(SVector{3, Float64}, length(cfg))
   cfg.Xs + 0.01 * Us
   F = t -> real(w' * evaluate(model, ACEConfig(cfg.Xs + t * Us)).val)
   F(0.01)
   dF = t -> sum( real(w' * gi.rr * Ui) for (Ui, gi) in zip(Us,
                           grad_config(model, ACEConfig(cfg.Xs + t * Us))) )
   dF(0.01)
   print_tf(@test all(ACEbase.Testing.fdtest(F, dF, 0.0, verbose=false)))
end
println()


text:
using Test
using QuantumLattices.Prerequisites.CompositeStructures

struct FT{S,N,T} <: CompositeNTuple{N,T}
    info::S
    contents::NTuple{N,T}
end

@testset "CompositeNTuple" begin
    t=FT("Info",(1,2,3,4))
    @test t|>length==4
    @test t|>typeof|>length==4
    @test t|>eltype==Int
    @test t|>typeof|>eltype==Int
    @test t==deepcopy(t)
    @test isequal(t,deepcopy(t))
    @test hash(t)==hash(deepcopy(t))
    @test t[1]==1
    @test t[end]==4
    @test t[1:3]==FT("Info",(1,2,3))
    @test collect(t)==[1,2,3,4]
    @test collect(t|>Iterators.reverse)==[4,3,2,1]
    @test t|>keys==(1,2,3,4)|>keys
    @test t|>values==(1,2,3,4)|>values
    @test t|>pairs==(1,2,3,4)|>pairs
    @test t|>reverse==FT("Info",(4,3,2,1))
    @test convert(Tuple,t)==(1,2,3,4)
end

struct FV{S,T} <: CompositeVector{T}
    info::S
    contents::Vector{T}
end

@testset "CompositeVector" begin
    v=FV("Info",[1,3,2,4])
    @test size(v)==(4,)
    @test size(v,1)==4
    @test length(v)==4
    @test v==deepcopy(v)
    @test isequal(v,deepcopy(v))
    @test v[end]==4
    @test v[2]==3
    @test v[CartesianIndex(2,)]==3
    @test v[[1,3,4]]==FV("Info",[1,2,4])
    @test v[1:3]==FV("Info",[1,3,2])
    @test (v[1]=10;v[2:3]=11:12;v==FV("Info",[10,11,12,4]))
    @test push!(v,1)==FV("Info",[10,11,12,4,1])
    @test push!(v,2,3)==FV("Info",[10,11,12,4,1,2,3])
    @test pushfirst!(v,20,21)==FV("Info",[20,21,10,11,12,4,1,2,3])
    @test insert!(v,2,30)==FV("Info",[20,30,21,10,11,12,4,1,2,3])
    @test append!(v,[0,-1])==FV("Info",[20,30,21,10,11,12,4,1,2,3,0,-1])
    @test prepend!(v,[8,9])==FV("Info",[8,9,20,30,21,10,11,12,4,1,2,3,0,-1])
    @test splice!(v,2)==9 && v==FV("Info",[8,20,30,21,10,11,12,4,1,2,3,0,-1])
    @test splice!(v,1,9)==8 && v==FV("Info",[9,20,30,21,10,11,12,4,1,2,3,0,-1])
    @test splice!(v,1:3)==FV("Info",[9,20,30]) && v==FV("Info",[21,10,11,12,4,1,2,3,0,-1])
    @test splice!(v,1:3,[8,7,6])==FV("Info",[21,10,11]) && v==FV("Info",[8,7,6,12,4,1,2,3,0,-1])
    @test deleteat!(v,4)==FV("Info",[8,7,6,4,1,2,3,0,-1])
    @test deleteat!(v,[1,2])==FV("Info",[6,4,1,2,3,0,-1])
    @test deleteat!(v,5:6)==FV("Info",[6,4,1,2,-1])
    @test pop!(v)==-1 && v==FV("Info",[6,4,1,2])
    @test popfirst!(v)==6 && v==FV("Info",[4,1,2])
    @test empty!(v)==FV("Info",Int[]) && v==FV("Info",Int[])
    v=FV("Info",[1,3,2,2,4])
    @test empty(v)==FV("Info",Int[])
    @test v|>collect==[1,3,2,2,4]
    @test v|>keys==[1,3,2,2,4]|>keys
    @test v|>values==[1,3,2,2,4]|>values
    @test v|>pairs==[1,3,2,2,4]|>pairs
    @test convert(Vector,v)==[1,3,2,2,4]
    @test v|>reverse==FV("Info",[4,2,2,3,1])
    @test sort(v)==FV("Info",[1,2,2,3,4]) && v==FV("Info",[1,3,2,2,4])
    @test sort!(v)==FV("Info",[1,2,2,3,4]) && v==FV("Info",[1,2,2,3,4])
    @test filter(x->x<=3,v)==FV("Info",[1,2,2,3]) && v==FV("Info",[1,2,2,3,4])
    @test filter!(x->x<=3,v)==FV("Info",[1,2,2,3]) && v==FV("Info",[1,2,2,3])
    @test findfirst(isequal(2),v)==2
    @test findlast(isequal(2),v)==3
    @test findall(isequal(2),v)==[2,3]
end

struct FD{S,P,I} <: CompositeDict{P,I}
    info::S
    contents::Dict{P,I}
end

@testset "CompositeDict" begin
    d=FD("Info",Dict("a"=>1,"b"=>2))
    @test d==deepcopy(d)
    @test isequal(d,deepcopy(d))
    @test isempty(d)==false
    @test length(d)==2
    @test haskey(d,"a")==true && haskey(d,"d")==false
    @test Pair("a",1) ∈ d && Pair("d",4) ∉ d
    @test get(d,"a",2)==1
    @test get(d,"d",4)==4
    @test get(()->4,d,"d")==4
    @test get!(d,"d",4)==4 && d==FD("Info",Dict("a"=>1,"b"=>2,"d"=>4))
    @test get!(()->5,d,"d")==4 && d==FD("Info",Dict("a"=>1,"b"=>2,"d"=>4))
    @test getkey(d,"e","e")=="e"
    @test d["d"]==4
    @test push!(d,Pair("e",4))==FD("Info",Dict("a"=>1,"b"=>2,"d"=>4,"e"=>4)) && d==FD("Info",Dict("a"=>1,"b"=>2,"d"=>4,"e"=>4))
    @test (d["d"]=4;d["e"]=5;d==FD("Info",Dict("a"=>1,"b"=>2,"d"=>4,"e"=>5)))
    @test pop!(d)==Pair("e",5) && d==FD("Info",Dict("a"=>1,"b"=>2,"d"=>4))
    @test pop!(d,"a")==1 && d==FD("Info",Dict("b"=>2,"d"=>4))
    @test pop!(d,"a",1)==1 && d==FD("Info",Dict("b"=>2,"d"=>4))
    @test delete!(d,"b")==FD("Info",Dict("d"=>4)) && d==FD("Info",Dict("d"=>4))
    @test empty!(d)==FD("Info",Dict{String,Int}()) && d==FD("Info",Dict{String,Int}())
    d=FD("Info",Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4))
    @test merge(FD("Info",Dict("a"=>1,"b"=>2)),FD("Info",Dict("c"=>3,"d"=>4)))==d
    @test merge(+,FD("Info",Dict("a"=>1,"b"=>2,"c"=>1)),FD("Info",Dict("c"=>2,"d"=>4)))==d
    @test empty(d)==FD("Info",Dict{String,Int}())
    @test d|>collect==Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4)|>collect
    @test d|>keys|>collect==Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4)|>keys|>collect
    @test d|>values|>collect==Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4)|>values|>collect
    @test d|>pairs|>collect==Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4)|>pairs|>collect
    @test convert(Dict,d)==Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4)
    @test filter(p->p.second<=3,d)==FD("Info",Dict("a"=>1,"b"=>2,"c"=>3)) && d==FD("Info",Dict("a"=>1,"b"=>2,"c"=>3,"d"=>4))
    @test filter!(p->p.second<=3,d)==FD("Info",Dict("a"=>1,"b"=>2,"c"=>3)) && d==FD("Info",Dict("a"=>1,"b"=>2,"c"=>3))
end

@testset "NamedContainer" begin
    @test NamedContainer{(:a,:b)}((1,'h'))==(a=1,b='h')
    @test NamedContainer{()}(())==NamedTuple()
end


text:
########################################################################
#                            OPERATIONS                                #
########################################################################

abstract type AbstractOp end

Base.getproperty(op::AbstractOp, f::Symbol) = f == :typ ? typeof(op.val) : getfield(op, f)

## Input

mutable struct Input <: AbstractOp
    id::Int
    val::Any
end

Base.show(io::IO, op::Input) = print(io, "inp %$(op.id)::$(op.typ)")

## Constant

mutable struct Constant <: AbstractOp
    id::Int
    typ::Type
    val
end

Constant(id::Int, val) = Constant(id, typeof(val), val)
Base.show(io::IO, op::Constant) = print(io, "const %$(op.id) = $(op.val)::$(op.typ)")

## Assign

mutable struct Assign <: AbstractOp
    id::Int
    src_id::Int
    val
end

Base.show(io::IO, op::Assign) = print(io, "%$(op.id) = %$(op.src_id)::$(op.typ)")

## Call

mutable struct Call <: AbstractOp
    id::Int
    val::Any
    fn::Union{Function, Type}
    args::Vector{Int}
end

# Call(id::Int, val::Any, fn::Union{Function, Type}, args::Vector{Int}) =
#     Call(id, val, fn, args)
Base.getproperty(op::Input, f::Call) = f == :typ ? typeof(op.val) : getfield(op, f)

function Base.show(io::IO, op::Call)
    arg_str = join(["%$(id)" for id in op.args], ", ")
    print(io, "%$(op.id) = $(op.fn)($arg_str)::$(op.typ)")
end

"""
Copy struct `x` replacing specified fields with new values
"""
function copy_with(x; kvs...)
    T = typeof(x)
    d = Dict(kvs)
    flds = fieldnames(T)
    new_flds = [get(d, f, getfield(x, f)) for f in flds]
    return T(new_flds...)
end

########################################################################
#                                 TAPE                                 #
########################################################################

const MaybeFunction = Union{Function, Nothing}

mutable struct Tape
    # linearized execution graph
    ops::Vector{<:AbstractOp}
    # id of result variable
    resultid::Int
    # derivs[var_id] == grad_id
    derivs::Dict{Int,Int}
    # mapping of argid -> Dict(struct field paths -> var id)
    fieldpaths::Dict{Int, Dict}
    # compiled tape or nothing
    compiled::MaybeFunction
    # device of the tape
    device::AbstractDevice
end

Tape(device::AbstractDevice) = Tape(AbstractOp[], -1, Dict(), Dict(), nothing, device)
Tape() = Tape(CPU())
Base.similar(tape::Tape) = Tape(AbstractOp[], tape.resultid, tape.derivs,
                                tape.fieldpaths, tape.compiled, tape.device)

function Base.show(io::IO, tape::Tape)
    println(io, "Tape")
    for op in tape.ops
        println(io, "  $op")
    end
end

Base.getindex(tape::Tape, i...) = getindex(tape.ops, i...)
# Base.getindex(tape::Tape, i::String...) =
#     getindex(tape.ops, [parse(Int, s[2:end]) for s in i]...)
# Base.getindex(tape::Tape, i::Symbol...) = getindex(tape, map(string, i)...)
Base.setindex!(tape::Tape, op::AbstractOp, i::Integer) = (tape.ops[i] = op)
Base.lastindex(tape::Tape) = lastindex(tape.ops)
Base.length(tape::Tape) = length(tape.ops)
Base.iterate(tape::Tape) = iterate(tape.ops)
Base.iterate(tape::Tape, s) = iterate(tape.ops, s)
Base.push!(tape::Tape, op::AbstractOp) = push!(tape.ops, op)

"""
Record a new operation (defined by its type and arguments for constructor)
to a tape and return ID of this new op.
"""
function record!(tape::Tape, optype::Type{<:AbstractOp}, args...)
    ret_id = length(tape) + 1
    push!(tape, optype(ret_id, args...))
    return ret_id
end

"""
Parse a complex call expression and record corresponding operations to a tape.

Optionally takes substitution table (st parameter) to replace known symbols with
provided values.

Keyword params:

 * st::Dict - substitution table, used to replace symbols in ex with tape op ids
 * bcast::Bool - replace all function calls with corresponding broadcasting
"""
function record_expr!(tape::Tape, ex::Expr; st=Dict(), bcast=false)
    # special case: x => identity(x)
    if ex isa Symbol
        ex = :($identity($ex))
    end
    # special case: x[i] => getindex(x, i)
    if Meta.isexpr(ex, :ref)
        ex = rewrite(ex, :(_x[_i]), :($getindex(_x, _i)))
    end
    @assert Meta.isexpr(ex, :call) "Expression isn't a call"
    new_op_args = Vector{Int}(undef, length(ex.args) - 1)
    for (i, x) in enumerate(ex.args[2:end])
        if haskey(st, x)
            # op ID comes from substitution table
            new_op_args[i] = st[x]
        elseif Meta.isexpr(x, :call)
            # recursively record arg expression
            arg_id = record_expr!(tape, x; st=st, bcast=bcast)
            new_op_args[i] = arg_id
        else
            # treat as constant
            arg_id = record!(tape, Constant, x)
            new_op_args[i] = arg_id
        end
    end
    fn = ex.args[1]
    fn = device_function(tape.device, fn)
    if bcast
        retval = broadcast(fn, [tape[id].val for id in new_op_args]...)
        fn_id = record!(tape, Constant, fn)
        return record!(tape, Call, retval, broadcast, [fn_id; new_op_args])
    elseif fn isa Symbol && fn in Set([:.+, :.*, :.-, :./, :.^])
        _fn = eval(Symbol(string(fn)[2:end]))
        retval = broadcast(_fn, [tape[id].val for id in new_op_args]...)
        fn_id = record!(tape, Constant, _fn)
        return record!(tape, Call, retval, broadcast, [fn_id; new_op_args])
    else
        retval = fn([tape[id].val for id in new_op_args]...)
        return record!(tape, Call, retval, fn, new_op_args)
    end
end

function record_expr!(tape::Tape, x::Symbol; st=Dict(), bcast=false)
    id = st[x]
    return record!(tape, Assign, id, tape[id].val)
end

function record_expr!(tape::Tape, x; st, bcast=false)
    return record!(tape, Constant, x)
end

########################################################################
#                           TRANSFORMATIONS                            #
########################################################################

"""
Replace all call arguments according to the provided dict
"""
function replace_in_args!(tape::Tape, st::Dict)
    for (i, op) in enumerate(tape)
        if op isa Call
            new_args = [get(st, x, x) for x in op.args]
            tape[i] = copy_with(op, args=new_args)
        end
    end
end

"""
Recover broadcast operation from Broadcast.broadcasted and Broadcast.materialize
"""
function recover_broadcast(tape::Tape)
    new_tape = copy_with(tape; ops=AbstractOp[])
    # TODO: seems like we don't need subs table any more
    # remove after squash_assigned is implemented
    st = Dict()
    for (id, op) in enumerate(tape)
        if op isa Call && op.fn === Broadcast.broadcasted
            # replace Broadcast.broadcasted with just broadcast & materialize value
            val = Broadcast.materialize(op.val)
            new_id = length(new_tape) + 1
            push!(new_tape.ops, copy_with(op; fn=broadcast, val=val))
            st[id] = length(new_tape)
        elseif op isa Call && op.fn === Broadcast.materialize
            # materialize becomes just assignment
            new_id = record!(new_tape, Assign, op.args[1], op.val)
            st[id] = new_id
        else
            # record any other operations as is
            new_id = length(new_tape) + 1
            push!(new_tape.ops, copy_with(op; id=new_id))
            st[id] = length(new_tape)
        end
    end
    replace_in_args!(new_tape, st)
    new_tape.resultid = get(st, tape.resultid, tape.resultid)
    return new_tape
end

# function squash_assigned(tape::Tape)
#     # note: after update to CuArrays v1.5.0 deepcopy(tape.ops) fails for Lilith.GRU if there's Assign op
#     # I don't know the reason, but [deepcopy(op) for op in tape.ops] fixes it
#     tape = copy_with(tape; ops=[deepcopy(op) for op in tape.ops])
#     # 1. compute subs table for chains of assignment operations
#     # and replace equivalent op ids
#     assign_st = Dict()
#     for (id, op) in enumerate(tape)
#         if op isa Assign
#             # propagate replacement through known st
#             src_id = op.src_id
#             while haskey(assign_st, src_id)
#                 src_id = assign_st[src_id]
#             end
#             assign_st[id] = src_id
#         end
#     end
#     replace_in_args!(tape, assign_st)
#     tape.resultid = get(assign_st, tape.resultid, tape.resultid)
#     # 2. create new ops w/o assignments
#     new_tape = copy_with(tape; ops=AbstractOp[])
#     reindex_st = Dict()
#     for (id, op) in enumerate(tape)
#         if !isa(op, Assign)
#             # record any other operations as is
#             new_id = length(new_tape) + 1
#             push!(new_tape.ops, copy_with(op; id=new_id))
#             if id != new_id
#                 reindex_st[id] = new_id
#             end
#         end
#     end
#     replace_in_args!(new_tape, reindex_st)
#     reindex_fields!(new_tape, reindex_st)
#     return new_tape
# end

function squash_assigned(tape::Tape)
    new_tape = copy_with(tape, ops=AbstractOp[])
    st = Dict{Int, Int}()  # substitution table for op indices
    for id=1:length(tape)
        # id == 228 && break
        op = reindex(tape[id], st)
        if op isa Assign
            # dig until not-assign argument is found
            root = op
            while root isa Assign
                root = new_tape[root.src_id]
            end
            st[op.id] = root.id   # replace this op id on _old_ tape with root id on the _new_ tape
            # note: not pushing this op into new tape
            # println("assign: $op; replacing $(op.id) => $(root.id)")
        else
            new_id = length(new_tape) + 1
            new_op = copy_with(op, id=new_id)
            push!(new_tape, new_op)
            if new_id != id
                st[id] = new_id
                # println("subs $id => $new_id")
            end
        end
    end
    reindex_fields!(new_tape, st)
    return new_tape
end

# function test_fuse_assigned()
#     tape = Tape()
#     record!(tape, Input, 3.0)
#     record!(tape, Input, 2.0)
#     record!(tape, Assign, 1, tape[1].val)
#     record!(tape, Assign, 3, tape[3].val)
#     record!(tape, Call, 5.0, +, [2, 4])
#     record!(tape, Assign, 5, tape[5].val)
#     tape.resultid = length(tape)
# end

"""
Unwind iterate() sequences into plain __getfield__ expressions.
unwind_iterate() doesn't remove unused elements for performance reasons,
so remove_unused() should be called after it.
"""
function unwind_iterate(tape::Tape)
    tape = copy_with(tape)
    for op in tape
        if (op isa Call && op.fn in (getfield, __getfield__)
            && tape[op.args[1]] isa Call && tape[op.args[1]].fn == Base.iterate
            && tape[op.args[2]] isa Constant && tape[op.args[2]].val == 1)
            iterate_op = tape[op.args[1]]
            iterable_op = tape[iterate_op.args[1]]
            idx = length(iterate_op.args) > 1 ? tape[iterate_op.args[2]].val : 1
            if iterable_op.val isa Tuple || iterable_op.val isa Vector || iterable_op.val isa UnitRange
                # 1. Replace iterable op with index in the original iterable
                tape[iterate_op.id] = Constant(iterate_op.id, idx)
                # 2. Replace __getfield__ on iterator with __getfield__ or getindex on original iterable
                idx_id = iterate_op.id
                obj_id = iterable_op.id
                # TODO: in which other cases getindex is better than __getfield__?
                get_op = iterable_op.val isa UnitRange ? getindex : __getfield__
                tape[op.id] = Call(op.id, op.val, get_op, [obj_id, idx_id])
            end
        end
    end
    return tape
end

function simplify(tape::Tape)
    tape = recover_broadcast(tape)
    tape = squash_assigned(tape)
    tape = unwind_iterate(tape)
    tape = eliminate_common(tape)
    tape = remove_unused(tape)
    return tape
end

########################################################################
#                              EXECUTION                               #
########################################################################

exec!(tape::Tape, op::Input) = ()
exec!(tape::Tape, op::Constant) = ()
exec!(tape::Tape, op::Assign) = (op.val = tape[op.src_id].val)
exec!(tape::Tape, op::Call) = (op.val = op.fn([tape[id].val for id in op.args]...))

function play!(tape::Tape, args...; use_compiled=true, debug=false)
    for (i, val) in enumerate(args)
        @assert(tape[i] isa Input, "More arguments than the original function had")
        tape[i].val = val
    end
    if use_compiled && tape.compiled != nothing
        Base.invokelatest(tape.compiled)
    else
        for op in tape
            if debug
                println(op)
            end
            exec!(tape, op)
        end
    end
    return tape[tape.resultid].val
end


text:
export howell_form2

# we assume that B has at least as many rows as columns
function howell_form2(A::nmod_mat)
  #A = deepcopy(B)
  n = A.r
  m = A.c
  if n < m
    A = vcat(A, MatrixSpace(base_ring(A), m-n, m)())
  end
  for j in 1:m
    for i in j+1:n
      #print("Computing xgcd of $(_raw_getindex(A,j,j)), $(_raw_getindex(A,i,j))\n")
      g,s,t,u,v = _xxgcd(_raw_getindex(A,j,j),_raw_getindex(A,i,j),A._n)
      #println("$g $s $t $u $v ")
      for k in 1:m
        t1 = s*_raw_getindex(A,j,k) + t*_raw_getindex(A, i, k)
        t2 = u*_raw_getindex(A,j,k) + v*_raw_getindex(A, i, k)
        t1 = ccall((:n_mod2_preinv, :libflint), Int, (Int, Int, UInt), t1, A._n, A._ninv)
        t2 = ccall((:n_mod2_preinv, :libflint), Int, (Int, Int, UInt), t2, A._n, A._ninv)
        ccall((:nmod_mat_set_entry, :libflint), Void, (Ptr{nmod_mat}, Int, Int, UInt), &A, j - 1, k - 1, t1)
        ccall((:nmod_mat_set_entry, :libflint), Void, (Ptr{nmod_mat}, Int, Int, UInt), &A, i - 1, k - 1, t2)
      end
    end
    #println(A)
  end
  #println("I have a triangular matrix:")
  #println(A)
  T = MatrixSpace(base_ring(A), 1, A.c)()
  for j in 1:m
    if _raw_getindex(A,j,j) != 0
      #println("nonzero case")
      u = _unit(_raw_getindex(A,j,j), A._n)
      for k in 1:m
        t1 = mod(u*_raw_getindex(A,j,k), A._n)
        ccall((:nmod_mat_set_entry, :libflint), Void, (Ptr{nmod_mat}, Int, Int, UInt), &A, j - 1, k - 1, t1)
      end
      #println("Multiplied row $j with $u and obtained \n$A\n")
      for i in 1:j-1
        #println("ich nehme zeile $j und reduziere damit zeile $i")
        #println("zeile $i: $(window(A, i:i, 1:m))")
        #println("zeile $j: $(window(A, j:j, 1:m))")
        #println("A[$i,$j]: $(_raw_getindex(A,i,j)) A[$j,$j]: $(_raw_getindex(A,j,j))")
        #println("A[$i,$j]: $(A[i,j]) A[$j,$j]: $(A[j,j])")
        q = UInt(prem(- Int(div(_raw_getindex(A, i, j), _raw_getindex(A, j, j))), Int(A._n)))
        #println("quotient ist $q")
        for k in j:m
          #println("adjusting row $i, divinding by $(_raw_getindex(A, j, j))")
          #println("I am replacing A[$i,$k] with")
          #println("A[$i,$j]: $(_raw_getindex(A,i,j)) A[$j,$j]: $(_raw_getindex(A,j,j))")
          #println("quotient is $(UInt(prem(- Int(div(_raw_getindex(A, i, j), _raw_getindex(A, j, j))), Int(A._n))))")
          #print(_raw_getindex(A, i, k) + UInt(prem(- Int(div(_raw_getindex(A, i, j), _raw_getindex(A, j, j))), Int(A._n))) *_raw_getindex(A, j, k))
          #println()
          _raw_setindex(A, i, k, mod(_raw_getindex(A, i, k) + q *_raw_getindex(A, j, k), A._n))
        end
      end
      #print("computing the annihilator for $j ... ")
      l = mod(div(A._n,gcd(_raw_getindex(A,j,j),A._n)),A._n)
      #println("annihilator is $l")
      if l == 0
        continue
      end
      for k in 1:m
        _raw_setindex(T, 1, k, mod(l*_raw_getindex(A,j,k), A._n))
      end
    else
      for k in 1:m
        _raw_setindex(T, 1, k, _raw_getindex(A,j,k))
      end
    end
    #println("T looks like: $T")
    #println("Augmented matrix looks like \n$A\n$T\n")
    for i in j+1:m
      g,s,t,u,v = _xxgcd(_raw_getindex(A,i,i),_raw_getindex(T,1,i),A._n)
      for k in 1:m
        t1 = s*_raw_getindex(A,i,k) + t*_raw_getindex(T, 1, k)
        t2 = u*_raw_getindex(A,i,k) + v*_raw_getindex(T, 1, k)
        t1 = ccall((:n_mod2_preinv, :libflint), Int, (Int, Int, UInt), t1, A._n, A._ninv)
        t2 = ccall((:n_mod2_preinv, :libflint), Int, (Int, Int, UInt), t2, A._n, A._ninv)
        ccall((:nmod_mat_set_entry, :libflint), Void, (Ptr{nmod_mat}, Int, Int, UInt), &A, i - 1, k - 1, t1)
        ccall((:nmod_mat_set_entry, :libflint), Void, (Ptr{nmod_mat}, Int, Int, UInt), &T, 0, k - 1, t2)
      end
    end
    #println("After reducing with T the augmented matrix looks like \n$A\n$T\n")

  end
  #return A
end

function _raw_setindex(A::nmod_mat, i::Int, j::Int, x::UInt)
  ccall((:nmod_mat_set_entry, :libflint), Void, (Ptr{nmod_mat}, Int, Int, UInt), &A, i - 1, j - 1, x)
end

function _stab(a::UInt, b::UInt, N::UInt)
  g = gcd(a,b,N)
  _, c = ppio(div(N,g),div(a,g))
  #c = prem(c,N)
  return c
end

function _unit(a::UInt, N::UInt)
  g,s,_ = gcdx(a,N)
  if g == 1
    return rem(s,N)
  else
    l = div(N,g)
    d = _stab(s,l,N)
    return rem(s + d*l,N)
  end
end

function _xxgcd(a::UInt, b::UInt, N::UInt)
  g,s,t = gcdx(Int(a),Int(b))
  if g == 0
    return (UInt(0), UInt(1), UInt(0), UInt(0), UInt(1))
  end
  u = UInt(prem(-div(Int(b),g),Int(N)))
  g = UInt(g)
  v = div(a,g)
  return (g,UInt(prem(s,Int(N))),UInt(prem(t,Int(N))),u,v)
end

################################################################################
#
#  positive remainder
#
################################################################################

function prem{T<:Integer}(a::Integer, m::T)
  b = a % m
  if b < 0
    return m+b
  else
    return b
  end
end

function prem{T}(a::fmpz, m::T)
  return prem(BigInt(a), m)
end


text:
function parseKillerJSON(json_sums)
    sums = []
    for s in json_sums
        indices = Tuple[]
        for ind in s["indices"]
            push!(indices, tuple(ind...))
        end

        if haskey(s, "color")
            push!(sums, (result=s["result"], indices=indices, color=s["color"]))
        else
            push!(sums, (result=s["result"], indices=indices, color="white"))
        end
    end
    return sums
end

@testset "Killer Sudoku" begin

@testset "Killer Sudoku from wikipedia" begin
    m = Model(with_optimizer(CS.Optimizer))
    @variable(m, 1 <= x[1:9,1:9] <= 9, Int)

    sums = parseKillerJSON(JSON.parsefile("data/killer_wikipedia"))

    for s in sums
        @constraint(m, sum([x[ind[1],ind[2]] for ind in s.indices]) == s.result)
    end

    # sudoku constraints
    for rc = 1:9
        @constraint(m, x[rc,:] in CS.AllDifferentSet(9))
        @constraint(m, x[:,rc] in CS.AllDifferentSet(9))
    end
    for br=0:2
        for bc=0:2
            @constraint(m, vec(x[br*3+1:(br+1)*3,bc*3+1:(bc+1)*3]) in CS.AllDifferentSet(9))
        end
    end

    optimize!(m)
    @test JuMP.termination_status(m) == MOI.OPTIMAL
    @test jump_fulfills_sudoku_constr(JuMP.value.(x))

    for s in sums
        @test s.result == sum(JuMP.value.([x[i[1],i[2]] for i in s.indices]))
    end
end

@testset "Killer Sudoku niallsudoku_5500 with coefficients" begin
    com = CS.init()

    grid = zeros(Int, (9,9))

    com_grid = Array{CS.Variable, 2}(undef, 9, 9)
    for (ind,val) in enumerate(grid)
        com_grid[ind] = CS.add_var!(com, 1, 9)
    end

    sums = parseKillerJSON(JSON.parsefile("data/killer_niallsudoku_5500"))

    # the upper left sum constraint is x+y+z = 10 and the solution is 2+1+7
    # here I change it to 5*x+7*y+z = 24
    CS.add_constraint!(com, 5*com_grid[CartesianIndex(1,1)]+7*com_grid[CartesianIndex(2,1)]+com_grid[CartesianIndex(2,2)] == 24)

    for s in sums[2:end]
        CS.add_constraint!(com, sum([com_grid[CartesianIndex(ind)] for ind in s.indices]) == s.result)
    end

    add_sudoku_constr!(com, com_grid)

    options = Dict{Symbol, Any}()
    options[:keep_logs] = true

    options = CS.combine_options(options)

    @test CS.solve!(com, options) == :Solved
    logs_1 = CS.get_logs(com)
    info_1 = com.info
    @test fulfills_sudoku_constr(com_grid)
    @test 5*CS.value(com_grid[CartesianIndex(1,1)])+7*CS.value(com_grid[CartesianIndex(2,1)])+CS.value(com_grid[CartesianIndex(2,2)]) == 24
    for s in sums[2:end]
        @test s.result == sum([CS.value(com_grid[CartesianIndex(i)]) for i in s.indices])
    end

    # test if deterministic by running it again
    com = CS.init()

    grid = zeros(Int, (9,9))

    com_grid = Array{CS.Variable, 2}(undef, 9, 9)
    for (ind,val) in enumerate(grid)
        com_grid[ind] = CS.add_var!(com, 1, 9)
    end

    sums = parseKillerJSON(JSON.parsefile("data/killer_niallsudoku_5500"))

    # the upper left sum constraint is x+y+z = 10 and the solution is 2+1+7
    # here I change it to 5*x+7*y+z = 24
    CS.add_constraint!(com, 5*com_grid[CartesianIndex(1,1)]+7*com_grid[CartesianIndex(2,1)]+com_grid[CartesianIndex(2,2)] == 24)

    for s in sums[2:end]
        CS.add_constraint!(com, sum([com_grid[CartesianIndex(ind)] for ind in s.indices]) == s.result)
    end

    add_sudoku_constr!(com, com_grid)

    options = Dict{Symbol, Any}()
    options[:keep_logs] = true

    options = CS.combine_options(options)
    status = CS.solve!(com, options)

    logs_2 = CS.get_logs(com)
    info_2 = com.info
    @test info_1.pre_backtrack_calls == info_2.pre_backtrack_calls
    @test info_1.backtrack_fixes == info_2.backtrack_fixes
    @test info_1.in_backtrack_calls == info_2.in_backtrack_calls
    @test info_1.backtrack_reverses == info_2.backtrack_reverses
    @test CS.same_logs(logs_1[:tree], logs_2[:tree])
end

function killer_negative()
    m = Model(with_optimizer(CS.Optimizer, keep_logs=true))

    @variable(m, -9 <= com_grid[1:9, 1:9] <= -1, Int)

    sums = parseKillerJSON(JSON.parsefile("data/killer_niallsudoku_5500"))

    # the upper left sum constraint is x+y+z = 10 and the solution is -2 + -1 + -7
    # here I change it to 5*x-7*y-z = 4
    @constraint(m, 5*com_grid[1,1]-7*com_grid[2,1]-com_grid[2,2] == 4)

    for s in sums[2:end]
        @constraint(m, sum([com_grid[ind...] for ind in s.indices]) == -s.result)
    end

    jump_add_sudoku_constr!(m, com_grid)
    optimize!(m)

    @test JuMP.termination_status(m) == MOI.OPTIMAL
    @test jump_fulfills_sudoku_constr(com_grid)
    @test 5*JuMP.value(com_grid[1,1])-7*JuMP.value(com_grid[2,1])-JuMP.value(com_grid[2,2]) == 4
    for s in sums[2:end]
        @test -s.result == sum([JuMP.value(com_grid[i...]) for i in s.indices])
    end
    return JuMP.backend(m).optimizer.model.inner
end

@testset "Killer Sudoku niallsudoku_5500 with negative coefficients and -9 to -1" begin
    com1 = killer_negative()
    com2 = killer_negative()
    info_1 = com1.info
    info_2 = com2.info
    @test info_1.pre_backtrack_calls == info_2.pre_backtrack_calls
    @test info_1.backtrack_fixes == info_2.backtrack_fixes
    @test info_1.in_backtrack_calls == info_2.in_backtrack_calls
    @test info_1.backtrack_reverses == info_2.backtrack_reverses
    logs_1 = CS.get_logs(com1)
    logs_2 = CS.get_logs(com2)
    @test CS.same_logs(logs_1[:tree], logs_2[:tree])
end

end


text:
include("counter.jl")
include("grid.jl")
include("cubic_interpolation.jl")
include("coord.jl")


text:
include("nodes_tests.jl")
include("binaryheaps_tests.jl")
include("priorityqueues_tests.jl")


text:
using Pipe: @pipe

mutable struct Coordinate
    x::Int
    y::Int
end

struct Vent
    from::Coordinate
    to::Coordinate
end

function input()
    s = read(stdin, String)
    return @pipe (split(s, "\n", keepempty = false)
        |> map( l -> match(r"(\d+),(\d+) -> (\d+),(\d+)", l), _ )
        |> map( m -> Vent( Coordinate(Base.parse(Int, m.captures[1]), Base.parse(Int, m.captures[2])), Coordinate(Base.parse(Int, m.captures[3]), Base.parse(Int, m.captures[4])) ), _)
    )
end
function expand_vent(v::Vent)
    result = Coordinate[]
    inc = Coordinate( v.from.x == v.to.x ? 0 : (v.to.x - v.from.x) / abs(v.to.x - v.from.x), v.from.y == v.to.y ? 0 : (v.to.y - v.from.y) / abs(v.to.y - v.from.y) )
    cursor = v.from
    while cursor.x != v.to.x || cursor.y != v.to.y
        push!(result, Coordinate(cursor.x, cursor.y))
        cursor.x += inc.x
        cursor.y += inc.y
    end
    push!(result, v.to)
    return result
end
@pipe (input()
    |> map( expand_vent, _ )
    |> Iterators.flatten( _ )
    |> collect( _ )
    |> sort( _, lt = (a,b) -> a.x < b.x || (a.x == b.x && a.y < b.y) )
    |> map( c -> (1, c), _)
    |> reduce( function(acc, c)
        if length(acc) > 0 && acc[1][2].x == c[2].x && acc[1][2].y == c[2].y
            acc[1] = (acc[1][1] + c[1], acc[1][2])
        else
            pushfirst!(acc, c)
        end
        return acc
    end, _, init=[])
    |> filter( t -> t[1] > 1, _)
    |> length(_)
    |> println(_) )


text:
should_precompile = true

# Don't edit the following! Instead change the script for `snoop_bot`.
ismultios = false
ismultiversion = false
# precompile_enclosure
@static if !should_precompile
    # nothing
elseif !ismultios && !ismultiversion
    @static if isfile(joinpath(
        @__DIR__,
        "../deps/SnoopCompile/precompile/precompile_Plots.jl",
    ))
        include("../deps/SnoopCompile/precompile/precompile_Plots.jl")
        _precompile_()
    end
else
end # precompile_enclosure


text:
"""
```
augment_variables!(m::Li2020, stategrid::StateGrid, funcvar::OrderedDict{Symbol, Vector{S}},
    endo::OrderedDict{Symbol, Vector{S}}) where {S <: Real}

augment_variables_nojump!(m::Li2020, stategrid::StateGrid, ode_f::Function, funcvar::OrderedDict{Symbol, Vector{S}},
    endo::OrderedDict{Symbol, Vector{S}}, odesol::ODESolution) where {S <: Real}

augment_variables_nojump!(m::Li2020, stategrid::StateGrid, funcvar::OrderedDict{Symbol, Vector{S}},
    endo::OrderedDict{Symbol, Vector{S}}, odesol::ODESolution) where {S <: Real}
```

calculates derivatives and additional endogenous variables that can be computed
after solving the equilibrium system of functional equations.
"""
function augment_variables!(m::Li2020, stategrid::StateGrid, funcvar::OrderedDict{Symbol, Vector{S}},
                            derivs::OrderedDict{Symbol, Vector{S}},
                            endo::OrderedDict{Symbol, Vector{S}}) where {S <: Real}

    # Unpack equilibrium endogenous variables and stategrid
    @unpack ψ, xK, yK, yg, σp, σ, σh, σw, μR_rd, rd_rg, rd_rg_H, μb_μh, μw, μp, μK, μR, rd, rd_rf, rg, μb, μh, invst, lvg, κp, κb, κd, κh, κfs, firesale_jump, κw, liq_prem, bank_liq_frac, δ_x, indic, rf, rh, K_growth, κK = endo
    @unpack p, xg, Q̂     = funcvar
    @unpack ∂p_∂w, ∂²p_∂w² = derivs
    w    = stategrid[:w]
    θ    = parameters_to_named_tuple(get_parameters(m))
    f_μK = get_setting(m, :μK)
    Φ    = get_setting(m, :Φ)
    yg_tol = get_setting(m, :yg_tol)
    firesale_bound = get_setting(m, :firesale_bound)
    firesale_interpolant = get_setting(m, :firesale_interpolant)

    # Calculate various quantities
    ∂p_∂w   .= differentiate(w, p)    # as Li (2020) does it
    ∂²p_∂w² .= differentiate(w, ∂p_∂w) # as Li (2020) does it
    invst   .= map(x -> Φ(x, θ[:χ], θ[:δ]), p)
    ψ       .= (invst .+ θ[:ρ] .* (p + Q̂) .- θ[:AL]) ./ (θ[:AH] - θ[:AL])
    ψ[ψ .> 1.] .= 1.

    # Portfolio choices
    xK   .= (ψ ./ w) .* (p ./ (p + Q̂))
    xK[1] = xK[2]
    lvg  .= xK + xg
    yK   .= (p ./ (p + Q̂) - w .* xK)  ./ (1. .- w)
    yK[1] = 1.
    yK[end] = yK[end - 1]
    Q     = get_setting(m, :avg_gdp) * get_setting(m, :gov_bond_gdp_level)
    yg   .= (Q ./ (p + Q̂) - w .* xg) ./ (1. .- w)
    yg[end] = yg[end - 1]
    yg[yg .< yg_tol] .= yg_tol # avoid yg < 0
    δ_x   .= max.(θ[:β] .* (xK + xg .- 1.) - xg, 0.)
    indic .= δ_x .> 0.

    # Volatilities
    σp .= ∂p_∂w .* w .* (1. .- w) .* (xK - yK) ./ (1. .- ∂p_∂w .* w .* (1. .- w) .* (xK - yK)) .* θ[:σK]
    σp[end] = 0. # no volatility at the end
    σ  .= xK .* (θ[:σK] .+ σp)
    σh .= yK .* (θ[:σK] .+ σp)
    σw .= (1. .- w) .* (σ - σh)

    # Deal with post jump issues
    firesale_jump .= xK .* κp + (θ[:α] / (1 - θ[:α])) .* δ_x
    firesale_ind   = firesale_jump .< firesale_bound
    firesale_spl   = extrapolate(interpolate((w[firesale_ind], ), firesale_jump[firesale_ind], firesale_interpolant), Line())
    firesale_jump .= firesale_spl(w)

    # Jumps
    κd  .= θ[:θ] .* (xK .+ θ[:ϵ] .- 1.) ./ (xK + xg .- 1.) .* (xK .+ θ[:ϵ] .> 1.)
    κb  .= θ[:θ] .* min.(1. - θ[:ϵ], xK) + (1. .- θ[:θ]) .* firesale_jump
    κfs .= (θ[:α] / (1 - θ[:α])) .* δ_x .* w ./ (1. .- w)
    κfs[end] = 0.
    κh  .= yK .* κp + (1. .- yK .- yg) .* κd - κfs
    κw  .= 1. .- (1. .- κb) ./ (1. .- κh .- w .* (κb - κh))
    endo[:κp][end] = 0. # Enforce zero at the end

    # Generate liquidity premium and price of risks
    liq_prem .= indic .* θ[:λ] .* ((1 - θ[:π]) * (1 - θ[:θ]) * θ[:α] / (1 - θ[:α])) ./ (1. .- firesale_jump)
    index     = convert(Int, max(argmax(liq_prem) - 6, 1))
    fit       = SLM(w[index:end], liq_prem[index:end], decreasing = true, right_value = 0., knots = 4)
    liq_prem .= vcat(liq_prem[1:index - 1], eval(fit, w[index:end]))

    # Liquidity holding across states
    bank_liq_frac .= xg .* w ./ (Q ./ (p + Q̂))

    # Main drifts and interest rates
    μR_rd   .= (θ[:σK] .+ σp) .^ 2 .* xK - θ[:AH] ./ p + (θ[:λ] * (1 - θ[:θ])) .*
        (κp + indic .* (θ[:α] / (1 - θ[:α]) * θ[:β])) ./ (1. .- firesale_jump) +
         ((xK .+ θ[:ϵ]) .< 1.) .* (θ[:λ] * θ[:θ]) ./ (1. .- xK)
    rd_rg   .= (θ[:λ] * (1 - θ[:θ]) * θ[:α] / (1 - θ[:α]) * (1 - θ[:β])) .* indic ./ (1. .- firesale_jump)
    rd_rg_H .= θ[:λ] .* κd ./ (1. .- yK .* κp - (1. .- yK .- yg) .* κd + κfs)
    index_xK = xK .<= 1. # In this scenario, the rd-rg difference must be solved from hh's FOC
    rd_rg[index_xK] = rd_rg_H[index_xK]
    μb_μh   .= (xK - yK) .* μR_rd + (xK .* θ[:AH] - yK .* θ[:AL]) ./ p - (xg - yg) .* rd_rg
    μw      .= (1. .- w) .* (μb_μh + σh .^ 2 - σ .* σh - w .* (σ - σh) .^ 2 -
                       θ[:η] ./ (1. .- w))
    μw[end]   = μw[end - 1] # drift should be similar at the end
    μp       .= ∂p_∂w .* w .* μw + (1. / 2.) .* ∂²p_∂w² .* (w .* (1. .- w) .* (σ - σh)) .^ 2
    μp[end]   = μp[end - 1] # drift should be similar at the end
    μK       .= map(x -> f_μK(x, θ[:χ], θ[:δ]), p)
    μR       .= μp .- θ[:δ] + μK + θ[:σK] .* σp - invst ./ p
    K_growth .= μK .- θ[:δ]
    κK       .= θ[:θ] .* ψ

    # Other interest rates
    rd    .= μR - μR_rd
    rg    .= rd - rd_rg
    rd_rf .= (θ[:λ] * (1. - θ[:θ]) * θ[:α] / (1 - θ[:α]) * (-θ[:β])) .* indic ./ (1. .- firesale_jump)
    rf    .= rd - rd_rf
    rh    .= rd - θ[:λ] .* κd ./ (1. .- yK .* κp - (1. .- yK .- yg) .* κd .+ κfs) # risk free rate for households

    # Other growth rates
    μb .= xK .* (μR + θ[:AH] ./ p - rd) - xg .* rd_rg + rd .- θ[:ρ]
    μh .= μb - μb_μh
end

"""
```
augment_variables_nojump!(m::Li2020, stategrid::StateGrid, ode_f::Function, funcvar::OrderedDict{Symbol, Vector{S}},
    endo::OrderedDict{Symbol, Vector{S}}, odesol::ODESolution) where {S <: Real}
```

calculates derivatives and additional endogenous variables that can be computed
after solving the equilibrium system of functional equations. The first derivative
of `p` is directly calculated using `ode_f`
"""
function augment_variables_nojump!(m::Li2020, stategrid::StateGrid, ode_f::Function, funcvar::OrderedDict{Symbol, Vector{S}},
                                   derivs::OrderedDict{Symbol, Vector{S}},
                                   endo::OrderedDict{Symbol, Vector{S}}, odesol::ODESolution) where {S <: Real}

    # Unpack equilibrium endogenous variables and stategrid
    p    = funcvar[:p]
    xg   = funcvar[:xg]
    derivs[:∂p_∂w] = similar(p)
    ∂p_∂w = derivs[:∂p_∂w]
    w    = stategrid[:w]
    @unpack ψ, xK, yK, yg, σp, σ, σh, μR_rd, μb_μh, μw = endo
    Q = get_setting(m, :avg_gdp) * get_setting(m, :gov_bond_gdp_level)

    # Interpolate remainder of the solution for q from odesol
    i = findfirst(stategrid[:w] .>= odesol.t[end])
    p[1:i] .= odesol(stategrid[:w][1:i])
    p[(i + 1):end] .= boundary_conditions(m)[:p][2]
    θ = parameters_to_named_tuple(map(x -> get_parameters(m)[get_keys(m)[x]], get_setting(m, :nojump_parameters)))
    ∂p_∂w[1:i] .= map(j -> ode_f(p[j], θ, stategrid[:w][j]), 1:i)
    ∂p_∂w[(i + 1):end] .= 0.

    # Solve for values not calculated during the calculation of equilibrium
    Φ = get_setting(m, :Φ)
    ψ[1:i] .= map(x -> max.(min.((Φ(x, θ[:χ], θ[:δ]) + θ[:ρ] * x - θ[:AL]) ./ (θ[:AH] - θ[:AL]), 1.), 0.), p[1:i])
    ψ[(i + 1):end] .= 1.
    xK .= (ψ ./ w)
    yK .= (1. .- ψ) ./ (1. .- w)
    xg .= zeros(eltype(w), length(w))
    yg .= fill(get_setting(m, :avg_gdp) * get_setting(m, :gov_bond_gdp_level), length(w)) ./ p ./ (1. .- w)
    ψ_is1 = (i + 1):length(w)
    ψ_no1 = 2:i
    σp[ψ_no1] .= sqrt.((θ[:AH] - θ[:AL]) ./ (p[ψ_no1] .* (xK[ψ_no1] - yK[ψ_no1]))) .- θ[:σK]
    σp[ψ_is1] .= θ[:σK] .* ∂p_∂w[ψ_is1] .* w[ψ_is1] .* (1. .- w[ψ_is1]) .*
        (xK[ψ_is1] - yK[ψ_is1]) ./ (1. .- ∂p_∂w[ψ_is1] .* w[ψ_is1] .* (1. .- w[ψ_is1]) .* (xK[ψ_is1] - yK[ψ_is1]))
    σ  .= xK .* (θ[:σK] .+ σp)
    σh .= yK .* (θ[:σK] .+ σp)

    if 0. in w
        ∂Φ        = get_setting(m, :∂Φ)
        w0        = findfirst(w .== 0.)
        p[w0 ]    = get_setting(m, :boundary_conditions)[:p][1]
        d0        = (θ[:AH] - θ[:AL]) / (∂Φ(p[w0], θ[:χ]) + θ[:ρ]) * ((θ[:AH] - θ[:AL]) / (p[w0] * θ[:σK])^2 * p[w0] + 1)
        ψ[w0]     = 0.
        ∂p_∂w[w0] = d0
        xK[w0]    = 0.
        yK[w0]    = 1.
        yg[w0]    = Q[w0] / p[w0]
        σp[w0]    = 0.
        σ[w0]     = 0.
        σh[w0]    = yK[w0] * θ[:σK]
    end

    if 1. in w
        wN        = findfirst(w .== 1.)
        ∂p_∂w[wN] = 0.
        xK[wN]    = 1.
        yK[wN]    = 0.
        xg[wN]    = 0.
        yg[wN]    = 1 / p[wN] * Q / (1 - w[wN - 1])
        σp[wN]    = 0.
        σ[wN]     = xK[wN] * θ[:σK]
        σh[wN]    = 0.
    end

    μR_rd  .= (θ[:σK] .+ σp) .^ 2 .* xK - θ[:AH] ./ p
    μb_μh  .= (xK - yK) .* μR_rd + (xK .* θ[:AH] - yK .* θ[:AL]) ./ p
    μw     .= (1. .- w) .* (μb_μh + σh .^ 2 - σ .* σh - w .* (σ - σh) .^ 2 -
                       θ[:η] ./ (1. .- w))
    μw[end] = μw[end - 1] # drift should be similar at the end
end

"""
```
augment_variables_nojump!(m::Li2020, stategrid::StateGrid, funcvar::OrderedDict{Symbol, Vector{S}},
    endo::OrderedDict{Symbol, Vector{S}}, odesol::ODESolution) where {S <: Real}
```

calculates derivatives and additional endogenous variables that can be computed
after solving the equilibrium system of functional equations and using
finite differences to approximate the derivatives.
"""
function augment_variables_nojump!(m::Li2020, stategrid::StateGrid, funcvar::OrderedDict{Symbol, Vector{S}},
                                   derivs::OrderedDict{Symbol, Vector{S}},
                                   endo::OrderedDict{Symbol, Vector{S}}, odesol::ODESolution) where {S <: Real}

    # Unpack equilibrium endogenous variables and stategrid
    @unpack p, xg, Q̂ = funcvar
    Q̂  .= zero(eltype(m))
    xg .= zero(eltype(m))
    derivs[:∂p_∂w] = similar(p)
    ∂p_∂w   = derivs[:∂p_∂w]
    ∂²p_∂w² = derivs[:∂²p_∂w²]
    @unpack ψ, xK, yK, yg, σp, σ, σh, μR_rd, rd_rg, μb_μh, μw, μp, μK, μR, rd, rg, μb, μh = endo
    w      = stategrid[:w]
    θ = parameters_to_named_tuple(map(x -> get_parameters(m)[get_keys(m)[x]], get_setting(m, :nojump_parameters)))
    f_μK = get_setting(m, :μK)

    # Interpolate remainder of the solution for q from odesol
    i = findfirst(w .>= odesol.t[end])
    p[1:i] .= odesol(w[1:i])
    p[(i + 1):end] .= boundary_conditions(m)[:p][2]
    ∂p_∂w   .= differentiate(w, p)    # as Li (2020) does it
    ∂p²_∂w² .= differentiate(w, ∂p_∂w) # as Li (2020) does it

    # Solve for values not calculated during the calculation of equilibrium
    Φ = get_setting(m, :Φ)
    ψ[1:i] .= map(x -> max.(min.((Φ(x, θ[:χ], θ[:δ]) + θ[:ρ] * x - θ[:AL]) ./ (θ[:AH] - θ[:AL]), 1.), 0.), p[1:i])
    ψ[(i + 1):end] .= 1.
    xK .= (ψ ./ w)
    yK .= (1. .- ψ) ./ (1. .- w)
    yg .= fill(get_setting(m, :avg_gdp) * get_setting(m, :gov_bond_gdp_level), length(w)) ./ p ./ (1. .- w)
    yg[end] = yg[end - 1]
    σp .= ∂p_∂w .* w .* (1. .- w) .* (xK - yK) ./ (1. .- ∂w .* w .* (1. .- w) .* (xK - yK)) .* θ[:σK]
    σp[end] = 0. # no volatility at the end
    σ  .= xK .* (θ[:σK] .+ σp)
    σh .= yK .* (θ[:σK] .+ σp)

    # Calculate remaining objects
    invst   = map(x -> Φ(x, θ[:χ], θ[:δ]), p)
    μR_rd  .= (θ[:σK] .+ σp) .^ 2 .* xK - θ[:AH] ./ p
    rd_rg  .= 0.
    μb_μh  .= (xK - yK) .* μR_rd + (xK .* θ[:AH] - yK .* θ[:AL]) ./ p - (xg - yg) .* rd_rg
    μw     .= (1. .- w) .* (μb_μh + σh .^ 2 - σ .* σh - w .* (σ - σh) .^ 2 -
                       θ[:η] ./ (1. .- w))
    μw[end] = μw[end - 1] # drift should be similar at the end
    μp     .= ∂p_∂w .* w .* μw + (1. / 2.) .* ∂p²_∂w² .* (w .* (1. .- w) .* (σ - σh)) .^ 2
    μp[end] = μp[end - 1] # drift should be similar at the end
    μK     .= map(x -> f_μK(x, θ[:χ], θ[:δ]), p)
    μR     .= μp .- θ[:δ] + μK + θ[:σK] .* σp - invst ./ p
    rd     .= μR - μR_rd
    rg      = copy(rd)
    μb     .= xK .* (μR + θ[:AH] ./ p - rd) - xg .* rd_rg + rd .- θ[:ρ]
    μh     .= μb - μb_μh
end


text:
include("solver_preamble.jl")
import SCS
factory = with_optimizer(SCS.Optimizer, verbose=0)
config = MOI.Test.TestConfig(atol=5e-4, rtol=5e-4)
@testset "Linear" begin
    Tests.linear_test(factory, config)
end
@testset "SOC" begin
    Tests.soc_test(factory, config, [
        "sdsos_options_pricing"
    ])
end
@testset "SDP" begin
    Tests.sd_test(factory, config, [
        "sos_options_pricing",
        # Expression: termination_status(model) == MOI.INFEASIBLE
        # Evaluated: MathOptInterface.OPTIMAL == MathOptInterface.INFEASIBLE
        "sos_horn"
    ])
end


text:
#printing that is not needed in Atom
import Base: show

show(io::IO, block::InputData) = df_show(io, block)

show(io::IO, band::Band) = df_show(io, band)

show(io::IO, job::Job) = df_show(io, job)
show(io::IO, c::Calculation) = df_show(io, c)
show(io::IO, flag_info::Calculations.QEFlagInfo) = df_show(io, flag_info)
show(io::IO, flag_info::Calculations.ElkFlagInfo) = df_show(io, flag_info)
show(io::IO, info::Calculations.ElkControlBlockInfo) = df_show(io, info)
show(io::IO, el::Element) = df_show(io, el)
show(io::IO, str::Structure) = df_show(io, str)
show(io::IO, at::Atom) = df_show(io, at)
show(io::IO, proj::Projection) = df_show(io, proj)


text:
module TestTOML

using TOML
using JSON
using Compat

function parsedate(str::AbstractString) #for the tests.
    d = match(TOML.date_pattern, str)
    if d == nothing ;; error("Invalid date string.") end
    TOML.DateTime(map(s -> parse(Int, s), d.captures)...)
end

testdir = Base.source_path()[1:end-8]

list = readdir("$testdir/valid")
valid = map(filter(a-> a[end-4:end]==".toml", list)) do s
    s[1:end-5]
end

invalid = readdir("$testdir//invalid")

tml = ""
jsn = ""

if isa(STDIN,Base.TTY)
    eval(Base,:(have_color = $(contains(get(ENV,"TERM",""), "xterm"))))
end

exitstatus = 0

function display_error(er)
    Base.with_output_color(:red, STDOUT) do io
        println(io, "ERROR: ", er)
    end
end

function display_success(msg)
    Base.with_output_color(:green, STDOUT) do io
        println(io, msg)
    end
end

# valid = filter(valid) do n ;; match(r"table-array", n) == nothing end

function test()

    # Inline test for line endings
    # (since git can mess it up in files we use strings)
    print("Valid crlf: ")
    try
        tml = TOML.parse("foo = 5\r\nbar = 6")
        jsn = jsn2data(JSON.parse("{ \"foo\": {\"type\":\"integer\", \"value\":\"5\"}, \"bar\":{\"type\":\"integer\", \"value\":\"6\"} }"))
        if jsn == tml
            display_success("Ok")
        else
            display_error("unexpected result.\nTOML:\n$tml\n\nJSON:\n$jsn\n")
            global exitstatus = 1
        end
    catch err
        if !isa(err, TOML.TOMLError)
            rethrow(err)
        end
        display_error("couldn't be parsed.\n" * repr(err))
        global exitstatus = 1
    end

    print("Valid lf: ")
    try
        tml = TOML.parse("foo = 5\nbar = 6")
        jsn = jsn2data(JSON.parse("{ \"foo\": {\"type\":\"integer\", \"value\":\"5\"}, \"bar\":{\"type\":\"integer\", \"value\":\"6\"} }"))
        if jsn == tml
            display_success("Ok")
        else
            display_error("unexpected result.\nTOML:\n$tml\n\nJSON:\n$jsn\n")
            global exitstatus = 1
        end
    catch err
        if !isa(err, TOML.TOMLError)
            rethrow(err)
        end
        display_error("couldn't be parsed.\n" * repr(err))
        global exitstatus = 1
    end

    for t in valid
        print("Valid $t: ")
            open(string(testdir, "/valid/", t, ".toml")) do tml
                open(string(testdir, "/valid/", t, ".json")) do jsn
                try
                    tml = TOML.parse(readall(tml))
                    jsn = jsn2data(JSON.parse(readall(jsn)))
                    if jsn == tml
                        display_success("Ok")
                    else
                        display_error("unexpected result.\nTOML:\n$tml\n\nJSON:\n$jsn\n")
                        global exitstatus = 1
                    end
                catch err
                    if !isa(err, TOML.TOMLError)
                        rethrow(err)
                    end
                    display_error("couldn't be parsed.\n" * repr(err))
                    global exitstatus = 1
                end
            end
        end
    end

    println("Count \\n end of lines:")
    try
        tml = TOML.parse("\ninvalid")
        display_error("should have failed but didn't.")
        print(tml)
        global exitstatus = 1
    catch err
        if !isa(err, TOML.TOMLError)
            rethrow(err)
        end
        if err.msg == "Unexpected end of file on line 2."
            display_success("  " * repr(err))
        else
            display_error("  " * repr(err) * "\nexpected\n  " * "Unexpected end of file on line 2.")
        end
    end

    println("Count \\r end of lines:")
    try
        tml = TOML.parse("\rinvalid")
        display_error("should have failed but didn't.")
        print(tml)
        global exitstatus = 1
    catch err
        if !isa(err, TOML.TOMLError)
            rethrow(err)
        end
        if err.msg == "Unexpected end of file on line 2."
            display_success("  " * repr(err))
        else
            display_error(repr(err) * "\n  expected\n" * "TOMLError: Unexpected end of file on line 2.")
        end
    end

    println("Count \\r\\n end of lines:")
    try
        tml = TOML.parse("\r\ninvalid")
        display_error("should have failed but didn't.")
        print(tml)
        global exitstatus = 1
    catch err
        if !isa(err, TOML.TOMLError)
            rethrow(err)
        end
        if err.msg == "Unexpected end of file on line 2."
            display_success("  " * repr(err))
        else
            display_error("  " * repr(err) * "\nexpected\n  " * "Unexpected end of file on line 2.")
        end
    end

    for t in invalid
        println("Invalid $t:")
        open(string(testdir, "/invalid/", t)) do tml
            try
                tml = TOML.parse(readall(tml))
                display_error("should have failed but didn't.")
                print(tml)
                global exitstatus = 1
            catch err
                if !isa(err, TOML.TOMLError)
                    rethrow(err)
                end
                display_success("  " * repr(err))
            end
        end
    end
end

jsnval = @compat Dict{ASCIIString,Function}(
    "string" =>identity,
    "float"  => (s -> parse(Float64, s)),
    "integer"=> (s -> parse(Int, s)),
    "datetime"   => parsedate,
    "array"  => (a -> map(jsn2data, a)),
    "bool"   => (b -> b == "true")
)

function jsn2data(jsn)
    # println(jsn)
    if "type" in keys(jsn)
        jsnval[jsn["type"]](jsn["value"])
    else
        @compat Dict{Any,Any}([k => jsn2data(v) for (k, v) in jsn])
    end
end

end # module TestTOML

TestTOML.test()

exit(TestTOML.exitstatus)


text:
@testset "Bistable System" begin

include("testutils.jl")

@testset "Time-series" begin

	de1, de2 = [ODEtime(func=bistable_ode!,
				u0 = u,
				p = p_,
				tspan=(0.,4.)) for u in [u_1, u_2]]

	sol1, sol2 = solve.([de1, de2])

	@test sol1[end].s1 > sol1[end].s2
	@test sol2[end].s2 > sol2[end].s1

	# Plotting
	#p1, p2 = plot.([sol1, sol2])
	#pp_ = plot(p1, p2)
	#plot!(pp_, title="Bistable System")
	#display(pp_)
end

@testset "Calculate Steady State from given initial state" begin
	de1, de2 = [ DEsteady(func = bistable_ode!,
							p=p_,
							u0=u,
							method= Tsit5())
							for u in [u_1, u_2] ]

	sol1, sol2 = solve.([de1, de2])

	@test sol1.u.s1 > sol1.u.s2
	@test sol2.u.s2 > sol2.u.s1
end

@testset "Find steady state with given region: Single core computation" begin

	param_gen = ParameterGrid([(0.1,5.,10), (0.1,5.,10)])

	de = DEsteady(func=bistable_ode!, p=p_, u0= u_1, method=Tsit5())

	u_ = de.u0

	sol = solve(de) # create `sol` struct
	sols = Vector{Any}(undef,length(param_gen)) # allocate the memory

	sols = map(param_gen) do u
		u_ .= u
		de_ = de(u_)
		return solve(de_)
	end # there is no significantly difference between using for loop and map.

	for s in sols
		println(s.u)
	end

end

@info "Grid search"

@testset "Type renew" begin
	param_gen = ParameterGrid([(0.,5.,100), (0.,5.,1000)])

	de = DEsteady(func=bistable_ode!, p=p_, u0= u_1, method=SSRootfind())
	de_new = de([100.0,200.0])

	@time sols = map(param_gen) do u
		de_i = de(u)
		return solve(de_i)
	end;

	@test length(sols) == length(param_gen)
	@test de_new.u0 != de.u0
end

@testset "Find steady state with given region: Multi-threading" begin

	param_gen = ParameterGrid([(0.1,5.,100), (0.1,5.,100)])

	de = DEsteady(func=bistable_ode!, p=p_, u0= u_1, method=SSRootfind())

	@time sols = solve(de, param_gen)

	@test length(sols) == length(param_gen)

end

@info "Random Search"

@testset "Random search: multi-threading ($(Threads.nthreads()))" begin

	param_gen = ParameterRandom([(0,1), (0,10)], Uniform(), 10000)

	de = DEsteady(func=bistable_ode!, p=p_, u0= u_1, method=SSRootfind())

	@time sols = solve(de, param_gen; ensemble_method=EnsembleThreads())

	@test length(sols) == length(param_gen)

end

end


text:
module MathOptCartesians

import MathOptInterface
const MOI = MathOptInterface

include("scalar_set.jl")
include("vector_set.jl")

end # module


text:
# Aliases (types)
const BitVector        = BitArray{1};
const BitMatrix        = BitArray{2};
const CoordinatesVector = Array{CartesianIndex{2},1};
const VectorsArray{T}   = Array{Vector{T},1}; # TBD: change to ArrayOfVectors

# Estimation structures
abstract type EstimSettings end

# Constructor for Γ
function build_vec_Γ(n::Int64, p::Int64, β::Float64)

    vec_Γ = Array{Float64,1}();
    for i=0:p-1
        for j=1:n
            push!(vec_Γ, β^i);
        end
    end

    return vec_Γ;
end

function build_Γ(n::Int64, p::Int64, λ::Float64, β::Float64)
    return Diagonal(λ.*build_vec_Γ(n, p, β))::DiagMatrix;
end

"""
    SmootherArrays(...)

Internal re-usable arrays for the Kalman smoother output.
"""
struct SmootherArrays
    J1::FloatVector
    J2::SymMatrix
    Xs_leading::FloatVector
    Ps_leading::SymMatrix
    F::Union{FloatMatrix, Nothing}
    G::Union{FloatMatrix, Nothing}
    H::Union{FloatMatrix, Nothing}
    M::Union{FloatMatrix, Nothing}
    N::Union{Vector{SparseMatrixCSC{Float64, Int64}}, Nothing}
    O::Union{Array{VectorsArray{Float64},1}, Nothing}
    buffer_M::Union{FloatVector, Nothing}
    buffer_O::Union{SymMatrix, Nothing}
end

# Constructor for SmootherArrays
function SmootherArrays(estim::EstimSettings, sspace::KalmanSettings, coordinates_measurement_states::Union{IntVector, Nothing}, coordinates_transition_current::Union{IntVector, Nothing}, coordinates_transition_lagged::Union{IntVector, Nothing}, coordinates_transition_PPs::Union{IntVector, Nothing})

    # Kalman smoother internal arrays
    J1 = zeros(sspace.m);
    J2 = Symmetric(zeros(sspace.m, sspace.m));
    Xs_leading = zeros(sspace.m);
    Ps_leading = Symmetric(zeros(sspace.m, sspace.m));

    # ECM statistics for the transition equation
    F, G, H = initialise_ecm_stats_transition(coordinates_transition_current, coordinates_transition_lagged, coordinates_transition_PPs);

    # ECM statistics for the measurement equation
    M, N, O, buffer_M, buffer_O = initialise_ecm_stats_measurement(estim, coordinates_measurement_states);

    # Return structure
    return SmootherArrays(J1, J2, Xs_leading, Ps_leading, F, G, H, M, N, O, buffer_M, buffer_O);
end

# Validation types

"""
    ValidationSettings(...)

Define an immutable structure used to define the validation settings.

The arguments are two dimensional arrays representing the bounds of the grid for each hyperparameter.

# Arguments
- `err_type`:
    - 1 In-sample error
    - 2 Out-of-sample error
    - 3 Block jackknife error
    - 4 Artificial jackknife error
- `Y`: observed measurements (`nxT`)
- `n`: Number of series
- `T`: Number of observations
- `is_stationary`: Boolean value
- `model_struct`: DataType identifying the estimation structure to use
- `model_args`: Tuple with the arguments required to setup the model specified in `model_struct` (irrelevant for VARs and VMAs)
- `model_kwargs`: Tuple with the keyword arguments required to setup the model specified in `model_struct` (default: nothing)
- `coordinates_params_rescaling`: Array of vectors including information on the parameters (if any) that require to be rescaled to match the data standardisation (default: nothing)
- `verb`: Verbose output (default: true)
- `verb_estim`: Further verbose output (default: false)
- `weights`: Weights for the forecast error. standardise_error has priority over weights. (default: ones(n))
- `t0`: weight associated to the LASSO component of the elastic-net penalty
- `subsample`: number of observations removed in the subsampling process, as a percentage of the original sample size. It is bounded between 0 and 1.
- `max_samples`: if `C(n*T,d)` is large, artificial_jackknife would generate `max_samples` jackknife samples. (used only for the artificial jackknife)
- `log_folder_path`: folder to store the log file. When this file is defined then the stdout is redirected to this file.
"""
struct ValidationSettings
    err_type::Int64
    Y::Union{FloatMatrix, JMatrix{Float64}}
    n::Int64
    T::Int64
    is_stationary::Bool
    model_struct::DataType
    model_args::Tuple
    model_kwargs::Union{NamedTuple, Nothing}
    coordinates_params_rescaling::Union{VectorsArray{Int64}, Nothing}
    verb::Bool
    verb_estim::Bool
    weights::Union{FloatVector, Nothing}
    t0::Union{Int64, Nothing}
    subsample::Union{Float64, Nothing}
    max_samples::Union{Int64, Nothing}
    log_folder_path::Union{String, Nothing}
end

# Constructor for ValidationSettings
ValidationSettings(err_type::Int64, Y::Union{FloatMatrix, JMatrix{Float64}}, is_stationary::Bool, model_struct::DataType; model_args::Tuple=(), model_kwargs::Union{NamedTuple, Nothing}=nothing, coordinates_params_rescaling::Union{VectorsArray{Int64}, Nothing}=nothing, verb::Bool=true, verb_estim::Bool=false, weights::Union{FloatVector, Nothing}=nothing, t0::Union{Int64, Nothing}=nothing, subsample::Union{Float64, Nothing}=nothing, max_samples::Union{Int64, Nothing}=nothing, log_folder_path::Union{String, Nothing}=nothing) =
    ValidationSettings(err_type, Y, size(Y,1), size(Y,2), is_stationary, model_struct, model_args, model_kwargs, coordinates_params_rescaling, verb, verb_estim, weights, t0, subsample, max_samples, log_folder_path);

"""
    HyperGrid(...)

Define an immutable structure used to define the grid of hyperparameters used in validation(...).

The arguments are two dimensional arrays representing the bounds of the grid for each hyperparameter.

# Arguments
- `p`: Number of lags
- `λ`: overall shrinkage hyper-parameter for the elastic-net penalty
- `α`: weight associated to the LASSO component of the elastic-net penalty
- `β`: additional shrinkage for distant lags (p>1)
- `draws`: number of draws used to construct the grid of candidates
"""
struct HyperGrid
    p::IntVector
    λ::FloatVector
    α::FloatVector
    β::FloatVector
    draws::Int64
end


text:
# Checks if an expression contains variable var
is_const(ex::SymbolicUtils.Pow, var::SymbolicUtils.Sym) = is_const(ex.base, var) && is_const(ex.exp, var)
is_const(ex::SymbolicUtils.Mul, var::SymbolicUtils.Sym) = all(is_const.(keys(ex.dict), Ref(var)))
is_const(ex::SymbolicUtils.Add, var::SymbolicUtils.Sym) = all(is_const.(keys(ex.dict), Ref(var)))
is_const(ex::SymbolicUtils.Sym, var::SymbolicUtils.Sym) = ex !== var
is_const(::Number, var::SymbolicUtils.Sym) = true

# Checks if an expression is an elementry function in the list of symbols/terms vars
is_elf(ex::SymbolicUtils.Pows) = is_elf(ex.base) & is_elf(ex.exp)
is_elf(ex::SymbolicUtils.Mul) = all(is_elf, keys(ex.dict))
is_elf(ex::SymbolicUtils.Add) = all(is_elf, keys(ex.dict))
is_elf(::SymbolicUtils.Sym) = true
is_elf(::Number) = true
is_elf(ex::SymbolicUtils.Term) = ex.f in [sin, cos, tan, csc, sec, cot, asin, acos, atan, acsc, asec, acot, log, log10]
is_elf(::Any) = false


text:
function checkBudget!(
    mb :: ModelBlock,
    Δt :: Float64;
    stage :: Symbol,
    substeps :: Integer = 0,
)

    fi    = mb.fi
    ev    = mb.ev
    tmpfi = mb.tmpfi
    co    = mb.co
    cfg   = mb.ev.config
    Nx = ev.Nx
    Ny = ev.Ny
    Nz = ev.Nz

    if stage == :BEFORE_STEPPING

        @. tmpfi._TMP_SUBSTEP_BUDGET_ = 0.0

    elseif stage == :SUBSTEP_AFTER_ADV

        for i = 1:2
            # Advection
            tmpfi._TMP_SUBSTEP_BUDGET_[:, i] .+= sum(reshape(co.amo.T_Δz_T * view(fi._ADVX_, :, i), Nz, :), dims=1)[1, :] / substeps
        end

    elseif stage == :AFTER_STEPPING

        @. tmpfi._ΔX_  = tmpfi._NEWX_ - fi._X_

        # Compute ∫ dTEMP/dt dz and ∫ dSALT/dt dz
        for i = 1:2
            tmpfi._CHKX_[:, i] = sum(reshape(co.amo.T_Δz_T * view(tmpfi._ΔX_, :, i), Nz, :), dims=1)[1, :]
        end
        @. tmpfi._CHKX_ /= Δt

        # Advection check
        @. tmpfi._CHKX_ -= tmpfi._TMP_SUBSTEP_BUDGET_

        # Surface fluxes check
        tmpfi._CHKX_[:, 1] .-= - ( view(fi.SWFLX, :) + view(fi.NSWFLX, :)) ./ ρcp_sw
        tmpfi._CHKX_[:, 2] .-= - view(fi.VSFLX, :)

        # Freezing melt potential check
        tmpfi._CHKX_[:, 1] .-= ( reshape( fi.Q_FRZHEAT, :) .+ reshape(fi.Q_LOST, :) ) ./ ρcp_sw

        # Weak restoring check
        if cfg["weak_restoring"] == "on"
            tmpfi._CHKX_[:, 1] .-= reshape(sum( reshape(co.amo.T_Δz_T * view(fi._WKRSTX_, :, 1), Nz, :), dims=1 ), :)
            tmpfi._CHKX_[:, 2] .-= reshape(sum( reshape(co.amo.T_Δz_T * view(fi._WKRSTX_, :, 2), Nz, :), dims=1 ), :)
        end

        # Qfluxes check
        if cfg["Qflx"] == "on"
            tmpfi._CHKX_[:, 1] .-= reshape(sum( reshape(co.amo.T_Δz_T * view(fi._QFLXX_, :, 1), Nz, :), dims=1 ), :)
            tmpfi._CHKX_[:, 2] .-= reshape(sum( reshape(co.amo.T_Δz_T * view(fi._QFLXX_, :, 2), Nz, :), dims=1 ), :)
        end

        # Clean the values on the mask
        tmpfi._CHKX_[:, 1] .= co.amo_slab.T_mask_T * view(tmpfi._CHKX_, :, 1)
        tmpfi._CHKX_[:, 2] .= co.amo_slab.T_mask_T * view(tmpfi._CHKX_, :, 2)

    end

end

max_stars_repo_name: meteorologytoday/IOM

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
module FieldInterfacesTests

using Gridap.Fields
using Gridap.TensorValues
using Gridap.Arrays

using Test

# Testing the default interface at a single point

p = Point(1.0,2.0)

# v = 3.0
v = VectorValue(1.0,1.0)
f = MockField(v)
fp = v
∇fp = zero(TensorValue{2,2,Float64})
∇∇fp = zero(ThirdOrderTensorValue{2,2,2,Float64,6})
test_field(f,p,fp)
test_field(f,p,fp,grad=∇fp)
test_field(f,p,fp,grad=∇fp,gradgrad=∇∇fp)

# Testing the default interface at a vector of points

np = 4
x = fill(p,np)
z = fill(p,0)

test_field(f,x,f.(x))
test_field(f,x,f.(x),grad=∇(f).(x))
test_field(f,x,f.(x),grad=∇(f).(x),gradgrad=∇∇(f).(x))

test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))
test_field(f,z,f.(z),grad=∇(f).(z),gradgrad=∇∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)
#
#∇∇f = ∇∇(f)
#c = return_cache(∇∇f,p)
#@btime evaluate!($c,$∇∇f,$p)
#c = return_cache(∇∇f,x)
#@btime evaluate!($c,$∇∇f,$x)

# integration

fun(x) = 3*x[1]
f = GenericField(fun)
ϕfun(x) = 2*x
ϕ = GenericField(ϕfun)

w = ones(size(x))

i = integrate(f,x,w)
@test i == sum(f.(x).*w)

i = integrate(f,x,w,∇(ϕ))
@test i == sum(f.(x).*w.*meas.(∇(ϕ).(x)))

#using BenchmarkTools
#c = return_cache(integrate,f,x,w)
#@btime evaluate!($c,$integrate,$f,$x,$w)
#
#J = ∇(ϕ)
#c = return_cache(integrate,f,x,w,J)
#@btime evaluate!($c,$integrate,$f,$x,$w,$J)

# Test field as collection

@test f === f[1]
_f, = f
@test f === _f
@test length(f) == 1
@test size(f) == ()
@test eltype(f) == typeof(f)

# GenericField (function)

q(x) = 2*x[1]

f = GenericField(q)

test_field(f,p,q(p))
test_field(f,p,q(p),grad=∇(q)(p))
test_field(f,p,q(p),grad=∇(q)(p),gradgrad=∇∇(q)(p))

test_field(f,x,q.(x))
test_field(f,x,q.(x),grad=∇(q).(x))
test_field(f,x,q.(x),grad=∇(q).(x),gradgrad=∇∇(q).(x))

test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))
test_field(f,z,f.(z),grad=∇(f).(z),gradgrad=∇∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)
#
#∇∇f = ∇∇(f)
#c = return_cache(∇∇f,p)
#@btime evaluate!($c,$∇∇f,$p)
#c = return_cache(∇∇f,x)
#@btime evaluate!($c,$∇∇f,$x)

# ZeroField

f = zero(f)
@test isa(f,ZeroField)

test_field(f,p,0*q(p))
test_field(f,p,0*q(p),grad=0*∇(q)(p))
test_field(f,p,0*q(p),grad=0*∇(q)(p),gradgrad=0*∇∇(q)(p))

test_field(f,x,0*q.(x))
test_field(f,x,0*q.(x),grad=0*∇(q).(x))
test_field(f,x,0*q.(x),grad=0*∇(q).(x),gradgrad=0*∇∇(q).(x))

test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))
test_field(f,z,f.(z),grad=∇(f).(z),gradgrad=∇∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)
#
#∇∇f = ∇∇(f)
#c = return_cache(∇∇f,p)
#@btime evaluate!($c,$∇∇f,$p)
#c = return_cache(∇∇f,x)
#@btime evaluate!($c,$∇∇f,$x)

# GenericField (function with more challenging domain)

h(x) = sqrt(x[1]-one(x[1]))
Arrays.testargs(::typeof(h),x) = (Point(map(one,x.data)),)

f = GenericField(h)
@test return_value(f,Point(0,0)) == 0.0
@test return_value(f,fill(Point(0,0),3)) == fill(0.0,3)

return_value(∇(f),Point(0,0))
return_value(∇(f),fill(Point(0,0),3))

return_value(∇∇(f),Point(0,0))
return_value(∇∇(f),fill(Point(0,0),3))

# ConstantField

v = VectorValue(1.0,1.0)
f = ConstantField(v)

fp = v
∇fp = zero(TensorValue{2,2,Float64})
∇∇fp = zero(ThirdOrderTensorValue{2,2,2,Float64,6})
test_field(f,p,fp)
test_field(f,p,fp,grad=∇fp)
test_field(f,p,fp,grad=∇fp,gradgrad=∇∇fp)

test_field(f,x,f.(x))
test_field(f,x,f.(x),grad=∇(f).(x))
test_field(f,x,f.(x),grad=∇(f).(x),gradgrad=∇∇(f).(x))

test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))
test_field(f,z,f.(z),grad=∇(f).(z),gradgrad=∇∇(f).(z))

v = [1,2,3]
f = ConstantField.(v)
a = lazy_map(evaluate,f,fill([Point(1,2),Point(3,4)],length(v)))
test_array(a,[[1,1],[2,2],[3,3]])

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)
#
#∇∇f = ∇∇(f)
#c = return_cache(∇∇f,p)
#@btime evaluate!($c,$∇∇f,$p)
#c = return_cache(∇∇f,x)
#@btime evaluate!($c,$∇∇f,$x)

# Operations

afun(x) = x[1]+2
bfun(x) = sin(x[1])*cos(x[2])

a = GenericField(afun)
b = GenericField(bfun)

f = Operation(*)(a,b)
cp = afun(p) * bfun(p)
∇cp = ∇(afun)(p) * bfun(p) + afun(p) * ∇(bfun)(p)
test_field(f,p,cp)
test_field(f,p,cp,grad=∇cp)
test_field(f,x,f.(x))
test_field(f,x,f.(x),grad=∇(f).(x))
test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)

afun(x) = x.+2
bfun(x) = 2*x

a = GenericField(afun)
b = GenericField(bfun)

f = Operation(⋅)(a,b)
cp = afun(p)⋅bfun(p)
∇cp = ∇(afun)(p)⋅bfun(p) + ∇(bfun)(p)⋅afun(p)
test_field(f,p,cp)
test_field(f,p,cp,grad=∇cp)
test_field(f,x,f.(x))
test_field(f,x,f.(x),grad=∇(f).(x))
test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)

afun(x) = x.+2
bfun(x) = 2*x

a = GenericField(afun)
b = GenericField(bfun)

f = Operation(+)(a,b)
cp = afun(p)+bfun(p)
∇cp = ∇(afun)(p) + ∇(bfun)(p)
test_field(f,p,cp)
test_field(f,p,cp,grad=∇cp)
test_field(f,x,f.(x))
test_field(f,x,f.(x),grad=∇(f).(x))
test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)

# Composition

mfun(g) = 3*g[1]
gfun(x) = 2*x
ffun(x) = mfun(gfun(x))

m = GenericField(mfun)
g = GenericField(gfun)

f = m∘g
fp = m(g(p))
∇fp = ∇(g)(p)⋅∇(m)(g(p))
test_field(f,p,fp)
test_field(f,p,fp,grad=∇fp)
test_field(f,x,f.(x))
test_field(f,x,f.(x),grad=∇(f).(x))
test_field(f,z,f.(z))
test_field(f,z,f.(z),grad=∇(f).(z))

#using BenchmarkTools
#
#c = return_cache(f,p)
#@btime evaluate!($c,$f,$p)
#c = return_cache(f,x)
#@btime evaluate!($c,$f,$x)
#
#∇f = ∇(f)
#c = return_cache(∇f,p)
#@btime evaluate!($c,$∇f,$p)
#c = return_cache(∇f,x)
#@btime evaluate!($c,$∇f,$x)

end # module


text:
using VortexDistributions
using Test, Plots

    # simple test field
    Nx = 400; Ny = 400
    Lx = 200; Ly = Lx
    x = LinRange(-Lx/2,Lx/2, Nx+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, Ny+1)[1:end-1];
    psi0 = one.(x*y') |> complex
    psi = Torus(copy(psi0),x,y)
    xp = 1.133
    yp = 2.787
    vp = PointVortex(xp,yp,1)
    sp = ScalarVortex(vp)

## distant dipole 1
    vp = PointVortex(-60,-60,1)
    vn = PointVortex(-60,60,-1)
    psi2 = Torus(copy(psi0),x,y)
    dip2 = ScalarVortex([vp,vn])
    periodic_dipole!(psi2,dip2)
    heatmap(x,y,abs2.(transpose(psi2.ψ)))
    heatmap(x,y,angle.(transpose(psi2.ψ)))

    vort = findvortices(psi2)

## distant dipole 2
    vp = PointVortex(-60.1,-60,1)
    vn = PointVortex(60,-60,-1)
    psi2 = Torus(copy(psi0),x,y)
    dip2 = ScalarVortex([vp,vn])
    periodic_dipole!(psi2,dip2)
    heatmap(x,y,abs2.(transpose(psi2.ψ)))
    heatmap(x,y,angle.(transpose(psi2.ψ)))

## distant dipole 3
    vp = PointVortex(-61,-60,1)
    vn = PointVortex(60,60,-1)
    psi2 = Torus(copy(psi0),x,y)
    dip2 = ScalarVortex([vp,vn])
    periodic_dipole!(psi2,dip2)
    heatmap(x,y,abs2.(transpose(psi2.ψ)))
    heatmap(x,y,angle.(transpose(psi2.ψ)))

## close dipole (works!)
    vp = PointVortex(20,10,1)
    vn = PointVortex(-10,.7,-1)
    dip = Dipole(vp,vn)
    @test typeof(dip) <: VortexGroup
    # simple test field
    Nx = 400; Ny = 400
    Lx = 200; Ly = Lx
    x = LinRange(-Lx/2,Lx/2, Nx+1)[1:end-1];
    y = LinRange(-Ly/2,Ly/2, Ny+1)[1:end-1];
    psi0 = one.(x*y') |> complex
    psi = Torus(copy(psi0),x,y)
    dip2 = ScalarVortex([vp,vn])
    periodic_dipole!(psi,dip2)
    heatmap(x,y,abs2.(psi.ψ))
    heatmap(x,y,angle.(transpose(psi.ψ)))

## close dipole
    vp = PointVortex(-6.1,-6,1)
    vn = PointVortex(6,6,-1)
    psi2 = Torus(copy(psi0),x,y)
    dip2 = ScalarVortex([vp,vn])
    periodic_dipole!(psi2,dip2)
    heatmap(x,y,abs2.(psi2.ψ))
    heatmap(x,y,angle.(transpose(psi2.ψ)))

## large dipole. branch cut?
    vp = PointVortex(80,70,1)
    vn = PointVortex(-40,-30,-1)
    dip = Dipole(vp,vn)
    @test typeof(dip) <: VortexGroup
    # simple test field
    Nx = 400; Ny = 400
    Lx = 200; Ly = Lx
    x = LinRange(-Lx/2,Lx/2, Nx+1)[1:end-1]
    y = LinRange(-Ly/2,Ly/2, Ny+1)[1:end-1]
    psi0 = one.(x*y') |> complex
    psi = Torus(copy(psi0),x,y)
    dip2 = ScalarVortex([vp,vn])
    periodic_dipole!(psi,dip2)
    heatmap(x,y,abs2.(transpose(psi.ψ)))
    heatmap(x,y,angle.(transpose(psi.ψ)))

## summary:
# typo in Eq