#
"""
$SIGNATURES

Represents a linear operator given by an `AbstractMatrix` that may be
applied to an `AbstractVecOrMat`. Its state is updated by the user-provided
`update_func` during operator evaluation (`L([w,], v, u, p, t)`), or by calls
to `update_coefficients[!](L, u, p, t)`. Both recursively call the
`update_function`, `update_func` which is assumed to have the signature

    update_func(A::AbstractMatrix, u, p, t; <accepted kwargs>) -> newA

or

    update_func!(A::AbstractMatrix, u, p, t; <accepted kwargs>) -> [modifies A]

The set of keyword-arguments accepted by `update_func[!]` should be provided
to `MatrixOperator` via the kwarg `accepted_kwargs` as a `Val` of a tuple of `Symbol`s
for zero-allocation kwarg filtering. For example, `accepted_kwargs = Val((:dtgamma,))`.
Plain tuples like `(:dtgamma,)` are deprecated but still supported.
`kwargs` cannot be passed down to `update_func[!]` if `accepted_kwargs`
are not provided.

$(UPDATE_COEFFS_WARNING)

# Arguments

  - `A::AbstractMatrix`: Matrix used for the current operator action.

# Keyword Arguments

  - `update_func`: Out-of-place update with signature
    `update_func(A, u, p, t; kwargs...) -> new_A`.
  - `update_func!`: In-place update with signature
    `update_func!(A, u, p, t; kwargs...)`.
  - `accepted_kwargs`: `Val` tuple of keyword names forwarded to the update function.

# Fields

  - `A`: Current matrix state.
  - `update_func`: Out-of-place matrix update function.
  - `update_func!`: In-place matrix update function.

# Interface Rules

Lazy matrix algebra is defined for `AbstractSciMLOperator`s. The Interface
supports lazy addition, subtraction, multiplication, inversion,
adjoints, and transposes. An update function must preserve the dimensions and
mathematical meaning of `A`; use `FunctionOperator` when the action itself is
nonlinear in the input vector.

# Returns

A `MatrixOperator` whose current action is `A`. When update functions are
provided, the returned operator stores them and applies them through the
documented update interface.

# Errors

Construction throws if `A` is not an `AbstractMatrix` or if a supplied update
function returns a value incompatible with the matrix action. Passing a
keyword not listed in `accepted_kwargs` is rejected by the keyword filter.

# Examples

Out-of-place update and usage

```
using LinearAlgebra, SciMLOperators

v = rand(4)
u = rand(4)
p = rand(4, 4)
t = rand()

mat_update = (A, u, p, t; scale = 0.0) -> t * p
M = MatrixOperator(zeros(4, 4); update_func = mat_update, accepted_kwargs = Val((:scale,)))

L = M * M + 3I
L = cache_operator(L, v)

# update and evaluate 
w = L(v, u, p, t; scale = 1.0)

# In-place evaluation
w = similar(v)
L(w, v, u, p, t; scale = 1.0)

# In-place with scaling
β = 0.5
L(w, v, u, p, t, 2.0, β; scale = 1.0) # w = 2.0*(L*v) + 0.5*w
```

In-place update and usage

```
w = zeros(4)
v = zeros(4)
u = rand(4)
p = rand(4) # Must be non-nothing
t = rand()

mat_update! = (A, u, p, t; scale = 0.0) -> (A .= t * p * u' * scale)
M = MatrixOperator(zeros(4, 4); update_func! = mat_update!, accepted_kwargs = Val((:scale,)))
L = M * M + 3I
L = cache_operator(L, v) 

# update L in-place and evaluate
update_coefficients!(L, u, p, t; scale = 1.0)
mul!(w, L, v)

# Or use the new interface that separates update and application
L(w, v, u, p, t; scale = 1.0)
```
"""
struct MatrixOperator{T, AT <: AbstractMatrix{T}, F, F!} <: AbstractSciMLOperator{T}
    A::AT
    update_func::F
    update_func!::F!

    function MatrixOperator(A, update_func, update_func!)
        return new{
            eltype(A),
            typeof(A),
            typeof(update_func),
            typeof(update_func!),
        }(
            A,
            update_func,
            update_func!
        )
    end
end

function MatrixOperator(
        A;
        update_func = nothing,
        update_func! = nothing,
        accepted_kwargs = nothing
    )
    update_func = preprocess_update_func(update_func, accepted_kwargs)
    update_func! = preprocess_update_func(update_func!, accepted_kwargs)

    return MatrixOperator(A, update_func, update_func!)
end

# constructors
function Base.similar(L::MatrixOperator, ::Type{T}, dims::Dims) where {T}
    return MatrixOperator(similar(L.A, T, dims))
end

# traits
# Method forwarding for MatrixOperator (previously using @forward from MacroTools)
LinearAlgebra.issymmetric(L::MatrixOperator) = LinearAlgebra.issymmetric(L.A)
LinearAlgebra.ishermitian(L::MatrixOperator) = LinearAlgebra.ishermitian(L.A)
LinearAlgebra.isposdef(L::MatrixOperator) = LinearAlgebra.isposdef(L.A)
issquare(L::MatrixOperator) = issquare(L.A)
has_ldiv(L::MatrixOperator) = has_ldiv(L.A)
has_ldiv!(L::MatrixOperator) = has_ldiv!(L.A)

isconvertible(::MatrixOperator) = true
islinear(::MatrixOperator) = true

function Base.show(io::IO, L::MatrixOperator)
    a, b = size(L)
    return print(io, "MatrixOperator($a × $b)")
end
Base.size(L::MatrixOperator) = size(L.A)
Base.iszero(L::MatrixOperator) = iszero(L.A)
for op in (
        :adjoint,
        :transpose,
    )
    @eval function Base.$op(L::MatrixOperator)
        isconstant(L) && return MatrixOperator($op(L.A))

        update_func = L.update_func === nothing ? nothing :
            (
                A, u, p, t; kwargs...,
            ) -> $op(
                L.update_func(
                    $op(L.A),
                    u,
                    p,
                    t;
                    kwargs...
                )
            )
        update_func! = L.update_func! === nothing ? nothing :
            (
                A, u, p, t; kwargs...,
            ) -> $op(
                L.update_func!(
                    $op(L.A),
                    u,
                    p,
                    t;
                    kwargs...
                )
            )

        return MatrixOperator(
            $op(L.A);
            update_func = update_func,
            update_func! = update_func!,
            accepted_kwargs = NoKwargFilter()
        )
    end
end

function Base.conj(L::MatrixOperator)
    isconstant(L) && return MatrixOperator(conj(L.A))

    update_func = L.update_func === nothing ? nothing :
        (
            A, u, p, t; kwargs...,
        ) -> conj(
            L.update_func(
                conj(L.A),
                u,
                p,
                t;
                kwargs...
            )
        )
    update_func! = L.update_func! === nothing ? nothing :
        (
            A, u, p, t; kwargs...,
        ) -> begin
            L.update_func!(conj!(L.A), u, p, t; kwargs...)
            conj!(L.A)
        end

    return MatrixOperator(
        conj(L.A);
        update_func = update_func,
        update_func! = update_func!,
        accepted_kwargs = NoKwargFilter()
    )
end

has_adjoint(L::MatrixOperator) = has_adjoint(L.A)
getops(L::MatrixOperator) = (L.A,)
function isconstant(L::MatrixOperator)
    return update_func_isconstant(L.update_func) & update_func_isconstant(L.update_func!)
end

function update_coefficients(L::MatrixOperator, u, p, t; kwargs...)
    if !isnothingfunc(L.update_func)
        @reset L.A = L.update_func(L.A, u, p, t; kwargs...)
    elseif !isnothingfunc(L.update_func!)
        L.update_func!(L.A, u, p, t; kwargs...)
    end
    return L
end

function update_coefficients!(L::MatrixOperator, u, p, t; kwargs...)
    if !isnothingfunc(L.update_func!)
        L.update_func!(L.A, u, p, t; kwargs...)
    elseif !isnothingfunc(L.update_func)
        if !ArrayInterface.ismutable(L.A)
            msg = """`update_coefficients!` cannot update a `MatrixOperator` holding
            an immutable matrix ($(typeof(L.A))) through an out-of-place
            `update_func`. Use `update_coefficients(L, u, p, t)` instead, or
            construct the operator with a mutating `update_func!`."""
            throw(ArgumentError(msg))
        end
        copyto!(L.A, L.update_func(L.A, u, p, t; kwargs...))
    end
    return nothing
end

# Out-of-place: v is action vector, u is update vector
function (L::MatrixOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
    L = update_coefficients(L, u, p, t; kwargs...)
    return L.A * v
end

# In-place: w is destination, v is action vector, u is update vector
function (L::MatrixOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
    update_coefficients!(L, u, p, t; kwargs...)
    return mul!(w, L.A, v)
end

# In-place with scaling: w = α*(L*v) + β*w
Base.@constprop :aggressive function (L::MatrixOperator)(
        w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...
    )
    update_coefficients!(L, u, p, t; kwargs...)
    mul!(w, L.A, v, α, β)
end

# TODO - add tests for MatrixOperator indexing
# propagate_inbounds here for the getindex fallback
Base.@propagate_inbounds Base.convert(::Type{AbstractMatrix}, L::MatrixOperator) = convert(
    AbstractMatrix,
    L.A
)
Base.@propagate_inbounds Base.setindex!(L::MatrixOperator, v, i::Int) = (L.A[i] = v)
Base.@propagate_inbounds Base.setindex!(
    L::MatrixOperator, v, I::Vararg{
        Int, N,
    }
) where {N} = (L.A[I...] = v)

Base.eachcol(L::MatrixOperator) = eachcol(L.A)
Base.eachrow(L::MatrixOperator) = eachrow(L.A)
Base.length(L::MatrixOperator) = length(L.A)
Base.iterate(L::MatrixOperator, args...) = iterate(L.A, args...)
Base.axes(L::MatrixOperator) = axes(L.A)
Base.eachindex(L::MatrixOperator) = eachindex(L.A)
function Base.IndexStyle(::Type{<:MatrixOperator{T, AType}}) where {T, AType}
    return Base.IndexStyle(AType)
end
Base.copyto!(L::MatrixOperator, rhs) = (copyto!(L.A, rhs); L)
# Unwrap a MatrixOperator `rhs` so the copy hits the underlying matrices' specialized
# `copyto!` (e.g. sparse→sparse copies `nzval`/`rowval`/`colptr`) instead of the generic
# element-wise fallback, which is O(n²) with sparse-index overhead.
Base.copyto!(L::MatrixOperator, rhs::MatrixOperator) = (copyto!(L.A, rhs.A); L)

Base.Broadcast.broadcastable(L::MatrixOperator) = L
Base.ndims(::Type{<:MatrixOperator{T, AType}}) where {T, AType} = ndims(AType)
ArrayInterface.issingular(L::MatrixOperator) = ArrayInterface.issingular(L.A)
function Base.copy(L::MatrixOperator)
    return MatrixOperator(
        copy(L.A);
        update_func = L.update_func,
        update_func! = L.update_func!,
        accepted_kwargs = NoKwargFilter()
    )
end

# operator application
Base.:*(L::MatrixOperator, v::AbstractVecOrMat) = L.A * v
Base.:\(L::MatrixOperator, v::AbstractVecOrMat) = L.A \ v

@inline function LinearAlgebra.mul!(
        w::AbstractVecOrMat, L::MatrixOperator, v::AbstractVecOrMat
    )
    return mul!(w, L.A, v)
end
@inline function LinearAlgebra.mul!(
        w::AbstractVecOrMat,
        L::MatrixOperator,
        v::AbstractVecOrMat,
        α,
        β
    )
    return mul!(w, L.A, v, α, β)
end

# These methods are used when the operator is applied from the right
# We do this to avoid to fall into generic methods that would be inefficient
@inline LinearAlgebra.mul!(w::AbstractTransposedVecOrMat, L::MatrixOperator, v::AbstractVecOrMat) = mul!(parent(w), transpose(v), transpose(L.A))
@inline LinearAlgebra.mul!(
    w::AbstractTransposedVecOrMat,
    L::MatrixOperator,
    v::AbstractVecOrMat,
    α,
    β
) = mul!(parent(w), transpose(v), transpose(L.A), α, β)

function LinearAlgebra.ldiv!(w::AbstractVecOrMat, L::MatrixOperator, v::AbstractVecOrMat)
    return ldiv!(w, L.A, v)
end
LinearAlgebra.ldiv!(L::MatrixOperator, v::AbstractVecOrMat) = ldiv!(L.A, v)

"""
$SIGNATURES

Represents an elementwise scaling (diagonal-scaling) operation that may
be applied to an `AbstractVecOrMat`. When `diag` is an `AbstractVector`
of length N, `L = DiagonalOperator(diag, ...)` can be applied to
`AbstractArray`s with `size(u, 1) == N`. Each column of the `v` will be
scaled by `diag`, as in `LinearAlgebra.Diagonal(diag) * v`.

When `diag` is a multidimensional array, `L = DiagonalOperator(diag, ...)` forms
an operator of size `(N, N)` where `N = size(diag, 1)` is the leading length of `diag`.
`L` then is the elementwise-scaling operation on arrays of `length(v) = length(diag)`
with leading length `size(u, 1) = N`.

Its state is updated by the user-provided `update_func` during operator
evaluation (`L([w,], v, u, p, t)`), or by calls to
`update_coefficients[!](L, u, p, t)`. Both recursively call the
`update_function`, `update_func` which is assumed to have the signature

    update_func(diag::AbstractVecOrMat, u, p, t; <accepted kwargs>) -> new_diag

or

    update_func!(diag::AbstractVecOrMat, u, p, t; <accepted kwargs>) -> [modifies diag]

The set of keyword-arguments accepted by `update_func[!]` should be provided
to `DiagonalOperator` via the kwarg `accepted_kwargs` as a `Val` of a tuple of `Symbol`s
for zero-allocation kwarg filtering. For example, `accepted_kwargs = Val((:dtgamma,))`.
Plain tuples like `(:dtgamma,)` are deprecated but still supported.
`kwargs` cannot be passed down to `update_func[!]` if `accepted_kwargs`
are not provided.

$(UPDATE_COEFFS_WARNING)

# Arguments

  - `diag::AbstractVector`: Diagonal entries of the operator.

# Keyword Arguments

  - `update_func`: Out-of-place update with signature
    `update_func(diag, u, p, t; kwargs...) -> new_diag`.
  - `update_func!`: In-place update with signature
    `update_func!(diag, u, p, t; kwargs...)`.
  - `accepted_kwargs`: `Val` tuple of keyword names forwarded to the update function.

# Interface Rules

Updates must preserve the diagonal's leading dimension. For multidimensional
`diag`, the operator acts elementwise while reporting a matrix size based on
its leading dimension.

# Returns

A `DiagonalOperator` backed by `diag` and the configured update functions.

# Errors

Construction or update throws when a supplied diagonal update changes the
leading dimension or returns an incompatible value.

# Examples

```julia
using SciMLOperators

D = DiagonalOperator([2.0, 3.0])
D * [4.0, 5.0] == [8.0, 15.0]
```
"""
function DiagonalOperator(
        diag::AbstractVector;
        update_func = nothing,
        update_func! = nothing,
        accepted_kwargs = nothing
    )
    diag_update_func = update_func_isconstant(update_func) ? update_func :
        (
            A, u, p, t; kwargs...,
        ) -> update_func(A.diag, u, p, t; kwargs...) |>
        Diagonal

    diag_update_func! = update_func_isconstant(update_func!) ? update_func! :
        (A, u, p, t; kwargs...) -> update_func!(A.diag, u, p, t; kwargs...)

    return MatrixOperator(
        Diagonal(diag);
        update_func = diag_update_func,
        update_func! = diag_update_func!,
        accepted_kwargs = accepted_kwargs
    )
end
LinearAlgebra.Diagonal(L::MatrixOperator) = MatrixOperator(Diagonal(L.A))

function Base.exp(L::MatrixOperator{T, <:Diagonal}) where {T}
    update_func = (
        A, u, p, t;
        kwargs...,
    ) -> update_coefficients(L, u, p, t; kwargs...).A |>
        exp
    return MatrixOperator(exp(L.A); update_func = update_func, accepted_kwargs = NoKwargFilter())
end

function Base.show(io::IO, L::MatrixOperator{T, <:Diagonal}) where {T}
    n = length(L.A.diag)
    return print(io, "DiagonalOperator($n × $n)")
end

"""
$SIGNATURES

    InvertibleOperator(L, F)

Pair an operator with a factorization or inverse object used for division.

# Arguments

  - `L`: Original operator used for multiplication and updates.
  - `F`: Factorization or inverse supporting `\\` or `ldiv!`.

# Fields

  - `L`: Wrapped operator.
  - `F`: Factorization or inverse representation.

# Interface Rules

Constructors such as `factorize`, `lu`, and `qr` create this wrapper for
concretizable operators. `L` and `F` must represent the same current action;
after a state update, extension code is responsible for keeping both current.
Use `ldiv!` or `\\` for the inverse action rather than assuming `F` is a
public field contract.
"""
struct InvertibleOperator{T, LT, FT} <: AbstractSciMLOperator{T}
    L::LT
    F::FT

    function InvertibleOperator(L, F)
        @assert has_ldiv(F) | has_ldiv!(F) "$F is not invertible"
        T = promote_type(eltype(L), eltype(F))

        return new{T, typeof(L), typeof(F)}(L, F)
    end
end

# constructor
function LinearAlgebra.factorize(L::AbstractSciMLOperator)
    fact = factorize(convert(AbstractMatrix, L))
    return InvertibleOperator(L, fact)
end

for fact in (
        :lu, :lu!,
        :qr, :qr!,
        :cholesky, :cholesky!,
        :ldlt, :ldlt!,
        :bunchkaufman, :bunchkaufman!,
        :lq, :lq!,
        :svd, :svd!,
    )
    @eval LinearAlgebra.$fact(
        L::AbstractSciMLOperator,
        args...
    ) = InvertibleOperator(
        L,
        LinearAlgebra.$fact(convert(AbstractMatrix, L), args...)
    )
    @eval LinearAlgebra.$fact(
        L::AbstractSciMLOperator;
        kwargs...
    ) = InvertibleOperator(
        L,
        LinearAlgebra.$fact(convert(AbstractMatrix, L); kwargs...)
    )
end

function Base.convert(
        ::Type{<:Factorization},
        L::InvertibleOperator{T, LT, <:Factorization}
    ) where {T, LT}
    return L.F
end

Base.convert(::Type{AbstractMatrix}, L::InvertibleOperator) = convert(AbstractMatrix, L.L)
has_concretization(L::InvertibleOperator) = has_concretization(L.L)

# traits
function Base.show(io::IO, L::InvertibleOperator)
    print(io, "InvertibleOperator(")
    show(io, L.L)
    print(io, ", ")
    show(io, L.F)
    return print(io, ")")
end

Base.size(L::InvertibleOperator) = size(L.L)
Base.transpose(L::InvertibleOperator) = InvertibleOperator(transpose(L.L), transpose(L.F))
Base.adjoint(L::InvertibleOperator) = InvertibleOperator(L.L', L.F')
Base.conj(L::InvertibleOperator) = InvertibleOperator(conj(L.L), conj(L.F))
Base.resize!(L::InvertibleOperator, n::Integer) = (resize!(L.L, n); resize!(L.F, n); L)
LinearAlgebra.opnorm(L::InvertibleOperator{T}, p::Real = 2) where {T} = one(T) / opnorm(L.F, p)
LinearAlgebra.issuccess(L::InvertibleOperator) = issuccess(L.F)

function update_coefficients(L::InvertibleOperator, u, p, t; kwargs...)
    @reset L.L = update_coefficients(L.L, u, p, t; kwargs...)
    @reset L.F = update_coefficients(L.F, u, p, t; kwargs...)
    return L
end
function update_coefficients!(L::InvertibleOperator, u, p, t; kwargs...)
    update_coefficients!(L.L, u, p, t; kwargs...)
    update_coefficients!(L.F, u, p, t; kwargs...)
    return nothing
end

getops(L::InvertibleOperator) = (L.L, L.F)

# Copy method to avoid aliasing
function Base.copy(L::InvertibleOperator)
    return InvertibleOperator(
        copy(L.L),
        copy(L.F)
    )
end

islinear(L::InvertibleOperator) = islinear(L.L)
isconvertible(L::InvertibleOperator) = isconvertible(L.L)

# Method forwarding for InvertibleOperator (previously using @forward from MacroTools)
# LinearAlgebra methods
LinearAlgebra.issymmetric(L::InvertibleOperator) = LinearAlgebra.issymmetric(L.L)
LinearAlgebra.ishermitian(L::InvertibleOperator) = LinearAlgebra.ishermitian(L.L)
LinearAlgebra.isposdef(L::InvertibleOperator) = LinearAlgebra.isposdef(L.L)

# SciML methods
isconstant(L::InvertibleOperator) = isconstant(L.L)
has_adjoint(L::InvertibleOperator) = has_adjoint(L.L)
has_mul(L::InvertibleOperator) = has_mul(L.L)
has_mul!(L::InvertibleOperator) = has_mul!(L.L)

has_ldiv(L::InvertibleOperator) = has_mul(L.F)
has_ldiv!(L::InvertibleOperator) = has_ldiv!(L.F)

function cache_internals(L::InvertibleOperator, v::AbstractVecOrMat)
    @reset L.L = cache_operator(L.L, v)
    @reset L.F = cache_operator(L.F, v)

    return L
end

# operator application
Base.:*(L::InvertibleOperator, v::AbstractVecOrMat) = L.L * v
Base.:\(L::InvertibleOperator, v::AbstractVecOrMat) = L.F \ v
function LinearAlgebra.mul!(w::AbstractVecOrMat, L::InvertibleOperator, v::AbstractVecOrMat)
    return mul!(w, L.L, v)
end
function LinearAlgebra.mul!(
        w::AbstractVecOrMat,
        L::InvertibleOperator,
        v::AbstractVecOrMat,
        α,
        β
    )
    return mul!(w, L.L, v, α, β)
end
function LinearAlgebra.ldiv!(
        w::AbstractVecOrMat,
        L::InvertibleOperator,
        v::AbstractVecOrMat
    )
    return ldiv!(w, L.F, v)
end
LinearAlgebra.ldiv!(L::InvertibleOperator, u::AbstractVecOrMat) = ldiv!(L.F, u)

# Out-of-place: v is action vector, u is update vector
function (L::InvertibleOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
    if !isconstant(L)
        L = update_coefficients(L, u, p, t; kwargs...)
    end
    return L.L * v
end

# In-place: w is destination, v is action vector, u is update vector
function (L::InvertibleOperator)(
        w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...
    )
    update_coefficients!(L, u, p, t; kwargs...)
    return mul!(w, L.L, v)
end

# In-place with scaling: w = α*(L*v) + β*w
function (L::InvertibleOperator)(
        w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...
    )
    update_coefficients!(L, u, p, t; kwargs...)
    return mul!(w, L.L, v, α, β)
end

"""
$SIGNATURES

Represents a generalized affine operation (`w = A * v + B * b`) that may
be applied to an `AbstractVecOrMat`. The user-provided update functions,
`update_func[!]` update the `AbstractVecOrMat` `b`, and are called
during operator evaluation (`L([w,], v, u, p, t)`), or by calls
to `update_coefficients[!](L, u, p, t)`. The update functions are
assumed to have the syntax

    update_func(b::AbstractVecOrMat, u, p, t; <accepted kwargs>) -> new_b

or

    update_func!(b::AbstractVecOrMat, u ,p , t; <accepted kwargs>) -> [modifies b]

and `B`, `b` are expected to have an appropriate size so that
`A * v + B * b` makes sense. Specifically, `size(A, 1) == size(B, 1)`, and
`size(v, 2) == size(b, 2)`.

The set of keyword-arguments accepted by `update_func[!]` should be provided
to `AffineOperator` via the kwarg `accepted_kwargs` as a `Val` of a tuple of `Symbol`s
for zero-allocation kwarg filtering. For example, `accepted_kwargs = Val((:dtgamma,))`.
Plain tuples like `(:dtgamma,)` are deprecated but still supported.
`kwargs` cannot be passed down to `update_func[!]` if `accepted_kwargs`
are not provided.

# Arguments

  - `A`: Matrix or SciML operator applied to the action vector.
  - `B`: Matrix or SciML operator applied to the additive term.
  - `b::AbstractArray`: Additive input to `B`.

# Keyword Arguments

  - `update_func`: Out-of-place update for `b`.
  - `update_func!`: In-place update for `b`.
  - `accepted_kwargs`: `Val` tuple of keyword names forwarded to the update function.

# Fields

  - `A`: Linear action component.
  - `B`: Additive-term action component.
  - `b`: Current additive input.
  - `update_func`: Out-of-place update for `b`.
  - `update_func!`: In-place update for `b`.

# Interface Rules

`AffineOperator` is not linear and cannot generally be converted to an
`AbstractMatrix`. Updates must preserve compatible dimensions for `A`, `B`,
and `b`; caching must preserve the affine action.

# Examples

```
v = rand(4)
u = rand(4)
p = rand(4)
t = rand()

A = MatrixOperator(rand(4, 4))
B = MatrixOperator(rand(4, 4))

vec_update_func = (b, u, p, t) -> p .* u * t
L = AffineOperator(A, B, zeros(4); update_func = vec_update_func)
L = cache_operator(L, v)

# update L and evaluate
w = L(v, u, p, t) # == A * v + B * (p .* u * t)
```
"""
struct AffineOperator{T, AT, BT, bT, F, F!} <: AbstractSciMLOperator{T}
    A::AT
    B::BT
    b::bT

    update_func::F # updates b
    update_func!::F! # updates b

    function AffineOperator(A, B, b, update_func, update_func!)
        T = promote_type(eltype.((A, B, b))...)

        return new{
            T,
            typeof(A),
            typeof(B),
            typeof(b),
            typeof(update_func),
            typeof(update_func!),
        }(
            A,
            B,
            b,
            update_func,
            update_func!
        )
    end
end

function AffineOperator(
        A::Union{AbstractMatrix, AbstractSciMLOperator},
        B::Union{AbstractMatrix, AbstractSciMLOperator},
        b::AbstractArray;
        update_func = nothing,
        update_func! = nothing,
        accepted_kwargs = nothing
    )
    @assert size(A, 1) == size(B, 1) "Dimension mismatch: A, B don't output vectors
  of same size"

    update_func = preprocess_update_func(update_func, accepted_kwargs)
    update_func! = preprocess_update_func(update_func!, accepted_kwargs)

    A = A isa AbstractMatrix ? MatrixOperator(A) : A
    B = B isa AbstractMatrix ? MatrixOperator(B) : B

    return AffineOperator(A, B, b, update_func, update_func!)
end

"""
$SIGNATURES

Construct the affine operation `v + b` as an `AffineOperator`.

# Arguments

  - `b::AbstractVecOrMat`: Additive term.

# Keyword Arguments

  - `update_func`: Out-of-place update for `b` with signature
    `update_func(b, u, p, t; kwargs...) -> new_b`.
  - `update_func!`: In-place update for `b` with signature
    `update_func!(b, u, p, t; kwargs...)`.
  - `accepted_kwargs`: `Val` tuple of forwarded update keyword names.

# Interface Rules

This is shorthand for `AffineOperator(I, I, b; kwargs...)`. The term `b` must
have a leading dimension compatible with the action vector. See
[`AffineOperator`](@ref) for update and caching rules.
"""
function AddVector(
        b::AbstractVecOrMat;
        update_func = nothing,
        update_func! = nothing,
        accepted_kwargs = nothing
    )
    N = size(b, 1)
    Id = IdentityOperator(N)

    return AffineOperator(
        Id, Id, b;
        update_func = update_func,
        update_func! = update_func!,
        accepted_kwargs = accepted_kwargs
    )
end

"""
$SIGNATURES

Construct the affine operation `v + B * b` as an `AffineOperator`.

# Arguments

  - `B`: Matrix or SciML operator acting on the additive term.
  - `b::AbstractVecOrMat`: Additive input to `B`.

# Keyword Arguments

  - `update_func`: Out-of-place update for `b`.
  - `update_func!`: In-place update for `b`.
  - `accepted_kwargs`: `Val` tuple of forwarded update keyword names.

# Interface Rules

`B * b` must have the same leading dimension as the action vector. This is
shorthand for `AffineOperator(I, B, b; kwargs...)`; see [`AffineOperator`](@ref)
for the update and caching contract.
"""
function AddVector(
        B, b::AbstractVecOrMat;
        update_func = nothing,
        update_func! = nothing,
        accepted_kwargs = nothing
    )
    N = size(B, 1)
    Id = IdentityOperator(N)

    return AffineOperator(
        Id, B, b;
        update_func = update_func,
        update_func! = update_func!,
        accepted_kwargs = accepted_kwargs
    )
end

function update_coefficients(L::AffineOperator, u, p, t; kwargs...)
    @reset L.A = update_coefficients(L.A, u, p, t; kwargs...)
    @reset L.B = update_coefficients(L.B, u, p, t; kwargs...)
    if !isnothingfunc(L.update_func)
        @reset L.b = L.update_func(L.b, u, p, t; kwargs...)
    end
    return L
end

function update_coefficients!(L::AffineOperator, u, p, t; kwargs...)
    if !isnothingfunc(L.update_func)
        L.update_func!(L.b, u, p, t; kwargs...)
    end
    for op in getops(L)
        update_coefficients!(op, u, p, t; kwargs...)
    end
    return nothing
end

# Copy method to avoid aliasing
function Base.copy(L::AffineOperator)
    return AffineOperator(
        copy(L.A),
        copy(L.B),
        copy(L.b);
        update_func = L.update_func,
        update_func! = L.update_func!
    )
end

function isconstant(L::AffineOperator)
    return update_func_isconstant(L.update_func) &
        update_func_isconstant(L.update_func!) &
        all(isconstant, (L.A, L.B))
end

getops(L::AffineOperator) = (L.A, L.B, L.b)

islinear(::AffineOperator) = false
isconvertible(::AffineOperator) = false

function Base.show(io::IO, L::AffineOperator)
    show(io, L.A)
    print(io, " + ")
    show(io, L.B)
    print(io, " * ")

    return if L.b isa AbstractVector
        n = size(L.b, 1)
        print(io, "AbstractVector($n)")
    elseif L.b isa AbstractMatrix
        n, k = size(L.b)
        print(io, "AbstractMatrix($n × $k)")
    end
end

Base.size(L::AffineOperator) = size(L.A)
Base.iszero(L::AffineOperator) = all(iszero, getops(L))
function Base.resize!(L::AffineOperator, n::Integer)
    resize!(L.A, n)
    resize!(L.B, n)
    resize!(L.b, n)

    return L
end

function Base.convert(::Type{AbstractMatrix}, L::AffineOperator)
    m, n = size(L)
    msg = """$L cannot be represented by an $m × $n AbstractMatrix"""
    throw(ArgumentError(msg))
end

has_adjoint(L::AffineOperator) = false
has_mul(L::AffineOperator) = has_mul(L.A)
has_mul!(L::AffineOperator) = has_mul!(L.A)
has_ldiv(L::AffineOperator) = has_ldiv(L.A)
has_ldiv!(L::AffineOperator) = has_ldiv!(L.A)

function cache_internals(L::AffineOperator, u::AbstractVecOrMat)
    @reset L.A = cache_operator(L.A, u)
    @reset L.B = cache_operator(L.B, L.b)
    return L
end

# operator application
function Base.:*(L::AffineOperator, v::AbstractVecOrMat)
    @assert size(L.b, 2) == size(v, 2)
    return (L.A * v) + (L.B * L.b)
end

function Base.:\(L::AffineOperator, v::AbstractVecOrMat)
    @assert size(L.b, 2) == size(v, 2)
    return L.A \ (v - (L.B * L.b))
end

function LinearAlgebra.mul!(w::AbstractVecOrMat, L::AffineOperator, v::AbstractVecOrMat)
    mul!(w, L.B, L.b)
    return mul!(w, L.A, v, true, true)
end

function LinearAlgebra.mul!(
        w::AbstractVecOrMat,
        L::AffineOperator,
        v::AbstractVecOrMat,
        α,
        β
    )
    mul!(w, L.B, L.b, α, β)
    return mul!(w, L.A, v, α, true)
end

function LinearAlgebra.ldiv!(w::AbstractVecOrMat, L::AffineOperator, v::AbstractVecOrMat)
    copy!(w, v)
    return ldiv!(L, w)
end

function LinearAlgebra.ldiv!(L::AffineOperator, v::AbstractVecOrMat)
    mul!(v, L.B, L.b, -1, 1)
    return ldiv!(L.A, v)
end
# Out-of-place: v is action vector, u is update vector
function (L::AffineOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
    L = update_coefficients(L, u, p, t; kwargs...)
    return (L.A * v) + (L.B * L.b)
end

# In-place: w is destination, v is action vector, u is update vector
function (L::AffineOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
    update_coefficients!(L, u, p, t; kwargs...)
    # First calculate A * v
    mul!(w, L.A, v)
    # Then add B * b
    return mul!(w, L.B, L.b, true, true)
end

# In-place with scaling: w = α*(L*v) + β*w
function (L::AffineOperator)(
        w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...
    )
    update_coefficients!(L, u, p, t; kwargs...)
    # Scale the existing w by β
    lmul!(β, w)
    # Add α * (A * v)
    mul!(w, L.A, v, α, true)
    # Add α * (B * b)
    return mul!(w, L.B, L.b, α, true)
end
#

has_concretization(::MatrixOperator) = true
