using Dates: AbstractTime

import Oceananigans: initialize!, prognostic_state, restore_prognostic_state!

"""
    AbstractSchedule

Supertype for objects that schedule `OutputWriter`s and `Diagnostics`.
Schedule must define the functor `Schedule(model)` that returns true or
false.
"""
abstract type AbstractSchedule end

# Default behavior is no alignment.
schedule_aligned_time_step(schedule, clock, Δt) = Δt

# Fallback initialization for schedule: call the schedule,
# then return `true`, indicating that the schedule "actuates" at
# initial call.
function initialize!(schedule::AbstractSchedule, model)
    schedule(model)

    # the default behavior `return true` dictates that by default,
    # schedules actuate at the initial call.
    return true
end

#####
##### TimeInterval
#####

mutable struct TimeInterval{IT, TT} <: AbstractSchedule
    interval :: IT
    first_actuation_time :: TT
    actuations :: Int
end

"""
$(TYPEDSIGNATURES)

Return a callable `TimeInterval` that schedules periodic output or diagnostic evaluation
on a `interval` of simulation time, as kept by `model.clock`.
"""
function TimeInterval(interval)
    IT = period_type(interval)
    interval = convert(IT, interval)
    TT = time_type(interval)
    first_actuation_time = zero(TT)
    return TimeInterval{IT, TT}(interval, first_actuation_time, 0)
end

initialize!(schedule::TimeInterval, model) = initialize_actuations!(schedule, model.clock.time)

function initialize_actuations!(schedule::TimeInterval, first_actuation_time)

    if schedule.first_actuation_time isa Number && first_actuation_time isa Dates.AbstractDateTime
        T = typeof(schedule.first_actuation_time)
        msg = "Cannot use $T TimeInterval times with DateTime clock. Use a Dates.Period instead."
        throw(ArgumentError(msg))
    end

    schedule.first_actuation_time = first_actuation_time
    schedule.actuations = 0

    return true
end

function next_actuation_time(schedule::TimeInterval)
    t₀ = schedule.first_actuation_time
    N = schedule.actuations
    T = schedule.interval
    return add_time_interval(t₀, T, N + 1)
end

#  Return `Inf` if we exhausted the actuations in a TimeInterval(::SpecifiedTimes)
function next_actuation_time(schedule::TimeInterval{<:AbstractArray})
    schedule.actuations ≥ length(schedule.interval) && return Inf
    return add_time_interval(schedule.first_actuation_time, schedule.interval, schedule.actuations + 1)
end

function (schedule::TimeInterval)(model)
    t = model.clock.time
    t★ = next_actuation_time(schedule)

    # Array-interval schedules return `Inf` once exhausted; never fire again.
    t★ === Inf && return false

    if t ≥ t★
        if schedule.actuations < typemax(Int)
            schedule.actuations += 1
            # Advance actuations so the next actuation is strictly in the future.
            while schedule.actuations < typemax(Int)
                tN = next_actuation_time(schedule)
                if tN === Inf || tN > t
                    break
                end
                schedule.actuations += 1
            end
        else # re-initialize the schedule to t★
            initialize!(schedule, t★)
        end
        return true
    else
        return false
    end
end

function schedule_aligned_time_step(schedule::TimeInterval, clock, Δt)
    t★ = next_actuation_time(schedule)
    t★ === Inf && return Δt
    t = clock.time
    δt = time_difference_seconds(t★, t)
    return min(Δt, δt)
end

function prognostic_state(schedule::TimeInterval)
    return (first_actuation_time = schedule.first_actuation_time,
            actuations = schedule.actuations,
            interval = schedule.interval)
end

function restore_prognostic_state!(restored::TimeInterval, from)
    if hasproperty(from, :interval) && from.interval == restored.interval
        restored.first_actuation_time = from.first_actuation_time
        restored.actuations = from.actuations
    end
    return restored
end

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

#####
##### IterationInterval
#####

struct IterationInterval <: AbstractSchedule
    interval :: Int
    offset :: Int
end

"""
    IterationInterval(interval; offset=0)

Return a callable `IterationInterval` that "actuates" (i.e., schedules output or callback execution)
whenever the model iteration (modified by `offset`) is a multiple of `interval`.

For example,

* `IterationInterval(100)` actuates at iterations `[100, 200, 300, ...]`.
* `IterationInterval(100, offset=-1)` actuates at iterations `[99, 199, 299, ...]`.
"""
IterationInterval(interval::Int; offset=0) = IterationInterval(interval, offset)
(schedule::IterationInterval)(model) = (model.clock.iteration - schedule.offset) % schedule.interval == 0

next_actuation_time(schedule::IterationInterval) = Inf

# IterationInterval has no state
prognostic_state(schedule::IterationInterval) = nothing
restore_prognostic_state!(restored::IterationInterval, from) = restored
restore_prognostic_state!(restored::IterationInterval, from::Nothing) = nothing

#####
##### WallTimeInterval
#####

mutable struct WallTimeInterval{FT} <: AbstractSchedule
    interval :: FT
    previous_actuation_time :: FT
end

"""
    WallTimeInterval(interval; start_time = time_ns() * 1e-9)

Return a callable `WallTimeInterval` that schedules periodic output or diagnostic evaluation
on a `interval` of "wall time" while a simulation runs, in units of seconds.

The "wall time" is the actual real world time in seconds, as kept by an actual
or hypothetical clock hanging on your wall.

The keyword argument `start_time` can be used to specify a starting wall time
other than the moment `WallTimeInterval` is constructed.
"""
function WallTimeInterval(interval; start_time = time_ns() * 1e-9)
    FT = Oceananigans.defaults.FloatType
    return WallTimeInterval(convert(FT, interval), convert(FT, start_time))
end

function (schedule::WallTimeInterval)(model)
    wall_time = time_ns() * 1e-9

    if wall_time >= schedule.previous_actuation_time + schedule.interval
        # Shave overshoot off previous_actuation_time to prevent overshoot from accumulating
        schedule.previous_actuation_time = wall_time - rem(wall_time, schedule.interval)
        return true
    else
        return false
    end
end

# WallTimeInterval uses absolute wall clock time, which doesn't make sense to checkpoint.
# On restore, the schedule starts fresh from the current wall time.
prognostic_state(::WallTimeInterval) = nothing
restore_prognostic_state!(restored::WallTimeInterval,  ::Nothing) = restored

#####
##### SpecifiedTimes
#####

mutable struct SpecifiedTimes{FT} <: AbstractSchedule
    times :: Vector{FT}
    previous_actuation :: Int
end

"""
$(TYPEDSIGNATURES)

Return a `schedule::SpecifiedTimes` that "actuates" (i.e., schedules output or callback execution)
whenever the model's clock equals the specified values in `times`. For example,

* `SpecifiedTimes([1, 15.3])` actuates when `model.clock.time` is `1` and `15.3`.

!!! info "Sorting specified times"
    The specified `times` need not be ordered as the `SpecifiedTimes` constructor
    will check and order them in ascending order if needed.
"""
function SpecifiedTimes(times...)
    length(times) == 0 && return SpecifiedTimes(Float64[], 0)

    first_time = times[1]

    if all(t -> t isa Number, times)
        FT = Oceananigans.defaults.FloatType
        return SpecifiedTimes{FT}(sort([convert(FT, t) for t in times]), 0)
    elseif all(t -> t isa AbstractTime, times)
        TT = typeof(first_time)
        return SpecifiedTimes{TT}(sort(collect(times)), 0)
    else
        throw(ArgumentError("SpecifiedTimes expects all times to be numbers or all to be Date/DateTime."))
    end
end

function SpecifiedTimes(times::AbstractVector)
    return SpecifiedTimes(Tuple(times)...)
end

function next_actuation_time(st::SpecifiedTimes)
    if st.previous_actuation >= length(st.times)
        return Inf
    else
        return st.times[st.previous_actuation+1]
    end
end

function (st::SpecifiedTimes)(model)
    current_time = model.clock.time
    next_time = next_actuation_time(st)

    if next_time === Inf
        return false
    end

    if current_time >= next_time
        st.previous_actuation += 1
        return true
    end

    return false
end

initialize!(st::SpecifiedTimes, model) = st(model)

function schedule_aligned_time_step(schedule::SpecifiedTimes, clock, Δt)
    t★ = next_actuation_time(schedule)
    δt = t★ == Inf ? Δt : time_difference_seconds(t★, clock.time)
    return min(Δt, δt)
end

function specified_times_str(st)
    str_elems = ["$(prettytime(t)), " for t in st.times]
    str = string("[", str_elems...)

    # Remove final separator ", "
    str = str[1:end-2]

    # Add closing bracket
    return string(str, "]")
end

Base.copy(st::SpecifiedTimes) = SpecifiedTimes(copy(st.times), st.previous_actuation)

function prognostic_state(schedule::SpecifiedTimes)
    return (; previous_actuation = schedule.previous_actuation)
end

function restore_prognostic_state!(restored::SpecifiedTimes, from)
    restored.previous_actuation = from.previous_actuation
    return restored
end

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

#####
##### ConsecutiveIterations
#####

mutable struct ConsecutiveIterations{S} <: AbstractSchedule
    parent :: S
    consecutive_iterations :: Int
    previous_parent_actuation_iteration :: Int
end

"""
    ConsecutiveIterations(parent_schedule, N=1)

Return a `schedule::ConsecutiveIterations` that actuates both when `parent_schedule`
actuates, and at iterations immediately following the actuation of `parent_schedule`.
This can be used, for example, when one wants to use output to reproduce a first-order approximation
of the time derivative of a quantity.
"""
ConsecutiveIterations(parent_schedule, N=1) = ConsecutiveIterations(parent_schedule, N, 0)

function (schedule::ConsecutiveIterations)(model)
    if schedule.parent(model)
        schedule.previous_parent_actuation_iteration = model.clock.iteration
        return true
    elseif model.clock.iteration - schedule.consecutive_iterations <= schedule.previous_parent_actuation_iteration
        return true # The iteration _after_ schedule.parent actuated!
    else
        return false
    end
end

schedule_aligned_time_step(schedule::ConsecutiveIterations, clock, Δt) =
    schedule_aligned_time_step(schedule.parent, clock, Δt)

function prognostic_state(schedule::ConsecutiveIterations)
    return (parent = prognostic_state(schedule.parent),
            previous_parent_actuation_iteration = schedule.previous_parent_actuation_iteration)
end

function restore_prognostic_state!(restored::ConsecutiveIterations, from)
    restore_prognostic_state!(restored.parent, from.parent)
    restored.previous_parent_actuation_iteration = from.previous_parent_actuation_iteration
    return restored
end

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

#####
##### Any and AndSchedule
#####

struct AndSchedule{S} <: AbstractSchedule
    schedules :: S
    AndSchedule(schedules::S) where S <: Tuple = new{S}(schedules)
end

"""
$(TYPEDSIGNATURES)

Return a schedule that actuates when all `child_schedule`s actuate.
"""
AndSchedule(schedules...) = AndSchedule(Tuple(schedules))

# Note that multiple schedules that have a "state" (like TimeInterval and WallTimeInterval)
# could cause the logic of AndSchedule to fail, due to the short-circuiting nature of `all`.
(as::AndSchedule)(model) = all(schedule(model) for schedule in as.schedules)

struct OrSchedule{S} <: AbstractSchedule
    schedules :: S
    OrSchedule(schedules::S) where S <: Tuple = new{S}(schedules)
end

"""
$(TYPEDSIGNATURES)

Return a schedule that actuates when any of the `child_schedule`s actuates.
"""
OrSchedule(schedules...) = OrSchedule(Tuple(schedules))

function (as::OrSchedule)(model)
    # Ensure that all `schedules` get queried
    actuations = Tuple(schedule(model) for schedule in as.schedules)
    return any(actuations)
end

schedule_aligned_time_step(any_or_all_schedule::Union{OrSchedule, AndSchedule}, clock, Δt) =
    minimum(schedule_aligned_time_step(schedule, clock, Δt)
            for schedule in any_or_all_schedule.schedules)

#####
##### Show methods
#####

function Base.summary(schedule::IterationInterval)
    summary = string("IterationInterval(", schedule.interval, ")")
    if schedule.offset != 0
        summary *= " with offset $(schedule.offset)"
    end
    return summary
end
Base.summary(schedule::TimeInterval) = string("TimeInterval(", prettytime(schedule.interval), ")")
Base.summary(schedule::SpecifiedTimes) = string("SpecifiedTimes(", specified_times_str(schedule), ")")
Base.show(io::IO, schedule::SpecifiedTimes) = print(io, summary(schedule))
Base.summary(schedule::ConsecutiveIterations) = string("ConsecutiveIterations(",
                                                       summary(schedule.parent), ", ",
                                                       schedule.consecutive_iterations, ")")

const StatefulSchedules = Union{TimeInterval, SpecifiedTimes, ConsecutiveIterations}

materialize_schedule(s) = s
materialize_schedule(ss::StatefulSchedules) = deepcopy(ss) # required to reuse a pre-defined schedule with a state
materialize_schedule(or::OrSchedule) = OrSchedule(Tuple(materialize_schedule(s) for s in or.schedules))
materialize_schedule(and::AndSchedule) = AndSchedule(Tuple(materialize_schedule(s) for s in and.schedules))
