using Oceananigans.Grids: Grids, constructor_arguments, rnode
using Oceananigans.Fields: Field, fill_halo_regions!, interior
using Oceananigans.BoundaryConditions: FBC
using OffsetArrays: OffsetArray

#####
##### GridFittedBottom (2.5D immersed boundary with modified bottom height)
#####

abstract type AbstractGridFittedBottom{H} <: AbstractGridFittedBoundary end

# To enable comparison with PartialCellBottom in the limiting case that
# fractional cell height is 1.0.
struct CenterImmersedCondition end
struct InterfaceImmersedCondition end

struct GridFittedBottom{H, I} <: AbstractGridFittedBottom{H}
    bottom_height :: H
    immersed_condition :: I
end

Base.summary(::CenterImmersedCondition) = "CenterImmersedCondition"
Base.summary(::InterfaceImmersedCondition) = "InterfaceImmersedCondition"

const GFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:GridFittedBottom}

"""
    GridFittedBottom(bottom_height, [immersed_condition=CenterImmersedCondition()])

Return a bottom immersed boundary.

Arguments
=========

* `bottom_height`: an array or function that gives the height of the
                   bottom in absolute ``z`` coordinates.

* `immersed_condition`: Determine whether the part of the domain that is
                        immersed are all the cell centers that lie below
                        `bottom_height` (`CenterImmersedCondition()`; default)
                        or all the cell faces that lie below `bottom_height`
                        (`InterfaceImmersedCondition()`). The only purpose of
                        `immersed_condition` to allow `GridFittedBottom` and
                        `PartialCellBottom` to have the same behavior when the
                        minimum fractional cell height for partial cells is set
                        to 0.
"""
GridFittedBottom(bottom_height) = GridFittedBottom(bottom_height, CenterImmersedCondition())

# 1-based interior view of a bare bottom-height array.
@inline function bottom_height_interior(bottom_height)
    parent_ranges = ntuple(Val(ndims(bottom_height))) do d
        H = 1 - first(axes(bottom_height, d))
        (1 + H):(size(bottom_height, d) - H)
    end
    return view(parent(bottom_height), parent_ranges...)
end

@inline bottom_heights_equal(h1, h2) = h1 == h2
@inline bottom_heights_equal(h1::AbstractArray, h2::AbstractArray) = bottom_height_interior(h1) == bottom_height_interior(h2)

set_bottom_height!(bottom_field, bottom_height) = set!(bottom_field, bottom_height)

function set_bottom_height!(bottom_field, bottom_height::OffsetArray)
    source = on_architecture(architecture(bottom_field), bottom_height_interior(bottom_height))
    copyto!(interior(bottom_field), source)
    return bottom_field
end

bottom_height_field(bottom_data, grid) = Field{Center, Center, Nothing}(grid; data=bottom_data)

"""
$(TYPEDSIGNATURES)

Return a `Field` at `(Center, Center, Nothing)` that wraps the bottom height of `grid`, which is
stored on `grid.immersed_boundary` as a bare `OffsetArray`. The returned `Field` shares its data
with `grid`, so mutating it mutates the bottom height of `grid`.
"""
bottom_height_field(grid::IBG) = bottom_height_field(grid.immersed_boundary.bottom_height, grid.underlying_grid)

function Base.summary(ib::GridFittedBottom)
    bottom_interior = bottom_height_interior(ib.bottom_height)
    zmax  = maximum(bottom_interior)
    zmin  = minimum(bottom_interior)
    zmean = sum(bottom_interior) / length(bottom_interior)

    summary1 = "GridFittedBottom("

    summary2 = string("mean(z)=", prettysummary(zmean),
                      ", min(z)=", prettysummary(zmin),
                      ", max(z)=", prettysummary(zmax))

    summary3 = ")"

    return summary1 * summary2 * summary3
end

Base.summary(ib::GridFittedBottom{<:Function}) = @sprintf("GridFittedBottom(%s)", ib.bottom_height)

function Base.show(io::IO, ib::GridFittedBottom)
    print(io, summary(ib), '\n')
    print(io, "└── bottom_height: ", prettysummary(ib.bottom_height), '\n')
end

Architectures.on_architecture(arch, ib::GridFittedBottom) = GridFittedBottom(on_architecture(arch, ib.bottom_height), ib.immersed_condition)

Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.bottom_height), adapt(to, ib.immersed_condition))

"""
$(TYPEDSIGNATURES)

Returns a new `ib` that holds the numerical `immersed_boundary`.
If `ib` is an `AbstractGridFittedBottom`, `ib.bottom_height` is an `OffsetArray` holding the
z-coordinate of the top-most interface of the last ``immersed`` cell in the column (wrap it as a
`Field` with [`bottom_height_field`](@ref)). If `ib` is a `GridFittedBoundary`, `ib.mask` is a `Field` of
booleans that indicates whether a cell is immersed or not.
"""
function materialize_immersed_boundary(grid, ib::GridFittedBottom)
    bottom_field = Field{Center, Center, Nothing}(grid)
    set_bottom_height!(bottom_field, ib.bottom_height)
    @apply_regionally compute_numerical_bottom_height!(bottom_field, grid, ib)
    fill_halo_regions!(bottom_field)
    return GridFittedBottom(bottom_field.data, ib.immersed_condition)
end

compute_numerical_bottom_height!(bottom_field, grid, ib) =
    launch!(architecture(grid), grid, :xy, _compute_numerical_bottom_height!, bottom_field, grid, ib)

@kernel function _compute_numerical_bottom_height!(bottom_field, grid, ib::GridFittedBottom)
    i, j = @index(Global, NTuple)
    zb = @inbounds bottom_field[i, j, 1]
    @inbounds bottom_field[i, j, 1] = rnode(i, j, 1, grid, c, c, f)
    condition = ib.immersed_condition
    for k in 1:grid.Nz
        z⁺ = rnode(i, j, k+1, grid, c, c, f)
        z  = rnode(i, j, k,   grid, c, c, c)
        immersed_cell = ifelse(condition isa CenterImmersedCondition, z ≤ zb, z⁺ ≤ zb)
        @inbounds bottom_field[i, j, 1] = ifelse(immersed_cell, z⁺, bottom_field[i, j, 1])
    end
end

@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom)
    # We use `rnode` for the `immersed_cell` because we do not want to have
    # wetting or drying that could happen for a moving grid if we use znode
    z  = rnode(i, j, k, underlying_grid, c, c, c)
    zb = @inbounds ib.bottom_height[i, j, 1]
    return z ≤ zb
end

@inline function _immersed_cell(i, j, k::AbstractArray, underlying_grid, ib::GridFittedBottom)
    # We use `rnode` for the `immersed_cell` because we do not want to have
    # wetting or drying that could happen for a moving grid if we use znode
    z  = rnode(i, j, k, underlying_grid, c, c, c)
    zb = @inbounds ib.bottom_height[i, j, 1]
    _zb = Base.stack(collect(zb for _ in k))
    return z .≤ _zb
end

#####
##### Static column depth
#####

# AbstractGridFittedBottomImmersedBoundaryGrid
const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom}

@inline static_column_depthᶜᶜᵃ(i, j, ibg::AGFBIBG) = @inbounds rnode(i, j, ibg.Nz+1, ibg, c, c, f) - ibg.immersed_boundary.bottom_height[i, j, 1]
@inline static_column_depthᶜᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i, j-1, ibg), static_column_depthᶜᶜᵃ(i, j, ibg))
@inline static_column_depthᶠᶜᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i-1, j, ibg), static_column_depthᶜᶜᵃ(i, j, ibg))
@inline static_column_depthᶠᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶠᶜᵃ(i, j-1, ibg), static_column_depthᶠᶜᵃ(i, j, ibg))

# Make sure column_height works for horizontally-Flat topologies.
XFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Flat, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom}
YFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:AbstractGridFittedBottom}

@inline static_column_depthᶠᶜᵃ(i, j, ibg::XFlatAGFIBG) = static_column_depthᶜᶜᵃ(i, j, ibg)
@inline static_column_depthᶜᶠᵃ(i, j, ibg::YFlatAGFIBG) = static_column_depthᶜᶜᵃ(i, j, ibg)
@inline static_column_depthᶠᶠᵃ(i, j, ibg::XFlatAGFIBG) = static_column_depthᶜᶠᵃ(i, j, ibg)
@inline static_column_depthᶠᶠᵃ(i, j, ibg::YFlatAGFIBG) = static_column_depthᶠᶜᵃ(i, j, ibg)

function Grids.constructor_arguments(grid::AGFBIBG)
    underlying_grid_args, underlying_grid_kwargs = constructor_arguments(grid.underlying_grid)
    grid_fitted_bottom_args = Dict(:bottom_height      => grid.immersed_boundary.bottom_height,
                                   :immersed_condition => grid.immersed_boundary.immersed_condition)
    return underlying_grid_args, underlying_grid_kwargs, grid_fitted_bottom_args
end

function Base.:(==)(gfb1::GridFittedBottom, gfb2::GridFittedBottom)
    return bottom_heights_equal(gfb1.bottom_height, gfb2.bottom_height) && gfb1.immersed_condition == gfb2.immersed_condition
end
