include("reactant_test_utils.jl")
include("reactant_correctness_utils.jl")

using Random
using CUDA
using Oceananigans.OrthogonalSphericalShellGrids: TripolarGrid, OrthogonalSphericalShellGrid
using SeawaterPolynomials: TEOS10EquationOfState

# Helper to generate all combinations
all_combos(xs...) = vec(collect(Iterators.product(xs...)))

@testset "Reactant correctness" begin
    @info "Testing Reactant correctness (comparing vanilla Oceananigans vs ReactantState)..."

    # Get vanilla architecture from TEST_ARCHITECTURE env var (set by reactant_test_utils.jl)
    vanilla_arch = get(ENV, "TEST_ARCHITECTURE", "CPU") == "GPU" ? GPU() : CPU()
    reactant_arch = ReactantState()

    # Field locations to test
    all_locations = all_combos((Center, Face), (Center, Face), (Center, Face))

    fill_halo_topologies = (
            (Periodic, Periodic, Bounded),
            (Periodic, Bounded, Bounded),
            (Bounded, Bounded, Bounded),
            # could also test these very uncommonly-used topologies:
            #(Bounded, Periodic, Bounded),
            #(Bounded, Bounded, Periodic),
            #(Periodic, Bounded, Periodic),
            #(Bounded, Periodic, Periodic),
            #(Periodic, Periodic, Periodic),
    )

    #####
    ##### Halo filling tests
    #####

    @testset "fill_halo_regions! correctness" begin

        # Test RectilinearGrid with all topologies
        for topo in fill_halo_topologies
            @testset "RectilinearGrid topology=$topo" begin
                @info "Testing fill_halo_regions! correctness on RectilinearGrid with topology=$topo..."
                kw = (size=(5, 6, 4), halo=(3, 3, 3), extent=(1, 1, 1), topology=topo)
                vanilla_grid = RectilinearGrid(vanilla_arch; kw...)
                reactant_grid = RectilinearGrid(reactant_arch; kw...)

                for loc in all_locations
                    LX, LY, LZ = loc
                    @testset "loc=$loc" begin
                        vanilla_field = Field{LX, LY, LZ}(vanilla_grid)
                        reactant_field = Field{LX, LY, LZ}(reactant_grid)

                        Random.seed!(12345)
                        data = randn(size(vanilla_field)...)
                        set!(vanilla_field, data)
                        set!(reactant_field, data)

                        fill_halo_regions!(vanilla_field)
                        rfill! = @compile raise=true raise_first=true fill_halo_regions!(reactant_field)
                        rfill!(reactant_field)

                        @test compare_parent("halo", vanilla_field, reactant_field)
                    end
                end
            end
        end

        # Test LatitudeLongitudeGrid with LX ∈ (Periodic, Bounded)
        for LX in (Periodic, Bounded)
            topo = (LX, Bounded, Bounded)
            @testset "LatitudeLongitudeGrid topology=$topo" begin
                @info "Testing fill_halo_regions! correctness on LatitudeLongitudeGrid with topology=$topo..."
                kw = (size=(6, 6, 4), halo=(3, 3, 3), longitude=(0, 10), latitude=(0, 10), z=(0, 1), topology=topo)
                vanilla_grid = LatitudeLongitudeGrid(vanilla_arch; kw...)
                reactant_grid = LatitudeLongitudeGrid(reactant_arch; kw...)

                for loc in all_locations
                    LXf, LYf, LZf = loc
                    @testset "loc=$loc" begin
                        vanilla_field = Field{LXf, LYf, LZf}(vanilla_grid)
                        reactant_field = Field{LXf, LYf, LZf}(reactant_grid)

                        Random.seed!(12345)
                        data = randn(size(vanilla_field)...)
                        set!(vanilla_field, data)
                        set!(reactant_field, data)

                        fill_halo_regions!(vanilla_field)
                        rfill! = @compile raise=true raise_first=true fill_halo_regions!(reactant_field)
                        rfill!(reactant_field)

                        @test compare_parent("halo", vanilla_field, reactant_field)
                    end
                end
            end
        end

        # TripolarGrid tests commented out: Reactant MLIR segfaults
        # in CanonicalizeLoopsPass when compiling Zipper boundary conditions
        # on TripolarGrid. This crash is non-deterministic and kills the
        # entire test process, preventing subsequent tests from running.
        @testset "TripolarGrid" begin
            @info "Testing fill_halo_regions! correctness on TripolarGrid..."
            kw = (size=(8, 10, 4), halo=(3, 3, 3), z=(0, 1))
            vanilla_grid = TripolarGrid(vanilla_arch; kw...)
            reactant_grid = TripolarGrid(reactant_arch; kw...)

            for loc in all_locations
                LX, LY, LZ = loc
                @testset "loc=$loc" begin
                    vanilla_field = Field{LX, LY, LZ}(vanilla_grid)
                    reactant_field = Field{LX, LY, LZ}(reactant_grid)

                    Random.seed!(12345)
                    data = randn(size(vanilla_field)...)
                    set!(vanilla_field, data)
                    set!(reactant_field, data)

                    fill_halo_regions!(vanilla_field)
                    @jit fill_halo_regions!(reactant_field)

                    @test compare_parent("halo", vanilla_field, reactant_field)
                end
            end
        end
    end

    #####
    ##### Tendency computation tests
    #####

    @testset "compute_simple_Gu! correctness" begin
        advection_schemes = (nothing, Centered(), WENO())

        # Helper to get advection name for testset
        adv_name(::Nothing) = "nothing"
        adv_name(a) = string(nameof(typeof(a)))

        # Topologies to test
        # Note: topologies with Periodic directions are commented out because
        # Reactant's MLIR raise pass fails on the periodic halo-filling kernel
        # with halo ≥ 2 ("cannot raise if yet" error). See MWE in
        # test/reactant_raise_periodic_halo_mwe.jl
        gu_topologies = (
            (Periodic, Periodic, Bounded),
            (Periodic, Bounded, Bounded),
            (Bounded, Bounded, Bounded),
            #(Bounded, Periodic, Bounded),
        )

        # RectilinearGrid tests
        @testset "RectilinearGrid" begin
            for topo in gu_topologies
                @testset "topology=$topo" begin
                    kw = (size=(4, 4, 4), halo=(3, 3, 3), extent=(1, 1, 1), topology=topo)
                    vanilla_grid = RectilinearGrid(vanilla_arch; kw...)
                    reactant_grid = RectilinearGrid(reactant_arch; kw...)

                    # Set up velocity fields
                    vanilla_velocities = (u=XFaceField(vanilla_grid), v=YFaceField(vanilla_grid), w=ZFaceField(vanilla_grid))
                    reactant_velocities = (u=XFaceField(reactant_grid), v=YFaceField(reactant_grid), w=ZFaceField(reactant_grid))

                    Random.seed!(54321)
                    for (vf, rf) in zip(vanilla_velocities, reactant_velocities)
                        data = randn(size(vf)...)
                        set!(vf, data)
                        set!(rf, data)
                        fill_halo_regions!(vf)
                        rfill! = @compile raise=true raise_first=true fill_halo_regions!(rf)
                        rfill!(rf)
                    end

                    vanilla_Gu = XFaceField(vanilla_grid)
                    reactant_Gu = XFaceField(reactant_grid)

                    # Test with coriolis=nothing
                    @testset "coriolis=nothing" begin
                        @info "Testing compute_simple_Gu! on RectilinearGrid($topo) with coriolis=nothing..."
                        for advection in advection_schemes
                            # We want to select differed setting based on reactant
                            advection = Oceananigans.Advection.materialize_advection(advection, reactant_grid)
                            @testset "advection=$(adv_name(advection))" begin
                                fill!(vanilla_Gu, 0)
                                fill!(reactant_Gu, 0)
                                compute_simple_Gu!(vanilla_Gu, advection, nothing, vanilla_velocities)
                                # Temporary workaround: raise=true is currently broken with WENO, Bounded x-axis topology here.
                                rcompute_Gu! = if advection isa WENO && topo[1] == Bounded
                                    @compile compute_simple_Gu!(reactant_Gu, advection, nothing, reactant_velocities)
                                else
                                    @compile raise=true raise_first=true compute_simple_Gu!(reactant_Gu, advection, nothing, reactant_velocities)
                                end
                                rcompute_Gu!(reactant_Gu, advection, nothing, reactant_velocities)
                                @test compare_interior("Gu", vanilla_Gu, reactant_Gu)
                            end
                        end
                    end

                    # Test with FPlane Coriolis
                    @testset "coriolis=FPlane" begin
                        @info "Testing compute_simple_Gu! on RectilinearGrid($topo) with FPlane Coriolis..."
                        coriolis = FPlane(f=1e-4)
                        for advection in advection_schemes
                            # We want to select differed setting based on reactant
                            advection = Oceananigans.Advection.materialize_advection(advection, reactant_grid)
                            @testset "advection=$(adv_name(advection))" begin
                                fill!(vanilla_Gu, 0)
                                fill!(reactant_Gu, 0)
                                compute_simple_Gu!(vanilla_Gu, advection, coriolis, vanilla_velocities)
                                # Temporary workaround: raise=true is currently broken with WENO, Bounded x-axis topology here.
                                rcompute_Gu! = if advection isa WENO && topo[1] == Bounded
                                    @compile compute_simple_Gu!(reactant_Gu, advection, coriolis, reactant_velocities)
                                else
                                    @compile raise=true raise_first=true compute_simple_Gu!(reactant_Gu, advection, coriolis, reactant_velocities)
                                end
                                rcompute_Gu!(reactant_Gu, advection, coriolis, reactant_velocities)
                                @test compare_interior("Gu", vanilla_Gu, reactant_Gu)
                            end
                        end
                    end
                end
            end
        end

        # LatitudeLongitudeGrid tests
        for TX in (Periodic, Bounded)
            @testset "LatitudeLongitudeGrid (TX=$TX)" begin
                topo = (TX, Bounded, Bounded)
                longitude = TX == Periodic ? (0, 360) : (0, 90)
                kw = (size=(4, 4, 4), halo=(3, 3, 3), longitude=longitude, latitude=(0, 10), z=(0, 1), topology=topo)
                vanilla_grid = LatitudeLongitudeGrid(vanilla_arch; kw...)
                reactant_grid = LatitudeLongitudeGrid(reactant_arch; kw...)

                # Set up velocity fields
                vanilla_velocities = (u=XFaceField(vanilla_grid), v=YFaceField(vanilla_grid), w=ZFaceField(vanilla_grid))
                reactant_velocities = (u=XFaceField(reactant_grid), v=YFaceField(reactant_grid), w=ZFaceField(reactant_grid))

                Random.seed!(54321)
                for (vf, rf) in zip(vanilla_velocities, reactant_velocities)
                    data = randn(size(vf)...)
                    set!(vf, data)
                    set!(rf, data)
                    fill_halo_regions!(vf)
                    rfill! = @compile raise=true raise_first=true fill_halo_regions!(rf)
                    rfill!(rf)
                end

                vanilla_Gu = XFaceField(vanilla_grid)
                reactant_Gu = XFaceField(reactant_grid)

                # Test with coriolis=nothing (advection-only)
                @testset "coriolis=nothing" begin
                    @info "Testing compute_simple_Gu! on LatitudeLongitudeGrid with coriolis=nothing..."
                    for advection in advection_schemes
                        # We want to select differed setting based on reactant
                        advection = Oceananigans.Advection.materialize_advection(advection, reactant_grid)
                        @testset "advection=$(adv_name(advection))" begin
                            fill!(vanilla_Gu, 0)
                            fill!(reactant_Gu, 0)
                            compute_simple_Gu!(vanilla_Gu, advection, nothing, vanilla_velocities)
                            # Temporary workaround: raise=true is currently broken with WENO here.
                            rcompute_Gu! = if advection isa WENO && TX == Bounded
                                @compile compute_simple_Gu!(reactant_Gu, advection, nothing, reactant_velocities)
                            else
                                @compile raise=true raise_first=true compute_simple_Gu!(reactant_Gu, advection, nothing, reactant_velocities)
                            end
                            rcompute_Gu!(reactant_Gu, advection, nothing, reactant_velocities)
                            @test compare_interior("Gu", vanilla_Gu, reactant_Gu)
                        end
                    end
                end

                # Test with HydrostaticSphericalCoriolis
                @testset "coriolis=HydrostaticSphericalCoriolis" begin
                    @info "Testing compute_simple_Gu! on LatitudeLongitudeGrid with HydrostaticSphericalCoriolis..."
                    coriolis = HydrostaticSphericalCoriolis()
                    for advection in advection_schemes
                        # We want to select differed setting based on reactant
                        advection = Oceananigans.Advection.materialize_advection(advection, reactant_grid)
                        @testset "advection=$(adv_name(advection))" begin
                            fill!(vanilla_Gu, 0)
                            fill!(reactant_Gu, 0)
                            compute_simple_Gu!(vanilla_Gu, advection, coriolis, vanilla_velocities)
                            # Temporary workaround: raise=true is currently broken with WENO here.
                            rcompute_Gu! = if advection isa WENO && TX == Bounded
                                @compile compute_simple_Gu!(reactant_Gu, advection, coriolis, reactant_velocities)
                            else
                                @compile raise=true raise_first=true compute_simple_Gu!(reactant_Gu, advection, coriolis, reactant_velocities)
                            end
                            rcompute_Gu!(reactant_Gu, advection, coriolis, reactant_velocities)
                            @test compare_interior("Gu", vanilla_Gu, reactant_Gu)
                        end
                    end
                end
            end
        end
    end

    # HydrostaticFreeSurfaceModel time-stepping tests commented out:
    # Reactant compilation fails with "'llvm.call' op incorrect number of operands"
    # during @jit update_state!(model). This affects all topologies and is a
    # Reactant MLIR code generation bug.
    @testset "HydrostaticFreeSurfaceModel time-stepping" begin

        # Topologies to test
        hfsm_topologies = (
            (Periodic, Periodic, Bounded),
            (Periodic, Bounded, Bounded),
            # (Bounded, Periodic, Bounded),
            # (Bounded, Bounded, Bounded),
        )

        # Grid setup
        # Note: WENOVectorInvariant requires halo >= (6, 6, 3)
        Nx, Ny, Nz = 8, 8, 4
        Lx, Ly, Lz = 1e5, 1e5, 100  # 100km x 100km x 100m domain

        # Common model parameters
        coriolis = FPlane(f=1e-4)
        advection = WENO()
        momentum_advection = WENOVectorInvariant()
        equation_of_state = TEOS10EquationOfState()
        buoyancy = SeawaterBuoyancy(; equation_of_state)
        tracers = (:T, :S)

        # Closures to test
        # Note: CATKEVerticalDiffusivity takes too long to compile with Reactant (>10 min)
        closures = (nothing,)

        closure_name(::Nothing) = "nothing"
        closure_name(c) = string(nameof(typeof(c)))

        for topo in hfsm_topologies
            @testset "topology=$topo" begin
                kw = (size=(Nx, Ny, Nz), halo=(6, 6, 3), x=(0, Lx), y=(0, Ly), z=(-Lz, 0), topology=topo)
                vanilla_grid = RectilinearGrid(vanilla_arch; kw...)
                reactant_grid = RectilinearGrid(reactant_arch; kw...)

                # Use SplitExplicitFreeSurface (FFT-based ImplicitFreeSurface doesn't work with Reactant)
                free_surface = SplitExplicitFreeSurface(vanilla_grid; substeps = 10)

                for closure in closures
                    @testset "closure=$(closure_name(closure))" begin
                        @info "Testing HydrostaticFreeSurfaceModel($topo) with closure=$(closure_name(closure))..."

                        model_kw = (; coriolis, buoyancy, free_surface, tracers, tracer_advection=advection, momentum_advection, closure)
                        vanilla_model = HydrostaticFreeSurfaceModel(vanilla_grid; model_kw...)
                        reactant_model = HydrostaticFreeSurfaceModel(reactant_grid; model_kw...)

                        # Set initial conditions
                        # Note: Use pre-computed arrays, NOT functions with randn(),
                        # because randn() would be called at different times for each model
                        Random.seed!(98765)

                        # Small velocity perturbations
                        u_init = 0.1 * randn(size(vanilla_model.velocities.u)...)
                        v_init = 0.1 * randn(size(vanilla_model.velocities.v)...)

                        # Realistic T/S with vertical gradient and small perturbations
                        T_init = [20.0 + 5.0 * (k - 0.5) / Nz + 0.01 * randn() for i=1:Nx, j=1:Ny, k=1:Nz]
                        S_init = [35.0 + 0.01 * randn() for i=1:Nx, j=1:Ny, k=1:Nz]

                        set!(vanilla_model, u=u_init, v=v_init, T=T_init, S=S_init)
                        set!(reactant_model, u=u_init, v=v_init, T=T_init, S=S_init)

                        # Small time-step for stability
                        Δt = 60.0  # 1 minute
                        stop_iteration = 3

                        # Run vanilla simulation
                        vanilla_simulation = Simulation(vanilla_model; Δt, stop_iteration, verbose=false)
                        run!(vanilla_simulation)

                        # Compile and run Reactant simulation
                        @jit Oceananigans.TimeSteppers.update_state!(reactant_model)

                        r_first_time_step! = @compile sync=true Oceananigans.TimeSteppers.first_time_step!(reactant_model, Δt)
                        r_time_step! = @compile sync=true Oceananigans.TimeSteppers.time_step!(reactant_model, Δt)

                        reactant_simulation = Simulation(reactant_model; Δt, stop_iteration, verbose=false)
                        r_run!(reactant_simulation, r_time_step!, r_first_time_step!)

                        # Compare results
                        @testset "velocity fields" begin
                            @test compare_interior("u", vanilla_model.velocities.u, reactant_model.velocities.u)
                            @test compare_interior("v", vanilla_model.velocities.v, reactant_model.velocities.v)
                        end

                        @testset "tracer fields" begin
                            @test compare_interior("T", vanilla_model.tracers.T, reactant_model.tracers.T)
                            @test compare_interior("S", vanilla_model.tracers.S, reactant_model.tracers.S)
                        end

                        @testset "simulation state" begin
                            @test iteration(vanilla_simulation) == iteration(reactant_simulation)
                            @test time(vanilla_simulation) ≈ time(reactant_simulation)
                        end
                    end
                end
            end
        end
    end
end
