using SciMLOperators, LinearAlgebra
using SparseArrays
using Random
using Test
using LoopVectorization

using SciMLOperators: InvertibleOperator, InvertedOperator, ⊗, AbstractSciMLOperator
using FFTW

Random.seed!(0)
N = 8
K = 19

@testset "MatrixOperator, InvertibleOperator" begin
    # Vectors for testing
    u = rand(N, K)  # Both update and action vector
    v = rand(N, K)  # Output/action vector
    w = zeros(N, K)  # Output vector

    p = nothing
    t = 0
    α = rand()
    β = rand()

    A = rand(N, N)
    At = A'

    AA = MatrixOperator(A)
    AAt = AA'

    @test AA isa MatrixOperator
    @test AAt isa MatrixOperator

    @test isconstant(AA)
    @test isconstant(AAt)

    @test issquare(AA)
    @test islinear(AA)

    FF = factorize(AA)
    FFt = FF'

    @test FF isa InvertibleOperator
    @test FFt isa InvertibleOperator

    @test isconstant(FF)
    @test isconstant(FFt)

    @test_throws MethodError resize!(AA, N)

    @test eachindex(A) === eachindex(AA)
    @test eachindex(A') === eachindex(AAt) === eachindex(MatrixOperator(At))

    @test A ≈ convert(AbstractMatrix, AA) ≈ convert(AbstractMatrix, FF)
    @test At ≈ convert(AbstractMatrix, AAt) ≈ convert(AbstractMatrix, FFt)

    @test A ≈ Matrix(AA) ≈ Matrix(FF)
    @test At ≈ Matrix(AAt) ≈ Matrix(FFt)

    # Test with new interface - same vector for update and action
    @test A * u ≈ AA(u, u, p, t)
    @test At * u ≈ AAt(u, u, p, t)

    # Test with different vectors for update and action
    @test A * v ≈ AA(v, u, p, t)
    @test At * v ≈ AAt(v, u, p, t)

    @test A \ u ≈ AA \ u ≈ FF \ u
    @test At \ u ≈ AAt \ u ≈ FFt \ u

    # Test in-place operations
    copy!(w, zeros(N, K))
    AA(w, v, u, p, t)
    @test w ≈ A * v

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    AA(w, v, u, p, t, α, β)
    @test w ≈ α * (A * v) + β * orig_w
end

@testset "InvertibleOperator test" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    p = nothing
    t = 0
    α = rand()
    β = rand()

    d = rand(N, K)
    D = DiagonalOperator(d)
    Di = DiagonalOperator(inv.(d)) |> InvertedOperator

    L = InvertibleOperator(D, Di)
    L = cache_operator(L, u)

    @test iscached(L)

    # Test with new interface
    @test L(v, u, p, t) ≈ d .* v
    @test L \ u ≈ d .\ u

    # Test in-place operations
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ d .* v

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * (d .* v) + β * orig_w

    # Test division operations
    copy!(w, zeros(N, K))
    ldiv!(w, L, u)
    @test w ≈ d .\ u

    # Existing test for in-place ldiv!
    v = copy(u)
    ldiv!(L, v)
    @test v ≈ d .\ u
end

@testset "MatrixOperator update test" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    p = rand(N)
    t = rand()
    α = rand()
    β = rand()

    L = MatrixOperator(
        zeros(N, N);
        update_func = (A, u, p, t) -> p * p',
        update_func! = (A, u, p, t) -> A .= p * p'
    )

    @test !isconstant(L)

    # Expected matrix after update
    A = p * p'

    # Test with new interface - same vector for update and action
    @test L(u, u, p, t) ≈ A * u

    # Test with different vectors for update and action
    @test L(v, u, p, t) ≈ A * v

    # Test in-place operation
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ A * v

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * (A * v) + β * orig_w

    # update_coefficients! with only an out-of-place update_func (#409)
    Loop = MatrixOperator(zeros(N, N); update_func = (A, u, p, t) -> p * p')
    update_coefficients!(Loop, u, p, t)
    @test convert(AbstractMatrix, Loop) ≈ p * p'

    # immutable storage cannot be updated in place through an out-of-place func
    Limm = MatrixOperator(
        reshape(1.0:(N * N), N, N); update_func = (A, u, p, t) -> p * p'
    )
    @test_throws ArgumentError update_coefficients!(Limm, u, p, t)

    A = [
        -2.0 1 0 0 0
        1 -2 1 0 0
        0 1 -2 1 0
        0 0 1 -2 1
        0 0 0 1 -2
    ]
    v = [3.0, 2.0, 1.0, 2.0, 3.0]
    opA = MatrixOperator(A)

    function update_function!(B, u, p, t)
        dt = p
        B .= A .* u + dt * I
    end

    u = Array(1:1.0:5)
    p = 0.1
    t = 0.0
    opB = MatrixOperator(copy(A); update_func! = update_function!)

    function Bfunc!(w, v, u, p, t)
        dt = p
        w[1] = -(2 * u[1] - dt) * v[1] + v[2] * u[1]
        for i in 2:4
            w[i] = v[i - 1] * u[i] - (2 * u[i] - dt) * v[i] + v[i + 1] * u[i]
        end
        w[5] = v[4] * u[5] - (2 * u[5] - dt) * v[5]
        nothing
    end

    function Bfunc!(v, u, p, t)
        w = zeros(5)
        Bfunc!(w, v, u, p, t)
        w
    end

    mfopB = FunctionOperator(Bfunc!, zeros(5), zeros(5); u, p, t, isconstant = false)

    @test iszero(
        opB(v, Array(2:1.0:6), 0.5, nothing) -
            mfopB(v, Array(2:1.0:6), 0.5, nothing)
    )
end

@testset "DiagonalOperator update test" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    p = rand(N)
    t = rand()
    α = rand()
    β = rand()

    D = DiagonalOperator(
        zeros(N);
        update_func = (diag, u, p, t) -> p * t,
        update_func! = (diag, u, p, t) -> diag .= p * t
    )

    @test !isconstant(D)
    @test issquare(D)
    @test islinear(D)

    # Expected result after update
    expected = (p * t) .* v

    # Test with new interface - different vectors for update and action
    @test D(v, u, p, t) ≈ expected

    # Test in-place operation
    copy!(w, zeros(N, K))
    D(w, v, u, p, t)
    @test w ≈ expected

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    D(w, v, u, p, t, α, β)
    @test w ≈ α * expected + β * orig_w

    expD = exp(D)
    @test expD isa MatrixOperator
    @test expD.A isa Diagonal
    @test expD(v, u, p, t) ≈ exp.(p * t) .* v
end

@testset "Batched Diagonal Operator" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    d = rand(N, K)
    α = rand()
    β = rand()
    p = nothing
    t = 0.0

    L = DiagonalOperator(d)
    @test isconstant(L)
    @test_throws MethodError resize!(L, N)

    @test issquare(L)
    @test islinear(L)

    # Test with new interface
    @test L(v, u, p, t) ≈ d .* v

    # Test in-place operation
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ d .* v

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * (d .* v) + β * orig_w

    # Test division operations
    @test L \ u ≈ d .\ u

    copy!(w, zeros(N, K))
    ldiv!(w, L, u)
    @test w ≈ d .\ u

    # Existing test for in-place ldiv!
    v = copy(u)
    ldiv!(L, u)
    @test u ≈ d .\ v
end

@testset "Batched DiagonalOperator update test" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    d = zeros(N, K)
    p = rand(N, K)
    t = rand()

    D = DiagonalOperator(
        d;
        update_func = (diag, u, p, t) -> p * t,
        update_func! = (diag, u, p, t) -> diag .= p * t
    )

    @test !isconstant(D)
    @test issquare(D)
    @test islinear(D)

    # Expected result after update
    expected = (p * t) .* v

    # Test with new interface - different vectors for update and action
    @test D(v, u, p, t) ≈ expected

    # Test in-place operation
    copy!(w, zeros(N, K))
    D(w, v, u, p, t)
    @test w ≈ expected

    # update_coefficients! with only an out-of-place update_func (#409)
    Doop = DiagonalOperator(zeros(N, K); update_func = (diag, u, p, t) -> p * t)
    update_coefficients!(Doop, u, p, t)
    copy!(w, zeros(N, K))
    mul!(w, Doop, v)
    @test w ≈ (p * t) .* v

    # immutable storage cannot be updated in place through an out-of-place func
    Dimm = DiagonalOperator(
        reshape(1.0:(N * K), N, K); update_func = (diag, u, p, t) -> p * t
    )
    @test_throws ArgumentError update_coefficients!(Dimm, u, p, t)
end

@testset "AffineOperator" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    A = rand(N, N)
    B = rand(N, N)
    D = Diagonal(A)
    b = rand(N, K)
    α = rand()
    β = rand()
    p = nothing
    t = 0.0

    L = AffineOperator(MatrixOperator(A), MatrixOperator(B), b)
    @test isconstant(L)
    @test issquare(L)
    @test !islinear(L)

    # Test with new interface
    expected = A * v + B * b
    @test L(v, u, p, t) ≈ expected

    # Test in-place operation
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ expected

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * expected + β * orig_w

    L = AffineOperator(MatrixOperator(D), MatrixOperator(B), b)
    @test issquare(L)
    @test !islinear(L)

    # Test division operations
    @test L \ u ≈ D \ (u - B * b)

    copy!(w, zeros(N, K))
    ldiv!(w, L, u)
    @test w ≈ D \ (u - B * b)

    # Existing test for in-place ldiv!
    v = copy(u)
    ldiv!(L, u)
    @test u ≈ D \ (v - B * b)

    L = AddVector(b)
    @test issquare(L)
    @test !islinear(L)

    # Test with new interface
    expected = v + b
    @test L(v, u, p, t) ≈ expected
    @test L \ u ≈ u - b

    # Test in-place operation
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ expected

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * expected + β * orig_w

    # Test division operations
    copy!(w, zeros(N, K))
    ldiv!(w, L, u)
    @test w ≈ u - b

    # Existing test for in-place ldiv!
    v = copy(u)
    ldiv!(L, u)
    @test u ≈ v - b

    L = AddVector(MatrixOperator(B), b)
    @test issquare(L)
    @test !islinear(L)

    # Test with new interface
    expected = v + B * b
    @test L(v, u, p, t) ≈ expected
    @test L \ u ≈ u - B * b

    # Test in-place operation
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ expected

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * expected + β * orig_w

    # Test division operations
    copy!(w, zeros(N, K))
    ldiv!(w, L, u)
    @test w ≈ u - B * b

    # Existing test for in-place ldiv!
    v = copy(u)
    ldiv!(L, u)
    @test u ≈ v - B * b
end

@testset "AffineOperator update test" begin
    # Vectors for testing
    u = rand(N, K)  # Update vector
    v = rand(N, K)  # Action vector
    w = zeros(N, K)  # Output vector

    A = rand(N, N)
    B = rand(N, N)
    p = rand(N, K)
    t = rand()
    α = rand()
    β = rand()

    L = AffineOperator(
        A, B, zeros(N, K);
        update_func = (b, u, p, t) -> p * t,
        update_func! = (b, u, p, t) -> b .= p * t
    )

    @test !isconstant(L)

    # Expected updated bias and result
    b = p * t
    expected = A * v + B * b

    # Test with new interface - different vectors for update and action
    @test L(v, u, p, t) ≈ expected

    # Test in-place operation
    copy!(w, zeros(N, K))
    L(w, v, u, p, t)
    @test w ≈ expected

    # Test in-place with scaling
    copy!(w, rand(N, K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * expected + β * orig_w
end

@testset "TensorProductOperator, square = $square" for square in [false, true]
    m1, n1 = 3, 5
    m2, n2 = 7, 11
    m3, n3 = 13, 17

    if square
        n1, n2, n3 = m1, m2, m3
    end

    A = rand(m1, n1)
    B = rand(m2, n2)
    C = rand(m3, n3)
    α = rand()
    β = rand()
    p = nothing
    t = 0.0

    AB = kron(A, B)
    ABC = kron(A, B, C)

    # test Base.kron overload
    # ensure kron(mat, mat) is not a TensorProductOperator
    @test !isa(AB, AbstractSciMLOperator)
    @test !isa(ABC, AbstractSciMLOperator)

    # test Base.kron overload
    _A = rand(N, N)
    @test kron(_A, MatrixOperator(_A)) isa TensorProductOperator{Float64}
    @test kron(MatrixOperator(_A), _A) isa TensorProductOperator{Float64}

    @test kron(MatrixOperator(_A), MatrixOperator(_A)) isa TensorProductOperator{Float64}

    # Inputs/Update vectors
    u2 = rand(n1 * n2, K)
    u3 = rand(n1 * n2 * n3, K)

    # Action vectors (same as update vectors initially)
    v2 = copy(u2)
    v3 = copy(u3)

    # Output vectors
    w2 = zeros(m1 * m2, K)
    w3 = zeros(m1 * m2 * m3, K)

    opAB = TensorProductOperator(A, B)
    opABC = TensorProductOperator(A, B, C)

    @test opAB isa TensorProductOperator{Float64}
    @test opABC isa TensorProductOperator{Float64}

    @test isconstant(opAB)
    @test isconstant(opABC)

    @test islinear(opAB)
    @test islinear(opABC)

    if square
        @test issquare(opAB)
        @test issquare(opABC)
    else
        @test !issquare(opAB)
        @test !issquare(opABC)
    end

    @test AB ≈ convert(AbstractMatrix, opAB)
    @test ABC ≈ convert(AbstractMatrix, opABC)

    # factorization tests
    opAB_F = factorize(opAB)
    opABC_F = factorize(opABC)

    @test isconstant(opAB_F)
    @test isconstant(opABC_F)

    @test opAB_F isa TensorProductOperator{Float64}
    @test opABC_F isa TensorProductOperator{Float64}

    @test AB ≈ convert(AbstractMatrix, opAB_F)
    @test ABC ≈ convert(AbstractMatrix, opABC_F)

    # Test with new interface
    @test AB * v2 ≈ opAB(v2, u2, p, t)
    @test ABC * v3 ≈ opABC(v3, u3, p, t)

    @test AB \ w2 ≈ opAB \ w2
    @test AB \ w2 ≈ opAB_F \ w2
    @test ABC \ w3 ≈ opABC \ w3
    @test ABC \ w3 ≈ opABC_F \ w3

    @test !iscached(opAB)
    @test !iscached(opABC)

    @test !iscached(opAB_F)
    @test !iscached(opABC_F)

    opAB = cache_operator(opAB, u2)
    opABC = cache_operator(opABC, u3)

    opAB_F = cache_operator(opAB_F, u2)
    opABC_F = cache_operator(opABC_F, u3)

    @test iscached(opAB)
    @test iscached(opABC)

    @test iscached(opAB_F)
    @test iscached(opABC_F)

    N2 = n1 * n2
    N3 = n1 * n2 * n3
    M2 = m1 * m2
    M3 = m1 * m2 * m3

    # Test in-place operations with new interface
    v2 = rand(n1 * n2, K)  # Action vector
    w2 = zeros(M2, K)      # Output vector
    opAB(w2, v2, u2, p, t)
    @test w2 ≈ AB * v2

    v3 = rand(n1 * n2 * n3, K)  # Action vector
    w3 = zeros(M3, K)           # Output vector
    opABC(w3, v3, u3, p, t)
    @test w3 ≈ ABC * v3

    # Test in-place with scaling
    v2 = rand(n1 * n2, K)    # Action vector
    w2 = rand(M2, K)         # Output vector
    orig_w2 = copy(w2)
    opAB(w2, v2, u2, p, t, α, β)
    @test w2 ≈ α * AB * v2 + β * orig_w2

    v3 = rand(n1 * n2 * n3, K)  # Action vector
    w3 = rand(M3, K)            # Output vector
    orig_w3 = copy(w3)
    opABC(w3, v3, u3, p, t, α, β)
    @test w3 ≈ α * ABC * v3 + β * orig_w3

    if square
        # Test division operations with new interface
        v2 = rand(M2, K)     # Action vector (size of output space)
        u2 = rand(N2, K)     # Update vector (size of input space)
        w2 = zeros(N2, K)    # Output vector (size of input space)

        # ldiv! with new interface
        ldiv!(w2, opAB_F, v2)
        @test w2 ≈ AB \ v2

        v3 = rand(M3, K)     # Action vector
        u3 = rand(N3, K)     # Update vector
        w3 = zeros(N3, K)    # Output vector
        ldiv!(w3, opABC_F, v3)
        @test w3 ≈ ABC \ v3

        # In-place ldiv! (original style)
        v2 = rand(M2, K)
        u2 = copy(v2)
        ldiv!(opAB_F, v2)
        @test v2 ≈ AB \ u2

        v3 = rand(M3, K)
        u3 = copy(v3)
        ldiv!(opABC_F, v3)
        @test v3 ≈ ABC \ u3
    else # TODO
        v2 = rand(M2, K)     # Action vector
        u2 = rand(N2, K)     # Update vector
        w2 = zeros(N2, K)    # Output vector

        if VERSION < v"1.9-"
            @test_broken ldiv!(w2, opAB_F, v2) ≈ AB \ v2
        else
            @test ldiv!(w2, opAB_F, v2) ≈ AB \ v2
        end

        v3 = rand(M3, K)     # Action vector
        u3 = rand(N3, K)     # Update vector
        w3 = zeros(N3, K)    # Output vector
        @test_broken ldiv!(w3, opABC_F, v3) ≈ ABC \ v3 # errors
    end

    @testset "Simplified Structure with IdentityOperator" begin
        Id1 = IdentityOperator(m1)
        Id2 = IdentityOperator(m2)
        Id3 = IdentityOperator(m3)
        A1 = MatrixOperator(rand(n1, n1))
        A2 = MatrixOperator(rand(n2, n2))

        op1 = kron(A1, Id1, Id2, Id3)
        op2 = kron(Id1, A1, Id2, Id3)
        op3 = kron(Id1, Id2, A1, Id3)
        op4 = kron(Id1, Id2, Id3, A1)

        op5 = kron(A1, A2, Id1, Id2)
        op6 = kron(Id1, A1, A2, Id2)
        op7 = kron(Id1, Id2, A1, A2)

        # Test the structure of the resulting operators
        # The nesting depth structure should be 2 at most
        @test op1.ops[1] isa MatrixOperator
        @test op1.ops[2] isa IdentityOperator

        @test op2.ops[1] isa TensorProductOperator
        @test op2.ops[2] isa IdentityOperator
        @test op2.ops[1].ops[1] isa IdentityOperator
        @test op2.ops[1].ops[2] isa MatrixOperator

        @test op3.ops[1] isa TensorProductOperator
        @test op3.ops[2] isa IdentityOperator
        @test op3.ops[1].ops[1] isa IdentityOperator
        @test op3.ops[1].ops[2] isa MatrixOperator

        @test op4.ops[1] isa IdentityOperator
        @test op4.ops[2] isa MatrixOperator

        @test op5.ops[1] isa MatrixOperator
        @test op5.ops[2] isa TensorProductOperator
        @test op5.ops[2].ops[1] isa MatrixOperator
        @test op5.ops[2].ops[2] isa IdentityOperator

        @test op6.ops[1] isa TensorProductOperator
        @test op6.ops[2] isa TensorProductOperator
        @test op6.ops[1].ops[1] isa IdentityOperator
        @test op6.ops[1].ops[2] isa MatrixOperator
        @test op6.ops[2].ops[1] isa MatrixOperator
        @test op6.ops[2].ops[2] isa IdentityOperator

        @test op7.ops[1] isa TensorProductOperator
        @test op7.ops[2] isa MatrixOperator
        @test op7.ops[1].ops[1] isa IdentityOperator
        @test op7.ops[1].ops[2] isa MatrixOperator

        @test convert(AbstractMatrix, op1) ≈ kron(convert(AbstractMatrix, A1), I(m1 * m2 * m3))
        @test convert(AbstractMatrix, op2) ≈ kron(I(m1), convert(AbstractMatrix, A1), I(m2 * m3))
        @test convert(AbstractMatrix, op3) ≈ kron(I(m1 * m2), convert(AbstractMatrix, A1), I(m3))
        @test convert(AbstractMatrix, op4) ≈ kron(I(m1 * m2 * m3), convert(AbstractMatrix, A1))

        @test convert(AbstractMatrix, op5) ≈ kron(
            convert(AbstractMatrix, A1), convert(AbstractMatrix, A2), I(m1 * m2)
        )
        @test convert(AbstractMatrix, op6) ≈ kron(
            I(m1), convert(AbstractMatrix, A1), convert(AbstractMatrix, A2), I(m2)
        )
        @test convert(AbstractMatrix, op7) ≈ kron(
            I(m1 * m2), convert(AbstractMatrix, A1), convert(AbstractMatrix, A2)
        )

        for op in (op1, op2, op3, op4, op5, op6, op7)
            v = rand(size(op, 2), K)
            @test convert(AbstractMatrix, op) * v ≈ op * v
        end

    end
end

@testset "TensorProductOperator cache size, square = $square" for square in [false, true]
    m1, m2 = 3, 7
    n1, n2 = square ? (m1, m2) : (5, 11)

    A = rand(m1, n1)
    B = rand(m2, n2)
    AB = kron(A, B)
    α, β = rand(), rand()

    op = TensorProductOperator(A, B)

    # Number of distinct buffers, since slots are deliberately aliased when square.
    nbuffers(L) = length(unique(objectid, filter(!isnothing, collect(L.cache))))

    # A vector input can only ever reach cache slots 1 and 5, so the other five are
    # left unallocated: 4 buffers -> 1 when both factors are square, 7 -> 2 otherwise.
    opv = cache_operator(op, rand(n1 * n2))
    @test iscached(opv)
    @test nbuffers(opv) == (square ? 1 : 2)

    # A matrix input still gets the full cache.
    opm = cache_operator(op, rand(n1 * n2, K))
    @test nbuffers(opm) == (square ? 4 : 7)

    # Specifically, the slots reached only from the `k > 1` branches stay empty.
    @test opv.cache[1] isa AbstractMatrix
    @test all(isnothing, (opv.cache[i] for i in (2, 3, 4, 6, 7)))
    @test all(!isnothing, opm.cache)

    # Every in-place path a vector input can take still gives the right answer.
    v = rand(n1 * n2)
    w = zeros(m1 * m2)
    @test mul!(w, opv, v) ≈ AB * v

    w = rand(m1 * m2)
    orig_w = copy(w)
    @test mul!(w, opv, v, α, β) ≈ α * (AB * v) + β * orig_w

    if square
        # `c5` is the one slot besides `c1` a vector input can reach.
        opv_F = cache_operator(factorize(op), rand(n1 * n2))
        y = rand(m1 * m2)
        @test ldiv!(zeros(n1 * n2), opv_F, y) ≈ AB \ y
    end
end

@testset "TensorSumOperator" begin
    A = rand(3, 3)
    B = rand(4, 4)
    α = rand()
    β = rand()
    p = nothing
    t = 0.0

    expected = kron(A, Matrix(I, 4, 4)) + kron(Matrix(I, 3, 3), B)
    L = kronsum(MatrixOperator(A), MatrixOperator(B))
    L_from_matrices = kronsum(A, B)

    @test L isa TensorSumOperator{Float64}
    @test L_from_matrices isa TensorSumOperator{Float64}
    @test size(L) == size(expected)
    @test islinear(L)
    @test isconstant(L)
    @test !iscached(L)
    @test convert(AbstractMatrix, L) ≈ expected
    @test convert(AbstractMatrix, L_from_matrices) ≈ expected

    u = rand(size(L, 2), K)
    v = copy(u)
    w = zeros(size(L, 1), K)

    @test L * v ≈ expected * v
    @test L(v, u, p, t) ≈ expected * v

    L = cache_operator(L, u)
    @test iscached(L)

    mul!(w, L, v)
    @test w ≈ expected * v

    copy!(w, rand(size(L, 1), K))
    orig_w = copy(w)
    mul!(w, L, v, α, β)
    @test w ≈ α * expected * v + β * orig_w

    copy!(w, zeros(size(L, 1), K))
    L(w, v, u, p, t)
    @test w ≈ expected * v

    copy!(w, rand(size(L, 1), K))
    orig_w = copy(w)
    L(w, v, u, p, t, α, β)
    @test w ≈ α * expected * v + β * orig_w

    x = rand(size(L, 2))
    @test L * x ≈ expected * x
    @test convert(AbstractMatrix, L') ≈ expected'
    @test convert(AbstractMatrix, transpose(L)) ≈ transpose(expected)
    @test_throws AssertionError kronsum(rand(3, 4), B)

    A_update = MatrixOperator(zeros(3, 3); update_func = (A, u, p, t) -> p[1] * Matrix(I, 3, 3))
    B_update = MatrixOperator(zeros(4, 4); update_func = (B, u, p, t) -> p[2] * Matrix(I, 4, 4))
    L_update = kronsum(A_update, B_update)
    @test L_update(v, u, (2.0, 3.0), t) ≈ 5v
end

@testset "Sparse conversion tests" begin
    A = rand(2, 2)
    opA = MatrixOperator(A)
    # Construct operator including TensorProductOperator, AddedOperator,
    # IdentityOperator, ComposedOperator and ScaledOperator:
    opB = kron(2 * opA + I, opA * opA)
    B = kron(2 * A + I, A * A)
    opB_sparse = sparse(opB)
    @test opB_sparse ≈ B
    @test issparse(opB_sparse)

    L = BlockDiagonalOperator(opA, opB, NullOperator(2, 3))
    sL = sparse(L)
    @test issparse(sL)
    @test sL ≈ Matrix(L)

    # TensorSumOperator
    Lkronsum = kronsum(opB, opB)
    sLkronsum = sparse(Lkronsum)
    @test issparse(sLkronsum)
    @test sLkronsum ≈ Matrix(Lkronsum)
end

@testset "copyto! MatrixOperator <- MatrixOperator delegates to .A" begin
    # A MatrixOperator rhs must be unwrapped so the copy hits the underlying matrices'
    # specialized copyto! (sparse->sparse copies nzval/rowval/colptr) rather than the
    # generic element-wise fallback.
    for Amk in (() -> rand(6, 6), () -> sparse(Tridiagonal(rand(5), rand(6), rand(5))))
        A = Amk()
        src = MatrixOperator(copy(A))
        dest = MatrixOperator(zero(A))
        copyto!(dest, src)
        @test dest.A == A
        @test typeof(dest.A) === typeof(A)   # structure preserved (sparse stays sparse)
    end
    # plain-array rhs still works (existing method)
    d = MatrixOperator(zeros(3, 3))
    copyto!(d, [1.0 2 3; 4 5 6; 7 8 9])
    @test d.A == [1.0 2 3; 4 5 6; 7 8 9]
end
