operators.jl

function _matmul(::Type{V}, A, B) where {V}
    return GenericMatrixExpr{V}(:*, Any[A, B], (size(A, 1), size(B, 2)), false)
end

function Base.:(*)(A::AbstractJuMPMatrix, B::Matrix)
    return _matmul(JuMP.variable_ref_type(A), A, B)
end
function Base.:(*)(A::Matrix, B::AbstractJuMPMatrix)
    return _matmul(JuMP.variable_ref_type(B), A, B)
end
function Base.:(*)(A::AbstractJuMPMatrix, B::AbstractJuMPMatrix)
    return _matmul(JuMP.variable_ref_type(A), A, B)
end

# Matrix-vector products: output is a 1-D `GenericArrayExpr` of length
# `size(A, 1)`. Allows users to write `W * x` for a vector variable `x`.
function _matvec(::Type{V}, A, b) where {V}
    return GenericArrayExpr{V,1}(:*, Any[A, b], (size(A, 1),), false)
end

function Base.:(*)(A::AbstractJuMPMatrix, b::Vector)
    return _matvec(JuMP.variable_ref_type(A), A, b)
end

function Base.:(*)(A::Matrix, b::AbstractJuMPVector{T}) where {T}
    return _matvec(JuMP.variable_ref_type(b), A, b)
end

function Base.:(*)(A::AbstractJuMPMatrix, b::AbstractJuMPVector{T}) where {T}
    return _matvec(JuMP.variable_ref_type(A), A, b)
end

function __broadcast(
    ::Type{V},
    axes::NTuple{N,Base.OneTo{Int}},
    op::Function,
    args::Vector{Any},
) where {V,N}
    return GenericArrayExpr{V,N}(Symbol(op), args, length.(axes), true)
end

function _broadcast(::Type{V}, op::Function, args...) where {V}
    return __broadcast(V, Broadcast.combine_axes(args...), op, Any[args...])
end

function Base.broadcasted(op::Function, x::AbstractJuMPArray)
    return _broadcast(JuMP.variable_ref_type(x), op, x)
end

function Base.broadcasted(op::Function, x::AbstractJuMPArray, y::AbstractArray)
    return _broadcast(JuMP.variable_ref_type(x), op, x, y)
end

function Base.broadcasted(op::Function, x::AbstractArray, y::AbstractJuMPArray)
    return _broadcast(JuMP.variable_ref_type(y), op, x, y)
end

function Base.broadcasted(
    op::Function,
    x::AbstractJuMPArray,
    y::AbstractJuMPArray,
)
    return _broadcast(JuMP.variable_ref_type(x), op, x, y)
end

function Base.broadcasted(op::Function, x::AbstractJuMPArray, y::Number)
    return _broadcast(JuMP.variable_ref_type(x), op, x, y)
end

function Base.broadcasted(op::Function, x::Number, y::AbstractJuMPArray)
    return _broadcast(JuMP.variable_ref_type(y), op, x, y)
end

function Base.broadcasted(
    ::typeof(Base.literal_pow),
    ::typeof(^),
    x::AbstractJuMPArray,
    ::Val{y},
) where {y}
    return Base.broadcasted(^, x, y)
end

function _transpose(x::AbstractJuMPArray{T,N}) where {T,N}
    V = JuMP.variable_ref_type(x)
    if N == 1
        return GenericArrayExpr{V,2}(:transpose, Any[x], (1, size(x, 1)), false)
    end
    @assert N == 2 "`transpose` only supports 1-D and 2-D arrays"
    return GenericArrayExpr{V,2}(
        :transpose,
        Any[x],
        (size(x, 2), size(x, 1)),
        false,
    )
end

LinearAlgebra.transpose(x::AbstractJuMPArray) = _transpose(x)
LinearAlgebra.adjoint(x::AbstractJuMPArray) = _transpose(x)

function Base.sum(x::AbstractJuMPArray; dims = Colon())
    V = JuMP.variable_ref_type(x)
    if dims === Colon()
        return JuMP.GenericNonlinearExpr{V}(:sum, Any[x])
    end
    sz = ntuple(i -> i in dims ? 1 : size(x, i), ndims(x))
    dims_vec = JuMP.value_type(V)[d for d in dims]
    return GenericArrayExpr{V,ndims(x)}(:sum_dims, Any[x, dims_vec], sz, false)
end

import LinearAlgebra

function _array_norm(x::AbstractJuMPArray)
    V = JuMP.variable_ref_type(x)
    return JuMP.GenericNonlinearExpr{V}(:norm, Any[x])
end

# Define norm for each concrete AbstractJuMPArray subtype to avoid
# ambiguity with JuMP's error-throwing
#   LinearAlgebra.norm(::AbstractArray{<:AbstractJuMPScalar})
function LinearAlgebra.norm(x::GenericArrayExpr)
    return _array_norm(x)
end

function LinearAlgebra.norm(x::ArrayOfVariables)
    return _array_norm(x)
end

# Subtraction between array expressions and constant arrays
function Base.:(-)(
    x::AbstractJuMPArray{T,N},
    y::AbstractArray{S,N},
) where {S,T,N}
    V = JuMP.variable_ref_type(x)
    @assert size(x) == size(y)
    return GenericArrayExpr{V,N}(:-, Any[x, y], size(x), false)
end

function Base.:(-)(
    x::AbstractArray{S,N},
    y::AbstractJuMPArray{T,N},
) where {S,T,N}
    V = JuMP.variable_ref_type(y)
    @assert size(x) == size(y)
    return GenericArrayExpr{V,N}(:-, Any[x, y], size(y), false)
end

function Base.:(-)(
    x::AbstractJuMPArray{T,N},
    y::AbstractJuMPArray{S,N},
) where {T,S,N}
    V = JuMP.variable_ref_type(x)
    @assert JuMP.variable_ref_type(y) == V
    @assert size(x) == size(y)
    return GenericArrayExpr{V,N}(:-, Any[x, y], size(x), false)
end

# Addition between array expressions and constant arrays
function Base.:(+)(
    x::AbstractJuMPArray{T,N},
    y::AbstractArray{S,N},
) where {S,T,N}
    V = JuMP.variable_ref_type(x)
    @assert size(x) == size(y)
    return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false)
end

function Base.:(+)(
    x::AbstractArray{S,N},
    y::AbstractJuMPArray{T,N},
) where {S,T,N}
    V = JuMP.variable_ref_type(y)
    @assert size(x) == size(y)
    return GenericArrayExpr{V,N}(:+, Any[x, y], size(y), false)
end

function Base.:(+)(
    x::AbstractJuMPArray{T,N},
    y::AbstractJuMPArray{S,N},
) where {T,S,N}
    V = JuMP.variable_ref_type(x)
    @assert JuMP.variable_ref_type(y) == V
    @assert size(x) == size(y)
    return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false)
end

# ── Evaluator helper ─────────────────────────────────────────────────────────

"""
    evaluator(f::Function, input_dim::Int; mode = Mode(), features = [:Grad, :Jac, :JacVec])

Compile a vector-valued Julia function `f` whose output is built with the
JuMP/ArrayDiff array operators into an [`Evaluator`](@ref) ready for
`eval_residual!` / `eval_residual_jtprod!`. The function is called once with a
fresh length-`input_dim` `ArrayOfVariables`, the resulting array expression
becomes the residual, and the evaluator is initialised with `features`.
"""
function evaluator(
    f::Function,
    input_dim::Int;
    mode = Mode(),
    features::Vector{Symbol} = Symbol[:Grad, :Jac, :JacVec],
)
    model = JuMP.Model()
    JuMP.@variable(model, x[1:input_dim], container = ArrayOfVariables,)
    residual_expr = f(x)
    ad_model = Model()
    set_residual!(ad_model, JuMP.moi_function(residual_expr))
    eval = Evaluator(ad_model, mode, JuMP.index.(JuMP.all_variables(model)))
    MOI.initialize(eval, features)
    return eval
end

# ── User-defined array operators ─────────────────────────────────────────────
#
# `add_operator(model, arity, f)` registers `f` on the `ArrayDiff.Model` via
# the [`UserDefinedArrayOperator`](@ref) attribute and returns a
# `JuMP.NonlinearOperator` wrapping `f`. When the returned operator is called
# with at least one `AbstractJuMPArray` argument, the dispatch methods below
# build either a `GenericArrayExpr` (when `f` returns an array) or a
# `JuMP.GenericNonlinearExpr` (scalar output) with that `name`. The reverse-
# mode derivative is pulled from `ChainRulesCore.rrule` at evaluator time.

"""
    add_operator(model::Model, arity::Int, f::Function; name::Symbol = Symbol(f))

Register `f` as a user-defined array operator on `model` and return a
`JuMP.NonlinearOperator` wrapping it. Mirrors `JuMP.add_nonlinear_operator`:
the call internally does `MOI.set(model, UserDefinedArrayOperator(name; arity), f)`.
The returned operator can be called with `AbstractJuMPArray` arguments to
build a `GenericArrayExpr` (array result) or `JuMP.GenericNonlinearExpr`
(scalar result). The output shape is determined by [`infer_sizes`](@ref).
"""
function add_operator(
    model::Model,
    arity::Int,
    f::Function;
    name::Symbol = Symbol(f),
)
    MOI.set(model, UserDefinedArrayOperator(name; arity), f)
    return JuMP.NonlinearOperator(f, name)
end

function _build_user_op_expr(op::JuMP.NonlinearOperator, V::Type, args::Tuple)
    shapes = map(size, args)
    out_sz = infer_sizes(op.func, shapes...)
    if isempty(out_sz)
        return JuMP.GenericNonlinearExpr{V}(op.head, Any[args...])
    end
    return GenericArrayExpr{V,length(out_sz)}(
        op.head,
        Any[args...],
        out_sz,
        false,
    )
end

function (op::JuMP.NonlinearOperator)(x::AbstractJuMPArray)
    V = JuMP.variable_ref_type(x)
    return _build_user_op_expr(op, V, (x,))
end

function (op::JuMP.NonlinearOperator)(
    x::AbstractJuMPArray,
    y::Union{Real,AbstractArray{<:Real},AbstractJuMPArray},
)
    V = JuMP.variable_ref_type(x)
    return _build_user_op_expr(op, V, (x, y))
end

function (op::JuMP.NonlinearOperator)(
    x::Union{Real,AbstractArray{<:Real}},
    y::AbstractJuMPArray,
)
    V = JuMP.variable_ref_type(y)
    return _build_user_op_expr(op, V, (x, y))
end