include("autotuning_utils.jl")
include("zygote_rules.jl")
include("forwarddiff_rules.jl")

function update_hyperparameters!(m::GP, state, x, y)
    μ₀ = pr_mean(m.f) # Get prior means
    k = kernel(m.f) # Get kernels
    if !isnothing(opt(m.f))
        if ADBACKEND[] == :Zygote
            hp_state = state.hyperopt_state
            # Compute gradients
            Δμ₀, Δk = Zygote.gradient(μ₀, k) do μ₀, k # Compute gradients for the whole model
                ELBO(m, x, y, μ₀, k)
            end

            # Optimize prior mean
            hp_state = state.hyperopt_state
            if !isnothing(Δμ₀)
                hp_state = update!(μ₀, Δμ₀, hp_state)
            end

            # Optimize kernel parameters
            if isnothing(Δk)
                @warn "Kernel gradients are equal to zero" maxlog = 1
            else
                state_k = update_kernel!(opt(m.f), kernel(m.f), Δk, hp_state)
                hp_state = merge(hp_state, (; state_k))
            end
            state = merge(state, (; hyperopt_state=hp_state))
        elseif ADBACKEND[] == :ForwardDiff
            θ, re = destructure((μ₀, k))
            Δ = ForwardDiff.gradient(θ) do θ
                ELBO(m, re(θ)...)
            end
        end
    end
    # end
    return state
end

# @traitfn function update_hyperparameters!(
#     m::TGP,
# ) where {TGP <: AbstractGPModel; IsFull{TGP}}
#     update_hyperparameters!.(m.f, [xview(m)])
#     setHPupdated!(m.inference, true)
# end

@traitfn function update_hyperparameters!(
    m::TGP, state, x, y
) where {TGP <: AbstractGPModel; IsFull{TGP}}
    if any((!) ∘ isnothing ∘ opt, m.f) # Check there is a least one optimiser
        hp_state = state.hyperopt_state
        μ₀ = pr_means(m) # Get prior means
        ks = kernels(m) # Get kernels
        if ADBACKEND[] == :Zygote
            Δμ₀, Δk = Zygote.gradient(μ₀, ks) do μ₀, ks # Compute gradients for the whole model
                ELBO(m, x, y, μ₀, ks, state)
            end
            # Optimize prior mean
            if !isnothing(Δμ₀)
                hp_state = update!.(μ₀, Δμ₀, hp_state)
            end
            if isnothing(Δk)
                @warn "Kernel gradients are equal to zero" maxlog = 1
            else
                hp_state = map(m.f, Δk, hp_state) do gp, Δ, hp_st
                    if isnothing(opt(gp))
                        return hp_st
                    else
                        state_k = update_kernel!(opt(gp), kernel(gp), Δ, hp_st.state_k)
                        return merge(hp_st, (; state_k))
                    end
                end
            end
        elseif ADBACKEND[] == :ForwardDiff
            θ, re = destructure((μ₀, ks))
            Δ = ForwardDiff.gradient(θ) do θ
                ELBO(m, re(θ)...)
            end
        end
        state = merge(state, (; hyperopt_state=hp_state))
    end
    return state
end

@traitfn function update_hyperparameters!(
    m::TGP, state, x, y
) where {TGP <: AbstractGPModel; !IsFull{TGP}}
    # Check that here is least one optimiser
    if any((!) ∘ isnothing ∘ opt, m.f) || any((!) ∘ isnothing ∘ Zopt, m.f)
        hp_state = state.hyperopt_state
        μ₀ = pr_means(m)
        ks = kernels(m)
        Zs = Zviews(m)
        if ADBACKEND[] == :Zygote
            Δμ₀, Δk, ΔZ = Zygote.gradient(μ₀, ks, Zs) do μ₀, ks, Zs
                ELBO(m, x, y, μ₀, ks, Zs, state)
            end
            # Optimize prior mean
            if !isnothing(Δμ₀)
                hp_state = update!.(μ₀, Δμ₀, hp_state)
            end

            # Optimize kernel parameters
            if isnothing(Δk)
                @warn "Kernel gradients are equal to zero" maxlog = 1
            else
                hp_state = map(m.f, Δk, hp_state) do gp, Δ, hp_st
                    if !isnothing(opt(gp))
                        state_k = update_kernel!(opt(gp), kernel(gp), Δ, hp_st.state_k)
                        return merge(hp_st, (; state_k))
                    else
                        return hp_st
                    end
                end
            end

            # Optimize inducing point locations
            if isnothing(ΔZ)
                @warn "Inducing point locations gradients are equal to zero" maxlog = 1
            else
                hp_state = map(m.f, ΔZ, hp_state) do gp, Δ, hp_st
                    if !isnothing(Zopt(gp))
                        state_Z = update_Z!(Zopt(gp), Zview(gp), Δ, hp_st.state_Z)
                        return merge(hp_st, (; state_Z))
                    else
                        return hp_st
                    end
                end
            end
        elseif ADBACKEND[] == :ForwardDiff
            θ, re = destructure((μ₀, ks, Zs))
            Δ = ForwardDiff.gradient(θ) do θ
                ELBO(m, re(θ)...)
            end
        end
        state = merge(state, (; hyperopt_state=hp_state))
    end
    return state
end

## Update all hyperparameters for the full batch GP models ##
function update_hyperparameters!(gp::AbstractLatent, X::AbstractVector)
    if !isnothing(gp.opt)
        f_l, f_μ₀ = hyperparameter_gradient_function(gp, X)
        ad_backend = K_ADBACKEND[] == :auto ? ADBACKEND[] : K_ADBACKEND[]
        Δμ₀ = f_μ₀()
        Δk = if ad_backend == :forward
            ∇L_ρ_forward(f_l, gp, X)
        elseif ad_backend == :zygote
            ∇L_ρ_zygote(f_l, gp, X)
        end
        apply_Δk!(gp.opt, kernel(gp), Δk) # Apply gradients to the kernel parameters
        apply_gradients_mean_prior!(pr_mean(gp), Δμ₀, X)
    end
end

## Update all hyperparameters for the sparse variational GP models ##
function update_hyperparameters!(
    gp::AbstractLatent,
    X::AbstractVector,
    ∇E_μ::AbstractVector,
    ∇E_Σ::AbstractVector,
    i::AbstractInference,
    vi_opt::InferenceOptimizer,
)
    Δμ₀, Δk = if !isnothing(gp.opt)
        f_ρ, f_Z, f_μ₀ = hyperparameter_gradient_function(gp)
        ad_backend = K_ADBACKEND[] == :auto ? ADBACKEND[] : K_ADBACKEND[]
        Δμ₀ = f_μ₀()
        Δk = if ad_backend == :forward
            ∇L_ρ_forward(f_ρ, gp, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        elseif ad_backend == :zygote
            ∇L_ρ_zygote(f_ρ, gp, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        end
        (Δμ₀, Δk)
    else
        nothing, nothing
    end
    if !isnothing(Zopt(gp))
        ad_backend = Z_ADBACKEND[] == :auto ? ADBACKEND[] : Z_ADBACKEND[]
        Z_grads = if ad_backend == :forward
            Z_gradient_forward(gp, f_Z, X, ∇E_μ, ∇E_Σ, i, vi_opt) #Compute the gradient given the inducing points location
        elseif ad_backend == :zygote
            Z_gradient_zygote(gp, f_Z, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        end
        update_Z!(opt(gp.Z), gp.Z, Z_grads) #Apply the gradients on the location
    end
    if !all([isnothing(Δk), isnothing(Δμ₀)])
        apply_Δk!(gp.opt, kernel(gp), Δk) # Apply gradients to the kernel parameters
        apply_gradients_mean_prior!(pr_mean(gp), Δμ₀, X)
    end
end

function update_hyperparameters!(
    gp::AbstractLatent,
    l::AbstractLikelihood,
    i::AbstractInference,
    X::AbstractVector,
    Y::AbstractVector,
)
    Δμ₀, Δk = if !isnothing(gp.opt)
        ad_backend = K_ADBACKEND[] == :auto ? ADBACKEND[] : K_ADBACKEND[]
        (Δμ₀, Δk) = if ad_backend == :forward
            ∇L_ρ_forward(f_ρ, gp, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        elseif ad_backend == :zygote
            ∇L_ρ_zygote(gp, l, i, X, Y)
        end
    else
        nothing, nothing
    end
    if !isnothing(opt(gp.Z))
        ad_backend = Z_ADBACKEND[] == :auto ? ADBACKEND[] : Z_ADBACKEND[]
        global Z_grads = if ad_backend == :forward
            Z_gradient_forward(gp, f_Z, X, ∇E_μ, ∇E_Σ, i, vi_opt) #Compute the gradient given the inducing points location
        elseif ad_backend == :zygote
            Z_gradient_zygote(gp, l, i, X, Y)
        end
        update!(opt(gp.Z), gp.Z.Z, Z_grads) #Apply the gradients on the location
    end
    if !all([isnothing(Δk), isnothing(Δμ₀)])
        apply!(kernel(gp), Δk, gp.opt) # Apply gradients to the kernel parameters
        update!(pr_mean(gp), Δμ₀, X)
    end
end

## Return the derivative of the KL divergence between the posterior and the GP prior ##
function hyperparameter_KL_gradient(J::AbstractMatrix, A::AbstractMatrix)
    return trace_ABt(J, A) / 2
end

function hyperparameter_gradient_function(gp::LatentGP{T}, ::AbstractVector) where {T}
    A = (inv(cov(gp)) - mean(gp) * transpose(mean(gp))) # μ = inv(K+σ²)*(y-μ₀)
    return (function (Jnn)
        return -hyperparameter_KL_gradient(Jnn, A)
    end, function ()
        return -mean(gp)
    end)
end

## Return functions computing gradients of the ELBO given the kernel hyperparameters for a non-sparse model ##
function hyperparameter_gradient_function(
    gp::VarLatent{T}, X::AbstractVector
) where {T<:Real}
    μ₀ = pr_mean(gp, X)
    A =
        (I - pr_cov(gp) \ (cov(gp) + (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))) /
        pr_cov(gp)
    return (
        function (Jnn)
            return -hyperparameter_KL_gradient(Jnn, A)
        end,
        function ()
            return pr_cov(gp) \ (mean(gp) - μ₀)
        end,
    )
end

## Return functions computing gradients of the ELBO given the kernel hyperparameters for a non-sparse latent GP ##
function hyperparameter_gradient_function(gp::SparseVarLatent{T}) where {T<:Real}
    μ₀ = pr_mean(gp, gp.Z)
    A =
        (I(dim(gp)) - pr_cov(gp) \ (cov(gp) + (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))) /
        cov(gp)
    κΣ = gp.κ * cov(gp)
    return (
        function (Jmm, Jnm, Jnn, ∇E_μ, ∇E_Σ, i, opt)
            return (
                hyperparameter_expec_gradient(gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, Jnn) -
                hyperparameter_KL_gradient(Jmm, A)
            )
        end,
        function (Jmm, Jnm, ∇E_μ, ∇E_Σ, i, opt)
            return hyperparameter_expec_gradient(
                gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, zero(gp.K̃)
            ) - hyperparameter_KL_gradient(Jmm, A)
        end,
        function ()
            return pr_cov(gp) \ (mean(gp) - μ₀)
        end,
    )
end

function hyperparameter_gradient_function(
    gp::TVarLatent{T}, X::AbstractVector
) where {T<:Real}
    μ₀ = pr_mean(gp, X)
    A =
        (
            I(dim(gp)) .-
            (pr_cov(gp)) \ (cov(gp) .+ (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))
        ) / (pr_cov(gp))
    return (
        function (Jnn)
            return -hyperparameter_KL_gradient(prior(gp).χ * Jnn, A)
        end,
        function ()
            return (pr_cov(gp)) \ (mean(gp) - μ₀)
        end,
    )
end

function hyperparameter_gradient_function(gp::OnlineVarLatent{T}) where {T<:Real}
    μ₀ = pr_mean(gp, gp.Z)
    A =
        (I(dim(gp)) - pr_cov(gp) \ (cov(gp) + (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))) /
        pr_cov(gp)
    κΣ = gp.κ * cov(gp)
    κₐΣ = gp.κₐ * cov(gp)
    return (
        function (Jmm, Jnm, Jnn, Jab, Jaa, ∇E_μ, ∇E_Σ, i, opt)
            ∇E = hyperparameter_expec_gradient(gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, Jnn)
            ∇KLₐ = hyperparameter_online_gradient(gp, κₐΣ, Jmm, Jab, Jaa)
            ∇KL = hyperparameter_KL_gradient(Jmm, A)
            return ∇E + ∇KLₐ - ∇KL
        end, # Function gradient given kernel parameters
        function (Jmm, Jnm, Jab, ∇E_μ, ∇E_Σ, i, opt)
            return hyperparameter_expec_gradient(
                gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, zero(gp.K̃)
            ) + hyperparameter_online_gradient(
                gp, κₐΣ, Jmm, Jab, zeros(T, length(gp.Zₐ), length(gp.Zₐ))
            ) - hyperparameter_KL_gradient(Jmm, A)
        end, # Function gradient given inducing points locations
        function ()
            return -(pr_cov(gp) \ (μ₀ - mean(gp)))
        end,
    ) # Function gradient given mean prior
end

## Gradient with respect to hyperparameter for analytical VI ##
function hyperparameter_expec_gradient(
    gp::AbstractLatent,
    ∇E_μ::AbstractVector{<:Real},
    ∇E_Σ::AbstractVector{<:Real},
    i::AnalyticVI,
    opt::AVIOptimizer,
    κΣ::AbstractMatrix{<:Real},
    Jmm::AbstractMatrix{<:Real},
    Jnm::AbstractMatrix{<:Real},
    Jnn::AbstractVector{<:Real},
)
    ι = (Jnm - gp.κ * Jmm) / pr_cov(gp)
    J̃ = Jnn - (diag_ABt(ι, gp.Knm) + diag_ABt(gp.κ, Jnm))
    dμ = dot(∇E_μ, ι * mean(gp))
    dΣ = -dot(∇E_Σ, J̃)
    dΣ += -dot(∇E_Σ, 2.0 * (diag_ABt(ι, κΣ)))
    dΣ += -dot(∇E_Σ, 2.0 * (ι * mean(gp)) .* (gp.κ * mean(gp)))
    return ρ(i) * (dμ + dΣ)
end

## Gradient with respect to hyperparameters for numerical VI ##
function hyperparameter_expec_gradient(
    gp::AbstractLatent,
    ∇E_μ::AbstractVector{<:Real},
    ∇E_Σ::AbstractVector{<:Real},
    i::NumericalVI,
    opt::NVIOptimizer,
    κΣ::AbstractMatrix{<:Real},
    Jmm::AbstractMatrix{<:Real},
    Jnm::AbstractMatrix{<:Real},
    Jnn::AbstractVector{<:Real},
)
    ι = (Jnm - gp.κ * Jmm) / pr_cov(gp)
    J̃ = Jnn - (diag_ABt(ι, gp.Knm) + diag_ABt(gp.κ, Jnm))
    dμ = dot(∇E_μ, ι * mean(gp))
    dΣ = dot(∇E_Σ, J̃ + 2.0 * diag_ABt(ι, κΣ))
    return ρ(i) * (dμ + dΣ)
end

function hyperparameter_online_gradient(
    gp::AbstractLatent,
    κₐΣ::Matrix{<:Real},
    Jmm::AbstractMatrix{<:Real},
    Jab::AbstractMatrix{<:Real},
    Jaa::AbstractMatrix{<:Real},
)
    ιₐ = (Jab - gp.κₐ * Jmm) / pr_cov(gp)
    trace_term =
        -sum(
            trace_ABt.(
                [gp.invDₐ],
                [
                    Jaa,
                    2 * ιₐ * transpose(κₐΣ),
                    -ιₐ * transpose(gp.Kab),
                    -gp.κₐ * transpose(Jab),
                ],
            ),
        ) / 2
    term_1 = dot(gp.prevη₁, ιₐ * mean(gp))
    term_2 = -dot(ιₐ * mean(gp), gp.invDₐ * gp.κₐ * mean(gp))
    return trace_term + term_1 + term_2
end


module NLSat

export NLSolver, Z3Solver, SMTSolverZ3, SMTSolverYices, SMTSolver, CFiniteSolver
export NLStatus, NLModel
export add_vars!, add!, add_soft!, solve

using DataStructures
using DelimitedFiles
using Distributed
using MacroTools: walk, postwalk, @capture, @match, replace
using Dates
using Z3

const NLModel = Dict{Symbol,<:Number}

# ------------------------------------------------------------------------------

abstract type AlgebraicNumber end
export AlgebraicNumber

function AlgebraicNumber(x::Expr, n::Int)
    @warn "Algebraic numbers not yet supported, got $((x, n)), returning $(NaN)"
    NaN
end

# ------------------------------------------------------------------------------

include("../utils.jl")
include("clauseset.jl")
include("smtlib.jl")
include("lisp.jl")
include("z3utils.jl")

# ------------------------------------------------------------------------------

@enum NLStatus sat unsat unknown timeout

abstract type NLSolver end

function add_vars!(s::NLSolver, d::Dict{Symbol,Type}) end
function add!(s::NLSolver, cs::ClauseSet) end
function add_soft!(s::NLSolver, cs::ClauseSet) end
function solve(s::NLSolver; timeout::Int = -1) end

# ------------------------------------------------------------------------------

# include("mathematica.jl")
include("z3solver.jl")
include("smtsolver.jl")
include("cfinitesolver.jl")

end # module


function _info(model::Optimizer, key::MOI.VariableIndex)
    if haskey(model.variable_info, key)
        return model.variable_info[key]
    end
    throw(MOI.InvalidIndex(key))
    return
end

function _make_var(model::Optimizer, variable::Variable)
    # Initialize `VariableInfo` with a dummy `VariableIndex` and a column,
    # because we need `add_item` to tell us what the `VariableIndex` is.
    index = CleverDicts.add_item(
        model.variable_info,
        VariableInfo(MOI.VariableIndex(0), variable),
    )
    _info(model, index).index = index
    return index
end

function _make_var(
    model::Optimizer,
    variable::Variable,
    set::MOI.AbstractScalarSet,
)
    index = _make_var(model, variable)
    return index,
    MOI.ConstraintIndex{MOI.VariableIndex, typeof(set)}(index.value)
end

function _make_vars(model::Optimizer, variables::Vector{<:Variable})
    # Barely used, because add_constrained_variables may have variable sets (except for AbstractVectorSet).
    # Only implemented in the unconstrained case.
    indices = Vector{MOI.VariableIndex}(undef, length(variables))
    for i in 1:length(variables)
        indices[i] = CleverDicts.add_item(
            model.variable_info,
            VariableInfo(MOI.VariableIndex(0), variables[i]),
        )
        _info(model, indices[i]).index = indices[i]
    end
    return indices
end

function _make_numvar(
    model::Optimizer,
    set::MOI.AbstractScalarSet;
    lb::Float64=-IloInfinity,
    ub::Float64=IloInfinity,
)
    return _make_var(model, cpo_java_numvar(model.inner, lb, ub), set)
end

function _make_intvar(
    model::Optimizer,
    set::MOI.AbstractScalarSet;
    lb::Int32=-IloMinInt,
    ub::Int32=IloMaxInt,
)
    vindex, cindex = _make_var(
        model,
        cpo_java_intvar(model.inner, Int32(lb), Int32(ub)),
        set,
    )
    _info(model, vindex).type = INTEGER
    return vindex, cindex
end

function _make_boolvar(model::Optimizer, set::MOI.AbstractScalarSet)
    vindex, cindex = _make_var(model, cpo_java_boolvar(model.inner), set)
    _info(model, vindex).type = BINARY
    return vindex, cindex
end

function supports_add_constrained_variables(
    ::Optimizer,
    ::Type{F},
) where {
    F <: Union{
        MOI.EqualTo{Float64},
        MOI.LessThan{Float64},
        MOI.GreaterThan{Float64},
        MOI.Interval{Float64},
        MOI.EqualTo{Int},
        MOI.LessThan{Int},
        MOI.GreaterThan{Int},
        MOI.Interval{Int},
        MOI.ZeroOne,
        MOI.Integer,
    },
}
    return true
end

function MOI.add_variable(model::Optimizer)
    return _make_var(
        model,
        cpo_java_numvar(model.inner, -IloInfinity, IloInfinity),
    )
end

function MOI.add_variables(model::Optimizer, N::Int)
    return _make_vars(
        model,
        cpo_java_numvararray(model.inner, Int32(N), -IloInfinity, IloInfinity),
    )
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.GreaterThan{T},
) where {T <: Real}
    return _make_numvar(model, set, lb=set.lower)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.LessThan{T},
) where {T <: Real}
    return _make_numvar(model, set, lb=set.upper)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.EqualTo{T},
) where {T <: Real}
    return _make_numvar(model, set, lb=set.value, ub=set.value)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.Interval{T},
) where {T <: Real}
    return _make_numvar(model, set, lb=set.lower, ub=set.upper)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.GreaterThan{T},
) where {T <: Integer}
    return _make_intvar(model, set, lb=set.lower)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.LessThan{T},
) where {T <: Integer}
    return _make_intvar(model, set, lb=set.upper)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.EqualTo{T},
) where {T <: Integer}
    return _make_intvar(model, set, lb=set.value, ub=set.value)
end

function MOI.add_constrained_variable(
    model::Optimizer,
    set::MOI.Interval{T},
) where {T <: Integer}
    return _make_intvar(model, set, lb=set.lower, ub=set.upper)
end

function MOI.add_constrained_variable(model::Optimizer, set::MOI.ZeroOne)
    return _make_boolvar(model, set)
end

function MOI.add_constrained_variable(model::Optimizer, set::MOI.Integer)
    return _make_intvar(model, set)
end

function MOI.is_valid(model::Optimizer, v::MOI.VariableIndex)
    return haskey(model.variable_info, v)
end

function MOI.delete(model::Optimizer, v::MOI.VariableIndex)
    # Actual deletion of the variable.
    info = _info(model, v)
    if info.n_socs > 0
        throw(MOI.DeleteNotAllowed(v))
    end
    cpo_java_remove(model.inner, info.variable)
    delete!(model.variable_info, v)

    # We throw away name_to_constraint so we will rebuild VariableIndex
    # constraint names (without v) on the next access to it.
    model.name_to_variable = nothing
    model.name_to_constraint = nothing

    return
end

function MOI.get(model::Optimizer, ::Type{MOI.VariableIndex}, name::String)
    if model.name_to_variable === nothing
        _rebuild_name_to_variable(model)
    end
    return get(model.name_to_variable, name, nothing)
end

function _rebuild_name_to_variable(model::Optimizer)
    model.name_to_variable = Dict{String, MOI.VariableIndex}()
    for (index, info) in model.variable_info
        if info.name == ""
            continue
        end

        if haskey(model.name_to_variable, info.name)
            error("Duplicate variable name detected: $(info.name)")
        end
        model.name_to_variable[info.name] = index
    end
    return
end

function MOI.get(model::Optimizer, ::MOI.VariableName, v::MOI.VariableIndex)
    return _info(model, v).name
end

function MOI.set(
    model::Optimizer,
    ::MOI.VariableName,
    v::MOI.VariableIndex,
    name::String,
)
    info = _info(model, v)
    info.name = name
    cpo_java_addable_setname(model.inner, info.variable, name)

    model.name_to_variable = nothing
    model.name_to_constraint = nothing

    return
end


mutable struct prior_zpr_t <: LCMType
  utime::Int64 # unix time associated with this pose
  id::Int64    # pose id (sequential, 0-based)

  z::Float64
  pitch::Float64
  roll::Float64

  var_z::Float64
  var_pitch::Float64
  var_roll::Float64
end

@lcmtypesetup(prior_zpr_t)


"""
    prepare_observables(m::HoppingBFModel)

Initializes observables for the Hopping BF model and returns a `Dict{String, Observable}`.

See also [`measure_observables!`](@ref) and [`finish_observables!`](@ref).
"""
@inline function prepare_observables(mc::DQMC_bond, m::HoppingBFModel)
    obs = Dict{String,Observable}()
    obs["confs"] = Observable(HBFConf, "Configurations")
    obs["greens"] = Observable(typeof(mc.s.greens), "Equal-times Green's function")
    obs["Eboson"] = Observable(Float64, "Bosonic energy")

    return obs
end

"""
    measure_observables!(m::HoppingBFModel, obs::Dict{String,Observable}, conf::HBFConf, E::Float64)

Measures observables and updates corresponding `Observable` objects in `obs`.

See also [`prepare_observables`](@ref) and [`finish_observables!`](@ref).
"""
@noinline function measure_observables!(mc::DQMC_bond, m::HoppingBFModel,
							obs::Dict{String,Observable}, conf::HBFConf)
    push!(obs["confs"], mc.conf)
    push!(obs["greens"], mc.s.greens)
    push!(obs["Eboson"], energy_boson(mc, m, conf))
    nothing
end

"""
    measure_observables!(mc::DQMC_bond, m::HoppingBFModel, obs::Dict{String,Observable})

Finish measurements of observables.

See also [`prepare_observables`](@ref) and [`measure_observables!`](@ref).
"""
@inline function finish_observables!(mc::DQMC_bond, m::HoppingBFModel,
							obs::Dict{String,Observable})
    nothing
end


@testset "simple workflow" begin
    env = CartPoleEnv{Float32}() |> StateOverriddenEnv(deepcopy)
    agent = Agent(;
        policy = RandomPolicy(env),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = 10_000,
            state_type = Float32,
            state_size = (4,),
        ),
    )
    N_EPISODE = 10000
    hook = TotalRewardPerEpisode()
    run(agent, env, StopAfterEpisode(N_EPISODE), hook)

    @test isapprox(sum(hook.rewards) / N_EPISODE, 21; atol = 2)
end


#! format: off

using PowerSystems
using PowerSimulations
using PowerSystemCaseBuilder

using HiGHS # mip solver
solver = optimizer_with_attributes(HiGHS.Optimizer, "mip_rel_gap" => 0.5)

sys = build_system(PSITestSystems, "modified_RTS_GMLC_DA_sys")

for line in get_components(Line, sys)
    if (get_base_voltage(get_from(get_arc(line))) >= 230.0) &&
       (get_base_voltage(get_to(get_arc(line))) >= 230.0)
        #if get_area(get_from(get_arc(line))) != get_area(get_to(get_arc(line)))
        @info "Changing $(get_name(line)) to MonitoredLine"
        convert_component!(MonitoredLine, line, sys)
    end
end

template = template_unit_commitment(network = PTDFPowerModel)

set_device_model!(template, MonitoredLine, StaticBranch)

set_device_model!(template, Line, StaticBranchUnbounded)

uc_prob = DecisionModel(template, sys, horizon = 24, optimizer = solver)
build!(uc_prob, output_dir = mktempdir())

solve!(uc_prob)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


mutable struct MultiPtsTurningDriver <: MultiPtsDriver
    Δt::Float64
    siji::AccSteerDriver
    N::Int #number of points
    Pts::Matrix{Float64}  # 2×N
    steermax::Float64 # pi/2;
    steerdotmax::Float64 # 20;
    accmax::Float64 # 10;

    commands::Array{Float64} #4x1
    index::Int
    subindex::Int
    laneNum_desire::Int
    acc_desire::Float64
    v_desire::Float64
    v_max::Float64 #20m/s
    v_min::Float64 #0
    turning_direction::Int

    color::Colorant

    MultiPtsTurningDriver(Δt::Float64; siji::AccSteerDriver=AccSteerDriver(0.0,0.0),N::Int=20,Pts::Matrix{Float64}=zeros(2,20),steermax::Float64=pi/4,steerdotmax::Float64=20.0,accmax::Float64=10.0,commands::Array{Float64}=zeros(4),index::Int=0,subindex::Int=0,laneNum_desire::Int=1,acc_desire::Float64=0.0,v_desire::Float64=0.0,v_max::Float64=20.0,v_min::Float64=0.0,turning_direction::Int=1,color::Colorant=RGB(1, 0, 0))=new(Δt,siji,N,Pts,steermax,steerdotmax,accmax,commands,index,subindex,laneNum_desire,acc_desire,v_desire,v_max,v_min,turning_direction,color)
end

AutomotiveDrivingModels.get_name(model::MultiPtsTurningDriver) = "MultiPtsTurningDriver"

Base.rand(model::MultiPtsTurningDriver) = rand(model.siji)
Distributions.pdf(model::MultiPtsTurningDriver, a::Float64) = a == model.siji.acc ? Inf : 0.0
Distributions.logpdf(model::MultiPtsTurningDriver, a::Float64) = a == model.siji.acc ? Inf : -Inf

function excute_action!(model::MultiPtsTurningDriver,acc::Float64,direction::Int,scene::Union{Scene,Frame{Entity{VehicleState, BicycleModel, Int}}}, roadway::Roadway, ego_index::Int)
    veh=scene[ego_index]
    laneNum=veh.state.posF.roadind.tag.lane
    model.acc_desire=acc
    model.laneNum_desire=laneNum+direction
    model.v_desire=veh.state.v
    model.index=0
    model.subindex=0
end

function set_pts!(model::MultiPtsTurningDriver,PtsIn::Matrix{Float64})
    model.Pts=PtsIn
    model.index=0
    model.subindex=0
    model
end

function AutomotiveDrivingModels.observe!(model::MultiPtsTurningDriver, scene::Union{Scene,Frame{Entity{VehicleState, BicycleModel, Int}}}, roadway::Roadway, egoid::Int)
    ego_index = findfirst(scene, egoid)
    get_acc_steer!(model,scene,roadway,ego_index)
    model
end

function get_acc_steer!(model::MultiPtsTurningDriver,scene::Union{Scene,Frame{Entity{VehicleState, BicycleModel, Int}}}, roadway::Roadway, ego_index::Int)
    if model.index==6 || model.index>model.N
    #if model.index>model.N
        #println("index reach ", model.index)
        set_lane_changing_pts_with_direction!(model, scene, roadway, ego_index,model.turning_direction)
    end
    if model.index==0
        #println("index is 0!")
        set_lane_changing_pts_with_direction!(model, scene, roadway, ego_index,model.turning_direction)
    else
        if model.subindex==0
            #compute new command
            get_commands!(model,scene, roadway,ego_index)
            model.siji.acc=model.commands[1]
            model.siji.steer=model.commands[3]
            model.siji.direction=model.turning_direction
            model.subindex=1
        elseif model.subindex==1
            #follow the commands
            model.siji.acc=model.commands[2]
            model.siji.steer=model.commands[4]
            model.siji.direction=model.turning_direction
            model.subindex=0
            model.index=model.index+1
        end
    end
end


"""

```
mddf_linkedcells_parallel(trajectory::Trajectory, options::Options,samples::Samples, RNG, mddf_compute!)
```

Computes the MDDF using linked cells, in parallel.

"""
function mddf_linkedcells_parallel(trajectory::Trajectory, options::Options,
                                   samples::Samples, RNG, mddf_compute!)

  # Simplify code by assigning some shortened names
  solute = trajectory.solute
  solvent = trajectory.solvent
  x_solute = trajectory.x_solute
  x_solvent = trajectory.x_solvent

  # Number of threads
  if options.nthreads < 0
    nspawn = Threads.nthreads() - 1
  else
    nspawn = options.nthreads - 1
  end
  if nspawn == 0
    error(" Parallel version must be executed only with more than 1 thread. ")
  end

  # Initializing the structure that carries the result per thread
  R0 = Result(trajectory,options)
  R = [ Result(trajectory,options,irefatom=R0.irefatom) for i in 1:nspawn ]

  # Safe passing of frame counter to the threads
  tframe = zeros(Int,nspawn)

  # Create data structures required for multithreading
  framedata = Vector{FrameData}(undef,nspawn)
  for ispawn in 1:nspawn
    framedata[ispawn] = FrameData(deepcopy(trajectory),R[1])
  end

  # Tasks
  t = Vector{Task}(undef,nspawn)
  free = ones(Bool,nspawn)

  # Skip initial frames if desired
  iframe = 0 # Counter for all frames of the input file
  while iframe < options.firstframe-1
    nextframe!(trajectory)
    iframe += 1
  end

  # Print some information about this run
  options.silent || title(R[1],solute,solvent,nspawn)

  ndone = 0
  iframe_read = 0 # Counter for the frames that are actually being considered
  if ! options.silent
    progress = Progress(R[1].nframes_read,1)
  end
  while ndone < R[1].nframes_read

    # Launch for each free thread the computation of a frame
    while iframe_read < R[1].nframes_read && count(free) > 0
      ifree = findfirst(x->x==true,free)
      # Skip frames that are not going to be considered
      while (iframe+1)%options.stride != 0
        nextframe!(trajectory)
        iframe += 1
      end
      # Actually read the data of the frame that will be considered
      nextframe!(trajectory)
      iframe += 1
      iframe_read += 1
      tframe[ifree] = iframe
      @. framedata[ifree].trajectory.x_solute = trajectory.x_solute
      @. framedata[ifree].trajectory.x_solvent = trajectory.x_solvent
      @. framedata[ifree].trajectory.sides = trajectory.sides
      # Spawn the calculations for this frame
      t[ifree] = ThreadPools.@tspawnat ifree+1 mddf_compute!(tframe[ifree],framedata[ifree],options,RNG,R[ifree])
      free[ifree] = false
    end

    # Wait a little bit before checking
    sleep(options.sleep)

    # Check thread status
    for ispawn in 1:nspawn
      if ! free[ispawn]
        if istaskfailed(t[ispawn])
          error(" Computation of MDDF failed in thread: $ispawn", fetch(t[ispawn]))
        end
        if istaskdone(t[ispawn])
          ndone += 1
          free[ispawn] = true
          options.silent || next!(progress)
          if options.GC && (Sys.free_memory() / Sys.total_memory() < options.GC_threshold)
            GC.gc() # why we need this anyway??? There should not be so much garbage.
          end
        end
      end
    end

  end
  closetraj(trajectory)

  # Sum up the results of all threads into the data of thread one (R1<-R1+R2)
  for ispawn in 2:nspawn
    sum!(R[1],R[ispawn])
  end

  # Setup the final data structure with final values averaged over the number of frames,
  # sampling, etc, and computes final distributions and integrals
  finalresults!(R[1],options,trajectory,samples)
  options.silent || println(bars)

  return R[1]

end


const OFFSET = 100.0 # offset in degrees!!!!
touch(string(getpid(),".pid"))
push!(LOAD_PATH, pwd())
using UnicodePlots, PyCall, Reactive, DataFrames, Light
@pyimport seabreeze
seabreeze.use("pyseabreeze")
@pyimport seabreeze.spectrometers as sb
devices = sb.list_devices()
s = sb.Spectrometer(devices[1])
maxy = s[:_dev][:interface][:_MAX_PIXEL_VALUE]
wl = s[:wavelengths]()
nwl = length(wl)
function watchit()
	watch_file("integrationtime")
	sleep(.1)
	put!(it, parse(Int,readchomp("integrationtime")))
end
function watchnr()
	watch_file("nreplicates")
	sleep(.1)
	take!(nr)
	put!(nr, parse(Int,readchomp("nreplicates")))
end
function watchr()
	watch_file("rotation")
	sleep(.1)
	put!(sp, readchomp("rotation"))
end
it = Channel{Int}(1)
nr = Channel{Int}(1)
put!(nr,1)
sp = Channel{AbstractString}(1)
yy = Input(wl)
X = DataFrame(wl = wl, l315c0 = zeros(nwl), l0c0 = zeros(nwl), l45c0 = zeros(nwl), l90c0 = zeros(nwl), l0c315 = zeros(nwl), l0c45 = zeros(nwl), dark = zeros(nwl), I = zeros(nwl), dolp = zeros(nwl), aop = zeros(nwl), docp = zeros(nwl))
writetable("data.csv", X)

@async map(readall, [`pdflatex plot.tex`, `evince plot.pdf`])

function plotit(y)
	p = scatterplot(wl,y,xlim=[300,800],ylim=[0,maxy])
	display(p)
	sleep(0.07)
end
lift(plotit,yy)
function normalize(I, dolp, aop, docp)
	I = I < 0 ? 0.0 : I > 10maxy ? 10maxy : I
	dolp = dolp < 0 ? 0.0 : dolp > 2 ? 2.0 : dolp
	docp = docp < -2 ? -2.0 : docp > 2 ? 2.0 : docp
	return (I, dolp, aop, docp)
end
function draw()
	if any(map(Bool,X[:l0c315])) && any(map(Bool,X[:l0c45]))
		s0, s1, s2, s3 = measurements2stokes(X[:l315c0] - X[:dark], X[:l0c0] - X[:dark], X[:l45c0] - X[:dark], X[:l90c0] - X[:dark], X[:l0c315] - X[:dark], X[:l0c45] - X[:dark])
	else
		s0, s1, s2, s3 = measurements2stokes(X[:l315c0] - X[:dark], X[:l0c0] - X[:dark], X[:l45c0] - X[:dark], X[:l90c0] - X[:dark])
	end
	for i = 1:nwl
		I, dolp, aop, docp = normalize(stokes2polar(s0[i], s1[i], s2[i], s3[i])...)
		X[i,:I] = I
		X[i,:dolp] = dolp
		X[i,:aop] = rad2deg(aop) + OFFSET
		X[i,:docp] = docp
	end
	writetable("data.csv", X)
	readall(`pdflatex plot.tex`)
end
function iter()
	wait(nr)
	nreps = fetch(nr)
	if isready(it)
		s[:integration_time_micros](take!(it))
	end
	y = zeros(nwl)
	for i = 1:nreps
		y .+= s[:intensities]()
	end
	y /= nreps
	push!(yy,y)
	if isready(sp)
		X[symbol(take!(sp))] = y
		draw()
	end
end
sleep(0)
s[:intensities]()
sleep(0)

@schedule begin
	while true
		watchit()
	end
end
@schedule begin
	while true
		watchr()
	end
end
@schedule begin
	while true
		watchnr()
	end
end
@schedule begin
	while true
		iter()
	end
end

sleep(60*60*24)


function GlobalMappingFunctionModel(arg0::jdouble, arg1::jdouble)
    return GlobalMappingFunctionModel((jdouble, jdouble), arg0, arg1)
end

function GlobalMappingFunctionModel(arg0::jdouble, arg1::jdouble, arg2::TimeScale)
    return GlobalMappingFunctionModel((jdouble, jdouble, TimeScale), arg0, arg1, arg2)
end

function get_parameters_drivers(obj::GlobalMappingFunctionModel)
    return jcall(obj, "getParametersDrivers", List, ())
end

function mapping_factors(obj::GlobalMappingFunctionModel, arg0::RealFieldElement, arg1::RealFieldElement, arg2::Vector{RealFieldElement}, arg3::FieldAbsoluteDate)
    return jcall(obj, "mappingFactors", Vector{RealFieldElement}, (RealFieldElement, RealFieldElement, Vector{RealFieldElement}, FieldAbsoluteDate), arg0, arg1, arg2, arg3)
end

function mapping_factors(obj::GlobalMappingFunctionModel, arg0::jdouble, arg1::jdouble, arg2::Vector{jdouble}, arg3::AbsoluteDate)
    return jcall(obj, "mappingFactors", Vector{jdouble}, (jdouble, jdouble, Vector{jdouble}, AbsoluteDate), arg0, arg1, arg2, arg3)
end


isdefined(Base, :__precompile__) && __precompile__()

"""
Module for constructing self-contained power system objects.
"""
module PowerSystems

#################################################################################
# Exports

export System
export Bus
export Arc
export AggregationTopology
export Area
export LoadZone
export get_aggregation_topology_accessor

export PowerSystemType
export Component
export Device
export Branch
export StaticInjection
export ACBranch
export Line
export MonitoredLine
export DCBranch
export HVDCLine
export VSCDCLine
export Transformer2W
export TapTransformer
export PhaseShiftingTransformer

export ThreePartCost
export TwoPartCost

export Generator
export HydroGen
export HydroDispatch
export HydroEnergyReservoir
export TechHydro

export RenewableGen
export TechRenewable
export RenewableFix
export RenewableDispatch

export ThermalGen
export TechThermal
export ThermalStandard

export ElectricLoad
export StaticLoad
export PowerLoad
export PowerLoadPF
export FixedAdmittance
export ControllableLoad
export InterruptibleLoad

export Storage
export GenericBattery

export DynamicGenerator
#AVR Exports
export AVR
export AVRFixed
export AVRSimple
export AVRTypeI
export AVRTypeII

#Machine Exports
export Machine
export BaseMachine
export OneDOneQMachine
export MarconatoMachine
export SimpleMarconatoMachine
export AndersonFouadMachine
export SimpleAFMachine
export FullMachine
export SimpleFullMachine

#PSS Exports
export PSS
export PSSFixed
export PSSFixed

#Shaft Exports
export SingleMass
export FiveMassShaft

#TG Exports
export TurbineGov
export TGFixed
export TGTypeI
export TGTypeII

export DynamicInverter
# Converter Exports
export Converter
export AvgCnvFixedDC

# DC Source Exports
export DCSource
export FixedDCSource

# Filter Exports
export Filter
export LCLFilter

# FrequencyEstimator Exports
export FrequencyEstimator
export PLL

# Outer Control Exports
export OuterControl
export VirtualInertiaQdroop
export VirtualInertia
export ReactivePowerDroop

# VSControl Export
export VSControl
export CombinedVIwithVZ

export Source

export Service
export Reserve
export ReserveDirection
export ReserveUp
export ReserveDown
export StaticReserve
export VariableReserve
export Transfer

export PTDF
export Ybus
export LODF
export GeneratorCostModels
export AngleUnits
export BusTypes
export LoadModels
export PrimeMovers
export ThermalFuels

export Forecast
export Deterministic
export Probabilistic
export ScenarioBased

export get_dynamic_components

export make_pf
export solve_powerflow!

export parse_file
export add_forecasts!
export add_forecast!
export remove_forecast!
export clear_forecasts!
export add_component!
export remove_component!
export remove_components!
export clear_components!
export add_service!
export remove_service!
export has_service
export get_buses
export get_components_in_aggregation_topology
export get_aggregation_topology_mapping
export get_contributing_devices
export get_contributing_device_mapping
export ServiceContributingDevices
export ServiceContributingDevicesKey
export ServiceContributingDevicesMapping
export are_forecasts_contiguous
export generate_initial_times
export get_component
export get_components
export get_components_by_name
export get_forecast_labels
export get_forecast_initial_times
export get_forecast_keys
export get_forecast
export get_forecast_values
export get_forecasts_horizon
export get_forecasts_initial_time
export get_forecasts_interval
export get_forecasts_resolution
export get_horizon
export get_initial_time
export get_resolution
export get_data
export iterate_components
export iterate_forecasts
export make_forecasts
export get_bus_numbers
export get_name
export get_basepower
export get_frequency
export to_json
export check_forecast_consistency
export validate_forecast_consistency
export clear_ext
export convert_component!

#################################################################################
# Imports

import SparseArrays
import LinearAlgebra: LAPACK.getri!
import LinearAlgebra: LAPACK.getrf!
import LinearAlgebra: BLAS.gemm
import LinearAlgebra
import Dates
import TimeSeries
import DataFrames
import JSON
import JSON2
import CSV
import YAML
import UUIDs
import Base.to_index

import InfrastructureSystems
import InfrastructureSystems:
    Components,
    Deterministic,
    Probabilistic,
    Forecast,
    ScenarioBased,
    InfrastructureSystemsType,
    InfrastructureSystemsInternal,
    FlattenIteratorWrapper,
    LazyDictFromIterator,
    DataFormatError,
    InvalidRange,
    InvalidValue,
    get_data,
    get_horizon,
    get_initial_time,
    get_resolution,
    get_name

const IS = InfrastructureSystems

#################################################################################

using DocStringExtensions

@template (FUNCTIONS, METHODS) = """
                                 $(TYPEDSIGNATURES)
                                 $(DOCSTRING)
                                 """

#################################################################################
# Includes

"""
Supertype for all PowerSystems types.
All subtypes must include a InfrastructureSystemsInternal member.
Subtypes should call InfrastructureSystemsInternal() by default, but also must
provide a constructor that allows existing values to be deserialized.
"""
abstract type PowerSystemType <: IS.InfrastructureSystemsType end

abstract type Component <: PowerSystemType end
# supertype for "devices" (bus, line, etc.)
abstract type Device <: Component end
abstract type DeviceParameter <: PowerSystemType end

include("common.jl")
include("models/static_models.jl")
include("models/dynamic_models.jl")

# Include utilities
include("utils/IO/base_checks.jl")

# PowerSystems models
include("models/topological_elements.jl")
include("models/branches.jl")
include("models/operational_cost.jl")
#include("models/network.jl")

# Static types
include("models/services.jl")
include("models/reserves.jl")
include("models/generation.jl")
include("models/storage.jl")
include("models/loads.jl")
include("models/devices.jl")
include("models/dynamic_generator_components.jl")
include("models/dynamic_inverter_components.jl")

# Include all auto-generated structs.
include("models/generated/includes.jl")
include("models/supplemental_constructors.jl")
include("models/supplemental_accessors.jl")

# Dynamic Composed types
include("models/dynamic_generator.jl")
include("models/dynamic_inverter.jl")

# Definitions of PowerSystem
include("base.jl")

#Data Checks
include("utils/IO/system_checks.jl")
include("utils/IO/branchdata_checks.jl")

# network calculations
include("utils/network_calculations/common.jl")
include("utils/network_calculations/ybus_calculations.jl")
include("utils/network_calculations/ptdf_calculations.jl")
include("utils/network_calculations/lodf_calculations.jl")

#PowerFlow
include("utils/power_flow/make_pf.jl")
include("utils/power_flow/power_flow.jl")

# Include Parsing files
include("parsers/common.jl")
include("parsers/enums.jl")
include("parsers/pm_io.jl")
include("parsers/im_io.jl")
include("parsers/power_system_table_data.jl")
include("parsers/power_models_data.jl")

# Better printing
include("utils/print.jl")

include("models/serialization.jl")

# Download test data
include("utils/data.jl")
import .UtilsData: TestData

end # module


function define_weights_voltages(strategy)
    # Create dictionary of spatial weights for different equipment types / areas
    spatial_weights = Dict{String, Any}()

    # Weights for AC OHL
    spatial_weights["ac_ohl"] =  Dict{String, Any}()
    spatial_weights["ac_ohl"]["nature"] = 2.5
    spatial_weights["ac_ohl"]["urban"] = 40
    spatial_weights["ac_ohl"]["mountain"] = 10
    spatial_weights["ac_ohl"]["sea"] = 40
    spatial_weights["ac_ohl"]["Natura2000"] = 40
    spatial_weights["ac_ohl"]["grid"] = 1
    spatial_weights["ac_ohl"]["agricultural"] = 1
    spatial_weights["ac_ohl"]["roads"] = 1
    spatial_weights["ac_ohl"]["railroads"] = 1

    if strategy == "OHL_on_existing_corridors"
        spatial_weights["ac_ohl"]["agricultural"] = 40
        spatial_weights["ac_ohl"]["nature"] = 40
        spatial_weights["ac_ohl"]["mountain"] = 40
    end

    if strategy == "cables_only"
        spatial_weights["ac_ohl"]["agricultural"] = 40
        spatial_weights["ac_ohl"]["nature"] = 40
        spatial_weights["ac_ohl"]["mountain"] = 40
        spatial_weights["ac_ohl"]["grid"] = 40
        spatial_weights["ac_ohl"]["roads"] = 40
        spatial_weights["ac_ohl"]["railroads"] = 40
    end

    # Weights for AC UGC
    spatial_weights["ac_ugc"] =  Dict{String, Any}()
    spatial_weights["ac_ugc"]["nature"] = 1.5
    spatial_weights["ac_ugc"]["urban"] = 2.5
    spatial_weights["ac_ugc"]["mountain"] = 4.5
    spatial_weights["ac_ugc"]["sea"] = 0.75
    spatial_weights["ac_ugc"]["Natura2000"] = 40
    spatial_weights["ac_ugc"]["grid"] = 1
    spatial_weights["ac_ugc"]["agricultural"] = 1
    spatial_weights["ac_ugc"]["roads"] = 1
    spatial_weights["ac_ugc"]["railroads"] = 1

    # Weights for DC OHL
    spatial_weights["dc_ohl"] =  Dict{String, Any}()
    spatial_weights["dc_ohl"]["nature"] = 2.5
    spatial_weights["dc_ohl"]["urban"] = 40
    spatial_weights["dc_ohl"]["mountain"] = 10
    spatial_weights["dc_ohl"]["sea"] = 40
    spatial_weights["dc_ohl"]["Natura2000"] = 10
    spatial_weights["dc_ohl"]["grid"] = 1
    spatial_weights["dc_ohl"]["agricultural"] = 1
    spatial_weights["dc_ohl"]["roads"] = 1
    spatial_weights["dc_ohl"]["railroads"] = 1

    if strategy == "OHL_on_existing_corridors"
        spatial_weights["dc_ohl"]["agricultural"] = 40
        spatial_weights["dc_ohl"]["nature"] = 40
        spatial_weights["dc_ohl"]["mountain"] = 40
    end

    if strategy == "cables_only"
        spatial_weights["dc_ohl"]["agricultural"] = 40
        spatial_weights["dc_ohl"]["nature"] = 40
        spatial_weights["dc_ohl"]["mountain"] = 40
        spatial_weights["dc_ohl"]["grid"] = 40
        spatial_weights["dc_ohl"]["roads"] = 40
        spatial_weights["dc_ohl"]["railroads"] = 40
    end

    # Weights for AC UGC
    spatial_weights["dc_ugc"] =  Dict{String, Any}()
    spatial_weights["dc_ugc"]["nature"] = 1.5
    spatial_weights["dc_ugc"]["urban"] = 2.5
    spatial_weights["dc_ugc"]["mountain"] = 4.5
    spatial_weights["dc_ugc"]["sea"] = 0.75
    spatial_weights["dc_ugc"]["Natura2000"] = 10
    spatial_weights["dc_ugc"]["grid"] = 1
    spatial_weights["dc_ugc"]["agricultural"] = 1
    spatial_weights["dc_ugc"]["roads"] = 1
    spatial_weights["dc_ugc"]["railroads"] = 1

    # Define transmission voltage levels
    voltages = Dict{String, Any}()
    voltages["ac_ohl"] = 400
    voltages["dc_ohl"] = 400
    voltages["ac_ugc"] = 400
    voltages["dc_ugc"] = 320

    #Minimum resolution
    resolution = Dict{String, Any}()
    resolution["delta_min"] = 4
    resolution["weight_min"] = Dict{String, Any}()
    resolution["weight_min"]["ac_ohl"] = 1
    resolution["weight_min"]["ac_ugc"] = 1
    resolution["weight_min"]["dc_ohl"] = 1
    resolution["weight_min"]["dc_ugc"] = 1

    impedances = Dict{String,Any}()
    impedances["ac"] = Dict{String,Any}()
    impedances["ac"]["ohl"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["16"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["16"]["r"] = 1.874
    impedances["ac"]["ohl"]["16"]["x"] = 0.347
    impedances["ac"]["ohl"]["16"]["c"] = 13.06

    impedances["ac"]["ohl"]["25"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["25"]["r"] = 1.118
    impedances["ac"]["ohl"]["25"]["x"] = 0.334
    impedances["ac"]["ohl"]["25"]["c"] = 13.06

    impedances["ac"]["ohl"]["50"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["50"]["r"] = 0.574
    impedances["ac"]["ohl"]["50"]["x"] = 0.312
    impedances["ac"]["ohl"]["50"]["c"] = 13.06

    impedances["ac"]["ohl"]["70"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["70"]["r"] = 0.404
    impedances["ac"]["ohl"]["70"]["x"] = 0.302
    impedances["ac"]["ohl"]["70"]["c"] = 13.06

    impedances["ac"]["ohl"]["95"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["95"]["r"] = 0.300
    impedances["ac"]["ohl"]["95"]["x"] = 0.291
    impedances["ac"]["ohl"]["95"]["c"] = 13.06

    impedances["ac"]["ohl"]["120"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["120"]["r"] = 0.252
    impedances["ac"]["ohl"]["120"]["x"] = 0.284
    impedances["ac"]["ohl"]["120"]["c"] = 13.06

    impedances["ac"]["ohl"]["150"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["150"]["r"] = 0.190
    impedances["ac"]["ohl"]["150"]["x"] = 0.279
    impedances["ac"]["ohl"]["150"]["c"] = 13.06

    impedances["ac"]["ohl"]["170"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["170"]["r"] = 0.165
    impedances["ac"]["ohl"]["170"]["x"] = 0.276
    impedances["ac"]["ohl"]["170"]["c"] = 13.06

    impedances["ac"]["ohl"]["185"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["185"]["r"] = 0.154
    impedances["ac"]["ohl"]["185"]["x"] = 0.274
    impedances["ac"]["ohl"]["185"]["c"] = 13.06

    impedances["ac"]["ohl"]["210"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["210"]["r"] = 0.133
    impedances["ac"]["ohl"]["210"]["x"] = 0.270
    impedances["ac"]["ohl"]["210"]["c"] = 13.06

    impedances["ac"]["ohl"]["230"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["230"]["r"] = 0.122
    impedances["ac"]["ohl"]["230"]["x"] = 0.268
    impedances["ac"]["ohl"]["230"]["c"] = 13.06

    impedances["ac"]["ohl"]["240"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["240"]["r"] = 0.116
    impedances["ac"]["ohl"]["240"]["x"] = 0.267
    impedances["ac"]["ohl"]["240"]["c"] = 13.06

    impedances["ac"]["ohl"]["265"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["265"]["r"] = 0.107
    impedances["ac"]["ohl"]["265"]["x"] = 0.264
    impedances["ac"]["ohl"]["265"]["c"] = 13.06

    impedances["ac"]["ohl"]["300"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["300"]["r"] = 0.093
    impedances["ac"]["ohl"]["300"]["x"] = 0.261
    impedances["ac"]["ohl"]["300"]["c"] = 13.06

    impedances["ac"]["ohl"]["340"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["340"]["r"] = 0.083
    impedances["ac"]["ohl"]["340"]["x"] = 0.258
    impedances["ac"]["ohl"]["340"]["c"] = 13.06

    impedances["ac"]["ohl"]["380"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["380"]["r"] = 0.074
    impedances["ac"]["ohl"]["380"]["x"] = 0.255
    impedances["ac"]["ohl"]["380"]["c"] = 13.06

    impedances["ac"]["ohl"]["435"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["435"]["r"] = 0.065
    impedances["ac"]["ohl"]["435"]["x"] = 0.252
    impedances["ac"]["ohl"]["435"]["c"] = 13.06

    impedances["ac"]["ohl"]["450"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["450"]["r"] = 0.063
    impedances["ac"]["ohl"]["450"]["x"] = 0.251
    impedances["ac"]["ohl"]["450"]["c"] = 13.06

    impedances["ac"]["ohl"]["490"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["490"]["r"] = 0.058
    impedances["ac"]["ohl"]["490"]["x"] = 0.249
    impedances["ac"]["ohl"]["490"]["c"] = 13.06

    impedances["ac"]["ohl"]["550"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["550"]["r"] = 0.051
    impedances["ac"]["ohl"]["550"]["x"] = 0.246
    impedances["ac"]["ohl"]["550"]["c"] = 13.06

    impedances["ac"]["ohl"]["560"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["560"]["r"] = 0.050
    impedances["ac"]["ohl"]["560"]["x"] = 0.26
    impedances["ac"]["ohl"]["560"]["c"] = 13.06

    impedances["ac"]["ohl"]["680"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["680"]["r"] = 0.042
    impedances["ac"]["ohl"]["680"]["x"] = 0.241
    impedances["ac"]["ohl"]["680"]["c"] = 13.06

    impedances["ac"]["ugc"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["300"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["300"]["r"] = 0.069
    impedances["ac"]["ugc"]["300"]["x"] = 0.110
    impedances["ac"]["ugc"]["300"]["c"] = 375

    impedances["ac"]["ugc"]["400"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["400"]["r"] = 0.057
    impedances["ac"]["ugc"]["400"]["x"] = 0.130
    impedances["ac"]["ugc"]["400"]["c"] = 170

    impedances["ac"]["ugc"]["500"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["500"]["r"] = 0.047
    impedances["ac"]["ugc"]["500"]["x"] = 0.125
    impedances["ac"]["ugc"]["500"]["c"] = 180

    impedances["ac"]["ugc"]["630"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["630"]["r"] = 0.038
    impedances["ac"]["ugc"]["630"]["x"] = 0.122
    impedances["ac"]["ugc"]["630"]["c"] = 200

    impedances["ac"]["ugc"]["800"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["800"]["r"] = 0.027
    impedances["ac"]["ugc"]["800"]["x"] = 0.102
    impedances["ac"]["ugc"]["800"]["c"] = 402

    impedances["ac"]["ugc"]["1200"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["1200"]["r"] = 0.0151
    impedances["ac"]["ugc"]["1200"]["x"] = 0.116
    impedances["ac"]["ugc"]["1200"]["c"] = 210

    impedances["ac"]["ugc"]["1600"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["1600"]["r"] = 0.0113
    impedances["ac"]["ugc"]["1600"]["x"] = 0.111
    impedances["ac"]["ugc"]["1600"]["c"] = 231

    impedances["ac"]["ugc"]["2000"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["2000"]["r"] = 0.010
    impedances["ac"]["ugc"]["2000"]["x"] = 0.1099
    impedances["ac"]["ugc"]["2000"]["c"] = 240

    impedances["ac"]["ugc"]["2500"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["2500"]["r"] = 0.0095
    impedances["ac"]["ugc"]["2500"]["x"] = 0.098
    impedances["ac"]["ugc"]["2500"]["c"] = 260

    impedances["dc"] = Dict{String,Any}()
    impedances["dc"]["ohl"] = Dict{String,Any}()
    for (cross_section, cs) in impedances["ac"]["ohl"]
        impedances["dc"]["ohl"][cross_section] = Dict{String,Any}()
        impedances["dc"]["ohl"][cross_section]["r"] = cs["r"]
        impedances["dc"]["ohl"][cross_section]["x"] = 0
        impedances["dc"]["ohl"][cross_section]["c"] = 0
    end

    impedances["dc"]["ugc"] = Dict{String,Any}()
    impedances["dc"]["ugc"]["95"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["95"]["r"] = 0.247
    impedances["dc"]["ugc"]["95"]["x"] = 0
    impedances["dc"]["ugc"]["95"]["c"] = 0

    impedances["dc"]["ugc"]["120"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["120"]["r"] = 0.196
    impedances["dc"]["ugc"]["120"]["x"] = 0
    impedances["dc"]["ugc"]["120"]["c"] = 0

    impedances["dc"]["ugc"]["150"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["150"]["r"] = 0.16
    impedances["dc"]["ugc"]["150"]["x"] = 0
    impedances["dc"]["ugc"]["150"]["c"] = 0

    impedances["dc"]["ugc"]["185"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["185"]["r"] = 0.128
    impedances["dc"]["ugc"]["185"]["x"] = 0
    impedances["dc"]["ugc"]["185"]["c"] = 0

    impedances["dc"]["ugc"]["240"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["240"]["r"] = 0.097
    impedances["dc"]["ugc"]["240"]["x"] = 0
    impedances["dc"]["ugc"]["240"]["c"] = 0

    impedances["dc"]["ugc"]["300"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["300"]["r"] = 0.08
    impedances["dc"]["ugc"]["300"]["x"] = 0
    impedances["dc"]["ugc"]["300"]["c"] = 0

    impedances["dc"]["ugc"]["400"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["400"]["r"] = 0.062
    impedances["dc"]["ugc"]["400"]["x"] = 0
    impedances["dc"]["ugc"]["400"]["c"] = 0

    impedances["dc"]["ugc"]["500"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["500"]["r"] = 0.0508
    impedances["dc"]["ugc"]["500"]["x"] = 0
    impedances["dc"]["ugc"]["500"]["c"] = 0

    impedances["dc"]["ugc"]["630"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["630"]["r"] = 0.039
    impedances["dc"]["ugc"]["630"]["x"] = 0
    impedances["dc"]["ugc"]["630"]["c"] = 0

    impedances["dc"]["ugc"]["800"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["800"]["r"] = 0.032
    impedances["dc"]["ugc"]["800"]["x"] = 0
    impedances["dc"]["ugc"]["800"]["c"] = 0

    impedances["dc"]["ugc"]["1000"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1000"]["r"] = 0.024
    impedances["dc"]["ugc"]["1000"]["x"] = 0
    impedances["dc"]["ugc"]["1000"]["c"] = 0

    impedances["dc"]["ugc"]["1200"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1200"]["r"] = 0.021
    impedances["dc"]["ugc"]["1200"]["x"] = 0
    impedances["dc"]["ugc"]["1200"]["c"] = 0

    impedances["dc"]["ugc"]["1400"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1400"]["r"] = 0.019
    impedances["dc"]["ugc"]["1400"]["x"] = 0
    impedances["dc"]["ugc"]["1400"]["c"] = 0

    impedances["dc"]["ugc"]["1600"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1600"]["r"] = 0.017
    impedances["dc"]["ugc"]["1600"]["x"] = 0
    impedances["dc"]["ugc"]["1600"]["c"] = 0

    impedances["dc"]["ugc"]["1800"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1800"]["r"] = 0.0155
    impedances["dc"]["ugc"]["1800"]["x"] = 0
    impedances["dc"]["ugc"]["1800"]["c"] = 0

    impedances["dc"]["ugc"]["2000"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2000"]["r"] = 0.014
    impedances["dc"]["ugc"]["2000"]["x"] = 0
    impedances["dc"]["ugc"]["2000"]["c"] = 0

    impedances["dc"]["ugc"]["2200"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2200"]["r"] = 0.013
    impedances["dc"]["ugc"]["2200"]["x"] = 0
    impedances["dc"]["ugc"]["2200"]["c"] = 0

    impedances["dc"]["ugc"]["2400"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2400"]["r"] = 0.012
    impedances["dc"]["ugc"]["2400"]["x"] = 0
    impedances["dc"]["ugc"]["2400"]["c"] = 0

    impedances["dc"]["ugc"]["2600"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2600"]["r"] = 0.010
    impedances["dc"]["ugc"]["2600"]["x"] = 0
    impedances["dc"]["ugc"]["2600"]["c"] = 0

    impedances["dc"]["ugc"]["2800"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2800"]["r"] = 0.09
    impedances["dc"]["ugc"]["2800"]["x"] = 0
    impedances["dc"]["ugc"]["2800"]["c"] = 0

    impedances["dc"]["ugc"]["3000"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["3000"]["r"] = 0.075
    impedances["dc"]["ugc"]["3000"]["x"] = 0
    impedances["dc"]["ugc"]["3000"]["c"] = 0

    return spatial_weights, voltages, resolution, impedances
end


const CLVSolver = StagedSolver{<: CLVProblem, <: CLVSolution}

"""
    CLVSolver(prob::CLVProblem; <keyword arguments>)
    CLVSolver(prob::LEProblem; <keyword arguments>)

The preferred and equivalent method to get a solver for a `CLVProblem`
is `init(prob::CLVProblem)`.  Note that `CLVSolver(prob::LEProblem)`
is equivalent to `init(CLVProblem(prob))`.

# Arguments
- `record::Vector{Symbol}`: Variables to be saved.  A subset of
  `[:G, :C, :D, :x]`.
"""
function CLVSolver(prob::CLVProblem;
                   record::Vector{Symbol} = Symbol[],
                   forward_relaxer::Type = ForwardRelaxer,
                   forward_dynamics::Type = ForwardDynamics,
                   backward_relaxer::Type = BackwardRelaxer,
                   backward_dynamics::Type = if :D in record
                       BackwardDynamicsWithD
                   else
                       BackwardDynamics
                   end,
                   )
    stage_types = [
        # PhaseRelaxer???
        forward_relaxer,
        forward_dynamics,
        backward_relaxer,
        backward_dynamics,
    ]
    return CLVSolver(prob, stage_types, record)
end

function CLVSolver(prob::CLVProblem, stage_types::AbstractVector,
                   record::Vector{Symbol})
    sol = CLVSolution(prob, record)
    args = (prob, sol)  # additional arguments to each `stage_types`
    # FIXME: Currently sol is used as "cache" (args[2]) as well.
    return StagedSolver(prob, sol, stage_types, args)
end

CLVSolver(prob::LEProblem; kwargs...) = CLVSolver(CLVProblem(prob); kwargs...)

init(prob::CLVProblem; kwargs...) = CLVSolver(prob; kwargs...)

"""
    forward_dynamics!(solver::CLVSolver) :: ForwardDynamics

Solve the CLV problem up to the forward dynamics stage and return an
iterator to step through the forward dynamics.
See also: [`backward_dynamics!`](@ref), [`goto!`](@ref).
"""
forward_dynamics!(solver::CLVSolver; kwargs...) =
    goto!(solver, ForwardDynamics; kwargs...)

"""
    indexed_forward_dynamics!(solver::CLVSolver)
    indexed_forward_dynamics!(stage::ForwardDynamics)

Just a short-hand for `enumerate(forward_dynamics!(solver))`.
It's for symmetry with [`indexed_backward_dynamics!`](@ref).
"""
indexed_forward_dynamics!(solver::CLVSolver; kwargs...) =
    indexed_forward_dynamics!(forward_dynamics!(solver; kwargs...))
indexed_forward_dynamics!(fitr::ForwardDynamics) = enumerate(fitr)

"""
    backward_dynamics!(solver::CLVSolver) :: BackwardDynamics

Solve the CLV problem up to the (final) backward dynamics stage and
return an iterator to step through the backward dynamics.
See also [`forward_dynamics!`](@ref), [`goto!`](@ref).

Note that finishing iteration of the returned iterator does not
finalize all the solver stages (namely, recording to `solver.sol`, if
non-default `backward_dynamics` is used).  In this case,
`solve!(solver)` has to be called after the iteration.

### Example
```julia
angles = [acos(abs(dot(C[:, 1], C[:, 2]))) * 2 / π for C
          in backward_dynamics!(solver)]
```
"""
backward_dynamics!(solver::CLVSolver; kwargs...) =
    goto!(solver, BackwardDynamics; kwargs...)

"""
    indexed_backward_dynamics!(solver::CLVSolver)
    indexed_backward_dynamics!(stage::BackwardDynamics)

It is equivalent to `zip(some_counter, backward_dynamics!(solver))`
where `some_counter` is an iterator over integers such that the same
indices returned by [`indexed_forward_dynamics!`](@ref) indicate that
those events are at the same time point of the backward and forward
passes.  This is useful when combining matrices [`CLV.G`](@ref) and
[`CLV.C`](@ref) to obtain the CLV in the (original) tangent space.

For such example, see:
[Covariant Lyapunov vectors on the Lorenz attractor](@ref)
in the online manual.

Note that `indexed_backward_dynamics!(solver::CLVSolver)` is
equivalent to
```julia
backward = backward_dynamics!(solver)
indexed_backward_dynamics!(backward)
```
Separately calling [`backward_dynamics!`](@ref) is useful when the
quantities other than [`CLV.G`](@ref) (e.g., [`CLV.R`](@ref)) are
required.

See also [`indexed_forward_dynamics!`](@ref).
"""
indexed_backward_dynamics!(solver::CLVSolver;
                           until::Int = 1,
                           kwargs...) =
    indexed_backward_dynamics!(backward_dynamics!(solver; kwargs...);
                               until = until)

function indexed_backward_dynamics!(bitr::BackwardDynamics;
                                    until::Int = 1)
    return zip(length(bitr)-1:-1:until, bitr)
end


using YaoArrayRegister
using YaoArrayRegister: matvec

export ControlBlock, control, cnot, cz

struct ControlBlock{N,BT<:AbstractBlock,C,M} <: AbstractContainer{BT,N}
    ctrl_locs::NTuple{C,Int}
    ctrl_config::NTuple{C,Int}
    content::BT
    locs::NTuple{M,Int}
    function ControlBlock{N,BT,C,M}(
        ctrl_locs,
        ctrl_config,
        block,
        locs,
    ) where {N,C,M,BT<:AbstractBlock}
        @assert_locs_safe N (ctrl_locs..., locs...)
        @assert nqubits(block) == M "number of locations doesn't match the size of block"
        @assert block isa AbstractBlock "expect a block, got $(typeof(block))"
        new{N,BT,C,M}(ctrl_locs, ctrl_config, block, locs)
    end
end

"""
    decode_sign(ctrls...)

Decode signs into control sequence on control or inversed control.
"""
decode_sign(ctrls::Int...) = decode_sign(ctrls)
decode_sign(ctrls::NTuple{N,Int}) where {N} =
    tuple(ctrls .|> abs, ctrls .|> sign .|> (x -> (1 + x) ÷ 2))

function ControlBlock{N}(
    ctrl_locs::NTuple{C},
    ctrl_config::NTuple{C},
    block::BT,
    locs::NTuple{K},
) where {N,M,C,K,BT<:AbstractBlock{M}}
    M == K || throw(DimensionMismatch("block position not maching its size!"))
    return ControlBlock{N,BT,C,M}(ctrl_locs, ctrl_config, block, locs)
end

function ControlBlock{N}(
    ctrl_locs::NTuple{C},
    ctrl_config::NTuple{C},
    block,
    locs::NTuple{K},
) where {N,M,C,K}
    error("expect a block, got $(typeof(block))")
end

# control bit configs are 1 by default, it use sign to encode control bit code
ControlBlock{N}(ctrl_locs::NTuple{C}, block::AbstractBlock, locs::NTuple) where {N,C} =
    ControlBlock{N}(decode_sign(ctrl_locs)..., block, locs)
ControlBlock{N}(ctrl_locs::NTuple{C}, block::Function, locs::NTuple) where {N,C} =
    ControlBlock{N}(decode_sign(ctrl_locs)..., parse_block(length(locs), block), locs)
ControlBlock{N}(ctrl_locs::NTuple{C}, block, locs::NTuple) where {N,C} =
    ControlBlock{N}(decode_sign(ctrl_locs)..., block, locs) # trigger error

# use pair to represent block under control in a compact way
ControlBlock{N}(ctrl_locs::NTuple{C}, target::Pair) where {N,C} =
    ControlBlock{N}(ctrl_locs, target.second, (target.first...,))

"""
    control(n, ctrl_locs, target)

Return a [`ControlBlock`](@ref) with number of active qubits `n` and control locs
`ctrl_locs`, and control target in `Pair`.

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control(4, (1, 2), 3=>X)
nqubits: 4
control(1, 2)
└─ (3,) X

julia> control(4, 1, 3=>X)
nqubits: 4
control(1)
└─ (3,) X
```
"""
control(total::Int, ctrl_locs, target::Pair) = ControlBlock{total}(Tuple(ctrl_locs), target)
control(total::Int, control_location::Int, target::Pair) = control(total, (control_location,), target)

"""
    control(ctrl_locs, target) -> f(n)

Return a lambda that takes the number of total active qubits as input. See also
[`control`](@ref).

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control((2, 3), 1=>X)
(n -> control(n, (2, 3), 1 => X))

julia> control(2, 1=>X)
(n -> control(n, 2, 1 => X))
```
"""
control(ctrl_locs, target::Pair) = @λ(n -> control(n, ctrl_locs, target))
control(control_location::Int, target::Pair) = @λ(n -> control(n, control_location, target))

"""
    control(target) -> f(ctrl_locs)

Return a lambda that takes a `Tuple` of control qubits locs as input. See also
[`control`](@ref).

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control(1=>X)
(ctrl_locs -> control(ctrl_locs, 1 => X))

julia> control((2, 3) => YaoBlocks.ConstGate.CNOT)
(ctrl_locs -> control(ctrl_locs, (2, 3) => CNOT))
```
"""
control(target::Pair) = @λ(ctrl_locs -> control(ctrl_locs, target))

"""
    control(ctrl_locs::Int...) -> f(target)

Return a lambda that takes a `Pair` of control target as input.
See also [`control`](@ref).

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control(1, 2)
(target -> control((1, 2), target))
```
"""
control(ctrl_locs::Int...) = @λ(target -> control(ctrl_locs, target))

"""
    cnot([n, ]ctrl_locs, location)

Return a speical [`ControlBlock`](@ref), aka CNOT gate with number of active qubits
`n` and locs of control qubits `ctrl_locs`, and `location` of `X` gate.

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> cnot(3, (2, 3), 1)
nqubits: 3
control(2, 3)
└─ (1,) X

julia> cnot(2, 1)
(n -> cnot(n, 2, 1))
```
"""
cnot(total::Int, ctrl_locs, locs::Int) = control(total, ctrl_locs, locs => X)
cnot(ctrl_locs, loc::Int) = @λ(n -> cnot(n, ctrl_locs, loc))

"""
    cz([n, ]ctrl_locs, location)

Return a speical [`ControlBlock`](@ref), aka CZ gate with number of active qubits
`n` and locs of control qubits `ctrl_locs`, and `location` of `Z` gate. See also
[`cnot`](@ref).
"""
cz(total::Int, ctrl_locs, locs::Int) = control(total, ctrl_locs, locs => Z)
cz(ctrl_locs, loc::Int) = @λ(n -> cz(n, ctrl_locs, loc))

mat(::Type{T}, c::ControlBlock{N,BT,C}) where {T,N,BT,C} =
    cunmat(N, c.ctrl_locs, c.ctrl_config, mat(T, c.content), c.locs)

function _apply!(r::AbstractRegister, c::ControlBlock)
    instruct!(r, mat_matchreg(r, c.content), c.locs, c.ctrl_locs, c.ctrl_config)
    return r
end

# specialization
for G in [:X, :Y, :Z, :S, :T, :Sdag, :Tdag]
    GT = Expr(:(.), :ConstGate, QuoteNode(Symbol(G, :Gate)))

    @eval function _apply!(r::AbstractRegister, c::ControlBlock{N,<:$GT}) where {N}
        instruct!(r, Val($(QuoteNode(G))), c.locs, c.ctrl_locs, c.ctrl_config)
        return r
    end
end

PropertyTrait(::ControlBlock) = PreserveAll()

occupied_locs(c::ControlBlock) = (c.ctrl_locs..., map(x -> c.locs[x], occupied_locs(c.content))...)
chsubblocks(pb::ControlBlock{N}, blk::AbstractBlock) where {N} =
    ControlBlock{N}(pb.ctrl_locs, pb.ctrl_config, blk, pb.locs)

# NOTE: ControlBlock will forward parameters directly without loop
cache_key(ctrl::ControlBlock) = cache_key(ctrl.content)

function Base.:(==)(lhs::ControlBlock{N,BT,C,M}, rhs::ControlBlock{N,BT,C,M}) where {BT,N,C,M}
    return (lhs.ctrl_locs == rhs.ctrl_locs) && (lhs.content == rhs.content) && (lhs.locs == rhs.locs)
end

Base.adjoint(blk::ControlBlock{N}) where {N} =
    ControlBlock{N}(blk.ctrl_locs, blk.ctrl_config, adjoint(blk.content), blk.locs)

# NOTE: we only copy one hierachy (shallow copy) for each block
function Base.copy(ctrl::ControlBlock{N,BT,C,M}) where {BT,N,C,M}
    return ControlBlock{N,BT,C,M}(ctrl.ctrl_locs, ctrl.ctrl_config, ctrl.content, ctrl.locs)
end

function YaoBase.iscommute(x::ControlBlock{N}, y::ControlBlock{N}) where {N}
    if x.locs == y.locs
        return iscommute(x.content, y.content)
    elseif !any(l -> l in y.ctrl_locs, x.locs) && !any(l -> l in x.ctrl_locs, y.locs)
        return true
    else
        return iscommute_fallback(x, y)
    end
end


Base.@kwdef mutable struct Solver <: AMGXObject
    handle::API.AMGX_solver_handle = API.AMGX_solver_handle(C_NULL)
    resources::Union{Resources, Nothing} = nothing
    config::Union{Config, Nothing} = nothing
    mode::Union{Mode, Nothing} = nothing
    bound_matrix::Union{AMGXMatrix, Nothing} = nothing
    function Solver(handle::API.AMGX_solver_handle, resources::Union{Resources, Nothing}, mode::Union{Mode, Nothing},
                    config::Union{Config, Nothing}, bound_matrix::Union{AMGXMatrix, Nothing})
        solver = new(handle, resources, config, mode, bound_matrix)
        finalizer(warn_not_destroyed_on_finalize, solver)
        return solver
    end
end
function dec_refcount_parents(solver::Solver)
    dec_refcount!(solver.resources)
    dec_refcount!(solver.config)
    solver.resources = nothing
    solver.config = nothing
    solver.bound_matrix = nothing
    solver.mode = nothing
    nothing
end
get_api_destroy_call(::Type{Solver}) = API.AMGX_solver_destroy

function create!(solver::Solver, res::Resources, mode::Mode, config::Config)
    solver_handle_ptr = Ref{API.AMGX_solver_handle}()
    amgx_mode = API.AMGX_Mode(Int(mode))
    @checked API.AMGX_solver_create(solver_handle_ptr, res.handle, amgx_mode, config.handle)
    solver.handle = solver_handle_ptr[]
    solver.resources = res
    solver.config = config
    solver.mode = mode
    inc_refcount!(solver.resources)
    inc_refcount!(solver.config)
    return solver
end
Solver(res::Resources, mode::Mode, config::Config) = create!(Solver(), res, mode, config)

function setup!(solver::Solver, matrix::AMGXMatrix)
    @checked API.AMGX_solver_setup(solver.handle, matrix.handle)
    solver.bound_matrix = matrix
    return solver
end

function solve!(sol::AMGXVector, solver::Solver, rhs::AMGXVector; zero_inital_guess::Bool=false)
    if solver.bound_matrix === nothing
        error("no matrix attached to solver")
    elseif solver.bound_matrix.handle == C_NULL
        error("matrix has already been destroyed")
    end
    if zero_inital_guess
        @checked API.AMGX_solver_solve_with_0_initial_guess(solver.handle, rhs.handle, sol.handle)
    else
        @checked API.AMGX_solver_solve(solver.handle, rhs.handle, sol.handle)
    end
    return sol
end

function get_iterations_number(solver::Solver)
    n_ptr = Ref{Cint}()
    @checked API.AMGX_solver_get_iterations_number(solver.handle, n_ptr)
    # for some reason AMGX returns 1 + number of iterations.
    return Int(n_ptr[]) - 1
end

function get_iteration_residual(solver::Solver, iter::Int=get_iterations_number(solver), block_idx::Int=0)
    res_ptr = Ref{Float64}()
    @checked API.AMGX_solver_get_iteration_residual(solver.handle, iter, block_idx, res_ptr)
    return res_ptr[]
end

@enum SolverStatus begin
   SUCCESS = Int(API.AMGX_SOLVE_SUCCESS)
   FAILED = Int(API.AMGX_SOLVE_FAILED)
   DIVERGED = Int(API.AMGX_SOLVE_DIVERGED)
end

function get_status(solver::Solver)
    status_ptr = Ref{API.AMGX_SOLVE_STATUS}()
    @checked API.AMGX_solver_get_status(solver.handle, status_ptr)
    return SolverStatus(Int(status_ptr[]))
end


using DelimitedFiles

struct CharChar
    a::Char
    b::Char
end

struct CharString
    a::Char
    b::String
end

mutable struct Cart
    id::Int8
    position::CartesianIndex
    symbol::Char
    mod_turns::Int8
    alive::Bool
end

function islesscart(c1::Cart,c2::Cart) :: Bool
    return c1.position[1]<c2.position[1] || (c1.position[1]==c2.position[1] && c1.position[2]<c2.position[2])
end

up = '^'
down = 'v'
left = '<'
right = '>'

vertical = '|'
horizontal = '-'
curve_backslash = '\\'
curve_slash = '/'
intersection = '+'
track_symbols = [vertical,horizontal,curve_slash,curve_backslash,intersection]
carts_symbols = [up,down,left,right]
directions = ["left","straight","right"]

straightMoves = Dict{Char,CartesianIndex}(
    up=>CartesianIndex(-1,0),
    down=>CartesianIndex(1,0),
    left=>CartesianIndex(0,-1),
    right=>CartesianIndex(0,1))
curves = Dict{CharChar,Char}(
    CharChar(up,curve_slash)=>right,
    CharChar(up,curve_backslash)=>left,
    CharChar(down,curve_slash)=>left,
    CharChar(down,curve_backslash)=>right,
    CharChar(left,curve_slash)=>down,
    CharChar(left,curve_backslash)=>up,
    CharChar(right,curve_slash)=>up,
    CharChar(right,curve_backslash)=>down)

turns = Dict{CharString,Char}(
    CharString(up,"left")=>left,
    CharString(up,"right")=>right,
    CharString(down,"left")=>right,
    CharString(down,"right")=>left,
    CharString(left,"left")=>down,
    CharString(left,"right")=>up,
    CharString(right,"left")=>up,
    CharString(right,"right")=>down
)

function solve(input)
    raw = readdlm(input,'\n',String)
    n_col = length(raw[1])
    n_row = length(raw)
    grid = Array{Char,2}(undef,n_row,n_col)
    for y in 1:n_row
        for x in 1:n_col
            grid[y,x] = raw[y][x]
        end
    end
    carts=Array{Cart,1}()
    cart_id=1
    for cart_position in findall(c->c in carts_symbols,grid)
        symbol = grid[cart_position]
        push!(carts,Cart(cart_id,cart_position,symbol,0,true))
        if symbol in carts_symbols[1:2]
            grid[cart_position] = vertical
        else
            grid[cart_position] = horizontal
        end
        cart_id+=1
    end
    #@show carts
    #display(grid,carts)
    tick = 0
    while length(carts)>1
        carts=sort(carts,lt=islesscart)
        #@show carts
        #@show tick
        for cart in carts
            #@show (cart.id,cart.position)
            move(cart,grid)
            carts = crash(carts)
        end
        tick += 1
    end
    if length(carts)==1
        pos = carts[1].position
        x=pos[2]-1
        y=pos[1]-1
        println("Last alive cart position: (X,Y)=($x,$y) id $(carts[1].id)")
    end
end

function display(grid,carts)
    displayed = copy(grid)
    for cart in carts
        if displayed[cart.position] in carts_symbols
            displayed[cart.position]='X'
        else
            displayed[cart.position]=cart.symbol
        end
    end
    for row in 1:size(displayed)[1]
        println(join(displayed[row,:]))
    end
    println("")
end

function move(cart::Cart, grid)
    if cart.alive
        cart.position += straightMoves[cart.symbol]
        next_position = grid[cart.position]
        if next_position in track_symbols[3:4]
            cart.symbol = curves[CharChar(cart.symbol,next_position)]
        elseif next_position == intersection
            cart.mod_turns = (cart.mod_turns % length(directions)) + 1
            next_direction = directions[cart.mod_turns]
            if next_direction != "straight"
                cart.symbol = turns[CharString(cart.symbol,next_direction)]
            end
        end
    end
end

function crash(carts)
    num_carts = length(carts)
    for i in 1:num_carts-1
        cart1 = carts[i]
        for j in i+1:num_carts
            cart2 = carts[j]
            if cart1.position == cart2.position
                cart1.alive = false
                cart2.alive = false
                x=cart1.position[2]-1
                y=cart1.position[1]-1
                println("new crash ! (X,Y)=($x,$y) cart ids $(cart1.id) $(cart2.id)")
                return filter(c->c.alive,carts)
            end
        end
    end
    return carts
end

#solve("input_test.txt")
#solve("input_test.2.txt")
solve("input.txt")


module Waterflow

using LightGraphs, MetaGraphs, Distributed, DataFrames
using Infiltrator

# Can't use string, DLL location has to be a const
# (which makes sense but still, many hours wasted!)
# https://github.com/JuliaLang/julia/issues/29602
const ihacres = "../../ihacres_nim/lib/ihacres.dll"

"""@def macro

Inline code to avoid repetitious declarations.
"""
macro def(name, definition)
    return quote
        macro $(esc(name))()
            esc($(Expr(:quote, definition)))
        end
    end
end

@def add_preprefix begin
    if !isnothing(id_prefix)
        prefix = id_prefix * prefix
    end
end

include("Node.jl")
include("StreamNode.jl")
include("DamNode.jl")
include("Climate.jl")

"""Run a model attached to a node.

water_order: Volume of water to be extracted (in ML/timestep)
exchange: Volume of flux (in ML/timestep), where negative values are losses to the groundwater system.
"""
function run_node!(mg::MetaGraph, g::AbstractGraph, node_id::Int, climate::Climate, timestep::Int;
                   water_order=nothing, exchange=nothing)

    curr_node = get_prop(mg, node_id, :node)
    if checkbounds(Bool, curr_node.outflow, timestep)
        # already ran for this time step so no need to recurse further
        return curr_node.outflow[timestep], curr_node.level[timestep]
    end

    outflow = 0.0
    inflow = 0.0

    ins = inneighbors(g, node_id)
    if length(ins) == 0
        inflow = 0.0
        src_name = "Out-of-Catchment"
    else
        inflow = 0.0
        for i in ins
            src_name = get_prop(mg, i, :name)
            # Get inflow from previous node
            upstream_flow, upstream_level = run_node!(mg, g, i, climate, timestep)
            inflow += upstream_flow
        end
    end

    gauge_id = get_prop(mg, node_id, :name)
    rain, et = climate_values(curr_node, climate, timestep)
    if ismissing(rain) | ismissing(et)
        @warn "No climate data found for node $(node_id) ($(gauge_id)) using column markers $(climate.rainfall_id) and $(climate.et_id)"
        return 0.0
    end

    wo = 0.0
    if !isnothing(water_order)
        release_col = filter(x -> occursin(gauge_id, string(x))
                                  & occursin("releases", string(x)),
                                  names(water_order))

        wo = checkbounds(Bool, water_order.Date, timestep) ? water_order[timestep, release_col][1] : 0.0
    end

    ex = 0.0
    if !isnothing(exchange)
        exchange_col = filter(x -> occursin(gauge_id, x)
                                  & occursin("exchange", x),
                                  names(exchange))
        ex = checkbounds(Bool, exchange.Date, timestep) ? exchange[timestep, exchange_col][1] : 0.0
    end

    # Calculate outflow for this node
    func = get_prop(mg, node_id, :nfunc)
    if curr_node isa StreamNode
        outflow, level = func(curr_node, rain, et, inflow, 0.0)
    elseif curr_node isa DamNode
        outflow, level = func(curr_node, rain, et, inflow, wo, ex)
    else
        throw(ArgumentError("Unknown node type!"))
    end

    return outflow, level
end

include("Network.jl")

export @def
export ihacres, StreamNode, DamNode, Climate
export find_inlets_and_outlets, create_network, create_node
export update_params!, run_node!, reset!, sim_length
export climate_values

end  # end module


"""
    A BoundaryCondition is a Kernel... it just computes differently

    A BoundaryCondition _must_ supply a bid::Array{Int64} variable
    that specifies which boundary IDs this BC is applied to
"""
abstract BoundaryCondition <: Kernel
  # bids::Array{Int64}

" A NodalBC gets applied _at_ nodes and overwrites residual / jacobian entries "
abstract NodalBC <: BoundaryCondition

" An IntegratedBC is integrated across sidesets to compute residual / jacobians "
abstract IntegratedBC <: BoundaryCondition

" Get the boundary IDs this BoundaryCondition should be applied to "
function boundaryIDs(bc::BoundaryCondition)
    return bc.bids
end

" Generic `computeJacobian()` for Float64.  Just returns 0 "
function computeJacobian(bc::BoundaryCondition, var::Variable{Float64})
    return 0.
end

"""
    Generic `computeJacobian()` for Dual

    Uses Automatic Differentiation to compute the Jacobian.

    NOTE!  This is sub-optimal if you want both the residual and the Jacobian!
    In that case use `computeResidualAndJacobian!()` instead!
"""
function computeJacobian{N,T}(bc::BoundaryCondition, var::Variable{Dual{N,T}})
    residual = computeResidual(bc)

    return partials(residual)[var.id]
end

" Compute the Residual and Jacobian using calls to `computeJacobian()` "
function computeResidualAndJacobian!(residual::Array{Float64},
                                     var_jacobians::Array{Float64},
                                     vars::Array{Variable{Float64}},
                                     bc::BoundaryCondition)
    residual[1] = computeResidual(bc)

    for v in vars
        var_jacobians[v.id] = computeJacobian(bc, v)
    end
end

" Compute the Residual and Jacobian using Automatic Differentiation "
function computeResidualAndJacobian!{N,T}(residual::Array{Dual{N,T}},
                                          var_jacobians::Array{Float64},
                                          vars::Array{Variable{Dual{N,T}}},
                                          bc::BoundaryCondition)
    residual[1] = computeResidual(bc)

    for v in vars
        var_jacobians[v.id] = partials(residual[1])[v.id]
    end
end


using SeaPearl
using SeaPearlExtras
using ReinforcementLearning
const RL = ReinforcementLearning
using Flux
using GeometricFlux
using JSON
using BSON: @save, @load
using Dates
using Random

# -------------------
# Experience variables
# -------------------
nbEpisodes = 100
restartPerInstances = 20
evalFreq = 50
nbInstances = 10
nbRandomHeuristics = 1

nbNodes = 20
nbMinColor = 5
density = 0.95
# -------------------
# Generator
# -------------------
coloring_generator = SeaPearl.ClusterizedGraphColoringGenerator(nbNodes, nbMinColor, density)

#include("rewards.jl")
include("features.jl")

# -------------------
# Internal variables
# -------------------
SR = SeaPearl.DefaultStateRepresentation{BetterFeaturization, SeaPearl.DefaultTrajectoryState}
numInFeatures = SeaPearl.feature_length(SR)
#numGlobalFeature = SeaPearl.global_feature_length(SR)

# -------------------
# Agent definition
# -------------------
include("agents.jl")

# -------------------
# Value Heuristic definition
# -------------------

learnedHeuristic=SeaPearl.LearnedHeuristic{SR, SeaPearl.CPReward, SeaPearl.FixedOutput}(agent)
# Basic value-selection heuristic
selectMin(x::SeaPearl.IntVar; cpmodel=nothing) = SeaPearl.minimum(x.domain)
heuristic_min = SeaPearl.BasicHeuristic(selectMin)
function select_random_value(x::SeaPearl.IntVar; cpmodel=nothing)
    selected_number = rand(1:length(x.domain))
    i = 1
    for value in x.domain
        if i == selected_number
            return value
        end
        i += 1
    end
    @assert false "This should not happen"
end

randomHeuristics = []
for i in 1:nbRandomHeuristics
    push!(randomHeuristics, SeaPearl.BasicHeuristic(select_random_value))
end

valueSelectionArray = [learnedHeuristic, heuristic_min]
append!(valueSelectionArray, randomHeuristics)
# -------------------
# Variable Heuristic definition
# -------------------
variableSelection = SeaPearl.MinDomainVariableSelection{false}()

# -------------------
# -------------------
# Core function
# -------------------
# -------------------

function trytrain(nbEpisodes::Int)
    experienceTime = now()
    dir = mkdir(string("exp_",Base.replace("$(round(experienceTime, Dates.Second(3)))",":"=>"-")))
    expParameters = Dict(
        :nbEpisodes => nbEpisodes,
        :restartPerInstances => restartPerInstances,
        :evalFreq => evalFreq,
        :nbInstances => nbInstances,
        :nbRandomHeuristics => nbRandomHeuristics,
        nbNodes => nbNodes,
        nbMinColor => nbMinColor,
        density => density
    )
    open(dir*"/params.json", "w") do file
        JSON.print(file, expParameters)
    end

    metricsArray, eval_metricsArray = SeaPearl.train!(
        valueSelectionArray=valueSelectionArray,
        generator=coloring_generator,
        nbEpisodes=nbEpisodes,
        strategy=SeaPearl.ILDSearch(0),
        variableHeuristic=variableSelection,
        out_solver=true,
        verbose = false,
        evaluator=SeaPearl.SameInstancesEvaluator(valueSelectionArray,coloring_generator; evalFreq = evalFreq, nbInstances = nbInstances),
        restartPerInstances = restartPerInstances
    )

    #saving model weights
    model = agent.policy.learner.approximator
    @save dir*"/model_gc"*string(coloring_generator.n)*".bson" model

    SeaPearlExtras.storedata(metricsArray[1]; filename=dir*"/graph_coloring_$(nbNodes)_training_learned")
    SeaPearlExtras.storedata(metricsArray[2]; filename=dir*"/graph_coloring_$(nbNodes)_training_greedy")
    for i = 1:nbRandomHeuristics
        SeaPearlExtras.storedata(metricsArray[2+i]; filename=dir*"/graph_coloring_$(nbNodes)_training_random$(i)")
    end
    SeaPearlExtras.storedata(eval_metricsArray[:,1]; filename=dir*"/graph_coloring_$(nbNodes)_learned")
    SeaPearlExtras.storedata(eval_metricsArray[:,2]; filename=dir*"/graph_coloring_$(nbNodes)_greedy")
    for i = 1:nbRandomHeuristics
        SeaPearlExtras.storedata(eval_metricsArray[:,i+2]; filename=dir*"/graph_coloring_$(nbNodes)_random$(i)")
    end

    return metricsArray, eval_metricsArray
end

metricsArray, eval_metricsArray = trytrain(nbEpisodes)
nothing


struct AutoTracker <: AbstractADType end

function instantiate_function(f, x, ::AutoTracker, p, num_cons = 0)
    num_cons != 0 && error("AutoTracker does not currently support constraints")
    _f = (θ, args...) -> first(f.f(θ, p, args...))

    if f.grad === nothing
        grad = (res, θ, args...) -> res isa DiffResults.DiffResult ? DiffResults.gradient!(res, Tracker.data(Tracker.gradient(x -> _f(x, args...), θ)[1])) : res .= Tracker.data(Tracker.gradient(x -> _f(x, args...), θ)[1])
    else
        grad = f.grad
    end

    if f.hess === nothing
        hess = (res, θ, args...) -> error("Hessian based methods not supported with Tracker backend, pass in the `hess` kwarg")
    else
        hess = f.hess
    end

    if f.hv === nothing
        hv = (res, θ, args...) -> error("Hessian based methods not supported with Tracker backend, pass in the `hess` and `hv` kwargs")
    else
        hv = f.hv
    end

    return OptimizationFunction{false,AutoTracker,typeof(f),typeof(grad),typeof(hess),typeof(hv),Nothing,Nothing,Nothing}(f,AutoTracker(),grad,hess,hv,nothing,nothing,nothing)
end


# Sequence-form integrated cyber air defense, G = (N, Σ, g, C)
using Pkg; Pkg.activate(pwd())
using Revise, LinearAlgebra, FileIO, JLD2, ProgressMeter
using JuMP, Clp
using Clp: ClpCInterface

# load functions
dir = pwd()
include(joinpath(dir, "src\\ops_utility.jl"))
include(joinpath(dir, "src\\ops_build.jl"))
include(joinpath(dir, "src\\ops_methods.jl"))

#############
# best response
#############

# fix player 2's realization plan as uniformly random
# how to make better realization plan? Maybe pick a few pure strategies?
# greedy plan: defender assigns cyber assets to IADS with most covered value
function build_rp2(gos::GameOptSet)
    ns2, ns2_stage, na_stage = gos.ns2, gos.ns2_stage, gos.na_stage
    rp2 = zeros(Float64, ns2)
    ns_cumsum = cumsum(ns2_stage)
    p0 = 1.0
    p3 = p0 / na_stage[3]
    p6 = p3 / na_stage[6]
    rp2[1] = p0
    rp2[(ns_cumsum[1] + 1):ns_cumsum[2]] .= p3
    rp2[(ns_cumsum[2] + 1):ns_cumsum[3]] .= p6
    return rp2
end

function build_rp1(gos::GameOptSet)
    ns1, ns1_stage, na_stage = gos.ns1, gos.ns1_stage, gos.na_stage
    rp = zeros(Float64, ns1)
    ns_cumsum = cumsum(ns1_stage)
    p0 = 1.0
    p5 = p0 / na_stage[5]
    rp[1] = p0
    rp[(ns_cumsum[1] + 1):ns_cumsum[2]] .= p5
    return rp
end

function lp_best_response(gos::GameOptSet, rfixed::Vector{Float64}; fixedplayer::Int64 = 2)
    ns1, ns2, ni1, ni2, reward = gos.ns1, gos.ns2, gos.ni1, gos.ni2, gos.reward
    seqI1, seqI2, nextseqI1, nextseqI2 = gos.seqI1, gos.seqI2, gos.nextseqI1, gos.nextseqI2
    ind_nz = findall(!iszero, gos.reward)
    nsi, nso = fixedplayer == 1 ? (ns2, ns1) : (ns1, ns2)
    nii, nio = fixedplayer == 1 ? (ni2, ni1) : (ni1, ni2)
    seqIi, seqIo = fixedplayer == 1 ? (seqI2, seqI1) : (seqI1, seqI2)
    nextseqIi, nextseqIo = fixedplayer == 1 ? (nextseqI2, nextseqI1) : (nextseqI1, nextseqI2)
    println("Building model...")
    m = Model(Clp.Optimizer)
    @variable(m, r[1:nsi] >= 0)
    if fixedplayer == 1 # reward is player 1's reward(seq1, seq2) so we can't switch the indexes
        @objective(m, Max, sum(-1 * gos.reward[I] * rfixed[I[1]] * r[I[2]] for I in ind_nz))
    else
        @objective(m, Max, sum(gos.reward[I] * r[I[1]] * rfixed[I[2]] for I in ind_nz))
    end
    @constraint(m, r[1] == 1.0)
    # @progress for i in 1:nii
    for i in 1:nii
        @constraint(m, sum(r[si] for si in nextseqIi[i]) == r[seqIi[i]])
    end
    println("Solving model...")
    @show m
    status = solve(m)
    return status, getobjectivevalue(m), getvalue(r)
end

#############
# Nash Equilibrium
#############

# TODO: make building the objective function faster in lp_nash
function lp_nash(gos::GameOptSet; timelimit = 5 * 60)
    ns2, ni1, ni2, reward = gos.ns2, gos.ni1, gos.ni2, gos.reward
    ns1_stage, na_stage = gos.ns1_stage, gos.na_stage
    seqI2, nextseqI2 = gos.seqI2, gos.nextseqI2
    println("Building model...")
    mne = Model(solver = ClpSolver(MaximumSeconds = timelimit))
    @variable(mne, rp2[1:ns2] >= 0)
    @variable(mne, v1[1:ni1])
    @variable(mne, v0)
    @objective(mne, Min, v0)
    @constraint(mne, v0 - sum(v1[j] for j in 1:ni1) >= sum(reward[1, s2] * rp2[s2] for s2 in 1:ns2))
    nvconstraint = 1
    # @progress for s1 in (1 + 1):(1 + ns1_stage[2])
    for s1 in (1 + 1):(1 + ns1_stage[2])
        iind = ceil(Int, (s1 - 1) / na_stage[5]) # Σ1 is in blocks na_stage[5] actions for each infoset
        @constraint(mne, v1[iind] >= sum(reward[s1, s2] * rp2[s2] for s2 in 1:ns2))
        nvconstraint += 1
    end
    @constraint(mne, rp2[1] == 1)
    # @progress for i in 1:ni2
    for i in 1:ni2
        @constraint(mne, sum(rp2[s2] for s2 in nextseqI2[i]) == rp2[seqI2[i]])
    end
    @show mne
    println("Solving model...")
    status, t_solve = @timed solve(mne)
    println("obj value: $(getobjectivevalue(mne)), status: $status")
    rp1 = ClpCInterface.dual_row_solution(getrawsolver(mne))[1:nvconstraint] # player 1 strategy is the dual of player 1 values for each sequence constraint
    return status, getobjectivevalue(mne), rp1, getvalue(rp2), t_solve
end

#############
# Robust Best Response
#############

function lp_robust_br(gos::GameOptSet, rp2lb, rp2ub; timelimit = 5 * 60)
    ns2, ni1, ni2, reward = gos.ns2, gos.ni1, gos.ni2, gos.reward
    ns1_stage, na_stage = gos.ns1_stage, gos.na_stage
    seqI2, nextseqI2 = gos.seqI2, gos.nextseqI2
    println("Building model...")
    mr = Model(solver = ClpSolver(MaximumSeconds = timelimit))
    @variable(mr, rp2[1:ns2])
    @variable(mr, v1[1:ni1])
    @variable(mr, v0)
    @objective(mr, Min, v0)
    @constraint(mr, v0 - sum(v1[j] for j in 1:ni1) >= sum(reward[1, s2] * rp2[s2] for s2 in 1:ns2))
    # @progress for s1 in (1 + 1):(1 + ns1_stage[2])
    nvconstraint = 1
    for s1 in (1 + 1):(1 + ns1_stage[2])
        iind = ceil(Int, (s1 - 1) / na_stage[5]) # Σ1 is in blocks na_stage[5] actions for each infoset
        @constraint(mr, v1[iind] >= sum(reward[s1, s2] * rp2[s2] for s2 in 1:ns2))
        nvconstraint += 1
    end
    @constraint(mr, rp2[1] == 1)
    # @progress for i in 1:ni2
    for i in 1:ni2
        @constraint(mr, sum(rp2[s2] for s2 in nextseqI2[i]) == rp2[seqI2[i]])
    end
    @constraint(mr, rp2 .>= rp2lb)
    @constraint(mr, rp2 .<= rp2ub)
    @show mr
    println("Solving model...")
    status, t_solve = @timed solve(mr)
    rp1 = ClpCInterface.dual_row_solution(getrawsolver(mr))[1:nvconstraint] # player 1 strategy is the dual of player 1 values for each sequence constraint
    return status, getobjectivevalue(mr), rp1, getvalue(rp2), t_solve
end


using CSyntax

@static if Sys.isapple()
    const VERSION_MAJOR = 4
    const VERSION_MINOR = 1
end

include(joinpath(@__DIR__, "glutils.jl"))

# init window
width, height = fb_width, fb_height = 640, 480
window = startgl(width, height)

glEnable(GL_DEPTH_TEST)
glDepthFunc(GL_LESS)

# vertex data
points = GLfloat[ 0.0, 0.75, 0.0,
                  0.5, 0.25, 0.0,
                 -0.5, 0.25, 0.0,
                  0.5,-0.25, 0.0,
                  0.0,-0.75, 0.0,
                 -0.5,-0.25, 0.0]

# create buffers located in the memory of graphic card
points_vbo = GLuint(0)
@c glGenBuffers(1, &points_vbo)
glBindBuffer(GL_ARRAY_BUFFER, points_vbo)
glBufferData(GL_ARRAY_BUFFER, sizeof(points), points, GL_STATIC_DRAW)

# create VAO
vao = GLuint(0)
@c glGenVertexArrays(1, &vao)
glBindVertexArray(vao)
glBindBuffer(GL_ARRAY_BUFFER, points_vbo)
glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, C_NULL)
glEnableVertexAttribArray(0)

# load and compile shaders from file
vert_shader = createshader(joinpath(@__DIR__, "tess.vert"), GL_VERTEX_SHADER)
tesc_shader = createshader(joinpath(@__DIR__, "tess.tesc"), GL_TESS_CONTROL_SHADER)
tese_shader = createshader(joinpath(@__DIR__, "tess.tese"), GL_TESS_EVALUATION_SHADER)
frag_shader = createshader(joinpath(@__DIR__, "tess.frag"), GL_FRAGMENT_SHADER)

# link program
shader_prog = createprogram(vert_shader, tesc_shader, tese_shader, frag_shader)

inner_tess_factor_loc = glGetUniformLocation(shader_prog, "tess_fac_inner")
outer_tess_factor_loc = glGetUniformLocation(shader_prog, "tess_fac_outer")

glEnable(GL_CULL_FACE)
glCullFace(GL_BACK)
glFrontFace(GL_CW)
glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)
glPatchParameteri(GL_PATCH_VERTICES, 3)

# set background color to gray
glClearColor(0.2, 0.2, 0.2, 1.0)

# render
let
inner_tess_factor = GLfloat(1)
outer_tess_factor = GLfloat(4)
qWasDown = false
aWasDown = false
wWasDown = false
sWasDown = false
updatefps = FPSCounter()
while !GLFW.WindowShouldClose(window)
    updatefps(window)
    # clear drawing surface
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glViewport(0, 0, GLFW.GetFramebufferSize(window)...)
    # drawing
    glUseProgram(shader_prog)
    glBindVertexArray(vao)
    glDrawArrays(GL_PATCHES, 0, 6)
    # check and call events
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_Q)
        if !qWasDown
            inner_tess_factor += 1
            println("inner tessellation factor = $inner_tess_factor.")
            glUniform1f(inner_tess_factor_loc, inner_tess_factor)
            qWasDown = true
        end
    else
        qWasDown = false
    end
    if GLFW.GetKey(window, GLFW.KEY_A)
        if !aWasDown
            inner_tess_factor -= 1
            println("inner tessellation factor = $inner_tess_factor.")
            glUniform1f(inner_tess_factor_loc, inner_tess_factor)
            aWasDown = true
        end
    else
        aWasDown = false
    end
    if GLFW.GetKey(window, GLFW.KEY_W)
        if !wWasDown
            outer_tess_factor += 1
            println("outer tessellation factor = $outer_tess_factor.")
            glUniform1f(outer_tess_factor_loc, outer_tess_factor)
            wWasDown = true
        end
    else
        wWasDown = false
    end
    if GLFW.GetKey(window, GLFW.KEY_S)
        if !sWasDown
            outer_tess_factor -= 1
            println("outer tessellation factor = $outer_tess_factor.")
            glUniform1f(outer_tess_factor_loc, outer_tess_factor)
            sWasDown = true
        end
    else
        sWasDown = false
    end
    # swap the buffers
    GLFW.SwapBuffers(window)
end
end # let
GLFW.DestroyWindow(window)