include("dependencies_for_runtests.jl")

using Statistics
using Oceananigans.Architectures: on_architecture
using Oceananigans.AbstractOperations: BinaryOperation
using Oceananigans.Fields: CenterField, ZFaceField, compute_at!, @compute, reverse_cumsum!
using Oceananigans.BoundaryConditions: fill_halo_regions!

trilinear(x, y, z) = x + y + z
interior_array(a, i, j, k) = Array(interior(a, i, j, k))

@testset "Fields computed by Reduction" begin
    @info "Testing Fields computed by reductions..."

    for arch in archs
        arch_str = string(typeof(arch))

        regular_grid = RectilinearGrid(arch, size = (2, 2, 2),
                                       x = (0, 2), y = (0, 2), z = (0, 2),
                                       topology = (Periodic, Periodic, Bounded))

        xy_regular_grid = RectilinearGrid(arch, size=(2, 2, 2),
                                          x=(0, 2), y=(0, 2), z=[0, 1, 2],
                                          topology = (Periodic, Periodic, Bounded))

        xz_regular_grid = RectilinearGrid(arch, size=(2, 2, 2),
                                          x=(0, 2), y=[0, 1, 2], z=(0, 2),
                                          topology = (Periodic, Periodic, Bounded))

        @testset "Averaged and integrated fields [$arch_str]" begin
            @info "  Testing averaged and integrated Fields [$arch_str]"

            for (name, grid) in [(:regular_grid => regular_grid),
                                 (:xy_regular_grid => xy_regular_grid),
                                 (:xz_regular_grid => xz_regular_grid)]

                @info "    Testing averaged and integrated Fields on $name..."

                Nx, Ny, Nz = size(grid)

                w = ZFaceField(grid)
                T = CenterField(grid)
                ζ = Field{Face, Face, Face}(grid)
                η = Field{Center, Center, Nothing}(grid)

                set!(T, trilinear)
                set!(w, trilinear)
                set!(ζ, trilinear)

                z_top = @allowscalar Oceananigans.Grids.znode(grid.Nz+1, grid, Face())
                set!(η, (x, y) -> trilinear(x, y, z_top))

                @compute Txyz = Field(Average(T, dims=(1, 2, 3)))

                # Note: halo regions must be *filled* prior to computing an average
                # if the average within halo regions is to be correct.
                fill_halo_regions!(T)
                @compute Txy = Field(Average(T, dims=(1, 2)))

                fill_halo_regions!(T)
                @compute Tx = Field(Average(T, dims=1))

                @compute wxyz = Field(Average(w, dims=(1, 2, 3)))
                @compute wxy = Field(Average(w, dims=(1, 2)))
                @compute wx = Field(Average(w, dims=1))

                @compute ζxyz = Field(Average(ζ, dims=(1, 2, 3)))
                @compute ζxy = Field(Average(ζ, dims=(1, 2)))
                @compute ζx = Field(Average(ζ, dims=1))

                @compute Θxyz = Field(Integral(T, dims=(1, 2, 3)))
                @compute Θxy = Field(Integral(T, dims=(1, 2)))
                @compute Θx = Field(Integral(T, dims=1))

                @compute Wxyz = Field(Integral(w, dims=(1, 2, 3)))
                @compute Wxy = Field(Integral(w, dims=(1, 2)))
                @compute Wx = Field(Integral(w, dims=1))

                @compute Zxyz = Field(Integral(ζ, dims=(1, 2, 3)))
                @compute Zxy = Field(Integral(ζ, dims=(1, 2)))
                @compute Zx = Field(Integral(ζ, dims=1))

                @compute ηxyz = Field(Average(η))
                @compute Ηxyz = Field(Integral(η))

                @test Field(Average(η)) == Field(Average(η, dims=(1, 2, 3)))
                @test Field(Integral(η)) == Field(Integral(η, dims=(1, 2, 3)))

                # CumulativeIntegral is only supported for Bounded dimensions
                # x and y are Periodic, so only test z direction
                @test_throws ArgumentError CumulativeIntegral(T, dims=1)
                @test_throws ArgumentError CumulativeIntegral(T, dims=2)
                @test_throws ArgumentError CumulativeIntegral(w, dims=1)
                @test_throws ArgumentError CumulativeIntegral(w, dims=2)
                @test_throws ArgumentError CumulativeIntegral(ζ, dims=1)
                @test_throws ArgumentError CumulativeIntegral(ζ, dims=2)

                # z is Bounded, so cumulative integrals work
                @compute Tcz = Field(CumulativeIntegral(T, dims=3))
                @compute wcz = Field(CumulativeIntegral(w, dims=3))
                @compute ζcz = Field(CumulativeIntegral(ζ, dims=3))

                @compute Trz = Field(CumulativeIntegral(T, dims=3, reverse=true))
                @compute wrz = Field(CumulativeIntegral(w, dims=3, reverse=true))
                @compute ζrz = Field(CumulativeIntegral(ζ, dims=3, reverse=true))

                if name ∈ (:regular_grid, :xy_regular_grid)
                    for T′ in (Tx, Txy)
                        @test T′.operand.operand === T
                    end

                    for w′ in (wx, wxy)
                        @test w′.operand.operand === w
                    end

                    for ζ′ in (ζx, ζxy)
                        @test ζ′.operand.operand === ζ
                    end
                end

                for f in (Txyz, wxyz, ζxyz, ηxyz, wx, wxy, Tx, Txy, ζx, ζxy, Wx, Wxy, Θx, Θxy, Zx, Zxy)
                    @test f.operand isa Reduction
                end

                for f in (Tcz, Trz, wcz, wrz, ζcz, ζrz)
                    @test f.operand isa Accumulation
                end

                for f in (Txyz, wxyz, ζxyz, ηxyz, wx, wxy, Tx, Txy, ζx, ζxy)
                    @test f.operand.scan! === Oceananigans.AbstractOperations.average!
                end

                @test Ηxyz.operand.scan! === Oceananigans.AbstractOperations.sum!

                for f in (Txyz, wxyz, ζxyz, ηxyz, wx, wxy, Tx, Txy, ζx, ζxy)
                    @test f.operand.scan! === Oceananigans.AbstractOperations.average!
                end

                for f in (Tcz, wcz, ζcz)
                    @test f.operand.scan! === cumsum!
                end

                for f in (Trz, wrz, ζrz)
                    @test f.operand.scan! === reverse_cumsum!
                end

                # Different behavior for regular grid z vs not.
                if grid === regular_grid
                    @test Txyz.operand.scan! === Oceananigans.AbstractOperations.average!
                    @test wxyz.operand.scan! === Oceananigans.AbstractOperations.average!
                    @test Txyz.operand.operand === T
                    @test wxyz.operand.operand === w
                else
                    @test Txyz.operand.scan! === Oceananigans.AbstractOperations.average!
                    @test wxyz.operand.scan! === Oceananigans.AbstractOperations.average!
                    @test Txyz.operand.operand isa BinaryOperation
                    @test wxyz.operand.operand isa BinaryOperation
                end

                @test Tx.operand.dims === tuple(1)
                @test wx.operand.dims === tuple(1)
                @test Txy.operand.dims === (1, 2)
                @test wxy.operand.dims === (1, 2)
                @test Txyz.operand.dims === (1, 2, 3)
                @test wxyz.operand.dims === (1, 2, 3)

                @test @allowscalar Txyz[1, 1, 1] ≈ 3
                @test interior_array(Txy, 1, 1, :) ≈ [2.5, 3.5]
                @test interior_array(Tx, 1, :, :) ≈ [[2, 3] [3, 4]]

                @test @allowscalar wxyz[1, 1, 1] ≈ 3
                @test interior_array(wxy, 1, 1, :) ≈ [2, 3, 4]
                @test interior_array(wx, 1, :, :) ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]]

                @test @allowscalar ηxyz[1, 1, 1] ≈ 4
                @test @allowscalar Ηxyz[1, 1, 1] ≈ 16

                averages_1d  = (Tx, wx, ζx)
                integrals_1d = (Θx, Wx, Zx)

                for (a, i) in zip(averages_1d, integrals_1d)
                    @test interior(i) == 2 .* interior(a)
                end

                averages_2d  = (Txy, wxy, ζxy)
                integrals_2d = (Θxy, Wxy, Zxy)

                for (a, i) in zip(averages_2d, integrals_2d)
                    @test interior(i) == 4 .* interior(a)
                end

                # T(x, y, z) = x + y + z
                # T at Center locations in z: T(0.5, 0.5, z_center) with z_center = [0.5, 1.5]
                # T = [1.5, 2.5], Δz = 1
                # Tcz is at Face locations in z (3 points for Nz=2 Bounded)
                # Tcz = [0, T[1]*Δz, T[1]*Δz + T[2]*Δz] = [0, 1.5, 4]
                @test interior_array(Tcz, 1, 1, :) ≈ [0, 1.5, 4]
                @test interior_array(Trz, 1, 1, :) ≈ [4, 2.5, 0]

                # w(x, y, z) = x + y + z
                # w at Face locations in z: w(0.5, 0.5, z_face) with z_face = [0, 1, 2]
                # w = [1, 2, 3], Δz = 1
                # wcz is at Center locations in z (2 points for Nz=2 Bounded)
                # wcz = [0, w[1]*Δz] = [0, 1]
                @test interior_array(wcz, 1, 1, :) ≈ [0, 1]
                # wrz (reverse): starts from end, accumulates backwards
                # wrz = [w[2]*Δz, 0] = [2, 0]
                @test interior_array(wrz, 1, 1, :) ≈ [2, 0]

                @compute Txyz = @allowscalar Field(Average(T, condition=T.>3))
                @compute Txy = @allowscalar Field(Average(T, dims=(1, 2), condition=T.>3))
                @compute Tx = @allowscalar Field(Average(T, dims=1, condition=T.>2))

                @test @allowscalar Txyz[1, 1, 1] ≈ 3.75
                @test interior_array(Txy, 1, 1, :) ≈ [3.5, 11.5/3]
                @test interior_array(Tx, 1, :, :) ≈ [[2.5, 3] [3, 4]]

                @compute wxyz = @allowscalar Field(Average(w, condition=w.>3))
                @compute wxy = @allowscalar Field(Average(w, dims=(1, 2), condition=w.>2))
                @compute wx = @allowscalar Field(Average(w, dims=1, condition=w.>1))

                @test @allowscalar wxyz[1, 1, 1] ≈ 4.25
                @test interior_array(wxy, 1, 1, :) ≈ [3, 10/3, 4]
                @test interior_array(wx, 1, :, :) ≈ [[2, 2.5] [2.5, 3.5] [3.5, 4.5]]

                # A bit more for cumulative integral (only z direction since x, y are Periodic)
                @compute T2cz = Field(CumulativeIntegral(2 * T, dims=3))
                @compute T2rz = Field(CumulativeIntegral(2 * T, dims=3, reverse=true))

                # T(x, y, z) = x + y + z
                # 2 * T(0.5, 0.5, z) at z_center = [0.5, 1.5] gives [3, 5], Δz = 1
                # T2cz = [0, 3*1, 3*1 + 5*1] = [0, 3, 8]
                @test interior_array(T2cz, 1, 1, :) ≈ [0, 3, 8]
                @test interior_array(T2rz, 1, 1, :) ≈ [8, 5, 0]
            end

            # Test whether a race condition gets hit for averages over large fields
            big_grid = RectilinearGrid(arch,
                                       topology = (Periodic, Periodic, Bounded),
                                       size = (256, 256, 128),
                                       x = (0, 2),
                                       y = (0, 2),
                                       z = (0, 2))

            c = CenterField(big_grid)
            c .= 1

            C = Field(Average(c, dims=(1, 2)))

            # Test that the mean consistently returns 1 at every z for many evaluations
            # Warm up
            for i = 1:10
                sum!(C, C.operand.operand)
            end

            results = []
            for i = 1:10
                mean!(C, C.operand.operand)
                push!(results, all(interior(C) .== 1))
            end

            @test mean(results) == 1.0
        end

        @testset "Unicode integral operators [$arch_str]" begin
            @info "  Testing Unicode integral operators [$arch_str]"

            grid = RectilinearGrid(arch, size = (2, 2, 2),
                                   x = (0, 2), y = (0, 2), z = (0, 2),
                                   topology = (Periodic, Periodic, Bounded))

            T = CenterField(grid)
            set!(T, trilinear)
            fill_halo_regions!(T)

            unicode_integrals = (∫dx, ∫dy, ∫dz, ∫∫dxdy, ∫∫dxdz, ∫∫dydz, ∫∫∫dxdydz)
            integrated_dims = (tuple(1), tuple(2), tuple(3), (1, 2), (1, 3), (2, 3), (1, 2, 3))

            for (∫dξ, dims) in zip(unicode_integrals, integrated_dims)
                ∫T = ∫dξ(T)
                @test ∫T isa Field
                @test ∫T.operand isa Reduction
                @test ∫T.operand.dims === dims
                @test ∫T == Field(Integral(T; dims))
            end

            @test ∫dV === ∫∫∫dxdydz

            # The returned Field is computed on construction: T = x + y + z integrates to 24 over [0, 2]³
            @test interior_array(∫dV(T), :, :, :)[1, 1, 1] ≈ 24

            # ... and still composes with other abstract operations
            @test Field(T / ∫dV(T)) ≈ Field(T / 24)

            # CumulativeIntegral is only supported over a single Bounded dimension
            @test_throws ArgumentError ∫dx(T, cumulative=true)
            @test_throws ArgumentError ∫dy(T, cumulative=true)
            @test_throws ArgumentError ∫∫dxdy(T, cumulative=true)

            Tcz = ∫dz(T, cumulative=true)
            Trz = ∫dz(T, cumulative=true, reverse=true)

            @test Tcz.operand isa Accumulation
            @test Tcz.operand.scan! === cumsum!
            @test Trz.operand.scan! === reverse_cumsum!

            @test interior_array(Tcz, 1, 1, :) ≈ [0, 1.5, 4]
            @test interior_array(Trz, 1, 1, :) ≈ [4, 2.5, 0]

            upper_half(i, j, k, grid, c) = Oceananigans.Grids.znode(k, grid, Center()) > 1

            ∫TdV = ∫dV(T, condition=upper_half)
            @test interior_array(∫TdV, :, :, :)[1, 1, 1] ≈ 14
            @test ∫TdV == Field(Integral(T, condition=upper_half))
        end

        @testset "Allocating reductions [$arch_str]" begin
            @info "  Testing allocating reductions [$arch_str]"

            grid = RectilinearGrid(arch, size = (2, 2, 2),
                                   x = (0, 2), y = (0, 2), z = (0, 2),
                                   topology = (Periodic, Periodic, Bounded))

            w = ZFaceField(grid)
            T = CenterField(grid)
            set!(T, trilinear)
            set!(w, trilinear)
            fill_halo_regions!(T)
            fill_halo_regions!(w)

            @compute Txyz = Field(Average(T, dims=(1, 2, 3)))
            @compute Txy = Field(Average(T, dims=(1, 2)))
            @compute Tx = Field(Average(T, dims=1))

            @compute wxyz = Field(Average(w, dims=(1, 2, 3)))
            @compute wxy = Field(Average(w, dims=(1, 2)))
            @compute wx = Field(Average(w, dims=1))

            # Mean
            @test @allowscalar Txyz[1, 1, 1] == mean(T)
            @test interior(Txy) == interior(mean(T, dims=(1, 2)))
            @test interior(Tx) == interior(mean(T, dims=1))

            @test @allowscalar wxyz[1, 1, 1] == mean(w)
            @test interior(wxy) == interior(mean(w, dims=(1, 2)))
            @test interior(wx) == interior(mean(w, dims=1))

            # Maximum and minimum
            @test maximum(T) == maximum(interior(T))
            @test minimum(T) == minimum(interior(T))

            for dims in (1, (1, 2))
                @test interior(minimum(T; dims)) == minimum(interior(T); dims)
                @test interior(minimum(T; dims)) == minimum(interior(T); dims)
            end
        end

        @testset "Conditional computation of averaged Fields [$(typeof(arch))]" begin
            @info "  Testing conditional computation of averaged Fields [$(typeof(arch))]"
            for FT in float_types
                grid = RectilinearGrid(arch, FT, size=(2, 2, 2), extent=(1, 1, 1))
                c = CenterField(grid)

                for dims in (1, 2, 3, (1, 2), (2, 3), (1, 3), (1, 2, 3))
                    C = Field(Average(c; dims))

                    @test !isnothing(C.status)

                    # Test conditional computation
                    set!(c, 1)
                    compute_at!(C, FT(1)) # will compute
                    @test all(interior(C) .== 1)
                    @test C.status.time == FT(1)

                    set!(c, 2)
                    compute_at!(C, FT(1)) # will not compute because status == 1
                    @test C.status.time == FT(1)
                    @test all(interior(C) .== 1)

                    compute_at!(C, FT(2))
                    @test C.status.time == FT(2)
                    @test all(interior(C) .== 2)
                end
            end
        end

        @testset "Immersed Fields reduction [$(typeof(arch))]" begin
            @info "  Testing reductions of immersed Fields [$(typeof(arch))]"
            underlying_grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(1, 1, 1))

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> y < 0.5 ? - 0.6 : 0))
            c = Field{Center, Center, Nothing}(grid)

            set!(c, (x, y) -> y)
            @test maximum(c) == grid.yᵃᶜᵃ[1]

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> y < 0.5 ? - 0.6 : -0.4))
            c = Field{Center, Center, Nothing}(grid)

            set!(c, (x, y) -> y)
            @test maximum(c) == grid.yᵃᶜᵃ[3]

            underlying_grid = RectilinearGrid(arch, size = (1, 1, 8), extent=(1, 1, 1))

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> -3/4))
            c = Field{Center, Center, Center}(grid)

            set!(c, (x, y, z) -> -z)
            @test maximum(c) == Array(interior(c))[1, 1, 3]
            c_condition = interior(c) .< 0.5
            avg_c_smaller_than_½ = Array(interior(compute!(Field(Average(c, condition=c_condition)))))
            @test avg_c_smaller_than_½[1, 1, 1] == 0.25

            zᶜᶜᶜ = KernelFunctionOperation{Center, Center, Center}(znode, grid, Center(), Center(), Center())
            ci = Array(interior(c)) # transfer to CPU
            bottom_half_average_manual = (ci[1, 1, 3] + ci[1, 1, 4]) / 2
            bottom_half_average = Average(c; condition=(zᶜᶜᶜ .< -1/2))
            bottom_half_average_field = Field(bottom_half_average)
            compute!(bottom_half_average_field)
            bottom_half_average_array = Array(interior(bottom_half_average_field))
            @test bottom_half_average_array[1, 1, 1] == bottom_half_average_manual

            # See: https://github.com/CliMA/Oceananigans.jl/issues/3948
            underlying_grid = RectilinearGrid(
                topology=(Periodic, Periodic, Periodic),
                size=(3, 3, 3),
                x=(0, 1), y=(0, 1), z=(0, 1)
            )

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> x + y))

            c = CenterField(grid)
            set!(c, (x, y, z) -> x + y + z)

            max_c² = Field(Reduction(maximum!, c^2, dims=3))
            ∫max_c² = Integral(max_c², dims=(1, 2))
            @test ∫max_c² isa Reduction
            ∫max_c²_field = Field(∫max_c²)
            @test ∫max_c²_field isa Field
            @test ∫max_c²_field.operand === ∫max_c²

            @info "  Testing conditional reductions of immersed Fields [$(typeof(arch))]"

            underlying_grid = LatitudeLongitudeGrid(arch,
                                                    topology = (Periodic, Bounded, Bounded),
                                                    size = (24, 16, 8),
                                                    longitude = (-10, 10),
                                                    latitude = (-55, -35),
                                                    z = (-1000, 0),
                                                    halo = (5, 5, 5))

            Lz_u = underlying_grid.Lz
            width = 0.5 # degrees
            bump(λ, φ) = - Lz_u * (1 - 2 * exp(-(λ^2 + φ^2) / 2width^2))

            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bump))

            u_ibg = XFaceField(grid)
            u_noibg = XFaceField(grid.underlying_grid)
            v_ibg = YFaceField(grid)
            v_noibg = YFaceField(grid.underlying_grid)
            w_ibg = ZFaceField(grid)
            w_noibg = ZFaceField(grid.underlying_grid)

            condition = trues(size(grid)) # should work for Periodic but not for Bounded directions

            ∫u_ibg = Integral(u_ibg; dims=1, condition)
            @test ∫u_ibg isa Reduction{<:Any, <:Any, <:ConditionalOperation}
            ∫u_noibg = Integral(u_noibg; dims=1, condition)
            @test ∫u_noibg isa Reduction{<:Any, <:Any, <:ConditionalOperation}
            @test_throws ArgumentError ∫v_ibg = Integral(v_ibg; dims=1, condition)
            @test_throws ArgumentError ∫v_noibg = Integral(v_noibg; dims=1, condition)
            @test_throws ArgumentError ∫w_ibg = Integral(w_ibg; dims=1, condition)
            @test_throws ArgumentError ∫w_noibg = Integral(w_noibg; dims=1, condition)
            ∫v_ibg = Integral(v_ibg; dims=1, condition=trues(size(v_ibg)))
            @test ∫v_ibg isa Reduction{<:Any, <:Any, <:ConditionalOperation}
            ∫v_noibg = Integral(v_noibg; dims=1, condition=trues(size(v_noibg)))
            @test ∫v_noibg isa Reduction{<:Any, <:Any, <:ConditionalOperation}
            ∫w_ibg = Integral(w_ibg; dims=1, condition=trues(size(w_ibg)))
            @test ∫w_ibg isa Reduction{<:Any, <:Any, <:ConditionalOperation}
            ∫w_noibg = Integral(w_noibg; dims=1, condition=trues(size(w_noibg)))
            @test ∫w_noibg isa Reduction{<:Any, <:Any, <:ConditionalOperation}
        end
    end
end
