using Oceananigans.Grids: topology, _node, φnode, φnodes, λnode, λnodes,
                          XFlatGrid, YFlatGrid, ZFlatGrid,
                          XYFlatGrid, YZFlatGrid, XZFlatGrid, XYZFlatGrid,
                          XRegularRG, YRegularRG, ZRegularRG,
                          XRegularLLG, YRegularLLG, ZRegularLLG,
                          ZRegOrthogonalSphericalShellGrid
using Oceananigans.Operators: Δx, Δy, Δz
using Oceananigans.Utils: interpolator, _interpolate

# GPU-compatile middle point calculation
@inline middle_point(l, h) = Base.unsafe_trunc(Int, (l + h) / 2)

"""
$(TYPEDSIGNATURES)

Return indices `low, high` of `vec`tor for which

```julia
vec[low] ≤ val && vec[high] ≥ val
```

using a binary search. The input array `vec` has to be monotonically increasing.

Code credit: https://gist.github.com/cuongld2/8e4fed9ba44ea2b4598f90e7d5b6c612/155f9cb595314c8db3a266c3316889443b068017
"""
@inline function index_binary_search(vec, val::Number, N)
    low = 0
    high = N - 1

    while low + 1 < high
        mid = middle_point(low, high)
        if @inbounds vec[mid + 1] == val
            return (mid + 1, mid + 1)
        elseif @inbounds vec[mid + 1] < val
            low = mid
        else
            high = mid
        end
    end

    return (low + 1, high + 1)
end

@inline function fractional_index(val, vec, N)
    i₁, i₂ = index_binary_search(vec, val, N)

    @inbounds x₁ = vec[i₁]
    @inbounds x₂ = vec[i₂]

    ii = (i₂ - i₁) / (x₂ - x₁) * (val - x₁) + i₁
    ii = ifelse(i₁ == i₂, i₁, ii)

    FT = typeof(val) # convert "fractional index" to type of `val`
    return convert(FT, ii)
end

#####
##### Note: interpolation on LatitudeLongitudeGrid assumes a thin shell
##### (i.e. curvature effects are not incorporated when interpolating).
##### Use other methods if a more accurate interpolation is required.
#####

@inline fractional_x_index(x, locs, grid::XFlatGrid) = zero(grid)

@inline function fractional_x_index(x, locs, grid::XRegularRG)
    x₀ = xnode(1, 1, 1, grid, locs...)
    dx = Δx(1, 1, 1, grid, locs...)
    FT = eltype(grid)
    return convert(FT, (x - x₀) / dx) + 1 # 1 - based indexing
end

@inline function fractional_x_index(x, locs, grid::RectilinearGrid)
    loc = @inbounds locs[1]
     Tx = topology(grid, 1)()
     Nx = length(loc, Tx, grid.Nx)
     xn = xnodes(grid, locs...)
     FT = eltype(grid)
    return convert(FT, fractional_index(x, xn, Nx))
end

# Because of precision errors with numbers close to 0,
# we need to make sure we approach the correct limit also from the left.
@inline function convert_to_0_360(x::FT) where FT
    x₀ = ((x % 360) + 360) % 360
    x⁻ = - eps(convert(FT, 360))
    return ifelse(x⁻ ≤ x < 0, 360 + x, x₀)
end

# No need to check precision for integers
@inline convert_to_0_360(x::Integer) = ((x % 360) + 360) % 360

# Find n for which 360 * n ≤ λ ≤ 360 * (n + 1)
@inline find_λ_range(λ) = ifelse((λ < 0) & (mod(λ, 360) != 0), λ ÷ 360 - 1, λ ÷ 360)

# Convert x to lie in the λ₀ : λ₀ + 360 range by accounting for the cyclic
# nature of the longitude coordinate.
@inline function convert_to_λ₀_λ₀_plus360(x, λ₀)
    x  = convert_to_0_360(x)
    n  = find_λ_range(λ₀)
    λ⁻ = convert_to_0_360(λ₀)
    n  = ifelse(x ≥ λ⁻, n, n+1)
    return x + 360 * n
end

# When interpolating longitude values, we convert the longitude to
# interpolate to lie in the λ₀ : λ₀ + 360 range, where λ₀ is the westernmost node
# of the interpolating grid.
#
# The spacing is read from the grid rather than differenced from two adjacent nodes:
# subtracting nodes of magnitude ~180 discards most of their significant digits in Float32,
# and the resulting relative error is multiplied by the cell index. Every longitude spacing
# on an x-regular grid is `Δλᶜᵃᵃ`.
@inline function fractional_x_index(λ, locs, grid::XRegularLLG)
    λ₀ = λnode(1, 1, 1, grid, locs...)
    Δλ = grid.Δλᶜᵃᵃ
    λc = convert_to_λ₀_λ₀_plus360(λ, λ₀ - Δλ/2) # Making sure we have the right range
    FT = eltype(grid)
    return convert(FT, (λc - λ₀) / Δλ) + 1 # 1 - based indexing
end

# When interpolating longitude values, we convert the longitude to
# interpolate to lie in the λ₀ : λ₀ + 360 range, where λ₀ is the westernmost node
# of the interpolating grid.
@inline function fractional_x_index(λ, locs, grid::LatitudeLongitudeGrid)
    loc = @inbounds locs[1]
     Tλ = topology(grid, 1)()
     Nλ = length(loc, Tλ, grid.Nx)
     λn = λnodes(grid, locs...)
     λ₀ = @inbounds λn[1]
     λ₁ = @inbounds λn[2]
     Δλ = λ₁ - λ₀
     λc = convert_to_λ₀_λ₀_plus360(λ, λ₀ - Δλ/2)
     FT = eltype(grid)
    return convert(FT, fractional_index(λc, λn, Nλ))
end

@inline fractional_y_index(y, locs, grid::YFlatGrid) = zero(grid)

@inline function fractional_y_index(y, locs, grid::YRegularRG)
    y₀ = ynode(1, 1, 1, grid, locs...)
    dy = Δy(1, 1, 1, grid, locs...)
    FT = eltype(grid)
    return convert(FT, (y - y₀) / dy) + 1 # 1 - based indexing
end

@inline function fractional_y_index(φ, locs, grid::YRegularLLG)
    φ₀ = φnode(1, 1, 1, grid, locs...)
    FT = eltype(grid)
    # From the grid, not `φnode(1, 2, 1) - φ₀`; see `fractional_x_index` above.
    return convert(FT, (φ - φ₀) / grid.Δφᵃᶜᵃ) + 1 # 1 - based indexing
end

@inline function fractional_y_index(y, locs, grid::RectilinearGrid)
    loc = @inbounds locs[2]
     Ty = topology(grid, 2)()
     Ny = length(loc, Ty, grid.Ny)
     yn = ynodes(grid, locs...)
     FT = eltype(grid)
    return convert(FT, fractional_index(y, yn, Ny))
end

@inline function fractional_y_index(y, locs, grid::LatitudeLongitudeGrid)
    loc = @inbounds locs[2]
    Ty = topology(grid, 2)()
    Ny = length(loc, Ty, grid.Ny)
    yn = φnodes(grid, locs...)
    FT = eltype(grid)
    return convert(FT, fractional_index(y, yn, Ny))
end

@inline fractional_z_index(z, locs, grid::ZFlatGrid) = zero(grid)

ZRegGrid = Union{ZRegularRG, ZRegularLLG, ZRegOrthogonalSphericalShellGrid}

@inline function fractional_z_index(z::FT, locs, grid::ZRegGrid) where FT
    z₀ = znode(1, 1, 1, grid, locs...)
    dz = Δz(1, 1, 1, grid, locs...)
    return convert(FT, (z - z₀) / dz) + 1 # 1 - based indexing
end

@inline function fractional_z_index(z, locs, grid)
    loc = @inbounds locs[3]
    Tz = topology(grid, 3)()
    Nz = length(loc, Tz, grid.Nz)
    zn = znodes(grid, loc)
    FT = eltype(grid)
    return convert(FT, fractional_index(z, zn, Nz))
end

struct FractionalIndices{I, J, K}
    i :: I
    j :: J
    k :: K
end

"""
$(TYPEDSIGNATURES)

Convert the coordinate tuple `at_node` to _fractional_ indices on a regular rectilinear grid
located at `(ℓx, ℓy, ℓz)`, a triplet of `Center` and `Face`. Fractional indices are
floats indicating a location between grid points.
"""
@inline FractionalIndices(at_node, grid, ℓx, ℓy, ℓz) = _fractional_indices(at_node, grid, ℓx, ℓy, ℓz)

@inline FractionalIndices(at_node, grid::XFlatGrid, ℓx, ℓy, ℓz)   = _fractional_indices(at_node, grid, nothing, ℓy, ℓz)
@inline FractionalIndices(at_node, grid::YFlatGrid, ℓx, ℓy, ℓz)   = _fractional_indices(at_node, grid, ℓx, nothing, ℓz)
@inline FractionalIndices(at_node, grid::ZFlatGrid, ℓx, ℓy, ℓz)   = _fractional_indices(at_node, grid, ℓx, ℓy, nothing)
@inline FractionalIndices(at_node, grid::XYFlatGrid, ℓx, ℓy, ℓz)  = _fractional_indices(at_node, grid, nothing, nothing, ℓz)
@inline FractionalIndices(at_node, grid::YZFlatGrid, ℓx, ℓy, ℓz)  = _fractional_indices(at_node, grid, ℓx, nothing, nothing)
@inline FractionalIndices(at_node, grid::XZFlatGrid, ℓx, ℓy, ℓz)  = _fractional_indices(at_node, grid, nothing, ℓy, nothing)
@inline FractionalIndices(at_node, grid::XYZFlatGrid, ℓx, ℓy, ℓz) = _fractional_indices(at_node, grid, nothing, nothing, nothing)

@inline function _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ℓx, ℓy, ℓz)
    ii = fractional_x_index(x, (ℓx, ℓy, ℓz), grid)
    jj = fractional_y_index(y, (ℓx, ℓy, ℓz), grid)
    kk = fractional_z_index(z, (ℓx, ℓy, ℓz), grid)
    return FractionalIndices(ii, jj, kk)
end

# All seven single/double/triple-`Nothing` combinations must be defined for 3-tuple
# input so that pairwise dispatch resolves via a strictly more-specific method.
@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ::Nothing, ℓy, ℓz) =
    _fractional_indices((y, z), grid, nothing, ℓy, ℓz)

@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ℓx, ::Nothing, ℓz) =
    _fractional_indices((x, z), grid, ℓx, nothing, ℓz)

@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ℓx, ℓy, ::Nothing) =
    _fractional_indices((x, y), grid, ℓx, ℓy, nothing)

@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ::Nothing, ::Nothing, ℓz) =
    _fractional_indices((z,), grid, nothing, nothing, ℓz)

@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ℓx, ::Nothing, ::Nothing) =
    _fractional_indices((x,), grid, ℓx, nothing, nothing)

@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ::Nothing, ℓy, ::Nothing) =
    _fractional_indices((y,), grid, nothing, ℓy, nothing)

@inline _fractional_indices((x, y, z)::NTuple{3, Any}, grid, ::Nothing, ::Nothing, ::Nothing) =
    FractionalIndices(nothing, nothing, nothing)

# The `::NTuple{N, Any}` constraints below are load-bearing: Julia's function
# arg destructuring is permissive, so a pattern like `(z,)` will silently
# match a 2- or 3-tuple by taking only the first element — picking the wrong
# coordinate. The explicit arity gives the dispatcher something to bind to.

# Recursive prunes for "extra" coordinates: the kernel computes `at_node` from
# `to_grid` (so its Flat dims are already dropped), but the source field can
# also be reduced (with `Nothing` locations). When the source has more
# `Nothing` locations than the destination has `Flat` dims, the caller's
# `at_node` ends up longer than the number of axes we actually need to
# interpolate over. Recurse, dropping the leading axis until the arity matches.

@inline _fractional_indices(at::NTuple{2, Any}, grid, ::Nothing, ::Nothing, ℓz) =
    _fractional_indices((at[2],), grid, nothing, nothing, ℓz)

@inline _fractional_indices(at::NTuple{2, Any}, grid, ::Nothing, ℓy, ::Nothing) =
    _fractional_indices((at[2],), grid, nothing, ℓy, nothing)

@inline _fractional_indices(at::NTuple{2, Any}, grid, ℓx, ::Nothing, ::Nothing) =
    _fractional_indices((at[1],), grid, ℓx, nothing, nothing)

@inline _fractional_indices(at::NTuple{2, Any}, grid, ::Nothing, ::Nothing, ::Nothing) =
    FractionalIndices(nothing, nothing, nothing)

@inline function _fractional_indices((y, z)::NTuple{2, Any}, grid, ::Nothing, ℓy, ℓz)
    jj = fractional_y_index(y, (nothing, ℓy, ℓz), grid)
    kk = fractional_z_index(z, (nothing, ℓy, ℓz), grid)
    return FractionalIndices(nothing, jj, kk)
end

@inline function _fractional_indices((x, z)::NTuple{2, Any}, grid, ℓx, ::Nothing, ℓz)
    ii = fractional_x_index(x, (ℓx, nothing, ℓz), grid)
    kk = fractional_z_index(z, (ℓx, nothing, ℓz), grid)
    return FractionalIndices(ii, nothing, kk)
end

@inline function _fractional_indices((x, y)::NTuple{2, Any}, grid, ℓx, ℓy, ::Nothing)
    ii = fractional_x_index(x, (ℓx, ℓy, nothing), grid)
    jj = fractional_y_index(y, (ℓx, ℓy, nothing), grid)
    return FractionalIndices(ii, jj, nothing)
end

@inline function _fractional_indices((x,)::NTuple{1, Any}, grid, ℓx, ::Nothing, ::Nothing)
    loc = (ℓx, nothing, nothing)
    ii = fractional_x_index(x, loc, grid)
    jj = nothing
    kk = nothing
    return FractionalIndices(ii, jj, kk)
end

@inline function _fractional_indices((y,)::NTuple{1, Any}, grid, ::Nothing, ℓy, ::Nothing)
    loc = (nothing, ℓy, nothing)
    ii = nothing
    jj = fractional_y_index(y, loc, grid)
    kk = nothing
    return FractionalIndices(ii, jj, kk)
end

@inline function _fractional_indices((z,)::NTuple{1, Any}, grid, ::Nothing, ::Nothing, ℓz)
    loc = (nothing, nothing, ℓz)
    ii = nothing
    jj = nothing
    kk = fractional_z_index(z, loc, grid)
    return FractionalIndices(ii, jj, kk)
end

@inline _fractional_indices((at,)::NTuple{1, Any}, grid, ::Nothing, ::Nothing, ::Nothing) = FractionalIndices(nothing, nothing, nothing)

@inline _fractional_indices(at_node, grid, ::Nothing, ::Nothing, ::Nothing) = FractionalIndices(nothing, nothing, nothing)

"""
$(TYPEDSIGNATURES)

Interpolate `from_field`, `at_node`, on `from_grid` and at `from_loc`ation,
where `at_node` is a tuple of coordinates and and `from_loc = (ℓx, ℓy, ℓz)`.

Note that this is a lower-level `interpolate` method defined for use in CPU/GPU kernels.
"""
@inline interpolate(at_node, from_field, from_loc, from_grid) =
    interpolate(identity, at_node, from_field, from_loc, from_grid)

@inline interpolate(fidx::FractionalIndices, from_field, from_loc, from_grid) =
    interpolate(identity, fidx, from_field, from_loc, from_grid)

# interpolator, _interpolate, and ϕ₁-ϕ₈ are imported from Oceananigans.Utils

#####
##### Interpolation of a mapped field: `interpolate(func, …)`
#####

# Like `mean(func, itr)` / `sum(func, itr)`, a leading function `func` is applied to each source value
# *before* the weighted blend; the result stays in `func`-space (no inverse is applied). With
# `func = log` this blends log-values, so `exp(interpolate(log, …))` is a positivity-preserving,
# geometric-mean interpolation — accurate for exponentially-varying fields (pressure, density).
# `func = identity` reproduces the plain `interpolate` exactly. `func` is routed through the shared
# low-level blend `_interpolate` (via a lazy mapped array) so `Field`s and `FieldTimeSeries` share
# one code path.

"""
    MappedData(func, data)

Lazily apply `func` elementwise to `data`: a read returns `func(data[I...])`. This lets the existing
`_interpolate` blend `func`-mapped values without copying. GPU-safe when `func` is (`log`/`exp` are).
"""
struct MappedData{T, N, F, A} <: AbstractArray{T, N}
    func :: F
    data :: A
end

@inline MappedData(func::F, data::A) where {F, A} =
    MappedData{eltype(A), ndims(A), F, A}(func, data)

@inline Base.size(m::MappedData) = size(m.data)
@inline Base.axes(m::MappedData) = axes(m.data)
@inline Base.getindex(m::MappedData, I::Vararg{Int}) = m.func(@inbounds m.data[I...])

Adapt.adapt_structure(to, m::MappedData) = MappedData(m.func, adapt(to, m.data))

# `func = identity` returns `data` unchanged, so the identity path is byte-identical to plain
# `interpolate` and pays no wrapper overhead.
@inline mapped_data(func, data) = MappedData(func, data)
@inline mapped_data(::typeof(identity), data) = data

"""
$(TYPEDSIGNATURES)

Interpolate `from_field` `at_node` after mapping each source value through `func`, in the spirit of
`mean(func, itr)`: `func` is applied to each value before the weighted blend and the result is returned
in `func`-space (no inverse). With `func = log`, `exp(interpolate(log, …))` is a geometric-mean
interpolation. `func = identity` reproduces `interpolate(at_node, from_field, …)`.
"""
# `at_node::Tuple` (it is always a coordinate tuple here) keeps this disjoint from the
# `FieldTimeSeries` method `interpolate(to_node, ::Time, fts, …)`, whose second argument is a `Time`.
@inline function interpolate(func::Base.Callable, at_node::Tuple, from_field, from_loc, from_grid)
    fidx = FractionalIndices(at_node, from_grid, from_loc...)
    return interpolate(func, fidx, from_field, from_loc, from_grid)
end

@inline function interpolate(func::Base.Callable, fidx::FractionalIndices, from_field, from_loc, from_grid)
    ix = interpolator(fidx.i)
    iy = interpolator(fidx.j)
    iz = interpolator(fidx.k)
    return _interpolate(mapped_data(func, from_field), ix, iy, iz)
end

"""
$(TYPEDSIGNATURES)

Interpolate the `from_field` `to_node`.

`to_node` is an N-tuple corresponding the dimensionality of `from_field`:

* `to_node = (x, y, z)` for 3D fields and trilinear interpolation,
* `to_node = (x, y)`, `(x, z)`, or `(y, z)` for 2D fields and bilinear interpolation,
* `to_node = (x,)`, `(y,)`, or `(z,)` for 1D fields and linear interpolation.

For 1D interpolation, `to_node` may also be a `Number`.
"""
@inline function interpolate(to_node, from_field)
    from_loc = Tuple(L() for L in location(from_field))
    return interpolate(to_node, from_field, from_loc, from_field.grid)
end

@inline interpolate(to_node::Number, from_field) = interpolate(tuple(to_node), from_field)

@inline flatten_node(x, y, z) = (x, y, z)

@inline flatten_node(::Nothing, y, z) = flatten_node(y, z)
@inline flatten_node(x, ::Nothing, z) = flatten_node(x, z)
@inline flatten_node(x, y, ::Nothing) = flatten_node(x, y)

@inline flatten_node(x, ::Nothing, ::Nothing) = tuple(x)
@inline flatten_node(::Nothing, y, ::Nothing) = tuple(y)
@inline flatten_node(::Nothing, ::Nothing, z) = tuple(z)

@inline flatten_node(::Nothing, ::Nothing, ::Nothing) = tuple()

@inline flatten_node(x, y) = (x, y)
@inline flatten_node(::Nothing, ::Nothing) = tuple()
@inline flatten_node(::Nothing, y) = flatten_node(y)
@inline flatten_node(x, ::Nothing) = flatten_node(x)

@inline flatten_node(x) = tuple(x)
@inline flatten_node(::Nothing) = tuple()

@kernel function _interpolate!(to_field, to_grid, to_location,
                               from_field, from_grid, from_location)

    i, j, k = @index(Global, NTuple)

    to_node = _node(i, j, k, to_grid, to_location...)
    to_node = flatten_node(to_node...)

    @inbounds to_field[i, j, k] = interpolate(to_node, from_field, from_location, from_grid)
end

"""
$(TYPEDSIGNATURES)

Interpolate `from_field` `to_field` and then fill the halo regions of `to_field`.
"""
function interpolate!(to_field::Field, from_field::AbstractField)
    to_grid   = to_field.grid
    from_grid = from_field.grid

    to_arch   = architecture(to_field)
    from_arch = architecture(from_field)

    # In case architectures are `Distributed` we
    # verify that the fields are on the same child architecture
    to_arch   = child_architecture(to_arch)
    from_arch = child_architecture(from_arch)

    if !isnothing(from_arch) && to_arch != from_arch
        msg = "Cannot interpolate! because from_field is on $from_arch while to_field is on $to_arch."
        throw(ArgumentError(msg))
    end

    # Make locations
    from_location = Tuple(L() for L in location(from_field))
    to_location   = Tuple(L() for L in location(to_field))

    params = KernelParameters(interior_indices(to_field))

    launch!(to_arch, to_grid, params,
            _interpolate!, to_field, to_grid, to_location,
            from_field, from_grid, from_location)

    fill_halo_regions!(to_field)

    return to_field
end
