"""
    Checkpointer{I, T, P, A} <: AbstractOutputWriter

An output writer for checkpointing models to a JLD2 file from which models can be restored.
"""
mutable struct Checkpointer{I, T, P, A} <: AbstractOutputWriter
         frequency :: I
          interval :: T
          previous :: Float64
               dir :: String
            prefix :: String
        properties :: P
    has_array_refs :: A
             force :: Bool
           verbose :: Bool
end

"""
    Checkpointer(model; frequency=nothing, interval=nothing, dir=".", prefix="checkpoint",
                        force=false, verbose=false,
                        properties = [:architecture, :boundary_conditions, :grid, :clock,
                                      :eos, :constants, :closure, :velocities, :tracers,
                                      :timestepper])

Construct a `Checkpointer` that checkpoints the model to a JLD2 file every so often as
specified by `frequency` or `interval`. The `model.clock.iteration` is included in the
filename to distinguish between multiple checkpoint files.

Note that extra model `properties` can be safely specified, but removing crucial properties
such as `:velocities` will make restoring from the checkpoint impossible.

The checkpoint file is generated by serializing model properties to JLD2. However,
functions cannot be serialized to disk (at least not with JLD2). So if a model property
contains a reference somewhere in its hierarchy it will not be included in the checkpoint
file (and you will have to manually restore them).

Keyword arguments
=================
- `frequency::Int`   : Save output every `n` model iterations.
- `interval::Int`    : Save output every `t` units of model clock time.
- `dir::String`      : Directory to save output to. Default: "." (current working directory).
- `prefix::String`   : Descriptive filename prefixed to all output files. Default: "checkpoint".
- `force::Bool`      : Remove existing files if their filenames conflict. Default: `false`.
- `verbose::Bool`    : Log what the output writer is doing with statistics on compute/write times and file sizes.
                       Default: `false`.
- `properties::Array`: List of model properties to checkpoint.
"""
function Checkpointer(model; frequency=nothing, interval=nothing, dir=".", prefix="checkpoint", force=false,
                      verbose=false, properties = [:architecture, :boundary_conditions, :grid, :clock, :coriolis,
                                                   :buoyancy, :closure, :velocities, :tracers, :timestepper])

    validate_interval(frequency, interval)

    # Grid needs to be checkpointed for restoring to work.
    :grid ∉ properties && push!(properties, :grid)

    has_array_refs = Symbol[]

    for p in properties
        isa(p, Symbol) || @error "Property $p to be checkpointed must be a Symbol."
        p ∉ propertynames(model) && @error "Cannot checkpoint $p, it is not a model property!"

        if has_reference(Function, getproperty(model, p))
            @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed."
            filter!(e -> e != p, properties)
        end

        has_reference(AbstractField, getproperty(model, p)) && push!(has_array_refs, p)
    end

    mkpath(dir)
    return Checkpointer(frequency, interval, 0.0, dir, prefix, properties, has_array_refs, force, verbose)
end

function write_output(model, c::Checkpointer)
    filepath = joinpath(c.dir, c.prefix * string(model.clock.iteration) * ".jld2")
    c.verbose && @info "Checkpointing to file $filepath..."

    t0 = time_ns()
    jldopen(filepath, "w") do file
        file["checkpointed_properties"] = c.properties
        file["has_array_refs"] = c.has_array_refs

        serializeproperties!(file, model, filter(e -> !(e ∈ c.has_array_refs), c.properties))
        saveproperties!(file, model, filter(e -> e ∈ c.has_array_refs, c.properties))
    end

    t1, sz = time_ns(), filesize(filepath)
    c.verbose && @info "Checkpointing done: time=$(prettytime((t1-t0)/1e9)), size=$(pretty_filesize(sz))"
end

defaultname(::Checkpointer, nelems) = :checkpointer
_arr(::CPU, a) = a
_arr(::GPU, a) = CuArray(a)

function restore_fields!(model, file, arch, fieldset; location="$fieldset")
    if fieldset == :timestepper
        restore_fields!(model.timestepper, file, arch, :Gⁿ; location="timestepper/Gⁿ")
        restore_fields!(model.timestepper, file, arch, :G⁻; location="timestepper/G⁻")
    else
        for p in propertynames(getproperty(model, fieldset))
            getproperty(getproperty(model, fieldset), p).data.parent .= _arr(arch, file[location * "/$p"])
        end
    end
end

"""
    restore_from_checkpoint(filepath; kwargs=Dict())

Restore a model from the checkpoint file stored at `filepath`. `kwargs` can be passed
to the `Model` constructor, which can be especially useful if you need to manually
restore forcing functions or boundary conditions that rely on functions.
"""
function restore_from_checkpoint(filepath; kwargs=Dict())
    file = jldopen(filepath, "r")

    cps = file["checkpointed_properties"]
    has_array_refs = file["has_array_refs"]

    # Restore model properties that were just serialized.
    # We skip fields that contain structs and restore them after model creation.
    for p in cps
        if p ∉ has_array_refs
            kwargs[p] = file["$p"]
        end
    end

    model = Model(; kwargs...)

    # Now restore fields.
    for p in cps
        if p in has_array_refs
            restore_fields!(model, file, model.architecture, p)
        end
    end

    close(file)

    return model
end
