include("reactant_test_utils.jl")

using CUDA
using Oceananigans.OutputReaders: cpu_interpolating_time_indices

arch = ReactantState()

ridge(λ, φ) = 0.1 * exp((λ - 2)^2 / 2)

@testset "Reactant unit tests" begin
    @testset "Checkbounds = auto" begin
        @test Base.JLOptions().check_bounds == 0
    end

    @info "Performing Reactanigans unit tests..."

    times = 0:1.0:4
    t = 2.1
    times = Reactant.to_rarray(times, track_numbers=Number)
    @test times isa Reactant.TracedStepRangeLen

    ñ, n₁, n₂ = @jit Oceananigans.OutputReaders.find_time_index(times, t)
    @test ñ ≈ 0.1
    @test n₁ == 3 # eg times = [0 1 2 ⟨⟨2.1⟩⟩ 3]
    @test n₂ == 4

    grid = RectilinearGrid(arch; size=(4, 4, 4), extent=(1, 1, 1))
    c = CenterField(grid)
    @test parent(c) isa Reactant.ConcreteRArray

    cpu_grid = on_architecture(CPU(), grid)
    @test architecture(cpu_grid) isa CPU

    cpu_c = on_architecture(CPU(), c)
    @test parent(cpu_c) isa Array
    @test architecture(cpu_c.grid) isa CPU

    @info "  Testing field set! with a number..."
    set!(c, 1)
    @test all(c .≈ 1)

    @info "  Testing field set! with a function..."
    set!(c, (x, y, z) -> 1)
    @test all(c .≈ 1)

    @info "  Testing field set! with an array..."
    a = ones(size(c)...)
    set!(c, a)
    @test all(c .≈ 1)

    @info "  Testing simple kernel launch!..."
    add_one!(c)
    @test all(c .≈ 2)

    @info "  Testing set!..."
    set!(c, (x, y, z) -> x + y * z)
    x, y, z = nodes(c)

    @allowscalar begin
        @test c[1, 1, 1] == x[1] + y[1] * z[1]
        @test c[1, 2, 1] == x[1] + y[2] * z[1]
        @test c[1, 2, 3] == x[1] + y[2] * z[3]
    end

    @info "  Testing fill halo regions!..."
    fill_halo_regions!(c)

    @allowscalar begin
        @test c[1, 1, 0] == c[1, 1, 1]
    end

    d = CenterField(grid)
    parent(d) .= 2

    @info "  Testing computed field..."
    cd = Field(c * d)
    compute!(cd)

    @allowscalar begin
        @test cd[1, 1, 1] == 2 * (x[1] + y[1] * z[1])
        @test cd[1, 2, 1] == 2 * (x[1] + y[2] * z[1])
        @test cd[1, 2, 3] == 2 * (x[1] + y[2] * z[3])
    end

    # Deconcretization / nondeconcretization
    @info "  Testing Field deconcretization..."
    c′ = OceananigansReactantExt.deconcretize(c)
    @test parent(c′) isa Array
    @test architecture(c′) isa ReactantState

    for FT in (Float64, Float32)
        @info "  Testing RectilinearGrid construction [$FT]..."
        sgrid = RectilinearGrid(arch, FT; size=(4, 4, 4), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 1))
        @test architecture(sgrid) isa ReactantState
        @test architecture(sgrid.xᶠᵃᵃ) isa ReactantState
        @test architecture(sgrid.xᶜᵃᵃ) isa ReactantState

        @info "  Testing LatitudeLongitudeGrid construction [$FT]..."
        llg = LatitudeLongitudeGrid(arch, FT; size = (4, 4, 4),
                                    longitude = [0, 1, 2, 3, 4],
                                    latitude = [0, 1, 2, 3, 4],
                                    z = (0, 1))

        @info "  Testing RectilinearGrid eltype..."
        RFT = @jit eltype(sgrid)
        @test RFT == Reactant.TracedRNumber{FT}

        @info "  Testing LatitudeLongitudeGrid eltype..."
        RFT = @jit eltype(llg)
        @test RFT == Reactant.TracedRNumber{FT}

        @test architecture(llg) isa ReactantState

        for name in propertynames(llg)
            p = getproperty(llg, name)
            if !(name ∈ (:architecture, :z))
                @test (p isa Number) || (p isa OffsetArray{FT, <:Any, <:Reactant.AbstractConcreteArray})
            end
        end

        @info "  Testing constantified LatitudeLongitudeGrid construction [$FT]..."
        cpu_llg = LatitudeLongitudeGrid(CPU(), FT; size = (4, 4, 4),
                                        longitude = [0, 1, 2, 3, 4],
                                        latitude = [0, 1, 2, 3, 4],
                                        z = (0, 1))

        constant_llg = OceananigansReactantExt.constant_with_arch(cpu_llg, ReactantState())

        for name in propertynames(constant_llg)
            p = getproperty(constant_llg, name)
            if !(name ∈ (:architecture, :z))
                @test (p isa Number) || (p isa OffsetArray{FT, <:Any, <:Array})
            end
        end

        @info "  Testing ImmersedBoundaryGrid construction [$FT]..."
        ibg = ImmersedBoundaryGrid(llg, GridFittedBottom(ridge))
        @test architecture(ibg) isa ReactantState
        @test architecture(ibg.immersed_boundary.bottom_height) isa ReactantState

        @info "  Testing constantified ImmersedBoundaryGrid construction [$FT]..."
        cpu_ibg = ImmersedBoundaryGrid(cpu_llg, GridFittedBottom(ridge))
        constant_ibg = OceananigansReactantExt.constant_with_arch(cpu_ibg, ReactantState())
        @test architecture(constant_ibg) isa ReactantState
        @test architecture(constant_ibg.immersed_boundary.bottom_height) isa CPU

        @info "  Testing constantified OrthogonalSphericalShellGrid construction [$FT]..."
        rllg = RotatedLatitudeLongitudeGrid(arch, FT; size = (4, 4, 4),
                                            north_pole = (0, 0),
                                            longitude = [0, 1, 2, 3, 4],
                                            latitude = [0, 1, 2, 3, 4],
                                            z = (0, 1))

        @test architecture(rllg) isa ReactantState

        for name in propertynames(rllg)
            p = getproperty(rllg, name)
            if !(name ∈ (:architecture, :z, :conformal_mapping))
                @test (p isa Number) || (p isa OffsetArray{FT, <:Any, <:Reactant.AbstractConcreteArray})
            end
        end

        @info "  Testing constantified OrthogonalSphericalShellGrid construction [$FT]..."
        @info "    Building CPU grid [$FT]..."
        cpu_rllg = RotatedLatitudeLongitudeGrid(CPU(), FT; size = (4, 4, 4),
                                                north_pole = (0, 0),
                                                longitude = [0, 1, 2, 3, 4],
                                                latitude = [0, 1, 2, 3, 4],
                                                z = (0, 1))

        @info "    Replacing architecture with ReactantState [$FT]..."
        constant_rllg = OceananigansReactantExt.constant_with_arch(cpu_rllg, ReactantState())

        for name in propertynames(constant_rllg)
            p = getproperty(constant_rllg, name)
            if !(name ∈ (:architecture, :z, :conformal_mapping))
                @test (p isa Number) || (p isa OffsetArray{FT, <:Any, <:Array})
            end
        end

        @info "  Testing constantified immersed OrthogonalSphericalShellGrid construction [$FT]..."
        cpu_ribg = ImmersedBoundaryGrid(cpu_rllg, GridFittedBottom(ridge))
        constant_ribg = OceananigansReactantExt.constant_with_arch(cpu_ribg, ReactantState())
        @test architecture(constant_ribg) isa ReactantState
        @test architecture(constant_ribg.immersed_boundary.bottom_height) isa CPU
    end

    @testset "KernelFunctionOperation reductions on LatitudeLongitudeGrid" begin
        grid = LatitudeLongitudeGrid(arch;
                                     size = (40, 40, 10),
                                     longitude = (0, 360),
                                     latitude = (-10, 10),
                                     z = (-1000, 0),
                                     halo = (5, 5, 5))

        returnone(i, j, k, grid) = 1
        kfo = KernelFunctionOperation{Center, Center, Center}(returnone, grid)

        @test minimum(kfo) == 1
        @test maximum(kfo) == 1

        @test minimum_xspacing(grid) > 0
        @test minimum_yspacing(grid) > 0
        @test minimum_zspacing(grid) > 0

        @test minimum_xspacing(grid) ≈ minimum(xspacings(grid))
        @test minimum_yspacing(grid) ≈ minimum(yspacings(grid))
        @test minimum_zspacing(grid) ≈ minimum(zspacings(grid))
    end

    @testset "KernelFunctionOperation reductions on RectilinearGrid" begin
        grid = RectilinearGrid(arch;
                               size = (10, 10, 10),
                               extent = (1, 1, 1))

        returnone(i, j, k, grid) = 1
        kfo = KernelFunctionOperation{Center, Center, Center}(returnone, grid)

        @test minimum(kfo) == 1
        @test maximum(kfo) == 1

        @test minimum_xspacing(grid) > 0
        @test minimum_yspacing(grid) > 0
        @test minimum_zspacing(grid) > 0

        @test minimum_xspacing(grid) ≈ minimum(xspacings(grid))
        @test minimum_yspacing(grid) ≈ minimum(yspacings(grid))
        @test minimum_zspacing(grid) ≈ minimum(zspacings(grid))
    end

    @testset "Field reductions on LatitudeLongitudeGrid" begin
        grid = LatitudeLongitudeGrid(arch;
                                     size = (10, 10, 10),
                                     longitude = (0, 360),
                                     latitude = (-60, 60),
                                     z = (-1000, 0),
                                     halo = (5, 5, 5))

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

        @test minimum(c) ≈ -950.0
        @test maximum(c) ≈ -50.0
    end

    @testset "AbstractOperation on LatitudeLongitudeGrid under trace" begin
        grid = LatitudeLongitudeGrid(arch;
                                     size = (4, 4, 4),
                                     longitude = (0, 360),
                                     latitude = (-60, 60),
                                     z = (0, 1))

        # The Reactant LatitudeLongitudeGrid `==` must return a plain `Bool`. The
        # generic grid `==` compares node values, which under Reactant trace to a
        # `TracedRNumber{Bool}` that cannot drive the `||` in validate_grid.
        @test (grid == grid) isa Bool
        @test grid == grid

        value_c = 2
        c  = CenterField(grid); set!(c, value_c)
        value_d = 3
        d  = CenterField(grid); set!(d, value_d)
        cd = CenterField(grid)

        # `c * d` builds a BinaryOperation -> validate_grid(c, d) -> (c.grid == d.grid).
        # Compiling the assignment exercises the Reactant grid `==` while materializing
        # into the preallocated `cd` (so we avoid the reduction allocation path).
        multiply!(cd, c, d) = (cd .= c * d; nothing)
        compiled_multiply! = @compile sync=true multiply!(cd, c, d)
        compiled_multiply!(cd, c, d)
        @test all(≈(value_c * value_d), Array(interior(cd)))
    end

    @testset "set! from a reduced field under trace" begin
        grid = LatitudeLongitudeGrid(arch;
                                     size = (4, 4, 4),
                                     longitude = (0, 360),
                                     latitude = (-60, 60),
                                     z = (0, 1))

        # A reduced column field (as used for reference states) must broadcast into
        # a full 3D field inside a compiled function: the interpolation fallback in
        # `set_to_field!` hops the grid to the CPU with `on_architecture(CPU(), grid)`,
        # which cannot work on a traced grid (its validating reconstruction drives
        # `||` with a `TracedRNumber{Bool}`, and traced data cannot be materialized
        # on the CPU mid-trace).
        u = CenterField(grid)
        v = Field{Nothing, Nothing, Center}(grid)
        column = reshape([1.0, 2.0, 3.0, 4.0], 1, 1, 4)
        set!(v, column)

        set_column!(u, v) = (set!(u, v); nothing)
        compiled_set_column! = @compile sync=true set_column!(u, v)
        compiled_set_column!(u, v)

        @test Array(interior(u)) == repeat(column, outer=(4, 4, 1))
    end

    @testset "Field reductions on RectilinearGrid" begin
        grid = RectilinearGrid(arch;
                               size = (10, 10, 10),
                               extent = (1, 1, 1))

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

        @test minimum(c) ≈ -0.95
        @test maximum(c) ≈ -0.05
    end
end

@testset "Reactant FieldTimeSeries tests" begin
    @info "Testing the use of a `FieldTimeSeries` on a `ReactantState` arch..."

    arch = ReactantState()
    Nx, Ny, Nz = (10, 10, 10) # number of cells
    grid = RectilinearGrid(arch; size=(Nx, Ny, Nz), extent=(1, 1, 1))
    fts  = FieldTimeSeries{Center, Center, Center}(grid, sort(rand(10)))

    # Test I can index into a Reactant FieldTimeSeries
    @test fts[5] isa Field
end

@testset "Field materialize traced" begin
    Nx, Nz = (128, 64)
    Lx, Lz = (1_000, 200)
    grid_ad = RectilinearGrid(ReactantState(); size = (Nx, Nz),
                              x = (-Lx/2, Lx/2), z = (0, Lz),
                              topology = (Periodic, Flat, Bounded))
    δρᵢ  = CenterField(grid_ad)
    @test maximum(δρᵢ) == maximum(Array(interior(δρᵢ)))
end
