using Oceananigans.Architectures: Architectures, on_architecture
using Oceananigans.BoundaryConditions: BoundaryConditions, DiscreteBoundaryFunction, getbc
using Oceananigans.Grids: AbstractGrid, topology

struct TKETopBoundaryConditionParameters{C, U}
    top_tracer_boundary_conditions :: C
    top_velocity_boundary_conditions :: U
end

const TKEBoundaryFunction = DiscreteBoundaryFunction{<:TKETopBoundaryConditionParameters}

@inline Adapt.adapt_structure(to, p::TKETopBoundaryConditionParameters) =
    TKETopBoundaryConditionParameters(adapt(to, p.top_tracer_boundary_conditions),
                                      adapt(to, p.top_velocity_boundary_conditions))

@inline Architectures.on_architecture(to, p::TKETopBoundaryConditionParameters) =
    TKETopBoundaryConditionParameters(on_architecture(to, p.top_tracer_boundary_conditions),
                                      on_architecture(to, p.top_velocity_boundary_conditions))

@inline BoundaryConditions.getbc(condition::TKEBoundaryFunction, i::Integer, j::Integer, grid::AbstractGrid, clock, fields, clo, buoyancy) =
    condition.func(i, j, grid, clock, fields, condition.parameters, clo, buoyancy)

@inline BoundaryConditions.getbc(condition::TKEBoundaryFunction, i::Integer, j::Integer, k::Integer, grid::AbstractGrid, clock, fields, clo, buoyancy) =
    condition.func(i, j, k, grid, clock, fields, condition.parameters, clo, buoyancy)

"""
$(TYPEDSIGNATURES)

Compute the flux of TKE through the surface / top boundary.
Designed to be used with TKETopBoundaryConditionParameters in a FluxBoundaryCondition, eg:

```
top_tracer_bcs = top_tracer_boundary_conditions(grid, tracer_names, user_bcs)
top_velocity_bcs = top_velocity_boundary_conditions(grid, user_bcs)
parameters = TKETopBoundaryConditionParameters(top_tracer_bcs, top_velocity_bcs)
top_tke_bc = FluxBoundaryCondition(top_tke_flux, discrete_form=true, parameters=parameters)
```

See the implementation in catke_equation.jl.
"""
@inline top_tke_flux(i, j, grid, clock, fields, parameters, closure, buoyancy) = zero(grid)

#####
##### For model constructors
#####

""" Infer tracer boundary conditions from user_bcs and tracer_names. """
function top_tracer_boundary_conditions(grid, tracer_names, user_bcs)
    default_tracer_bcs = NamedTuple(c => FieldBoundaryConditions(grid, (Center(), Center(), Center())) for c in tracer_names)
    bcs = merge(default_tracer_bcs, user_bcs)
    return NamedTuple(c => bcs[c].top for c in tracer_names)
end

""" Infer velocity boundary conditions from `user_bcs` and `tracer_names`. """
function top_velocity_boundary_conditions(grid, user_bcs)
    default_top_bc = default_prognostic_bc(topology(grid, 3)(), Center(), DefaultBoundaryCondition())

    user_bc_names = keys(user_bcs)
    u_top_bc = :u ∈ user_bc_names ? user_bcs.u.top : default_top_bc
    v_top_bc = :v ∈ user_bc_names ? user_bcs.v.top : default_top_bc

    return (u=u_top_bc, v=v_top_bc)
end

""" Computes the friction velocity u★ based on fluxes of u and v. """
@inline function friction_velocity(i, j, grid, clock, fields, velocity_bcs)
    FT = eltype(grid)
    τx = getbc(velocity_bcs.u, i, j, grid, clock, fields)
    τy = getbc(velocity_bcs.v, i, j, grid, clock, fields)
    return sqrt(sqrt(τx^2 + τy^2))
end

""" Computes the convective velocity w★. """
@inline function top_convective_turbulent_velocity_cubed(i, j, grid, clock, fields, buoyancy, tracer_bcs)
    Jᵇ = top_buoyancy_flux(i, j, grid, buoyancy, tracer_bcs, clock, fields)
    Δz = Δzᶜᶜᶜ(i, j, grid.Nz, grid)
    return clip(Jᵇ) * Δz
end

@inline top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple::Tuple{<:Any}, buoyancy) =
    top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple[1], buoyancy)

@inline top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple::Tuple{<:Any, <:Any}, buoyancy) =
    top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple[1], buoyancy) +
    top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple[2], buoyancy)

@inline top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple::Tuple{<:Any, <:Any, <:Any}, buoyancy) =
    top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple[1], buoyancy) +
    top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple[2], buoyancy) +
    top_tke_flux(i, j, grid, clock, fields, parameters, closure_tuple[3], buoyancy)
