text:
#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

using Base.Meta

_is_sum(s::Symbol) = (s == :sum) || (s == :∑) || (s == :Σ)
_is_prod(s::Symbol) = (s == :prod) || (s == :∏)

"""
    _add_kw_args(call, kw_args)

Add the keyword arguments `kw_args` to the function call expression `call`,
escaping the expressions. The elements of `kw_args` should be expressions of the
form `:(key = value)`. The `kw_args` vector can be extracted from the arguments
of a macro with [`Containers._extract_kw_args`](@ref).

## Examples

```jldoctest; setup = :(using JuMP)
julia> call = :(f(1, a=2))
:(f(1, a=2))

julia> JuMP._add_kw_args(call, [:(b=3), :(c=4)])

julia> call
:(f(1, a=2, $(Expr(:escape, :(b=3))), $(Expr(:escape, :(c=4)))))
```
"""
function _add_kw_args(call, kw_args)
    for kw in kw_args
        @assert isexpr(kw, :(=))
        push!(call.args, esc(Expr(:kw, kw.args...)))
    end
end

_valid_model(m::AbstractModel, name) = nothing
function _valid_model(m, name)
    return error(
        "Expected $name to be a JuMP model, but it has type ",
        typeof(m),
    )
end

"""
    _finalize_macro(model, code, source::LineNumberNode)

Wraps the `code` generated by a macro in a code block with the first argument as
`source`, the `LineNumberNode` of where the macro was called from in the user's
code. This results in better stacktraces in error messages.

In addition, it checks that `model` is a valid `AbstractModel`.

Note: assumes `model` is already escaped.
"""
function _finalize_macro(model, code, source::LineNumberNode)
    return Expr(
        :block,
        source,
        :(_valid_model($model, $(quot(model.args[1])))),
        code,
    )
end

function _error_if_cannot_register(model::AbstractModel, name::Symbol)
    obj_dict = object_dictionary(model)
    if haskey(obj_dict, name)
        error(
            """An object of name $name is already attached to this model. If this
          is intended, consider using the anonymous construction syntax, e.g.,
          `x = @variable(model, [1:N], ...)` where the name of the object does
          not appear inside the macro.

          Alternatively, use `unregister(model, :$(name))` to first unregister
          the existing name from the model. Note that this will not delete the
          object; it will just remove the reference at `model[:$(name)]`.
      """,
        )
    end
    return
end

function _error_if_cannot_register(model::AbstractModel, name)
    return error("Invalid name $name.")
end

"""
    _macro_assign_and_return(code, variable, name;
                             model_for_registering=nothing)

Return `code` which returns the value of `variable` and then assigns
`variable` to `name`. If `model_for_registering` is given, the generated code
assigns the resulting object to the model dictionary.
"""
function _macro_assign_and_return(
    code,
    variable,
    name;
    model_for_registering = nothing,
)
    return quote
        $(
            if model_for_registering !== nothing
                :(_error_if_cannot_register(
                    $model_for_registering,
                    $(quot(name)),
                ))
            end
        )
        $variable = $code
        $(
            if model_for_registering !== nothing
                :($model_for_registering[$(quot(name))] = $variable)
            end
        )
        # This assignment should be in the scope calling the macro
        $(esc(name)) = $variable
    end
end

function _check_vectorized(sense::Symbol)
    sense_str = string(sense)
    if sense_str[1] == '.'
        Symbol(sense_str[2:end]), true
    else
        sense, false
    end
end

# two-argument build_constraint is used for one-sided constraints.
# Right-hand side is zero.

"""
    sense_to_set(_error::Function, ::Val{sense_symbol})

Converts a sense symbol to a set `set` such that
`@constraint(model, func sense_symbol 0)` is equivalent to
`@constraint(model, func in set)` for any `func::AbstractJuMPScalar`.

## Example

Once a custom set is defined you can directly create a JuMP constraint with it:
```jldoctest sense_to_set; setup = :(using JuMP)
julia> struct CustomSet{T} <: MOI.AbstractScalarSet
           value::T
       end

julia> Base.copy(x::CustomSet) = CustomSet(x.value)

julia> model = Model();

julia> @variable(model, x)
x

julia> cref = @constraint(model, x in CustomSet(1.0))
x ∈ CustomSet{Float64}(1.0)
```

However, there might be an appropriate sign that could be used in order to
provide a more convenient syntax:
```jldoctest sense_to_set
julia> JuMP.sense_to_set(::Function, ::Val{:⊰}) = CustomSet(0.0)

julia> MOIU.shift_constant(set::CustomSet, value) = CustomSet(set.value + value)

julia> cref = @constraint(model, x ⊰ 1)
x ∈ CustomSet{Float64}(1.0)
```
Note that the whole function is first moved to the right-hand side, then the
sign is transformed into a set with zero constant and finally the constant is
moved to the set with `MOIU.shift_constant`.
"""
function sense_to_set end

function sense_to_set(_error::Function, ::Union{Val{:(<=)},Val{:(≤)}})
    return MOI.LessThan(0.0)
end
function sense_to_set(_error::Function, ::Union{Val{:(>=)},Val{:(≥)}})
    return MOI.GreaterThan(0.0)
end
sense_to_set(_error::Function, ::Val{:(==)}) = MOI.EqualTo(0.0)
function sense_to_set(_error::Function, ::Val{S}) where {S}
    return _error("Unrecognized sense $S")
end

function _build_call(_error::Function, vectorized::Bool, func, set)
    return if vectorized
        :(build_constraint.($_error, $func, Ref($(esc(set)))))
    else
        :(build_constraint($_error, $func, $(esc(set))))
    end
end
function parse_one_operator_constraint(
    _error::Function,
    vectorized::Bool,
    ::Union{Val{:in},Val{:∈}},
    func,
    set,
)
    variable, parse_code = _MA.rewrite(func)
    return parse_code, _build_call(_error, vectorized, variable, set)
end
function parse_one_operator_constraint(_error::Function, args...)
    return _unknown_constraint_expr(_error)
end

_functionize(v::VariableRef) = convert(AffExpr, v)
_functionize(v::AbstractArray{VariableRef}) = _functionize.(v)
_functionize(x) = x
_functionize(::MutableArithmetics.Zero) = 0.0
function parse_one_operator_constraint(
    _error::Function,
    vectorized::Bool,
    sense::Val,
    lhs,
    rhs,
)
    # Simple comparison - move everything to the LHS.
    #
    # `_functionize` deals with the pathological case where the `lhs` is a `VariableRef`
    # and the `rhs` is a summation with no terms. `_build_call` should be passed a
    # `GenericAffExpr` or a `GenericQuadExpr`, and not a `VariableRef` as would be the case
    # without `_functionize`.
    if vectorized
        func = :($lhs .- $rhs)
    else
        func = :($lhs - $rhs)
    end
    set = sense_to_set(_error, sense)
    variable, parse_code = _MA.rewrite(func)
    return parse_code,
    _build_call(_error, vectorized, :(_functionize($variable)), set)
end

function parse_constraint_expr(_error::Function, expr::Expr)
    return parse_constraint_head(_error, Val(expr.head), expr.args...)
end
function parse_constraint_head(_error::Function, ::Val{:call}, args...)
    return parse_constraint(_error, args...)
end

function parse_constraint(_error::Function, sense::Symbol, lhs, rhs)
    (sense, vectorized) = _check_vectorized(sense)
    return vectorized,
    parse_one_operator_constraint(_error, vectorized, Val(sense), lhs, rhs)...
end

function parse_ternary_constraint(
    _error::Function,
    vectorized::Bool,
    lb,
    ::Union{Val{:(<=)},Val{:(≤)}},
    aff,
    rsign::Union{Val{:(<=)},Val{:(≤)}},
    ub,
)
    newaff, parseaff = _MA.rewrite(aff)
    newlb, parselb = _MA.rewrite(lb)
    newub, parseub = _MA.rewrite(ub)
    if vectorized
        buildcall = :(build_constraint.($_error, $newaff, $newlb, $newub))
    else
        buildcall = :(build_constraint($_error, $newaff, $newlb, $newub))
    end
    return parseaff, parselb, parseub, buildcall
end

function parse_ternary_constraint(
    _error::Function,
    vectorized::Bool,
    ub,
    ::Union{Val{:(>=)},Val{:(≥)}},
    aff,
    rsign::Union{Val{:(>=)},Val{:(≥)}},
    lb,
)
    return parse_ternary_constraint(
        _error,
        vectorized,
        lb,
        Val(:(<=)),
        aff,
        Val(:(<=)),
        ub,
    )
end

function parse_ternary_constraint(_error::Function, args...)
    return _error(
        "Only two-sided rows of the form lb <= expr <= ub or ub >= expr >= lb are supported.",
    )
end

function parse_constraint_head(
    _error::Function,
    ::Val{:comparison},
    lb,
    lsign::Symbol,
    aff,
    rsign::Symbol,
    ub,
)
    return parse_constraint(_error, lb, lsign, aff, rsign, ub)
end

function parse_constraint(
    _error::Function,
    lb,
    lsign::Symbol,
    aff,
    rsign::Symbol,
    ub,
)
    (lsign, lvectorized) = _check_vectorized(lsign)
    (rsign, rvectorized) = _check_vectorized(rsign)
    ((vectorized = lvectorized) == rvectorized) ||
        _error("Signs are inconsistently vectorized")
    parseaff, parselb, parseub, buildcall = parse_ternary_constraint(
        _error,
        vectorized,
        lb,
        Val(lsign),
        aff,
        Val(rsign),
        ub,
    )
    parsecode = quote
        $parseaff
        $parselb
        $parseub
    end
    return vectorized, parsecode, buildcall
end

function _unknown_constraint_expr(_error::Function)
    # Unknown
    return _error(
        "Constraints must be in one of the following forms:\n" *
        "       expr1 <= expr2\n" *
        "       expr1 >= expr2\n" *
        "       expr1 == expr2\n" *
        "       lb <= expr <= ub",
    )
end

function parse_constraint_head(_error::Function, ::Val, args...)
    return _unknown_constraint_expr(_error)
end
function parse_constraint(_error::Function, args...)
    # Define this as the last fallback: either this is a function call that may
    # be overridden by extensions, or a syntax that is not recognised.
    # Multiple dispatch does not work here, due to ambiguity with:
    #     parse_constraint(_error::Function, lb, lsign::Symbol, aff, rsign::Symbol, ub)
    if args[1] isa Symbol
        (sense, vectorized) = _check_vectorized(args[1])
        vectorized,
        parse_one_operator_constraint(
            _error,
            vectorized,
            Val(sense),
            args[2:end]...,
        )...
    else
        _unknown_constraint_expr(_error)
    end
end

# Generic fallback.
function build_constraint(
    _error::Function,
    func,
    set::Union{MOI.AbstractScalarSet,MOI.AbstractVectorSet},
)
    return _error(
        "unable to add the constraint because we don't recognize " *
        "$(func) as a valid JuMP function.",
    )
end

function build_constraint(
    _error::Function,
    v::AbstractJuMPScalar,
    set::MOI.AbstractScalarSet,
)
    return ScalarConstraint(v, set)
end
function build_constraint(
    _error::Function,
    expr::Union{GenericAffExpr,GenericQuadExpr},
    set::MOI.AbstractScalarSet,
)
    offset = constant(expr)
    add_to_expression!(expr, -offset)
    return ScalarConstraint(expr, MOIU.shift_constant(set, -offset))
end
function build_constraint(
    _error::Function,
    α::Number,
    set::MOI.AbstractScalarSet,
)
    return build_constraint(_error, convert(AffExpr, α), set)
end
function build_constraint(
    _error::Function,
    ::MutableArithmetics.Zero,
    set::MOI.AbstractScalarSet,
)
    return build_constraint(_error, zero(AffExpr), set)
end

function build_constraint(
    ::Function,
    x::AbstractVector{<:AbstractJuMPScalar},
    set::MOI.AbstractVectorSet,
)
    return VectorConstraint(x, set)
end
function build_constraint(
    _error::Function,
    a::Vector{<:Number},
    set::MOI.AbstractVectorSet,
)
    return build_constraint(_error, convert(Vector{AffExpr}, a), set)
end

function build_constraint(
    _error::Function,
    x::AbstractArray,
    set::MOI.AbstractScalarSet,
)
    return _error(
        "Unexpected vector in scalar constraint. Did you mean to use",
        " the dot comparison operators like .==, .<=, and .>=",
        " instead?",
    )
end

function build_constraint(
    _error::Function,
    x::Matrix,
    set::MOI.AbstractVectorSet,
)
    return _error(
        "unexpected matrix in vector constraint. Do you need to flatten the " *
        "matrix into a vector using `vec()`?",
    )
end

function build_constraint(
    _error::Function,
    ::Matrix,
    T::Union{
        MOI.PositiveSemidefiniteConeSquare,
        MOI.PositiveSemidefiniteConeTriangle,
    },
)
    return _error("instead of `$(T)`, use `JuMP.PSDCone()`.")
end

# three-argument build_constraint is used for two-sided constraints.
function build_constraint(
    _error::Function,
    func::AbstractJuMPScalar,
    lb::Real,
    ub::Real,
)
    if isnan(lb) || isnan(ub)
        _error("Invalid bounds, cannot contain NaN: [$(lb), $(ub)].")
    end
    return build_constraint(
        _error,
        func,
        MOI.Interval(Float64(lb), Float64(ub)),
    )
end

# This method intercepts `@constraint(model, lb <= var <= ub)` and promotes
# `var` to an `AffExpr` to form a `ScalarAffineFunction-in-Interval` instead of
# `SingleVariable-in-Interval`. To create a
# `MOI.SingleVariable`-in-`MOI.Interval`, use
# `@constraint(model, var in MOI.Interval(lb, ub))`. We do this for consistency
# with how one-sided (in)equality constraints are parsed.
function build_constraint(
    _error::Function,
    func::AbstractVariableRef,
    lb::Real,
    ub::Real,
)
    return build_constraint(_error, 1.0func, lb, ub)
end

function build_constraint(_error::Function, expr, lb, ub)
    lb isa Number || _error(string("Expected $lb to be a number."))
    ub isa Number || _error(string("Expected $ub to be a number."))
    if lb isa Number && ub isa Number
        _error("Range constraint is not supported for $expr.")
    end
end

function build_constraint(
    ::Function,
    x::AbstractVector{<:AbstractJuMPScalar},
    set::MOI.SOS1,
)
    return VectorConstraint(x, MOI.SOS1{Float64}(set.weights))
end

function build_constraint(
    ::Function,
    x::AbstractVector{<:AbstractJuMPScalar},
    set::MOI.SOS2,
)
    return VectorConstraint(x, MOI.SOS2{Float64}(set.weights))
end

# TODO: update 3-argument @constraint macro to pass through names like @variable

"""
    _constraint_macro(
        args, macro_name::Symbol, parsefun::Function, source::LineNumberNode
    )

Returns the code for the macro `@constraint_like args...` of syntax
```julia
@constraint_like con     # Single constraint
@constraint_like ref con # group of constraints
```
where `@constraint_like` is either `@constraint` or `@SDconstraint`.

The expression `con` is parsed by `parsefun` which returns a `build_constraint`
call code that, when executed, returns an `AbstractConstraint`. The macro
keyword arguments (except the `container` keyword argument which is used to
determine the container type) are added to the `build_constraint` call. The
returned value of this call is passed to `add_constraint` which returns a
constraint reference.

`source` is a `LineNumberNode` that should refer to the line that the macro was
called from in the user's code. One way of generating this is via the hidden
variable `__source__`.
"""
function _constraint_macro(
    args,
    macro_name::Symbol,
    parsefun::Function,
    source::LineNumberNode,
)
    _error(str...) = _macro_error(macro_name, args, source, str...)

    args, kw_args, requestedcontainer = Containers._extract_kw_args(args)

    if length(args) < 2
        if length(kw_args) > 0
            _error("Not enough positional arguments")
        else
            _error("Not enough arguments")
        end
    end
    m = args[1]
    x = args[2]
    extra = args[3:end]

    m = esc(m)
    # Two formats:
    # - @constraint_like(m, a*x <= 5)
    # - @constraint_like(m, myref[a=1:5], a*x <= 5)
    length(extra) > 1 && _error("Too many arguments.")
    # Canonicalize the arguments
    c = length(extra) == 1 ? x : gensym()
    x = length(extra) == 1 ? extra[1] : x

    anonvar = isexpr(c, :vect) || isexpr(c, :vcat) || length(extra) != 1
    variable = gensym()
    name = Containers._get_name(c)
    base_name = anonvar ? "" : string(name)
    # TODO: support the base_name keyword argument

    if isa(x, Symbol)
        _error(
            "Incomplete constraint specification $x. Are you missing a comparison (<=, >=, or ==)?",
        )
    end

    (x.head == :block) && _error(
        "Code block passed as constraint. Perhaps you meant to use @constraints instead?",
    )

    # Strategy: build up the code for add_constraint, and if needed we will wrap
    # in a function returning `ConstraintRef`s and give it to `Containers.container`.
    idxvars, indices = Containers._build_ref_sets(_error, c)
    if args[1] in idxvars
        _error(
            "Index $(args[1]) is the same symbol as the model. Use a " *
            "different name for the index.",
        )
    end
    vectorized, parsecode, buildcall = parsefun(_error, x)
    _add_kw_args(buildcall, kw_args)
    if vectorized
        # TODO: Pass through names here.
        constraintcall = :(add_constraint.($m, $buildcall))
    else
        constraintcall =
            :(add_constraint($m, $buildcall, $(_name_call(base_name, idxvars))))
    end
    code = quote
        $parsecode
        $constraintcall
    end

    creation_code =
        Containers.container_code(idxvars, indices, code, requestedcontainer)

    if anonvar
        # Anonymous constraint, no need to register it in the model-level
        # dictionary nor to assign it to a variable in the user scope.
        # We simply return the constraint reference
        macro_code = creation_code
    else
        # We register the constraint reference to its name and
        # we assign it to a variable in the local scope of this name
        macro_code = _macro_assign_and_return(
            creation_code,
            variable,
            name,
            model_for_registering = m,
        )
    end
    return _finalize_macro(m, macro_code, source)
end

# This function needs to be implemented by all `AbstractModel`s
constraint_type(m::Model) = ConstraintRef{typeof(m)}

"""
    @constraint(m::Model, expr)

Add a constraint described by the expression `expr`.

    @constraint(m::Model, ref[i=..., j=..., ...], expr)

Add a group of constraints described by the expression `expr` parametrized by
`i`, `j`, ...

The expression `expr` can either be

* of the form `func in set` constraining the function `func` to belong to the
  set `set` which is either a [`MOI.AbstractSet`](https://jump.dev/MathOptInterface.jl/v0.6.2/apireference.html#Sets-1)
  or one of the JuMP shortcuts [`SecondOrderCone`](@ref),
  [`RotatedSecondOrderCone`](@ref) and [`PSDCone`](@ref), e.g.
  `@constraint(model, [1, x-1, y-2] in SecondOrderCone())` constrains the norm
  of `[x-1, y-2]` be less than 1;
* of the form `a sign b`, where `sign` is one of `==`, `≥`, `>=`, `≤` and
  `<=` building the single constraint enforcing the comparison to hold for the
  expression `a` and `b`, e.g. `@constraint(m, x^2 + y^2 == 1)` constrains `x`
  and `y` to lie on the unit circle;
* of the form `a ≤ b ≤ c` or `a ≥ b ≥ c` (where `≤` and `<=` (resp. `≥` and
  `>=`) can be used interchangeably) constraining the paired the expression
  `b` to lie between `a` and `c`;
* of the forms `@constraint(m, a .sign b)` or
  `@constraint(m, a .sign b .sign c)` which broadcast the constraint creation to
  each element of the vectors.

## Note for extending the constraint macro

Each constraint will be created using
`add_constraint(m, build_constraint(_error, func, set))` where
* `_error` is an error function showing the constraint call in addition to the
  error message given as argument,
* `func` is the expression that is constrained
* and `set` is the set in which it is constrained to belong.

For `expr` of the first type (i.e. `@constraint(m, func in set)`), `func` and
`set` are passed unchanged to `build_constraint` but for the other types, they
are determined from the expressions and signs. For instance,
`@constraint(m, x^2 + y^2 == 1)` is transformed into
`add_constraint(m, build_constraint(_error, x^2 + y^2, MOI.EqualTo(1.0)))`.

To extend JuMP to accept new constraints of this form, it is necessary to add
the corresponding methods to `build_constraint`. Note that this will likely mean
that either `func` or `set` will be some custom type, rather than e.g. a
`Symbol`, since we will likely want to dispatch on the type of the function or
set appearing in the constraint.
"""
macro constraint(args...)
    return _constraint_macro(
        args,
        :constraint,
        parse_constraint_expr,
        __source__,
    )
end

function parse_SD_constraint_expr(_error::Function, expr::Expr)
    return parse_SD_constraint(_error, expr.args...)
end
function parse_SD_constraint(_error::Function, sense::Symbol, lhs, rhs)
    # Simple comparison - move everything to the LHS
    aff = :()
    if sense == :⪰ || sense == :(≥) || sense == :(>=)
        succ = lhs
        prec = rhs
    elseif sense == :⪯ || sense == :(≤) || sense == :(<=)
        succ = rhs
        prec = lhs
    else
        _error("Invalid sense $sense in SDP constraint")
    end
    if prec == 0
        aff = succ
    elseif succ == 0
        aff = :(-$prec)
    else
        aff = :($succ - $prec)
    end
    vectorized = false
    parsecode, buildcall = parse_one_operator_constraint(
        _error,
        false,
        Val(:in),
        aff,
        :(JuMP.PSDCone()),
    )
    return vectorized, parsecode, buildcall
end

function parse_SD_constraint(_error::Function, args...)
    return _error(
        "Constraints must be in one of the following forms:\n" *
        "       expr1 <= expr2\n" *
        "       expr1 >= expr2",
    )
end

"""
    @SDconstraint(model::Model, expr)

Add a semidefinite constraint described by the expression `expr`.

    @SDconstraint(model::Model, ref[i=..., j=..., ...], expr)

Add a group of semidefinite constraints described by the expression `expr`
parametrized by `i`, `j`, ...

The expression `expr` needs to be of the form `a sign b` where `sign` is `⪰`,
`≥`, `>=`, `⪯`, `≤` or `<=` and `a` and `b` are `square` matrices. It constrains
the square matrix `x` (or `-x` if the sign is `⪯`, `≤` or `<=`) to be symmetric
and positive semidefinite where

* `x = a`, if `b` is the symbol `0`,
* `x = -b`, if `a` is the symbol `0`,
* otherwise, `x = a - b`.

By default, we check numerical symmetry of the matrix `x`, and if symmetry is
violated by some arbitrary amount, we add explicit equality constraints.
You can use `Symmetric(x) in PSDCone()` with the [`@constraint`](@ref) macro to
skip these checks if you know the matrix must be symmetric; see
[`PSDCone`](@ref) for more information.

## Examples

The following constrains the matrix `[x-1 2x-2; -3 x-4]` to be symmetric and
positive semidefinite, that is, it constrains `2x-2` to be equal to `-3` and
constrains all eigenvalues of the matrix to be nonnegative.
```jldoctest; setup = :(using JuMP)
julia> model = Model();

julia> @variable(model, x)
x

julia> a = [x 2x
            0  x];

julia> b = [1 2
            3 4];

julia> cref = @SDconstraint(model, a ⪰ b)
[x - 1  2 x - 2;
 -3     x - 4  ] ∈ PSDCone()

julia> jump_function(constraint_object(cref))
4-element Array{GenericAffExpr{Float64,VariableRef},1}:
 x - 1
 -3
 2 x - 2
 x - 4

julia> moi_set(constraint_object(cref))
MathOptInterface.PositiveSemidefiniteConeSquare(2)
```
In the set `PositiveSemidefiniteConeSquare(2)` in the last output, `Square`
means that the matrix is passed as a square matrix as the corresponding
off-diagonal entries need to be constrained to be equal. A similar set
`PositiveSemidefiniteConeTriangle` exists which only uses the upper triangular
part of the matrix assuming that it is symmetric, see [`PSDCone`](@ref) to see
how to use it.
"""
macro SDconstraint(args...)
    return _constraint_macro(
        args,
        :SDconstraint,
        parse_SD_constraint_expr,
        __source__,
    )
end

"""
    @build_constraint(constraint_expr)

Constructs a `ScalarConstraint` or `VectorConstraint` using the same
machinery as [`@constraint`](@ref) but without adding the constraint to a model.

Constraints using broadcast operators like `x .<= 1` are also supported and will
create arrays of `ScalarConstraint` or `VectorConstraint`.

## Examples

```jldoctest; setup = :(using JuMP)
model = Model();
@variable(model, x);
@build_constraint(2x >= 1)

# output
ScalarConstraint{GenericAffExpr{Float64,VariableRef},MathOptInterface.GreaterThan{Float64}}(2 x, MathOptInterface.GreaterThan{Float64}(1.0))
```
"""
macro build_constraint(constraint_expr)
    function _error(str...)
        return _macro_error(
            :build_constraint,
            (constraint_expr,),
            __source__,
            str...,
        )
    end

    if isa(constraint_expr, Symbol)
        _error(
            "Incomplete constraint specification $constraint_expr. " *
            "Are you missing a comparison (<=, >=, or ==)?",
        )
    end

    is_vectorized, parse_code, build_call =
        parse_constraint_expr(_error, constraint_expr)
    result_variable = gensym()
    code = quote
        $parse_code
        $result_variable = $build_call
    end

    return code
end

_add_JuMP_prefix(s::Symbol) = Expr(:., JuMP, :($(QuoteNode(s))))

for (mac, sym) in [
    (:constraints, Symbol("@constraint")),
    (:NLconstraints, Symbol("@NLconstraint")),
    (:SDconstraints, Symbol("@SDconstraint")),
    (:variables, Symbol("@variable")),
    (:expressions, Symbol("@expression")),
    (:NLexpressions, Symbol("@NLexpression")),
]
    @eval begin
        macro $mac(m, x)
            if typeof(x) != Expr || x.head != :block
                # We do a weird string interpolation here so that it gets
                # interpolated at compile time, not run-time.
                error("Invalid syntax for @" * $(string(mac)))
            end
            @assert isa(x.args[1], LineNumberNode)
            lastline = x.args[1]
            code = Expr(:block)
            for it in x.args
                if isa(it, LineNumberNode)
                    lastline = it
                elseif isexpr(it, :tuple) # line with commas
                    args = []
                    # Keyword arguments have to appear like:
                    # x, (start = 10, lower_bound = 5)
                    # because of the precedence of "=".
                    for ex in it.args
                        if isexpr(ex, :tuple) # embedded tuple
                            append!(args, ex.args)
                        else
                            push!(args, ex)
                        end
                    end
                    mac = esc(
                        Expr(
                            :macrocall,
                            $(_add_JuMP_prefix(sym)),
                            lastline,
                            m,
                            args...,
                        ),
                    )
                    push!(code.args, mac)
                else # stand-alone symbol or expression
                    push!(
                        code.args,
                        esc(
                            Expr(
                                :macrocall,
                                $(_add_JuMP_prefix(sym)),
                                lastline,
                                m,
                                it,
                            ),
                        ),
                    )
                end
            end
            push!(code.args, :(nothing))
            return code
        end
    end
end

# Doc strings for the auto-generated macro pluralizations
@doc """
    @constraints(model, args...)

Adds groups of constraints at once, in the same fashion as the
[`@constraint`](@ref) macro.

The model must be the first argument, and multiple constraints can be added on
multiple lines wrapped in a `begin ... end` block.

# Examples

```julia
@constraints(model, begin
    x >= 1
    y - w <= 2
    sum_to_one[i=1:3], z[i] + y == 1
end)
```
""" :(@constraints)

@doc """
    @variables(model, args...)

Adds multiple variables to model at once, in the same fashion as the
[`@variable`](@ref) macro.

The model must be the first argument, and multiple variables can be added on
multiple lines wrapped in a `begin ... end` block.

# Examples

```julia
@variables(model, begin
    x
    y[i = 1:2] >= 0, (start = i)
    z, Bin, (start = 0, base_name = "Z")
end)
```

!!! note
    Keyword arguments must be contained within parentheses (refer to the example
    above).
""" :(@variables)

@doc """
    @expressions(model, args...)

Adds multiple expressions to model at once, in the same fashion as the
[`@expression`](@ref) macro.

The model must be the first argument, and multiple variables can be added on
multiple lines wrapped in a `begin ... end` block.

# Examples

```julia
@expressions(model, begin
    my_expr, x^2 + y^2
    my_expr_1[i = 1:2], a[i] - z[i]
end)
```
""" :(@expressions)

@doc """
    @SDconstraints(model, args...)

Adds multiple semi-definite constraints to model at once, in the same fashion as
the [`@SDconstraint`](@ref) macro.

The model must be the first argument, and multiple constraints can be added on
multiple lines wrapped in a `begin ... end` block.

# Examples

```julia
@SDconstraints(model, begin
    A * x >= b
    b - C * y >= 0
end)
```
""" :(@SDconstraints)

@doc """
    @NLconstraints(model, args...)

Adds multiple nonlinear constraints to model at once, in the same fashion as
the [`@NLconstraint`](@ref) macro.

The model must be the first argument, and multiple variables can be added on
multiple lines wrapped in a `begin ... end` block.

# Examples

```julia
@NLconstraints(model, begin
    t >= sqrt(x^2 + y^2)
    [i = 1:2], z[i] <= log(a[i])
end)
```
""" :(@NLconstraints)

@doc """
    @NLexpressions(model, args...)

Adds multiple nonlinear expressions to model at once, in the same fashion as the
[`@NLexpression`](@ref) macro.

The model must be the first argument, and multiple variables can be added on
multiple lines wrapped in a `begin ... end` block.

# Examples

```julia
@NLexpressions(model, begin
    my_expr, sqrt(x^2 + y^2)
    my_expr_1[i = 1:2], log(a[i]) - z[i]
end)
```
""" :(@NLexpressions)

"""
    _moi_sense(_error::Function, sense)

Return an expression whose value is an `MOI.OptimizationSense` corresponding
to `sense`. Sense is either the symbol `:Min` or `:Max`, corresponding
respectively to `MOI.MIN_SENSE` and `MOI.MAX_SENSE` or it is another symbol,
which should be the name of a variable or expression whose value is an
`MOI.OptimizationSense`.
In the last case, the expression throws an error using the `_error`
function in case the value is not an `MOI.OptimizationSense`.
"""
function _moi_sense(_error::Function, sense)
    if sense == :Min
        expr = MOI.MIN_SENSE
    elseif sense == :Max
        expr = MOI.MAX_SENSE
    else
        # Refers to a variable that holds the sense.
        # TODO: Better document this behavior
        expr = esc(sense)
    end
    return :(_throw_error_for_invalid_sense($_error, $expr))
end

function _throw_error_for_invalid_sense(_error::Function, sense)
    return _error(
        "Unexpected sense `$value`. The sense must be an",
        " `MOI.OptimizatonSense`, `Min` or `Max`.",
    )
end
function _throw_error_for_invalid_sense(
    _error::Function,
    sense::MOI.OptimizationSense,
)
    return sense
end

"""
    @objective(model::Model, sense, func)

Set the objective sense to `sense` and objective function to `func`. The
objective sense can be either `Min`, `Max`, `MathOptInterface.MIN_SENSE`,
`MathOptInterface.MAX_SENSE` or `MathOptInterface.FEASIBILITY_SENSE`; see
[`MathOptInterface.ObjectiveSense`](https://jump.dev/MathOptInterface.jl/v0.8/apireference.html#MathOptInterface.ObjectiveSense).
In order to set the sense programatically, i.e., when `sense` is a Julia
variable whose value is the sense, one of the three
`MathOptInterface.ObjectiveSense` values should be used. The function `func` can
be a single JuMP variable, an affine expression of JuMP variables or a quadratic
expression of JuMP variables.

## Examples

To minimize the value of the variable `x`, do as follows:
```jldoctest @objective; setup = :(using JuMP)
julia> model = Model()
A JuMP Model
Feasibility problem with:
Variables: 0
Model mode: AUTOMATIC
CachingOptimizer state: NO_OPTIMIZER
Solver name: No optimizer attached.

julia> @variable(model, x)
x

julia> @objective(model, Min, x)
x
```

To maximize the value of the affine expression `2x - 1`, do as follows:
```jldoctest @objective
julia> @objective(model, Max, 2x - 1)
2 x - 1
```

To set a quadratic objective and set the objective sense programatically, do
as follows:
```jldoctest @objective
julia> sense = MOI.MIN_SENSE
MIN_SENSE::OptimizationSense = 0

julia> @objective(model, sense, x^2 - 2x + 1)
x² - 2 x + 1
```
"""
macro objective(model, args...)
    function _error(str...)
        return _macro_error(:objective, (model, args...), __source__, str...)
    end

    # We don't overwrite `model` as it is used in `_error`
    esc_model = esc(model)
    if length(args) != 2
        # Either just an objective sense, or just an expression.
        _error(
            "needs three arguments: model, objective sense (Max or Min) and expression.",
        )
    end
    sense, x = args
    sense_expr = _moi_sense(_error, sense)
    newaff, parsecode = _MA.rewrite(x)
    code = quote
        $parsecode
        set_objective($esc_model, $sense_expr, $newaff)
        $newaff
    end
    return _finalize_macro(esc_model, code, __source__)
end

"""
    @expression(args...)

Efficiently builds a linear or quadratic expression but does not add to model
immediately. Instead, returns the expression which can then be inserted in other
constraints. For example:

```julia
@expression(m, shared, sum(i*x[i] for i=1:5))
@constraint(m, shared + y >= 5)
@constraint(m, shared + z <= 10)
```

The `ref` accepts index sets in the same way as `@variable`, and those indices
can be used in the construction of the expressions:

```julia
@expression(m, expr[i=1:3], i*sum(x[j] for j=1:3))
```

Anonymous syntax is also supported:

```julia
expr = @expression(m, [i=1:3], i*sum(x[j] for j=1:3))
```
"""
macro expression(args...)
    _error(str...) = _macro_error(:expression, args, __source__, str...)
    args, kw_args, requestedcontainer = Containers._extract_kw_args(args)
    if length(args) == 3
        m = esc(args[1])
        c = args[2]
        x = args[3]
    elseif length(args) == 2
        m = esc(args[1])
        c = gensym()
        x = args[2]
    else
        _error("needs at least two arguments.")
    end
    length(kw_args) == 0 || _error("unrecognized keyword argument")

    anonvar = isexpr(c, :vect) || isexpr(c, :vcat) || length(args) == 2
    variable = gensym()

    idxvars, indices = Containers._build_ref_sets(_error, c)
    if args[1] in idxvars
        _error(
            "Index $(args[1]) is the same symbol as the model. Use a " *
            "different name for the index.",
        )
    end
    code = _MA.rewrite_and_return(x)
    code = Containers.container_code(idxvars, indices, code, requestedcontainer)
    # don't do anything with the model, but check that it's valid anyway
    if anonvar
        macro_code = code
    else
        macro_code = _macro_assign_and_return(
            code,
            variable,
            Containers._get_name(c),
            model_for_registering = m,
        )
    end
    return _finalize_macro(m, macro_code, __source__)
end

_esc_non_constant(x::Number) = x
_esc_non_constant(x::Expr) = isexpr(x, :quote) ? x : esc(x)
_esc_non_constant(x) = esc(x)

# Returns the type of what `add_variable(::Model, build_variable(...))` would return where `...` represents the positional arguments.
# Example: `@variable m [1:3] foo` will allocate an vector of element type `variable_type(m, foo)`
# Note: it needs to be implemented by all `AbstractModel`s
variable_type(m::Model) = VariableRef

"""
    build_variable(_error::Function, info::VariableInfo; extra_kw_args...)

Returns a new variable.

Extensions should define a method with additional positional arguments to
dispatch the call to a different method. The return type should only depend on
the positional arguments for `variable_type` to make sense.

As an example, `@variable(model, x, foo)` foo will call
`build_variable(_error, info, foo)`

See the [`@variable`](@ref) macro doc for more details.
"""
function build_variable(_error::Function, info::VariableInfo; extra_kw_args...)
    for (kwarg, _) in extra_kw_args
        _error("Unrecognized keyword argument $kwarg")
    end
    return ScalarVariable(info)
end

function build_variable(_error::Function, ::VariableInfo, args...; kwargs...)
    return _error(
        "Unrecognized arguments: $(join(args, ", ")). (You may have passed " *
        "these as positional arguments, or as a keyword value to " *
        "`variable_type`.)\n\nIf you're trying to create a JuMP extension, " *
        "you need to implement `build_variable`.",
    )
end

function build_variable(
    _error::Function,
    variable::ScalarVariable,
    set::MOI.AbstractScalarSet,
)
    return VariableConstrainedOnCreation(variable, set)
end
function build_variable(
    _error::Function,
    variables::Vector{<:ScalarVariable},
    set::MOI.AbstractVectorSet,
)
    return VariablesConstrainedOnCreation(variables, set)
end

function _macro_error(macroname, args, source, str...)
    return error(
        "At $(source.file):$(source.line): " *
        "`@$macroname($(join(args, ", ")))`: ",
        str...,
    )
end

# Given a base_name and idxvars, returns an expression that constructs the name
# of the object. For use within macros only.
function _name_call(base_name, idxvars)
    if isempty(idxvars) || base_name == ""
        return base_name
    end
    ex = Expr(:call, :string, base_name, "[")
    for i in 1:length(idxvars)
        # Converting the arguments to strings before concatenating is faster:
        # https://github.com/JuliaLang/julia/issues/29550.
        esc_idxvar = esc(idxvars[i])
        push!(ex.args, :(string($esc_idxvar)))
        i < length(idxvars) && push!(ex.args, ",")
    end
    push!(ex.args, "]")
    return ex
end

reverse_sense(::Val{:<=}) = Val(:>=)
reverse_sense(::Val{:≤}) = Val(:≥)
reverse_sense(::Val{:>=}) = Val(:<=)
reverse_sense(::Val{:≥}) = Val(:≤)
reverse_sense(::Val{:(==)}) = Val(:(==))

"""
    parse_one_operator_variable(_error::Function, infoexpr::_VariableInfoExpr, sense::Val{S}, value) where S

Update `infoexr` for a variable expression in the `@variable` macro of the form `variable name S value`.
"""
function parse_one_operator_variable end

function parse_one_operator_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Val{:in},
    set,
)
    return set
end
function parse_one_operator_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Union{Val{:<=},Val{:≤}},
    upper,
)
    _set_upper_bound_or_error(_error, infoexpr, upper)
    return
end
function parse_one_operator_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Union{Val{:>=},Val{:≥}},
    lower,
)
    _set_lower_bound_or_error(_error, infoexpr, lower)
    return
end
function parse_one_operator_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Val{:(==)},
    value,
)
    _fix_or_error(_error, infoexpr, value)
    return
end
function parse_one_operator_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Val{S},
    value,
) where {S}
    return _error("Unknown sense $S.")
end

# There is not way to determine at parsing time which of lhs or rhs is the
# variable name and which is the value if both are symbols. For instance,
# lhs could be the Symbol `:x` and rhs could be the Symbol `:a` where a
# variable `a` is assigned to 1 in the local scope. Knowing this, we know
# that `x` is the variable name but at parse time there is now way to know
# that `a` has a value.
# In that case we assume the variable is the lhs.
function parse_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    sense::Symbol,
    var,
    value,
)
    set = parse_one_operator_variable(
        _error,
        infoexpr,
        Val(sense),
        _esc_non_constant(value),
    )
    return var, set
end

# If the lhs is a number and not the rhs, we can deduce that the rhs is
# the variable.
function parse_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    sense::Symbol,
    value::Number,
    var,
)
    set = parse_one_operator_variable(
        _error,
        infoexpr,
        reverse_sense(Val(sense)),
        _esc_non_constant(value),
    )
    return var, set
end

function parse_ternary_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Union{Val{:<=},Val{:≤}},
    lower,
    ::Union{Val{:<=},Val{:≤}},
    upper,
)
    _set_lower_bound_or_error(_error, infoexpr, lower)
    _set_upper_bound_or_error(_error, infoexpr, upper)
    return
end
function parse_ternary_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Union{Val{:>=},Val{:≥}},
    upper,
    ::Union{Val{:>=},Val{:≥}},
    lower,
)
    return parse_ternary_variable(
        _error,
        infoexpr,
        Val(:≤),
        lower,
        Val(:≤),
        upper,
    )
end
function parse_ternary_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    ::Val,
    lvalue,
    ::Val,
    rvalue,
)
    return _error("Use the form lb <= ... <= ub.")
end
function parse_variable(
    _error::Function,
    infoexpr::_VariableInfoExpr,
    lvalue,
    lsign::Symbol,
    var,
    rsign::Symbol,
    rvalue,
)
    # lvalue lsign var rsign rvalue
    set = parse_ternary_variable(
        _error,
        infoexpr,
        Val(lsign),
        _esc_non_constant(lvalue),
        Val(rsign),
        _esc_non_constant(rvalue),
    )
    return var, set
end

function _reorder_parameters(args)
    if !isexpr(args[1], :parameters)
        return args
    end
    args = collect(args)
    p = popfirst!(args)
    for arg in p.args
        @assert arg.head == :kw
        push!(args, Expr(:(=), arg.args[1], arg.args[2]))
    end
    return args
end

"""
    @variable(model, kw_args...)

Add an *anonymous* variable to the model `model` described by the keyword
arguments `kw_args` and returns the variable.

    @variable(model, expr, args..., kw_args...)

Add a variable to the model `model` described by the expression `expr`, the
positional arguments `args` and the keyword arguments `kw_args`. The expression
`expr` can either be (note that in the following the symbol `<=` can be used
instead of `≤` and the symbol `>=`can be used instead of `≥`)

* of the form `varexpr` creating variables described by `varexpr`;
* of the form `varexpr ≤ ub` (resp. `varexpr ≥ lb`) creating variables described by
  `varexpr` with upper bounds given by `ub` (resp. lower bounds given by `lb`);
* of the form `varexpr == value` creating variables described by `varexpr` with
  fixed values given by `value`; or
* of the form `lb ≤ varexpr ≤ ub` or `ub ≥ varexpr ≥ lb` creating variables
  described by `varexpr` with lower bounds given by `lb` and upper bounds given
  by `ub`.
* of the form `varexpr in set` creating variables described by
  `varexpr` constrained to belong to `set`, see [Variables constrained on creation](@ref).

The expression `varexpr` can either be

* of the form `varname` creating a scalar real variable of name `varname`;
* of the form `varname[...]` or `[...]` creating a container of variables.

The recognized positional arguments in `args` are the following:

* `Bin`: Sets the variable to be binary, i.e. either 0 or 1.
* `Int`: Sets the variable to be integer, i.e. one of ..., -2, -1, 0, 1, 2, ...
* `Symmetric`: Only available when creating a square matrix of variables, i.e.
  when `varexpr` is of the form `varname[1:n,1:n]` or `varname[i=1:n,j=1:n]`.
  It creates a symmetric matrix of variable, that is, it only creates a
  new variable for `varname[i,j]` with `i ≤ j` and sets `varname[j,i]` to the
  same variable as `varname[i,j]`. It is equivalent to using
  `varexpr in SymMatrixSpace()` as `expr`.
* `PSD`: The square matrix of variable is both `Symmetric` and constrained to be
  positive semidefinite. It is equivalent to using `varexpr in PSDCone()` as
  `expr`.

The recognized keyword arguments in `kw_args` are the following:

* `base_name`: Sets the name prefix used to generate variable names. It
  corresponds to the variable name for scalar variable, otherwise, the
  variable names are set to `base_name[...]` for each index `...` of the axes
  `axes`.
* `lower_bound`: Sets the value of the variable lower bound.
* `upper_bound`: Sets the value of the variable upper bound.
* `start`: Sets the variable starting value used as initial guess in optimization.
* `binary`: Sets whether the variable is binary or not.
* `integer`: Sets whether the variable is integer or not.
* `variable_type`: See the "Note for extending the variable macro" section below.
* `set`: Equivalent to using `varexpr in value` as `expr` where `value` is the
  value of the keyword argument.
* `container`: Specify the container type.

## Examples

The following are equivalent ways of creating a variable `x` of name `x` with
lower bound 0:
```julia
# Specify everything in `expr`
@variable(model, x >= 0)
# Specify the lower bound using a keyword argument
@variable(model, x, lower_bound=0)
# Specify everything in `kw_args`
x = @variable(model, base_name="x", lower_bound=0)
```

The following are equivalent ways of creating a `DenseAxisArray` of index set
`[:a, :b]` and with respective upper bounds 2 and 3 and names `x[a]` and `x[b]`.
The upper bound can either be specified in `expr`:
```jldoctest variable_macro; setup = :(using JuMP; model = Model())
ub = Dict(:a => 2, :b => 3)
@variable(model, x[i=keys(ub)] <= ub[i])

# output
1-dimensional DenseAxisArray{VariableRef,1,...} with index sets:
    Dimension 1, Symbol[:a, :b]
And data, a 2-element Array{VariableRef,1}:
 x[a]
 x[b]
```
or it can be specified with the `upper_bound` keyword argument:
```jldoctest variable_macro
@variable(model, y[i=keys(ub)], upper_bound=ub[i])

# output
1-dimensional DenseAxisArray{VariableRef,1,...} with index sets:
    Dimension 1, Symbol[:a, :b]
And data, a 2-element Array{VariableRef,1}:
 y[a]
 y[b]
```

## Note for extending the variable macro

The single scalar variable or each scalar variable of the container are created
using `add_variable(model, build_variable(_error, info, extra_args...;
extra_kw_args...))` where

* `model` is the model passed to the `@variable` macro;
* `_error` is an error function with a single `String` argument showing the
  `@variable` call in addition to the error message given as argument;
* `info` is the `VariableInfo` struct containing the information gathered in
  `expr`, the recognized keyword arguments (except `base_name` and
  `variable_type`) and the recognized positional arguments (except `Symmetric`
  and `PSD`);
* `extra_args` are the unrecognized positional arguments of `args` plus the
  value of the `variable_type` keyword argument if present. The `variable_type`
  keyword argument allows the user to pass a position argument to
  `build_variable` without the need to give a positional argument to
  `@variable`. In particular, this allows the user to give a positional
  argument to the `build_variable` call when using the anonymous single variable
  syntax `@variable(model, kw_args...)`; and
* `extra_kw_args` are the unrecognized keyword argument of `kw_args`.

## Examples

The following creates a variable `x` of name `x` with `lower_bound` 0 as with the first
example above but does it without using the `@variable` macro
```julia
info = VariableInfo(true, 0, false, NaN, false, NaN, false, NaN, false, false)
JuMP.add_variable(model, JuMP.build_variable(error, info), "x")
```

The following creates a `DenseAxisArray` of index set `[:a, :b]` and with respective
upper bounds 2 and 3 and names `x[a]` and `x[b]` as with the second example
above but does it without using the `@variable` macro
```jldoctest variable_macro
# Without the `@variable` macro
x = JuMP.Containers.container(i -> begin
        info = VariableInfo(false, NaN, true, ub[i], false, NaN, false, NaN, false, false)
        x[i] = JuMP.add_variable(model, JuMP.build_variable(error, info), "x[\$i]")
    end, JuMP.Containers.vectorized_product(keys(ub)))

# output
1-dimensional DenseAxisArray{VariableRef,1,...} with index sets:
    Dimension 1, Symbol[:a, :b]
And data, a 2-element Array{VariableRef,1}:
 x[a]
 x[b]
```

The following are equivalent ways of creating a `Matrix` of size
`N x N` with variables custom variables created with a JuMP extension using
the `Poly(X)` positional argument to specify its variables:
```julia
# Using the `@variable` macro
@variable(model, x[1:N,1:N], Symmetric, Poly(X))
# Without the `@variable` macro
x = Matrix{JuMP.variable_type(model, Poly(X))}(N, N)
info = VariableInfo(false, NaN, false, NaN, false, NaN, false, NaN, false, false)
for i in 1:N, j in i:N
    x[i,j] = x[j,i] = JuMP.add_variable(model, build_variable(error, info, Poly(X)), "x[\$i,\$j]")
end
```
"""
macro variable(args...)
    _error(str...) = _macro_error(:variable, args, __source__, str...)
    # We need to re-order the parameters here to account for cases like
    # `@variable(model; integer = true)`, since Julia handles kwargs by placing
    # them first(!) in the list of arguments.
    args = _reorder_parameters(args)
    model = esc(args[1])

    extra, kw_args, requestedcontainer =
        Containers._extract_kw_args(args[2:end])

    # if there is only a single non-keyword argument, this is an anonymous
    # variable spec and the one non-kwarg is the model
    if length(extra) == 0
        x = gensym()
        anon_singleton = true
    else
        x = popfirst!(extra)
        if x == :Int
            _error(
                "Ambiguous variable name $x detected. To specify an anonymous integer " *
                "variable, use `@variable(model, integer = true)` instead.",
            )
        elseif x == :Bin
            _error(
                "Ambiguous variable name $x detected. To specify an anonymous binary " *
                "variable, use `@variable(model, binary = true)` instead.",
            )
        elseif x == :PSD
            _error(
                "Size of anonymous square matrix of positive semidefinite anonymous variables is not specified. To specify size of square matrix " *
                "use `@variable(model, [1:n, 1:n], PSD)` instead.",
            )
        end
        anon_singleton = false
    end

    info_kw_args = filter(_is_info_keyword, kw_args)
    extra_kw_args = filter(
        kw -> begin
            kw.args[1] != :base_name &&
                kw.args[1] != :variable_type &&
                kw.args[1] != :set &&
                !_is_info_keyword(kw)
        end,
        kw_args,
    )
    base_name_kw_args = filter(kw -> kw.args[1] == :base_name, kw_args)
    variable_type_kw_args = filter(kw -> kw.args[1] == :variable_type, kw_args)
    set_kw_args = filter(kw -> kw.args[1] == :set, kw_args)
    infoexpr = _VariableInfoExpr(; _keywordify.(info_kw_args)...)

    # There are four cases to consider:
    # x                                       | type of x | x.head
    # ----------------------------------------+-----------+------------
    # var                                     | Symbol    | NA
    # var[1:2]                                | Expr      | :ref
    # var <= ub or var[1:2] <= ub             | Expr      | :call
    # var in set or var[1:2] in set           | Expr      | :call
    # lb <= var <= ub or lb <= var[1:2] <= ub | Expr      | :comparison
    # In the three last cases, we call parse_variable
    explicit_comparison = isexpr(x, :comparison) || isexpr(x, :call)
    if explicit_comparison
        var, set = parse_variable(_error, infoexpr, x.args...)
    else
        var = x
        set = nothing
    end

    anonvar = isexpr(var, :vect) || isexpr(var, :vcat) || anon_singleton
    if anonvar && explicit_comparison && set === nothing
        _error(
            "Cannot use explicit bounds via >=, <= with an anonymous variable",
        )
    end
    variable = gensym()
    # TODO: Should we generate non-empty default names for variables?
    name = Containers._get_name(var)
    if isempty(base_name_kw_args)
        base_name = anonvar ? "" : string(name)
    else
        base_name = esc(base_name_kw_args[1].args[2])
    end

    if !isa(name, Symbol) && !anonvar
        Base.error(
            "Expression $name should not be used as a variable name. Use the \"anonymous\" syntax $name = @variable(model, ...) instead.",
        )
    end

    if !isempty(set_kw_args)
        if length(set_kw_args) > 1
            _error(
                "`set` keyword argument was given $(length(set_kw_args)) times.",
            )
        end
        if set !== nothing
            _error(
                "Cannot specify set twice, it was already set to `$set` so the `set` keyword argument is not allowed.",
            )
        end
        set = esc(set_kw_args[1].args[2])
    end

    # process keyword arguments
    if any(t -> (t == :PSD), extra)
        if set !== nothing
            _error(
                "Cannot specify set twice, it was already set to `$set` so the `PSD` argument is not allowed.",
            )
        end
        set = :(JuMP.PSDCone())
    end
    if any(t -> (t == :Symmetric), extra)
        if set !== nothing
            _error(
                "Cannot specify `Symmetric` when the set is already specified, the variable is constrained to belong to `$set`.",
            )
        end
        set = :(JuMP.SymMatrixSpace())
    end
    extra = filter(x -> (x != :PSD && x != :Symmetric), extra) # filter out PSD and sym tag
    for ex in extra
        if ex == :Int
            _set_integer_or_error(_error, infoexpr)
        elseif ex == :Bin
            _set_binary_or_error(_error, infoexpr)
        end
    end
    extra = esc.(filter(ex -> !(ex in [:Int, :Bin]), extra))
    if !isempty(variable_type_kw_args)
        push!(extra, esc(variable_type_kw_args[1].args[2]))
    end

    info = _constructor_expr(infoexpr)
    if isa(var, Symbol)
        # Easy case - a single variable
        name_code = base_name
    else
        isa(var, Expr) || _error("Expected $var to be a variable name")
        # We now build the code to generate the variables (and possibly the
        # SparseAxisArray to contain them)
        idxvars, indices = Containers._build_ref_sets(_error, var)
        if args[1] in idxvars
            _error(
                "Index $(args[1]) is the same symbol as the model. Use a " *
                "different name for the index.",
            )
        end
        name_code = _name_call(base_name, idxvars)
        if set !== nothing
            name_code = Containers.container_code(
                idxvars,
                indices,
                name_code,
                requestedcontainer,
            )
        end
    end

    # Code to be used to create each variable of the container.
    buildcall = :(build_variable($_error, $info, $(extra...)))
    _add_kw_args(buildcall, extra_kw_args)
    if set !== nothing
        if isa(var, Symbol)
            scalar_variables = buildcall
        else
            scalar_variables = Containers.container_code(
                idxvars,
                indices,
                buildcall,
                requestedcontainer,
            )
        end
        buildcall = :(build_variable($_error, $scalar_variables, $set))
    end

    variablecall = :(add_variable($model, $buildcall, $name_code))
    if isa(var, Symbol) || set !== nothing
        # The looped code is trivial here since there is a single variable
        creation_code = variablecall
    else
        creation_code = Containers.container_code(
            idxvars,
            indices,
            variablecall,
            requestedcontainer,
        )
    end

    if anonvar
        # Anonymous variable, no need to register it in the model-level
        # dictionary nor to assign it to a variable in the user scope.
        # We simply return the variable
        macro_code = creation_code
    else
        # We register the variable reference to its name and
        # we assign it to a variable in the local scope of this name
        macro_code = _macro_assign_and_return(
            creation_code,
            variable,
            name,
            model_for_registering = model,
        )
    end
    return _finalize_macro(model, macro_code, __source__)
end

"""
    @NLobjective(model, sense, expression)

Add a nonlinear objective to `model` with optimization sense `sense`.
`sense` must be `Max` or `Min`.

# Example

    @NLobjective(model, Max, 2x + 1 + sin(x))
"""
macro NLobjective(model, sense, x)
    function _error(str...)
        return _macro_error(:NLobjective, (model, sense, x), __source__, str...)
    end
    sense_expr = _moi_sense(_error, sense)
    ex = gensym()
    code = quote
        $ex = $(_process_NL_expr(model, x))
        set_objective($(esc(model)), $sense_expr, $ex)
    end
    return _finalize_macro(esc(model), code, __source__)
end

"""
    @NLconstraint(m::Model, expr)

Add a constraint described by the nonlinear expression `expr`. See also
[`@constraint`](@ref). For example:

```julia
@NLconstraint(model, sin(x) <= 1)
@NLconstraint(model, [i = 1:3], sin(i * x) <= 1 / i)
```
"""
macro NLconstraint(m, x, args...)
    function _error(str...)
        return _macro_error(:NLconstraint, (m, x, args...), __source__, str...)
    end
    esc_m = esc(m)
    # Two formats:
    # - @NLconstraint(m, a*x <= 5)
    # - @NLconstraint(m, myref[a=1:5], sin(x^a) <= 5)
    extra, kw_args, requestedcontainer = Containers._extract_kw_args(args)
    (length(extra) > 1 || length(kw_args) > 0) && _error("too many arguments.")
    # Canonicalize the arguments
    c = length(extra) == 1 ? x : gensym()
    con = length(extra) == 1 ? extra[1] : x

    anonvar = isexpr(c, :vect) || isexpr(c, :vcat) || length(extra) != 1
    variable = gensym()

    # Strategy: build up the code for non-macro add_constraint, and if needed
    # we will wrap in loops to assign to the ConstraintRefs
    idxvars, indices = Containers._build_ref_sets(_error, c)
    if m in idxvars
        _error(
            "Index $(m) is the same symbol as the model. Use a different " *
            "name for the index.",
        )
    end
    # Build the constraint
    if isexpr(con, :call) # one-sided constraint
        # Simple comparison - move everything to the LHS
        op = con.args[1]
        if op == :(==)
            lb = 0.0
            ub = 0.0
        elseif op == :(<=) || op == :(≤)
            lb = -Inf
            ub = 0.0
        elseif op == :(>=) || op == :(≥)
            lb = 0.0
            ub = Inf
        else
            _error("expected comparison operator (<=, >=, or ==).")
        end
        lhs = :($(con.args[2]) - $(con.args[3]))
        code = quote
            c = _NonlinearConstraint($(_process_NL_expr(m, lhs)), $lb, $ub)
            push!($esc_m.nlp_data.nlconstr, c)
            ConstraintRef(
                $esc_m,
                NonlinearConstraintIndex(length($esc_m.nlp_data.nlconstr)),
                ScalarShape(),
            )
        end
    elseif isexpr(con, :comparison)
        # ranged row
        if (con.args[2] != :<= && con.args[2] != :≤) ||
           (con.args[4] != :<= && con.args[4] != :≤)
            _error(
                "only ranged rows of the form lb <= expr <= ub are supported.",
            )
        end
        lb = con.args[1]
        ub = con.args[5]
        code = quote
            if !isa($(esc(lb)), Number)
                _error("expected ", $(string(lb)), " to be a number.")
            elseif !isa($(esc(ub)), Number)
                _error("expected ", $(string(ub)), " to be a number.")
            end
            c = _NonlinearConstraint(
                $(_process_NL_expr(m, con.args[3])),
                $(esc(lb)),
                $(esc(ub)),
            )
            push!($esc_m.nlp_data.nlconstr, c)
            ConstraintRef(
                $esc_m,
                NonlinearConstraintIndex(length($esc_m.nlp_data.nlconstr)),
                ScalarShape(),
            )
        end
    else
        # Unknown
        _error(
            "constraints must be in one of the following forms:\n" *
            "       expr1 <= expr2\n" *
            "       expr1 >= expr2\n" *
            "       expr1 == expr2",
        )
    end
    looped =
        Containers.container_code(idxvars, indices, code, requestedcontainer)
    creation_code = quote
        _init_NLP($esc_m)
        $looped
    end
    if anonvar
        macro_code = creation_code
    else
        macro_code = _macro_assign_and_return(
            creation_code,
            variable,
            Containers._get_name(c),
            model_for_registering = esc_m,
        )
    end
    return _finalize_macro(esc_m, macro_code, __source__)
end

"""
    @NLexpression(args...)

Efficiently build a nonlinear expression which can then be inserted in other
nonlinear constraints and the objective. See also [`@expression`]. For example:

```julia
@NLexpression(model, my_expr, sin(x)^2 + cos(x^2))
@NLconstraint(model, my_expr + y >= 5)
@NLobjective(model, Min, my_expr)
```

Indexing over sets and anonymous expressions are also supported:
```julia
@NLexpression(m, my_expr_1[i=1:3], sin(i * x))
my_expr_2 = @NLexpression(m, log(1 + sum(exp(x[i])) for i in 1:2))
```
"""
macro NLexpression(args...)
    _error(str...) = _macro_error(:NLexpression, args, __source__, str...)
    args, kw_args, requestedcontainer = Containers._extract_kw_args(args)
    if length(args) <= 1
        _error(
            "To few arguments ($(length(args))); must pass the model and nonlinear expression as arguments.",
        )
    elseif length(args) == 2
        m, x = args
        c = gensym()
    elseif length(args) == 3
        m, c, x = args
    end
    if length(args) > 3 || length(kw_args) > 0
        _error("To many arguments ($(length(args))).")
    end

    anonvar = isexpr(c, :vect) || isexpr(c, :vcat) || length(args) == 2
    variable = gensym()

    idxvars, indices = Containers._build_ref_sets(_error, c)
    if args[1] in idxvars
        _error(
            "Index $(args[1]) is the same symbol as the model. Use a " *
            "different name for the index.",
        )
    end
    code = :(NonlinearExpression($(esc(m)), $(_process_NL_expr(m, x))))
    creation_code =
        Containers.container_code(idxvars, indices, code, requestedcontainer)
    if anonvar
        macro_code = creation_code
    else
        macro_code = _macro_assign_and_return(
            creation_code,
            variable,
            Containers._get_name(c),
            model_for_registering = esc(m),
        )
    end
    return _finalize_macro(esc(m), macro_code, __source__)
end

"""
    @NLparameter(model, param == value)

Create and return a nonlinear parameter `param` attached to the model `model`
with initial value set to `value`. Nonlinear parameters may be used only in
nonlinear expressions.

# Example
```jldoctest; setup=:(using JuMP)
model = Model()
@NLparameter(model, x == 10)
value(x)

# output
10.0
```

    @NLparameter(model, param_collection[...] == value_expr)

Create and return a collection of nonlinear parameters `param_collection`
attached to the model `model` with initial value set to `value_expr` (may
depend on index sets).
Uses the same syntax for specifying index sets as [`@variable`](@ref).

# Example
```jldoctest; setup=:(using JuMP)
model = Model()
@NLparameter(model, y[i = 1:10] == 2 * i)
value(y[9])

# output
18.0
```
"""
macro NLparameter(m, ex, extra...)
    function _error(str...)
        return _macro_error(:NLparameter, (m, ex, extra...), __source__, str...)
    end

    extra, kw_args, requestedcontainer = Containers._extract_kw_args(extra)
    (length(extra) == 0 && length(kw_args) == 0) ||
        _error("Too many arguments.")
    if !isexpr(ex, :call) || length(ex.args) != 3 || ex.args[1] != :(==)
        _error("Syntax error.")
    end
    c = ex.args[2]
    x = ex.args[3]
    anonvar = isexpr(c, :vect) || isexpr(c, :vcat)
    if anonvar
        _error(
            "Anonymous nonlinear parameter syntax is not currently supported.",
        )
    end
    esc_m = esc(m)
    variable = gensym()

    idxvars, indices = Containers._build_ref_sets(_error, c)
    if m in idxvars
        _error(
            "Index $(m) is the same symbol as the model. Use a different " *
            "name for the index.",
        )
    end
    code = quote
        if !isa($(esc(x)), Number)
            _error("Expected ", $(string(x)), " to be a number.")
        end
        _new_parameter($esc_m, $(esc(x)))
    end
    creation_code = Containers.container_code(idxvars, indices, code, :Auto)

    # TODO: NLparameters are not registered in the model because we don't yet
    # have an anonymous version.
    macro_code = _macro_assign_and_return(
        creation_code,
        variable,
        Containers._get_name(c),
    )
    return _finalize_macro(esc_m, macro_code, __source__)
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# MLP: Convenience type for an array of layers
# This is deprecated, please use the new Net().

typealias MLP Vector{Op}

forw(n::MLP, x; o...)=(for l in n; x=forw(l, x; o...); end; x)
back(n::MLP, dy; returndx=false, o...)=(for i=length(n):-1:1; dy=back(n[i],dy; returndx=(i>1||returndx), o...); end; dy)
loss(n::MLP, dy; y=n[end].y)=loss(n[end], dy; y=y)
update!(n::MLP; o...)=(for l in n; update!(l; o...); end)

function params(r::MLP)
    p = Any[]
    for o in r
        append!(p, params(o))
    end
    return p
end

# All of this should be inherited from Model() and/or implemented with the new Net():
# TODO: batching needs to be done outside?
# TODO: predict with y=nothing
# TODO: train with shuffle, iters
# TODO: strip, savenet, loadnet implement for the new Net().
# TODO: accuracy needs to go somewhere

# The backprop algorithm

function backprop(net::MLP, x, y; o...)
    forw(net, x; o...) # calculate network output given input x
    back(net, y; o...) # calculate derivatives dx,dw given desired output y
end

# Predict implements forw with minibatches.

function predict(net::MLP, x; y=nothing, batch=128, o...)
    ninst = size(x, ndims(x))
    (batch == 0 || batch > ninst) && (batch = ninst)
    (xx,yy) = (xbatch(x, batch), nothing)
    gpu() && (gpumem() < (1<<28)) && gc() # need this until julia triggers gc() when gpumem is low
    for b = 1:batch:ninst
        e  = min(ninst, b + batch - 1)
        xx = cslice!(xx, x, b:e) # 1114
        yy = forw(net, xx; train=false, o...) # 11587
        (y == nothing) && (y = dsimilar(x, (clength(yy), ccount(x))))
        y = ccopy!(y, b, yy)
    end
    return y
end

# Train implements backprop with updates and minibatches.
# It runs for one epoch by default, iters can be specified to stop earlier.

function train(net::MLP, x, y; batch=128, shuffle=false, iters=0, o...)
    shuffle && ((x,y)=shufflexy!(x,y))
    ninst = ccount(x)
    ninst==0 && (return warn("No instances"))
    (batch == 0 || batch > ninst) && (batch = ninst)
    (xx,yy) = (xbatch(x, batch), ybatch(y, batch))
    for b = 1:batch:ninst
        e = min(ninst, b + batch - 1)
        (xx,yy) = (cslice!(xx, x, b:e), cslice!(yy, y, b:e))
        backprop(net, xx, yy; o...)
        update!(net; o...)
        (iters > 0) && (e/batch >= iters) && break
        gpu() && (gpumem() < (1<<28)) && gc() # need this until julia triggers gc() when gpumem is low
    end
    # strip!(net) # this is only for one epoch, why do I strip here?
end

# To shuffle data before each epoch:

function shufflexy!(x,y)
    nx = size(x, ndims(x))
    ny = size(y, ndims(y))
    @assert nx == ny
    r = randperm(nx)
    x = x[map(n->1:n,size(x)[1:end-1])...,r]
    y = y[map(n->1:n,size(y)[1:end-1])...,r]
    return (x,y)
end

# X batches preserve the sparsity of the input, they use the KU
# versions for resizeability (cslice!).

xbatch(x,b)=(issparse(x) ?
             (gpu() ? CudaSparseMatrixCSC(spzeros(eltype(x), size(x)...)) :
              spzeros(eltype(x), size(x)...)) :
             KUdense(barray(), eltype(x), csize(x,b)))

# Y batches are always dense, because Y should be always dense.  We
# use KUdense for resizeability (cslice!):

ybatch(y,b)=KUdense(barray(), eltype(y), csize(y,b))

# Minibatches get created on GPU if gpu() is true:

barray()=(gpu()?CudaArray:Array)

function strip!(l::Op)
    for f in fieldnames(l)
        isdefined(l,f) || continue
        isa(l.(f), KUparam) && strip!(l.(f))
        in(f, (:x, :x2, :y, :dx, :dy, :xdrop)) && (l.(f)=nothing)
    end
    return l
end

strip!(p::KUparam)=(p.init=p.diff=nothing;p)
strip!(n::MLP)=(for l in n; strip!(l); end; gc(); n)

using JLD

function savenet(filename::String, net::MLP)
    net = strip!(net)
    GPU && (net = cpucopy(net))
    save(filename, "kunet", net)
end

function loadnet(filename::String)
    net = load(filename, "kunet")
    net = strip!(net)
    gpu() ? gpucopy(net) : net
end

accuracy(y,z)=mean(findmax(convert(Array,y),1)[2] .== findmax(convert(Array,z),1)[2])

ext: jl


text:
function apply(x::AbstractArray, dims, fun::Function)
  mapslices(fun, x, dims = dims)
end

export apply


text:
abstract type Rule end

"""
    RealPolitics()

Nobody ever changes their mind.
"""
struct RealPolitics <: Rule
    update!::Function
    function RealPolitics()
        new(x -> x)
    end
end

"""
    BoundedConfidence(; bc::AbstractFloat, inertia::AbstractFloat)

Update rule for a `Meeting` based on bounded confidence model, but with additional `inertia`.
"""
struct BoundedConfidence <: Rule
    update!::Function
    function BoundedConfidence(; bc::AbstractFloat, inertia::AbstractFloat)
        function update!(meeting::Meeting; bc = bc, inertia = inertia)
            for topic in 1:length(meeting.participants[1].opinions)
                for _ in 1:100
                    negotiators = StatsBase.sample(meeting.participants, 2)
                    # TODO: refactor if more negotiators in consideration
                    if abs(negotiators[1].opinions[topic] - negotiators[2].opinions[topic]) <= bc
                        negotiators_opinions = [agent.opinions[topic] for agent in negotiators]
                        for (i, agent) in enumerate(negotiators)
                            w = ones(Float64, length(negotiators_opinions))
                            w[i] = inertia
                            agent.opinions[topic] = StatsBase.mean(negotiators_opinions, StatsBase.weights(w))
                        end
                    end
                end
            end
            return meeting
        end
        new(update!)
    end
end

"""
    ContinuousHomophily(; inertia::AbstractFloat)

Update rule based on homophily principle with `inertia` parameter.
"""
struct ContinuousHomophily <: Rule
    update!::Function
    function ContinuousHomophily(; inertia::AbstractFloat)
        function update!(meeting::Meeting; inertia=inertia)
            for topic in 1:length(meeting.participants[1].opinions)
                for _ in 1:100
                    negotiators = StatsBase.sample(meeting.participants, 2)
                    # TODO: refactor if more negotiators in consideration
                    negotiators_opinions = [agent.opinions[topic] for agent in negotiators]
                    for (i, agent) in enumerate(negotiators)
                        w = ones(Float64, length(negotiators_opinions))
                        w[i] = inertia * abs(negotiators[1].opinions[topic] - negotiators[2].opinions[topic])
                        agent.opinions[topic] = StatsBase.mean(negotiators_opinions, StatsBase.weights(w))
                    end
                end
            end
            return meeting
        end
        new(update!)
    end
end


text:
# script to parse CUDA headers and generate Julia wrappers

using Crayons

#
# Parsing
#

using Clang

function wrap(name, headers...; wrapped_headers=headers, library="lib$name", defines=[], include_dirs=[])
    clang_args = String[]
    append!(clang_args, map(dir->"-I$dir", include_dirs))
    for define in defines
        if isa(define, Pair)
            append!(clang_args, ["-D", "$(first(define))=$(last(define))"])
        else
            append!(clang_args, ["-D", "$define"])
        end
    end

    output_file = "lib$(name).jl"
    common_file = "lib$(name)_common.jl"

    # only wrap headers that are wanted, and make sure we only process each header once
    included_from = Dict()
    function wrap_header(root, current)
        return any(header->endswith(current, header), wrapped_headers) &&
               get!(included_from, current, root) == root
    end

    context = init(;
                    headers = [headers...],
                    output_file = output_file,
                    common_file = common_file,
                    clang_includes = [include_dirs..., CLANG_INCLUDE],
                    clang_args = clang_args,
                    header_wrapped = wrap_header,
                    header_library = x->library,
                    clang_diagnostics = true,
                  )
    run(context)

    return output_file, common_file
end

#
# Pointer type database
#

using JSON, DataStructures

const db_path = joinpath(@__DIR__, "pointers.json")

function load_db()
    global db
    db = if isfile(db_path)
        JSON.parsefile(db_path; dicttype=DataStructures.OrderedDict)
    else
        Dict{String, Any}()
    end
end

function save_db()
    global db
    open(db_path, "w") do io
        JSON.print(io, db, 4)
    end
end

#
# Fixing-up
#

using CSTParser, Tokenize

## pass infrastructure

struct Edit{T}
    loc::T
    text::String
end

function pass(x, state, f = (x, state)->nothing)
    f(x, state)
    if x.args isa Vector
        for a in x.args
            pass(a, state, f)
        end
    else
        state.offset += x.fullspan
    end
    state
end

function apply(text, edit::Edit{Int})
    string(text[1:edit.loc], edit.text, text[nextind(text, edit.loc):end])
end
function apply(text, edit::Edit{UnitRange{Int}})
    # println("Rewriting '$(text[edit.loc])' to '$(edit.text)'")
    string(text[1:prevind(text, first(edit.loc))], edit.text, text[nextind(text, last(edit.loc)):end])
end

## rewrite passes

mutable struct State
    offset::Int
    edits::Vector{Edit}
end

# insert `@checked` before each function with a `ccall` returning a checked type`
checked_types = [
    "CUresult",
    "CUptiResult",
    "NVPA_Status",
    "nvmlReturn_t",
    "cublasStatus_t",
    "cudnnStatus_t",
    "cufftResult",
    "curandStatus_t",
    "cusolverStatus_t",
    "cusparseStatus_t",
    "cutensorStatus_t",
]
function insert_check_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.FunctionDef
        _, def, body, _ = x.args
        @assert body isa CSTParser.EXPR && body.typ == CSTParser.Block
        @assert length(body.args) == 1

        # Clang.jl-generated ccalls should be directly part of a function definition
        call = body.args[1]
        @assert call isa CSTParser.EXPR && call.typ == CSTParser.Call && call.args[1].val == "ccall"

        # get the ccall return type
        rv = call.args[5]

        if rv.val in checked_types
            push!(state.edits, Edit(state.offset, "@checked "))
        end
    end
end

# rewrite ordinary `ccall`s to `@runtime_ccall`
function rewrite_ccall_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.Call && x.args[1].val == "ccall"
        push!(state.edits, Edit(state.offset, "@runtime_"))

        # rewrite the function handle `libcuda` into a call `libcuda()`
        # FIXME: this used to work by setting `libname="lib$name()"`,
        #        but starting with Julia 1.4 this resulted in `var"libcuda()"`
        handle = x.args[3]
        if handle.typ == CSTParser.TupleH
            fun, lib = handle.args[2], handle.args[4]
            @assert lib.typ == CSTParser.IDENTIFIER

            offset = state.offset + sum(x->x.fullspan, x.args[1:2]) +
                                    sum(x->x.fullspan, handle.args[1:4])
            push!(state.edits, Edit(offset, "()"))
        end
    end
end

# insert `initialize_api()` before each function with a `ccall` calling non-whitelisted fns
preinit_apicalls = Set{String}([
    # CUDAdrv
    ## error handling
    "cuGetErrorString",
    "cuGetErrorName",
    ## initialization
    "cuInit",
    ## version management
    "cuDriverGetVersion",
    ## device management
    "cuDeviceGet",
    "cuDeviceGetAttribute",
    "cuDeviceGetCount",
    "cuDeviceGetLuid",
    "cuDeviceGetName",
    "cuDeviceGetUuid",
    "cuDeviceTotalMem",
    "cuDeviceGetProperties",     # deprecated
    "cuDeviceComputeCapability", # deprecated
    ## context management
    "cuCtxCreate",
    "cuCtxDestroy",
    "cuCtxSetCurrent",
    "cuCtxGetCurrent",
    "cuCtxPushCurrent",
    "cuCtxPopCurrent",
    "cuCtxGetDevice", # this actually does require a context, but we use it unsafely
    ## primary context management
    "cuDevicePrimaryCtxGetState",
    "cuDevicePrimaryCtxRelease",
    "cuDevicePrimaryCtxReset",
    "cuDevicePrimaryCtxRetain",
    "cuDevicePrimaryCtxSetFlags",
    # CUPTI
    "cuptiGetVersion",
    "cuptiGetResultString",
    # NVML
    "nvmlInit",
    "nvmlInitWithFlags",
    "nvmlShutdown",
    "nvmlErrorString",
    # CUBLAS
    "cublasGetVersion",
    "cublasGetProperty",
    "cublasGetCudartVersion",
    # CURAND
    "curandGetVersion",
    "curandGetProperty",
    # CUFFT
    "cufftGetVersion",
    "cufftGetProperty",
    # CUSPARSE
    "cusparseGetVersion",
    "cusparseGetProperty",
    "cusparseGetErrorName",
    "cusparseGetErrorString",
    # CUSOLVER
    "cusolverGetVersion",
    "cusolverGetProperty",
    # CUDNN
    "cudnnGetVersion",
    "cudnnGetProperty",
    "cudnnGetCudartVersion",
    "cudnnGetErrorString",
    # CUTENSOR
    "cutensorGetVersion",
    "cutensorGetCudartVersion",
    "cutensorGetErrorString",
])
function insert_init_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.Call && x.args[1].val == "ccall"
        fun = x.args[3].args[2].args[2].val

        # strip the version tag
        if occursin(r"_v\d$", fun)
            fun = fun[1:end-3]
        end

        # call the API initializer
        if !in(fun, preinit_apicalls)
            push!(state.edits, Edit(state.offset, "initialize_api()\n    "))
        end
    end
end

# change ::Ptr arguments to ::CuPtr / ::PtrOrCuPtr based on user input
function rewrite_pointers_pass(x, state, headers)
    if x isa CSTParser.EXPR && x.typ == CSTParser.Call && x.args[1].val == "ccall"
        # get the ccall arguments, skipping comma's and parentheses
        handle = x.args[3]
        fn = handle.args[2].args[2].val
        rv = x.args[5]
        tt = x.args[7]
        args = x.args[9:2:end]

        # check for pointer types
        types = tt.args[2:2:end-1]
        is_pointer = Bool[x.typ == CSTParser.Curly && x.args[1].val == "Ptr" for x in types]
        offset = state.offset + sum(x->x.fullspan, x.args[1:6])
        if any(is_pointer)
            if haskey(db, fn)
                replacements = db[fn]

                # regenerate replacements with the new argument names
                changed = false
                old_replacements = collect(replacements)
                replacements = OrderedDict{String,Any}()
                for (i, arg) in enumerate(args)
                    if arg.val != old_replacements[i].first
                        changed = true
                    end
                    replacements[arg.val] = old_replacements[i].second
                end

                if changed
                    db[fn] = replacements
                    save_db()
                end
            else
                # print some context from the header (some mention host/device pointer)
                print(Crayon(foreground=:yellow))
                run(`awk "/\<$fn\>/,/;/" $headers`)
                println(Crayon(reset=true))

                # print pointer arguments and their types
                println(Crayon(foreground = :red), fn, Crayon(reset=true))
                for (i, arg) in enumerate(args)
                    if is_pointer[i]
                        println("- argument $i: $(arg.val)::$(Expr(types[i]))")
                    end
                end
                println()

                # prompt
                run(pipeline(`echo -n $fn`, `xclip -i -selection clipboard`));
                print("GPU pointers> ")
                gpu_pointers = parse.(Int, split(readline(stdin)))
                if gpu_pointers == [0]
                    # 0 is special match for all pointers
                    gpu_pointers = findall(is_pointer)
                elseif !isempty(gpu_pointers) && all(i->i<0, gpu_pointers)
                    # negative indicates all but these
                    gpu_pointers = map(i->-i, gpu_pointers)
                    @assert all(i->is_pointer[i], gpu_pointers) "You selected non-pointer arguments"
                    gpu_pointers = setdiff(findall(is_pointer), gpu_pointers)
                end
                print("Dual GPU/CPU pointers> ")
                dual_pointers = parse.(Int, split(readline(stdin)))
                @assert all(i->is_pointer[i], gpu_pointers ∪ dual_pointers) "You selected non-pointer arguments"

                # generate replacements
                replacements = OrderedDict{String,Any}()
                for (i, arg) in enumerate(args)
                    replacements[arg.val] = if is_pointer[i]
                        if i in gpu_pointers
                            "CuPtr"
                        elseif i in dual_pointers
                            "PtrOrCuPtr"
                        else
                            "Ptr"
                        end
                    else
                        nothing
                    end
                end

                db[fn] = replacements
                save_db()
            end

            # generate edits
            for (i, (_,replacement)) in enumerate(replacements)
                offset += tt.args[2*i-1].fullspan
                if replacement !== nothing && replacement != "Ptr"
                    ptr = types[i].args[1]
                    push!(state.edits, Edit(offset+1:offset+ptr.span, replacement))
                end
                offset += types[i].fullspan
            end

            println()
        end
    end
end

# replace handles from the CUDA runtime library with CUDA driver API equivalents
function rewrite_runtime_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.IDENTIFIER && x.val == "cudaStream_t"
        offset = state.offset
        push!(state.edits, Edit(offset+1:offset+x.span, "CUstream"))
    end
end

# remove aliases that are used to version functions
#
# when NVIDIA changes the behavior of an API, they version the function (`cuFunction_v2`).
# To maintain backwards compatibility, they ship aliases with their headers such that
# compiled binaries will keep using the old version, and newly-compiled ones will use the
# developer's CUDA version. remove those, since we target multiple CUDA versions.
#
# remove this if we ever decide to support a single supported version of CUDA.
cuda_version_alias(lhs, rhs) = occursin(Regex("$(lhs)_v\\d"), rhs)
nvtx_string_alias(lhs, rhs) = occursin(Regex("$(lhs)(A|W|Ex)"), rhs)
function remove_aliases_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.Const
        @assert x.args[2].typ == CSTParser.BinaryOpCall
        lhs, op, rhs = x.args[2].args
        @assert op.typ == CSTParser.OPERATOR && op.kind == Tokenize.Tokens.EQ
        if lhs.typ == CSTParser.IDENTIFIER && rhs.typ == CSTParser.IDENTIFIER &&
            (cuda_version_alias(lhs.val, rhs.val) || nvtx_string_alias(lhs.val, rhs.val))
            offset = state.offset
            push!(state.edits, Edit(offset+1:offset+x.span, ""))
        end
    end
end

## indenting passes

mutable struct IndentState
    offset::Int
    lines
    edits::Vector{Edit}
end

function get_lines(text)
    lines = Tuple{Int,Int}[]
    pt = Tokens.EMPTY_TOKEN(Tokens.Token)
    for t in CSTParser.Tokenize.tokenize(text)
        if pt.endpos[1] != t.startpos[1]
            if t.kind == Tokens.WHITESPACE
                nl = findfirst("\n", t.val) != nothing
                if !nl
                    push!(lines, (length(t.val), 0))
                else
                end
            else
                push!(lines, (0, 0))
            end
        elseif t.startpos[1] != t.endpos[1] && t.kind == Tokens.TRIPLE_STRING
            nls = findall(x->x == '\n', t.val)
            for nl in nls
                push!(lines, (t.startpos[2] - 1, nl + t.startbyte))
            end
        elseif t.startpos[1] != t.endpos[1] && t.kind == Tokens.WHITESPACE
            push!(lines, (t.endpos[2], t.endbyte - t.endpos[2] + 1))
        end
        pt = t
    end
    lines
end

function wrap_at_comma(x, state, indent, offset, column)
    comma = nothing
    for y in x
        if column + y.fullspan > 92 && comma !== nothing
            column = indent
            push!(state.edits, Edit(comma, ",\n" * " "^column))
            column += offset - comma[1] - 1 # other stuff might have snuck between the comma and the current expr
            comma = nothing
        elseif y.typ == CSTParser.PUNCTUATION && y.kind == Tokens.COMMA
            comma = (offset+1):(offset+y.fullspan)
        end
        offset += y.fullspan
        column += y.fullspan
    end
end

function indent_ccall_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.MacroCall && x.args[1].args[2].val == "runtime_ccall"
        # figure out how much to indent by looking at where the expr starts
        line = findlast(y -> state.offset >= y[2], state.lines) # index, not the actual number
        line_indent, line_offset = state.lines[line]
        expr_indent = state.offset - line_offset
        indent = expr_indent + sum(x->x.fullspan, x.args[1:2])

        if length(x.args[7]) > 2    # non-empty tuple type
            # break before the tuple type
            offset = state.offset + sum(x->x.fullspan, x.args[1:6])
            push!(state.edits, Edit(offset:offset, "\n" * " "^indent))

            # wrap tuple type
            wrap_at_comma(x.args[7], state, indent+1, offset, indent+1)
        end

        if length(x.args) > 9
            # break before the arguments
            offset = state.offset + sum(x->x.fullspan, x.args[1:8])
            push!(state.edits, Edit(offset:offset, "\n" * " "^indent))

            # wrap arguments
            wrap_at_comma(x.args[9:end], state, indent, offset, indent)
        end
    end
end

function indent_definition_pass(x, state)
    if x isa CSTParser.EXPR && x.typ == CSTParser.FunctionDef
        # figure out how much to indent by looking at where the expr starts
        line = findlast(y -> state.offset >= y[2], state.lines) # index, not the actual number
        line_indent, line_offset = state.lines[line]
        expr_indent = state.offset - line_offset
        indent = expr_indent + x.args[1].fullspan + sum(x->x.fullspan, x.args[2].args[1:2])

        if length(x.args[2]) > 2    # non-empty args
            offset = state.offset + x.args[1].fullspan + sum(x->x.fullspan, x.args[2].args[1:2])
            wrap_at_comma(x.args[2].args[3:end-1], state, indent, offset, indent)
        end
    end
end

#
# Main application
#

using CUDA_full_jll, CUDNN_CUDA102_jll, CUTENSOR_CUDA102_jll

function process(name, headers...; libname=name, rewrite_pointers=true, kwargs...)
    output_file, common_file = wrap(libname, headers...; kwargs...)

    for file in (output_file, common_file)
        text = read(file, String)

        ## rewriting passes

        state = State(0, Edit[])
        ast = CSTParser.parse(text, true)

        state.offset = 0
        pass(ast, state, insert_check_pass)

        state.offset = 0
        pass(ast, state, rewrite_ccall_pass)

        state.offset = 0
        pass(ast, state, insert_init_pass)

        if rewrite_pointers
            state.offset = 0
            pass(ast, state, (x,state)->rewrite_pointers_pass(x,state,headers))
        end

        state.offset = 0
        pass(ast, state, rewrite_runtime_pass)

        state.offset = 0
        pass(ast, state, remove_aliases_pass)

        # apply
        state.offset = 0
        sort!(state.edits, lt = (a,b) -> first(a.loc) < first(b.loc), rev = true)
        for i = 1:length(state.edits)
            text = apply(text, state.edits[i])
        end

        ## indentation passes

        lines = get_lines(text)
        state = IndentState(0, lines, Edit[])
        ast = CSTParser.parse(text, true)

        state.offset = 0
        pass(ast, state, indent_definition_pass)

        state.offset = 0
        pass(ast, state, indent_ccall_pass)

        # apply
        state.offset = 0
        sort!(state.edits, lt = (a,b) -> first(a.loc) < first(b.loc), rev = true)
        for i = 1:length(state.edits)
            text = apply(text, state.edits[i])
        end

        ## header

        squeezed = replace(text, "\n\n\n"=>"\n\n")
        while length(text) != length(squeezed)
            text = squeezed
            squeezed = replace(text, "\n\n\n"=>"\n\n")
        end

        text = """
            # Automatically-generated headers for $name
            #
            # DO NOT EDIT THIS FILE DIRECTLY.

            """ * squeezed

        write(file, text)
    end

    ## move to destination

    dstdir = joinpath(dirname(dirname(@__DIR__)), "lib", name)
    for src in (output_file, common_file)
        dst = joinpath(dstdir, src)
        cp(src, dst; force=true)
    end

    ## manual patches

    patchdir = joinpath(@__DIR__, "patches", name)
    if isdir(patchdir)
        for entry in readdir(patchdir)
            if endswith(entry, ".patch")
                path = joinpath(patchdir, entry)
                run(`patch -p1 -d $dstdir -i $path`)
            end
        end
    end

    return
end

function main()
    load_db()

    cuda = joinpath(CUDA_full_jll.artifact_dir, "cuda", "include")
    cupti = joinpath(CUDA_full_jll.artifact_dir, "cuda", "extras", "CUPTI", "include")
    cudnn = joinpath(CUDNN_CUDA102_jll.artifact_dir, "include")
    cutensor = joinpath(CUTENSOR_CUDA102_jll.artifact_dir, "include")

    process("cudadrv", "$cuda/cuda.h", "$cuda/cudaProfiler.h";
            include_dirs=[cuda], libname="cuda", rewrite_pointers=false)

    process("nvtx", "$cuda/nvtx3/nvToolsExt.h", "$cuda/nvtx3/nvToolsExtCuda.h";
            include_dirs=[cuda], rewrite_pointers=false)

    process("nvml", "$cuda/nvml.h"; include_dirs=[cuda], rewrite_pointers=false)

    process("cupti", "$cupti/cupti.h", "$cupti/cupti_profiler_target.h";
            include_dirs=[cuda, cupti],
            wrapped_headers=["cupti_result.h", "cupti_version.h", "cupti_activity.h",
                             "cupti_callbacks.h", "cupti_events.h",
                             "cupti_profiler_target.h",
                             "cupti_metrics.h"], # deprecated, but still required
            defines=["__packed__"=>"", "aligned"=>""],
            rewrite_pointers=false)
    # NOTE: libclang (the C API) doesn't support/expose the __packed__/aligned attributes,
    #       so disable them (Julia doesn't support packed structs anyway)

    process("cublas", "$cuda/cublas_v2.h", "$cuda/cublasXt.h";
            wrapped_headers=["cublas_v2.h", "cublas_api.h", "cublasXt.h"],
            defines=["CUBLASAPI"=>""],
            include_dirs=[cuda])

    process("cufft", "$cuda/cufft.h"; include_dirs=[cuda])

    process("curand", "$cuda/curand.h"; include_dirs=[cuda])

    process("cusparse", "$cuda/cusparse.h"; include_dirs=[cuda])

    process("cusolver", "$cuda/cusolverDn.h", "$cuda/cusolverSp.h";
             wrapped_headers=["cusolver_common.h", "cusolverDn.h", "cusolverSp.h"],
             include_dirs=[cuda])

    process("cudnn", "$cudnn/cudnn.h"; include_dirs=[cuda, cudnn])

    process("cutensor", "$cutensor/cutensor.h";
            wrapped_headers=["cutensor.h", "cutensor/types.h"],
            include_dirs=[cuda, cutensor])
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


text:
using PyPlot,EE

close("all")
figure(figsize=(3.3,2.5))
rc("text",usetex=true);
rc("text.latex",preamble="\\usepackage{fourier}\\usepackage[scaled=0.95]{inconsolata}");
rc("font", family="serif")

U = 100+0im
refU = 100
Z = 30+40im
I = U/Z
Ur = real(Z)*I
Ux = 1im*imag(Z)*I
refI=abs(I)*0.8

phasor(U,label=L"$\underline{U}$",tlabel=-0.1,ref=refU,relrot=true)
phasor(Ur,label=L"$\underline{U}_r$",tlabel=-0.1,ref=refU,relrot=true)
phasor(Ux,origin=Ur,label=L"$\underline{U}_x$",tlabel=0.15,ref=refU,relrot=true)
phasor(I,label=L"$\underline{I}$",tlabel=0.15,rlabel=0.7,ref=refI,relrot=true,par=0.05)

axis("square")
xlim(-1,1)
ylim(-1,1)

# Remove axis
ax=gca()
ax[:set_axis_off]()