using Test

import StaticArrays
import ClimaCore
using ClimaCore: Geometry, Domains, Meshes, Topologies, Spaces, Fields

function toy_sphere(::Type{FT}) where {FT}
    radius = FT(1e7)
    zmax = FT(1e4)
    helem = npoly = 2
    velem = 4

    hdomain = Domains.SphereDomain(radius)
    hmesh = Meshes.EquiangularCubedSphere(hdomain, helem)
    htopology = Topologies.Topology2D(hmesh)
    quad = Spaces.Quadratures.GLL{npoly + 1}()
    hspace = Spaces.SpectralElementSpace2D(htopology, quad)

    vdomain = Domains.IntervalDomain(
        Geometry.ZPoint{FT}(zero(FT)),
        Geometry.ZPoint{FT}(zmax);
        boundary_tags = (:bottom, :top),
    )
    vmesh = Meshes.IntervalMesh(vdomain, nelems = velem)
    # center_space = Spaces.CenterFiniteDifferenceSpace(vmesh)

    vspace = Spaces.CenterFiniteDifferenceSpace(vmesh)

    # TODO: Replace this with a space that includes topography.
    center_space = Spaces.ExtrudedFiniteDifferenceSpace(hspace, vspace)
    # center_coords = Fields.coordinate_field(center_space)
    face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(center_space)

    # face_space = Spaces.FaceFiniteDifferenceSpace(center_space)
    return (;center_space,face_space)
end

function field_vec(center_space, face_space)
    Y = Fields.FieldVector(
        c = map(Fields.coordinate_field(center_space)) do coord
                FT = Spaces.undertype(center_space)
                (;
                    uₕ = Geometry.Covariant12Vector(FT(0), FT(0)),
                    uₕ_phys = StaticArrays.SVector(FT(0), FT(0)),
                    uₕ_phys_tup = (FT(0), FT(0)),
                )
            end
    )
    return Y
end

run_svec!(Y) = run_svec!(Y.c.uₕ_phys, Y.c.uₕ)
function run_svec!(uₕ_phys, uₕ)
    @. uₕ_phys = StaticArrays.SVector(
        Geometry.UVVector(uₕ).components.data.:1,
        Geometry.UVVector(uₕ).components.data.:2,
    )
    return nothing
end

run_tup!(Y) = run_tup!(Y.c.uₕ_phys_tup, Y.c.uₕ)
function run_tup!(uₕ_phys_tup, uₕ)
    @. uₕ_phys_tup = tuple(
        Geometry.UVVector(uₕ).components.data.:1,
        Geometry.UVVector(uₕ).components.data.:2,
    )
    return nothing
end

function test_alloc(Y)
    run_svec!(Y) # compile first
    p = @allocated run_svec!(Y)
    @test p == 0

    run_tup!(Y) # compile first
    p = @allocated run_tup!(Y)
    @test p == 0
end

(;center_space,face_space) = toy_sphere(Float32)
Y = field_vec(center_space,face_space)

@testset "Allocations" begin
    test_alloc(Y)
end
nothing