using Oceananigans: AbstractModel

using Oceananigans.AbstractOperations: KernelFunctionOperation
using Oceananigans.Advection: VectorInvariant
using Oceananigans.Architectures: AbstractArchitecture
using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
using Oceananigans.DistributedComputations
using Oceananigans.Fields: Field, tracernames, TracerFields, XFaceField, YFaceField, CenterField, compute!
using Oceananigans.Forcings: model_forcing
using Oceananigans.Grids: topology, Flat, architecture, RectilinearGrid, Center
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
using Oceananigans.Models: validate_model_halo, validate_tracer_advection
using Oceananigans.Advection: materialize_advection
using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!
using Oceananigans.TurbulenceClosures: with_tracers, build_closure_fields
using Oceananigans.Utils: tupleit

import Oceananigans: prognostic_state, restore_prognostic_state!
import Oceananigans.Architectures: architecture

const RectilinearGrids = Union{RectilinearGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:RectilinearGrid}}

function shallow_water_tendency_fields(grid, tracer_names, prognostic_names)
    u =  XFaceField(grid)
    v =  YFaceField(grid)
    h = CenterField(grid)
    tracers = TracerFields(tracer_names, grid)

    return NamedTuple{prognostic_names}((u, v, h, Tuple(tracers)...))
end

function ShallowWaterSolutionFields(grid, bcs, prognostic_names)
    u =  XFaceField(grid, boundary_conditions = getproperty(bcs, prognostic_names[1]))
    v =  YFaceField(grid, boundary_conditions = getproperty(bcs, prognostic_names[2]))
    h = CenterField(grid, boundary_conditions = getproperty(bcs, prognostic_names[3]))

    return NamedTuple{prognostic_names[1:3]}((u, v, h))
end

mutable struct ShallowWaterModel{G, A<:AbstractArchitecture, CL, GR, V, U, R, F, E, B, Q, C, K, TS, FR} <: AbstractModel{TS, A}
                          grid :: G         # Grid of physical points on which `Model` is solved
                  architecture :: A         # Computer `Architecture` on which `Model` is run
                         clock :: CL        # Tracks iteration number and simulation time of `Model`
    gravitational_acceleration :: GR        # Gravitational acceleration, full, or reduced
                     advection :: V         # Advection scheme for velocities, mass and tracers
                    velocities :: U         # Velocities in the shallow water model
                      coriolis :: R         # Set of parameters for the background rotation rate of `Model`
                       forcing :: F         # Container for forcing functions defined by the user
                       closure :: E         # Diffusive 'turbulence closure' for all model fields
                    bathymetry :: B         # Bathymetry/Topography for the model
                      solution :: Q         # Container for transports `uh`, `vh`, and height `h`
                       tracers :: C         # Container for tracer fields
                closure_fields :: K         # Container for closure auxiliary fields
                   timestepper :: TS        # Object containing timestepper fields and parameters
                   formulation :: FR        # Either conservative or vector-invariant
end

supported_timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)

struct ConservativeFormulation end

struct VectorInvariantFormulation end

"""
    ShallowWaterModel(grid;
                      gravitational_acceleration,
                      clock = Clock(grid),
                      momentum_advection = UpwindBiased(order=5),
                      tracer_advection = WENO(),
                      mass_advection = WENO(),
                      coriolis = nothing,
                      forcing::NamedTuple = NamedTuple(),
                      closure = nothing,
                      bathymetry = nothing,
                      tracers = (),
                      closure_fields = nothing,
                      boundary_conditions::NamedTuple = NamedTuple(),
                      timestepper::Symbol = :RungeKutta3,
                      formulation = ConservativeFormulation())

Construct a shallow water model on `grid` with `gravitational_acceleration` constant.

Arguments
=========

  - `grid`: (required) The resolution and discrete geometry on which `model` is solved. The
            architecture (CPU/GPU) that the model is solve is inferred from the architecture
            of the grid.

Keyword arguments
=================

  - `gravitational_acceleration`: (required) The gravitational acceleration constant.
  - `clock`: The `clock` for the model.
  - `momentum_advection`: The scheme that advects velocities. See `Oceananigans.Advection`.
    Default: `UpwindBiased(order=5)`.
  - `tracer_advection`: The scheme that advects tracers. See `Oceananigans.Advection`. Default: `WENO()`.
  - `mass_advection`: The scheme that advects the mass equation. See `Oceananigans.Advection`. Default:
    `WENO()`.
  - `coriolis`: Parameters for the background rotation rate of the model.
  - `forcing`: `NamedTuple` of user-defined forcing functions that contribute to solution tendencies.
  - `closure`: The turbulence closure for `model`. See `Oceananigans.TurbulenceClosures`.
  - `bathymetry`: The bottom bathymetry.
  - `tracers`: A tuple of symbols defining the names of the modeled tracers, or a `NamedTuple` of
               preallocated `CenterField`s.
  - `closure_fields`: Stores diffusivity fields when the closures require a diffusivity to be
                          calculated at each timestep.
  - `boundary_conditions`: `NamedTuple` containing field boundary conditions.
  - `timestepper`: A symbol or a `TimeStepper` object that specifies the time-stepping method.
                   Supported symbols include $(join("`" .* repr.(supported_timesteppers) .* "`", ", ")).
                   Default: `:RungeKutta3`.
  - `formulation`: Whether the dynamics are expressed in conservative form (`ConservativeFormulation()`;
                   default) or in non-conservative form with a vector-invariant formulation for the
                   non-linear terms (`VectorInvariantFormulation()`).

!!! warning "Formulation-grid compatibility requirements"
    The `ConservativeFormulation()` requires `RectilinearGrid`.
    Use `VectorInvariantFormulation()` with `LatitudeLongitudeGrid`.
"""
function ShallowWaterModel(grid;
                           gravitational_acceleration,
                           clock = Clock(grid),
                           momentum_advection = UpwindBiased(order=5),
                           tracer_advection = WENO(),
                           mass_advection = WENO(),
                           coriolis = nothing,
                           forcing::NamedTuple = NamedTuple(),
                           closure = nothing,
                           bathymetry = nothing,
                           tracers = (),
                           closure_fields = nothing,
                           boundary_conditions::NamedTuple = NamedTuple(),
                           timestepper::Symbol = :RungeKutta3,
                           formulation = ConservativeFormulation())

    @warn "The ShallowWaterModel is currently unvalidated, subject to change, and should not be used for scientific research without adequate validation."

    arch = architecture(grid)

    tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)

    topology(grid, 3) === Flat ||
        throw(ArgumentError("ShallowWaterModel requires `topology(grid, 3) === Flat`. " *
                            "Use `topology = ($(topology(grid, 1)), $(topology(grid, 2)), Flat)` " *
                            "when constructing `grid`."))

    (typeof(grid) <: RectilinearGrids || formulation == VectorInvariantFormulation()) ||
        throw(ArgumentError("`ConservativeFormulation()` requires a rectilinear `grid`. \n" *
                            "Use `VectorInvariantFormulation()` or change your grid to a rectilinear one."))

    if timestepper isa Symbol && timestepper ∉ supported_timesteppers
        msg = """
        timestepper = :$timestepper is not supported.
        Supported timesteppers are: $(join(repr.(supported_timesteppers), ", ")).
        You can also construct your own TimeStepper and pass it to the constructor.
        """
        throw(ArgumentError(msg))
    end

    # Check halos and throw an error if the grid's halo is too small
    validate_model_halo(grid, momentum_advection, tracer_advection, closure)

    prognostic_field_names = formulation isa ConservativeFormulation ? (:uh, :vh, :h, tracers...) :  (:u, :v, :h, tracers...)
    default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions()
                                                                           for name in prognostic_field_names))

    momentum_advection = validate_momentum_advection(momentum_advection, formulation)

    if isnothing(tracer_advection)
        tracer_advection_tuple = NamedTuple{tracernames(tracers)}(nothing for tracer in 1:length(tracers))
    else
        default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid)

        # Advection schemes
        tracer_advection_tuple = with_tracers(tracernames(tracers),
                                            tracer_advection,
                                            (name, tracer_advection) -> default_tracer_advection,
                                            with_velocities=false)
    end

    advection = merge((momentum=momentum_advection, mass=mass_advection), tracer_advection_tuple)

    # Fill any settings in advection scheme that might have been deferred until
    # the grid and backend is known
    advection = NamedTuple(name => materialize_advection(scheme, grid) for (name, scheme) in pairs(advection))

    bathymetry_field = CenterField(grid)
    if !isnothing(bathymetry)
        set!(bathymetry_field, bathymetry)
        fill_halo_regions!(bathymetry_field)
    else
        fill!(bathymetry_field, 0)
    end

    boundary_conditions = merge(default_boundary_conditions, boundary_conditions)
    boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, prognostic_field_names)

    solution = ShallowWaterSolutionFields(grid, boundary_conditions, prognostic_field_names)
    tracers  = TracerFields(tracers, grid, boundary_conditions)
    closure_fields = build_closure_fields(closure_fields, grid, clock, tracernames(tracers), boundary_conditions, closure)

    # Instantiate timestepper if not already instantiated
    timestepper = TimeStepper(timestepper, grid, tracernames(tracers);
                              Gⁿ = shallow_water_tendency_fields(grid, tracernames(tracers), prognostic_field_names),
                              G⁻ = shallow_water_tendency_fields(grid, tracernames(tracers), prognostic_field_names))

    # Regularize forcing and closure for model tracer and velocity fields.
    model_fields = merge(solution, tracers)
    forcing = model_forcing(forcing, model_fields)
    closure = with_tracers(tracernames(tracers), closure)

    model = ShallowWaterModel(grid,
                              arch,
                              clock,
                              eltype(grid)(gravitational_acceleration),
                              advection,
                              shallow_water_velocities(formulation, solution),
                              coriolis,
                              forcing,
                              closure,
                              bathymetry_field,
                              solution,
                              tracers,
                              closure_fields,
                              timestepper,
                              formulation)

    update_state!(model)

    return model
end

validate_momentum_advection(momentum_advection, formulation) = momentum_advection
validate_momentum_advection(momentum_advection, ::VectorInvariantFormulation) =
    throw(ArgumentError("VectorInvariantFormulation requires a vector invariant momentum advection scheme. \n"*
                        "Use `momentum_advection = VectorInvariant()`."))
validate_momentum_advection(momentum_advection::Union{VectorInvariant, Nothing}, ::VectorInvariantFormulation) = momentum_advection

formulation(model::ShallowWaterModel)  = model.formulation
architecture(model::ShallowWaterModel) = model.architecture
timestepper(model::ShallowWaterModel)  = model.timestepper

# The w velocity is needed to use generic TurbulenceClosures methods, therefore it is set to nothing
shallow_water_velocities(::VectorInvariantFormulation, solution) = (u = solution.u, v = solution.v, w = nothing)

# TODO: convert u and v into binary operations
function shallow_water_velocities(::ConservativeFormulation, solution)
    u = compute!(Field(solution.uh / solution.h))
    v = compute!(Field(solution.vh / solution.h))
    return (; u, v, w=nothing)
end

shallow_water_velocities(model::ShallowWaterModel) = shallow_water_velocities(model.formulation, model.solution)

shallow_water_fields(velocities, solution, tracers, ::ConservativeFormulation) = merge(velocities, solution, tracers)
shallow_water_fields(velocities, solution, tracers, ::VectorInvariantFormulation) = merge(solution, (; w = velocities.w), tracers)

#####
##### Checkpointing
#####

function prognostic_state(model::ShallowWaterModel)
    return (clock = prognostic_state(model.clock),
            solution = prognostic_state(model.solution),
            velocities = prognostic_state(model.velocities),
            tracers = prognostic_state(model.tracers),
            closure_fields = prognostic_state(model.closure_fields),
            timestepper = prognostic_state(model.timestepper))
end

function restore_prognostic_state!(restored::ShallowWaterModel, from)
    restore_prognostic_state!(restored.clock, from.clock)
    restore_prognostic_state!(restored.solution, from.solution)
    restore_prognostic_state!(restored.velocities, from.velocities)
    restore_prognostic_state!(restored.timestepper, from.timestepper)
    restore_prognostic_state!(restored.tracers, from.tracers)
    restore_prognostic_state!(restored.closure_fields, from.closure_fields)
    return restored
end

restore_prognostic_state!(::ShallowWaterModel, ::Nothing) = nothing
