#  Copyright 2016, 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 http://mozilla.org/MPL/2.0/.

function tryParseIdxSet(arg::Expr)
    # :in appears as arg.head only prior to 0.5
    if arg.head === :(=) || arg.head === :in
        @assert length(arg.args) == 2
        return true, arg.args[1], arg.args[2]
    elseif isexpr(arg, :call) && arg.args[1] === :in
        return true, arg.args[2], arg.args[3]
    else
        return false, nothing, nothing
    end
end

function parseIdxSet(arg::Expr)
    parse_done, idxvar, idxset = tryParseIdxSet(arg)
    if parse_done
        return idxvar, idxset
    end
    error("Invalid syntax: $arg")
end

# addtoexpr(ex::Number, c::Number) = ex + c

addtoexpr(ex::Number, c::Number, x::Number) = ex + c*x

addtoexpr(ex::Number, c::Number, x::Variable) = AffExpr([x],[c],ex)

function addtoexpr{T<:GenericAffExpr}(ex::Number, c::Number, x::T)
    # It's only safe to mutate the first argument.
    if c == 0
        T(ex)
    else
        x = copy(x)
        scale!(x.coeffs, c)
        x.constant *= c
        x.constant += ex
        x
    end
end

function addtoexpr{T<:GenericQuadExpr}(ex::Number, c::Number, x::T)
    # It's only safe to mutate the first argument.
    if c == 0
        T(ex)
    else
        x = copy(x)
        scale!(x.qcoeffs, c)
        scale!(x.aff.coeffs, c)
        x.aff.constant *= c
        x.aff.constant += ex
        x
    end
end

addtoexpr(ex::Number, c::Variable, x::Variable) = QuadExpr([c],[x],[1.0],zero(AffExpr))

function addtoexpr{T<:GenericAffExpr}(ex::Number, c::T, x::T)
    q = c*x
    q.aff.constant += ex
    q
end

function addtoexpr{C,V}(ex::Number, c::GenericAffExpr{C,V}, x::V)
    q = c*x
    q.aff.constant += ex
    q
end

function addtoexpr{T<:GenericQuadExpr}(ex::Number, c::T, x::Number)
    if x == 0
        T(ex)
    else
        q = c*x
        q.aff.constant += ex
        q
    end
end

function addtoexpr(aff::GenericAffExpr, c::Number, x::Number)
    aff.constant += c*x
    aff
end

function addtoexpr{C,V}(aff::GenericAffExpr{C,V}, c::Number, x::V)
    if c != 0
        push!(aff.vars,   x)
        push!(aff.coeffs, c)
    end
    aff
end

function addtoexpr{C,V}(aff::GenericAffExpr{C,V},c::Number,x::GenericAffExpr{C,V})
    if c != 0
        append!(aff.vars, x.vars)
        sizehint!(aff.coeffs, length(aff.coeffs)+length(x.coeffs))
        for i in 1:length(x.coeffs)
            push!(aff.coeffs, c*x.coeffs[i])
        end
        aff.constant += c*x.constant
    end
    aff
end

# help w/ ambiguity
addtoexpr{C,V<:Number}(aff::GenericAffExpr{C,V}, c::Number, x::Number) = aff + c*x

function addtoexpr{C,V}(aff::GenericAffExpr{C,V}, c::V, x::Number)
    if x != 0
        push!(aff.vars,   c)
        push!(aff.coeffs, x)
    end
    aff
end

addtoexpr{C,V}(aff::GenericAffExpr{C,V},c::V,x::V) =
    GenericQuadExpr{C,V}([c],[x],[one(C)],aff)

# TODO: add generic versions of following two methods
addtoexpr(aff::AffExpr,c::AffExpr,x::Variable) =
    QuadExpr(c.vars,
             fill(x,length(c.vars)),
             c.coeffs,
             addtoexpr(aff,c.constant,x))

addtoexpr(aff::AffExpr,c::Variable,x::AffExpr) =
    QuadExpr(fill(c,length(x.vars)),
             x.vars,
             x.coeffs,
             addtoexpr(aff,c,x.constant))

function addtoexpr{C,V}(aff::GenericAffExpr{C,V},c::GenericAffExpr{C,V},x::Number)
    if x != 0
        append!(aff.vars, c.vars)
        append!(aff.coeffs, c.coeffs * x)
        aff.constant += c.constant * x
    end
    aff
end

function addtoexpr{C,V}(aff::GenericAffExpr{C,V}, c::GenericQuadExpr{C,V}, x::Number)
    if x == 0
        GenericQuadExpr{C,V}(aff)
    else
        GenericQuadExpr{C,V}(copy(c.qvars1),
                            copy(c.qvars2),
                            c.qcoeffs*x,
                            addtoexpr(aff,c.aff,x))
    end
end

function addtoexpr{C,V}(aff::GenericAffExpr{C,V}, c::Number, x::GenericQuadExpr{C,V})
    if c == 0
        GenericQuadExpr{C,V}(aff)
    else
        GenericQuadExpr{C,V}(copy(x.qvars1),
                            copy(x.qvars2),
                            c*x.qcoeffs,
                            addtoexpr(aff,c,x.aff))
    end
end

function addtoexpr{C,V}(ex::GenericAffExpr{C,V}, c::GenericAffExpr{C,V}, x::GenericAffExpr{C,V})
    q = convert(GenericQuadExpr{C,V}, ex)
    addtoexpr(q, one(C), c*x)
    q
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::Number,x::V)
    if c != 0
        push!(quad.aff, convert(C,c), x)
    end
    quad
end

function addtoexpr(quad::GenericQuadExpr,c::Number,x::Number)
    quad.aff.constant += c*x
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},x::V,y::V)
    push!(quad.qvars1, x)
    push!(quad.qvars2, y)
    push!(quad.qcoeffs, one(C))
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::Number,x::GenericAffExpr{C,V})
    if c != 0
        append!(quad.aff.vars, x.vars)
        sizehint!(quad.aff.coeffs, length(quad.aff.coeffs)+length(x.coeffs))
        for i in 1:length(x.coeffs)
            push!(quad.aff.coeffs, c*x.coeffs[i])
        end
        quad.aff.constant += c*x.constant
    end
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::GenericAffExpr{C,V},x::Number)
    if x != 0
        addtoexpr(quad.aff,c,x)
    end
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::GenericAffExpr{C,V},x::V)
    append!(quad.qvars1, c.vars)
    append!(quad.qvars2, fill(x,length(c.vars)))
    append!(quad.qcoeffs, c.coeffs)
    addtoexpr(quad.aff,c.constant,x)
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::V,x::GenericAffExpr{C,V})
    append!(quad.qvars1, fill(c,length(x.vars)))
    append!(quad.qvars2, x.vars)
    append!(quad.qcoeffs, x.coeffs)
    addtoexpr(quad.aff,c,x.constant)
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::GenericQuadExpr{C,V},x::Number)
    if x != 0
        append!(quad.qvars1,c.qvars1)
        append!(quad.qvars2,c.qvars2)
        sizehint!(quad.qcoeffs, length(quad.qcoeffs)+length(c.qcoeffs))
        for i in 1:length(c.qcoeffs)
            push!(quad.qcoeffs, c.qcoeffs[i]*x)
        end
        addtoexpr(quad,c.aff,x)
    end
    quad
end

function addtoexpr{C,V}(quad::GenericQuadExpr{C,V},c::Number,x::GenericQuadExpr{C,V})
    if c != 0
        append!(quad.qvars1,x.qvars1)
        append!(quad.qvars2,x.qvars2)
        sizehint!(quad.qcoeffs, length(quad.qcoeffs)+length(x.qcoeffs))
        for i in 1:length(x.qcoeffs)
            push!(quad.qcoeffs, c*x.qcoeffs[i])
        end
        addtoexpr(quad,c,x.aff)
    end
    quad
end

function addtoexpr{C,V}(ex::GenericQuadExpr{C,V}, c::GenericAffExpr{C,V}, x::GenericAffExpr{C,V})
    q = c*x
    addtoexpr(ex, 1.0, q)
    ex
end

# Catch nonlinear expressions and parameters being used in addconstraint, etc.

typealias _NLExpr Union{NonlinearExpression,NonlinearParameter}
_nlexprerr() = error("""Cannot use nonlinear expression or parameter in @constraint or @objective.
                        Use @NLconstraint or @NLobjective instead.""")
# Following three definitions avoid ambiguity warnings
addtoexpr{C,V<:_NLExpr}(expr::GenericQuadExpr{C,V}, c::GenericAffExpr{C,V}, x::V) = _nlexprerr()
addtoexpr{C,V<:_NLExpr}(expr::GenericQuadExpr{C,V}, c::Number, x::V) = _nlexprerr()
addtoexpr{C,V<:_NLExpr}(expr::GenericQuadExpr{C,V}, c::V, x::GenericAffExpr{C,V}) = _nlexprerr()
for T1 in (GenericAffExpr,GenericQuadExpr), T2 in (Number,Variable,GenericAffExpr,GenericQuadExpr)
    @eval addtoexpr(::$T1, ::$T2, ::_NLExpr) = _nlexprerr()
    @eval addtoexpr(::$T1, ::_NLExpr, ::$T2) = _nlexprerr()
end

addtoexpr(ex, c, x) = ex + c*x

@generated addtoexpr_reorder(ex, arg) = :(addtoexpr(ex, 1.0, arg))

@generated function addtoexpr_reorder(ex, x, y)
    if x <: Union{Variable,AffExpr} && y <: Number
        :(addtoexpr(ex, y, x))
    else
        :(addtoexpr(ex, x, y))
    end
end

@generated function addtoexpr_reorder(ex, args...)
    n = length(args)
    @assert n ≥ 3
    varidx = find(t -> (t == Variable || t == AffExpr), collect(args))
    allscalar = all(t -> (t <: Number), args[setdiff(1:n, varidx)])
    idx = (allscalar && length(varidx) == 1) ? varidx[1] : n
    coef = Expr(:call, :*, [:(args[$i]) for i in setdiff(1:n,idx)]...)
    :(addtoexpr(ex, $coef, args[$idx]))
end

function parseCurly(x::Expr, aff::Symbol, lcoeffs, rcoeffs, newaff=gensym())
    header = x.args[1]
    if length(x.args) < 3
        error("Need at least two arguments for $header")
    end
    if issum(header)
        parseSum(x, aff, lcoeffs, rcoeffs, newaff)
    elseif header ∈ [:norm1, :norm2, :norminf, :norm∞]
        parseNorm(header, x, aff, lcoeffs, rcoeffs, newaff)
    else
        error("Expected sum or norm2 outside curly braces; got $header")
    end
end

function parseSum(x::Expr, aff::Symbol, lcoeffs, rcoeffs, newaff)
    # we have a filter condition
    if isexpr(x.args[2],:parameters)
        cond = x.args[2]
        if length(cond.args) != 1
            error("No commas after semicolon allowed in sum expression, use && for multiple conditions")
        end
        # generate inner loop code first and then wrap in for loops
        inneraff, innercode = parseExpr(x.args[3], aff, lcoeffs, rcoeffs, aff)
        code = quote
            if $(esc(cond.args[1]))
                $innercode
            end
        end
        for level in length(x.args):-1:4
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
        end
    else # no condition
        inneraff, code = parseExpr(x.args[2], aff, lcoeffs, rcoeffs, aff)
        for level in length(x.args):-1:3
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
        end
        len = :len
        # precompute the number of elements to add
        # this is unncessary if we're just summing constants
        _, lastidxset = parseIdxSet(x.args[length(x.args)]::Expr)
        preblock = :($len += length($(esc(lastidxset))))
        for level in (length(x.args)-1):-1:3
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            preblock = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $preblock
                end
            end)
        end
        preblock = quote
            $len = 0
            $preblock
            _sizehint_expr!($aff,$len)
        end
        code = :($preblock;$code)
    end
    :($code; $newaff=$aff)
end

function parseNorm(normp::Symbol, x::Expr, aff::Symbol, lcoeffs, rcoeffs, newaff)
    @assert string(x.args[1])[1:4] == "norm"
    # we have a filter condition
    finalaff = gensym()
    gennorm  = gensym()
    len      = gensym()
    normexpr = gensym()
    if isexpr(x.args[2],:parameters)
        cond = x.args[2]
        if length(cond.args) != 1
            error("No commas after semicolon allowed in sum expression, use && for multiple conditions")
        end
        # generate inner loop code first and then wrap in for loops
        inneraff, innercode = parseExprToplevel(x.args[3], :normaff)
        code = quote
            if $(esc(cond.args[1]))
                normaff = 0.0
                $innercode
                push!($normexpr, $inneraff)
            end
        end
        for level in length(x.args):-1:4
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
        end
        preblock = :($normexpr = GenericAffExpr[])
    else # no condition
        inneraff, code = parseExprToplevel(x.args[2], :normaff)
        code = :(normaff = 0.0; $code; push!($normexpr, $inneraff))
        _, lastidxset = parseIdxSet(x.args[length(x.args)]::Expr)
        preblock = :($len += length($(esc(lastidxset))))
        for level in length(x.args):-1:3
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
            preblock = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $preblock
                end
            end)
        end
        preblock = quote
            $len = 0
            $preblock
            $normexpr = GenericAffExpr[]
            _sizehint_expr!($normexpr, $len)
        end
    end
    if normp == :norm2
        param = 2
    elseif normp == :norm1
        param = 1
    elseif normp ∈ [:norminf, :norm∞]
        param = Inf
    else
        error("Unrecognized norm: $normp")
    end
    quote
        $preblock
        $code
        $gennorm = _build_norm($param,$normexpr)
        $newaff = $(Expr(:call, :addtoexpr_reorder, aff,lcoeffs...,gennorm,rcoeffs...))
    end
end

function parseGenerator(x::Expr, aff::Symbol, lcoeffs, rcoeffs, newaff=gensym())
    @assert isexpr(x,:call)
    @assert length(x.args) > 1
    @assert isexpr(x.args[2],:generator)
    header = x.args[1]
    if issum(header)
        parseGeneratorSum(x.args[2], aff, lcoeffs, rcoeffs, newaff)
    elseif header ∈ [:norm1, :norm2, :norminf, :norm∞]
        parseGeneratorNorm(header, x.args[2], aff, lcoeffs, rcoeffs, newaff)
    else
        error("Expected sum or norm2 outside generator expression; got $header")
    end
end

function parseGeneratorSum(x::Expr, aff::Symbol, lcoeffs, rcoeffs, newaff)
    # we have a filter condition
    if isexpr(x.args[2],:filter)
        cond = x.args[2].args[1]
        # generate inner loop code first and then wrap in for loops
        inneraff, innercode = parseExpr(x.args[1], aff, lcoeffs, rcoeffs, aff)
        code = quote
            if $(esc(cond))
                $innercode
            end
        end
        for level in length(x.args[2].args):-1:2
            _idxvar, idxset = parseIdxSet(x.args[2].args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
        end
    else # no condition
        inneraff, code = parseExpr(x.args[1], aff, lcoeffs, rcoeffs, aff)
        for level in length(x.args):-1:2
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
        end
        len = :len
        # precompute the number of elements to add
        # this is unncessary if we're just summing constants
        _, lastidxset = parseIdxSet(x.args[length(x.args)]::Expr)
        preblock = :($len += length($(esc(lastidxset))))
        for level in (length(x.args)-1):-1:2
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            preblock = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $preblock
                end
            end)
        end
        preblock = quote
            $len = 0
            $preblock
            _sizehint_expr!($aff,$len)
        end
        code = :($preblock;$code)
    end
    :($code; $newaff=$aff)
end

function parseGeneratorNorm(normp::Symbol, x::Expr, aff::Symbol, lcoeffs, rcoeffs, newaff)
    @assert string(normp)[1:4] == "norm"
    finalaff = gensym()
    gennorm  = gensym()
    len      = gensym()
    normexpr = gensym()
    # we have a filter condition
    if isexpr(x.args[2],:filter)
        cond = x.args[2].args[1]
        # generate inner loop code first and then wrap in for loops
        inneraff, innercode = parseExprToplevel(x.args[1], :normaff)
        code = quote
            if $(esc(cond))
                normaff = 0.0
                $innercode
                push!($normexpr, $inneraff)
            end
        end
        for level in length(x.args[2].args):-1:2
            _idxvar, idxset = parseIdxSet(x.args[2].args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
        end
        preblock = :($normexpr = GenericAffExpr[])
    else # no condition
        inneraff, code = parseExprToplevel(x.args[1], :normaff)
        code = :(normaff = 0.0; $code; push!($normexpr, $inneraff))
        _, lastidxset = parseIdxSet(x.args[length(x.args)]::Expr)
        preblock = :($len += length($(esc(lastidxset))))
        for level in length(x.args):-1:2
            _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
            idxvar = esc(_idxvar)
            code = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $code
                end
            end)
            preblock = :(let
                $(localvar(idxvar))
                for $idxvar in $(esc(idxset))
                    $preblock
                end
            end)
        end
        preblock = quote
            $len = 0
            $preblock
            $normexpr = GenericAffExpr[]
            _sizehint_expr!($normexpr, $len)
        end
    end
    if normp == :norm2
        param = 2
    elseif normp == :norm1
        param = 1
    elseif normp ∈ [:norminf, :norm∞]
        param = Inf
    else
        error("Unrecognized norm: $normp")
    end
    quote
        $preblock
        $code
        $gennorm = _build_norm($param,$normexpr)
        $newaff = $(Expr(:call, :addtoexpr_reorder, aff,lcoeffs...,gennorm,rcoeffs...))
    end
end

is_complex_expr(ex) = isa(ex,Expr) && !isexpr(ex,:ref)

parseExprToplevel(x, aff::Symbol) = parseExpr(x, aff, [], [])

# output is assigned to newaff
function parseExpr(x, aff::Symbol, lcoeffs::Vector, rcoeffs::Vector, newaff::Symbol=gensym())
    if !isa(x,Expr)
        # at the lowest level
        callexpr = Expr(:call,:addtoexpr_reorder,aff,lcoeffs...,esc(x),rcoeffs...)
        return newaff, :($newaff = $callexpr)
    else
        if x.head == :call && x.args[1] == :+
            b = Expr(:block)
            aff_ = aff
            for arg in x.args[2:(end-1)]
                aff_, code = parseExpr(arg, aff_, lcoeffs, rcoeffs)
                push!(b.args, code)
            end
            newaff, code = parseExpr(x.args[end], aff_, lcoeffs, rcoeffs, newaff)
            push!(b.args, code)
            return newaff, b
        elseif x.head == :call && x.args[1] == :-
            if length(x.args) == 2 # unary subtraction
                return parseExpr(x.args[2], aff, vcat(-1.0, lcoeffs), rcoeffs, newaff)
            else # a - b - c ...
                b = Expr(:block)
                aff_, code = parseExpr(x.args[2], aff, lcoeffs, rcoeffs)
                push!(b.args, code)
                for arg in x.args[3:(end-1)]
                    aff_,code = parseExpr(arg, aff_, vcat(-1.0, lcoeffs), rcoeffs)
                    push!(b.args, code)
                end
                newaff,code = parseExpr(x.args[end], aff_, vcat(-1.0, lcoeffs), rcoeffs, newaff)
                push!(b.args, code)
                return newaff, b
            end
        elseif x.head == :call && x.args[1] == :*
            # we might need to recurse on multiple arguments, e.g.,
            # (x+y)*(x+y)
            n_expr = mapreduce(is_complex_expr, +, x.args)
            if n_expr == 1 # special case, only recurse on one argument and don't create temporary objects
                which_idx = 0
                for i in 2:length(x.args)
                    if is_complex_expr(x.args[i])
                        which_idx = i
                    end
                end
                return parseExpr(x.args[which_idx], aff, vcat(lcoeffs, [esc(x.args[i]) for i in 2:(which_idx-1)]),
                                                         vcat(rcoeffs, [esc(x.args[i]) for i in (which_idx+1):length(x.args)]),
                                                         newaff)
            else
                blk = Expr(:block)
                for i in 2:length(x.args)
                    if is_complex_expr(x.args[i])
                        s = gensym()
                        newaff_, parsed = parseExprToplevel(x.args[i], s)
                        push!(blk.args, :($s = 0.0; $parsed))
                        x.args[i] = newaff_
                    else
                        x.args[i] = esc(x.args[i])
                    end
                end
                callexpr = Expr(:call,:addtoexpr_reorder,aff,lcoeffs...,x.args[2:end]...,rcoeffs...)
                push!(blk.args, :($newaff = $callexpr))
                return newaff, blk
            end
        elseif x.head == :call && x.args[1] == :^ && is_complex_expr(x.args[2])
            if x.args[3] == 2
                blk = Expr(:block)
                s = gensym()
                newaff_, parsed = parseExprToplevel(x.args[2], s)
                push!(blk.args, :($s = 0.0; $parsed))
                push!(blk.args, :($newaff = $aff + $(Expr(:call,:*,lcoeffs...,newaff_,newaff_,rcoeffs...))))
                return newaff, blk
            elseif x.args[3] == 1
                return parseExpr(x.args[2], aff, lcoeffs, rcoeffs)
            elseif x.args[3] == 0
                return parseExpr(1, aff, lcoeffs, rcoeffs)
            else
                blk = Expr(:block)
                s = gensym()
                newaff_, parsed = parseExprToplevel(x.args[2], s)
                push!(blk.args, :($s = 0.0; $parsed))
                push!(blk.args, :($newaff = $aff + $(Expr(:call,:*,lcoeffs...,Expr(:call,:^,newaff_,esc(x.args[3])),rcoeffs...))))
                return newaff, blk
            end
        elseif x.head == :call && x.args[1] == :/
            @assert length(x.args) == 3
            numerator = x.args[2]
            denom = x.args[3]
            return parseExpr(numerator, aff, lcoeffs,vcat(esc(:(1/$denom)),rcoeffs),newaff)
        elseif isexpr(x,:call) && length(x.args) >= 2 && isexpr(x.args[2],:generator)
            return newaff, parseGenerator(x,aff,lcoeffs,rcoeffs,newaff)
        elseif isexpr(x,:call) && length(x.args) >= 2 && isexpr(x.args[2],:flatten)
            flatten_error(x.args[2])
        elseif x.head == :curly
            return newaff, parseCurly(x,aff,lcoeffs,rcoeffs,newaff)
        else # at lowest level?
            !isexpr(x,:comparison) || error("Unexpected comparison in expression $x")
            callexpr = Expr(:call,:addtoexpr_reorder,aff,lcoeffs...,esc(x),rcoeffs...)
            return newaff, :($newaff = $callexpr)
        end
    end
end

# Semi-automatically generated precompilation hints
include("precompile.jl")


h5p_set_virtual(dcpl_id, vspace_id, src_file_name, src_dset_name, src_space_id) =
    ccall((:H5Pset_virtual, libhdf5), Herr, (Hid, Hid, Ptr{UInt8}, Ptr{UInt8}, Hid), dcpl_id, vspace_id, src_file_name, src_dset_name, src_space_id)

mutable struct VirtualLayout{S, D, SS}
    shape::S
    dtype::D
    sources::SS
end

VirtualLayout(shape, dtype) = VirtualLayout(tuple(shape...), dtype, [])

repcolon(is, shape) = ifelse.(is .== Colon(), Base.OneTo.(shape), is)

function Base.setindex!(layout::VirtualLayout, v, is...)
    push!(layout.sources, (repcolon(is, layout.shape), v))
    shape = collect(layout.shape)
    for (vis, vs) in layout.sources
        layout.shape = max.(layout.shape, last.(vis))
    end
end

struct VirtualSource{P, N, S, D, I}
    path::P
    name::N
    shape::S
    dtype::D
    is::I
end

function VirtualSource(path, name)
    dset = h5open(path, "r")[name]
    shape, dtype = size(dset), eltype(dset)
    VirtualSource(path, name, shape, dtype, nothing)
end

VirtualSource(dset::HDF5Dataset) =
    VirtualSource(filename(dset), name(dset), size(dset), eltype(dset), nothing)

Base.getindex(vs::VirtualSource, is...) = VirtualSource(vs.path, vs.name, vs.shape, vs.dtype, repcolon(is, vs.shape))

function select_hyperslab(dspace, is)
    isnothing(is) && return dspace
    start = [i[1] - 1 for i in reverse(is)]
    count = [1 for i in is]
    stride = [step(i) for i in is]
    block = [length(i) for i in reverse(is)]
    h5s_select_hyperslab(dspace, H5S_SELECT_SET, start, stride, count, block)
    return dspace
end

function d_create_virtual(parent::Union{HDF5File, HDF5Group}, name::String, layout, virtual = true)
    checkvalid(parent)
    dcpl = p_create(H5P_DATASET_CREATE)
    dspace = dataspace(layout.shape)
    dtype = datatype(layout.dtype)
    if virtual
        for (is, vs) in layout.sources
            src_dspace = dataspace(vs.shape)
            select_hyperslab(dspace, is)
            select_hyperslab(src_dspace, vs.is)
            h5p_set_virtual(dcpl, dspace, vs.path, vs.name, src_dspace)
        end
        d_create(parent, name, dtype, dspace, HDF5Properties(), dcpl)
    else
        d_create(parent, name, dtype, dspace, HDF5Properties(), dcpl)
        for (is, vs) in layout.sources
            vsis = something(vs.is, ntuple(i -> (:), length(vs.shape)))
            parent[name][is...] = h5read(vs.path, vs.name, vsis)
        end
    end
end


@testset "disjunctive.jl" begin
    @testset "getEST(task::Task)" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 3, "t1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "t2", trailer)
        p1 = 2
        p2 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2])
        processing_time = [p1,p2]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        @test constraint.active.value
        @test SeaPearl.getEST(constraint.tasks[1]) == 1
        @test SeaPearl.getEST(constraint.tasks[2]) == 2
    end

    @testset "getLCT(task::Task)" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 3, "t1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "t2", trailer)
        p1 = 2
        p2 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2])
        processing_time = [p1,p2]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        @test constraint.active.value
        @test SeaPearl.getLCT(constraint.tasks[1]) == 5
        @test SeaPearl.getLCT(constraint.tasks[2]) == 6
    end

    @testset "getECT(task::Task)" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 3, "t1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "t2", trailer)
        p1 = 2
        p2 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2])
        processing_time = [p1,p2]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

       @test constraint.active.value
        @test SeaPearl.getECT(constraint.tasks[1]) == 3
        @test SeaPearl.getECT(constraint.tasks[2]) == 5
    end

    @testset "getLST(task::Task)" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 3, "t1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "t2", trailer)
        p1 = 2
        p2 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2])
        processing_time = [p1,p2]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        @test constraint.active.value
        @test SeaPearl.getLST(constraint.tasks[1]) == 3
        @test SeaPearl.getLST(constraint.tasks[2]) == 3
    end

    @testset "function Disjunctive(earliestStartingTime::Array{IntVar},
                    processingTime::Array{Int}, trailer)" begin

        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(8, 8, "task2", trailer)
        task3 = SeaPearl.IntVar(2, 3, "task3", trailer)
        task4 = SeaPearl.IntVar(1, 9, "task4", trailer)

        p1 = 2
        p2 = 1
        p3 = 3
        p4 = 3

        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3, task4])
        processing_time = [p1, p2, p3, p4]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        @test size(constraint.tasks)[1] == 4
        @test constraint.active.value
        @test constraint in task1.onDomainChange
        @test constraint in task2.onDomainChange
        @test constraint in task3.onDomainChange
        @test constraint in task4.onDomainChange

     end

    @testset "propagate example from article" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(0, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(1, 5, "task2", trailer)
        p1 = 4
        p2 = 1
        tasks = Vector{SeaPearl.IntVar}([task1, task2])
        processing_time = [p1,p2]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        prunedDomains = SeaPearl.CPModification()

        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test res
        @test length(task1.domain) == 2
        @test length(task2.domain) == 2
        @test 0 in task1.domain
        @test 1 in task1.domain
        @test 4 in task2.domain
        @test 5 in task2.domain
    end

    @testset "propagate! with 3 tasks" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 3, "task1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(1, 5, "task3", trailer)

        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        prunedDomains = SeaPearl.CPModification()

        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test res
        @test length(task1.domain) == 3
        @test length(task2.domain) == 2
        @test length(task3.domain) == 1
        @test 1 in task1.domain
        @test 2 in task1.domain
        @test 3 in task1.domain
        @test 2 in task2.domain
        @test 3 in task2.domain
        @test 5 in task3.domain
    end

    @testset "propagate! with 4 tasks" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(8, 8, "task2", trailer)
        task3 = SeaPearl.IntVar(2, 3, "task3", trailer)
        task4 = SeaPearl.IntVar(1, 9, "task4", trailer)

        p1 = 2
        p2 = 1
        p3 = 3
        p4 = 3

        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3, task4])
        processing_time = [p1, p2, p3, p4]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        prunedDomains = SeaPearl.CPModification()

        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test res
        @test length(task1.domain) == 1
        @test length(task2.domain) == 1
        @test length(task3.domain) == 1
        @test length(task4.domain) == 1

        @test 1 in task1.domain
        @test 8 in task2.domain
        @test 3 in task3.domain
        @test 9 in task4.domain
        @test !constraint.active.value
    end

    @testset "propagate! failed" begin
        trailer = SeaPearl.Trailer()
        task1 = SeaPearl.IntVar(1, 3, "task1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(1, 4, "task3", trailer)
        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]
        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        prunedDomains = SeaPearl.CPModification()

        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test !res
    end

    @testset "propagate! failed" begin
        trailer = SeaPearl.Trailer()

        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(3, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(5, 5, "task3", trailer)

        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]

        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        prunedDomains = SeaPearl.CPModification()

        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)
        @test !res
    end

    @testset "propagate! full solving with fixed 3 tasks and Infeasible status" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(3, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(5, 5, "task3", trailer)

        p1 = 2
        p2 = 3
        p3 = 3
        SeaPearl.addVariable!(model, task1)
        SeaPearl.addVariable!(model, task2)
        SeaPearl.addVariable!(model, task3)

        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]
        SeaPearl.addConstraint!(model,SeaPearl.Disjunctive(tasks, processing_time, trailer))

        variableSelection = SeaPearl.MinDomainVariableSelection{false}()
        status = SeaPearl.solve!(model; variableHeuristic=variableSelection)

        @test status == :Infeasible
        @test length(model.statistics.solutions) == 1 # one infeasible solution
    end

    @testset "propagate! full solving with 3 tasks and Infeasible status" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(1, 5, "task3", trailer)

        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]
        SeaPearl.addConstraint!(model,SeaPearl.Disjunctive(tasks, processing_time, trailer))
        SeaPearl.addVariable!(model, task1)
        SeaPearl.addVariable!(model, task2)
        SeaPearl.addVariable!(model, task3)

        variableSelection = SeaPearl.MinDomainVariableSelection{false}()
        status = @time SeaPearl.solve!(model; variableHeuristic=variableSelection)

        @test status == :Infeasible
        @test model.statistics.numberOfSolutions == 0

    end

    @testset "propagate! chain" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        task1 = SeaPearl.IntVar(1, 3, "task1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(1, 7, "task3", trailer)

        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]

        constraint = SeaPearl.Disjunctive(tasks, processing_time, trailer)

        toPropagate = Set{SeaPearl.Constraint}([constraint])
        prunedDomains = SeaPearl.CPModification()
        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test res
        @test length(task1.domain) == 3
        @test length(task2.domain) == 2
        @test length(task3.domain) == 3
        @test 1 in task1.domain
        @test 2 in task1.domain
        @test 3 in task1.domain
        @test 2 in task2.domain
        @test 3 in task2.domain
        @test 5 in task3.domain
        @test 6 in task3.domain
        @test 7 in task3.domain

        SeaPearl.remove!(task1.domain, 3)
        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test res
        @test length(task1.domain) == 2
        @test length(task2.domain) == 1
        @test length(task3.domain) == 2
        @test 1 in task1.domain
        @test 2 in task1.domain
        @test 3 in task2.domain
        @test 6 in task3.domain
        @test 7 in task3.domain

        SeaPearl.remove!(task1.domain, 2)
        res = SeaPearl.propagate!(constraint, toPropagate, prunedDomains)
        @test res
        @test length(task1.domain) == 1
        @test length(task2.domain) == 1
        @test length(task3.domain) == 2
        @test 1 in task1.domain
        @test 3 in task2.domain
        @test 6 in task3.domain
        @test 7 in task3.domain
    end

    @testset "full solving with 3 tasks and optimal status" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(1, 6, "task3", trailer)
        SeaPearl.addVariable!(model, task1)
        SeaPearl.addVariable!(model, task2)
        SeaPearl.addVariable!(model, task3)

        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]
        SeaPearl.addConstraint!(model,SeaPearl.Disjunctive(tasks, processing_time, trailer))

        variableSelection = SeaPearl.MinDomainVariableSelection{false}()
        status = @time SeaPearl.solve!(model; variableHeuristic=variableSelection)

        @test status == :Optimal
        @test model.statistics.numberOfSolutions == 1
        @test 1 in task1.domain
        @test 3 in task2.domain
        @test 6 in task3.domain
        @test 1 == length(task1.domain)
        @test 1 == length(task2.domain)
        @test 1 == length(task3.domain)
    end

    @testset "full solving with 3 tasks, optimal status and multiple solution" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        task1 = SeaPearl.IntVar(1, 1, "task1", trailer)
        task2 = SeaPearl.IntVar(2, 3, "task2", trailer)
        task3 = SeaPearl.IntVar(1, 7, "task3", trailer)
        SeaPearl.addVariable!(model, task1)
        SeaPearl.addVariable!(model, task2)
        SeaPearl.addVariable!(model, task3)

        p1 = 2
        p2 = 3
        p3 = 3
        tasks = Vector{SeaPearl.IntVar}([task1, task2, task3])
        processing_time = [p1, p2, p3]
        SeaPearl.addConstraint!(model,SeaPearl.Disjunctive(tasks, processing_time, trailer))

        variableSelection = SeaPearl.MinDomainVariableSelection{false}()
        status = @time SeaPearl.solve!(model; variableHeuristic=variableSelection)

        @test status == :Optimal
        @test model.statistics.numberOfSolutions == 2
        @test 1 in task1.domain
        @test 3 in task2.domain
        @test 6 in task3.domain
        @test 7 in task3.domain
        @test 1 == length(task1.domain)
        @test 1 == length(task2.domain)
        @test 2 == length(task3.domain)
    end

    @testset "job shop test wiht two tasks" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)
        nbMachine = 2
        nbTask = 2
        timeLimit = 10
        tasks = Matrix(undef, nbTask, nbMachine)
        objectif = SeaPearl.IntVar(0, timeLimit, "obj", trailer)
        SeaPearl.addVariable!(model, objectif)
        processingTime = [1 4; 4 1]
        endTask = Matrix(undef, nbTask, nbMachine)

        for i in 1:nbTask
            for j in 1:nbMachine
                tasks[i,j] = SeaPearl.IntVar(0, timeLimit, "task_"*string(i)*"_"*string(j), trailer)
                endTask[i,j] = SeaPearl.IntVarViewOffset(tasks[i,j], processingTime[i,j], "end_"*string(i)*"_"*string(j))
                SeaPearl.addVariable!(model, tasks[i,j])
                SeaPearl.addConstraint!(model, SeaPearl.LessOrEqual(endTask[i,j], objectif, trailer))
            end
        end

        for i in 1:nbTask
            SeaPearl.addConstraint!(model, SeaPearl.Disjunctive([tasks[i,j] for j in 1:nbMachine],
                                                            [processingTime[i,j] for j in 1:nbMachine],
                                                            trailer))
        end
        for j in 1:nbMachine
            SeaPearl.addConstraint!(model, SeaPearl.Disjunctive([tasks[i,j] for i in 1:nbTask],
                                                            [processingTime[i,j] for i in 1:nbTask],
                                                            trailer))
        end

        SeaPearl.addObjective!(model, objectif)
        variableSelection = SeaPearl.MinDomainVariableSelection{false}()
        status = @time SeaPearl.solve!(model; variableHeuristic=variableSelection,)

        @test status == :Optimal
        @test (model.statistics.solutions[end-1]["obj"]) == 5
    end
end


const tips_file_path = String(@__DIR__) * "/TIPS_2017.nc"
const TIPS_2017_ISOT_HASH_CONST = ncread(tips_file_path, "TIPS_2017_T")
const TIPS_2017_ISOQ_HASH_CONST = ncread(tips_file_path, "TIPS_2017_Q")

T_end_idxs = zeros(Int64, size(TIPS_2017_ISOT_HASH_CONST)[1], size(TIPS_2017_ISOT_HASH_CONST)[2])
Q_end_idxs = zeros(Int64, size(TIPS_2017_ISOQ_HASH_CONST)[1], size(TIPS_2017_ISOQ_HASH_CONST)[2])

for M in 1:size(TIPS_2017_ISOT_HASH_CONST)[1]
    for I in 1:size(TIPS_2017_ISOT_HASH_CONST)[2]
        t_end_idx = findfirst(x->x==-1, TIPS_2017_ISOT_HASH_CONST[M,I,:])
        q_end_idx = findfirst(x->x==-1, TIPS_2017_ISOQ_HASH_CONST[M,I,:])
        T_end_idxs[M,I] = isnothing(t_end_idx) ? length(TIPS_2017_ISOT_HASH_CONST[M,I,:]) : t_end_idx-1
        Q_end_idxs[M,I] = isnothing(q_end_idx) ? length(TIPS_2017_ISOQ_HASH_CONST[M,I,:]) : q_end_idx-1
    end
end

get_TT(M, I) = TIPS_2017_ISOT_HASH_CONST[M, I, 1:T_end_idxs[M,I]]
get_TQ(M, I) = TIPS_2017_ISOQ_HASH_CONST[M, I, 1:Q_end_idxs[M,I]]


using SciMLBase, EllipsisNotation

"""
    NODEDataloader{T,N} <: AbstractNODEDataloader{T,N}

Struct containing batched data for sequence learning of ODEs. Can be indexed and interated over. Each batch returns a tuple (t, data(t)).

# Initialized with

    NODEData(sol::SciMLBase.AbstractTimeseriesSolution, N_length::Integer; dt=nothing, valid_set=nothing)

* `sol`: DE solution
* `N_length`: length of each batch
* `dt`: time increment that `sol` is interpolated at. If `nothing` then the `sol.t` is used as the time steps of the data
* `valid_set` if valid_set ∈ [0,1] splits the data into a train and valid set with `valid_set` of the share of the data belonging to the valid_set.

    NODEDataloader(data::AbstractArray{T,N}, t::AbstractArray{T,1}, N_length::Integer)

* `data`:: Data that is already in N_dim_1 x ... x N_t format
* `t`: time axis
* `N_length`: length of each batch

"""
struct NODEDataloader{T,U,N} <: AbstractNODEDataloader{T,U,N}
    data::AbstractArray{T,N}
    t::AbstractArray{U,1}
    N::Integer
    N_length::Integer
end

function NODEDataloader(sol::Union{SciMLBase.AbstractTimeseriesSolution, SciMLBase.AbstractDiffEqArray}, N_length::Integer; dt=nothing, valid_set=nothing)

    if isnothing(dt)
        data = DeviceArray(sol)
        t = sol.t
    else
        t = sol.t[1]:eltype(sol)(dt):sol.t[end]
        data = DeviceArray(sol(t))
    end

    if isnothing(valid_set)
        N_t = length(t)
        N = N_t - N_length

        return NODEDataloader(DeviceArray(data), t, N, N_length)
    else
        @assert 0<valid_set<1

        N_t = length(t)
        N_t_valid = Int(floor(valid_set*N_t))
        N_t_train = N_t - N_t_valid

        return NODEDataloader(DeviceArray(data[..,1:N_t_train]), t[1:N_t_train], N_t_train - N_length, N_length), NODEDataloader(DeviceArray(data[..,N_t_train+1:N_t]), t[N_t_train+1:N_t], N_t_valid - N_length, N_length)
    end
end

NODEDataloader(data::AbstractArray{T,N}, t::AbstractArray{U,1}, N_length::Integer) where {T,U,N} = NODEDataloader(DeviceArray(data), t, length(t) - N_length, N_length)
NODEDataloader(data::NODEDataloader, N_length::Integer) = NODEDataloader(data.data, data.t, N_length)

function Base.getindex(iter::NODEDataloader{T,U,N}, i::Integer) where {T,U,N}
    @assert 0 < i <= iter.N
    return (iter.t[i:i+iter.N_length-1] ,iter.data[..,i:i+iter.N_length-1])
end

function Base.iterate(iter::AbstractNODEDataloader, state=1)
    if state>iter.N
        return nothing
    else
        return (iter[state], state+1)
    end
end

Base.length(iter::AbstractNODEDataloader) = iter.N
Base.eltype(iter::AbstractNODEDataloader) = eltype(iter.data)

Base.firstindex(iter::AbstractNODEDataloader) = 1
Base.lastindex(iter::AbstractNODEDataloader) = iter.N
Base.show(io::IO,seq::NODEDataloader{T,U,N}) where {T,U,N} = print(io, "NODEData{",T,",",N,"} with ",seq.N," batches with length ",seq.N_length)

"""
    get_trajectory(data::NODEDataloader, N)

Returns a (t, x(t)) tuple of length `N`.
"""
function get_trajectory(data::NODEDataloader, N)
    ind = 1:N
    (data.t[ind], data.data[..,ind])
end

cpu(data::NODEDataloader) = NODEDataloader(Array(data.data), data.t, data.N, data.N_length)
gpu(data::NODEDataloader) = NODEDataloader(DeviceArray(data.data), data.t, data.N, data.N_length)


"""
Shorthand for `isnothing(optional) ? fallback : optional`
"""
@inline ifnothing(optional, fallback) = isnothing(optional) ? fallback : optional

function round_to_IRect2D(r::Rect{2})
    newori = round.(Int, minimum(r))
    othercorner = round.(Int, maximum(r))
    newwidth = othercorner .- newori
    Rect{2, Int}(newori, newwidth)
end

function sceneareanode!(finalbbox, limits, aspect)

    scenearea = Node(IRect(0, 0, 100, 100))

    onany(finalbbox, limits, aspect) do bbox, limits, aspect

        w = width(bbox)
        h = height(bbox)
        # as = mw / mh
        as = w / h
        mw, mh = w, h

        if aspect isa AxisAspect
            aspect = aspect.aspect
        elseif aspect isa DataAspect
            aspect = limits.widths[1] / limits.widths[2]
        end

        if !isnothing(aspect)
            if as >= aspect
                # too wide
                mw *= aspect / as
            else
                # too high
                mh *= as / aspect
            end
        end

        restw = w - mw
        resth = h - mh

        # l = left(bbox) + alignment[1] * restw
        # b = bottom(bbox) + alignment[2] * resth
        l = left(bbox) + 0.5f0 * restw
        b = bottom(bbox) + 0.5f0 * resth

        newbbox = BBox(l, l + mw, b, b + mh)

        # only update scene if pixel positions change
        new_scenearea = round_to_IRect2D(newbbox)
        if new_scenearea != scenearea[]
            scenearea[] = new_scenearea
        end
    end

    scenearea
end

function roundedrectvertices(rect, cornerradius, cornersegments)
    cr = cornerradius
    csegs = cornersegments

    cr = min(width(rect) / 2, height(rect) / 2, cr)

    # inner corners
    ictl = topleft(rect) .+ Point2(cr, -cr)
    ictr = topright(rect) .+ Point2(-cr, -cr)
    icbl = bottomleft(rect) .+ Point2(cr, cr)
    icbr = bottomright(rect) .+ Point2(-cr, cr)

    # check if corners touch so we can remove one vertex that is doubled
    wtouching = width(rect) / 2 == cr
    htouching = height(rect) / 2 == cr

    cstr = if wtouching
        anglepoint.(Ref(ictr), LinRange(0, pi/2, csegs), cr)
    else
        anglepoint.(Ref(ictr), LinRange(0, pi/2, csegs)[1:end-1], cr)
    end
    cstl = if htouching
        anglepoint.(Ref(ictl), LinRange(pi/2, pi, csegs), cr)
    else
        anglepoint.(Ref(ictl), LinRange(pi/2, pi, csegs)[1:end-1], cr)
    end
    csbl = if wtouching
        anglepoint.(Ref(icbl), LinRange(pi, 3pi/2, csegs), cr)
    else
        anglepoint.(Ref(icbl), LinRange(pi, 3pi/2, csegs)[1:end-1], cr)
    end
    csbr = if htouching
        anglepoint.(Ref(icbr), LinRange(3pi/2, 2pi, csegs), cr)
    else
        anglepoint.(Ref(icbr), LinRange(3pi/2, 2pi, csegs)[1:end-1], cr)
    end
    arr = [cstr; cstl; csbl; csbr]
end

"""
    tightlimits!(la::Axis)

Sets the autolimit margins to zero on all sides.
"""
function tightlimits!(la::Axis)
    la.xautolimitmargin = (0, 0)
    la.yautolimitmargin = (0, 0)
    autolimits!(la)
end

"""
    tightlimits!(la::Axis, sides::Union{Left, Right, Bottom, Top}...)

Sets the autolimit margins to zero on all given sides.

Example:

```
tightlimits!(laxis, Bottom())
```
"""
function tightlimits!(la::Axis, sides::Union{Left, Right, Bottom, Top}...)
    for s in sides
        tightlimits!(la, s)
    end
end

function tightlimits!(la::Axis, ::Left)
    la.xautolimitmargin = Base.setindex(la.xautolimitmargin[], 0.0, 1)
    autolimits!(la)
end

function tightlimits!(la::Axis, ::Right)
    la.xautolimitmargin = Base.setindex(la.xautolimitmargin[], 0.0, 2)
    autolimits!(la)
end

function tightlimits!(la::Axis, ::Bottom)
    la.yautolimitmargin = Base.setindex(la.yautolimitmargin[], 0.0, 1)
    autolimits!(la)
end

function tightlimits!(la::Axis, ::Top)
    la.yautolimitmargin = Base.setindex(la.yautolimitmargin[], 0.0, 2)
    autolimits!(la)
end

"""
    layoutscene(padding = 30; kwargs...)

Create a `Scene` in `campixel!` mode and a `GridLayout` aligned to the scene's pixel area with `alignmode = Outside(padding)`.
"""
function layoutscene(padding = 30; kwargs...)
    scene = Scene(; camera = campixel!, kwargs...)
    gl = GridLayout(scene, alignmode = Outside(padding))
    scene, gl
end

"""
    layoutscene(nrows::Int, ncols::Int, padding = 30; kwargs...)

Create a `Scene` in `campixel!` mode and a `GridLayout` aligned to the scene's pixel area with size `nrows` x `ncols` and `alignmode = Outside(padding)`.
"""
function layoutscene(nrows::Int, ncols::Int, padding = 30; kwargs...)
    scene = Scene(; camera = campixel!, kwargs...)
    gl = GridLayout(scene, nrows, ncols, alignmode = Outside(padding))
    scene, gl
end

GridLayoutBase.GridLayout(scene::Scene, args...; kwargs...) = GridLayout(args...; bbox = lift(x -> FRect2D(x), pixelarea(scene)), kwargs...)

function axislines!(scene, rect, spinewidth, topspinevisible, rightspinevisible,
    leftspinevisible, bottomspinevisible, topspinecolor, leftspinecolor,
    rightspinecolor, bottomspinecolor)

    bottomline = lift(rect, spinewidth) do r, sw
        y = bottom(r)
        p1 = Point2(left(r) - 0.5sw, y)
        p2 = Point2(right(r) + 0.5sw, y)
        [p1, p2]
    end

    leftline = lift(rect, spinewidth) do r, sw
        x = left(r)
        p1 = Point2(x, bottom(r) - 0.5sw)
        p2 = Point2(x, top(r) + 0.5sw)
        [p1, p2]
    end

    topline = lift(rect, spinewidth) do r, sw
        y = top(r)
        p1 = Point2(left(r) - 0.5sw, y)
        p2 = Point2(right(r) + 0.5sw, y)
        [p1, p2]
    end

    rightline = lift(rect, spinewidth) do r, sw
        x = right(r)
        p1 = Point2(x, bottom(r) - 0.5sw)
        p2 = Point2(x, top(r) + 0.5sw)
        [p1, p2]
    end

    (lines!(scene, bottomline, linewidth = spinewidth, show_axis = false,
        visible = bottomspinevisible, color = bottomspinecolor),
    lines!(scene, leftline, linewidth = spinewidth, show_axis = false,
        visible = leftspinevisible, color = leftspinecolor),
    lines!(scene, rightline, linewidth = spinewidth, show_axis = false,
        visible = rightspinevisible, color = rightspinecolor),
    lines!(scene, topline, linewidth = spinewidth, show_axis = false,
        visible = topspinevisible, color = topspinecolor))
end

function interleave_vectors(vec1::Vector{T}, vec2::Vector{T}) where T
    n = length(vec1)
    @assert n == length(vec2)

    vec = Vector{T}(undef, 2 * n)
    @inbounds for i in 1:n
        k = 2(i - 1)
        vec[k + 1] = vec1[i]
        vec[k + 2] = vec2[i]
    end
    vec
end

"""
Take a sequence of variable definitions with docstrings above each and turn
them into Attributes, a Dict of varname => docstring pairs and a Dict of
varname => default_value pairs.

# Example

    attrs, docdict, defaultdict = @documented_attributes begin
        "The width."
        width = 10
        "The height."
        height = 20 + x
    end

    attrs == Attributes(
        width = 10,
        height = 20
    )

    docdict == Dict(
        width => "The width.",
        height => "The height."
    )

    defaultdict == Dict(
        width => "10",
        height => "20 + x"
    )
"""
macro documented_attributes(exp)
    if exp.head != :block
        error("Not a block")
    end

    expressions = filter(x -> !(x isa LineNumberNode), exp.args)

    vars_and_exps = map(expressions) do e
        if e.head == :macrocall && e.args[1] == GlobalRef(Core, Symbol("@doc"))
            varname = e.args[4].args[1]
            var_exp = e.args[4].args[2]
            str_exp = e.args[3]
        elseif e.head == Symbol("=")
            varname = e.args[1]
            var_exp = e.args[2]
            str_exp = "no description"
        else
            error("Neither docstringed variable nor normal variable: $e")
        end
        varname, var_exp, str_exp
    end

    # make a dictionary of :variable_name => docstring_expression
    exp_docdict = Expr(:call, :Dict,
        (Expr(:call, Symbol("=>"), QuoteNode(name), strexp)
            for (name, _, strexp) in vars_and_exps)...)

    # make a dictionary of :variable_name => docstring_expression
    defaults_dict = Expr(:call, :Dict,
        (Expr(:call, Symbol("=>"), QuoteNode(name), exp isa String ? "\"$exp\"" : string(exp))
            for (name, exp, _) in vars_and_exps)...)

    # make an Attributes instance with of variable_name = variable_expression
    exp_attrs = Expr(:call, :Attributes,
        (Expr(:kw, name, exp)
            for (name, exp, _) in vars_and_exps)...)

    esc(quote
        ($exp_attrs, $exp_docdict, $defaults_dict)
    end)
end

"""
Turn a combination of docdict and defaultdict from `@documented_attributes`
into a string to insert into a docstring.
"""
function docvarstring(docdict, defaultdict)
    buffer = IOBuffer()
    maxwidth = maximum(length ∘ string, keys(docdict))
    for (var, doc) in sort(collect(pairs(docdict)))
        print(buffer, "`$var`\\\nDefault: `$(defaultdict[var])`\\\n$doc\n\n")
    end
    String(take!(buffer))
end

function subtheme(scene, key::Symbol)
    sub = haskey(theme(scene), key) ? theme(scene, key) : Attributes()
    if !(sub isa Attributes)
        error("Subtheme is not of type Attributes but is $sub")
    end
    sub
end

"""
    xaxis_top!(la::Axis)

Move the x-axis to the top, while correctly aligning the tick labels at the bottom.
"""
function xaxis_top!(la::Axis)
    la.xaxisposition = :top
    la.xticklabelalign = (la.xticklabelalign[][1], :bottom)
    nothing
end

"""
    xaxis_bottom!(la::Axis)

Move the x-axis to the bottom, while correctly aligning the tick labels at the top.
"""
function xaxis_bottom!(la::Axis)
    la.xaxisposition = :bottom
    la.xticklabelalign = (la.xticklabelalign[][1], :top)
    nothing
end

"""
    yaxis_left!(la::Axis)

Move the y-axis to the left, while correctly aligning the tick labels at the right.
"""
function yaxis_left!(la::Axis)
    la.yaxisposition = :left
    la.yticklabelalign = (:right, la.yticklabelalign[][2])
    nothing
end

"""
    yaxis_right!(la::Axis)

Move the y-axis to the right, while correctly aligning the tick labels at the left.
"""
function yaxis_right!(la::Axis)
    la.yaxisposition = :right
    la.yticklabelalign = (:left, la.yticklabelalign[][2])
    nothing
end

"""
    labelslider!(scene, label, range; format = string, sliderkw = Dict(), labelkw = Dict(), valuekw = Dict(), layoutkw...)

Construct a horizontal GridLayout with a label, a slider and a value label in `scene`.

Returns a `NamedTuple`:

`(slider = slider, label = label, valuelabel = valuelabel, layout = layout)`

Specify a format function for the value label with the `format` keyword.
The slider is forwarded the keywords from `sliderkw`.
The label is forwarded the keywords from `labelkw`.
The value label is forwarded the keywords from `valuekw`.
All other keywords are forwarded to the `GridLayout`.

Example:

```
ls = labelslider!(scene, "Voltage:", 0:10; format = x -> "\$(x)V")
layout[1, 1] = ls.layout
```
"""
function labelslider!(scene, label, range; format = string,
        sliderkw = Dict(), labelkw = Dict(), valuekw = Dict(), layoutkw...)
    slider = Slider(scene; range = range, sliderkw...)
    label = Label(scene, label; labelkw...)
    valuelabel = Label(scene, lift(format, slider.value); valuekw...)
    layout = hbox!(label, slider, valuelabel; layoutkw...)
    (slider = slider, label = label, valuelabel = valuelabel, layout = layout)
end

"""
    labelslidergrid!(scene, labels, ranges; formats = [string],
        sliderkw = Dict(), labelkw = Dict(), valuekw = Dict(), layoutkw...)

Construct a GridLayout with a column of label, a column of sliders and a column of value labels in `scene`.
The argument values are broadcast, so you can use scalars if you want to keep labels, ranges or formats constant across rows.

Returns a `NamedTuple`:

`(sliders = sliders, labels = labels, valuelabels = valuelabels, layout = layout)`

Specify format functions for the value labels with the `formats` keyword.
The sliders are forwarded the keywords from `sliderkw`.
The labels are forwarded the keywords from `labelkw`.
The value labels are forwarded the keywords from `valuekw`.
All other keywords are forwarded to the `GridLayout`.

Example:

```
ls = labelslidergrid!(scene, ["Voltage", "Ampere"], Ref(0:0.1:100); format = x -> "\$(x)V")
layout[1, 1] = ls.layout
```
"""
function labelslidergrid!(scene, labels, ranges; formats = [string],
        sliderkw = Dict(), labelkw = Dict(), valuekw = Dict(), layoutkw...)

    elements = broadcast(labels, ranges, formats) do label, range, format
        slider = Slider(scene; range = range, sliderkw...)
        label = Label(scene, label; halign = :left, labelkw...)
        valuelabel = Label(scene, lift(format, slider.value); halign = :right, valuekw...)
        (; slider = slider, label = label, valuelabel = valuelabel)
    end

    sliders = map(x -> x.slider, elements)
    labels = map(x -> x.label, elements)
    valuelabels = map(x -> x.valuelabel, elements)

    layout = grid!(hcat(labels, sliders, valuelabels); layoutkw...)

    (sliders = sliders, labels = labels, valuelabels = valuelabels, layout = layout)
end

# helper function to create either h or vlines depending on `direction`
# this works only with LAxes because it needs to react to limit changes
function hvlines!(ax::Axis, direction::Int, datavals, axmins, axmaxs; attributes...)

    datavals, axmins, axmaxs = map(x -> x isa Observable ? x : Observable(x), (datavals, axmins, axmaxs))

    linesegs = lift(ax.limits, ax.scene.px_area, datavals, axmins, axmaxs) do lims, pxa,
            datavals, axmins, axmaxs

        xlims = (minimum(lims)[direction], maximum(lims)[direction])
        xfrac(f) = xlims[1] + f * (xlims[2] - xlims[1])
        segs = broadcast(datavals, axmins, axmaxs) do dataval, axmin, axmax
            if direction == 1
                (Point2f0(xfrac(axmin), dataval), Point2f0(xfrac(axmax), dataval))
            elseif direction == 2
                (Point2f0(dataval, xfrac(axmin)), Point2f0(dataval, xfrac(axmax)))
            else
                error("direction must be 1 or 2")
            end
        end
        # handle case that none of the inputs is an array, but we need an array for linesegments!
        if segs isa Tuple
            segs = [segs]
        end
        segs
    end

    linesegments!(ax, linesegs; xautolimits = direction == 2, yautolimits = direction == 1, attributes...)
end

"""
    hlines!(ax::Axis, ys; xmin = 0.0, xmax = 1.0, attrs...)

Create horizontal lines across `ax` at `ys` in data coordinates and `xmin` to `xmax`
in axis coordinates (0 to 1). All three of these can have single or multiple values because
they are broadcast to calculate the final line segments.
"""
hlines!(ax::Axis, ys; xmin = 0.0, xmax = 1.0, attrs...) =
    hvlines!(ax, 1, ys, xmin, xmax; attrs...)

"""
    vlines!(ax::Axis, xs; ymin = 0.0, ymax = 1.0, attrs...)

Create vertical lines across `ax` at `xs` in data coordinates and `ymin` to `ymax`
in axis coordinates (0 to 1). All three of these can have single or multiple values because
they are broadcast to calculate the final line segments.
"""
vlines!(ax::Axis, xs; ymin = 0.0, ymax = 1.0, attrs...) =
    hvlines!(ax, 2, xs, ymin, ymax; attrs...)


import ACEbase: Discrete1pBasis
import ACE: evaluate, evaluate!

export Categorical1pBasis

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

struct SList{N, T}
   list::SVector{N, T}

   function SList{N, T}(list::SVector{N, T})  where {N, T}
      if isabstracttype(T)
         error("`SList` can only contain a single type")
      end
      return new(list)
   end
end

SList(list::AbstractArray) = SList(SVector(list...))
SList(args...) = SList(SVector(args...))
SList(list::SVector{N, T}) where {N, T} = SList{N, T}(list)

Base.length(list::SList) = length(list.list)
Base.rand(list::SList) = rand(list.list)

i2val(list::SList, i::Integer) = list.list[i]

Base.iterate(list::SList, args...) = iterate(list.list, args...)

function val2i(list::SList, val)
   for j = 1:length(list)
      if list.list[j] == val
         return j
      end
   end
   error("val = $val not found in this list")
end

write_dict(list::SList{N,T}) where {N, T} =
      Dict( "__id__" => "ACE_SList",
                 "T" => write_dict(T),
              "list" => list.list )

function read_dict(::Val{:ACE_SList}, D::Dict)
   list = D["list"]
   T = read_dict(D["T"])
   svector = SVector{length(list), T}((T.(list))...)
   return SList(svector)
end

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

@doc raw"""
`Categorical1pBasis` : defines the discrete 1p basis
```math
   \phi_q(u) = \delta(u - U_q),
```
where ``U_q, q = 1, \dots, Q`` are finitely many possible values that the
variable ``u`` may take. Suppose, e.g., we allow the values `[:a, :b, :c]`,
then
```julia
P = Categorical1pBasis([:a, :b, :c]; varsym = :u, idxsym = :q)
evaluate(P, State(u = :a))     # Bool[1, 0, 0]
evaluate(P, State(u = :b))     # Bool[0, 1, 0]
evaluate(P, State(u = :c))     # Bool[0, 0, 1]
```
If we evaluate it with an unknown state we get an error:
```julia
evaluate(P, State(u = :x))
# Error : val = x not found in this list
```
"""
struct Categorical1pBasis{VSYM, ISYM, LEN, T} <: Discrete1pBasis{Bool}
   categories::SList{LEN, T}
   label::String
end

_varsym(::Categorical1pBasis{VSYM, ISYM}) where {VSYM, ISYM} = VSYM
_isym(::Categorical1pBasis{VSYM, ISYM}) where {VSYM, ISYM} = ISYM

_val(X, B::Categorical1pBasis) = getproperty(X, _varsym(B))
_idx(b, B::Categorical1pBasis) = getproperty(b, _isym(B))

Base.length(B::Categorical1pBasis) = length(B.categories)

Categorical1pBasis(categories::AbstractArray;
              varsym::Symbol = nothing, idxsym::Symbol = nothing,
              label = "C$(idxsym)") =
      Categorical1pBasis(categories, varsym, idxsym, label)

Categorical1pBasis(categories::AbstractArray, varsym::Symbol, isym::Symbol, label::String) =
      Categorical1pBasis(SList(categories), varsym, isym, label)

Categorical1pBasis(categories::SList{LEN, T}, varsym::Symbol, isym::Symbol, label::String) where {LEN, T} =
      Categorical1pBasis{varsym, isym, LEN, T}(categories, label)

function evaluate(basis::Categorical1pBasis, X::AbstractState)
   A = Vector{Bool}(undef, length(basis))
   return evaluate!(A, basis, X)
end

function ACE.evaluate!(A, basis::Categorical1pBasis, X::AbstractState)
   fill!(A, false)
   A[val2i(basis.categories, _val(X, basis))] = true
   return A
end

symbols(basis::Categorical1pBasis) = [ _isym(basis), ]

indexrange(basis::Categorical1pBasis) = Dict( _isym(basis) => basis.categories.list )

isadmissible(b, basis::Categorical1pBasis) = (_idx(b, basis) in basis.categories)

get_index(B::Categorical1pBasis, b) = val2i(B.categories, _idx(b, B) )

degree(b, basis::Categorical1pBasis, args...) = 0

Base.rand(basis::Categorical1pBasis) = rand(basis.list)

function get_spec(basis::Categorical1pBasis, i)
   return NamedTuple{(_isym(basis),)}((i2val(basis.categories, i),))
end

get_spec(basis::Categorical1pBasis) = [ get_spec(basis, i) for i = 1:length(basis) ]

write_dict(B::Categorical1pBasis) =
      Dict( "__id__" => "ACE_Categorical1pBasis",
            "categories" => write_dict(B.categories),
            "VSYM" => String(_varsym(B)),
            "ISYM" => String(_isym(B)),
            "label" => B.label)

read_dict(::Val{:ACE_Categorical1pBasis}, D::Dict)  =
   Categorical1pBasis( read_dict(D["categories"]),
                  Symbol(D["VSYM"]), Symbol(D["ISYM"]),
                  D["label"] )


export RandomVariable, randomvar, simplerandomvar

## Random variable implementation

mutable struct RandomVariable <: AbstractVariable
    name                :: Symbol
    inputmsgs           :: Vector{LazyObservable{AbstractMessage}}
    marginal            :: Union{Nothing, MarginalObservable}
    pipeline            :: AbstractPipelineStage
    local_constraint
    prod_constraint
    prod_strategy
    form_constraint
    form_check_strategy
end

function randomvar(name::Symbol; pipeline = EmptyPipelineStage(), local_constraint = Marginalisation(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    return RandomVariable(name, Vector{LazyObservable{AbstractMessage}}(), nothing, pipeline, local_constraint, prod_constraint, prod_strategy, form_constraint, form_check_strategy)
end

function randomvar(name::Symbol, dims::Tuple; pipeline = EmptyPipelineStage(), local_constraint = Marginalisation(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    return randomvar(name, dims...; pipeline = pipeline, local_constraint = local_constraint, prod_constraint = prod_constraint, prod_strategy = prod_strategy, form_constraint = form_constraint, form_check_strategy = form_check_strategy)
end

function randomvar(name::Symbol, dims::Vararg{Int}; pipeline = EmptyPipelineStage(), local_constraint = Marginalisation(), prod_constraint = ProdAnalytical(), prod_strategy = FoldLeftProdStrategy(), form_constraint = UnspecifiedFormConstraint(), form_check_strategy = FormConstraintCheckPickDefault())
    vars = Array{RandomVariable}(undef, dims)
    for i in CartesianIndices(axes(vars))
        @inbounds vars[i] = randomvar(Symbol(name, :_, Symbol(join(i.I, :_))); pipeline = pipeline, local_constraint = local_constraint, prod_constraint = prod_constraint, prod_strategy = prod_strategy, form_constraint = form_constraint, form_check_strategy = form_check_strategy)
    end
    return vars
end

degree(randomvar::RandomVariable)              = length(randomvar.inputmsgs)
name(randomvar::RandomVariable)                = randomvar.name
local_constraint(randomvar::RandomVariable)    = randomvar.local_constraint
prod_constraint(randomvar::RandomVariable)     = randomvar.prod_constraint
prod_strategy(randomvar::RandomVariable)       = randomvar.prod_strategy
form_constraint(randomvar::RandomVariable)     = randomvar.form_constraint
form_check_strategy(randomvar::RandomVariable) = _form_check_strategy(randomvar.form_check_strategy, randomvar)

_form_check_strategy(::FormConstraintCheckPickDefault, randomvar::RandomVariable) = default_form_check_strategy(form_constraint(randomvar))
_form_check_strategy(form_check_strategy, randomvar::RandomVariable)              = form_check_strategy

messages_prod_fn(randomvar::RandomVariable) = messages_prod_fn(prod_strategy(randomvar), prod_constraint(randomvar), UnspecifiedFormConstraint(), default_form_check_strategy(UnspecifiedFormConstraint()))
marginal_prod_fn(randomvar::RandomVariable) = marginal_prod_fn(prod_strategy(randomvar), prod_constraint(randomvar), form_constraint(randomvar), form_check_strategy(randomvar))

getlastindex(randomvar::RandomVariable) = degree(randomvar) + 1

messagein(randomvar::RandomVariable, index::Int)  = @inbounds randomvar.inputmsgs[index]
messageout(randomvar::RandomVariable, index::Int) = collectLatest(AbstractMessage, Message, skipindex(randomvar.inputmsgs, index), messages_prod_fn(randomvar))

get_pipeline_stages(randomvar::RandomVariable)        = randomvar.pipeline
add_pipeline_stage!(randomvar::RandomVariable, stage) = randomvar.pipeline = (randomvar.pipeline + stage)

_getmarginal(randomvar::RandomVariable)                                = randomvar.marginal
_setmarginal!(randomvar::RandomVariable, marginal::MarginalObservable) = randomvar.marginal = marginal
_makemarginal(randomvar::RandomVariable)                               = collectLatest(AbstractMessage, Marginal, randomvar.inputmsgs, marginal_prod_fn(randomvar))

function setmessagein!(randomvar::RandomVariable, index::Int, messagein)
    if index === degree(randomvar) + 1
        push!(randomvar.inputmsgs, messagein)
    else
        throw_setmessagein!(randomvar, index)
    end
end

# https://github.com/JuliaLang/julia/issues/29688
throw_setmessagein!(randomvar, index) = throw("Inconsistent state in setmessagein! function for random variable $(randomvar). `index` should be equal to `degree(randomvar) + 1 = $(degree(randomvar) + 1)`, $(index) is given instead")


@testset "Left Division" begin
    mat = diagm([1,2,3])
    v = [2,3,4]
    tv = v.'

    @test_throws Exception mat \ tv
end


module OpenBLAS32

using LinearAlgebra.BLAS
using OpenBLAS32_jll

__init__() = BLAS.lbt_forward(OpenBLAS32_jll.libopenblas_path)

end # module


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

module PowerModelsACDC

# import Compat
import JuMP
import Memento
import PowerModels
const _PM = PowerModels
import InfrastructureModels
# import InfrastructureModels: ids, ref, var, con, sol, nw_ids, nws, optimize_model!, @im_fields
const _IM = InfrastructureModels

import JuMP: with_optimizer
export with_optimizer

# Create our module level logger (this will get precompiled)
const _LOGGER = Memento.getlogger(@__MODULE__)

# Register the module level logger at runtime so that folks can access the logger via `getlogger(PowerModels)`
# NOTE: If this line is not included then the precompiled `_PM._LOGGER` won't be registered at runtime.
__init__() = Memento.register(_LOGGER)

include("prob/acdcopf.jl")
include("prob/acdcpf.jl")
include("prob/acdcopf_bf.jl")
include("prob/tnepopf.jl")
include("prob/acdctnepopf.jl")
include("prob/tnepopf_bf.jl")
include("prob/acdctnepopf_bf.jl")
include("prob/mp_tnepopf.jl")
include("prob/mp_acdctnepopf.jl")
include("prob/mp_tnepopf_bf.jl")
include("prob/mp_acdctnepopf_bf.jl")

include("core/solution.jl")
include("core/data.jl")
include("core/variabledcgrid.jl")
include("core/variableconv.jl")
include("core/base.jl")
include("core/constraint.jl")
include("core/objective.jl")
include("core/relaxation_scheme.jl")

include("formdcgrid/acp.jl")
include("formdcgrid/dcp.jl")
include("formdcgrid/wr.jl")
include("formdcgrid/wrm.jl")
include("formdcgrid/bf.jl")
include("formdcgrid/lpac.jl")
include("formdcgrid/shared.jl")

include("formconv/acp.jl")
include("formconv/dcp.jl")
include("formconv/wr.jl")
include("formconv/wrm.jl")
include("formconv/bf.jl")
include("formconv/lpac.jl")
include("formconv/shared.jl")

include("core/constraint_template.jl")
include("io/multinetwork.jl")
include("io/results.jl")
end


module ParametricOptInterface

using MathOptInterface

const MOI = MathOptInterface

const PARAMETER_INDEX_THRESHOLD = 1_000_000_000_000_000_000
const SUPPORTED_SETS =
    (MOI.LessThan{Float64}, MOI.EqualTo{Float64}, MOI.GreaterThan{Float64})
const SUPPORTED_VECTOR_SETS = (MOI.Nonnegatives, MOI.SecondOrderCone)

"""
    Parameter(::Float64)

The `Parameter` structure stores the numerical value associated to a given
parameter.

## Example

```julia-repl
julia> ParametricOptInterface.Parameter(5)
ParametricOptInterface.Parameter(5)
```
"""
struct Parameter <: MOI.AbstractScalarSet
    val::Float64
end

"""
    Optimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer

Declares a `Optimizer`, which allows the handling of parameters in a
optimization model.

## Example

```julia-repl
julia> ParametricOptInterface.Optimizer(GLPK.Optimizer())
ParametricOptInterface.Optimizer{Float64,GLPK.Optimizer}
```
"""
mutable struct Optimizer{T,OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
    optimizer::OT
    parameters::Dict{MOI.VariableIndex,T}
    parameters_name::Dict{MOI.VariableIndex,String}
    updated_parameters::Dict{MOI.VariableIndex,T}
    variables::Dict{MOI.VariableIndex,MOI.VariableIndex}
    last_variable_index_added::Int64
    last_parameter_index_added::Int64
    affine_constraint_cache::MOI.Utilities.DoubleDicts.DoubleDict{
        Vector{MOI.ScalarAffineTerm{Float64}},
    }
    quadratic_constraint_cache_pv::Dict{
        MOI.ConstraintIndex,
        Vector{MOI.ScalarQuadraticTerm{Float64}},
    }
    quadratic_constraint_cache_pp::Dict{
        MOI.ConstraintIndex,
        Vector{MOI.ScalarQuadraticTerm{Float64}},
    }
    quadratic_constraint_cache_pc::MOI.Utilities.DoubleDicts.DoubleDict{
        Vector{MOI.ScalarAffineTerm{Float64}},
    }
    quadratic_constraint_variables_associated_to_parameters_cache::Dict{
        MOI.ConstraintIndex,
        Vector{MOI.ScalarAffineTerm{T}},
    }
    quadratic_added_cache::Dict{MOI.ConstraintIndex,MOI.ConstraintIndex}
    last_quad_add_added::Int64
    vector_constraint_cache::MOI.Utilities.DoubleDicts.DoubleDict{
        Vector{MOI.VectorAffineTerm{Float64}},
    }
    affine_objective_cache::Vector{MOI.ScalarAffineTerm{T}}
    quadratic_objective_cache_pv::Vector{MOI.ScalarQuadraticTerm{T}}
    quadratic_objective_cache_pp::Vector{MOI.ScalarQuadraticTerm{T}}
    quadratic_objective_cache_pc::Vector{MOI.ScalarAffineTerm{T}}
    quadratic_objective_variables_associated_to_parameters_cache::Vector{
        MOI.ScalarAffineTerm{T},
    }
    multiplicative_parameters::Set{Int64}
    dual_value_of_parameters::Vector{Float64}
    evaluate_duals::Bool
    number_of_parameters_in_model::Int64
    function Optimizer(optimizer::OT; evaluate_duals::Bool = true) where {OT}
        return new{Float64,OT}(
            optimizer,
            Dict{MOI.VariableIndex,Float64}(),
            Dict{MOI.VariableIndex,String}(),
            Dict{MOI.VariableIndex,Float64}(),
            Dict{MOI.VariableIndex,MOI.VariableIndex}(),
            0,
            PARAMETER_INDEX_THRESHOLD,
            MOI.Utilities.DoubleDicts.DoubleDict{
                Vector{MOI.ScalarAffineTerm{Float64}},
            }(),
            Dict{MOI.ConstraintIndex,Vector{MOI.ScalarQuadraticTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex,Vector{MOI.ScalarQuadraticTerm{Float64}}}(),
            MOI.Utilities.DoubleDicts.DoubleDict{
                Vector{MOI.ScalarAffineTerm{Float64}},
            }(),
            Dict{MOI.ConstraintIndex,Vector{MOI.ScalarAffineTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex,MOI.ConstraintIndex}(),
            0,
            MOI.Utilities.DoubleDicts.DoubleDict{
                Vector{MOI.VectorAffineTerm{Float64}},
            }(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Vector{MOI.ScalarQuadraticTerm{Float64}}(),
            Vector{MOI.ScalarQuadraticTerm{Float64}}(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Set{Int64}(),
            Vector{Float64}(),
            evaluate_duals,
            0,
        )
    end
end

function ParametricOptimizer(args...; kwargs...)
    @warn("ParametricOptimizer is deprecated. Use `POI.Optimizer` instead.")
    return Optimizer(args...; kwargs...)
end

include("utils.jl")
include("duals.jl")
include("update_parameters.jl")

function MOI.is_empty(model::Optimizer)
    return MOI.is_empty(model.optimizer) &&
           isempty(model.parameters) &&
           isempty(model.parameters_name) &&
           isempty(model.variables) &&
           isempty(model.updated_parameters) &&
           isempty(model.variables) &&
           model.last_variable_index_added == 0 &&
           model.last_parameter_index_added == PARAMETER_INDEX_THRESHOLD &&
           isempty(model.affine_constraint_cache) &&
           isempty(model.quadratic_constraint_cache_pv) &&
           isempty(model.quadratic_constraint_cache_pp) &&
           isempty(model.quadratic_constraint_cache_pc) &&
           isempty(
               model.quadratic_constraint_variables_associated_to_parameters_cache,
           ) &&
           isempty(model.quadratic_added_cache) &&
           model.last_quad_add_added == 0 &&
           isempty(model.affine_objective_cache) &&
           isempty(model.quadratic_objective_cache_pv) &&
           isempty(model.quadratic_objective_cache_pp) &&
           isempty(model.quadratic_objective_cache_pc) &&
           isempty(
               model.quadratic_objective_variables_associated_to_parameters_cache,
           ) &&
           isempty(model.dual_value_of_parameters) &&
           model.number_of_parameters_in_model == 0
end

function MOI.supports_constraint(
    model::Optimizer,
    F::Union{
        Type{MOI.SingleVariable},
        Type{MOI.ScalarAffineFunction{T}},
        Type{MOI.VectorOfVariables},
        Type{MOI.VectorAffineFunction{T}},
    },
    S::Type{<:MOI.AbstractSet},
) where {T}
    return MOI.supports_constraint(model.optimizer, F, S)
end

function MOI.supports_constraint(
    model::Optimizer,
    ::Type{MOI.ScalarQuadraticFunction{T}},
    S::Type{<:MOI.AbstractSet},
) where {T}
    return MOI.supports_constraint(
        model.optimizer,
        MOI.ScalarAffineFunction{T},
        S,
    )
end

function MOI.supports(
    model::Optimizer,
    attr::Union{
        MOI.ObjectiveSense,
        MOI.ObjectiveFunction{MOI.SingleVariable},
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}},
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}},
    },
) where {T}
    return MOI.supports(model.optimizer, attr)
end

function MOI.empty!(model::Optimizer{T}) where {T}
    MOI.empty!(model.optimizer)
    empty!(model.parameters)
    empty!(model.parameters_name)
    empty!(model.updated_parameters)
    empty!(model.variables)
    model.last_variable_index_added = 0
    model.last_parameter_index_added = PARAMETER_INDEX_THRESHOLD
    empty!(model.affine_constraint_cache)
    empty!(model.quadratic_constraint_cache_pv)
    empty!(model.quadratic_constraint_cache_pp)
    empty!(model.quadratic_constraint_cache_pc)
    empty!(model.quadratic_constraint_variables_associated_to_parameters_cache)
    empty!(model.quadratic_added_cache)
    model.last_quad_add_added = 0
    empty!(model.vector_constraint_cache)
    empty!(model.affine_objective_cache)
    empty!(model.quadratic_objective_cache_pv)
    empty!(model.quadratic_objective_cache_pp)
    empty!(model.quadratic_objective_cache_pc)
    empty!(model.quadratic_objective_variables_associated_to_parameters_cache)
    empty!(model.dual_value_of_parameters)
    model.number_of_parameters_in_model = 0
    return
end

function MOI.set(
    model::Optimizer,
    attr::MOI.VariableName,
    v::MOI.VariableIndex,
    name::String,
)
    if is_parameter_in_model(model, v)
        model.parameters_name[v] = name
    else
        MOI.set(model.optimizer, attr, v, name)
    end
    return
end

function MOI.get(model::Optimizer, attr::MOI.VariableName, v::MOI.VariableIndex)
    if is_parameter_in_model(model, v)
        return get(model.parameters_name, v, "")
    else
        return MOI.get(model.optimizer, attr, v)
    end
end

function MOI.supports(
    model::Optimizer,
    attr::MOI.VariableName,
    tp::Type{MOI.VariableIndex},
)
    return MOI.supports(model.optimizer, attr, tp)
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ConstraintName,
    c::MOI.ConstraintIndex,
    name::String,
)
    MOI.set(model.optimizer, attr, c, name)
    return
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintName,
    c::MOI.ConstraintIndex,
)
    return MOI.get(model.optimizer, attr, c)
end

function MOI.get(model::Optimizer, ::MOI.NumberOfVariables)
    return length(model.parameters) + length(model.variables)
end

function MOI.supports(
    model::Optimizer,
    attr::MOI.ConstraintName,
    tp::Type{<:MOI.ConstraintIndex},
)
    return MOI.supports(model.optimizer, attr, tp)
end

# TODO
# This is not correct, you need to put the parameters back into the function
# function MOI.get(model::Optimizer, attr::MOI.ConstraintFunction, ci::MOI.ConstraintIndex{F, S}) where {F, S}
#     MOI.get(model.optimizer, attr, ci)
# end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    MOI.throw_if_not_valid(model, ci)
    return MOI.get(model.optimizer, attr, ci)
end

function MOI.get(model::Optimizer, attr::MOI.ObjectiveSense)
    return MOI.get(model.optimizer, attr)
end

function MOI.get(model::Optimizer, attr::MOI.ObjectiveFunctionType)
    if !isempty(model.quadratic_constraint_cache_pv) ||
       !isempty(model.quadratic_constraint_cache_pc)
        return MOI.ScalarQuadraticFunction{Float64}
    end
    return MOI.get(model.optimizer, attr)
end

# TODO
# Same as ConstraintFunction getter. And you also need to convert to F
# function MOI.get(
#     model::Optimizer,
#     attr::MOI.ObjectiveFunction{F}) where F <: Union{MOI.SingleVariable,MOI.ScalarAffineFunction{T}} where T

#     MOI.get(model.optimizer, attr)
# end

# TODO
# You might have transformed quadratic functions into affine functions so this is incorrect
# function MOI.get(model::Optimizer, ::MOI.ListOfConstraints)
#     constraints = Set{Tuple{DataType, DataType}}()
#     inner_ctrs = MOI.get(model.optimizer, MOI.ListOfConstraints())
#     for (F, S) in inner_ctrs
#         push!(constraints, (F,S))
#     end

#     collect(constraints)
# end

function MOI.get(
    model::Optimizer,
    attr::MOI.ListOfConstraintIndices{F,S},
) where {S,F<:Union{MOI.VectorOfVariables,MOI.SingleVariable}}
    return MOI.get(model.optimizer, attr)
end

# TODO
# You might have transformed quadratic functions into affine functions so this is incorrect
# function MOI.get(
#     model::Optimizer,
#     ::MOI.ListOfConstraintIndices{F, S}
# ) where {S<:MOI.AbstractSet, F<:Union{
#     MOI.ScalarAffineFunction{T},
#     MOI.VectorAffineFunction{T},
# }} where T
#     MOI.get(model.optimizer, MOI.ListOfConstraintIndices{F, S}())
# end

function MOI.supports_add_constrained_variable(::Optimizer, ::Type{Parameter})
    return true
end

function MOI.supports_add_constrained_variables(
    model::Optimizer,
    ::Type{MOI.Reals},
)
    return MOI.supports_add_constrained_variables(model.optimizer, MOI.Reals)
end

function MOI.add_variable(model::Optimizer)
    next_variable_index!(model)
    v_p = MOI.VariableIndex(model.last_variable_index_added)
    v = MOI.add_variable(model.optimizer)
    model.variables[v_p] = v
    return v_p
end

function MOI.add_constrained_variable(model::Optimizer, set::Parameter)
    next_parameter_index!(model)
    p = MOI.VariableIndex(model.last_parameter_index_added)
    model.parameters[p] = set.val
    cp = MOI.ConstraintIndex{MOI.SingleVariable,Parameter}(
        model.last_parameter_index_added,
    )
    update_number_of_parameters!(model)
    return p, cp
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.SingleVariable,
    set::MOI.AbstractScalarSet,
)
    if is_parameter_in_model(model, f.variable)
        error("Cannot constrain a parameter")
    elseif !is_variable_in_model(model, f.variable)
        error("Variable not in the model")
    end
    return MOI.add_constraint(model.optimizer, f, set)
end

function add_constraint_with_parameters_on_function(
    model::Optimizer,
    f::MOI.ScalarAffineFunction{T},
    set::MOI.AbstractScalarSet,
) where {T}
    vars, params, param_constant =
        separate_possible_terms_and_calculate_parameter_constant(model, f.terms)
    ci = MOI.Utilities.normalize_and_add_constraint(
        model.optimizer,
        MOI.ScalarAffineFunction(vars, f.constant + param_constant),
        set,
    )
    model.affine_constraint_cache[ci] = params
    return ci
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.ScalarAffineFunction{T},
    set::MOI.AbstractScalarSet,
) where {T}
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function MOI.get(
    model::Optimizer,
    attr::MOI.VariablePrimal,
    v::MOI.VariableIndex,
)
    if is_parameter_in_model(model, v)
        return model.parameters[v]
    elseif is_variable_in_model(model, v)
        return MOI.get(model.optimizer, attr, model.variables[v])
    else
        error("Variable not in the model")
    end
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintSet,
    cp::MOI.ConstraintIndex{MOI.SingleVariable,Parameter},
    set::Parameter,
)
    p = MOI.VariableIndex(cp.value)
    if !is_parameter_in_model(model, p)
        error("Parameter not in the model")
    end
    return model.updated_parameters[p] = set.val
end

struct ParameterValue <: MOI.AbstractVariableAttribute end

"""
    MOI.set(
        model::Optimizer,
        ::ParameterValue,
        x::MOI.VariableIndex,
        value::Float64,
    )

Sets the parameter to a given value, using its `MOI.VariableIndex` as reference.

# Example

```julia-repl
julia> MOI.set(model, ParameterValue(), w, 2.0)
2.0
```
"""
function MOI.set(
    model::Optimizer,
    ::ParameterValue,
    vi::MOI.VariableIndex,
    val::Float64,
)
    if !is_parameter_in_model(model, vi)
        error("Parameter not in the model")
    end
    return model.updated_parameters[vi] = val
end

function MOI.set(
    model::Optimizer,
    ::ParameterValue,
    vi::MOI.VariableIndex,
    val::Real,
)
    return MOI.set(model, ParameterValue(), vi, convert(Float64, val))
end

function empty_objective_function_caches!(model::Optimizer)
    empty!(model.affine_objective_cache)
    empty!(model.quadratic_objective_cache_pv)
    empty!(model.quadratic_objective_cache_pp)
    empty!(model.quadratic_objective_cache_pc)
    return
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction,
    f::MOI.ScalarAffineFunction{T},
) where {T}
    # clear previously defined objetive function cache
    empty_objective_function_caches!(model)
    if !function_has_parameters(model, f)
        MOI.set(model.optimizer, attr, f)
    else
        vars, params, param_constant =
            separate_possible_terms_and_calculate_parameter_constant(
                model,
                f.terms,
            )
        MOI.set(
            model.optimizer,
            attr,
            MOI.ScalarAffineFunction(vars, f.constant + param_constant),
        )
        model.affine_objective_cache = params
    end
    return
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction,
    v::MOI.SingleVariable,
)
    if haskey(model.parameters, v)
        error("Cannot use a parameter as objective function alone")
    elseif !haskey(model.variables, v)
        error("Variable not in the model")
    end
    return MOI.set(
        model.optimizer,
        attr,
        MOI.SingleVariable(model.variables[v.variable]),
    )
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveSense,
    sense::MOI.OptimizationSense,
)
    MOI.set(model.optimizer, attr, sense)
    return
end

function MOI.get(
    model::Optimizer,
    attr::T,
) where {
    T<:Union{
        MOI.TerminationStatus,
        MOI.ObjectiveValue,
        MOI.DualObjectiveValue,
        MOI.PrimalStatus,
        MOI.DualStatus,
    },
}
    return MOI.get(model.optimizer, attr)
end

function MOI.get(
    model::Optimizer,
    attr::T,
    c::MOI.ConstraintIndex,
) where {T<:Union{MOI.ConstraintPrimal,MOI.ConstraintDual}}
    return MOI.get(model.optimizer, attr, c)
end

function MOI.set(model::Optimizer, ::MOI.Silent, bool::Bool)
    MOI.set(model.optimizer, MOI.Silent(), bool)
    return
end

function MOI.set(model::Optimizer, attr::MOI.RawParameter, val::Any)
    MOI.set(model.optimizer, attr, val)
    return
end

# TODO(odow): remove this.
function MOI.set(model::Optimizer, attr::String, val::Any)
    return MOI.set(model.optimizer, MOI.RawParameter(attr), val)
end

function MOI.get(model::Optimizer, ::MOI.SolverName)
    name = MOI.get(model.optimizer, MOI.SolverName())
    return "Optimizer with $(name) attached"
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.VectorOfVariables,
    set::MOI.AbstractVectorSet,
) where {T}
    if function_has_parameters(model, f)
        error("VectorOfVariables does not allow parameters")
    end
    return MOI.add_constraint(model.optimizer, f, set)
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.VectorAffineFunction{T},
    set::MOI.AbstractVectorSet,
) where {T}
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function add_constraint_with_parameters_on_function(
    model::Optimizer,
    f::MOI.VectorAffineFunction{T},
    set::MOI.AbstractVectorSet,
) where {T}
    vars, params, param_constants =
        separate_possible_terms_and_calculate_parameter_constant(model, f, set)
    ci = MOI.add_constraint(
        model.optimizer,
        MOI.VectorAffineFunction(vars, f.constants + param_constants),
        set,
    )
    model.vector_constraint_cache[ci] = params
    return ci
end

function add_constraint_with_parameters_on_function(
    model::Optimizer,
    f::MOI.ScalarQuadraticFunction{T},
    set::S,
) where {T,S<:MOI.AbstractScalarSet}
    if function_quadratic_terms_has_parameters(model, f.quadratic_terms)
        (
            quad_vars,
            quad_aff_vars,
            quad_params,
            aff_terms,
            variables_associated_to_parameters,
            quad_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.quadratic_terms,
        )
    else
        quad_vars = f.quadratic_terms
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[]
        quad_params = MOI.ScalarQuadraticTerm{T}[]
        aff_terms = MOI.ScalarAffineTerm{T}[]
        quad_param_constant = zero(T)
        variables_associated_to_parameters = MOI.VariableIndex[]
    end
    if function_affine_terms_has_parameters(model, f.affine_terms)
        (
            aff_vars,
            aff_params,
            terms_with_variables_associated_to_parameters,
            aff_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.affine_terms,
            variables_associated_to_parameters,
        )
    else
        aff_vars = f.affine_terms
        aff_params = MOI.ScalarAffineTerm{T}[]
        terms_with_variables_associated_to_parameters =
            MOI.ScalarAffineTerm{T}[]
        aff_param_constant = zero(T)
    end
    aff_terms = vcat(aff_terms, aff_vars)
    const_term = f.constant + aff_param_constant + quad_param_constant
    quad_terms = quad_vars
    f_quad = if !isempty(quad_vars)
        MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term)
    else
        MOI.ScalarAffineFunction(aff_terms, const_term)
    end
    model.last_quad_add_added += 1
    ci =
        MOI.Utilities.normalize_and_add_constraint(model.optimizer, f_quad, set)
    # This part is used to remember that ci came from a quadratic function
    # It is particularly useful because sometimes the constraint mutates
    new_ci = MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{T},S}(
        model.last_quad_add_added,
    )
    model.quadratic_added_cache[new_ci] = ci
    fill_quadratic_constraint_caches!(
        model,
        new_ci,
        quad_aff_vars,
        quad_params,
        aff_params,
        terms_with_variables_associated_to_parameters,
    )
    return new_ci
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.ScalarQuadraticFunction{T},
    set::MOI.AbstractScalarSet,
) where {T}
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction,
    f::MOI.ScalarQuadraticFunction{T},
) where {T}
    # clear previously defined objetive function cache
    empty_objective_function_caches!(model)
    if !function_has_parameters(model, f)
        MOI.set(model.optimizer, attr, f)
        return
    end
    if function_quadratic_terms_has_parameters(model, f.quadratic_terms)
        (
            quad_vars,
            quad_aff_vars,
            quad_params,
            aff_terms,
            variables_associated_to_parameters,
            quad_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.quadratic_terms,
        )
    else
        quad_vars = f.quadratic_terms
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[]
        quad_params = MOI.ScalarQuadraticTerm{T}[]
        aff_terms = MOI.ScalarAffineTerm{T}[]
        quad_param_constant = zero(T)
        variables_associated_to_parameters = MOI.VariableIndex[]
    end

    if function_affine_terms_has_parameters(model, f.affine_terms)
        (
            aff_vars,
            aff_params,
            terms_with_variables_associated_to_parameters,
            aff_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.affine_terms,
            variables_associated_to_parameters,
        )
    else
        aff_vars = f.affine_terms
        aff_params = MOI.ScalarAffineTerm{T}[]
        terms_with_variables_associated_to_parameters =
            MOI.ScalarAffineTerm{T}[]
        aff_param_constant = zero(T)
    end

    aff_terms = vcat(aff_terms, aff_vars)
    const_term = f.constant + aff_param_constant + quad_param_constant
    quad_terms = quad_vars

    if !isempty(quad_vars)
        MOI.set(
            model.optimizer,
            attr,
            MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term),
        )
    else
        MOI.set(
            model.optimizer,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
            MOI.ScalarAffineFunction(aff_terms, const_term),
        )
    end

    if !isempty(quad_terms)
        f_quad = MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term)

        MOI.set(model.optimizer, attr, f_quad)
    else
        f_quad = MOI.ScalarAffineFunction(aff_terms, const_term)

        MOI.set(
            model.optimizer,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
            f_quad,
        )
    end

    model.quadratic_objective_cache_pv = quad_aff_vars
    model.quadratic_objective_cache_pp = quad_params
    model.quadratic_objective_cache_pc = aff_params
    model.quadratic_objective_variables_associated_to_parameters_cache =
        terms_with_variables_associated_to_parameters

    return
end

function MOI.optimize!(model::Optimizer)
    if !isempty(model.updated_parameters)
        update_parameters!(model)
    end
    MOI.optimize!(model.optimizer)
    if model.evaluate_duals
        calculate_dual_of_parameters(model)
    end
    return
end

end # module


using KernelAbstractions
using KernelAbstractions.Extras
using CUDA
using CUDAKernels
using Random123
using WormlikeChain
using StaticArrays

const N= 256

function getrow(A,row)
    ntuple(i -> A[row,i], Val(3))
end

@kernel function foo(A,B)
    I = @index(Global)
    myA= getrow(A,I)
    myB= getrow(B,I)
    myA= myA .+ myB
    @inbounds A[I,1] +=  myA[1]

end

fookh = foo(CPU(), 8)
fookd = foo(CUDADevice(), N)

A_h = ones(Float32,N,3)
A_d = CUDA.ones(Float32,N,3)
B_h = ones(Float32,N,3)
B_d = CUDA.ones(Float32,N,3)

#@time wait(fookd(A_d, B_d, ndrange=N))
#@time wait(fookh(A_h, B_h, ndrange=N))
show(A_d)