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}())


        indices = [1, 6, 15]

        v = rand(size(g₁)...)
        u = fill(1.)

        neumannSATₗ = H⁻¹∘eₗ'∘dₗ # TODO: Remove
        neumannSATᵣ = H⁻¹∘eᵣ'∘dᵣ # TODO: Remove

        @testset for i ∈ indices
            @test (@ballocated LazyTensors.apply($D₂, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($H, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($(H∘H), $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($H⁻¹, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($eₗ', $u, $i)) == 0
            @test (@ballocated LazyTensors.apply($eᵣ', $u, $i)) == 0
            @test (@ballocated LazyTensors.apply($dₗ', $u, $i)) == 0
            @test (@ballocated LazyTensors.apply($dᵣ', $u, $i)) == 0
            @test (@ballocated LazyTensors.apply($neumannSATₗ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($neumannSATᵣ, $v, $i)) == 0
        end

        @testset "boundary operators" begin
            @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
    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)


        # neumannSAT₁ₗ = H⁻¹∘e₁ₗ'∘H₁ₗ∘d₁ₗ
        # neumannSAT₂ᵤ = H⁻¹∘e₂ᵤ'∘H₂ᵤ∘d₂ᵤ

        neumannSAT₁ₗ = e₁ₗ'∘d₁ₗ # TODO: Remove
        neumannSAT₂ᵤ = e₂ᵤ'∘d₂ᵤ # TODO: Remove

        indices = [1,6,15]

        v = rand(size(g₂)...)
        u = rand(first(size(g₂)))

        @testset for i ∈ indices
            @testset for j ∈ indices
                @test (@ballocated LazyTensors.apply($D₂x, $v,  $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($D₂y, $v,  $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($(D₂x∘D₂y), $v, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($H, $v, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($(H∘H), $v, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($H⁻¹, $v, $i, $j)) == 0

                @test (@ballocated LazyTensors.apply($e₁ₗ', $u, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($e₁ᵤ', $u, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($e₂ₗ', $u, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($e₂ᵤ', $u, $i, $j)) == 0

                @test (@ballocated LazyTensors.apply($d₁ₗ', $u, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($d₁ᵤ', $u, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($d₂ₗ', $u, $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($d₂ᵤ', $u, $i, $j)) == 0

                @test (@ballocated LazyTensors.apply($neumannSAT₁ₗ, $v,  $i, $j)) == 0
                @test (@ballocated LazyTensors.apply($neumannSAT₂ᵤ, $v,  $i, $j)) == 0



                ## Experiments
                @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₁ₗ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($e₁ᵤ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($e₂ₗ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($e₂ᵤ, $v, $i)) == 0

            @test (@ballocated LazyTensors.apply($d₁ₗ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($d₁ᵤ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($d₂ₗ, $v, $i)) == 0
            @test (@ballocated LazyTensors.apply($d₂ᵤ, $v, $i)) == 0

            ## Experiments
            @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
    end
end