using Test
using BenchmarkTools

using Sbplib.Grids
using Sbplib.SbpOperators

using Sbplib.LazyTensors
using Sbplib.RegionIndices

@testset "Allocations" begin
    stencil_set = read_stencil_set(sbp_operators_path()*"standard_diagonal.toml", order=4)

    @testset "1D" begin
        g₁ = EquidistantGrid(15, 0.,1.)

        H = inner_product(g₁, stencil_set)
        H⁻¹ = inverse_inner_product(g₁, stencil_set)
        D₂ = second_derivative(g₁, stencil_set, 1)

        eₗ = boundary_restriction(g₁, stencil_set, CartesianBoundary{1,Lower}())
        eᵣ = boundary_restriction(g₁, stencil_set, CartesianBoundary{1,Upper}())

        dₗ = normal_derivative(g₁, stencil_set, CartesianBoundary{1,Lower}())
        dᵣ = normal_derivative(g₁, stencil_set, CartesianBoundary{1,Upper}())



        @testset "Derivative operator" begin
            v = rand(size(g₁)...)
            @test (@ballocated LazyTensors.apply($D₂, $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($D₂, $v, 6)) == 0
            @test (@ballocated LazyTensors.apply($D₂, $v, 15)) == 0
        end

        @testset "inner_product operator" begin
            v = rand(size(g₁)...)

            @test (@ballocated LazyTensors.apply($H, $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 6)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 15)) == 0

            @test (@ballocated LazyTensors.apply($(H∘H), $v, 5)) == 0
        end

        @testset "inverse_inner_product operator" begin
              v = rand(size(g₁)...)
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 6)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 15)) == 0
        end

        @testset "boundary operators" begin
            v = rand(size(g₁)...)
            @test (@ballocated LazyTensors.apply($eₗ, $v)) == 0
            @test (@ballocated LazyTensors.apply($eᵣ, $v)) == 0
            @test (@ballocated LazyTensors.apply($dₗ, $v)) == 0
            @test (@ballocated LazyTensors.apply($dᵣ, $v)) == 0
        end

        @testset "boundary operator transposes" begin
            v = fill(1.)
            @test (@ballocated LazyTensors.apply($eₗ', $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($eₗ', $v, 7)) == 0
            @test (@ballocated LazyTensors.apply($eₗ', $v, 15)) == 0

            @test (@ballocated LazyTensors.apply($eᵣ', $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($eᵣ', $v, 7)) == 0
            @test (@ballocated LazyTensors.apply($eᵣ', $v, 15)) == 0

            @test (@ballocated LazyTensors.apply($dₗ', $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($dₗ', $v, 7)) == 0
            @test (@ballocated LazyTensors.apply($dₗ', $v, 15)) == 0

            @test (@ballocated LazyTensors.apply($dᵣ', $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($dᵣ', $v, 7)) == 0
            @test (@ballocated LazyTensors.apply($dᵣ', $v, 15)) == 0
        end

        @testset "sat terms" begin
            v = rand(size(g₁)...)
            neumannSATₗ = H⁻¹∘eₗ'∘dₗ
            neumannSATᵣ = H⁻¹∘eᵣ'∘dᵣ


            @test (@ballocated LazyTensors.apply($neumannSATₗ, $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSATₗ, $v, 6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSATₗ, $v, 15)) == 0

            @test (@ballocated LazyTensors.apply($neumannSATᵣ, $v, 1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSATᵣ, $v, 6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSATᵣ, $v, 15)) == 0
        end
    end


    @testset "2D" begin
        g₂ = EquidistantGrid((15,15), (0.,0.),(1.,1.))

        H = inner_product(g₂, stencil_set)
        H⁻¹ = inverse_inner_product(g₂, stencil_set)
        D₂x = second_derivative(g₂, stencil_set, 1)
        D₂y = second_derivative(g₂, stencil_set, 2)

        e₁ₗ = boundary_restriction(g₂, stencil_set, CartesianBoundary{1,Lower}())
        e₁ᵤ = boundary_restriction(g₂, stencil_set, CartesianBoundary{1,Upper}())
        e₂ₗ = boundary_restriction(g₂, stencil_set, CartesianBoundary{2,Lower}())
        e₂ᵤ = boundary_restriction(g₂, stencil_set, CartesianBoundary{2,Upper}())

        d₁ₗ = normal_derivative(g₂, stencil_set, CartesianBoundary{1,Lower}())
        d₁ᵤ = normal_derivative(g₂, stencil_set, CartesianBoundary{1,Upper}())
        d₂ₗ = normal_derivative(g₂, stencil_set, CartesianBoundary{2,Lower}())
        d₂ᵤ = normal_derivative(g₂, stencil_set, CartesianBoundary{2,Upper}())

        H₁ₗ = inner_product(boundary_grid(g₂, CartesianBoundary{1,Lower}()), stencil_set)
        H₁ᵤ = inner_product(boundary_grid(g₂, CartesianBoundary{1,Upper}()), stencil_set)
        H₂ₗ = inner_product(boundary_grid(g₂, CartesianBoundary{2,Lower}()), stencil_set)
        H₂ᵤ = inner_product(boundary_grid(g₂, CartesianBoundary{2,Upper}()), stencil_set)


        @testset "Derivative operator" begin
            v = rand(size(g₂)...)
            @test (@ballocated LazyTensors.apply($D₂x, $v,  1,  7)) == 0
            @test (@ballocated LazyTensors.apply($D₂x, $v,  6,  7)) == 0
            @test (@ballocated LazyTensors.apply($D₂x, $v, 15,  7)) == 0

            @test (@ballocated LazyTensors.apply($D₂y, $v,  7,  1)) == 0
            @test (@ballocated LazyTensors.apply($D₂y, $v,  7,  6)) == 0
            @test (@ballocated LazyTensors.apply($D₂y, $v,  7, 15)) == 0

            @test (@ballocated LazyTensors.apply($(D₂x∘D₂y), $v,   1,  1)) == 0
            @test (@ballocated LazyTensors.apply($(D₂x∘D₂y), $v,   7,  6)) == 0
            @test (@ballocated LazyTensors.apply($(D₂x∘D₂y), $v,  15, 15)) == 0
        end

        @testset "inner_product operator" begin
            v = rand(size(g₂)...)
            @test (@ballocated LazyTensors.apply($H, $v, 1,  1)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 1,  6)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 1, 15)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 6,  1)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 6,  6)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 6, 15)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 15,  1)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 15,  6)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, 15, 15)) == 0


            @test (@ballocated LazyTensors.apply($(H∘H), $v, 5, 5)) == 0
        end

        @testset "inverse_inner_product operator" begin
            v = rand(size(g₂)...)
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 1,  1)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 1,  6)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 1, 15)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 6,  1)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 6,  6)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 6, 15)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 15,  1)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 15,  6)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, 15, 15)) == 0
        end

        @testset "boundary operators" begin
            v = rand(size(g₂)...)
            @test (@ballocated LazyTensors.apply($e₁ₗ, $v,  1)) == 0
            @test (@ballocated LazyTensors.apply($e₁ᵤ, $v,  5)) == 0
            @test (@ballocated LazyTensors.apply($e₂ₗ, $v, 15)) == 0
            @test (@ballocated LazyTensors.apply($e₂ᵤ, $v,  3)) == 0

            @test (@ballocated LazyTensors.apply($d₁ₗ, $v,  5)) == 0
            @test (@ballocated LazyTensors.apply($d₁ᵤ, $v, 15)) == 0
            @test (@ballocated LazyTensors.apply($d₂ₗ, $v,  1)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ, $v,  5)) == 0
        end

        @testset "boundary operator transposes" begin
            v = rand(first(size(g₂)))

            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 1,  1)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 1,  6)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 1, 15)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 6,  1)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 6,  6)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 6, 15)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 15,  1)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 15,  6)) == 0
            @test (@ballocated LazyTensors.apply($e₁ₗ', $v, 15, 15)) == 0

            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 1,  1)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 1,  6)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 1, 15)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 6,  1)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 6,  6)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 6, 15)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 15,  1)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 15,  6)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ', $v, 15, 15)) == 0
        end

        @testset "sat terms" begin
            v = rand(size(g₂)...)
            u = rand(size(g₂)[1])
            # neumannSAT₁ₗ = H⁻¹∘e₁ₗ'∘H₁ₗ∘d₁ₗ
            # neumannSAT₂ᵤ = H⁻¹∘e₂ᵤ'∘H₂ᵤ∘d₂ᵤ

            neumannSAT₁ₗ = e₁ₗ'∘d₁ₗ
            neumannSAT₂ᵤ = e₂ᵤ'∘d₂ᵤ

            # indices = [1,6,15]
            indices = [1]

            @testset for i ∈ indices
                @testset for j ∈ indices
                    @test        (@ballocated LazyTensors.apply($(e₁ₗ'∘d₁ₗ),         $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(d₁ₗ'∘e₁ₗ),         $v, $i, $j)) == 0

                    @test_broken (@ballocated LazyTensors.apply($(H⁻¹∘e₁ₗ'∘H₁ₗ∘d₁ₗ), $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(e₁ₗ'∘H₁ₗ∘d₁ₗ∘H⁻¹), $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(e₁ₗ'∘H₁ₗ∘d₁ₗ),     $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(H⁻¹∘e₁ₗ'∘d₁ₗ),     $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(H⁻¹∘D₂x),          $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(H⁻¹∘D₂y),          $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(H⁻¹∘D₂x∘D₂y),      $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(D₂x∘e₁ₗ'∘H₁ₗ∘d₁ₗ), $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(D₂y∘e₁ₗ'∘H₁ₗ∘d₁ₗ), $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(D₂x∘e₁ₗ'∘d₁ₗ),     $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(D₂y∘e₁ₗ'∘d₁ₗ),     $v, $i, $j)) == 0

                    @test (@ballocated LazyTensors.apply($(H⁻¹∘e₁ₗ'), $u, $i, $j)) == 0
                    @test (@ballocated LazyTensors.apply($(e₁ₗ'∘H₁ₗ), $u, $i, $j)) == 0
                    @test (@ballocated LazyTensors.apply($(H₁ₗ∘d₁ₗ),  $v, $i)) == 0

                    @test_broken (@ballocated LazyTensors.apply($(H⁻¹∘e₁ₗ'∘H₁ₗ), $u, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(H₁ₗ∘e₁ₗ∘H⁻¹),  $v, $i)) == 0
                    @test        (@ballocated LazyTensors.apply($(H₁ₗ∘d₁ₗ∘H⁻¹),  $v, $i)) == 0

                    @test_broken (@ballocated LazyTensors.apply($(e₁ₗ'∘e₁ₗ∘e₁ₗ'∘e₁ₗ), $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(e₁ₗ'∘e₁ₗ∘e₁ₗ'),     $u, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(e₁ₗ∘e₁ₗ'∘e₁ₗ),      $v, $i)) == 0
                    @test        (@ballocated LazyTensors.apply($(e₁ₗ'∘e₁ₗ),          $v, $i, $j)) == 0
                    @test        (@ballocated LazyTensors.apply($(e₁ₗ∘e₁ₗ'),          $u, $i)) == 0


                    @test (@ballocated LazyTensors.apply($H, $v, $i, $j)) == 0
                    @test (@ballocated LazyTensors.apply($(H∘H), $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(H∘H∘H), $v, $i, $j)) == 0
                    @test_broken (@ballocated LazyTensors.apply($(H∘H∘H∘H), $v, $i, $j)) == 0
                end
                @test (@ballocated LazyTensors.apply($(e₁ₗ∘d₁ₗ'), $u, $i)) == 0
                @test (@ballocated LazyTensors.apply($(d₁ₗ∘e₁ₗ'), $u, $i)) == 0

                @test (@ballocated LazyTensors.apply($(e₁ₗ∘H∘d₁ₗ'), $u, $i)) == 0
                @test (@ballocated LazyTensors.apply($(d₁ₗ∘H∘e₁ₗ'), $u, $i)) == 0

                @test (@ballocated LazyTensors.apply($(H₁ₗ∘e₁ₗ∘d₁ₗ'), $u, $i)) == 0
                @test (@ballocated LazyTensors.apply($(H₁ₗ∘d₁ₗ∘e₁ₗ'), $u, $i)) == 0

                @test_broken (@ballocated LazyTensors.apply($(H₁ₗ∘e₁ₗ∘H∘d₁ₗ'), $u, $i)) == 0
                @test_broken (@ballocated LazyTensors.apply($(H₁ₗ∘d₁ₗ∘H∘e₁ₗ'), $u, $i)) == 0

                @test_broken (@ballocated LazyTensors.apply($(e₁ₗ∘H∘d₁ₗ'∘H₁ₗ), $u, $i)) == 0
                @test_broken (@ballocated LazyTensors.apply($(d₁ₗ∘H∘e₁ₗ'∘H₁ₗ), $u, $i)) == 0
            end
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  1,  1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  1,  6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  1, 15)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  6,  1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  6,  6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  6, 15)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v, 15,  1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v, 15,  6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v, 15, 15)) == 0

            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  1,  1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  1,  6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  1, 15)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  6,  1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  6,  6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  6, 15)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v, 15,  1)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v, 15,  6)) == 0
            @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v, 15, 15)) == 0
        end

    end
end