module TestBlockBanded

using ArrayLayouts
using BandedMatrices
using BlockArrays
using BlockBandedMatrices
using FillArrays
using LinearAlgebra
using Test

import BlockBandedMatrices: MemoryLayout, ColumnMajor, BroadcastStyle,
                            BlockBandedStyle, blockrowsupport, blockcolsupport

@testset "BlockBandedMatrix" begin
    @test BroadcastStyle(BlockBandedMatrix) == BlockBandedStyle()

    @testset "BlockBandedMatrix constructors" begin
        l , u = 1,1
        N = M = 4
        cols = rows = 1:N

        @test Matrix(BlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows, cols, (l,u))) ==
            Array(BlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows, cols, (l,u))) ==
            zeros(Float64, 10, 10)

        @test Matrix(BlockBandedMatrix{Int}(Zeros(sum(rows),sum(cols)), rows,cols, (l,u))) ==
            zeros(Int, 10, 10)

        @test Matrix(BlockBandedMatrix(Eye(sum(rows)), rows,cols, (l,u))) ==
            Matrix{Float64}(I, 10, 10)

        @test Matrix(BlockBandedMatrix{Int}(Eye(sum(rows)), rows,cols, (l,u))) ==
            Matrix{Int}(I, 10, 10)

        @test Matrix(BlockBandedMatrix(I, rows,cols, (l,u))) ==
            Matrix{Float64}(I, 10, 10)

        @test Matrix(BlockBandedMatrix{Int}(I, rows,cols, (l,u))) ==
            Matrix{Int}(I, 10, 10)

        BS = BlockBandedMatrices.BlockBandedSizes(1:3, 1:3, 1, 1)
        B = BlockBandedMatrix{Float64}(undef, BS)
        @test blockbandwidths(B) == (1,1)
        blks = blocks(B)
        for (ind, blk) in zip(CartesianIndices(blks), blks)
            @test size(blk) == Tuple(ind)
        end

        B2 = BlockBandedMatrices._BlockBandedMatrix(1:30, 1:3, 1:3, (1, 1))
        @test @view(B2[Block(1,1)]) == reshape(1:1, 1, 1)
        @test @view(B2[Block(2,2)]) == hcat(5:6, 11:12)
        @test @view(B2[Block(3,3)]) == hcat(18:20, 23:25, 28:30)
    end

    @testset "BlockBandedMatrix block indexing" begin
        l , u = 1,1
        N = M = 4
        cols = rows = 1:N
        A = BlockBandedMatrix{Int}(undef, rows,cols, (l,u))
        A.data .= 1:70

        @test A[1,1] == 1
        @test A[1,3] == 10

        @test blockbandwidth(A,1)  == 1
        @test blockbandwidths(A) == (l,u)

        # check views of blocks are indexing correctly


        @test A[Block(1), Block(1)] isa Matrix
        @test A[Block(1), Block(1)] == A[Block(1,1)] == view(A, Block(1, 1)) == Matrix(view(A, Block(1,1)))
        @test A[1,1] == view(A,Block(1),Block(1))[1,1] == view(A,Block(1,1))[1,1] == A[Block(1,1)][1,1]  == A[Block(1),Block(1)][1,1] == 1
        @test A[2,1] == view(A,Block(2),Block(1))[1,1] == view(A,Block(2,1))[1,1] == 2
        @test A[3,1] == view(A,Block(2),Block(1))[2,1] == 3
        @test A[4,1] == 0
        @test A[1,2] == view(A,Block(1,2))[1,1] == 4
        @test A[1,3] == view(A,Block(1,2))[1,2] == view(A,Block(1,2))[2] == 10

        @test view(A, Block(3),Block(1)) ≈ [0,0,0]
        @test_throws BandError view(A, Block(3),Block(1))[1,1] = 4
        @test_throws BlockBoundsError view(A, Block(5,1))


        # test blocks
        V = view(A, Block(1,1))
        @test_throws BoundsError V[2,1]

        V = view(A, Block(3,4))
        @test V[3,1] == 45
        V[3,1] = -7
        @test V[3,1] == -7
        @test Matrix(V) isa Matrix{Int}
        @test Matrix{Float64}(V) isa Matrix{Float64}
        @test Matrix{Float64}(Matrix(V)) == Matrix{Float64}(V)
        @test A[4:6,7:10] ≈ Matrix(V)

        A[1,1] = -5
        @test A[1,1] == -5
        # setindex! should return the array
        @test setindex!(A, -6, 1, 3) === A
        @test A[1,3] == -6

        A[Block(3,4)] = Matrix(Ones{Int}(3,4))
        @test A[Block(3,4)] == Matrix(Ones{Int}(3,4))

        ## Bug in setindex!
        ret = BlockBandedMatrix(Zeros{Float64}((4,6)), [2,2], [2,2,2], (0,2))


        V = view(ret, Block(1), Block(2))
        V[1,1] = 2
        @test ret[1,2] == 0

        # setindex! should return the array
        @test setindex!(V, 4, 1, 2) === V
        @test V[1,2] == 4
    end

    @testset "blockcol/rowsupport" begin
        l , u = 1,2
        N = M = 4
        cols = rows = 1:N
        A = BlockBandedMatrix{Int}(undef, rows,cols, (l,u))
        A.data .= 1:length(A.data)

        @test @inferred(blockrowsupport(A, Block(1))) == Block.(1:3)
        @test blockrowsupport(A, Block(2)) == Block.(1:4)
        @test blockrowsupport(A, Block(3)) == Block.(2:4)

        @test @inferred(blockcolsupport(A, Block(1))) == Block.(1:2)
        @test blockcolsupport(A, Block(2)) == Block.(1:3)
        @test blockcolsupport(A, Block(3)) == Block.(1:4)
        @test blockcolsupport(A, Block(4)) == Block.(2:4)

        @test @inferred(blockrowsupport(A, Block.(3:4))) == Block.(2:4)
        @test @inferred(blockcolsupport(A, Block.(1:2))) == Block.(1:3)

        @test @inferred(blockcolsupport(A, Block.(1:0))) == Block.(1:0)
    end

    @testset "block-banded matrix interface for blockranges" begin
        l , u = 1,1
        N = M = 4
        cols = rows = 1:N
        A = BlockBandedMatrix{Int}(undef, rows,cols, (l,u))
        A.data .= 1:length(A.data)

        V = view(A, Block.(2:3), Block.(3:4))
        @test isblockbanded(V)

        B = BlockBandedMatrix(V)
        @test B isa BlockBandedMatrix
        @test blockbandwidths(V) == blockbandwidths(B) == (2,0)
        @test B == V == A[Block.(2:3), Block.(3:4)]

        @test A[Block.(2:3), Block.(3:4)] isa BlockBandedMatrix

        x = randn(size(B,2))
        y = similar(x, size(B,1))
        @test (similar(y) .= MulAdd(B, x)) == (similar(y) .= MulAdd(V,x))
    end

    @testset "BlockBandedMatrix indexing" begin
        l , u = 1,1
        N = M = 10
        cols = rows = 1:N
        A = BlockBandedMatrix{Float64}(undef, rows,cols, (l,u))
        A.data .= 1:length(A.data)

        A[1,1] = 5
        @test A[1,1] == 5

        @test_throws BandError A[1,4] = 5
        A[1,4] = 0
        @test A[1,4] == 0

        @test A[1:10,1:10] ≈ Matrix(A)[1:10,1:10]

        l , u = 2,1
        N = M = 5
        cols = rows = 1:N

        A = BlockBandedMatrix{Int}(undef, rows,cols, (l,u))
        A.data .= 1:length(A.data)

        @test A[1,2] == 7
        A[1,2] = -5
        @test A[1,2] == -5
    end

    @testset "BlockBandedMatrix BLAS arithmetic" begin
        l , u = 1,1
        N = M = 10
        cols = rows = fill(100,N)
        A = BlockBandedMatrix{Float64}(undef, rows,cols, (l,u))
            A.data .= 1:length(A.data)

        V = view(A, Block(N,N))
        @test MemoryLayout(typeof(V)) == ColumnMajor()

        Y = zeros(cols[N], cols[N])
        @time axpy!(2.0, V, Y)
        @test Y ≈ 2A[Block(N,N)]

        Y = BandedMatrix(Zeros(cols[N], cols[N]), (0, 0))
        @test_throws BandError axpy!(2.0, V, Y)

        AN = A[Block(N,N)]
        @time axpy!(2.0, V, V)
        @test A[Block(N,N)] ≈ 3AN

        A = BlockBandedMatrix(Ones{Float64}((4,6)), [2,2], [2,2,2], (0,2))
        B = BlockBandedMatrix(Ones{Float64}((6,6)), [2,2,2], [2,2,2], (0,1))
        @test sum(A) == 20
        @test sum(B) == 20
        C = BlockBandedMatrix{Float64}(undef, [2,2], [2,2,2], (0,3))
        @test mul!(C,A,B) == ArrayLayouts.materialize!(MulAdd(1.0,A,B,0.0,similar(C))) == A*B
        AB = A*B
        @test AB isa BlockBandedMatrix
        @test Matrix(AB) ≈ Matrix(A)*Matrix(B)
    end

    @testset "BlockBandedMatrix fill and copy" begin
        l , u = 1,1
        N = M = 10
        cols = rows = 1:N
        A = BlockBandedMatrix{Float64}(undef, rows,cols, (l,u))
            A.data .= randn(length(A.data))

        fill!(view(A, Block(2,1)), 2.0)
        @test A[Block(2,1)] ≈ fill(2.0, 2)

        @test_throws BandError fill!(view(A, Block(3,1)), 2.0)

        fill!(view(A, Block(3,1)), 0.0)
        @test A[Block(3,1)] ≈ zeros(3)


        @test_throws BandError fill!(A, 2.0)
        fill!(A, 0.0)
        @test Matrix(A) == zeros(size(A))
    end

    @testset "BlockBandedMatrix type inferrence bug (#9)" begin
        s = BlockBandedMatrices.BlockBandedSizes([1, 2], [1, 2], 0, 0)
        f(s) = s.block_starts.data
        @inferred(f(s))
    end

    @testset "BlockVector" begin
        s = BlockVector(1:21, 1:6)
        @test BlockBandedMatrices.subblockbandwidths(s) == (5, 0)
    end
end

end # module
