include("dependencies_for_runtests.jl")

using Oceananigans.Diagnostics

using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant

TestModel_VerticallyStrectedRectGrid(arch, FT, ν=1.0, Δx=0.5) =
    NonhydrostaticModel(RectilinearGrid(arch, FT, size=(3, 3, 3), x=(0, 3Δx), y=(0, 3Δx), z=0:Δx:3Δx,);
                        closure = ScalarDiffusivity(FT, ν=ν, κ=ν))


TestModel_RegularRectGrid(arch, FT, ν=1.0, Δx=0.5) =
    NonhydrostaticModel(RectilinearGrid(arch, FT, topology=(Periodic, Periodic, Periodic), size=(3, 3, 3), extent=(3Δx, 3Δx, 3Δx));
                        closure = ScalarDiffusivity(FT, ν=ν, κ=ν))

function diffusive_cfl_diagnostic_is_correct(arch, FT)
    Δt = FT(1.3e-6)
    Δx = FT(0.5)
    ν = FT(1.2)
    CFL_by_hand = Δt * ν / Δx^2

    model = TestModel_RegularRectGrid(arch, FT, ν, Δx)
    cfl = DiffusiveCFL(FT(Δt))

    return cfl(model) ≈ CFL_by_hand
end

function advective_cfl_diagnostic_is_correct_on_regular_grid(arch, FT)
    model = TestModel_RegularRectGrid(arch, FT)

    Δt = FT(1.3e-6)
    Δx = minimum_xspacing(model.grid)
    u₀ = FT(1.2)
    CFL_by_hand = Δt * u₀ / Δx

    set!(model, u=u₀)
    cfl = AdvectiveCFL(FT(Δt))

    return cfl(model) ≈ CFL_by_hand
end

function advective_cfl_diagnostic_is_correct_on_vertically_stretched_grid(arch, FT)
    model = TestModel_VerticallyStrectedRectGrid(arch, FT)

    Δt = FT(1.3e-6)
    Δx = FT(model.grid.Δxᶜᵃᵃ)
    u₀ = FT(1.2)
    CFL_by_hand = Δt * u₀ / Δx

    set!(model, u=u₀)
    cfl = AdvectiveCFL(FT(Δt))

    return cfl(model) ≈ CFL_by_hand
end

function advective_timescale_cfl_on_regular_grid(arch, FT)
    model = TestModel_RegularRectGrid(arch, FT)

    Δt = FT(1.7)

    Δx = model.grid.Δxᶜᵃᵃ
    Δy = model.grid.Δyᵃᶜᵃ
    Δz = model.grid.z.Δᵃᵃᶜ

    u₀ = FT(1.2)
    v₀ = FT(-2.5)
    w₀ = FT(3.9)

    set!(model, u=u₀, v=v₀, w=w₀)

    CFL_by_hand = Δt * (abs(u₀) / Δx + abs(v₀) / Δy + abs(w₀) / Δz)

    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)

    return cfl(model) ≈ CFL_by_hand
end

function advective_timescale_cfl_on_stretched_grid(arch, FT)
    grid = RectilinearGrid(arch, size=(4, 4, 8), x=(0, 100), y=(0, 100), z=[k^2 for k in 0:8])
    model = NonhydrostaticModel(grid)

    Δt = FT(15.5)

    Δx = model.grid.Δxᶜᵃᵃ
    Δy = model.grid.Δyᵃᶜᵃ

    # At k = 1, w = 0 so the CFL constraint happens at the second face (k = 2).
    Δz_min = @allowscalar Oceananigans.Operators.Δzᵃᵃᶠ(1, 1, 2, grid)

    u₀ = FT(1.2)
    v₀ = FT(-2.5)
    w₀ = FT(3.9)

    set!(model, u=u₀, v=v₀, w=w₀, enforce_incompressibility=false)

    CFL_by_hand = Δt * (abs(u₀) / Δx + abs(v₀) / Δy + abs(w₀) / Δz_min)

    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)

    return cfl(model) ≈ CFL_by_hand
end

function advective_timescale_cfl_on_lat_lon_grid(arch, FT)
    grid = LatitudeLongitudeGrid(arch, size=(8, 8, 8), longitude=(-10, 10), latitude=(0, 45), z=(-1000, 0))
    model = HydrostaticFreeSurfaceModel(grid; momentum_advection=VectorInvariant())

    Δt = FT(1000)

    Nx, Ny, Nz = size(grid)

    # Will be the smallest at higher latitudes.
    Δx_min = @allowscalar Oceananigans.Operators.Δxᶠᶜᵃ(1, Ny, 1, grid)

    # Will be the same at every grid point.
    Δy_min = @allowscalar Oceananigans.Operators.Δyᶜᶠᵃ(1, 1, 1, grid)

    Δz = model.grid.z.Δᵃᵃᶠ

    u₀ = FT(1.2)
    v₀ = FT(-2.5)
    w₀ = FT(-0.1)

    # Avoid recomputation of w
    set!(model.velocities.u, u₀)
    set!(model.velocities.v, v₀)
    set!(model.velocities.w, w₀)

    CFL_by_hand = Δt * (abs(u₀) / Δx_min + abs(v₀) / Δy_min + abs(w₀) / Δz)

    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)

    return cfl(model) ≈ CFL_by_hand
end

function advective_timescale_cfl_on_flat_2d_grid(arch, FT)
    Δx = 0.5
    topo = (Periodic, Flat, Bounded)
    grid = RectilinearGrid(arch, FT, topology=topo, size=(3, 3), x=(0, 3Δx), z=(0, 3Δx))

    model = NonhydrostaticModel(grid)
    set!(model, v=1)

    Δt = FT(1.7)
    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)

    return cfl(model) == 0
end

get_iteration(model) = model.clock.iteration
get_time(model) = model.clock.time

@testset "Diagnostics" begin
    @info "Testing diagnostics..."

    for arch in archs
        @testset "CFL [$(typeof(arch))]" begin
            @info "  Testing CFL diagnostics [$(typeof(arch))]..."
            for FT in float_types
                @test diffusive_cfl_diagnostic_is_correct(arch, FT)
                @test advective_cfl_diagnostic_is_correct_on_regular_grid(arch, FT)
                @test advective_cfl_diagnostic_is_correct_on_vertically_stretched_grid(arch, FT)
                @test advective_timescale_cfl_on_regular_grid(arch, FT)
                @test advective_timescale_cfl_on_stretched_grid(arch, FT)
                @test advective_timescale_cfl_on_lat_lon_grid(arch, FT)
                @test advective_timescale_cfl_on_flat_2d_grid(arch, FT)
            end
        end
    end
end
