export CirculationField, DissipationField, EnstrophyField

using LinearAlgebra: ldiv!

abstract type AbstractScalarField{divide_by_area} end

@inline Base.fieldname(f::AbstractScalarField, suffix) =
    Symbol(fieldname(f), suffix)

divide_by_area(::AbstractScalarField{D}) where {D} = D

struct CirculationField{divide_by_area} <: AbstractScalarField{divide_by_area}
    @inline CirculationField(; divide_by_area::Bool = false) =
        new{divide_by_area}()
end

Base.fieldname(::CirculationField) = :Î“

metadata(f::CirculationField) = (
    "divided_by_area" => divide_by_area(f),
)

abstract type DissipationLikeField{D} <: AbstractScalarField{D} end
Base.fieldname(::DissipationLikeField) = :Îµ

struct DissipationField{divide_by_area} <: DissipationLikeField{divide_by_area}
    inplane :: Bool
    Î½ :: Float64
    @inline function DissipationField(;
            divide_by_area::Bool = true,
            inplane_only::Bool = false,
            Î½ = 1.0,
        )
        new{divide_by_area}(inplane_only, Î½)
    end
end

compute_inplane(f::DissipationField) = f.inplane

metadata(f::DissipationField) = (
    "divided_by_area" => divide_by_area(f),
    "inplane (2D)" => compute_inplane(f),
    "viscosity" => f.Î½,
)

"""
    EnstrophyField

Represents the in-plane enstrophy field ``Î© = Ï‰_z^2 / 2``.
"""
struct EnstrophyField{divide_by_area} <: DissipationLikeField{divide_by_area}
    @inline EnstrophyField(; divide_by_area::Bool = false) =
        new{divide_by_area}()
end

# This is always computed in-plane (it's not the full enstrophy!).
compute_inplane(::EnstrophyField) = true

metadata(f::EnstrophyField) = (
    "divided_by_area" => divide_by_area(f),
    "inplane (2D)" => compute_inplane(f),
)

"""
    compute_from_velocity!(
        field::DissipationField, Îµ::AbstractArray{<:Real,2},
        v_hat::NTuple{2, AbstractArray{<:Complex,2}};
        ks, fft_plan, buf, buf_hat,
    )

Compute in-plane dissipation field, ``Îµ_z``, from in-plane velocity field in
Fourier space, `(\\hat{v}_x, \\hat{v}_y)``.

Definition:

```math
Îµ_z = 2Î½ (S_{xx}^2 + 2 S_{xy}^2 + S_{xz}^2)
```

where ``S_{ij} = (ð�œ•_i v_j + ð�œ•_j v_i) / 2``.
"""
function compute_from_velocity!(
        field::DissipationField, Îµ::AbstractArray{<:Real,2},
        v_hat::NTuple{2, AbstractArray{<:Complex,2}};
        ks, fft_plan, buf, buf_hat,
    )
    @assert compute_inplane(field)
    @assert size(buf) == size(Îµ)
    @assert size(buf_hat) == size(v_hat[1]) == length.(ks)

    Î½ = field.Î½

    # 1. Sxx = ð�œ•_x v_x
    @inbounds for (I, vx) in pairs(IndexCartesian(), v_hat[1])
        i = Tuple(I)[1]
        kx = ks[1][i]
        buf_hat[I] = im * kx * vx
    end
    ldiv!(buf, fft_plan, buf_hat)
    Îµ .= buf.^2

    # 2. Syy = ð�œ•_y v_y
    @inbounds for (I, vy) in pairs(IndexCartesian(), v_hat[2])
        j = Tuple(I)[2]
        ky = ks[2][j]
        buf_hat[I] = im * ky * vy
    end
    ldiv!(buf, fft_plan, buf_hat)
    Îµ .+= buf.^2

    # 3. Sxy = (ð�œ•_x v_y + ð�œ•_y v_x) / 2
    @inbounds for I in CartesianIndices(buf_hat)
        kx, ky = getindex.(ks, Tuple(I))
        buf_hat[I] = im * (kx * v_hat[2][I] + ky * v_hat[1][I]) / 2
    end
    ldiv!(buf, fft_plan, buf_hat)

    Îµ .= (Îµ .+ 2 .* buf.^2) .* 2Î½

    Îµ
end

function compute_from_velocity!(
        field::EnstrophyField, Î©::AbstractArray{<:Real,2},
        v_hat::NTuple{2, AbstractArray{<:Complex,2}};
        ks, fft_plan, buf_hat,
        buf = nothing,  # not needed; just for compatibility with other method
    )
    @assert compute_inplane(field)
    @assert size(buf_hat) == size(v_hat[1]) == length.(ks)

    # 1. Compute Ï‰_z
    Ï‰_hat = buf_hat
    Ï‰ = Î©
    @inbounds for I in CartesianIndices(Ï‰_hat)
        kx, ky = getindex.(ks, Tuple(I))
        Ï‰_hat[I] = im * (kx * v_hat[2][I] - ky * v_hat[1][I])
    end
    ldiv!(Ï‰, fft_plan, Ï‰_hat)

    # 2. Enstrophy
    Î© .= Ï‰.^2 ./ 2

    Î©
end

function find_field(
        ::Type{F},
        fields::Tuple{Vararg{AbstractScalarField}},
    ) where {F <: AbstractScalarField}
    _find_field(F, fields...)
end

function _find_field(::Type{F}, field, etc...) where {F}
    if field isa F
        field
    else
        _find_field(F, etc...)
    end
end

_find_field(::Type) = nothing
