# Fill in the gradients . No defaults since almost always called internally to custom rule
function generate_perturbation_derivatives!(m, p_d, p_f, cache::AbstractSolverCache{1};
                                            settings = PerturbationSolverSettings())
    @assert cache.p_d_symbols == collect(Symbol.(keys(p_d)))

    p = isnothing(p_f) ? p_d : order_vector_by_symbols(merge(p_d, p_f), m.mod.p_symbols)

    # Fill in derivatives
    ret = evaluate_first_order_functions_p!(m, cache, settings, p)
    (ret == :Success) || return ret
    ret = solve_first_order_p!(m, cache, settings)
    maybe_call_function(settings.solve_first_order_p_callback, ret, m, cache, settings)
    return ret
end

function generate_perturbation_derivatives!(m, p_d, p_f, cache::AbstractSolverCache{2};
                                            settings = PerturbationSolverSettings())
    @assert cache.p_d_symbols == collect(Symbol.(keys(p_d)))

    p = isnothing(p_f) ? p_d : order_vector_by_symbols(merge(p_d, p_f), m.mod.p_symbols)

    # Fill in derivatives, first by calling the first-order
    # Fill in derivatives
    ret = evaluate_first_order_functions_p!(m, cache, settings, p)
    (ret == :Success) || return ret
    ret = solve_first_order_p!(m, cache, settings)
    maybe_call_function(settings.solve_first_order_p_callback, ret, m, cache, settings)

    # 2nd Order calculations
    ret = evaluate_second_order_functions_p!(m, cache, settings, p)
    (ret == :Success) || return ret
    ret = solve_second_order_p!(m, cache, settings)
    maybe_call_function(settings.solve_second_order_p_callback, ret, m, cache, settings)
    return ret
end

function evaluate_first_order_functions_p!(m, c, settings, p)
    (settings.print_level > 2) &&
        println("Evaluating first-order function derivatives into cache")
    try
        @unpack y, x = c  # Precondition: valid (y, x) steady states
        isnothing(c.H_p) ||
            fill_array_by_symbol_dispatch(m.mod.H_p!, c.H_p, c.p_d_symbols, y, x, p) #not required if steady_state_p!
        fill_array_by_symbol_dispatch(m.mod.H_yp_p!, c.H_yp_p, c.p_d_symbols, y, x, p)
        fill_array_by_symbol_dispatch(m.mod.H_y_p!, c.H_y_p, c.p_d_symbols, y, x, p)
        fill_array_by_symbol_dispatch(m.mod.H_xp_p!, c.H_xp_p, c.p_d_symbols, y, x, p)
        fill_array_by_symbol_dispatch(m.mod.H_x_p!, c.H_x_p, c.p_d_symbols, y, x, p)
        fill_array_by_symbol_dispatch(m.mod.Γ_p!, c.Γ_p, c.p_d_symbols, p)
        isnothing(c.Ω_p) ||
            fill_array_by_symbol_dispatch(m.mod.Ω_p!, c.Ω_p, c.p_d_symbols, p)
    catch e
        if !is_linear_algebra_exception(e)
            (settings.print_level > 2) && println("Rethrowing exception")
            rethrow(e)
        else
            settings.print_level == 0 || display(e)
            return :Failure # generic failure
        end
    end
    return :Success  # no failing code-paths yet.
end

function evaluate_second_order_functions_p!(m, c, settings, p)
    (settings.print_level > 2) &&
        println("Evaluating second-order function derivatives into cache")
    try
        @unpack y, x = c  # Precondition: valid (y, x) steady states
        fill_array_by_symbol_dispatch(m.mod.Ψ_p!, c.Ψ_p, c.p_d_symbols, y, x, p)
    catch e
        if !is_linear_algebra_exception(e)
            (settings.print_level > 2) && println("Rethrowing exception")
            rethrow(e)
        else
            settings.print_level == 0 || display(e)
            return :Failure # generic failure
        end
    end
    return :Success  # no failing code-paths yet.
end

function solve_first_order_p!(m, c, settings)
    @unpack n_x, n_y, n_ϵ = m
    n = n_x + n_y
    n_p = length(c.p_d_symbols)
    (settings.print_level > 2) && println("Solving first order derivatives of perturbation")

    buff = c.first_order_solver_p_buffer
    try
        if isnothing(m.mod.ȳ_p!) && isnothing(m.mod.x̄_p!)
            # Zeroth-order derivatives if not provided
            # Calculating c.y_p, c.x_p
            A_zero = [c.H_y + c.H_yp c.H_x + c.H_xp]
            # H_p is a vector of vectors
            x_zeroth = A_zero \ -hcat(c.H_p...) # (47)
            c.y_p .= x_zeroth[1:n_y, :]
            c.x_p .= x_zeroth[(n_y + 1):n, :]
        end

        # The derivatives
        for i in 1:n_p
            c.Σ_p[i] .= Symmetric(c.Γ_p[i] * c.Γ' + c.Γ * c.Γ_p[i]')
        end

        # Write equation (52) as E + AX + CXD = 0, a generalized Sylvester equation
        # first-order derivatives
        buff.R .= vcat(c.g_x * c.h_x, c.g_x, c.h_x, c.I_x)
        buff.A .= [c.H_y c.H_xp + c.H_yp * c.g_x]

        # i.e. C = [c.H_yp zeros(n, n_x)]
        buff.C[:, 1:n_y] .= c.H_yp
        ws = GeneralizedSylvesterWs(n, n, n_x, 1)

        # Initialize
        for i in 1:n_p
            # p-specific Sylvester preparation
            buff.bar[1:n_y] = c.y_p[i]
            buff.bar[(n_y + 1):(2 * n_y)] = c.y_p[i]
            buff.bar[(2 * n_y + 1):(2 * n_y + n_x)] = c.x_p[i]
            buff.bar[(2 * n_y + n_x + 1):end] = c.x_p[i]

            buff.dH[:, 1:n_y] = c.H_yp_p[i]
            buff.dH[:, (n_y + 1):(2 * n_y)] = c.H_y_p[i]
            buff.dH[:, (2 * n_y + 1):(2 * n_y + n_x)] = c.H_xp_p[i]
            buff.dH[:, (2 * n_y + n_x + 1):end] = c.H_x_p[i]

            for j in 1:n
                buff.dH[j, :] += c.Ψ[j] * buff.bar
            end
            mul!(buff.E, buff.dH, buff.R, -1.0, 0.0)

            # solves AXB + CXD = E
            # Sylvester
            generalized_sylvester_solver!(buff.A, buff.C, c.h_x, buff.E, 1, ws)
            c.g_x_p[i] .= buff.E[1:n_y, :]
            c.h_x_p[i] .= buff.E[(n_y + 1):n, :]

            # Q weighted derivatives
            c.C_1_p[i] .= c.Q * vcat(c.g_x_p[i], zeros(n_x, n_x))
            c.A_1_p[i] .= c.h_x_p[i]

            # V derivatives
            tmp = c.h_x_p[i] * Array(c.V) * c.h_x'
            c.V_p[i] .= lyapd(c.h_x, c.η * c.Σ_p[i] * c.η' + tmp + tmp')
            # B derivatives
            c.B_p[i] .= c.η * c.Γ_p[i]
        end

        #@exfiltrate  # flip on to see intermediate calculations.  TURN OFF BEFORE PROFILING

    catch e
        if !is_linear_algebra_exception(e)
            (settings.print_level > 2) && println("Rethrowing exception")
            rethrow(e)
        else
            settings.print_level == 0 || display(e)
            return :Failure # generic failure
        end
    end
    return :Success
end

function solve_second_order_p!(m, c, settings)
    @unpack n_x, n_y, n_ϵ, n_z = m
    n = n_x + n_y
    n_p = length(c.p_d_symbols)
    (settings.print_level > 2) &&
        println("Solving second order derivatives of perturbation")
    buff = c.second_order_solver_p_buffer

    try
        # General Prep
        buff.A .= [c.H_y c.H_xp + c.H_yp * c.g_x]
        buff.C[:, 1:n_y] .= c.H_yp
        ws = GeneralizedSylvesterWs(n, n, n_x, 2)
        buff.E .= zeros(n, n_x * n_x)
        buff.R .= vcat(c.g_x * c.h_x, c.g_x, c.h_x, c.I_x)
        buff.A_σ .= [c.H_yp + c.H_y c.H_xp + c.H_yp * c.g_x]
        buff.R_σ .= vcat(c.g_x, zeros(n_y, n_x), c.I_x, zeros(n_x, n_x))

        buff.gh_stack[1:n_y, :] .= reshape(c.g_xx, n_y, n_x * n_x)
        buff.gh_stack[n_y+1:end, :] .= reshape(c.h_xx, n_x, n_x * n_x)
        buff.g_xx_flat .= reshape(c.g_xx, n_y, n_x * n_x)
        for i in 1:n
            for j in 1:n_x
                buff.Ψ_x_sum[j][i] .= c.Ψ_xp[j][i] .+ c.Ψ_x[j][i]
            end
            for j in 1:n_y
                buff.Ψ_y_sum[j][i] .= c.Ψ_yp[j][i] .+ c.Ψ_y[j][i]
            end
        end

        tmp1 = similar(buff.R')
        tmp2 = similar(c.h_x)
        for i in 1:n_p
            # Prep for _xx_p
            # Compute the total derivatives
            buff.bar[1:n_y] = c.y_p[i]
            buff.bar[(n_y + 1):(2 * n_y)] = c.y_p[i]
            buff.bar[(2 * n_y + 1):(2 * n_y + n_x)] = c.x_p[i]
            buff.bar[(2 * n_y + n_x + 1):end] = c.x_p[i]
            buff.dH[:, 1:n_y] = c.H_yp_p[i]
            buff.dH[:, (n_y + 1):(2 * n_y)] = c.H_y_p[i]
            buff.dH[:, (2 * n_y + 1):(2 * n_y + n_x)] = c.H_xp_p[i]
            buff.dH[:, (2 * n_y + n_x + 1):end] = c.H_x_p[i]
            for j in 1:n
                buff.dH[j, :] += c.Ψ[j] * buff.bar
            end
            for j in 1:n
                buff.dΨ[j] .= c.Ψ_p[i][j]
            end
            for j in 1:n_y
                if (c.y_p[i][j] != 0)
                    for k in 1:n
                        buff.dΨ[k] .+= buff.Ψ_y_sum[j][k] .* c.y_p[i][j]
                    end
                end
            end
            for j in 1:n_x
                if (c.x_p[i][j] != 0)
                    for k in 1:n
                        buff.dΨ[k] .+= buff.Ψ_x_sum[j][k] .* c.x_p[i][j]
                    end
                end
            end

            # Constants: (60)
            buff.R_p[1:n_y, :] = c.g_x_p[i] * c.h_x + c.g_x * c.h_x_p[i]
            buff.R_p[(n_y + 1):(2 * n_y), :] = c.g_x_p[i]
            buff.R_p[(2 * n_y + 1):(2 * n_y + n_x), :] = c.h_x_p[i]
            # Flip the sign of E for Sylvester input
            fill!(buff.E, 0.0)
            for j in 1:n
                mul!(tmp1, buff.R_p', c.Ψ[j])
                mul!(tmp2, tmp1, buff.R)
                buff.E[j, :] .-= vec(tmp2)
                buff.E[j, :] .-= vec(tmp2')
                mul!(tmp1, buff.R', buff.dΨ[j])
                mul!(tmp2, tmp1, buff.R)
                buff.E[j, :] .-= vec(tmp2)
                # buff.E[j, :] .= -vec(buff.R_p' * c.Ψ[j] * buff.R + buff.R' * c.Ψ[j] * buff.R_p + buff.R' * buff.dΨ[j] * buff.R)
            end
            # Constants: (56)
            kron!(buff.kron_h_x, c.h_x, c.h_x)
            buff.E .-= buff.dH[:, 1:n_y] * buff.g_xx_flat * buff.kron_h_x # Plug (57) in (56)
            tmp = c.H_yp * buff.g_xx_flat
            kron!(buff.kron_h_x, c.h_x_p[i], c.h_x)
            mul!(buff.E, tmp, buff.kron_h_x, -1.0, 1.0) # Plug (58) in (56), step 1
            kron!(buff.kron_h_x, c.h_x, c.h_x_p[i])
            mul!(buff.E, tmp, buff.kron_h_x, -1.0, 1.0) # Plug (58) in (56), step 2
            buff.E .-= hcat(buff.dH[:, (n_y + 1):(2 * n_y)],
            buff.dH[:, 1:n_y] * c.g_x +
            c.H_yp * c.g_x_p[i] +
            buff.dH[:, (2 * n_y + 1):(2 * n_y + n_x)]) * buff.gh_stack # Plug (59) in (56)

            # Solve the Sylvester equations (56)
            generalized_sylvester_solver!(buff.A, buff.C, c.h_x, buff.E, 2, ws)
            copyto!(c.g_xx_p[i], buff.E[1:n_y, :]) # Reshaping into n_y * n_x * n_x
            copyto!(c.h_xx_p[i], buff.E[(n_y + 1):end, :]) # # Reshaping into n_x * n_x * n_x

            # Prep for _σσ_p
            # Solve _σσ_p
            R_σ_p = vcat(c.g_x_p[i], zeros(n_y + n_x * 2, n_x))
            η_sq_p = c.η * c.Σ_p[i] * c.η'
            C_σ = -(buff.dH[:, 1:n_y] + buff.dH[:, (n_y + 1):(2 * n_y)]) * c.g_σσ
            C_σ -= (buff.dH[:, 1:n_y] * c.g_x + c.H_yp * c.g_x_p[i] + buff.dH[:, (2 * n_y + 1):(2 * n_y + n_x)]) * c.h_σσ
            # C_σ = -hcat(buff.dH[:, 1:n_y] + buff.dH[:, (n_y + 1):(2 * n_y)],
            #             buff.dH[:, 1:n_y] * c.g_x +
            #             c.H_yp * c.g_x_p[i] +
            #             buff.dH[:, (2 * n_y + 1):(2 * n_y + n_x)]) * vcat(c.g_σσ, c.h_σσ) # Plug (65) in (64), flip the sign to solve (64)
            C_σ -= (buff.dH[:, 1:n_y] * buff.g_xx_flat + c.H_yp * buff.E[1:n_y, :]) * vec(c.η_Σ_sq)# (67), 2nd line
            C_σ -= (c.H_yp * buff.g_xx_flat) * vec(η_sq_p) # (67), 3rd line, second part
            for j in 1:n
                mul!(tmp1, R_σ_p', c.Ψ[j])
                mul!(tmp2, tmp1, buff.R_σ)
                C_σ[j] -= dot(tmp2, c.η_Σ_sq)
                C_σ[j] -= dot(tmp2', c.η_Σ_sq)
                mul!(tmp1, buff.R_σ', buff.dΨ[j])
                mul!(tmp2, tmp1, buff.R_σ)
                C_σ[j] -= dot(tmp2, c.η_Σ_sq)
                # C_σ[j] -= dot((R_σ_p' * c.Ψ[j] * buff.R_σ + buff.R_σ' * c.Ψ[j] * R_σ_p + buff.R_σ' * buff.dΨ[j] * buff.R_σ), c.η_Σ_sq) # (67), 1st line
                mul!(tmp1, buff.R_σ', c.Ψ[j])
                mul!(tmp2, tmp1, buff.R_σ)
                C_σ[j] -= dot(tmp2, η_sq_p)
                # C_σ[j] -= dot((buff.R_σ' * c.Ψ[j] * buff.R_σ), η_sq_p) # (67), 3rd line, first part
            end

            # Solve _σσ_p
            X_σ = buff.A_σ \ C_σ # solve (64)
            c.g_σσ_p[:, i] .= X_σ[1:n_y]
            c.h_σσ_p[:, i] .= X_σ[(n_y + 1):end]
            fill!(c.C_2_p[i], 0.0) # reset as we need to use `+=`
            for j in 1:n_z
                for k in 1:n_y
                    if (c.Q[j, k] != 0)
                        c.C_2_p[i][j, :, :] .+= 0.5 * c.Q[j, k] * c.g_xx_p[i][k, :, :]
                    end
                end
            end
            c.A_2_p[i] .= 0.5 * c.h_xx_p[i]
        end

        c.C_0_p .= 0.5 * c.Q * vcat(c.g_σσ_p, zeros(n_x, n_p))
        c.A_0_p .= 0.5 * c.h_σσ_p

    catch e
        settings.print_level == 0 || display(e)
        return :Failure # generic failure
    end
    return :Success
end

function ChainRulesCore.rrule(::typeof(generate_perturbation), m::PerturbationModel,
                              p_d::NamedTuple{DFieldsType,DTupleType}, p_f, order::Val{1};
                              cache = SolverCache(m, Val(1), p_d),
                              settings = PerturbationSolverSettings()) where {DFieldsType,
                                                                              DTupleType}
    (settings.print_level > 2) && println("Calculating generate_perturbation primal ")
    sol = generate_perturbation(m, p_d, p_f, Val(1); cache, settings)
    if (sol.retcode == :Success)
        grad_ret = generate_perturbation_derivatives!(m, p_d, p_f, cache)
        if (grad_ret != :Success)
            sol = FirstOrderPerturbationSolution(:GradientFailure, m, cache)
        end
    end

    function generate_perturbation_pb(Δsol)
        (settings.print_level > 2) && println("Calculating generate_perturbation pullback")
        Δp = (p_d === nothing) ? nothing : zeros(length(p_d))
        if (sol.retcode == :Success) & (p_d !== nothing)
            c = cache # temp to avoid renaming everything
            n_p_d = length(p_d)
            if (~iszero(Δsol.A))
                for i in 1:n_p_d
                    Δp[i] += dot(c.h_x_p[i], Δsol.A)
                end
            end
            if (~iszero(Δsol.g_x))
                for i in 1:n_p_d
                    Δp[i] += dot(c.g_x_p[i], Δsol.g_x)
                end
            end
            if (~iszero(Δsol.C))
                for i in 1:n_p_d
                    Δp[i] += dot(c.C_1_p[i], Δsol.C)
                end
            end
            if (~isnothing(Δsol.x_ergodic))
                if ((Δsol.x_ergodic != NoTangent()) & (Δsol.x_ergodic != ZeroTangent()))
                    for i in 1:n_p_d
                        tmp = c.V.L \ c.V_p[i] / c.V.U
                        tmp[diagind(tmp)] /= 2.0
                        t1 = c.V.L * LowerTriangular(tmp)
                        # Cholesky by default stores the U part, but that information is lost when passing back
                        # from MvNormal. Therefore, we have to extract the components out manually
                        Δp[i] += dot(t1', UpperTriangular(Δsol.x_ergodic.C.factors)) # dot(c.V_p[i], Δsol.x_ergodic) is the original logic
                    end
                end
            end
            if (~iszero(Δsol.Γ))
                for i in 1:n_p_d
                    Δp[i] += dot(c.Γ_p[i], Δsol.Γ)
                end
            end
            if (~iszero(Δsol.B))
                for i in 1:n_p_d
                    Δp[i] += dot(c.B_p[i], Δsol.B)
                end
            end
            if (~isnothing(Δsol.D)) # D is a Distribution object which complicates stuff here
                if ((Δsol.D != NoTangent()) & (Δsol.D != ZeroTangent()))
                    # Only supports diagonal matrices for now.
                    for i in 1:n_p_d
                        Δp[i] += dot(c.Ω_p[i], Δsol.D.σ)
                    end
                end
            end
            if (~iszero(Δsol.x))
                for i in 1:n_p_d
                    Δp[i] += dot(c.x_p[i], Δsol.x)
                end
            end
            if (~iszero(Δsol.y))
                for i in 1:n_p_d
                    Δp[i] += dot(c.y_p[i], Δsol.y)
                end
            end
        end
        Δp_nt = NamedTuple{DFieldsType,DTupleType}(tuple(Δp...))  # turn tuple into named tuple in the same order
        return NoTangent(), NoTangent(),
               Tangent{NamedTuple{DFieldsType,DTupleType},
                       NamedTuple{DFieldsType,DTupleType}}(Δp_nt), NoTangent(), NoTangent()
    end
    # keep the named tuple the same
    return sol, generate_perturbation_pb
end

function ChainRulesCore.rrule(::typeof(generate_perturbation), m::PerturbationModel,
                              p_d::NamedTuple{DFieldsType,DTupleType}, p_f, order::Val{2};
                              cache = SolverCache(m, Val(2), p_d),
                              settings = PerturbationSolverSettings()) where {DFieldsType,
                                                                              DTupleType}
    (settings.print_level > 2) && println("Calculating generate_perturbation primal ")
    sol = generate_perturbation(m, p_d, p_f, Val(2); cache, settings)
    if (sol.retcode == :Success)
        grad_ret = generate_perturbation_derivatives!(m, p_d, p_f, cache)
        if (grad_ret != :Success)
            sol = SecondOrderPerturbationSolution(:GradientFailure, m, cache)
        end
    end
    function generate_perturbation_pb(Δsol)
        (settings.print_level > 2) && println("Calculating generate_perturbation pullback")
        Δp = (p_d === nothing) ? nothing : zeros(length(p_d))
        if (sol.retcode == :Success) & (p_d !== nothing)
            c = cache # temp to avoid renaming everything
            n_p_d = length(p_d)
            if (~iszero(Δsol.A_1))
                for i in 1:n_p_d
                    Δp[i] += dot(c.A_1_p[i], Δsol.A_1)
                end
            end
            if (~iszero(Δsol.g_x))
                for i in 1:n_p_d
                    Δp[i] += dot(c.g_x_p[i], Δsol.g_x)
                end
            end
            if (~iszero(Δsol.C_1))
                for i in 1:n_p_d
                    Δp[i] += dot(c.C_1_p[i], Δsol.C_1)
                end
            end
            if (~iszero(Δsol.Γ))
                for i in 1:n_p_d
                    Δp[i] += dot(c.Γ_p[i], Δsol.Γ)
                end
            end
            if (~iszero(Δsol.B))
                for i in 1:n_p_d
                    Δp[i] += dot(c.B_p[i], Δsol.B)
                end
            end
            if (~iszero(Δsol.A_2))
                for i in 1:n_p_d
                    Δp[i] += dot(c.A_2_p[i], Δsol.A_2)
                end
            end
            if (~iszero(Δsol.g_xx))
                for i in 1:n_p_d
                    Δp[i] += dot(c.g_xx_p[i], Δsol.g_xx)
                end
            end
            if (~iszero(Δsol.C_2))
                for i in 1:n_p_d
                    Δp[i] += dot(c.C_2_p[i], Δsol.C_2)
                end
            end
            if (~isnothing(Δsol.D)) # D is a Distribution object which complicates stuff here
                if ((Δsol.D != NoTangent()) & (Δsol.D != ZeroTangent()))
                    # Only supports diagonal matrices for now.
                    for i in 1:n_p_d
                        Δp[i] += dot(c.Ω_p[i], Δsol.D.σ)
                    end
                end
            end
            if (~iszero(Δsol.x))
                for i in 1:n_p_d
                    Δp[i] += dot(c.x_p[i], Δsol.x)
                end
            end
            if (~iszero(Δsol.y))
                for i in 1:n_p_d
                    Δp[i] += dot(c.y_p[i], Δsol.y)
                end
            end
            # Currently the results for the second-order correction terms are not vector of vectors, but just arrays
            if (~iszero(Δsol.C_0))
                Δp += c.C_0_p' * Δsol.C_0
            end
            if (~iszero(Δsol.g_σσ))
                Δp += c.g_σσ_p' * Δsol.g_σσ
            end
            if (~iszero(Δsol.A_0))
                Δp += c.A_0_p' * Δsol.A_0
            end
        end

        Δp_nt = NamedTuple{DFieldsType,DTupleType}(tuple(Δp...)) # turn tuple into named tuple in the same order
        return NoTangent(), NoTangent(),
               Tangent{NamedTuple{DFieldsType,DTupleType},
                       NamedTuple{DFieldsType,DTupleType}}(Δp_nt), NoTangent(), NoTangent()
    end

    return sol, generate_perturbation_pb
end


using Ipopt
using Test

@testset "C API" begin
    # First of all, test that hs071 example works
    include("hs071_test.jl")

    # Test setting some options
    # String option
    println("\nString option")
    addOption(prob, "hessian_approximation", "exact")
    @test_throws ErrorException addOption(prob, "hessian_approximation", "badoption")
    println("\nInt option")
    # Int option
    addOption(prob, "file_print_level", 3) == nothing
    @test_throws ErrorException addOption(prob, "file_print_level", -1)
    # Double option
    println("\nFloat option")
    addOption(prob, "derivative_test_tol", 0.5)
    @test_throws ErrorException addOption(prob, "derivative_test_tol", -1.0)

    # Test opening an output file
    openOutputFile(prob, "blah.txt", 5)
    Ipopt.freeProblem(prob) # Needed before the `rm` on Windows.
    # unlink the output file
    rm("blah.txt")

    # Test that the ipopt binary works
    @test Ipopt.amplexefun("-v") == 0
end

@testset "MathProgBase" begin
    include("MPB_wrapper.jl")
end

@testset "MathOptInterface" begin
    include("MOI_wrapper.jl")
end


# HZAM_Sym_Julia_V1.jl
# First public release of HZAM_Sym_Julia
# by Darren Irwin, begun in early July 2021
# Re-written from the R version of HZAM_Sym

# Please note that this version is not intended to be super user-friendly software
# for people who aren't familiar with the Julia language. Rather, I am making this available
# so that those who know the basics of Julia can run this and tinker if they desire,
# and so that the community can see exactly what code produced the results in the paper.

# If you use this code, please cite the paper below, and possibly my GitHub or Dryad repository
# where you obtained this code:

# Irwin, D., and Schluter, D. Hybridization and the coexistence of species. bioRxiv 2021.04.04.438369; doi: https://doi.org/10.1101/2021.04.04.438369

# You need to make sure these packages below are added to your Julia environment.
# to add, type "]" to get the "pkg>" prompt, then type e.g. "add Distributions";
# or do this:
# import Pkg;
# Pkg.add("Distributions")
# Pkg.add("Statistics")
# Pkg.add("JLD2")
# Pkg.add("CSV")
# Pkg.add("DataFrames")
# Pkg.add("Plots")
# Pkg.add("CategoricalArrays")
# Pkg.add("Colors")
# Pkg.add("ColorSchemes")
# Pkg.add("CategoricalArrays")

using Distributions # needed for "Poisson" function
using Statistics  # needed for "mean" function
using JLD2 # needed for saving / loading data in Julia format
using CSV # for saving in csv format
using DataFrames # for converting data to save as csv

# for plotting:
using Plots
gr()  # use GR backend for graphs
using CategoricalArrays
using Colors, ColorSchemes
import ColorSchemes.plasma
using Plots.PlotMeasures  # needed for plot margin adjustment

# to start Julia with multiple threads, type in terminal e.g.:
# julia --threads 4
# To check, type in Julia: Threads.nthreads()
# 4

# set up functions (should not define functions repeatedly in loop, as causes re-compilation, slows things)

function generate_genotype_array(N_pop0,N_pop1,loci)
    total_N = N_pop0 + N_pop1
    genotypes = Array{Int8, 3}(undef, 2, loci, total_N) # The "Int8" is the type (8-bit integer), and "undef" means an unitialized array, so values are meaningless
    genotypes[:,:,1:N_pop0] .= 0  # assigns genotypes of pop1
    genotypes[:,:,(N_pop0+1):total_N] .= 1  # assigns genotypes of pop
    return genotypes
end

function calc_traits_additive(genotypes)
    N = size(genotypes, 3)
    traits = Array{Float32, 1}(undef, N) # Float32 should be enough precision; memory saving compared to Float64
    for i in 1:N
        traits[i] = mean(genotypes[:,:,i])
    end
    return traits
end

function get_survival_fitnesses_epistasis(genotypes, w_hyb, beta=1)
    survival_HI = calc_traits_additive(genotypes)
    epistasis_fitnesses = 1 .- (1 - w_hyb) .* (4 .* survival_HI .* (1 .- survival_HI)).^beta
    return epistasis_fitnesses
end

function get_survival_fitnesses_hetdisadvantage(genotypes, w_hyb)
    N = size(genotypes, 3)
    num_loci = size(genotypes, 2)
    s_per_locus = 1 - w_hyb ^ (1/num_loci)  # loss in fitness due to each heterozygous locus
    num_hetloci = Array{Int16, 1}(undef, N)
    for ind in 1:N  # count number of het loci per individual
        num_hetloci[ind] = sum(genotypes[1,:,ind] .!= genotypes[2,:,ind])
    end
    hetdisadvantage_fitnesses = (1 - s_per_locus) .^ num_hetloci
    return hetdisadvantage_fitnesses
end

function run_HZAM(set_name::String, ecolDiff, intrinsic_R, replications;  # the semicolon makes the following optional keyword arguments
    K_total::Int = 1000, max_generations::Int = 1000,
    total_loci::Int = 6, female_mating_trait_loci = 1:3, male_mating_trait_loci = 1:3,
    competition_trait_loci = 1:3, hybrid_survival_loci = 1:3, neutral_loci = 4:6,
    survival_fitness_method::String = "epistasis", per_reject_cost = 0,
    starting_pop_ratio = 1.0)
    # ecolDiff should be from 0 to 1 (parameter "E" in the paper)
    # intrinsic_R is called "R" in the paper
    # replications should be somehting like "1:10" or just "1" for 1 replicate, or something like "2:5" to add replicates after 1 is done
    # K_total (default 1000) is the total carrying capacity, divided equally between K_alpha and K_beta
    # max_generations (default 1000) is the number of generations each simulation will run
    # total_loci (default 6) is the total number of loci, regardless of their role (with indices 1:total_loci, referred to below)
    # Need to specify indices (columns) of four types of functional loci (can be the same). At least one should begin with index 1:
    # female_mating_trait_loci (default 1:3) is indices of the loci that determine the female mating trait
    # male_mating_trait_loci (default 1:3) is indices of the loci that determine the male mating trait
    # competition_trait_loci (default 1:3) is indices of the loci that determine the ecological trait (used in fitness related to resource use)
    # hybrid_survival_loci (default 1:3) is indices of the loci that determine survival probability of offspring to adulthood (can be viewed as incompatibilities and/or fitness valley based on ecology)
    # Specify indices (columns) of neutral loci (which have no effect on anything, just along for the ride):
    # neutral_loci (default 4:6) is indices of neutral loci (used for neutral measure of hybrid index; not used in the HZAM-sym paper)
    # per_reject_cost (default 0, can take values of 0 to 1) is fitness loss of female per male rejected (due to search time, etc.)
    # starting_pop_ratio (default 1) is the starting ratio of pop A to pop B

    save_outcomes_JL = true
    save_outcomes_csv = true  # whether to save the whole outcome array as csv files (with each rep as separate file)
    save_each_sim = false  # whether to save detailed data for each simulation

    # the set of hybrid fitnesses (w_hyb) values that will be run
    w_hyb_set = [1, 0.98, 0.95, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0] # for just one run, just put one number in this and next line
    # the set of assortative mating strengths (S_AM) that will be run
    S_AM_set = [1, 3, 10, 30, 100, 300, 1000, Inf]  # ratio of: probably of accepting homospecific vs. prob of accepting heterospecific

    # get the chosen survival fitness function
    if survival_fitness_method == "epistasis"
        get_survival_fitnesses = get_survival_fitnesses_epistasis
        short_survFitnessMethod = "Ep"
    elseif survival_fitness_method == "hetdisadvantage"
        get_survival_fitnesses = get_survival_fitnesses_hetdisadvantage
        short_survFitnessMethod = "Het"
    else
        println("ERROR--no survival fitness method chosen--should be either epistasis or hetdisadvantage")
    end

    total_functional_loci = max(maximum(female_mating_trait_loci), maximum(male_mating_trait_loci), maximum(competition_trait_loci), maximum(hybrid_survival_loci))
    functional_loci_range = 1:total_functional_loci
    num_neutral_loci = length(Vector(neutral_loci))
    if total_functional_loci + num_neutral_loci ≠ total_loci
        println("#### WARNING: Please examine your loci numbers and indices, as they don't all match up ####")
    end

    # intrinsic_R = 1.05  # Intrinsic growth rate, this is the average maximum expected number of offspring per individual, when pop size far below K
    K_A = K_total / 2  # EVEN NUMBER; carrying capacity (on resource alpha) of entire range (for two sexes combined), regardless of species
    K_B = K_total / 2   # EVEN NUMBER; carrying capacity (on resource beta) of entire range (for two sexes combined), regardless of species

    pop0_starting_N = K_A   # starting N of species 0
    pop0_starting_N_half = Int(pop0_starting_N/2)
    pop1_starting_N = Int(round(starting_pop_ratio * K_B))   # starting N of species 1, which can be lower if starting_pop_ratio is below 1)
    pop1_starting_N_half = Int(pop1_starting_N/2)

    beta = 1  # the epistasis parameter beta

    # specify ecological resource competitive abilities for two resources A and B
    # ecolDiff = 1.0 # this is "E" in the paper
    competAbility_useResourceA_species0 = (1 + ecolDiff)/2    # equals 1 when ecolDiff = 1
    competAbility_useResourceB_species0 = 1 - competAbility_useResourceA_species0
    competAbility_useResourceA_species1 = (1 - ecolDiff)/2   # equals 0 when ecolDiff = 1
    competAbility_useResourceB_species1 = 1 - competAbility_useResourceA_species1

    # set up array of strings to record outcomes
    outcome_array = Array{String, 3}(undef, length(w_hyb_set), length(S_AM_set), length(replications))

    for k in 1:length(replications)  # loop through the replicate runs
        replicate_ID = replications[k]

        run_set_name = string(set_name,"_rep", replicate_ID)

        # Loop through the different simulation sets
        Threads.@threads for i in 1:length(w_hyb_set)
            for j in 1:length(S_AM_set)
                w_hyb = w_hyb_set[i]
                S_AM = S_AM_set[j]
                println("w_hyb = ",w_hyb,"; S_AM = ",S_AM)

                run_name = string("HZAM_animation_run",run_set_name,"_surv",short_survFitnessMethod,"_ecolDiff",ecolDiff,"_growthrate",intrinsic_R,"_K",K_total,"_FL",total_functional_loci,"_NL",num_neutral_loci,"_gen",max_generations,"_SC",per_reject_cost,"_Whyb",w_hyb,"_SAM",S_AM)

                # convert S_AM to pref_ratio (for use in math below)
                if S_AM == 1
                    S_AM_for_math = 1 + 10^(-15)
                elseif S_AM == Inf
                    S_AM_for_math = 10^(15)
                else
                    S_AM_for_math = S_AM
                end
                pref_SD = sqrt( -1 / (2 * log(1/S_AM_for_math)))  # width of female acceptance curve for male trait

                # set up initial values for one simulation
                extinction = false
                outcome = []
                final_distribution = []
                # Generate genotype array for population of females:
                # this is a 3D array, where rows (D1) are alleles (row 1 from mother, row 2 from father),
                # columns (D2) are loci, and pages (D3) are individuals
                genotypes_F = generate_genotype_array(pop0_starting_N_half, pop1_starting_N_half, total_loci)
                genotypes_M = generate_genotype_array(pop0_starting_N_half, pop1_starting_N_half, total_loci)
                # functional loci are first, followed by neutral loci

                for generation in 1:max_generations

                    # Prepare for mating and reproduction
                    N_F = size(genotypes_F, 3)
                    N_M = size(genotypes_M, 3)

                    # println("generation: ", generation, "; individuals: ", N_F + N_M)

                    # calculate mating trait values (T) from genotypes
                    female_mating_traits = calc_traits_additive(genotypes_F[:,female_mating_trait_loci,:])
                    male_mating_traits = calc_traits_additive(genotypes_M[:,male_mating_trait_loci,:])

                    # calculate ecological competition trait values from genotypes
                    competition_traits_F = calc_traits_additive(genotypes_F[:,competition_trait_loci,:])
                    competition_traits_M = calc_traits_additive(genotypes_M[:,competition_trait_loci,:])

                    # calculate individual contributions to resource use, according to linear gradient between use of species 0 and species 1
                    ind_useResourceA_F = competAbility_useResourceA_species1 .+ ((1 .- competition_traits_F) .* (competAbility_useResourceA_species0 - competAbility_useResourceA_species1))
                    ind_useResourceB_F = competAbility_useResourceB_species0 .+ (competition_traits_F .* (competAbility_useResourceB_species1 - competAbility_useResourceB_species0))
                    ind_useResourceA_M = competAbility_useResourceA_species1 .+ ((1 .- competition_traits_M) .* (competAbility_useResourceA_species0 - competAbility_useResourceA_species1))
                    ind_useResourceB_M = competAbility_useResourceB_species0 .+ (competition_traits_M .* (competAbility_useResourceB_species1 - competAbility_useResourceB_species0))
                    # sum up the resource use over all individuals
                    total_useResourceA = sum(ind_useResourceA_F) + sum(ind_useResourceA_M)
                    total_useResourceB = sum(ind_useResourceB_F) + sum(ind_useResourceB_M)
                    # calculate growth rates due to each resource (according to discrete time logistic growth equation)
                    growth_rate_ResourceA = (intrinsic_R * K_A) / (K_A + ((total_useResourceA)*(intrinsic_R - 1)))
                    growth_rate_ResourceB = (intrinsic_R * K_B) / (K_B + ((total_useResourceB)*(intrinsic_R - 1)))

                    # Set up structure to record number of matings per male (and female, although almost always 1 for females),
                    # to determine sexual selection due to HI class:
                    matings_per_male = zeros(Int8, N_M)
                    matings_per_female = zeros(Int8, N_F)

                    # create empty data structures for keeping track of numbers offspring of parents
                    daughters_per_mother = zeros(Int16, N_F)
                    sons_per_mother = zeros(Int16, N_F)
                    daughters_per_father = zeros(Int16, N_M)
                    sons_per_father = zeros(Int16, N_M)

                    # make empty arrays for storing genotypes of daughters and sons
                    genotypes_daughters = Array{Int8, 3}(undef, 2, total_loci, 0)
                    genotypes_sons = Array{Int8, 3}(undef, 2, total_loci, 0)

                    # create structures for recording indices of mother and father (for error checking, and potentially for tracking genealogies)
                    # first column for mother index (3rd dim of genotypes_F) and second column for father index (3rd dim of genotypes_M)
                    daughter_parent_IDs = Array{Int}(undef, 0, 2)
                    son_parent_IDs = Array{Int}(undef, 0, 2)

                    # loop through mothers, mating and reproducing
                    for mother in 1:N_F
                        # initialize tracking variables
                        mate = false # becomes true when female is paired with male
                        rejects = 0 # will track number of rejected males (in cases there is cost--which there isn't in main HZAM-Sym paper)
                        father = [] # will contain the index of the male mate
                        # make vector of indices of eligible males
                        elig_M = Vector(1:N_M)
                        # determine male mate of female
                        while mate == false
                            # present female with random male, and remove him from list using "splice!"
                            focal_male = splice!(elig_M, rand(eachindex(elig_M)))
                            # compare male trait with female's trait (preference), and determine
                            # whether she accepts; note that match_strength is determined by a
                            # Gaussian, with a maximum of 1 and minimum of zero.
                            match_strength = (exp(1) ^ ((-(male_mating_traits[focal_male] - female_mating_traits[mother])^2) / (2 * (pref_SD ^2))))
                            if rand() < match_strength
                                # she accepts male, and mates
                                father = focal_male
                                matings_per_male[focal_male] += 1 # this adds 1 to the matings for that male
                                matings_per_female[mother] += 1
                                mate = true
                            else
                                # she rejects male
                                rejects += 1
                                if length(elig_M) == 0
                                    break
                                end
                            end
                        end
                        # determine fitness cost due to mate search (number of rejected males)
                        search_fitness = (1-per_reject_cost) ^ rejects    # (in most of HZAM-Sym paper, per_reject_cost = 0)

                        # determine fitness due to female use of available resources
                        growth_rate_of_focal_female = (ind_useResourceA_F[mother] * growth_rate_ResourceA) + (ind_useResourceB_F[mother] * growth_rate_ResourceB)
                        #combine for total fitness:
                        reproductive_fitness = 2 * growth_rate_of_focal_female * search_fitness  # the 2 is because only females, not males, produce offspring
                        # now draw the number of offspring from a poisson distribution with a mean of reproductive_fitness
                        if isempty(father)
                            offspring = 0  # if no mate (because all males were rejected), then no offspring
                        else
                            offspring = rand(Poisson(reproductive_fitness))
                        end
                        # if offspring, generate their genotypes and sexes
                        if offspring >= 1
                            for kid in 1:offspring
                                kid_info = Array{Int8, 2}(undef, 2, total_loci) # place to store genotype of one offspring
                                # generate genotypes; for each locus (column), first row for allele from mother, second row for allele from father
                                for locus in 1:total_loci
                                    kid_info[1,locus] = genotypes_F[rand([1 2]), locus, mother]  # for this locus, pick a random allele from the mother
                                    kid_info[2,locus] = genotypes_M[rand([1 2]), locus, father]
                                end
                                # determine sex of kid and add to table
                                if rand() > 0.5 # kid is daughter
                                    genotypes_daughters = cat(genotypes_daughters, kid_info, dims=3)
                                    daughter_parent_IDs = cat(daughter_parent_IDs, [mother father], dims=1)
                                    daughters_per_mother[mother] += 1
                                    daughters_per_father[father] += 1
                                else # kid is son
                                    genotypes_sons = cat(genotypes_sons, kid_info, dims=3)
                                    son_parent_IDs = cat(son_parent_IDs, [mother father], dims=1)
                                    sons_per_mother[mother] += 1
                                    sons_per_father[father] += 1
                                end
                            end
                        end
                    end # of loop through mothers

                    # check if either no daughters or no sons, and end the simulation if so
                    if (size(genotypes_daughters, 3) == 0) || (size(genotypes_sons, 3) == 0)
                        extinction = true # record an extinction of whole population (both "species")
                        break # break out of current loop (this simulation)
                    end

                    # For someday: add in here the option of tracking fitness?

                    # determine survival fitnesses of daughters due to epistasis

                    survival_fitness_daughters = get_survival_fitnesses(genotypes_daughters[:,hybrid_survival_loci,:], w_hyb)
                    daughters_survive = survival_fitness_daughters .> rand(length(survival_fitness_daughters))
                    # same for sons:
                    survival_fitness_sons = get_survival_fitnesses(genotypes_sons[:,hybrid_survival_loci,:], w_hyb)
                    sons_survive = survival_fitness_sons .> rand(length(survival_fitness_sons))

                    # check if either no surviving daughters or no surviving sons, and end the simulation if so
                    if (sum(daughters_survive) == 0) || (sum(sons_survive) == 0)
                        extinction = true # record an extinction of whole population (both "species")
                        break # break out of current loop (this simulation)
                    end

                    # assign surviving offspring to new adult population
                    genotypes_F = genotypes_daughters[:,:,daughters_survive]
                    genotypes_M = genotypes_sons[:,:,sons_survive]

                end # of loop through generations

                # Record results of the one simulation
                functional_HI_all_inds = []
                species0_proportion = []
                species1_proportion = []
                HI_NL_all_inds = []
                species0_proportion_NL = []
                species1_proportion_NL = []
                if extinction  # whole simulation went extinct
                    outcome = "extinction"
                    if save_each_sim
                        @save string("simulation_data.",run_name,".jld2") outcome
                    end
                else  # no complete extinction
                    # use trait loci to calculate HI of each individual
                    functional_HI_all_inds = [calc_traits_additive(genotypes_F[:,functional_loci_range,:]); calc_traits_additive(genotypes_M[:,functional_loci_range,:])]
                    # calculate proportion of all individuals who are species0 or species1 (defined as low and high 10% of HI distribution, respectively)
                    species0_proportion = sum(functional_HI_all_inds .< 0.1) / length(functional_HI_all_inds)
                    species1_proportion = sum(functional_HI_all_inds .> 0.9) / length(functional_HI_all_inds)
                    if species0_proportion >= 0.85 || species1_proportion >= 0.85
                        outcome = "one_species"
                    elseif (species0_proportion + species1_proportion >= 0.85) && (species0_proportion >= 0.15) && (species1_proportion >= 0.15)
                        outcome = "two_species"
                    else
                        outcome = "blended"
                    end
                    HI_NL_all_inds = [calc_traits_additive(genotypes_F[:,neutral_loci,:]); calc_traits_additive(genotypes_M[:,neutral_loci,:])]
                    species0_proportion_NL = sum(HI_NL_all_inds .== 0) / length(HI_NL_all_inds)
                    species1_proportion_NL = sum(HI_NL_all_inds .== 1) / length(HI_NL_all_inds)
                    if save_each_sim
                        @save string("HZAM_Sym_Julia_results_GitIgnore/simulation_data.",run_name,".jld2") outcome functional_HI_all_inds species0_proportion species1_proportion HI_NL_all_inds species0_proportion_NL species1_proportion_NL
                    end
                end
                println(run_name, "  outcome was: ", outcome)
                outcome_array[i,j,k] = outcome
            end # of S_AM loop
        end # of w_hyb loop
    end # of replicate loop

    if save_outcomes_JL
        filename = string("HZAM_Sym_Julia_results_GitIgnore/outcomeArray_set",set_name,"_surv",short_survFitnessMethod,"_ecolDiff",ecolDiff,"_growthrate",intrinsic_R,"_K",K_total,"_FL",total_functional_loci,"_NL",num_neutral_loci,"_gen",max_generations,"_SC",per_reject_cost,".jld2")
        save_object(filename, outcome_array)
    end

    if save_outcomes_csv
        for i in 1:size(outcome_array, 3)
            filename = string("HZAM_Sym_Julia_results_GitIgnore/outcomeArray_set",set_name,"_surv",short_survFitnessMethod,"_ecolDiff",ecolDiff,"_growthrate",intrinsic_R,"_K",K_total,"_FL",total_functional_loci,"_NL",num_neutral_loci,"_gen",max_generations,"_SC",per_reject_cost,"_rep",replications[i])
            CSV.write(filename, Tables.table(outcome_array[:,:,i]), writeheader=false)
        end
    end
    return outcome_array
end

#### functions for summarizing and plotting results

# convert outcome array (an array of strings) to categorical array
function convert_to_cat_array(outcome_array)
    cat_outcome_array = compress(CategoricalArray(outcome_array))
    levels!(cat_outcome_array, ["extinction", "blended", "one_species", "two_species"])
    return cat_outcome_array
end

# function for plotting grid of pie charts showing distribution of four outcomes, using categorical array as input
function plot_all_outcomes(cat_outcome_array)
    num_outcome_types = length(levels(cat_outcome_array))
    outcome_counts = Array{Int16, 2}(undef, num_outcome_types, (size(cat_outcome_array, 1)*size(cat_outcome_array, 2) ))
    for i in 1:size(cat_outcome_array, 1)
        for j in 1:size(cat_outcome_array, 2)
            for outcome_num in 1:num_outcome_types
                outcome_counts[outcome_num, j + (i-1)*size(cat_outcome_array, 2)] = sum(cat_outcome_array[i,j,:] .== levels(cat_outcome_array)[outcome_num])
            end
        end
    end
    colors_of_outcomes = [RGB(0,0,0), plasma[0.525], plasma[0.2], plasma[0.9]] # colors for 4 outcome categories
    pie(outcome_counts, layout = grid(size(cat_outcome_array, 1), size(cat_outcome_array, 2)), legend = false, palette = colors_of_outcomes, margin = -2.0mm)
    plot!(size=(800,1300))
end

function get_most_common_outcomes(cat_outcome_array)
    levels_of_outcomes = levels(cat_outcome_array)
    num_outcome_types = length(levels_of_outcomes)
    most_common_outcomes = CategoricalArray{String, 2}(undef, size(cat_outcome_array, 1), size(cat_outcome_array, 2))
    levels!(most_common_outcomes, ["extinction", "blended", "one_species", "two_species"])
    for i in 1:size(cat_outcome_array, 1)
        for j in 1:size(cat_outcome_array, 2)
            outcome_counts = Vector{Int}(undef, num_outcome_types)
            for outcome_num in 1:num_outcome_types
                outcome_counts[outcome_num] = sum(cat_outcome_array[i,j,:] .== levels(cat_outcome_array)[outcome_num])
            end
            outcomes_with_max_count = findall(outcome_counts .== maximum(outcome_counts))
            if length(outcomes_with_max_count) == 1
                most_common_outcomes[i,j] = levels_of_outcomes[outcomes_with_max_count[1]]
            elseif length(outcomes_with_max_count) >= 2
                # if tie in outcome count, choose one randomly
                most_common_outcomes[i,j] = levels_of_outcomes[sample(outcomes_with_max_count)]
            end
        end
    end
    return most_common_outcomes
end

function plot_common_outcomes(common_outcome_array)
    # make heat map of outcomes
    colors_of_outcomes = [RGB(0,0,0), plasma[0.525], plasma[0.2], plasma[0.9]] # colors for 4 outcome categories
    one_outcome_array = reverse(common_outcome_array, dims=1)
    x_midpoints = log10.([1, 3, 10, 30, 100, 300, 1000, 5000])  # the S_AM values, with Inf convert to 3000 for graphing
    w_hyb_set = [1, 0.98, 0.95, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0] # for just one run, just put one number in this and next line
    y_midpoints = reverse(w_hyb_set)
    min_color = minimum(levelcode.(one_outcome_array)) # this and next line needed to choose the proper colors for the figure
    max_color = maximum(levelcode.(one_outcome_array))
    heatmap(x_midpoints, y_midpoints, one_outcome_array, c = colors_of_outcomes[min_color:max_color], yflip = false, tick_direction = :out, colorbar = false, size = (440,310), framestyle = :box)
    xaxis!("Strength of conspecific mate preference")
    xticklabels = ["1", "3", "10", "30", "100", "300", "1000", "complete"]
    plot!(xticks=(x_midpoints, xticklabels))
    yaxis!("Hybrid fitness")
    yticklabels = ["0.0", "0.1", "0.2", "0.3", "0.4", "0.5", "0.6", "0.7", "0.8", "0.9", "", "", "1.0"]
    plot!(yticks=(y_midpoints, yticklabels), tick_direction = :out)
    # add white lines
    plot!([xlims()[1], xlims()[2]], [0.05, 0.05], linecolor = :white, widen = false, legend = false, linewidth=3)
    x_for_line = mean(x_midpoints[[length(x_midpoints)-1 length(x_midpoints)]])
    plot!([x_for_line, x_for_line], [ylims()[1], ylims()[2]], linecolor = :white, widen = false, legend = false, linewidth=3)
end

function make_and_save_figs(ResultsFolder, RunName, RunOutcomes)
    cat_RunOutcomes = convert_to_cat_array(RunOutcomes)
    display(plot_all_outcomes(cat_RunOutcomes))
    savefig(string(ResultsFolder,"/",RunName,"_AllOutcomes.png"))
    savefig(string(ResultsFolder,"/",RunName,"_AllOutcomes.pdf"))
    most_common_outcomes = get_most_common_outcomes(cat_RunOutcomes)
    display(plot_common_outcomes(most_common_outcomes))
    savefig(string(ResultsFolder,"/",RunName,"_MostCommonOutcomes.png"))
    savefig(string(ResultsFolder,"/",RunName,"_MostCommonOutcomes.pdf"))
end

#### Run the actual simulation by calling the above functions:

# If wanting to run a test to make sure things are working (after above functions are loaded into memory):
# RunOutcomes = run_HZAM("TEST", 1.0, 1.05, 1:3, 100, 50)

ResultsFolder = "HZAM_Sym_Julia_results_GitIgnore" # or change to another folder where results will be saved

# To run simulations under the parameters in the paper, you can execute the following,
# under the heading for each Figure panel in the paper. Because these are stochastic simulations,
# results will differ a little bit (but not much) from those shown in the paper.
# Also, the "make_and_save_figs" function has saved the figure panels under each set of conditions run below,
# and the main script saved the Julia object where the results are stored--this starts with "outcomeArray_set"
# and will be available on Dryad.

# (Note that if the below have "Het" in their name, they were done using the heterozygote disadvantage (underdominance)
# option for survival fitness. If they do not have "Het" in their name, they were done using the epistasis option.)

# Figure 3A
RunName = "JL_fig3aHet"
RunOutcomes = run_HZAM(RunName, 0.0, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage")
# took about 5 hours to run, with 4 threads on a Mac desktop with 4 cores
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 3B
RunName = "JL_fig3bHet"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 4A
RunName = "JL_Fig4aHet"
RunOutcomes = run_HZAM(RunName, 1.0, 1.025, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 4B
RunName = "JL_Fig4bHet"
RunOutcomes = run_HZAM(RunName, 1.0, 1.2, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 4C
RunName = "JL_Fig4cHet"
RunOutcomes = run_HZAM(RunName, 1.0, 2.6, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 1 panel 2
RunName = "JL_Fig6(1,2)Het"
RunOutcomes = run_HZAM(RunName, 0.25, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 1 panel 3
RunName = "JL_Fig6(1,3)Het"
RunOutcomes = run_HZAM(RunName, 0.5, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 1 panel 4
RunName = "JL_Fig6(1,4)Het"
RunOutcomes = run_HZAM(RunName, 0.75, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 2 panel 1
RunName = "JL_Fig6(2,1)Het"
RunOutcomes = run_HZAM(RunName, 0.0, 2.6, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 2 panel 2
RunName = "JL_Fig6(2,2)Het"
RunOutcomes = run_HZAM(RunName, 0.25, 2.6, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 2 panel 3
RunName = "JL_Fig6(2,3)Het"
RunOutcomes = run_HZAM(RunName, 0.5, 2.6, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure 6 row 2 panel 4
RunName = "JL_Fig6(2,4)Het"
RunOutcomes = run_HZAM(RunName, 0.75, 2.6, 1:25;
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S2A (epistasis):
RunName = "JL_fig3a"
RunOutcomes = run_HZAM(RunName, 0.0, 1.05, 1:25)
# this run (which does 2600 simulations) takes about 4 hours (with 4 threads on a somewhat powerful desktop Mac)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S2B (epistasis):
RunName = "JL_fig3b"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S7A
RunName = "JL_fig3bHet125gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    max_generations = 125, survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S7B
RunName = "JL_fig3bHet250gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    max_generations = 250, survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S7C
RunName = "JL_fig3bHet500gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    max_generations = 500, survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S7E
RunName = "JL_fig3bHet2000gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    max_generations = 2000, survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S7F
RunName = "JL_fig3bHet4000gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    max_generations = 4000, survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S8A
RunName = "JL_LikeFig3bHet_butFL1"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    K_total = 1000, max_generations = 1000,
    total_loci = 2, female_mating_trait_loci = 1:1, male_mating_trait_loci = 1:1,
    competition_trait_loci = 1:1, hybrid_survival_loci = 1:1, neutral_loci = 2:2,
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S8C
RunName = "JL_LikeFig3bHet_butFL9"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    K_total = 1000, max_generations = 1000,
    total_loci = 18, female_mating_trait_loci = 1:9, male_mating_trait_loci = 1:9,
    competition_trait_loci = 1:9, hybrid_survival_loci = 1:9, neutral_loci = 10:18,
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S8D
RunName = "JL_LikeFig3bHet_butFL27"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    K_total = 1000, max_generations = 1000,
    total_loci = 54, female_mating_trait_loci = 1:27, male_mating_trait_loci = 1:27,
    competition_trait_loci = 1:27, hybrid_survival_loci = 1:27, neutral_loci = 28:54,
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S9A
RunName = "JL_fig3bHetDiffLociTraitPref"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    K_total = 1000, max_generations = 1000,
    total_loci = 9, female_mating_trait_loci = 1:3, male_mating_trait_loci = 4:6,
    competition_trait_loci = 1:6, hybrid_survival_loci = 1:6, neutral_loci = 7:9,
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S9B
RunName = "JL_fig3bHetDiffLociTraitPrefEcol"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    K_total = 1000, max_generations = 1000,
    total_loci = 12, female_mating_trait_loci = 1:3, male_mating_trait_loci = 4:6,
    competition_trait_loci = 7:9, hybrid_survival_loci = 7:9, neutral_loci = 10:12,
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S10C
RunName = "JL_LikeFig3bHet_but_starting_pop_ratio_0.5"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage",
    starting_pop_ratio = 0.5)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S10D
RunName = "JL_LikeFig4bHet_but_starting_pop_ratio_0.5"
RunOutcomes = run_HZAM(RunName, 1.0, 1.2, 1:25;
    survival_fitness_method = "hetdisadvantage",
    starting_pop_ratio = 0.5)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S10E
RunName = "JL_LikeFig3bHet_but_starting_pop_ratio_0.1"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25;
    survival_fitness_method = "hetdisadvantage",
    starting_pop_ratio = 0.1)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Figure S10F
RunName = "JL_LikeFig4bHet_but_starting_pop_ratio_0.1"
RunOutcomes = run_HZAM(RunName, 1.0, 1.2, 1:25;
    survival_fitness_method = "hetdisadvantage",
    starting_pop_ratio = 0.1)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results):
RunName = "JL_LikeFig3bHet_butFL2"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    K_total = 1000, max_generations = 1000,
    total_loci = 4, female_mating_trait_loci = 1:2, male_mating_trait_loci = 1:2,
    competition_trait_loci = 1:2, hybrid_survival_loci = 1:2, neutral_loci = 3:4,
    survival_fitness_method = "hetdisadvantage")
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_LikeFig3b_butFL1"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 1000,
    2, 1:1, 1:1,
    1:1, 1:1, 2:2)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_LikeFig3b_butFL9"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 1000,
    18, 1:9, 1:9,
    1:9, 1:9, 10:18)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_LikeFig3b_butFL27"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 1000,
    54, 1:27, 1:27,
    1:27, 1:27, 28:54)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_LikeFig4c_butFL27"
RunOutcomes = run_HZAM(RunName, 1.0, 2.6, 1:25,
    1000, 1000,
    54, 1:27, 1:27,
    1:27, 1:27, 28:54)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig4a"
RunOutcomes = run_HZAM(RunName, 1.0, 1.025, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig4b"
RunOutcomes = run_HZAM(RunName, 1.0, 1.2, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig4c"
RunOutcomes = run_HZAM(RunName, 1.0, 2.6, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(1,2)"
RunOutcomes = run_HZAM(RunName, 0.25, 1.05, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(1,3)"
RunOutcomes = run_HZAM(RunName, 0.5, 1.05, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(1,4)"
RunOutcomes = run_HZAM(RunName, 0.75, 1.05, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(2,1)"
RunOutcomes = run_HZAM(RunName, 0.0, 2.6, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(2,2)"
RunOutcomes = run_HZAM(RunName, 0.25, 2.6, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(2,3)"
RunOutcomes = run_HZAM(RunName, 0.5, 2.6, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_Fig6(2,4)"
RunOutcomes = run_HZAM(RunName, 0.75, 2.6, 1:25)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_fig3b125gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 125)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_fig3b250gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 250)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_fig3b500gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000,500)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_fig3b2000gen"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000,2000)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_fig3bDiffLociTraitPref"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 1000,
    9, 1:3, 4:6,
    1:6, 1:6, 7:9)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# Not shown in paper (but available in saved results), with epistasis:
RunName = "JL_fig3bDiffLociTraitPrefEcol"
RunOutcomes = run_HZAM(RunName, 1.0, 1.05, 1:25,
    1000, 1000,
    12, 1:3, 4:6,
    7:9, 7:9, 10:12)
make_and_save_figs(ResultsFolder, RunName, RunOutcomes)

# make graph of epistasis fitness (for Figure 1Cii):
survival_HI = collect(0.0:0.01:1.0)
beta = 1
w_hyb = 0.8
epistasis_fitnesses = 1 .- (1 - w_hyb) .* (4 .* survival_HI .* (1 .- survival_HI)).^beta
using Plots
display(plot(survival_HI, epistasis_fitnesses, ylims = (0,1)))
savefig("plot.pdf")


"""
DR_BlockModel

Block model struture contrains a set of `JuMP.Model` objects, each of which
represents a sub-block model with the information of how to couple these block
models.
"""

mutable struct DR_BlockModel <: DD.AbstractBlockModel
    model::Dict{Int,JuMP.Model} # Dictionary of block models
    coupling_variables::Vector{DD.CouplingVariableRef} # array of variables that couple block models
    variables_by_couple::Dict{Any,Vector{DD.CouplingVariableKey}} # maps `couple_id` to `CouplingVariableKey`

    dual_bound::Float64
    dual_solution::Vector{Float64}

    P_solution::Dict{Int,Float64}

    # TODO: These may be available with heuristics.
    primal_bound::Float64
    primal_solution::Dict{Any, Float64} #coupling_id : value
    combined_weights::Dict{Int, Float64} # block_id : value
    record::Dict{Any, Any}

    function DR_BlockModel()
        return new(
            Dict(),
            [],
            Dict(),
            0.0,
            [],
            Dict(),
            +Inf,
            Dict(),
            Dict(),
            Dict())
    end
end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
"""
    ElectricUtility

ElectricUtility data struct with inner constructor:
```julia
function ElectricUtility(;
    outage_start_time_step::Int=0,  # for modeling a single outage, with critical load spliced into the baseline load ...
    outage_end_time_step::Int=0,  # ... utiltity production_factor = 0 during the outage
    allow_simultaneous_export_import::Bool = true,  # if true the site has two meters (in effect)
    # variables below used for minimax the expected outage cost,
    # with max taken over outage start time, expectation taken over outage duration
    outage_start_time_steps::Array{Int,1}=Int[],  # we minimize the maximum outage cost over outage start times
    outage_durations::Array{Int,1}=Int[],  # one-to-one with outage_probabilities, outage_durations can be a random variable
    outage_probabilities::Array{R,1} where R<:Real = [1.0],
    outage_time_steps::Union{Missing, UnitRange} = isempty(outage_durations) ? missing : 1:maximum(outage_durations),
    scenarios::Union{Missing, UnitRange} = isempty(outage_durations) ? missing : 1:length(outage_durations),
    net_metering_limit_kw::Real = 0,
    interconnection_limit_kw::Real = 1.0e9
)
```julia

!!! note Outage indexing begins at 1 (not 0) and the outage is inclusive of the outage end time step.
    For instance, to model a 3-hour outage from 12AM to 3AM on Jan 1, outage_start_time_step = 1 and outage_end_time_step = 3.
    To model a 1-hour outage from 6AM to 7AM on Jan 1, outage_start_time_step = 7 and outage_end_time_step = 7.

"""
 struct ElectricUtility
    outage_start_time_step::Int  # for modeling a single outage, with critical load spliced into the baseline load ...
    outage_end_time_step::Int  # ... utiltity production_factor = 0 during the outage
    allow_simultaneous_export_import::Bool  # if true the site has two meters (in effect)
    # variables below used for minimax the expected outage cost,
    # with max taken over outage start time, expectation taken over outage duration
    outage_start_time_steps::Array{Int,1}  # we minimize the maximum outage cost over outage start times
    outage_durations::Array{Int,1}  # one-to-one with outage_probabilities, outage_durations can be a random variable
    outage_probabilities::Array{R,1} where R<:Real
    outage_time_steps::Union{Missing, UnitRange}
    scenarios::Union{Missing, UnitRange}
    net_metering_limit_kw::Real
    interconnection_limit_kw::Real

    function ElectricUtility(;
        outage_start_time_step::Int=0,  # for modeling a single outage, with critical load spliced into the baseline load ...
        outage_end_time_step::Int=0,  # ... utiltity production_factor = 0 during the outage
        allow_simultaneous_export_import::Bool = true,  # if true the site has two meters (in effect)
        # variables below used for minimax the expected outage cost,
        # with max taken over outage start time, expectation taken over outage duration
        outage_start_time_steps::Array{Int,1}=Int[],  # we minimize the maximum outage cost over outage start times
        outage_durations::Array{Int,1}=Int[],  # one-to-one with outage_probabilities, outage_durations can be a random variable
        outage_probabilities::Array{R,1} where R<:Real = [1.0],
        outage_time_steps::Union{Missing, UnitRange} = isempty(outage_durations) ? missing : 1:maximum(outage_durations),
        scenarios::Union{Missing, UnitRange} = isempty(outage_durations) ? missing : 1:length(outage_durations),
        net_metering_limit_kw::Real = 0,
        interconnection_limit_kw::Real = 1.0e9
    )

        # Error if outage_start/end_time_step is provided and outage_start_time_steps not empty
        if (outage_start_time_step != 0 || outage_end_time_step !=0) && outage_start_time_steps != []
            throw(@error "Cannot supply singular outage_start(or end)_time_step and multiple outage_start_time_steps. Please use one or the other.")
        end

        return new(
            outage_start_time_step,
            outage_end_time_step,
            allow_simultaneous_export_import,
            outage_start_time_steps,
            outage_durations,
            outage_probabilities,
            outage_time_steps,
            scenarios,
            net_metering_limit_kw,
            interconnection_limit_kw
        )
    end
end


using OpenVX

# FIXME: There's a problem with channel combine node that causes periodic artifacts
# in the output image. After some testing and checking against reference C API
# examples this looks to be a bug in AMD's implementation.

include("imageio.jl")

"""
Creates a graph with one input image and one output image.
You supply the input and output images, it is assumed that the input and output
images are RGB.
"""
function make_filter_graph(ctx, input::VX.vx_image, output::VX.vx_image, width, height)

    num_virt = 6 # number of virtual UInt8/grayscale images that we need

    graph = VX.CreateGraph(ctx)

    virtu8 = Vector{Ptr{Cvoid}}(undef, num_virt)
    nodes = Vector{Ptr{Cvoid}}(undef, 7)
    for i in 1:num_virt
        virtu8[i] = VX.CreateVirtualImage(graph, width, height, VX.DF_IMAGE_U8)
    end

    # Do custom filtering on the input image
    # First, extract R, G, and B channels to individual virtual images
    nodes[1] = VX.ChannelExtractNode(graph, input, VX.CHANNEL_R, virtu8[1])
    nodes[2] = VX.ChannelExtractNode(graph, input, VX.CHANNEL_G, virtu8[2])
    nodes[3] = VX.ChannelExtractNode(graph, input, VX.CHANNEL_B, virtu8[3])

    # Now, run box filter on each of gray scale images
    for i in 1:3
        nodes[3+i] = VX.Gaussian3x3Node(graph, virtu8[i], virtu8[i + 3])
    end

    # Now combine images together in the output color image
    nodes[7] = VX.ChannelCombineNode(graph, virtu8[4], virtu8[5], virtu8[6], C_NULL, output)

    for i in 1:num_virt
        VX.ReleaseImage(Ref(virtu8[i]))
    end

    for i in 1:length(nodes)
        VX.ReleaseNode(Ref(nodes[i]))
    end

    return graph
end

# Main

ctx = VX.CreateContext()

filename = "assets/airplane.bmp"
width = 1200
height = 1200

image = image_from_file(filename, width, height)
output = VX.CreateImage(ctx, width, height, VX.DF_IMAGE_RGB)
graph = make_filter_graph(ctx, image, output, width, height)
VX.VerifyGraph(graph)
VX.ProcessGraph(graph)
file_from_image(output, "airplane_gauss.bmp", width, height)


# Test the creation of SACtr's
using Test
using SAC
@static if VERSION >= v"0.7-"
    import Statistics: mean
end

@testset "Construction" begin
    # Construction from array
    let npts = 1000, delta = 1, b = -5
        local t = rand(SAC.SACFloat, npts)
        local s = SACtr(t, delta)
        @test s.t == t
        @test s.delta == delta
        @test s.npts == npts

        s = SACtr(t, delta, b)
        @test s.b == b
        @test s.e ≈ b + (npts - 1)*delta
        @test s.depmax == maximum(t)
        @test s.depmin == minimum(t)
        @test s.depmen == mean(t)
    end

    # Construction without array
    let npts = rand(1:1000), b = 2, delta = 0.1
        @test SACtr(delta, npts).npts == npts
        @test SACtr(delta, npts).b ≈ 0.0
        @test SACtr(delta, npts, b).b ≈ b
    end

    ## Field modification and access
    # Single traces
    let a = rand(SAC.SACFloat, 100), delta = rand(), s = SACtr(a, delta)
        s[:kevnm] = "ABCD"
        @test s.kevnm == "ABCD"
        @test s[:kevnm] == "ABCD"
        s[:user0] = 1
        @test s.user0 ≈ 1
        @test s[:user0] ≈ 1
        s[:t] = 1:100
        @test s.t ≈ 1:100
        @test s.t ≈ s[:t]
        s[:t] = 2:101
        @test s.t ≈ 2:101
        @test s[:depmin] ≈ 2
        @test s[:depmax] ≈ 101
        @test s[:depmen] ≈ mean(2:101)
        local k = 1
        s[:t] = k
        @test s[:depmin] ≈ k
        @test s[:depmax] ≈ k
        @test all(s[:t] .== SAC.SACFloat(k))
        k = 1:100
        s[:t] = k
        @test all(s[:t] .== SAC.SACFloat.(k))
    end

    # Arrays of traces
    let N = 5, b = -10, delta = 2, npts = 100
        local a = [SACtr(delta, npts, b) for _ in 1:N]
        @test all(a[:b] .≈ b)
        @test all(a[:delta] .≈ delta)
        @test all(a[:npts] .== npts)
        a[:a] = 1:N
        @test all(a[:a] .== SAC.SACFloat.(1:N))
        a[:t] = [1:npts for _ in 1:N]
        @test all(a[:t] .== [SAC.SACFloat.(1:npts) for _ in 1:N])
        local k = -1
        a[:t] = k
        @test all(a[:t] .== [[SAC.SACFloat(k) for j in 1:npts] for i in 1:N])
        a[:kcmpnm] = "ABC"
        @test all(a[:kcmpnm] .== "ABC")
        a[:kcmpnm] = ["a", "b", "c", "d", "e"]
        @test all(a[:kcmpnm] .== ["a", "b", "c", "d", "e"])
    end
end


module mION

JUPL_PATH = ARGS[1] #get path from command line

push!(LOAD_PATH, "$JUPL_PATH/init")
push!(LOAD_PATH, "$JUPL_PATH/src")
using Glob

include("$JUPL_PATH/src/io.jl")
output_files = glob("output/*.csv")
for file in output_files
    rm(file)
end
jpg_files = glob("output/*.jpg")
for file in jpg_files
    rm(file)
end

using Constants
using Get_Sim_Params

include("$JUPL_PATH/init/process_params.jl")
const dr, mi_g, mi, qi, nquiet, cs, CFL, dtm,
        lm, nsmooth, vsmooth, tsmooth = process_params()

include("$JUPL_PATH/src/types.jl")
hydro = hydro_vars(nz, nspec, neqi)

include("$JUPL_PATH/init/init_sim.jl")
const r = init_spacegrid(dr)
const eps_visc = init_art_viscosity(geom)
if !restart
    hydro = init_variables(hydro)
end
hydro = do_smoothing(hydro)
hydro = init_fluxes(hydro)
init_predictor_corrector(hydro)

erf_table = get_erf_integral(JUPL_PATH)
write_sim_parameters()
open_files()
write_data(hydro)
close_files()

include("$JUPL_PATH/src/utils.jl")
include("$JUPL_PATH/src/boundary.jl")
include("$JUPL_PATH/src/numerics.jl")
include("$JUPL_PATH/src/collisions.jl")

nq = 0
j = 0
while tm <= maxTime
    j += 1
    tm += dtm
    is_print = is_print_time(tm, dt_print)

    nq, time1, time2 = is_quiet_time!(time1, time2, tm, tm_quiet, geom, nz, nq, dt_print, dtm)

    apply_BC!(hydro)

    #----- predictor -----------
    #---------------------------
    predictor!(hydro, nq, j)
    update_variables!(hydro.U1D_p, hydro)
    hydro, k_DT, ke, Qextra = calculate_collisions!(hydro, erf_table)
    source_terms!(hydro, k_DT, ke, Qextra, nq)
      # #----start debug
      # if j==1
      #     open_files()
      #     write_data(hydro)
      #     close_files()
      #     write_all_data(hydro)
      #     hydro, k_DT, ke, Qextra = calculate_collisions!(hydro, erf_table)
      #     systemerror(0)
      # end
      # #----end debug

    #----- corrector -----------
    #---------------------------
    corrector!(hydro, nq, j)

    #----- final step ----------
    hydro.U1D[2:nq-1,1:neqi+1] = 0.5 * (hydro.U1D[2:nq-1,1:neqi+1] + hydro.U1D_c[2:nq-1,1:neqi+1]) +
        eps_visc[1:nq-2,1:neqi+1] .*
        (hydro.U1D[3:nq,1:neqi+1] - 2 * hydro.U1D[2:nq-1,1:neqi+1] + hydro.U1D[1:nq-2,1:neqi+1])

    update_variables!(hydro.U1D, hydro)
    hydro, k_DT, ke, Qextra = calculate_collisions!(hydro, erf_table)
    source_terms!(hydro, k_DT, ke, Qextra, nq)

    if is_print
        open_files()
        write_data(hydro)
        close_files()
        is_print = false
    end

end

println("---Simulation ended---")

close_files()

end


# classification for different kinds of hamiltonian.
# algorithm designer will need these type to adopt/reject hamiltonian.

abstract type AbstractHamiltonian end
abstract type AbstractFermionicHamiltonian <: AbstractHamiltonian end
abstract type AbstractSpinlessFermionicHamiltonian <: AbstractFermionicHamiltonian end
abstract type AbstractSpinfulFermionicHamiltonian <: AbstractFermionicHamiltonian end
abstract type AbstractElectronHamiltonian <: AbstractSpinfulFermionicHamiltonian end

abstract type AbstractSpinHamiltonian <: AbstractHamiltonian end

abstract type AbstractBosonicHamiltonian <: AbstractHamiltonian end

export Hamiltonian
# Notice: Hamiltonian is a universal hamiltonian.
struct Hamiltonian{HType<:AbstractHamiltonian,N}
    ex :: HamiltonianExpression
end

const Electron = AbstractSpinfulFermionicHamiltonian
Hamiltonian(ex,htype,N::Int) = Hamiltonian{htype,N}(ex)
Hamiltonian(ex,htype::String,N::Int) = Hamiltonian{eval(Meta.parse(htype)),N}(ex)


using Murmur3
using Base.Test
using JSON

# Test JSON hash examples
f = open(joinpath(dirname(@__FILE__()), "tests.json"))
tests = JSON.parse(readall(f))
close(f)

for test in tests

    if "seed" in keys(test)
        x86_32 = lpad(hex(Murmur3.x86.hash32(test["text"], test["seed"])), 8, "0")
        x86_128 = lpad(hex(Murmur3.x86.hash128(test["text"], test["seed"])), 32, "0")
        x64_128 = lpad(hex(Murmur3.x64.hash128(test["text"], test["seed"])), 32, "0")
    else
        x86_32 = lpad(hex(Murmur3.x86.hash32(test["text"])), 8, "0")
        x86_128 = lpad(hex(Murmur3.x86.hash128(test["text"])), 32, "0")
        x64_128 = lpad(hex(Murmur3.x64.hash128(test["text"])), 32, "0")
    end

    @test test["x86_32"] == x86_32
    @test test["x86_128"] == x86_128
    @test test["x64_128"] == x64_128

end

# Test multiple dispatch
@test Murmur3.x86.hash32(UInt8[1, 2, 3], 5) == UInt32(0x4c298f63)
@test Murmur3.x86.hash32(UInt8[1, 2, 3]) == UInt32(0x80d1d204)
@test Murmur3.x86.hash32("Test", UInt32(5)) == UInt32(0x578de544)
@test Murmur3.x86.hash32("Test", 5) == UInt32(0x578de544)
@test Murmur3.x86.hash32("Test") == UInt32(0x07556ca6)

@test Murmur3.x86.hash128(UInt8[1, 2, 3], 5) == UInt128(0x5c845b3cf61b38eef61b38eef61b38ee)
@test Murmur3.x86.hash128(UInt8[1, 2, 3]) == UInt128(0xf60164e1b5134233b5134233b5134233)
@test Murmur3.x86.hash128("Test", UInt32(5)) == UInt128(0xfca9493ccf980d40cf980d40cf980d40)
@test Murmur3.x86.hash128("Test", 5) == UInt128(0xfca9493ccf980d40cf980d40cf980d40)
@test Murmur3.x86.hash128("Test") == UInt128(0x1bccc51416e53d6c16e53d6c16e53d6c)

@test Murmur3.x64.hash128(UInt8[1, 2, 3], 5) == UInt128(0x88451581d843d60150f5d4a75a40b6c2)
@test Murmur3.x64.hash128(UInt8[1, 2, 3]) == UInt128(0x0e1337a91a643eef3c239a65494e4a40)
@test Murmur3.x64.hash128("Test", UInt64(5)) == UInt128(0x536b3b8cfcdf124c3ea8439527849576)
@test Murmur3.x64.hash128("Test", 5) == UInt128(0x536b3b8cfcdf124c3ea8439527849576)
@test Murmur3.x64.hash128("Test") == UInt128(0x9eac3743ee31bc6fbf2185a737a1b11a)


using POMDPs
using ContinuousPOMDPTreeSearchExperiments
using POMDPToolbox
using BasicPOMCP
using POMCPOW
using ParticleFilters
using CPUTime
using VDPTag
using DataFrames
using ProfileView

pomdp = VDPTagPOMDP()

solvers = Dict{String, Union{Solver,Policy}}(

    "pomcpow" => begin
        ro = ToNextML(mdp(pomdp))
        solver = POMCPOWSolver(tree_queries=10_000_000,
                               criterion=MaxUCB(40.0),
                               final_criterion=MaxTries(),
                               max_depth=20,
                               max_time=0.1,
                               k_action=8.0,
                               alpha_action=1/20,
                               k_observation=4.0,
                               alpha_observation=1/20,
                               estimate_value=0.0,
                               check_repeat_act=true,
                               check_repeat_obs=true,
                               rng=MersenneTwister(13)
                              )
    end,

    "pomcp" => begin
        ro = ToNextML(mdp(pomdp))
        POMCPSolver(max_depth=20,
                    max_time=0.1,
                    c=40.0,
                    tree_queries=typemax(Int),
                    # default_action=TagAction(false, 0.0),
                    estimate_value=0.0,
                    rng=MersenneTwister(17)
                   )
    end
)

n_angles = 10
dpomdp = AODiscreteVDPTagPOMDP(n_angles=n_angles, n_obs_angles=n_angles)

solver = solvers["pomcp"]

planner = solve(solver, dpomdp)
cplanner = translate_policy(planner, dpomdp, pomdp, dpomdp)

filter = ObsAdaptiveParticleFilter(deepcopy(pomdp),
                                   LowVarianceResampler(10_000),
                                   0.05, MersenneTwister(90_000))

@show simulate(RolloutSimulator(max_steps=100), pomdp, cplanner, filter)

Profile.clear()
@profile simulate(RolloutSimulator(max_steps=100), pomdp, cplanner, filter)
ProfileView.view()


module PV_TSSOS

using DynamicPolynomials

using JuMP

using MosekTools

using LinearAlgebra

using LightGraphs

using RowEchelon

export block_compact_POP, block_noncompact_POP, block_noncompact_POP_with_lower_bound, adding_spherical_constraints

include("functions.jl")

end


###############################################################################
#
# Initialize SPACMono structure
#
###############################################################################
"""
    initialize_spac_canopy!(node::SPACMono{FT}) where {FT<:AbstractFloat}

Initialize the RT parameters for a given
- `node` [`SPACMono`](@ref) type struct
"""
function initialize_spac_canopy!(
            node::SPACMono{FT}
) where {FT<:AbstractFloat}
    # 0.1 create variables required
    @unpack angles, envirs, in_rad, leaves_rt, n_canopy, plant_ps, photo_set,
            rt_con, soil_opt, wl_set = node;
    canopy_rt = node.canopy_rt;
    can_opt = node.can_opt;
    can_rad = node.can_rad;
    plant_hs = node.plant_hs;
    fraction_sl::Array{FT,1} = repeat(canopy_rt.lidf, outer=[canopy_rt.nAzi]) /
                               length(canopy_rt.lazitab);
    n_sl = length(canopy_rt.lidf) * length(canopy_rt.lazitab);

    # fluspect the canopy layers
    for ican in 1:n_canopy
        fluspect!(leaves_rt[ican], wl_set);
    end

    # Four Different steps to compute Short-Wave RT
    canopy_geometry!(canopy_rt, angles, can_opt, rt_con)
    canopy_matrices!(leaves_rt, can_opt);
    short_wave!(canopy_rt, can_opt, can_rad, in_rad,
                soil_opt, rt_con);
    canopy_fluxes!(canopy_rt, can_opt, can_rad, in_rad,
                   soil_opt, leaves_rt, wl_set, rt_con);

    # Compute Long Wave (Last term is LW incoming in W m^-2)
    thermal_fluxes!(leaves_rt, can_opt, can_rad, canopy_rt,
                    soil_opt, [FT(400.0)], wl_set);

    # update the canopy leaf area partition information
    for i_can in 1:n_canopy
        rt_layer = n_canopy + 1 - i_can;
        iPS      = plant_ps[i_can];

        # calculate the fraction of sunlit and shaded leaves
        f_view = (can_opt.Ps[rt_layer] + can_opt.Ps[rt_layer+1]) / 2;

        for iLF in 1:n_sl
            iPS.APAR[iLF] = can_rad.absPAR_sunCab[(rt_layer-1)*n_sl + iLF] *
                            FT(1e6);
            iPS.LAIx[iLF] = f_view * fraction_sl[iLF];
        end
        iPS.APAR[end] = can_rad.absPAR_shadeCab[rt_layer] * FT(1e6);
        iPS.LAIx[end] = 1 - f_view;

        update_leaf_TP!(photo_set, iPS, plant_hs.leaves[i_can], envirs[i_can]);

        envirs[i_can].t_air = T_25(FT);
        envirs[i_can].p_sat = saturation_vapor_pressure(T_25(FT));
        envirs[i_can].p_H₂O = envirs[i_can].p_sat / 2;
    end

    return nothing
end

###############################################################################
#
# As the land model is modularized, there could be multiple places that conatin
#     exact the same information. Thus, these information needs to be synced.
#     This type of information includes leaf area, hydraulic traits, and
#     photosynthetic traits.
#
# The following function is meant for leaf area
#
###############################################################################
"""
    update_LAI!(node::SPACMono{FT}, lai::FT) where {FT<:AbstractFloat}

Update SPAC leaf area index
"""
function update_LAI!(node::SPACMono{FT}, lai::FT) where {FT<:AbstractFloat}
    # 1. Update the general information of the SPAC
    node.la = node.ga * lai;

    # 2. Update the LAI for canopy radiation model
    node.canopy_rt.LAI = lai;

    # 3. Update the LA and LAI for stomatal and photosynthesis models
    for _iPS in node.plant_ps
        _iPS.LA   = node.la / node.n_canopy;
        _iPS.LAI  = lai / node.n_canopy;
        _iPS.tLAI = lai;
    end

    # 4. Update the LA for hydraulic model
    #    TODO abstract this to avoid memory allocation
    for _iHS in node.plant_hs.leaves
        _iHS.area = node.la / node.n_canopy;
    end

    return nothing
end

"""
    update_Kmax!(node::SPACMono{FT}, kmax::FT) where {FT<:AbstractFloat}

Update the maximal hydraulic conductance for SPAC

TODO need to abstractize this to PlantHydraulics.jl
"""
function update_Kmax!(node::SPACMono{FT}, kmax::FT) where {FT<:AbstractFloat}
    # Root:Stem:Leaf conductance ratio is set as 1:2:2.
    #    For trees, roots: 2X, trunk: 8X, branch: 8X, leaves: 4X
    #    For palms, roots: 2X, trunk: 4X, leaves: 4X
    #    For grasses, roots: 2X, leaves: 2X

    # 1. update the hydraulic conductance in Roots
    for root in node.plant_hs.roots
        root.k_max      = 2 * kmax / node.plant_hs.n_root;
        root.k_element .= root.k_max * root.N;
    end

    # 2. If is a tree
    if typeof(node.plant_hs) <: TreeLikeOrganism
        node.plant_hs.trunk.k_max      = 8 * kmax;
        node.plant_hs.trunk.k_element .= node.plant_hs.trunk.k_max *
                                         node.plant_hs.trunk.N;
        for stem in node.plant_hs.branch
            stem.k_max      = 8 * kmax / node.plant_hs.n_canopy;
            stem.k_element .= stem.k_max * stem.N;
        end
        for leaf in node.plant_hs.leaves
            leaf.k_sla      = 4 * kmax / node.la;
            leaf.k_element .= leaf.k_sla * leaf.N;
        end
    end

    # 3. If is a palm
    if typeof(node.plant_hs) <: PalmLikeOrganism
        node.plant_hs.trunk.k_max      = 4 * kmax;
        node.plant_hs.trunk.k_element .= node.plant_hs.trunk.k_max *
                                         node.plant_hs.trunk.N;
        for leaf in node.plant_hs.leaves
            leaf.k_sla      = 4 * kmax / node.la;
            leaf.k_element .= leaf.k_sla * leaf.N;
        end
    end

    # 4. If is a grass
    if typeof(node.plant_hs) <: GrassLikeOrganism
        for leaf in node.plant_hs.leaves
            leaf.k_sla      = 2 * kmax / node.la;
            leaf.k_element .= leaf.k_sla * leaf.N;
        end
    end

    return nothing
end

"""
    update_VJRWW!(node::SPACMono{FT}, vcmax::FT) where {FT<:AbstractFloat}

Update Vcmax25(WW), Jmax25(WW), and Rd25(WW) from a given Vcmax25
"""
function update_VJRWW!(node::SPACMono{FT}, vcmax::FT) where {FT<:AbstractFloat}
    # TODO change the ratio accordingly to photo_set
    # TODO add another ratio V2J in photo_set
    # Update Vcmax25, Jmax25 (1.67 Vcmax), and Rd25 (0.015 Vcmax)
    for _iPS in node.plant_ps
        _iPS.ps.Vcmax     = vcmax;
        _iPS.ps.Vcmax25   = vcmax;
        _iPS.ps.Vcmax25WW = vcmax;
        _iPS.ps.Jmax      = vcmax * 1.67;
        _iPS.ps.Jmax25    = vcmax * 1.67;
        _iPS.ps.Jmax25WW  = vcmax * 1.67;
        _iPS.ps.Rd        = vcmax * 0.015;
        _iPS.ps.Rd25      = vcmax * 0.015;
        _iPS.ps.Rd25WW    = vcmax * 0.015;
    end

    return nothing
end

"""
    update_VJR!(node::SPACMono{FT}, ratio::FT) where {FT<:AbstractFloat}

Update effective Vcmax25, Jmax25, and Rd25 from a given tune factor
"""
function update_VJR!(node::SPACMono{FT}, ratio::FT) where {FT<:AbstractFloat}
    # TODO change the ratio accordingly to photo_set
    # TODO add another ratio V2J in photo_set
    # Update Vcmax25, Jmax25 (1.67 Vcmax), and Rd25 (0.015 Vcmax)
    for _iPS in node.plant_ps
        _iPS.ps.Vcmax   = _iPS.ps.Vcmax25WW * ratio;
        _iPS.ps.Vcmax25 = _iPS.ps.Vcmax25WW * ratio;
        _iPS.ps.Jmax    = _iPS.ps.Jmax25WW * ratio;
        _iPS.ps.Jmax25  = _iPS.ps.Jmax25WW * ratio;
        _iPS.ps.Rd      = _iPS.ps.Rd25WW * ratio;
        _iPS.ps.Rd25    = _iPS.ps.Rd25WW * ratio;
    end

    return nothing
end

"""
    update_Cab!(node::SPACMono{FT}, cab::FT) where {FT<:AbstractFloat}

Update leaf chlorophyll content, and then rerun `fluspect!`
"""
function update_Cab!(node::SPACMono{FT}, cab::FT) where {FT<:AbstractFloat}
    for _leaf in node.leaves_rt
        _leaf.Cab = cab;
        fluspect!(_leaf, node.wl_set);
    end

    return nothing
end

"""
    update_Weibull!(node::SPACMono{FT}, b::FT, c::FT) where {FT<:AbstractFloat}

Update Weibull B and C for all components
"""
function update_Weibull!(node::SPACMono{FT}, b::FT) where {FT<:AbstractFloat}
    # 1. update B and C for roots
    for _root in node.plant_hs.roots
        _root.vc.b = b;
    end

    # 2. update B and C for trunk (if exists)
    if !(typeof(node.plant_hs) <: GrassLikeOrganism)
        node.plant_hs.trunk.vc.b = b;
    end

    # 3. update B and C for branches (if exist)
    if typeof(node.plant_hs) <: TreeLikeOrganism
        for _stem in node.plant_hs.branch
            _stem.vc.b = b;
        end
    end

    # 4. update B and C for leaves
    for _leaf in node.plant_hs.leaves
        _leaf.vc.b = b;
    end

    return nothing
end

function update_Weibull!(node::SPACMono{FT}, b::FT, c::FT) where {FT<:AbstractFloat}
    # 1. update B and C for roots
    for _root in node.plant_hs.roots
        _root.vc.b = b;
        _root.vc.c = c;
    end

    # 2. update B and C for trunk (if exists)
    if !(typeof(node.plant_hs) <: GrassLikeOrganism)
        node.plant_hs.trunk.vc.b = b;
        node.plant_hs.trunk.vc.c = c;
    end

    # 3. update B and C for branches (if exist)
    if typeof(node.plant_hs) <: TreeLikeOrganism
        for _stem in node.plant_hs.branch
            _stem.vc.b = b;
            _stem.vc.c = c;
        end
    end

    # 4. update B and C for leaves
    for _leaf in node.plant_hs.leaves
        _leaf.vc.b = b;
        _leaf.vc.c = c;
    end

    return nothing
end


# Create an instance of FAIR used in optimzation and assign it exogenous forcing scenario from NICE.
fair_opt = create_fair_for_opt()
update_param!(fair_opt, :F_exogenous, interpolate_to_annual(bau_model[:radiativeforcing, :forcoth], 10))

# Calculate indices from NICE corresponding to FAIR annual years
nice_decadal_indices = findall((in)(collect(2005:10:2595)), (collect(2005:2595)))

#--------------------------------------------------------------------------#
#------ Carry out optimization for NICE+FAIR with revenue recycling. ------#
#--------------------------------------------------------------------------#

# Create NICE+FAIR with Revenue Recycling objective function with user-specified settings.
fair_recycle_objective = construct_fair_recycle_objective(nice, n_fair_loops, revenue_recycling = true)

# Set upper and lower bounds for global CO₂ tax (upper bound = maximum RICE backstop prices).
upper_bound = maximum(rice_backstop_prices, dims=2)[2:(n_objectives+1)]

# Create an NLopt optimization object for local optimization.
fair_recycle_opt = Opt(local_opt_algorithm, n_objectives)

# Set upper and lower bounds for global CO₂ tax (upper bound = maximum RICE backstop prices).
lower_bounds!(fair_recycle_opt, zeros(n_objectives))
upper_bounds!(fair_recycle_opt, upper_bound)

# Set maximum run time.
maxtime!(fair_recycle_opt, local_stop_time)

# Set convergence tolerance.
ftol_rel!(fair_recycle_opt, local_tolerance)

# Set objective function.
max_objective!(fair_recycle_opt, (x, grad) -> fair_recycle_objective(x))

# Optimize model using local optimization algorithm.
fair_recycle_max_welfare, fair_recycle_opt_tax, fair_recycle_convergence = optimize(fair_recycle_opt, upper_bound .* 0.1)

# ----- Run Model Under Optimal Policy and Save Results ----- #

# Get optimal regional CO₂ mitigation rates from final optimal tax.
fair_recycle_full_opt_tax, fair_recycle_opt_regional_mitigation = mu_from_tax(fair_recycle_opt_tax, rice_backstop_prices)

# Run NICE model under optimal policy.
update_param!(nice, :MIU, fair_recycle_opt_regional_mitigation)
update_param!(nice, :global_carbon_tax, fair_recycle_full_opt_tax)
run(nice)

# Set placeholder

# Now itrate through optimal policy version using FAIR.
for loops = 1:n_fair_loops

    # Run FAIR with annualized NICE CO₂ emissions to calculate temperature.
    update_param!(fair_opt, :E, interpolate_to_annual(nice[:emissions, :E], 10))
    run(fair_opt)

    # Set FAIR temperature projections in NICE.
    set_param!(nice, :TATM, fair_opt[:temperature, :T][nice_decadal_indices])
    run(nice)
end

# Save results (set global optimzation tax to -9999.99 because only doing local optimization).
save_nice_recycle_results(nice, bau_model, fair_recycle_full_opt_tax, ones(length(fair_recycle_full_opt_tax)) .* -9999.99, output_directory, revenue_recycling=true)

#-----------------------------------------------------------------------------#
#------ Carry out optimization for NICE+FAIR without revenue recycling. ------#
#-----------------------------------------------------------------------------#

# Create NICE+FAIR reference objective function with user-specified settings.
fair_reference_objective = construct_fair_recycle_objective(nice, n_fair_loops, revenue_recycling = false)

# Create an NLopt optimization object for local optimization.
fair_reference_opt = Opt(local_opt_algorithm, n_objectives)

# Set upper and lower bounds for global CO₂ tax (upper bound = maximum RICE backstop prices).
lower_bounds!(fair_reference_opt, zeros(n_objectives))
upper_bounds!(fair_reference_opt, upper_bound)

# Set maximum run time.
maxtime!(fair_reference_opt, local_stop_time)

# Set convergence tolerance.
ftol_rel!(fair_reference_opt, local_tolerance)

# Set objective function.
max_objective!(fair_reference_opt, (x, grad) -> fair_reference_objective(x))

# Optimize model using local optimization algorithm.
fair_reference_max_welfare, fair_reference_opt_tax, fair_reference_convergence = optimize(fair_reference_opt, upper_bound .* 0.1)

# ----- Run Model Under Optimal Policy and Save Results ----- #

# Get optimal regional CO₂ mitigation rates from final optimal tax.
fair_reference_full_opt_tax, fair_reference_opt_regional_mitigation = mu_from_tax(fair_reference_opt_tax, rice_backstop_prices)

# Run NICE model under optimal policy.
update_param!(nice, :MIU, fair_reference_opt_regional_mitigation)
update_param!(nice, :global_carbon_tax, fair_reference_full_opt_tax)
run(nice)

# Now itrate through optimal policy version using FAIR.
for loops = 1:n_fair_loops

    # Run FAIR with annualized NICE CO₂ emissions to calculate temperature.
    update_param!(fair_opt, :E, interpolate_to_annual(nice[:emissions, :E], 10))
    run(fair_opt)

    # Set FAIR temperature projections in NICE.
    update_param!(nice, :TATM, fair_opt[:temperature, :T][nice_decadal_indices])
    run(nice)
end

# Save results (set global optimzation tax to -9999.99 because only doing local optimization).
save_nice_recycle_results(nice, bau_model, fair_reference_full_opt_tax, ones(length(fair_recycle_full_opt_tax)) .* -9999.99, output_directory, revenue_recycling=false)


immutable DropoutWork
    states
    statessize
    reservespace
    reservesize
end

function dropout(x, droprate)
    h = handle(x)
    y = similar(x)
    p = Ptr{Void}[0]
    cudnnCreateDropoutDescriptor(p)
    dropoutdesc = p[1]

    statessize_p = Cint[0]
    cudnnDropoutGetStatesSize(h, statessize_p)
    statessize = statessize_p[1]
    states = CuArray{Int8}(Int(statessize))

    xdesc = tensor_desc(x)
    ydesc = tensor_desc(y)
    reservesize_p = Cint[0]
    cudnnDropoutGetReserveSpaceSize(xdesc, reservesize_p)
    reservesize = reservesize_p[1]
    reservespace = CuArray{Int8}(Int(reservesize))

    states_p = Ptr{Void}[0]
    cudnnSetDropoutDescriptor(dropoutdesc, h, Cfloat(droprate), states, statessize, 0)
    cudnnDropoutForward(h, dropoutdesc, xdesc, x, ydesc, y, reservespace, reservesize)

    cudnnDestroyDropoutDescriptor(dropoutdesc)
    cudnnDestroyTensorDescriptor(xdesc)
    cudnnDestroyTensorDescriptor(ydesc)
    y, DropoutWork(states, statessize, reservespace, reservesize)
end

function ∇dropout!(dy, dx, droprate, work::DropoutWork)
    h = handle(dy)
    p = Ptr{Void}[0]
    cudnnCreateDropoutDescriptor(p)
    dropoutdesc = p[1]

    dydesc = tensor_desc(dy)
    dxdesc = tensor_desc(dx)
    cudnnSetDropoutDescriptor(dropoutdesc, h, Cfloat(droprate), work.states, work.statessize, 0)
    cudnnDropoutBackward(h, dropoutdesc, dydesc, dy, dxdesc, dx, work.reservespace, work.reservesize)

    cudnnDestroyDropoutDescriptor(dropoutdesc)
    cudnnDestroyTensorDescriptor(dydesc)
    cudnnDestroyTensorDescriptor(dxdesc)
    dx
end


module SteadyStateDiffEq

using Reexport
@reexport using DiffEqBase

using NLsolve, DiffEqCallbacks
using LinearAlgebra

include("algorithms.jl")
include("solve.jl")

export SSRootfind, DynamicSS

end # module


l = @lens _[begin]
@test l isa Setfield.DynamicIndexLens
obj = (1,2,3)
@test get(obj, l) == 1
@test set(obj, l, true) == (true,2,3)

l = @lens _[2*begin]
@test l isa Setfield.DynamicIndexLens
obj = (1,2,3)
@test get(obj, l) == 2
@test set(obj, l, true) == (1,true,3)

one = 1
plustwo(x) = x + 2
l = @lens _.a[plustwo(begin) - one].b
obj = (a=(1, (a=10, b=20), 3), b=4)
@test get(obj, l) == 20
@test set(obj, l, true) == (a=(1, (a=10, b=true), 3), b=4)