text:
################################################################################
#                               BASE EXTENSIONS
################################################################################
# Extend Base.copy for new variable types
Base.copy(v::GeneralVariableRef) = v
Base.copy(v::DispatchVariableRef) = v

# Extend Base.:(==) for GeneralVariableRef
function Base.:(==)(v::T, w::T)::Bool where {T <: GeneralVariableRef}
    return v.model === w.model && v.raw_index == w.raw_index &&
           v.index_type == w.index_type && v.param_index == w.param_index
end

# Extend Base.:(==) for DispatchVariableRef
function Base.:(==)(v::T, w::U)::Bool where {T <: DispatchVariableRef,
                                             U <: DispatchVariableRef}
    return v.model === w.model && v.index == w.index
end

# Extend Base.broadcastable
Base.broadcastable(v::GeneralVariableRef) = Ref(v)
Base.broadcastable(v::DispatchVariableRef) = Ref(v)

# Extend Base.length
Base.length(v::GeneralVariableRef)::Int = 1
Base.length(v::DispatchVariableRef)::Int = 1

# Extend JuMP functions
JuMP.isequal_canonical(v::GeneralVariableRef, w::GeneralVariableRef)::Bool = v == w
JuMP.isequal_canonical(v::DispatchVariableRef, w::DispatchVariableRef)::Bool = v == w
JuMP.variable_type(model::InfiniteModel)::DataType = GeneralVariableRef

# Extract the root name of a variable reference (removes the bracketed container indices)
function _remove_name_index(vref::GeneralVariableRef)::String
    name = JuMP.name(vref)
    first_bracket = findfirst(isequal('['), name)
    if first_bracket === nothing
        return name
    else
        # Hacky fix to handle invalid Unicode
        try
            return name[1:first_bracket-1]
        catch
            return name[1:first_bracket-2]
        end
    end
end

# Define basic attribute getters
_index_type(vref::GeneralVariableRef)::DataType = vref.index_type
_raw_index(vref::GeneralVariableRef)::Int = vref.raw_index
_param_index(vref::GeneralVariableRef)::Int = vref.param_index

################################################################################
#                          BASIC REFERENCE ACCESSERS
################################################################################
"""
    JuMP.index(vref::GeneralVariableRef)::AbstractInfOptIndex

Extend [`JuMP.index`](@ref JuMP.index(::JuMP.VariableRef)) to return the
appropriate index of `vref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; m = InfiniteModel(); @hold_variable(m, vref))
julia> index(vref)
HoldVariableIndex(1)
```
"""
function JuMP.index(vref::GeneralVariableRef)::AbstractInfOptIndex
    index_type = _index_type(vref)
    if index_type == DependentParameterIndex
        return index_type(DependentParametersIndex(_raw_index(vref)),
                           _param_index(vref))
    else
        return index_type(_raw_index(vref))
    end
end

"""
    JuMP.index(vref::DispatchVariableRef)::AbstractInfOptIndex

Extend [`JuMP.index`](@ref JuMP.index(::JuMP.VariableRef)) to return the
appropriate index of `vref`.
"""
function JuMP.index(vref::DispatchVariableRef)::AbstractInfOptIndex
    return vref.index
end

"""
    JuMP.owner_model(vref::GeneralVariableRef)::InfiniteModel

Extend [`JuMP.owner_model`](@ref JuMP.owner_model(::JuMP.AbstractVariableRef)) to
return the model where `vref` is stored.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; m = InfiniteModel(); @hold_variable(m, 0 <= vref <= 1))
julia> owner_model(vref)
An InfiniteOpt Model
Feasibility problem with:
Finite Parameters: 0
Infinite Parameters: 0
Variable: 1
Derivatives: 0
Measures: 0
`HoldVariableRef`-in-`MathOptInterface.GreaterThan{Float64}`: 1 constraint
`HoldVariableRef`-in-`MathOptInterface.LessThan{Float64}`: 1 constraint
Names registered in the model: vref
Optimizer model backend information:
Model mode: AUTOMATIC
CachingOptimizer state: NO_OPTIMIZER
Solver name: No optimizer attached.
```
"""
function JuMP.owner_model(vref::GeneralVariableRef)::InfiniteModel
    return vref.model
end

"""
    JuMP.owner_model(vref::DispatchVariableRef)::InfiniteModel

Extend [`JuMP.owner_model`](@ref JuMP.owner_model(::JuMP.AbstractVariableRef)) to
return the model where `vref` is stored.
"""
function JuMP.owner_model(vref::DispatchVariableRef)::InfiniteModel
    return vref.model
end

################################################################################
#                          DISPATCH VARIABLE MAKERS
################################################################################
"""
    dispatch_variable_ref(model::InfiniteModel, index::AbstractInfOptIndex)

Return the variable reference associated the type of `index`. This needs to be
defined for each variable reference type.
"""
function dispatch_variable_ref end

"""
    dispatch_variable_ref(vef::GeneralVariableRef)::DispatchVariableRef

Return the concrete [`DispatchVariableRef`](@ref) this associated with `vref`.
This relies on `dispatch_variable_ref` being extended for the index type,
otherwise an `MethodError` is thrown.
"""
function dispatch_variable_ref(vref::GeneralVariableRef)::DispatchVariableRef
    model = JuMP.owner_model(vref)
    idx = JuMP.index(vref)
    return dispatch_variable_ref(model, idx)
end

################################################################################
#                          CORE DATA METHODS
################################################################################
"""
    _add_data_object(model::InfiniteModel, object::AbstractDataObject)::ObjectIndex

Add `object` to the appropriate `CleverDict` in `model` and return the its
index. This needs to be defined for the type of `object`. These definitions
need to use `MOIUC.add_item` to add the object to the `CleverDict`.
"""
function _add_data_object end

"""
    _data_dictionary(vref::DispatchVariableRef)::MOIUC.CleverDict

Return the `CleverDict` that stores data objects for the type of `vref`. This
needs to be defined for the type of `vref`.
"""
function _data_dictionary end

"""
    _data_dictionary(vref::GeneralVariableRef)::MOIUC.CleverDict

Return the `CleverDict` that stores data objects for the type of `vref`. It
relies on `_data_dictionary` being defined
for the underlying `DispatchVariableRef`, otherwise an `MethodError` is thrown.
"""
function _data_dictionary(vref::GeneralVariableRef)::MOIUC.CleverDict
    return _data_dictionary(dispatch_variable_ref(vref))
end

"""
    _data_object(vref::DispatchVariableRef)::AbstractDataObject

Return the data object associated with `vref`, in other words the object its
index points to in the `InfiniteModel`. This needs to be defined for the type
of `vref`. This should use `_data_dictionary` to access the `CleverDict` that
the object is stored in.
"""
function _data_object end

"""
    _data_object(vref::GeneralVariableRef)::AbstractDataObject

Return the data object associated with `vref`, in other words the object its
index points to in the `InfiniteModel`. It relies on `_data_object` being defined
for the underlying `DispatchVariableRef`, otherwise an `MethodError` is thrown.
"""
function _data_object(vref::GeneralVariableRef)::AbstractDataObject
    return _data_object(dispatch_variable_ref(vref))
end

"""
    _delete_data_object(vref::DispatchVariableRef)::Nothing

Delete the concrete `AbstractDataObject` associated with `vref`.
"""
function _delete_data_object(vref::DispatchVariableRef)::Nothing
    delete!(_data_dictionary(vref), JuMP.index(vref))
    return
end

################################################################################
#                                 NAME METHODS
################################################################################
# Dispatch fallback
function JuMP.name(vref::DispatchVariableRef)
    throw(ArgumentError("`JuMP.name` not defined for variable reference type " *
                        "`$(typeof(vref))`."))
end

"""
    JuMP.name(vref::GeneralVariableRef)::String

Extend [`JuMP.name`](@ref JuMP.name(::JuMP.VariableRef)) to
return the name of `vref`. It relies on `JuMP.name` being defined
for the underlying `DispatchVariableRef`, otherwise an `ArugmentError` is thrown.
"""
function JuMP.name(vref::GeneralVariableRef)::String
    return JuMP.name(dispatch_variable_ref(vref))
end

# Dispatch fallback
function JuMP.set_name(vref::DispatchVariableRef, name::String)
    throw(ArgumentError("`JuMP.set_name` not defined for variable reference type " *
                        "`$(typeof(vref))`."))
end

"""
    JuMP.set_name(vref::GeneralVariableRef, name::String)::Nothing

Extend [`JuMP.set_name`](@ref JuMP.name(::JuMP.VariableRef, ::String)) to
set the name of `vref`. It relies on `JuMP.set_name` being defined
for the underlying `DispatchVariableRef`, otherwise an `ArugmentError` is thrown.
"""
function JuMP.set_name(vref::GeneralVariableRef, name::String)::Nothing
    return JuMP.set_name(dispatch_variable_ref(vref), name)
end

################################################################################
#                             VALIDITY METHODS
################################################################################
"""
    JuMP.is_valid(model::InfiniteModel, vref::DispatchVariableRef)::Bool

Extend [`JuMP.is_valid`](@ref JuMP.is_valid(::JuMP.Model, ::JuMP.VariableRef)) to
return `Bool` if `vref` is a valid reference.
"""
function JuMP.is_valid(model::InfiniteModel, vref::DispatchVariableRef)::Bool
    return model === JuMP.owner_model(vref) &&
           haskey(_data_dictionary(vref), JuMP.index(vref))
end

# DependentParameterRef
function JuMP.is_valid(model::InfiniteModel, vref::DependentParameterRef)::Bool
    return model === JuMP.owner_model(vref) &&
           haskey(_data_dictionary(vref), JuMP.index(vref).object_index)
end

"""
    JuMP.is_valid(model::InfiniteModel, vref::GeneralVariableRef)::Bool

Extend [`JuMP.is_valid`](@ref JuMP.is_valid(::JuMP.Model, ::JuMP.VariableRef)) to
return `Bool` if `vref` is a valid reference.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @hold_variable(model, vref))
julia> is_valid(model, vref)
true
```
"""
function JuMP.is_valid(model::InfiniteModel, vref::GeneralVariableRef)::Bool
    return JuMP.is_valid(model, dispatch_variable_ref(vref))
end

################################################################################
#                             CORE OBJECT METHODS
################################################################################
"""
    _core_variable_object(vref::DispatchVariableRef)::Union{InfOptParameter, InfOptVariable, Measure}

Return the core object that `vref` points to. This needs to be extended for type
of `vref`. This should use `_data_object` to access the data object where the
variable object is stored.
"""
function _core_variable_object end

"""
    _core_variable_object(vref::GeneralVariableRef)::Union{InfOptParameter, InfOptVariable, Measure}

Return the core object that `vref` points to. This is enabled
with appropriate definitions of `_core_variable_object` for the
underlying `DispatchVariableRef`, otherwise an `MethodError` is thrown.
"""
function _core_variable_object(vref::GeneralVariableRef)
    return _core_variable_object(dispatch_variable_ref(vref))
end

"""
    _set_core_variable_object(vref::DispatchVariableRef, object)::Nothing

Sets the core object that `vref` points to `object`. This needs to be extended
for types of `vref` and `object`. This should use `_data_object` to access the
data object where the variable object is stored.
"""
function _set_core_variable_object end

################################################################################
#                               DEPENDENCY METHODS
################################################################################
# Define wrappers for internal usage methods and their templates
for op = (:_infinite_variable_dependencies, :_reduced_variable_dependencies,
          :_point_variable_dependencies, :_measure_dependencies,
          :_constraint_dependencies, :_derivative_dependencies,
          :_derivative_constraint_dependencies, :_parameter_function_dependencies,
          :_generative_measures)
    @eval begin
        # define the api template
        func = $op
        """
            $func(vref::DispatchVariableRef)::Vector{AbstractInfOptIndex}

        Return the indices of these entities that depend on `vref`. This needs to
        be extended for type of `vref`. This should use `_data_object` to access the
        data object where the name is stored if appropriate.
        """
        function $(op) end
        # define the dispatch version
        """
            $func(vref::GeneralVariableRef)::Vector{AbstractInfOptIndex}

        Return the indices of these entities that depend on `vref`. This is enabled
        with appropriate definitions of `$func` for the
        underlying `DispatchVariableRef`, otherwise an `MethodError` is thrown.
        """
        function $op(vref::GeneralVariableRef)
            return $op(dispatch_variable_ref(vref))
        end
    end
end

################################################################################
#                              USED BY METHODS
################################################################################
# Define the usage method wrappers and their fallbacks
for op = (:used_by_infinite_variable, :used_by_reduced_variable,
          :used_by_point_variable, :used_by_measure, :used_by_constraint,
          :used_by_objective, :used_by_derivative, :is_used, :has_derivative_constraints,
          :used_by_parameter_function)
    @eval begin
        # define the fallback method
        func = $op
        function $op(vref::DispatchVariableRef)
            str = string("`", func, "` not defined for variable reference type " *
                         "`$(typeof(vref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            $func(vref::GeneralVariableRef)::Bool

        Define `$func` for general variable references. It relies on `$func`
        being defined for the underlying `DispatchVariableRef`, otherwise an
        `ArugmentError` is thrown. See the underlying docstrings for more
        information.
        """
        function $op(vref::GeneralVariableRef)::Bool
            return $op(dispatch_variable_ref(vref))
        end
    end
end

################################################################################
#                              DELETE METHODS
################################################################################
# Dispatch fallback
function JuMP.delete(model::InfiniteModel, vref)
    throw(ArgumentError("`JuMP.delete` not defined for variable reference type " *
                        "`$(typeof(vref))`."))
end

"""
    JuMP.delete(model::InfiniteModel, vref::GeneralVariableRef)::Nothing

Extend [`JuMP.delete`](@ref JuMP.delete(::JuMP.Model, ::JuMP.VariableRef)) to
delete `vref` and its dependencies. It relies on `JuMP.delete`
being defined for the underlying `DispatchVariableRef`, otherwise an
`ArugmentError` is thrown.
"""
function JuMP.delete(model::InfiniteModel, vref::GeneralVariableRef)::Nothing
    return JuMP.delete(model, dispatch_variable_ref(vref))
end

"""
    JuMP.delete(model::InfiniteModel,
                prefs::AbstractArray{<:GeneralVariableRef})::Nothing

Extend `JuMP.delete` to delete a group of dependent infinite parameters and
their dependencies. An `ArugmentError` is thrown if `prefs` are not dependent
infinite parameters.
"""
function JuMP.delete(model::InfiniteModel,
                     prefs::AbstractArray{<:GeneralVariableRef})::Nothing
    return JuMP.delete(model, dispatch_variable_ref.(prefs))
end

################################################################################
#                              PARAMETER METHODS
################################################################################
"""
    _parameter_number(pref::DispatchVariableRef)::Int

Return the parameter creation number for `pref` assuming it is an infinite
parameter. This needs to be defined for the type of `pref`. This should use
the `_data_object` to get the number.
"""
function _parameter_number end

"""
    _parameter_number(pref::GeneralVariableRef)::Int

Return the parameter creation number for `pref` assuming it is an infinite
parameter. It relies on `_parameter_number` being properly defined for the
underlying `DispatchVariableRef`, otherwise an `MethodError` is thrown.
"""
function _parameter_number(pref::GeneralVariableRef)::Int
    return _parameter_number(dispatch_variable_ref(pref))
end

"""
    _object_number(pref::DispatchVariableRef)::Int

Return the object number for `pref` assuming it is an infinite
parameter. This needs to be defined for the type of `pref`. This should use
the `_data_object` to get the number.
"""
function _object_number end

"""
    _object_number(pref::GeneralVariableRef)::Int

Return the object number for `pref` assuming it is an infinite
parameter. It relies on `_object_number` being properly defined for the
underlying `DispatchVariableRef`, otherwise an `MethodError` is thrown.
"""
function _object_number(pref::GeneralVariableRef)::Int
    return _object_number(dispatch_variable_ref(pref))
end

# Define 1 argument user method wrappers and their fallbacks
for op = (:infinite_set, :num_supports, :significant_digits, :has_supports,
          :supports, :delete_supports, :fill_in_supports!, :parameter_value,
          :derivative_method, :has_generative_supports, :has_internal_supports,
          :add_generative_supports, :raw_function, :generative_support_info)
    @eval begin
        # define the fallback method
        func = $op
        function $op(pref; kwargs...)
            str = string("`", func, "` not defined for variable reference type " *
                         "`$(typeof(pref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            $func(prefs; [kwargs...])

        Define `$func` for general variable references. It relies on `$func`
        being defined for the underlying `DispatchVariableRef`, otherwise an
        `ArugmentError` is thrown. See the underlying docstrings for more
        information. Note that this is a auto generated wrapper and the underlying
        method may or may not use `kwargs`.
        """
        function $op(prefs::Union{GeneralVariableRef, AbstractArray{<:GeneralVariableRef}}; kwargs...)
            return $op(dispatch_variable_ref.(prefs); kwargs...)
        end
    end
end

# Dispatch fallback
function set_infinite_set(pref, set::AbstractInfiniteSet)
    throw(ArgumentError("`set_infinite_set` not defined for variable reference type(s) " *
                        "`$(typeof(pref))`."))
end

"""
    set_infinite_set(pref::GeneralVariableRef, set::InfiniteScalarSet)::Nothing

Specify the scalar infinite set of the infinite parameter `pref` to `set`. Note
this will reset/delete all the supports contained in the
underlying parameter object. Also, errors if `pref` is used
by a measure. An `ArgumentError` is thrown if `pref` is not an infinite parameter.
"""
function set_infinite_set(pref::GeneralVariableRef,
                          set::InfiniteScalarSet)::Nothing
    return set_infinite_set(dispatch_variable_ref(pref), set)
end

"""
    set_infinite_set(prefs::AbstractArray{<:GeneralVariableRef},
                     set::InfiniteArraySet)::Nothing

Specify the multi-dimensional infinite set of the dependent infinite parameters
`prefs` to `set`. Note this will reset/delete all the supports contained in the
underlying [`DependentParameters`](@ref) object. This will error if the not all
of the dependent infinite parameters are included or if any of them are used by
measures. An `ArgumentError` is thrown if `prefs` are not dependent infinite
parameters.
"""
function set_infinite_set(prefs::AbstractArray{<:GeneralVariableRef},
                          set::InfiniteArraySet)::Nothing
    return set_infinite_set(dispatch_variable_ref.(prefs), set)
end

# Better fallbacks for supports
function supports(pref::DispatchVariableRef; kwargs...)
    throw(ArgumentError("`supports` not defined for variable reference type(s) " *
                        "`$(typeof(pref))`."))
end
function supports(prefs::AbstractArray; kwargs...)
    throw(ArgumentError("`supports` not defined for variable reference type(s) " *
                        "`$(typeof(prefs))`."))
end

# Dispatch fallback
function set_supports(pref, supports; kwargs...)
    throw(ArgumentError("`set_supports` not defined for variable reference type(s) " *
                        "`$(typeof(pref))`."))
end

"""
    set_supports(pref::GeneralVariableRef, supports::Union{Real, Vector{<:Real}};
                 [force::Bool = false])::Nothing

Set the support points associated with a single infinite
parameter `pref`. An `ArgumentError` is thrown if `pref` is not an independent
infinite parameter.
"""
function set_supports(pref::GeneralVariableRef,
                      supports::Union{Real, Vector{<:Real}};
                      force::Bool = false,
                      label::Type{<:AbstractSupportLabel} = UserDefined
                      )::Nothing
    return set_supports(dispatch_variable_ref(pref), supports,
                        force = force, label = label)
end

"""
    set_supports(
        prefs::Union{Vector{GeneralVariableRef}, AbstractArray{<:GeneralVariableRef}},
        supports::Union{Array{<:Real, 2}, AbstractArray{<:Vector{<:Real}}};
        [force::Bool = false]
        )::Nothing

Set the support points associated with dependent infinite
parameters `prefs`. An `ArgumentError` is thrown if `prefs` is are not
dependent infinite parameters.
"""
function set_supports(prefs::AbstractArray{<:GeneralVariableRef},
                      supports::Union{Array{<:Real, 2}, AbstractArray{<:Vector{<:Real}}};
                      label::Type{<:AbstractSupportLabel} = UserDefined,
                      force::Bool = false
                      )::Nothing
    return set_supports(dispatch_variable_ref.(prefs), supports, label = label,
                        force = force)
end

# Dispatch fallback
function add_supports(pref, supports; kwargs...)
    throw(ArgumentError("`add_supports` not defined for variable reference type(s) " *
                        "`$(typeof(pref))`."))
end

"""
    add_supports(pref::GeneralVariableRef,
                 supports::Union{Real, Vector{<:Real}})::Nothing

Add the support points `supports` to a single infinite
parameter `pref`. An `ArgumentError` is thrown if `pref` is not an independent
infinite parameter.
"""
function add_supports(pref::GeneralVariableRef,
                      supports::Union{Real, Vector{<:Real}};
                      check::Bool = true,
                      label::Type{<:AbstractSupportLabel} = UserDefined
                      )::Nothing
    return add_supports(dispatch_variable_ref(pref), supports,
                        check = check, label = label)
end

"""
    add_supports(
        prefs::Union{Vector{GeneralVariableRef}, AbstractArray{<:GeneralVariableRef}},
        supports::Union{Array{<:Real, 2}, AbstractArray{<:Vector{<:Real}}}
        )::Nothing

Add the support points `supports` to the dependent infinite
parameters `prefs`. An `ArgumentError` is thrown if `prefs` is are not
dependent infinite parameters.
"""
function add_supports(prefs::AbstractArray{<:GeneralVariableRef},
                      supports::Union{Array{<:Real, 2}, AbstractArray{<:Vector{<:Real}}};
                      label::Type{<:AbstractSupportLabel} = UserDefined,
                      check::Bool = true
                      )::Nothing
    return add_supports(dispatch_variable_ref.(prefs), supports, label = label,
                        check = check)
end

# Fallback
function JuMP.set_value(vref::DispatchVariableRef, value::Real)
    throw(ArgumentError("`JuMP.set_value` not defined for variable reference type " *
                        "`$(typeof(vref))`."))
end

"""
    JuMP.set_value(vref::DispatchVariableRef, value::Real)::Nothing

Extend `JuMP.set_value` to
set the value of `vref`. It relies on `JuMP.set_value`
being defined for the underlying `DispatchVariableRef`, otherwise an
`ArugmentError` is thrown.
"""
function JuMP.set_value(vref::GeneralVariableRef, value::Real)::Nothing
    return JuMP.set_value(dispatch_variable_ref(vref), value)
end

# Dispatch fallback
function set_derivative_method(pref::DispatchVariableRef, method)
    throw(ArgumentError("`set_derivative_method` not defined for variable reference type(s) " *
                        "`$(typeof(pref))`."))
end

"""
    set_derivative_method(pref::GeneralVariableRef,
                          method::AbstractDerivativeMethod
                          )::Nothing

Specify the numerical derivative evaluation technique associated with `pref`.
An `ArgumentError` is thrown if `pref` is not an infinite parameter.
"""
function set_derivative_method(pref::GeneralVariableRef,
                               method::AbstractDerivativeMethod
                               )::Nothing
    return set_derivative_method(dispatch_variable_ref(pref), method)
end

# Define parameter status setters
for op = (:_set_has_generative_supports, :_set_has_internal_supports,
          :_set_has_derivative_constraints)
    @eval begin
        # define the fallback method
        func = $op
        function $op(vref::DispatchVariableRef, status)
            str = string("`", func, "` not defined for variable reference type " *
                         "`$(typeof(vref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        function $op(vref::GeneralVariableRef, status::Bool)::Nothing
            return $op(dispatch_variable_ref(vref), status)
        end
    end
end

################################################################################
#                              VARIABLE METHODS
################################################################################
# Define single argument variable method wrappers and their fallbacks
for op = (:raw_parameter_refs, :parameter_refs, :parameter_list,
          :start_value_function, :reset_start_value_function,
          :infinite_variable_ref, :eval_supports, :raw_parameter_values,
          :parameter_values, :parameter_bounds, :delete_parameter_bounds)
    @eval begin
        # define the fallback method
        func = $op
        function $op(vref::DispatchVariableRef)
            str = string("`", func, "` not defined for variable reference type " *
                         "`$(typeof(vref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            $func(vref::GeneralVariableRef)

        Define `$func` for general variable references. It relies on `$func`
        being defined for the underlying `DispatchVariableRef`, otherwise an
        `ArugmentError` is thrown. See the underlying docstrings for more
        information.
        """
        function $op(vref::GeneralVariableRef)
            return $op(dispatch_variable_ref(vref))
        end
    end
end

# Dispatch fallback
function set_start_value_function(vref::DispatchVariableRef, start)
    throw(ArgumentError("`set_start_value_function` not defined for variable reference type " *
                        "`$(typeof(vref))`."))
end

"""
    set_start_value_function(vref::GeneralVariableRef, start::Union{Real, Function})::Nothing

Set the start value function of `vref`. It relies on `set_start_value_function`
being defined for the underlying `DispatchVariableRef`, otherwise an
`ArugmentError` is thrown.
"""
function set_start_value_function(vref::GeneralVariableRef, start)::Nothing
    return set_start_value_function(dispatch_variable_ref(vref), start)
end

"""
    has_parameter_bounds(vref::GeneralVariableRef)::Bool

Return a `Bool` indicating if `vref` is limited to a sub-domain as defined
by parameter bound.
"""
function has_parameter_bounds(vref::GeneralVariableRef)::Bool
    return has_parameter_bounds(dispatch_variable_ref(vref))
end

# Dispatch fallback
function set_parameter_bounds(vref::DispatchVariableRef, bounds; kwargs...)
    throw(ArgumentError("`set_parameter_bounds` not defined for variable reference type(s) " *
                        "`$(typeof(vref))`."))
end

"""
    set_parameter_bounds(vref::GeneralVariableRef,
                         bounds::ParameterBounds{GeneralVariableRef};
                         [force::Bool = false])::Nothing

Specify a new set of parameter bounds for a hold variable `vref`.
An `ArgumentError` is thrown if `vref` is not a hold variable.
"""
function set_parameter_bounds(vref::GeneralVariableRef,
                              bounds::ParameterBounds{GeneralVariableRef};
                              force::Bool = false, _error::Function = error
                              )::Nothing
    return set_parameter_bounds(dispatch_variable_ref(vref), bounds,
                                force = force, _error = _error)
end

# Dispatch fallback
function add_parameter_bounds(vref::DispatchVariableRef, bounds; kwargs...)
    throw(ArgumentError("`add_parameter_bounds` not defined for variable reference type(s) " *
                        "`$(typeof(vref))`."))
end

"""
    add_parameter_bounds(vref::GeneralVariableRef,
                         bounds::ParameterBounds{GeneralVariableRef}
                         )::Nothing

Specify more parameter bounds for a hold variable `vref`.
An `ArgumentError` is thrown if `vref` is not a hold variable.
"""
function add_parameter_bounds(vref::GeneralVariableRef,
                              bounds::ParameterBounds{GeneralVariableRef};
                              _error::Function = error
                              )::Nothing
    return add_parameter_bounds(dispatch_variable_ref(vref), bounds,
                                _error = _error)
end

################################################################################
#                               MEASURE METHODS
################################################################################
# Define measure queries and their fallbacks
for op = (:measure_function, :measure_data, :is_analytic, :expand)
    @eval begin
        # define the fallback method
        func = $op
        function $op(mref::DispatchVariableRef)
            str = string("`", func, "` not defined for variable reference type " *
                         "`$(typeof(mref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            $func(mref::GeneralVariableRef)

        Define `$func` for general variable references. Errors if `mref` does
        not correspond to a `MeasureRef`. See the underlying docstrings for more
        information.
        """
        function $op(mref::GeneralVariableRef)
            return $op(dispatch_variable_ref(mref))
        end
    end
end

################################################################################
#                               DERIVATIVE METHODS
################################################################################
# Define measure queries and their fallbacks
for op = (:derivative_argument, :operator_parameter, :evaluate,
          :derivative_constraints, :delete_derivative_constraints)
    @eval begin
        # define the fallback method
        func = $op
        function $op(dref::DispatchVariableRef)
            str = string("`", func, "` not defined for variable reference type " *
                         "`$(typeof(dref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            $func(dref::GeneralVariableRef)

        Define `$func` for general variable references. Errors if `dref` does
        not correspond to a `DerivativeRef`. See the underlying docstrings for more
        information.
        """
        function $op(dref::GeneralVariableRef)
            return $op(dispatch_variable_ref(dref))
        end
    end
end

################################################################################
#                            VARIABLE INFO METHODS
################################################################################
# Define the 1 argument JuMP variable info methods
for op = (:has_lower_bound, :has_upper_bound, :is_fixed, :is_binary, :is_integer,
          :lower_bound, :upper_bound, :fix_value, :start_value,
          :set_binary, :set_integer,
          :LowerBoundRef, :UpperBoundRef, :FixRef, :BinaryRef, :IntegerRef,
          :delete_lower_bound, :delete_upper_bound, :unfix, :unset_binary,
          :unset_integer)
    @eval begin
        # define the fallback method
        func = JuMP.$op
        function JuMP.$op(vref::DispatchVariableRef)
            str = string("`JuMP.", func, "` not defined for variable reference type " *
                         "`$(typeof(vref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            JuMP.$func(vref::GeneralVariableRef)

        Define `JuMP.$func` for general variable references. It relies on `JuMP.$func`
        being defined for the underlying `DispatchVariableRef`, otherwise an
        `ArugmentError` is thrown. See the underlying docstrings for more
        information.
        """
        function JuMP.$op(vref::GeneralVariableRef)
            return JuMP.$op(dispatch_variable_ref(vref))
        end
    end
end

# Define the 2 argument setting methods (except for fix)
for op = (:set_lower_bound, :set_upper_bound, :set_start_value)
    @eval begin
        # define the fallback method
        func = JuMP.$op
        function JuMP.$op(vref::DispatchVariableRef, value)
            str = string("`JuMP.", func, "` not defined for variable reference type " *
                         "`$(typeof(vref))`.")
            throw(ArgumentError(str))
        end
        # define the dispatch version
        """
            JuMP.$func(vref::GeneralVariableRef, value::Real)::Nothing

        Define `JuMP.$func` for general variable references. It relies on `JuMP.$func`
        being defined for the underlying `DispatchVariableRef`, otherwise an
        `ArugmentError` is thrown. See the underlying docstrings for more
        information.
        """
        function JuMP.$op(vref::GeneralVariableRef, value::Real)::Nothing
            return JuMP.$op(dispatch_variable_ref(vref), value)
        end
    end
end

# Dispatch fallback for JuMP.fix
function JuMP.fix(vref::DispatchVariableRef, value::Real; force::Bool = false)
    throw(ArgumentError("`JuMP.fix` not defined for variable reference type " *
                        "`$(typeof(vref))`."))
end

"""
    JuMP.fix(vref::GeneralVariableRef, value::Real; force::Bool = false)::Nothing

Define `JuMP.fix` for general variable references. It relies on `JuMP.fix`
being defined for the underlying `DispatchVariableRef`, otherwise an
`ArugmentError` is thrown. See the underlying docstrings for more
information.
"""
function JuMP.fix(vref::GeneralVariableRef, value::Real;
                  force::Bool = false)::Nothing
    return JuMP.fix(dispatch_variable_ref(vref), value, force = force)
end


text:
using TMLEEpistasis
using ArgParse

function parse_commandline()
    s = ArgParseSettings(description="Generate queries that can be used for TMLE, see "*
                                     "https://github.com/olivierlabayle/GenesInteraction.jl/")

    @add_arg_table s begin
        "--mode", "-m"
            help = "`given` or `asb`. If `given` then the option `--query-prefix` must be"*
                   " provided. If `asb` options `--asb-prefix` and `--trans-actors must be`"*
                   " provided. Queries will be generated as MAJOR:MAJOR -> MAJOR:MINOR for each SNP."
            arg_type = String
            default = "given"
        "--outdir", "-o"
            help = "Out directory where the queries will be saved"
            arg_type = String
            default = "."
        "--query-prefix", "-q"
            help = "Prefix path to queries, must be provided if mode is: `given`"
            arg_type = String
            required = false
        "--call-threshold", "-c"
            arg_type = Float64
            help = "This is written down on the query file, it is the threshold that"*
                    " is later used to hard call genotypes based on their probabilities"
            default = 0.9
        "--exclude", "-e"
            arg_type = String
            help = "A file containing SNPs to be excluded from the analysis (one per line)"
            required = false
        "--minor-genotype-freq", "-f"
            arg_type = Float64
            help = "Minor genotype frequency filter"
            default = 0.001
            required = false
        "--asb-prefix", "-a"
            arg_type = String
            help = "Prefix path to filtered allelic-specific binding SNPS output by baal-nf "*
                   "see https://git.ecdf.ed.ac.uk/oalmelid/baal-nf"
            required = false
        "--trans-actors", "-t"
            arg_type = String
            help = "Path to trans-acting SNPS"
            required = false
        "chr-prefix"
            help = "Prefix path to BGEN chromosomes from the UKBB."
            arg_type = String
            required = true
    end

    return parse_args(s)
end

parsed_args = parse_commandline()

build_genotypes_and_queries(parsed_args)

ext: jl


text:
# misc easy to make mistakes

function lintrange(ex::Expr, ctx::LintContext)
    if length(ex.args) == 2 && typeof(ex.args[1]) <: Real &&
            typeof(ex.args[2]) <: Real && ex.args[2] < ex.args[1] &&
            ( ex.args[2] != ex.args[1]-1 || !( ex.args[1] in [0,1] ) )
        msg(ctx, :E433, "for a decreasing range, use a negative step e.g. 10:-1:1")
    else
        for a in ex.args
            lintexpr(a, ctx)
        end
    end
end

function lintvcat(ex::Expr, ctx::LintContext)
    for a in ex.args
        if isexpr(a, :vcat)
            msg(ctx, :W444, "nested vcat is treated as a 1-dimensional array")
        end
        lintexpr(a, ctx)
    end
end

function lintvect(ex::Expr, ctx::LintContext)
    for a in ex.args
        lintexpr(a, ctx)
    end
end

function linthcat(ex::Expr, ctx::LintContext)
    for a in ex.args
        if isexpr(a, :hcat)
            msg(ctx, :W445, "nested hcat is treated as a 1-row horizontal array of dim=2")
        end
        lintexpr(a, ctx)
    end
end

function linttyped_hcat(ex::Expr, ctx::LintContext)
    #dump(ex)
    if length(ex.args) == 3
        if ex.args[3] == QuoteNode(Symbol("end"))
            msg(ctx, :I681, "ambiguity of :end as a symbol vs as part of a range")
        elseif ex.args[2] == Symbol("end") && typeof(ex.args[3]) <: Integer &&
            ex.args[3] < 0
            msg(ctx, :I682, "ambiguity of `[end -n]` as a matrix row vs index [end-n]")
        end
    end
    for a in ex.args
        lintexpr(a, ctx)
    end
end

function lintcell1d(ex::Expr, ctx::LintContext)
    for a in ex.args
        lintexpr(a, ctx)
    end
end


text:
using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using StaticArrays

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DQuad
using UniformQuadMesh

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
include("../EntropyStableEuler.jl/src/logmean.jl")
using EntropyStableEuler
using EntropyStableEuler.Fluxes2D

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,rhou,E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end

function pfun_nd(rho, rhou, rhov, E)
    rhoUnorm2 = (rhou^2+rhov^2)/rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function pfun_nd(rho, rhou, E)
    rhoUnorm2 = (rhou^2)/rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,u,v,p)
    Unorm = u^2+v^2
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rho*u,rho*v,E)
end

function euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
    rhoR,uR,vR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg)
    FxS4 = (@. f4aux*uavg)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3)
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. f4aux*vavg)
    return (FxS1,FxS2,FxS3,FxS4),(FyS1,FyS2,FyS3,FyS4)
end

function euler_fluxes(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
                           rhoR,uR,vR,betaR,rhologR,betalogR)
end

const TOL = 1e-14
const Nc = 4 # number of components

"Approximation parameters"
N = 2
K1D = 8
T = 0.1

# # Initial condition 2D shocktube
# const γ = 1.4
# const rhoL = 6.0#120.0
# const uL   = 0.0
# const vL   = 0.0
# const pL   = rhoL/γ
# const rhoR = 1.2
# const uR   = 0.0
# const vR   = 0.0
# const pR   = rhoR/γ

# const mu = 0.01
# const lambda = 0.0
# const Pr = 0.73
# const cp = γ/(γ-1)
# const cv = 1/(γ-1)
# const kappa = mu*cp/Pr

# Becker viscous shocktube
const γ = 1.4
const M_0 = 3.0
const mu = 0.01
const lambda = 2/3*mu
const Pr = 3/4
const cp = γ/(γ-1)
const cv = 1/(γ-1)
const kappa = mu*cp/Pr

const v_inf = 0.2
const rho_0 = 1.0
const m_0 = 1.0
const v_0 = 1.0
const v_1 = (γ-1+2/M_0^2)/(γ+1)
const v_01 = sqrt(v_0*v_1)

const uL = v_0+v_inf
const uR = v_1+v_inf
const vL = 0.0
const vR = 0.0
const rhoL = m_0/v_0
const rhoR = m_0/v_1
const eL = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_0^2)
const eR = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_1^2)
const pL = (γ-1)*rhoL*eL
const pR = (γ-1)*rhoR*eR
const EL = pL/(γ-1)+0.5*rhoL*uL^2
const ER = pR/(γ-1)+0.5*rhoR*uR^2

# # Becker viscous shocktube 2
# const γ = 1.4
# const M_0 = 20.0
# const mu = 0.001
# const lambda = 2/3*mu
# const Pr = 3/4
# const cp = γ/(γ-1)
# const cv = 1/(γ-1)
# const kappa = mu*cp/Pr

# const v_inf = 0.2
# const rho_0 = 1.0
# const v_0 = 1.0
# const m_0 = rho_0*v_0
# const v_1 = (γ-1+2/M_0^2)/(γ+1)
# const v_01 = sqrt(v_0*v_1)

# const uL = v_0+v_inf
# const uR = v_1+v_inf
# const vL = 0.0
# const vR = 0.0
# const rhoL = m_0/v_0
# const rhoR = m_0/v_1
# const eL = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_0^2)
# const eR = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_1^2)
# const pL = (γ-1)*rhoL*eL
# const pR = (γ-1)*rhoR*eR
# const EL = pL/(γ-1)+0.5*rhoL*uL^2
# const ER = pR/(γ-1)+0.5*rhoR*uR^2

function Riemann_2D(x,y)
    if x >= 0 && y >= 0
        return .5313, 0.0, 0.0, 0.4
    elseif x < 0 && y >= 0
        return 1.0, .7276, 0.0, 1.0
    elseif x < 0 && y < 0
        return .8, 0.0, 0.0, 1.0
    else
        return 1.0, 0.0, .7276, 1.0
    end
end

"Mesh related variables"
VX, VY, EToV = uniform_quad_mesh(Int(K1D/2*3),K1D)
# @. VX = (VX+1)/2
# @. VY = (VY+1)/4
@. VX = (VX/4*3)+1/4
@. VY = (VY+1)/2
# VX, VY, EToV = uniform_quad_mesh(K1D,K1D)

rd = init_reference_quad(N,gauss_lobatto_quad(0,0,N))
"Initialize reference element"
r,_ = gauss_lobatto_quad(0,0,N)   # Reference nodes
VDM = vandermonde_1D(N,r)         # modal to nodal
Dr = grad_vandermonde_1D(N,r)/VDM # nodal differentiation
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1]) # nodal linear interpolation

Nq = N
rq,wq = gauss_lobatto_quad(0,0,Nq)
M = diagm(wq)
Mlump = zeros(size(M))
Mlump_inv = zeros(size(M))
for i = 1:Nq+1
    Mlump[i,i] = sum(M[i,:])
    Mlump_inv[i,i] = 1.0/Mlump[i,i]
end

# 1D operators
M1D = Mlump
M1D_inv = Mlump_inv
Q1D = M*Dr
B1D = zeros(N+1,N+1)
B1D[1,1] = -1
B1D[end,end] = 1
L = Array(spdiagm(0=>-2*ones(N+1), 1=>ones(N), -1=>ones(N)))
L[1,1] = -1
L[end,end] = -1
psi = pinv(L)*-1/2*B1D*ones(N+1)
S0 = zeros(N+1,N+1)
for i = 1:N+1
    for j = 1:N+1
        if L[i,j] != 0
            S0[i,j] = psi[j] - psi[i]
        end
    end
end
Q01D = S0+1/2*B1D

# Drop zeros
Q1D = Matrix(droptol!(sparse(Q1D),TOL))
Q01D = Matrix(droptol!(sparse(Q01D),TOL))
B1D = Matrix(droptol!(sparse(B1D),TOL))
S01D = Matrix(droptol!(sparse(S0),TOL))
S1D = Matrix(droptol!(sparse((Q1D-Q1D')/2),TOL))

# Tensor product operators
M = droptol!(sparse(kron(M1D,M1D)),TOL)
Minv = spdiagm(0 => 1 ./ diag(M))
Qr = droptol!(sparse(kron(M1D,Q1D)),TOL)
Qs = droptol!(sparse(kron(Q1D,M1D)),TOL)
Qr0 = droptol!(sparse(kron(M1D,Q01D)),TOL)
Qs0 = droptol!(sparse(kron(Q01D,M1D)),TOL)
Sr = droptol!(sparse((Qr-Qr')/2),TOL)
Ss = droptol!(sparse((Qs-Qs')/2),TOL)
S0r = droptol!(sparse(kron(M1D,S01D)),TOL)
S0s = droptol!(sparse(kron(S01D,M1D)),TOL)
Br = droptol!(sparse(Qr+Qr'),TOL)
Bs = droptol!(sparse(Qs+Qs'),TOL)

@unpack Vf,Dr,Ds,LIFT = rd

md = init_mesh((VX,VY),EToV,rd)
@unpack rxJ,sxJ,ryJ,syJ,sJ,J,xf,yf,mapM,mapP,mapB,nxJ,nyJ,x,y = md
xb,yb = (x->x[mapB]).((xf,yf))

# 2D shocktube
topwall  = mapB[findall(@. abs(yb-1/2) < TOL)]
wall     = mapB[findall(@. abs(yb-1/2) > TOL)]
boundary = mapB
nx_top  = nxJ[topwall]./sJ[topwall]
ny_top  = nyJ[topwall]./sJ[topwall]
nx_wall = nxJ[wall]./sJ[wall]
ny_wall = nyJ[wall]./sJ[wall]
nx_b    = nxJ[boundary]./sJ[boundary]
ny_b    = nyJ[boundary]./sJ[boundary]

# Becker's viscous shocktube
leftwall     = mapB[findall(@. abs(xb+0.5) < 1e-12)]
rightwall    = mapB[findall(@. abs(xb-1.0) < 1e-12)]
topwall      = mapB[findall(@. abs(yb-1.0) < 1e-12)]
bottomwall   = mapB[findall(@. abs(yb) < 1e-12)]
vwall = [leftwall;rightwall]

# Make domain periodic
@unpack Nfaces,Vf,wq = rd
@unpack xf,yf,K,mapM,mapP,mapB = md
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB
@pack! md = mapP

function impose_BCs_inviscid_Ub!(QP,Qf,leftwall,rightwall)
    # # No-slip at walls
    # u_1 = Qf[2][boundary]
    # u_2 = Qf[3][boundary]
    # n_1 = nx_b
    # n_2 = ny_b

    # Un = @. u_1*n_1 + u_2*n_2
    # Ut = @. u_1*n_2 - u_2*n_1

    # # ρ^+ = ρ, p^+ = p (beta^+ = beta)
    # @. QP[1][boundary] = Qf[1][boundary]
    # @. QP[4][boundary] = Qf[4][boundary]

    # # u_n^+ = -u_n, u_t^+ = u_t
    # @. QP[2][boundary] = 1/(-n_1^2-n_2^2)*(n_1*Un-n_2*Ut)
    # @. QP[3][boundary] = 1/(-n_1^2-n_2^2)*(n_2*Un+n_1*Ut)

    # QP[1][leftwall] .= rhoL
    # QP[2][leftwall] .= uL
    # QP[3][leftwall] .= vL
    # QP[4][leftwall] .= rhoL/(2*pL)
    #
    # QP[1][rightwall] .= Qf[1][rightwall]
    # QP[2][rightwall] .= Qf[2][rightwall]
    # QP[3][rightwall] .= Qf[3][rightwall]
    # QP[4][rightwall] .= Qf[4][rightwall]

    for li = leftwall
        QP[1][li] = rhoL
        QP[2][li] = uL
        QP[3][li] = vL
        QP[4][li] = rhoL/(2*pL)
    end

    for ri = rightwall
        QP[1][ri] = Qf[1][ri]
        QP[2][ri] = Qf[2][ri]
        QP[3][ri] = Qf[3][ri]
        QP[4][ri] = Qf[4][ri]
    end
end

function impose_BCs_inviscid_U!(QP,Qf,leftwall,rightwall)
    # QP[1][leftwall] .= rhoL
    # QP[2][leftwall] .= rhoL*uL
    # QP[3][leftwall] .= rhoL*vL
    # QP[4][leftwall] .= EL
    #
    # QP[1][rightwall] .= Qf[1][rightwall]
    # QP[2][rightwall] .= Qf[2][rightwall]
    # QP[3][rightwall] .= Qf[3][rightwall]
    # QP[4][rightwall] .= Qf[4][rightwall]

    for li = leftwall
        QP[1][li] = rhoL
        QP[2][li] = rhoL*uL
        QP[3][li] = rhoL*vL
        QP[4][li] = EL
    end

    for ri = rightwall
        QP[1][ri] = Qf[1][ri]
        QP[2][ri] = Qf[2][ri]
        QP[3][ri] = Qf[3][ri]
        QP[4][ri] = Qf[4][ri]
    end
end

function impose_BCs_flux!(flux_x_P,flux_y_P,leftwall,rightwall)
    # flux_x_P[1][leftwall] .= rhoL*uL
    # flux_x_P[2][leftwall] .= rhoL*uL^2+pL
    # flux_x_P[3][leftwall] .= rhoL*uL*vL
    # flux_x_P[4][leftwall] .= uL*(EL+pL)
    #
    # flux_y_P[1][leftwall] .= rhoL*vL
    # flux_y_P[2][leftwall] .= rhoL*uL*vL
    # flux_y_P[3][leftwall] .= rhoL*vL^2+pL
    # flux_y_P[4][leftwall] .= vL*(EL+pL)
    #
    # flux_x_P[1][rightwall] .= rhoR*uR
    # flux_x_P[2][rightwall] .= rhoR*uR^2+pR
    # flux_x_P[3][rightwall] .= rhoR*uR*vR
    # flux_x_P[4][rightwall] .= uR*(ER+pR)
    #
    # flux_y_P[1][rightwall] .= rhoR*vR
    # flux_y_P[2][rightwall] .= rhoR*uR*vR
    # flux_y_P[3][rightwall] .= rhoR*vR^2+pR
    # flux_y_P[4][rightwall] .= vR*(ER+pR)

    for li = leftwall
        flux_x_P[1][li] = rhoL*uL
        flux_x_P[2][li] = rhoL*uL^2+pL
        flux_x_P[3][li] = rhoL*uL*vL
        flux_x_P[4][li] = uL*(EL+pL)

        flux_y_P[1][li] = rhoL*vL
        flux_y_P[2][li] = rhoL*uL*vL
        flux_y_P[3][li] = rhoL*vL^2+pL
        flux_y_P[4][li] = vL*(EL+pL)
    end

    for ri = rightwall
        flux_x_P[1][ri] = rhoR*uR
        flux_x_P[2][ri] = rhoR*uR^2+pR
        flux_x_P[3][ri] = rhoR*uR*vR
        flux_x_P[4][ri] = uR*(ER+pR)

        flux_y_P[1][ri] = rhoR*vR
        flux_y_P[2][ri] = rhoR*uR*vR
        flux_y_P[3][ri] = rhoR*vR^2+pR
        flux_y_P[4][ri] = vR*(ER+pR)
    end
end

function impose_BCs_lam!(lamP,lam,leftwall,rightwall)
    # lam[leftwall]   .= 0.0
    # lam[rightwall]  .= 0.0
    # lamP[leftwall]  .= 0.0
    # lamP[rightwall] .= 0.0

    for li = leftwall
        lam[li]  = 0.0
        lamP[li] = 0.0
    end

    for ri = leftwall
        lam[ri]  = 0.0
        lamP[ri] = 0.0
    end
end

function impose_BCs_entropyvars!(VUP,VUf,leftwall,rightwall)
    # # adiabatic on wall
    # @. VUP[2][wall] = -VUf[2][wall]
    # @. VUP[3][wall] = -VUf[3][wall]
    # @. VUP[4][wall] =  VUf[4][wall]

    # # Reflective on top
    # v_1 = VUf[2][topwall]
    # v_2 = VUf[3][topwall]
    # n_1 = nx_top
    # n_2 = ny_top

    # VUn = @. v_1*n_1 + v_2*n_2
    # VUt = @. v_1*n_2 - v_2*n_1

    # # v_4^+ = v_4
    # @. VUP[4][topwall] = VUf[4][topwall]
    # # v_n^+ = -v_n, v_t^+ = v_t
    # @. VUP[2][topwall] = VUf[2][topwall] - 2*VUn*n_1
    # @. VUP[3][topwall] = VUf[3][topwall] - 2*VUn*n_2

    VL = v_ufun(rhoL,rhoL*uL,rhoL*vL,EL)
    VR = v_ufun(rhoR,rhoR*uR,rhoR*vR,ER)

    # VUP[1][leftwall] .= VL[1]
    # VUP[2][leftwall] .= VL[2]
    # VUP[3][leftwall] .= VL[3]
    # VUP[4][leftwall] .= VL[4]
    #
    # VUP[1][rightwall] .= VUf[1][rightwall]
    # VUP[2][rightwall] .= VUf[2][rightwall]
    # VUP[3][rightwall] .= VUf[3][rightwall]
    # VUP[4][rightwall] .= VUf[4][rightwall]

    for li = leftwall
        VUP[1][li] = VL[1]
        VUP[2][li] = VL[2]
        VUP[3][li] = VL[3]
        VUP[4][li] = VL[4]
    end

    for ri = rightwall
        VUP[1][ri] = VUf[1][ri]
        VUP[2][ri] = VUf[2][ri]
        VUP[3][ri] = VUf[3][ri]
        VUP[4][ri] = VUf[4][ri]
    end
end

function impose_BCs_stress!(σxP,σyP,σxf,σyf,vwall)
    # # Adiabatic no-slip BC
    # @. σxP[2][wall] = σxf[2][wall]
    # @. σyP[2][wall] = σyf[2][wall]
    # @. σxP[3][wall] = σxf[3][wall]
    # @. σyP[3][wall] = σyf[3][wall]
    # @. σxP[4][wall] = -σxf[4][wall]
    # @. σyP[4][wall] = -σyf[4][wall]

    # # Reflective on top
    # sigma_x_1 = σxf[2][topwall]
    # sigma_x_2 = σxf[3][topwall]
    # sigma_y_1 = σyf[2][topwall]
    # sigma_y_2 = σyf[3][topwall]
    # n_1 = nx_top
    # n_2 = ny_top

    # σn_x = @. sigma_x_1*n_1 + sigma_x_2*n_2
    # σn_y = @. sigma_y_1*n_1 + sigma_y_2*n_2

    # @. σxP[2][topwall] = -σxf[2][topwall]+2*n_1*σn_x
    # @. σyP[2][topwall] = -σyf[2][topwall]+2*n_1*σn_y
    # @. σxP[3][topwall] = -σxf[3][topwall]+2*n_2*σn_x
    # @. σyP[3][topwall] = -σyf[3][topwall]+2*n_2*σn_y
    # @. σxP[4][topwall] = -σxf[4][topwall]
    # @. σyP[4][topwall] = -σyf[4][topwall]

    @. σxP[1][vwall] = σxf[1][vwall]
    @. σyP[1][vwall] = σyf[1][vwall]
    @. σxP[2][vwall] = σxf[2][vwall]
    @. σyP[2][vwall] = σyf[2][vwall]
    @. σxP[3][vwall] = σxf[3][vwall]
    @. σyP[3][vwall] = σyf[3][vwall]
    @. σxP[4][vwall] = σxf[4][vwall]
    @. σyP[4][vwall] = σyf[4][vwall]

    for vi = vwall
        σxP[1][vi] = σxf[1][vi]
        σyP[1][vi] = σyf[1][vi]
        σxP[2][vi] = σxf[2][vi]
        σyP[2][vi] = σyf[2][vi]
        σxP[3][vi] = σxf[3][vi]
        σyP[3][vi] = σyf[3][vi]
        σxP[4][vi] = σxf[4][vi]
        σyP[4][vi] = σyf[4][vi]
    end
end

function init_visc_fxn(λ,μ,Pr)
    let λ=-λ,μ=μ,Pr=Pr
        function viscous_matrices!(Kxx,Kxy,Kyy,v)
            v1,v2,v3,v4 = v
            inv_v4_cubed = @. 1/(v4^3)
            λ2μ = (λ+2.0*μ)
            Kxx[2,2] = inv_v4_cubed*-λ2μ*v4^2
            Kxx[2,4] = inv_v4_cubed*λ2μ*v2*v4
            Kxx[3,3] = inv_v4_cubed*-μ*v4^2
            Kxx[3,4] = inv_v4_cubed*μ*v3*v4
            Kxx[4,2] = inv_v4_cubed*λ2μ*v2*v4
            Kxx[4,3] = inv_v4_cubed*μ*v3*v4
            Kxx[4,4] = inv_v4_cubed*-(λ2μ*v2^2 + μ*v3^2 - γ*μ*v4/Pr)

            Kxy[2,3] = inv_v4_cubed*-λ*v4^2
            Kxy[2,4] = inv_v4_cubed*λ*v3*v4
            Kxy[3,2] = inv_v4_cubed*-μ*v4^2
            Kxy[3,4] = inv_v4_cubed*μ*v2*v4
            Kxy[4,2] = inv_v4_cubed*μ*v3*v4
            Kxy[4,3] = inv_v4_cubed*λ*v2*v4
            Kxy[4,4] = inv_v4_cubed*(λ+μ)*(-v2*v3)

            Kyy[2,2] = inv_v4_cubed*-μ*v4^2
            Kyy[2,4] = inv_v4_cubed*μ*v2*v4
            Kyy[3,3] = inv_v4_cubed*-λ2μ*v4^2
            Kyy[3,4] = inv_v4_cubed*λ2μ*v3*v4
            Kyy[4,2] = inv_v4_cubed*μ*v2*v4
            Kyy[4,3] = inv_v4_cubed*λ2μ*v3*v4
            Kyy[4,4] = inv_v4_cubed*-(λ2μ*v3^2 + μ*v2^2 - γ*μ*v4/Pr)
        end
        return viscous_matrices!
    end
end

viscous_matrices! = init_visc_fxn(lambda,mu,Pr)

# # Initial condition 2D shocktube
# rho_init(x) = (x < 0.5) ? rhoL : rhoR
# p_init(x)   = (x < 0.5) ? pL   : pR

# rho = @. rho_init(x)
# u   = zeros(size(x))
# v   = zeros(size(x))
# p   = @. p_init(x)

# U = primitive_to_conservative_hardcode.(rho,u,v,p)
# rho = [x[1] for x in U]
# rhou = [x[2] for x in U]
# rhov = [x[3] for x in U]
# E = [x[4] for x in U]
# U = (rho,rhou,rhov,E)

# Initial condition Becker's viscous shocktube
function bisection_solve_velocity(x,max_iter,tol)
    v_L = v_1
    v_R = v_0
    num_iter = 0

    L_k = kappa/m_0/cv
    f(v) = -x+2*L_k/(γ+1)*(v_0/(v_0-v_1)*log((v_0-v)/(v_0-v_01))-v_1/(v_0-v_1)*log((v-v_1)/(v_01-v_1)))

    v_new = (v_L+v_R)/2
    while num_iter < max_iter
        v_new = (v_L+v_R)/2

        if abs(f(v_new)) < tol
            return v_new
        elseif sign(f(v_L)) == sign(f(v_new))
            v_L = v_new
        else
            v_R = v_new
        end
        num_iter += 1
    end

    return v_new
end

const max_iter = 100
const tol = 1e-14

function exact_sol_viscous_shocktube(x,t)
    u   = bisection_solve_velocity(x-v_inf*t,max_iter,tol)
    rho = m_0/u
    e   = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-u^2)
    return rho, rho*(v_inf+u), 0.0, rho*(e+1/2*(v_inf+u)^2)
end

U = exact_sol_viscous_shocktube.(x,0.0)
U = ([x[1] for x in U], [x[2] for x in U], [x[3] for x in U], [x[4] for x in U])
# U = @. Riemann_2D(x,y)
# rho,u,v,p = [(x->x[i]).(U) for i = 1:Nc]
# U = @. primitive_to_conservative_hardcode(rho,u,v,p)
# rho,rhou,rhov,E = [(x->x[i]).(U) for i = 1:Nc]
# U = (rho,rhou,rhov,E)

Np = (N+1)*(N+1)
face_idx = [1:N+1; (N+1):(N+1):Np; Np:-1:Np-N; Np-N:-(N+1):1]
x_idx    = [(N+2):(2*N+2); (3*N+4):(4*N+4)]
y_idx    = [1:(N+1); (2*N+3):(3*N+3)]
S0r1 = sum(S0r,dims=2)
S0s1 = sum(S0s,dims=2)
K      = size(U[1],2)
Np     = (N+1)*(N+1)
Nfaces = 4
Nfp    = Nfaces*(N+1)

# Preallocation
p      = zero(U[1])
VU     = zero.(U)
flux_x = zero.(U)
flux_y = zero.(U)
lam    = zeros(Nfp,K)
LFc    = zeros(Nfp,K)
Ub     = zero.(U)
sigma_x = zero.(U)
sigma_y = zero.(U)

function zhang_wavespd(rhoL,rhouL,rhovL,EL,sigma2L,sigma3L,sigma4L,pL)
    uL   = rhouL/rhoL
    eL   = (EL - .5*rhoL*uL^2)/rhoL
    tauL = sigma2L
    qL   = uL*tauL-sigma3L
    wavespdL = abs(uL)+1/(2*rhoL^2*eL)*(sqrt(rhoL^2*qL^2+2*rhoL^2*eL*abs(tauL-pL)^2)+rhoL*abs(qL))
    return wavespdL
end

function limiting_param(U_low, P_ij)
    l = 1.0
    # Limit density
    if U_low[1] + P_ij[1] < -TOL
        l = min(abs(U_low[1])/(abs(P_ij[1])+1e-14), 1.0)
    end

    # limiting internal energy (via quadratic function)
    a = P_ij[1]*P_ij[4]-1.0/2.0*(P_ij[2]^2+P_ij[3]^2)
    b = U_low[4]*P_ij[1]+U_low[1]*P_ij[4]-U_low[2]*P_ij[2]-U_low[3]*P_ij[3]
    c = U_low[4]*U_low[1]-1.0/2.0*(U_low[2]^2+U_low[3]^2)

    l_eps_ij = 1.0
    if b^2-4*a*c >= 0
        r1 = (-b+sqrt(b^2-4*a*c))/(2*a)
        r2 = (-b-sqrt(b^2-4*a*c))/(2*a)
        if r1 > TOL && r2 > TOL
            l_eps_ij = min(r1,r2)
        elseif r1 > TOL && r2 < -TOL
            l_eps_ij = r1
        elseif r2 > TOL && r1 < -TOL
            l_eps_ij = r2
        end
    end

    l = min(l,l_eps_ij)
    return l
end

function rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y)
    # TODO: hardcoded!
    K      = size(U[1],2)
    Np     = (N+1)*(N+1)
    Nfaces = 4
    Nfp    = Nfaces*(N+1)

    # J = (1/K1D)^2 # hardcoded!!
    # Jf = 1/K1D
    # rxJ = 1/K1D
    # syJ = 1/K1D
    # sJ = 1/K1D

    J   = (1/K1D/2)^2
    Jf  = 1/K1D/2
    rxJ = 1/K1D/2
    syJ = 1/K1D/2
    sJ  = 1/K1D/2

    p = pfun_nd.(U[1],U[2],U[3],U[4])
    #@. p = (γ-1)*(U[4]-.5*(U[2]^2+U[3]^2)/U[1])
    @. flux_x[1] = U[2]
    @. flux_x[2] = U[2]^2/U[1]+p
    @. flux_x[3] = U[2]*U[3]/U[1]
    @. flux_x[4] = U[4]*U[2]/U[1]+p*U[2]/U[1]
    @. flux_y[1] = U[3]
    @. flux_y[2] = U[2]*U[3]/U[1]
    @. flux_y[3] = U[3]^2/U[1]+p
    @. flux_y[4] = U[4]*U[3]/U[1]+p*U[3]/U[1]

    Uf = (x->Vf*x).(U)
    UP = (x->x[mapP]).(Uf)
    impose_BCs_inviscid_U!(UP,Uf,leftwall,rightwall)

    # simple lax friedrichs dissipation
    (rhoM,rhouM,rhovM,EM) = Uf
    rhoUM_n = @. (rhouM*nxJ + rhovM*nyJ)/sJ
    @. lam  = abs(sqrt(abs(rhoUM_n/rhoM))+sqrt(γ*(γ-1)*(EM-.5*rhoUM_n^2/rhoM)/rhoM))
    lamP = lam[mapP]
    impose_BCs_lam!(lamP,lam,leftwall,rightwall)
    @. LFc = max(lam,lamP)*sJ

    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[3]/U[1]
    @. Ub[4] = U[1]/(2*p)
    Ubf = (x->Vf*x).(Ub)
    UbP = (x->x[mapP]).(Ubf)
    impose_BCs_inviscid_Ub!(UbP,Ubf,leftwall,rightwall)

    flux_x_f = (x->Vf*x).(flux_x)
    flux_x_P = (x->x[mapP]).(flux_x_f)
    flux_y_f = (x->Vf*x).(flux_y)
    flux_y_P = (x->x[mapP]).(flux_y_f)
    impose_BCs_flux!(flux_x_P,flux_y_P,leftwall,rightwall)

    U_low    = [zeros(Float64,Np) for i = 1:Nc]
    U_low_i  = zeros(Float64,Nc)
    F_low    = [zeros(Float64,Np,Np) for i = 1:Nc]
    F_high   = [zeros(Float64,Np,Np) for i = 1:Nc]
    F_P      = [zeros(Float64,Nfp) for i = 1:Nc]
    rhsU  = [zeros(Float64,Np,K) for i = 1:Nc]
    L     = zeros(Float64,Np,Np)
    P_ij  = zeros(Float64,Nc)

    # =======================
    # Calculate viscous part
    # =======================
    # VU  = v_ufun(U...)
    # VUf = (x->Vf*x).(VU)
    # VUP = (x->x[mapP]).(VUf)
    vector_norm(U) = sum((x->x.^2).(U))
    @. VU[3] = U[2]^2+U[3]^2 # rhoUnorm
    @. VU[4] = U[4]-.5*VU[3]/U[1] # rhoe
    @. VU[1] = log((γ-1)*VU[4]/(U[1]^γ)) # sU
    @. VU[1] = (-U[4]+VU[4]*(γ+1-VU[1]))/VU[4]
    @. VU[2] = U[2]/VU[4]
    @. VU[3] = U[3]/VU[4]
    @. VU[4] = -U[1]/VU[4]
    VUf = (x->Vf*x).(VU)
    VUP = (x->x[mapP]).(VUf)
    impose_BCs_entropyvars!(VUP,VUf,leftwall,rightwall)

    surfx(uP,uf) = LIFT*(@. .5*(uP-uf)*nxJ)
    surfy(uP,uf) = LIFT*(@. .5*(uP-uf)*nyJ)
    VUx = rxJ.*(x->Dr*x).(VU) .+ surfx.(VUP,VUf)
    VUy = syJ.*(x->Ds*x).(VU) .+ surfy.(VUP,VUf)
    # VUx = (x->x./J).(VUx)
    # VUy = (x->x./J).(VUy)
    for c = 1:Nc
        @. VUx[c] = VUx[c]/J
        @. VUy[c] = VUy[c]/J
    end

    # initialize sigma_x,sigma_y = viscous rhs
    Kxx,Kxy,Kyy = ntuple(x->MMatrix{4,4}(zeros(Nc,Nc)),3)
    sigma_x_e = zero.(getindex.(VU,:,1))
    sigma_y_e = zero.(getindex.(VU,:,1))
    for e = 1:K
        fill!.(sigma_x_e,0.0)
        fill!.(sigma_y_e,0.0)

        # mult by matrices and perform local projections
        for i = 1:Np
            vxi = getindex.(VUx,i,e)
            vyi = getindex.(VUy,i,e)
            viscous_matrices!(Kxx,Kxy,Kyy,getindex.(VU,i,e))

            for col = 1:Nc
                vxi_col = vxi[col]
                vyi_col = vyi[col]
                for row = 1:Nc
                    sigma_x_e[row][i] += Kxx[row,col]*vxi_col + Kxy[row,col]*vyi_col
                    sigma_y_e[row][i] += Kxy[col,row]*vxi_col + Kyy[row,col]*vyi_col
                end
            end
        end
        setindex!.(sigma_x,sigma_x_e,:,e)
        setindex!.(sigma_y,sigma_y_e,:,e)
    end

    sxf = (x->Vf*x).(sigma_x)
    syf = (x->Vf*x).(sigma_y)
    sxP = (x->x[mapP]).(sxf)
    syP = (x->x[mapP]).(syf)
    impose_BCs_stress!(sxP,syP,sxf,syf,vwall)

    # =====================
    # Loop through elements
    # =====================
    for k = 1:K
        for c = 1:Nc
            F_low[c]  .= 0.0
            F_high[c] .= 0.0
            F_P[c]    .= 0.0
            U_low[c]  .= 0.0
        end
        L .= 0.0

        # Calculate low order algebraic flux
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(rxJ^2*S0r[i,j]^2+syJ^2*S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij_x = rxJ*S0r[i,j]/c_ij_norm
                    n_ij_y = syJ*S0s[i,j]/c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],n_ij_x*U[2][i,k]+n_ij_y*U[3][i,k],U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],n_ij_x*U[2][j,k]+n_ij_y*U[3][j,k],U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm

                    for c = 1:Nc
                        F_low[c][i,j] = (rxJ*S0r[i,j]*(flux_x[c][i,k]+flux_x[c][j,k])
                                        +syJ*S0s[i,j]*(flux_y[c][i,k]+flux_y[c][j,k])
                                        -rxJ*S0r[i,j]*(sigma_x[c][i,k]+sigma_x[c][j,k])
                                        -syJ*S0s[i,j]*(sigma_y[c][i,k]+sigma_y[c][j,k])
                                        -d_ij*(U[c][j,k]-U[c][i,k]))
                    end
                end
            end
        end

        # Calculate high order algebraic flux
        for i = 1:Np-1
            for j = i+1:Np
                # TODO: can preallocate nonzero entries
                if Sr[i,j] != 0.0 || Ss[i,j] != 0.0
                    F1,F2 = euler_fluxes(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[4][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],Ub[4][j,k])
                    for c = 1:Nc
                        val = (2*rxJ*Sr[i,j]*F1[c]+2*syJ*Ss[i,j]*F2[c]
                                -rxJ*S0r[i,j]*(sigma_x[c][i,k]+sigma_x[c][j,k])
                                -syJ*S0s[i,j]*(sigma_y[c][i,k]+sigma_y[c][j,k]))
                        F_high[c][i,j] = val
                        F_high[c][j,i] = -val
                    end
                end
            end
        end

        # Calculate interface fluxes
        for i = 1:Nfp
            S0r_ij = -S0r1[face_idx[i]]
            S0s_ij = -S0s1[face_idx[i]]

            # flux in x direction
            if i in x_idx
                wavespd_M = wavespeed_1D(Uf[1][i,k],Uf[2][i,k],Uf[4][i,k])
                wavespd_P = wavespeed_1D(UP[1][i,k],UP[2][i,k],UP[4][i,k])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0r_ij)
                for c = 1:Nc
                    F_P[c][i] = (Jf*S0r_ij*(flux_x_f[c][i,k]+flux_x_P[c][i,k])
                                -Jf*S0r_ij*(sxf[c][i,k]+sxP[c][i,k])
                                -LFc[i,k]*abs(S0r_ij)*(UP[c][i,k]-Uf[c][i,k]))
                end
            end

            # flux in y direction
            if i in y_idx
                wavespd_M = wavespeed_1D(Uf[1][i,k],Uf[3][i,k],Uf[4][i,k])
                wavespd_P = wavespeed_1D(UP[1][i,k],UP[3][i,k],UP[4][i,k])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0s_ij)
                for c = 1:Nc
                    F_P[c][i] = (Jf*S0s_ij*(flux_y_f[c][i,k]+flux_y_P[c][i,k])
                                -Jf*S0s_ij*(syf[c][i,k]+syP[c][i,k])
                                -LFc[i,k]*abs(S0s_ij)*(UP[c][i,k]-Uf[c][i,k]))
                end
            end
        end

        # Calculate low order solution
        for i = 1:Np
            for j = 1:Np
                for c = 1:Nc
                    U_low[c][i] += F_low[c][i,j]
                end
            end
        end
        for i = 1:Nfp
            for c = 1:Nc
                U_low[c][face_idx[i]] += F_P[c][i]
            end
        end
        for c = 1:Nc
            for i = 1:Np
                U_low[c][i] = U[c][i,k]-dt*Minv[i,i]/J*U_low[c][i]
            end
        end

        # Calculate limiting parameters
        for i = 1:Np
            m_i = J/Minv[i,i]
            lambda_j = 1/(Np-1)
            for c = 1:Nc
                U_low_i[c] = U_low[c][i]
            end
            for j = 1:Np
                if i != j
                    for c = 1:Nc
                        P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                    L[i,j] = limiting_param(U_low_i, P_ij)
                end
            end
        end

        # Symmetrize limiting parameters
        for i = 1:Np
            for j = i+1:Np
                l = min(L[i,j],L[j,i])
                L[i,j] = l
                L[j,i] = l
            end
        end

        # L = ones(size(L))

        for c = 1:Nc
            for i = 1:Np
                for j = 1:Np
                    rhsU[c][i,k] += (L[i,j]-1)*F_low[c][i,j]-L[i,j]*F_high[c][i,j]
                end
            end
            for i = 1:Nfp
                rhsU[c][face_idx[i],k] -= F_P[c][i]
            end
        end
    end

    for c = 1:Nc
        rhsU[c] .= 1/J*Minv*rhsU[c]
    end

    return rhsU
end

# Time stepping
"Time integration"
t = 0.0
U = collect(U)
resU = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]
resW = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]
resZ = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]

gr(aspect_ratio=:equal,legend=false,
   markerstrokewidth=0,markersize=2,axis=nothing)
const GIFINTERVAL = 100
#plot()

# dt = 1e-4
# rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
# @btime rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
# @profiler rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);

dt_hist = []
anim = Animation()
i = 1

while t < T
    # SSPRK(3,3)
    dt = min(1e-4,T-t)
    rhsU = rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
    @. resW = U + dt*rhsU
    rhsU = rhs_IDP!(resW,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
    @. resZ = resW+dt*rhsU
    @. resW = 3/4*U+1/4*resZ
    rhsU = rhs_IDP!(resW,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,leftwall,rightwall,vwall,dt,p,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
    @. resZ = resW+dt*rhsU
    @. U = 1/3*U+2/3*resZ

    push!(dt_hist,dt)
    global t = t + dt
    println("Current time $t with time step size $dt, and final time $T, at step $i")
    global i = i + 1
    # if mod(i,50) == 1
    #     scatter(Vp*x,Vp*y,Vp*U[1],zcolor=Vp*U[1],camera=(0,90))
    #     frame(anim)
    # end
end

################
### Plotting ###
################

#plotting nodes
@unpack VDM = rd
rp,sp = equi_nodes_2D(10)
Vp = vandermonde_2D(N,rp,sp)/VDM
gr(aspect_ratio=:equal,legend=false,
   markerstrokewidth=0,markersize=2)
#vv = Vp*(@. (Q[2]/Q[1])^2 + (Q[3]/Q[1])^2)

exact_U = @. exact_sol_viscous_shocktube.(x,T)
exact_rho = [x[1] for x in exact_U]
exact_rhou = [x[2] for x in exact_U]
exact_rhov = [x[3] for x in exact_U]
exact_E = [x[4] for x in exact_U]

rho = U[1]
rhou = U[2]
rhov = U[3]
E = U[4]

exact_rho_p = Vp*exact_rho

Linferr = maximum(abs.(exact_rho-rho))/maximum(abs.(rho)) +
          maximum(abs.(exact_rhou-rhou))/maximum(abs.(rhou)) +
          #maximum(abs.(exact_rhov-rhov))/maximum(abs.(rhov).+1e-14) +
          maximum(abs.(exact_E-E))/maximum(abs.(E))

J = J[1] # TODO: assume uniform mesh
L1err = sum(J*wq.*abs.(exact_rho-rho))/sum(J*wq.*abs.(rho)) +
        sum(J*wq.*abs.(exact_rhou-rhou))/sum(J*wq.*abs.(rhou)) +
        #sum(J.*abs.(exact_rhov-rhov))/sum(J.*abs.(rhov).+1e-14) +
        sum(J*wq.*abs.(exact_E-E))/sum(J*wq.*abs.(E))
println("N = $N, K = $K")
println("L1 error is $L1err")
println("Linf error is $Linferr")

scatter(x,y,U[1],zcolor=U[1],camera=(0,90),colorbar=:right)
xp = Vp*x
yp = Vp*y
vv = Vp*U[1]
scatter(xp,yp,vv,zcolor=vv,camera=(0,90),colorbar=:right)
#scatter(xp,yp,exact_rho_p,zcolor=exact_rho_p,camera=(0,90),colorbar=:right)

max_stars_repo_name: yiminllin/ESDG-PosLimit

max_issues_repo_name: yiminllin/ESDG-PosLimit

max_forks_repo_name: yiminllin/ESDG-PosLimit


text:
using ExoplanetsSysSim
include("christiansen_func.jl")

global sim_param_closure = setup_sim_param_christiansen()

#p_lim_arr_num = [0.5, 1.25, 2.5, 5., 10., 20., 40., 80., 160., 320.]
#r_lim_arr_num = [1., 1.25, 1.5, 1.75, 2.]

p_lim_arr_num = [10., 20.]
r_lim_arr_num = [1., 1.25]

add_param_fixed(sim_param_closure, "p_lim_arr", p_lim_arr_num)
add_param_fixed(sim_param_closure, "r_lim_arr", r_lim_arr_num*ExoplanetsSysSim.earth_radius)

@time inv_det_prob(sim_param_closure)


text:
"""
KernelFunctions. [Github](https://github.com/JuliaGaussianProcesses/KernelFunctions.jl)
[Documentation](https://juliagaussianprocesses.github.io/KernelFunctions.jl/stable/)
"""
module KernelFunctions

if !isfile(joinpath(@__DIR__, "update_v0.8.0"))
    printstyled(
        stdout,
        """
        WARNING: SqExponentialKernel changed convention in version 0.8.0.
        This kernel now divides the squared distance by 2 to align with standard practice.
        This warning will be removed in 0.9.0.
        """;
        color = Base.info_color(),
    )
    touch(joinpath(@__DIR__, "update_v0.8.0"))
end

export kernelmatrix, kernelmatrix!, kerneldiagmatrix, kerneldiagmatrix!
export transform
export duplicate, set! # Helpers

export Kernel, MOKernel
export ConstantKernel, WhiteKernel, EyeKernel, ZeroKernel, WienerKernel
export CosineKernel
export SqExponentialKernel, RBFKernel, GaussianKernel, SEKernel
export LaplacianKernel, ExponentialKernel, GammaExponentialKernel
export ExponentiatedKernel
export FBMKernel
export MaternKernel, Matern32Kernel, Matern52Kernel
export LinearKernel, PolynomialKernel
export RationalQuadraticKernel, GammaRationalQuadraticKernel
export MahalanobisKernel, GaborKernel, PiecewisePolynomialKernel
export PeriodicKernel, NeuralNetworkKernel
export KernelSum, KernelProduct
export TransformedKernel, ScaledKernel
export TensorProduct

export Transform, SelectTransform, ChainTransform, ScaleTransform, LinearTransform,
    ARDTransform, IdentityTransform, FunctionTransform, PeriodicTransform

export NystromFact, nystrom

export spectral_mixture_kernel, spectral_mixture_product_kernel

export ColVecs, RowVecs

export MOInput
export IndependentMOKernel, LatentFactorMOKernel

using Compat
using Requires
using Distances, LinearAlgebra
using Functors
using SpecialFunctions: loggamma, besselk, polygamma
using ZygoteRules: @adjoint, pullback
using StatsFuns: logtwo
using InteractiveUtils: subtypes
using StatsBase

abstract type Kernel end
abstract type SimpleKernel <: Kernel end

include("utils.jl")
include(joinpath("distances", "pairwise.jl"))
include(joinpath("distances", "dotproduct.jl"))
include(joinpath("distances", "delta.jl"))
include(joinpath("distances", "sinus.jl"))

include(joinpath("transform", "transform.jl"))
include(joinpath("transform", "scaletransform.jl"))
include(joinpath("transform", "ardtransform.jl"))
include(joinpath("transform", "lineartransform.jl"))
include(joinpath("transform", "functiontransform.jl"))
include(joinpath("transform", "selecttransform.jl"))
include(joinpath("transform", "chaintransform.jl"))
include(joinpath("transform", "periodic_transform.jl"))

include(joinpath("basekernels", "constant.jl"))
include(joinpath("basekernels", "cosine.jl"))
include(joinpath("basekernels", "exponential.jl"))
include(joinpath("basekernels", "exponentiated.jl"))
include(joinpath("basekernels", "fbm.jl"))
include(joinpath("basekernels", "gabor.jl"))
include(joinpath("basekernels", "maha.jl"))
include(joinpath("basekernels", "matern.jl"))
include(joinpath("basekernels", "nn.jl"))
include(joinpath("basekernels", "periodic.jl"))
include(joinpath("basekernels", "piecewisepolynomial.jl"))
include(joinpath("basekernels", "polynomial.jl"))
include(joinpath("basekernels", "rationalquad.jl"))
include(joinpath("basekernels", "sm.jl"))
include(joinpath("basekernels", "wiener.jl"))

include(joinpath("kernels", "transformedkernel.jl"))
include(joinpath("kernels", "scaledkernel.jl"))
include(joinpath("matrix", "kernelmatrix.jl"))
include(joinpath("kernels", "kernelsum.jl"))
include(joinpath("kernels", "kernelproduct.jl"))
include(joinpath("kernels", "tensorproduct.jl"))
include(joinpath("approximations", "nystrom.jl"))
include("generic.jl")

include(joinpath("mokernels", "mokernel.jl"))
include(joinpath("mokernels", "moinput.jl"))
include(joinpath("mokernels", "independent.jl"))
include(joinpath("mokernels", "slfm.jl"))

include("zygote_adjoints.jl")

include("test_utils.jl")

function __init__()
    @require Kronecker="2c470bb0-bcc8-11e8-3dad-c9649493f05e" begin
        include(joinpath("matrix", "kernelkroneckermat.jl"))
    end
    @require PDMats="90014a1f-27ba-587c-ab20-58faa44d9150" begin
        include(joinpath("matrix", "kernelpdmat.jl"))
    end
end

end


text:
include("../common.jl")
build_julia(v"1.5.3")

max_stars_repo_name: Affie/Yggdrasil

max_issues_repo_name: Affie/Yggdrasil

max_forks_repo_name: Affie/Yggdrasil


text:
using Plot3D: read_multi_block, Block, find_block_connectivity
using JLD

## Cell 1 Reads the block data
if !isfile("blocks.jld")
    blocks = read_multi_block("../../../../testfiles/8block.p3d")
    jldopen("blocks.jld", "w") do file
        write(file, "blocks", blocks)  # alternatively, say "@write file A"
    end
end
##
blocks = jldopen("blocks.jld", "r") do file
    read(file, "blocks")
end

## Cell 2 Read the Block data
if !isfile("block_matches.jld")
    print("Block read, Finding matching faces\n")
    block_match_indices, block_match_corners = find_block_connectivity(blocks[1],blocks[2])
    jldopen("block_matches.jld", "w") do file
        write(file, "block_match_indices", block_match_indices)  # alternatively, say "@write file A"
        write(file, "block_match_corners", block_match_corners)  # alternatively, say "@write file A"
    end
end

block_matches = jldopen("block_matches.jld", "r") do file
    read(file, "block_match_indices")
end
print("check")
##


text:
@testset "tensorproduct" begin
    rng = MersenneTwister(123456)
    u1 = rand(rng, 10)
    u2 = rand(rng, 10)
    v1 = rand(rng, 5)
    v2 = rand(rng, 5)

    # kernels
    k1 = SqExponentialKernel()
    k2 = ExponentialKernel()
    kernel1 = TensorProduct(k1, k2)
    kernel2 = TensorProduct([k1, k2])

    @test kernel1.kernels === (k1, k2) === TensorProduct((k1, k2)).kernels
    @test length(kernel1) == length(kernel2) == 2
    @test_throws DimensionMismatch kernel1(rand(3), rand(3))

    @testset "val" begin
        for (x, y) in (((v1, u1), (v2, u2)), ([v1, u1], [v2, u2]))
            val = k1(x[1], y[1]) * k2(x[2], y[2])

            @test kernel1(x, y) == kernel2(x, y) == val
        end
    end

    @testset "kernelmatrix and kerneldiagmatrix" begin
        X = rand(rng, 2, 10)
        x_cols = ColVecs(X)
        x_rows = RowVecs(X')
        Y = rand(rng, 2, 10)
        y_cols = ColVecs(Y)
        y_rows = RowVecs(Y')

        trueX = kernelmatrix(k1, X[1, :]) .* kernelmatrix(k2, X[2, :])
        trueXY = kernelmatrix(k1, X[1, :], Y[1, :]) .* kernelmatrix(k2, X[2, :], Y[2, :])
        tmp = Matrix{Float64}(undef, 10, 10)
        tmp_diag = Vector{Float64}(undef, 10)

        for kernel in (kernel1, kernel2), (x, y) in ((x_cols, y_cols), (x_rows, y_rows))
            @test kernelmatrix(kernel, x) ≈ trueX

            @test kernelmatrix(kernel, x, y) ≈ trueXY

            fill!(tmp, 0)
            kernelmatrix!(tmp, kernel, x)
            @test tmp ≈ trueX

            fill!(tmp, 0)
            kernelmatrix!(tmp, kernel, x, y)
            @test tmp ≈ trueXY

            @test kerneldiagmatrix(kernel, x) ≈ diag(kernelmatrix(kernel, x))

            fill!(tmp_diag, 0)
            kerneldiagmatrix!(tmp_diag, kernel, x)
            @test tmp_diag ≈ diag(kernelmatrix(kernel, x))
        end
    end

    @testset "single kernel" begin
        kernel = TensorProduct(k1)
        @test length(kernel) == 1

        @testset "eval" begin
            for (x, y) in (((v1,), (v2,)), ([v1], [v2]))
                val = k1(x[1], y[1])

                @test kernel(x, y) == val
            end
        end

        @testset "kernelmatrix" begin
            N = 10

            x = randn(rng, N)
            y = randn(rng, N)
            vectors = (x, y)

            X = reshape(x, 1, :)
            x_cols = ColVecs(X)
            x_rows = RowVecs(X')
            Y = reshape(y, 1, :)
            y_cols = ColVecs(Y)
            y_rows = RowVecs(Y')

            trueX = kernelmatrix(k1, x)
            trueXY = kernelmatrix(k1, x, y)
            tmp = Matrix{Float64}(undef, N, N)
            tmp_diag = Vector{Float64}(undef, N)

            for (x, y) in ((x, y), (x_cols, y_cols), (x_rows, y_rows))

                @test kernelmatrix(kernel, x) ≈ trueX

                @test kernelmatrix(kernel, x, y) ≈ trueXY

                fill!(tmp, 0)
                kernelmatrix!(tmp, kernel, x)
                @test tmp ≈ trueX

                fill!(tmp, 0)
                kernelmatrix!(tmp, kernel, x, y)
                @test tmp ≈ trueXY

                @test kerneldiagmatrix(kernel, x) ≈ diag(kernelmatrix(kernel, x))

                fill!(tmp_diag, 0)
                kerneldiagmatrix!(tmp_diag, kernel, x)
                @test tmp_diag ≈ diag(kernelmatrix(kernel, x))
            end
        end
    end
    test_ADs(()->TensorProduct(SqExponentialKernel(), LinearKernel()), dims = [2, 2]) # ADs = [:ForwardDiff, :ReverseDiff])
end


text:
## generic rrule for broadcasted that we DON'T use can be found here:
## https://github.com/JuliaDiff/ChainRules.jl/issues/531
## below are implementations of rrule for broadcasting of the most popular functions

function rrule(::typeof(Broadcast.broadcasted), ::typeof(tanh), x)
  y = tanh.(x)
  function bcast_tanh_pullback(dy)
    dx = @. (1 - y ^ 2) * dy
    return NoTangent(), NoTangent(), dx
  end
  return y, bcast_tanh_pullback
end

function rrule(::typeof(Broadcast.broadcasted), ::typeof(-), a, b)
  y = a .- b
  function bcast_minus_pullback(dy)
    return NoTangent(), NoTangent(), dy, -dy
  end
  return y, bcast_minus_pullback
end


text:
# ------------ uv.h ------------
# Documentation:
# https://github.com/clibs/uv/blob/master/docs/src/threading.rst
const uv_thread_t = Culong # = pthread_t

"""
Returns the maximum size of the mask used for process/thread affinities,
or `UV_ENOTSUP` if affinities are not supported on the current platform.

Ref: [docs](https://github.com/clibs/uv/blob/d0240ce496fcd86d45e8a6b211732220fdb27eac/docs/src/misc.rst#L285)
"""
uv_cpumask_size() = @ccall uv_cpumask_size()::Cint

"""
Ref: [docs](https://github.com/clibs/uv/blob/d0240ce496fcd86d45e8a6b211732220fdb27eac/docs/src/threading.rst#L130)
"""
uv_thread_self() = @ccall uv_thread_self()::uv_thread_t

"""
    uv_thread_getaffinity(self_ref, cpumask, masksize)
Gets the specified thread's affinity setting. On Unix, this maps the
`cpu_set_t` returned by `pthread_getaffinity_np(3)` to bytes in `cpumask`.

The `masksize` specifies the number of entries (bytes) in `cpumask`,
and must be greater-than-or-equal-to `uv_cpumask_size`.

**Note:** Thread affinity getting is not atomic on Windows and unsupported on macOS.

Ref: [docs](https://github.com/clibs/uv/blob/master/docs/src/threading.rst)
"""
function uv_thread_getaffinity(self_ref, cpumask, masksize)
    @ccall uv_thread_getaffinity(
        self_ref::Ptr{uv_thread_t}, cpumask::Ptr{Cchar}, masksize::Cssize_t
    )::Cint
end

"""
    uv_thread_getaffinity()
Query the calling thread's affinity.
"""
function uv_thread_getaffinity()
    # inspired by https://github.com/JuliaLang/julia/pull/42340
    masksize = uv_cpumask_size()
    self = uv_thread_self()
    self_ref = Ref(self)
    cpumask = zeros(Cchar, masksize)
    err = uv_thread_getaffinity(self_ref, cpumask, masksize)
    @assert err == 0
    n = findlast(isone, cpumask)
    resize!(cpumask, n)
    return cpumask
end

"""
    uv_thread_setaffinity(self_ref, cpumask, oldmask, masksize)
Sets the specified thread's affinity to `cpumask`, which is specified in
bytes. Optionally returning the previous affinity setting in `oldmask`.
On Unix, uses `pthread_getaffinity_np(3)` to get the affinity setting
and maps the `cpu_set_t` to bytes in `oldmask`. Then maps the bytes in `cpumask`
to a `cpu_set_t` and uses `pthread_setaffinity_np(3)`. On Windows, maps
the bytes in `cpumask` to a bitmask and uses `SetThreadAffinityMask()` which
returns the previous affinity setting.

The `masksize` specifies the number of entries (bytes) in `cpumask / oldmask`,
and must be greater-than-or-equal-to `uv_cpumask_size()`.

**Note:** Thread affinity setting is not atomic on Windows and unsupported on macOS.

Ref: [docs](https://github.com/clibs/uv/blob/master/docs/src/threading.rst)
"""
function uv_thread_setaffinity(self_ref, cpumask, oldmask, masksize)
    @ccall uv_thread_setaffinity(
        self_ref::Ptr{uv_thread_t},
        cpumask::Ptr{Cchar},
        oldmask::Ptr{Cchar},
        masksize::Csize_t,
    )::Cint
end

"""
    uv_thread_setaffinity(procid::Integer)
Set the calling thread's affinity to `procid`.
"""
function uv_thread_setaffinity(procid::Integer)
    masksize = uv_cpumask_size()
    0 ≤ procid ≤ masksize ||
        throw(ArgumentError("Invalid procid. It must hold 0 ≤ procid ≤ masksize."))
    self = uv_thread_self()
    self_ref = Ref(self)
    cpumask = zeros(Cchar, masksize)
    cpumask[procid + 1] = 1
    err = uv_thread_setaffinity(self_ref, cpumask, C_NULL, masksize)
    return err == 0
end


text:
using RecursiveArrayTools, DataFrames, GLM
struct Regression_Sensitivity_Coefficients
    Pearson
    Standard_Regression
    Partial_Correlation
    Spearman
    Standard_Rank_Regression
    Partial_Rank_Correlation
end

function regression_sensitivity(f,p_range,p_fixed,n;coeffs=:rank)

    pearson_coeffs = []
    src_coeffs = []
    partial_coeffs = []
    spearman_coeffs = []
    srr_coeffs = []
    partial_rank_coeffs = []

    for i in 1:length(p_range)
        pcs = []
        srcs = []
        params = []
        for j in 1:length(p_range)
            if i != j
                push!(params,p_fixed[j])
            else
                push!(params,0)
            end
        end
        xis = []
        yis = Array{Float64}[]
        for j in 1:n
            params[i] = (p_range[i][2]-p_range[i][1])*rand() + p_range[i][1]
            x = params
            y = f(params)
            push!(xis,copy(x))
            push!(yis,copy(y))
        end
        yis = VectorOfArray(yis)
        x_mean = mean(xis)
        x_vrs = [sum(xis[i] - x_mean) for o in 1:length(xis)]
        x_lm = [j[i] for j in xis]
        x_rnk = sortperm(x_lm)
        x_rnk_mean = mean(x_rnk)
        x_rnk_vrs = [x_rnk[i] - x_rnk_mean for o in 1:length(x_rnk)]
        x_var = var(x_lm)
        x_rnk_var = var(x_rnk)
        pcs,srcs = pcs_and_srcs(yis,x_lm,x_vrs,x_var)
        pcc = pcc_f(x_lm,yis)
        yis_rnk = zero(yis)
        for k in 1:size(yis)[1]
            for j in 1:size(yis)[2]
                yis_rnk[k,j,:] = sortperm(yis[k,j,:])
            end
        end
        if coeffs == :rank
            prcc = pcc_f(x_rnk,yis_rnk)
            spear_rcc, srrc = pcs_and_srcs(yis_rnk,x_rnk,x_rnk_vrs,x_rnk_var)
        end
        push!(pearson_coeffs,pcs)
        push!(src_coeffs,srcs)
        push!(partial_coeffs,pcc)
        if coeffs == :rank
            push!(partial_rank_coeffs,prcc)
            push!(spearman_coeffs,spear_rcc)
            push!(srr_coeffs,srrc)
        end
    end
    if coeffs == :rank
        regre_coeff = Regression_Sensitivity_Coefficients(pearson_coeffs,src_coeffs,partial_coeffs,spearman_coeffs,partial_rank_coeffs,srr_coeffs)
    else
        regre_coeff = Regression_Sensitivity_Coefficients(pearson_coeffs,src_coeffs,partial_coeffs,nothing,nothing,nothing)
    end
    regre_coeff
end

function pcs_and_srcs(yis,x_lm,x_vrs,x_var)
    pcs = []
    srcs = []
    for k in 1:size(yis)[1]
        pc = []
        src = []
        for j in 2:size(yis)[2]
            pear_coeff_num = sum(x_vrs .* (yis[k,j,:] - mean(yis[k,j,:])))
            pear_coeff_deno = sqrt(sum(x_vrs.^2)) * sqrt(sum((yis[k,j,:] - mean(yis[k,j,:])).^2))
            pear_coeff = pear_coeff_num/pear_coeff_deno
            push!(pc,pear_coeff)
            df = DataFrame(X=x_lm,Y=yis[k,j,:])
            lin_model = lm(@formula(Y ~ X),df)
            push!(src,coef(lin_model)[2].*sqrt.(x_var ./ var(yis[k,j,:])))
        end
        push!(pcs,pc)
        push!(srcs,src)
    end
    pcs,srcs
end

function pcc_f(x_lm,yis)
    df_arr = Array{Float64}[Float64[] for i in 1:length(x_lm)]
    for o in 1:length(x_lm)
        for j in 1:length(x_lm)
            push!(df_arr[o],x_lm[(j+o)%length(x_lm)+1])
        end
    end
    df = DataFrame(df_arr)
    last_ind = length(x_lm)
    formula = "@formula(x$last_ind ~ "
    for j in 1:length(x_lm)-1
        formula *= "x$j +"
    end
    formula = formula[1:end-1] * ")"
    ols_x = lm(eval(Meta.parse(formula)),df)
    x_cap = x_lm .- predict(ols_x)
    y_caps = []
    for k in 1:size(yis)[1]
        y_cps = []
        for j in 2:size(yis)[2]
            df_arr = Array{Float64}[Float64[] for i in 1:length(yis[k,j,:])]
            for o in 1:length(yis[k,j,:])
                for l in 1:length(yis[k,j,:])
                    push!(df_arr[o],yis[k,j,:][(l+o)%length(yis[k,j,:])+1])
                end
            end
            df = DataFrame(df_arr)
            last_ind = length(yis[k,j,:])
            formula = "@formula(x$last_ind ~ "
            for j in 1:length(yis[k,j,:])-1
                formula *= "x$j +"
            end
            formula = formula[1:end-1] * ")"
            ols_y = lm(eval(Meta.parse(formula)),df)
            y_cap = yis[k,j,:] .- predict(ols_y)
            push!(y_cps,y_cap)
        end
        push!(y_caps,y_cps)
    end
    pcc = []
    for k in 1:length(y_caps)
        pcc_ = []
        for j in 1:length(y_caps[k])
            num = sum(x_cap .* y_caps[k][j])
            deno = sqrt(sum(x_cap.^2)*sum(y_caps[k][j].^2))
            push!(pcc_,num/deno)
        end
        push!(pcc,pcc_)
    end
    pcc
end

function regression_sensitivity(prob::DiffEqBase.DEProblem,alg,t,p_range,p_fixed,n;coeffs=:rank)
    f = function (p)
        prob1 = remake(prob;p=p)
        Array(solve(prob1,alg;saveat=t))
    end
    regression_sensitivity(f,p_range,p_fixed,n,coeffs=coeffs)
end


text:
"""
    AutoCorrelationFunction{T}

A representation of the autocorrelation function.

It supports accessing the coeficients starting from lag 0 by
- `coef(acf::AutoCorrelationFunction)`: implements StatsBase.coef
- `coef(acf::AutoCorrelationFunction, lag::Integer)`: coefficient for lag

Wrapping the vector of coefficients into its own type helps avoiding
method ambiguities.

# Examples
```jldoctest am; output = false
using StatsBase: coef
acf = AutoCorrelationFunction([1,0.4,0.1])
coef(acf) == [1,0.4,0.1]
coef(acf,1) == 0.4
# output
true
```
"""
struct AutoCorrelationFunction{T}
    coef::T
end
# implements StatsBase coef
AutoCorrelationFunction(coef::AbstractVector{<:Number}) =
    AutoCorrelationFunction{typeof(coef)}(coef)
coef(acf::AutoCorrelationFunction) = acf.coef
coef(acf::AutoCorrelationFunction, lag::Integer) = acf.coef[lag+1]

cormatrix_for_acf(n::Int,acf::AutoCorrelationFunction) =
    cormatrix_for_acf(n, coef(acf))

function cormatrix_for_acf(n::Int,acf::AbstractVector)
    nacf::Int = length(acf)
    corrM = BandedMatrix{Float64}(undef, (n,n), (nacf-1,nacf-1))
    corrM[band(0)] .= acf[1]
    for i in 1:(nacf-1)
      corrM[band(i)] .= corrM[band(-i)] .= acf[i+1]
    end
    corrM
end


text:
include("helper.jl")

my_slow_tests = [
    "problems/test_optimize_single_objective_problems.jl",
    "test_bboptimize.jl"
]

@testset "BlackBoxOptim long-running test suite $t" for t in my_slow_tests
    include(t)
end


text:
module steadystate

using ..states
using ..operators
using ..superoperators
using ..timeevolution
using ..metrics

type ConvergenceReached <: Exception end

"""
Calculate steady state using long time master equation evolution.

Arguments
---------

H
    Operator specifying the Hamiltonian.
J
    Vector of jump operators.

Keyword Arguments
-----------------

rho0
    Initial density operator. If not given the :math:`|0 \\rangle\\langle0|`
    state in respect to the choosen basis is used.
eps
    Tracedistance used as termination criterion.
hmin
    Minimal time step used in the time evolution.
Gamma
    Vector or matrix specifying the coefficients for the jump operators.
Jdagger (optional)
    Vector containing the hermitian conjugates of the jump operators. If they
    are not given they are calculated automatically.
fout (optional)
    If given this function fout(t, rho) is called every time an output should
    be displayed. To limit copying to a minimum the given density operator rho
    is further used and therefore must not be changed.
kwargs
    Further arguments are passed on to the ode solver.
"""
function master(H::Operator, J::Vector;
                rho0::DenseOperator=tensor(basis_ket(H.basis_l, 1), basis_bra(H.basis_r, 1)),
                eps::Float64=1e-3, hmin=1e-7,
                Gamma::Union{Vector{Float64}, Matrix{Float64}}=ones(Float64, length(J)),
                Jdagger::Vector=map(dagger, J),
                fout::Union{Function,Void}=nothing,
                tmp::DenseOperator=deepcopy(rho0),
                kwargs...)
    t0 = 0.
    rho0 = deepcopy(rho0)
    function fout_steady(t, rho)
        if fout!=nothing
            fout(t, rho)
        end
        dt = t - t0
        drho = metrics.tracedistance(rho0, rho)
        t0 = t
        rho0.data[:] = rho.data
        if drho/dt < eps
            throw(ConvergenceReached())
        end
    end
    try
        timeevolution.master([0., Inf], rho0, H, J; Gamma=Gamma, Jdagger=Jdagger,
                            hmin=hmin, hmax=Inf,
                            tmp=tmp,
                            display_initialvalue=false,
                            display_finalvalue=false,
                            display_intermediatesteps=true,
                            fout=fout_steady, kwargs...)
    catch e
        if !isa(e, ConvergenceReached)
            rethrow(e)
        end
    end
    return rho0
end

"""
Find steady state by calculating the eigenstate of the Liouvillian matrix.

Arguments
---------

L
    Dense or sparse super-operator.
"""
function eigenvector(L::DenseSuperOperator)
    d, v = Base.eig(L.data)
    data = reshape(v[:,1], length(L.basis_r[1]), length(L.basis_r[2]))
    return DenseOperator(L.basis_r[1], L.basis_r[2], data)
end

function eigenvector(L::SparseSuperOperator)
    d, v, nconv, niter, nmult, resid = Base.eigs(L.data; nev=1, sigma=1e-30)
    data = reshape(v[:,1], length(L.basis_r[1]), length(L.basis_r[2]))
    op = DenseOperator(L.basis_r[1], L.basis_r[2], data)
    return op/trace(op)
end

"""
Find steady state by calculating the eigenstate of the Liouvillian matrix.

Arguments
---------

H
    Operator specifying the Hamiltonian.
J
    Vector of jump operators.
"""
eigenvector(H::Operator, J::Vector) = eigenvector(liouvillian(H, J))

end # module


text:
#
# meta.jl --
#
# Functions to generate code for interpolation methods.
#
#------------------------------------------------------------------------------
#
# This file is part of the LinearInterpolators package licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2016-2021, Éric Thiébaut.
#

module Meta

const Arg = Union{Number,Symbol,Expr}

#"""
#```julia
#@getcoefs N j w ker lim pos
#```
#
#yields the code for extracting the interpolation coefficients, `N` is the
#interpolation kernel size, `j` and `w` are prefixes for the variables to store
#the indices and weights for interpolation (the actual variable names are given
#by appending a `_i` suffix for `i` in `1:N`), `ker` is the interpolation
#kernel, `lim` are the limits and `pos` is the position of interpolation.
#
#"""
#macro getcoefs(N, j, w, ker, lim, pos)
#    generate_getcoefs(N, j, w, ker, lim, pos)
#end
#
#macro nonsep_interp(alpha, src, beta, dst,
#                    s1, ker1, lim1, pos1,
#                    s2, ker2, lim2, pos2)
#    return generate_nonsep_interp(alpha, src, beta, dst,
#                          s1, ker1, lim1, pos1,
#                          s2, ker2, lim2, pos2)
#end
#
#macro nonsep_interp_adj(src, dst,
#                        s1, ker1, lim1, pos1,
#                        s2, ker2, lim2, pos2)
#    return generate_nonsep_interp_adj(src, dst,
#                              s1, ker1, lim1, pos1,
#                              s2, ker2, lim2, pos2)
#end

"""
    generate_getcoefs(n, j, w, ker, lim, pos)

generates the code for getting the interpolation coefficients.  Here `n` is the
size of `ker` the interpolation kernel, `j` and `w` are symbols or strings used
as prefixes for local variables to store the interpolation indices and weights
respectively, `ker` and `lim` are symbols with the name of the local variables
which store the interpolation kernel and the limits along the dimension of
interpolation, `pos` is a symbol or an expression which gives the position to
interpolate.  For instance:

    generate_getcoefs(2, :i, :c, :kr, :lm, :(x[i]))
        -> :((i_1, i_2, c_1, c_2) = getcoefs(kr, lm, x[i]))

Another possibility is:

    generate_getcoefs(J, W, ker, lim, pos)

where `J` and `W` are vectors of length `n`, the size of `ker` the
interpolation kernel, with the symbolic names of the local variables to store
the interpolation indices and weights respectively.

"""
function generate_getcoefs(n::Integer,
                           j::Union{AbstractString,Symbol},
                           w::Union{AbstractString,Symbol},
                           ker::Symbol, lim::Symbol, pos::Arg)
    return generate_getcoefs([Symbol(j,"_",s) for s in 1:n],
                             [Symbol(w,"_",s) for s in 1:n], ker, lim, pos)
end

function generate_getcoefs(J::AbstractVector{Symbol},
                           W::AbstractVector{Symbol},
                           ker::Symbol, lim::Symbol, pos::Arg)
    @assert length(J) == length(W)
    vars =  Expr(:tuple, J..., W...)
    return :($vars = getcoefs($ker, $lim, $pos))
end

"""
    generate_sum(ex)

generates an expression whose result is the sum of the elements of the vector
`ex` which are symbols (being interpreted as the name of variables) or
expressions.

"""
generate_sum(args::Arg...) = generate_sum(args)
generate_sum(ex::Union{AbstractVector,Tuple{Vararg{Arg}}}) =
    (length(ex) == 1 ? ex[1] : Expr(:call, :+, ex...))

"""
    group_expressions(ex...) -> code

generates a single expression from all expressions given in argument.  The
result may be the same as the input if it is a single expression or a block of
expressions if several expressions are specified.

To insert the result as a block of statements (like a `begin` ... `end` block)
in a quoted expression, write something like:

    quote
        some_pre_code
        \$code
        some_post_code
    end

To strip the surrounding `begin` ... `end` keywords, write instead:

    quote
        some_pre_code
        \$(code.args...)
        some_post_code
    end

"""
group_expressions(args::Expr...) = group_expressions(args)
group_expressions(ex::Union{AbstractVector{Expr},Tuple{Vararg{Expr}}}) =
    (length(ex) == 1 ? ex[1] : Expr(:block, ex...))

function generate_interp_expr(arr::Symbol,
                              J::AbstractVector{Symbol},
                              W::AbstractVector{Symbol})
    @assert length(J) == length(W)
    return generate_sum([:($arr[$(J[i])]*$(W[i])) for i in 1:length(J)])
end

function generate_interp_expr(arr::Symbol,
                              J::AbstractVector{Symbol},
                              W::AbstractVector{Symbol},
                              inds::Arg)
    @assert length(J) == length(W)
    return generate_sum([:($arr[$(J[i]),$inds]*$(W[i])) for i in 1:length(J)])
end

function generate_interp_expr(arr::Symbol,
                              J1::AbstractVector{Symbol},
                              W1::AbstractVector{Symbol},
                              J2::AbstractVector{Symbol},
                              W2::AbstractVector{Symbol})
    @assert length(J2) == length(W2)
    n = length(J2)
    ex = Array{Expr}(undef, n)
    for i in 1:n
        sub = generate_interp_expr(arr, J1, W1, J2[i])
        ex[i] = :($sub*$(W2[i]))
    end
    return generate_sum(ex)
end

end # module


text:
using Pkg
Pkg.activate("TowardsNotebook")
Pkg.instantiate()

import TowardsNotebook
TowardsNotebook.notebook()

max_stars_repo_name: Jblew/towards-julia-notebook

max_issues_repo_name: Jblew/towards-julia-notebook

max_forks_repo_name: Jblew/towards-julia-notebook