using Oceananigans.Diagnostics: AbstractDiagnostic
using Oceananigans.OutputWriters: fetch_output
using Oceananigans.Utils: AbstractSchedule, prettytime, period_type, time_type
using Oceananigans.TimeSteppers: Clock
using Dates: Period, Second, value, AbstractDateTime

import Oceananigans: run_diagnostic!, prognostic_state, restore_prognostic_state!, initialize!
import Oceananigans.Utils: TimeInterval, SpecifiedTimes, initialize_actuations!
import Oceananigans.Grids: grid
import Oceananigans.Fields: location, indices, set!

"""
    mutable struct AveragedTimeInterval <: AbstractSchedule

Container for parameters that configure and handle time-averaged output.
"""
mutable struct AveragedTimeInterval{I, T} <: AbstractSchedule
    interval :: I
    window :: I
    stride :: Int
    first_actuation_time :: T
    actuations :: Int
    collecting :: Bool
end

"""
    AveragedTimeInterval(interval; window=interval, stride=1)

Returns a `schedule` that specifies periodic time-averaging of output.
The time `window` specifies the extent of the time-average, which
reoccurs every `interval`.

`output` is computed and accumulated into the average every `stride` iterations
during the averaging window. For example, `stride=1` computes output every iteration,
whereas `stride=2` computes output every other iteration. Time-averages with
longer `stride`s are faster to compute, but less accurate.

The time-average of ``a`` is a left Riemann sum corresponding to

```math
⟨a⟩ = T⁻¹ \\int_{tᵢ-T}^{tᵢ} a \\mathrm{d} t \\, ,
```

where ``⟨a⟩`` is the time-average of ``a``, ``T`` is the time-window for averaging,
and the ``tᵢ`` are discrete times separated by the time `interval`. The ``tᵢ`` specify
both the end of the averaging window and the time at which output is written.

Example
=======

```jldoctest averaged_time_interval
using Oceananigans.OutputWriters: AveragedTimeInterval
using Oceananigans.Units: days

schedule = AveragedTimeInterval(4days, window=2days)

# output
AveragedTimeInterval(window=2 days, stride=1, interval=4 days)
```

An `AveragedTimeInterval` schedule directs an output writer
to time-average its outputs before writing them to disk:

```jldoctest averaged_time_interval
using Oceananigans
using Oceananigans.Units

grid = RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))

model = NonhydrostaticModel(grid)

simulation = Simulation(model, Δt=10minutes, stop_time=30days)

simulation.output_writers[:velocities] = JLD2Writer(model, model.velocities,
                                                    filename= "averaged_velocity_data.jld2",
                                                    schedule = AveragedTimeInterval(4days, window=2days, stride=2))

# output

JLD2Writer scheduled on TimeInterval(4 days):
├── filepath: averaged_velocity_data.jld2
├── 3 outputs: (u, v, w) averaged on AveragedTimeInterval(window=2 days, stride=2, interval=4 days)
├── array_type: Array{Float32}
├── including: [:coriolis, :buoyancy, :closure]
├── file_splitting: NoFileSplitting
└── file size: 0 bytes (file not yet created)
```
"""
function AveragedTimeInterval(interval; window=interval, stride=1)
    window > interval && throw(ArgumentError("Averaging window $window is greater than the output interval $interval."))
    IT = period_type(interval)
    interval = convert(IT, interval)
    window = convert(IT, window)
    TT = time_type(interval)
    first_actuation_time = zero(TT)
    return AveragedTimeInterval{IT, TT}(interval, window, stride, first_actuation_time, 0, false)
end

function next_actuation_time(sch::AveragedTimeInterval)
    t₀ = sch.first_actuation_time
    N = sch.actuations
    interval = sch.interval
    return t₀ + (N + 1) * interval
    # the next actuation time is the end of the time averaging window
end

# Schedule actuation
function (sch::AveragedTimeInterval)(model)
    scheduled = sch.collecting || model.clock.time > next_actuation_time(sch) - sch.window
    return scheduled
end

initialize!(sch::AveragedTimeInterval, model) = initialize_actuations!(sch, model.clock.time)

function initialize_actuations!(schedule::AveragedTimeInterval, first_actuation_time)
    if schedule.first_actuation_time isa Number && first_actuation_time isa AbstractDateTime
        T = typeof(schedule.first_actuation_time)
        msg = "Cannot use $T AveragedTimeInterval 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

outside_window(sch::AveragedTimeInterval, clock) = clock.time <= next_actuation_time(sch) - sch.window
initialize_schedule!(sch::AveragedTimeInterval, clock) = nothing

# Accumulated clock time from repeated Δt additions may fall just short of target times
# due to floating point arithmetic. For example, 10 iterations of Δt=0.1 yields
# clock.time = 0.9999999999999999 rather than 1.0, causing (0.999... >= 1.0) = false.
# Using eps(t★) as tolerance ensures times within one floating point increment are accepted.
function end_of_window(sch::AveragedTimeInterval, clock)
    t★ = next_actuation_time(sch)
    return clock.time >= t★ - eps(t★)
end

TimeInterval(sch::AveragedTimeInterval) = TimeInterval(sch.interval)
Base.copy(sch::AveragedTimeInterval) = AveragedTimeInterval(sch.interval, window=sch.window, stride=sch.stride)

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

function prognostic_state(schedule::AveragedTimeInterval)
    return (first_actuation_time = schedule.first_actuation_time,
            actuations = schedule.actuations,
            collecting = schedule.collecting)
end

function restore_prognostic_state!(restored::AveragedTimeInterval, from)
    restored.first_actuation_time = from.first_actuation_time
    restored.actuations = from.actuations
    restored.collecting = from.collecting
    return restored
end

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

#####
##### WindowedTimeAverage
#####

mutable struct WindowedTimeAverage{OP, R, T, S} <: AbstractDiagnostic
                      result :: R
                     operand :: OP
           window_start_time :: T
      window_start_iteration :: Int
    previous_collection_time :: T
                    schedule :: S
               fetch_operand :: Bool
end

const IntervalWindowedTimeAverage = WindowedTimeAverage{<:Any, <:Any, <:Any, <:AveragedTimeInterval}
const SpecifiedWindowedTimeAverage = WindowedTimeAverage{<:Any, <:Any, <:Any, <:AveragedSpecifiedTimes}

stride(wta::IntervalWindowedTimeAverage) = wta.schedule.stride
stride(wta::SpecifiedWindowedTimeAverage) = wta.schedule.stride

"""
    WindowedTimeAverage(operand, model=nothing; schedule, fetch_operand=true)

Returns an object for computing running averages of `operand` over `schedule.window` and
recurring on `schedule.interval`, where `schedule` is an `AveragedTimeInterval` or `AveragedSpecifiedTimes`.
During the collection period, averages are computed every `schedule.stride` iteration.

`operand` may be a `Oceananigans.Field` or a function that returns an array or scalar.

Calling `wta(model)` for `wta::WindowedTimeAverage` object returns `wta.result`.
"""
function WindowedTimeAverage(operand, model=nothing; schedule, fetch_operand=true)

    if fetch_operand
        output = fetch_output(operand, model)
        result = similar(output)
        result .= output
    else
        result = similar(operand)
        result .= operand
    end

    time = isnothing(model) ? get_default_time(schedule) : model.clock.time

    return WindowedTimeAverage(result, operand, time, 0, time, schedule, fetch_operand)
end

# Helper functions to get default time based on schedule type
get_default_time(schedule::AveragedTimeInterval) = zero(typeof(schedule.interval))
get_default_time(schedule::AveragedSpecifiedTimes) = zero(eltype(schedule.specified_times.times))

# Time-averaging doesn't change spatial location or grid
grid(wta::WindowedTimeAverage) = grid(wta.operand)
location(wta::WindowedTimeAverage) = location(wta.operand)
indices(wta::WindowedTimeAverage) = indices(wta.operand)
set!(u::Field, wta::WindowedTimeAverage) = set!(u, wta.result)
Base.parent(wta::WindowedTimeAverage) = parent(wta.result)

# This is called when output is requested.
function (wta::WindowedTimeAverage)(model)

    # For the paranoid
    wta.schedule.collecting &&
        model.clock.iteration > 0 &&
        @warn "Returning a WindowedTimeAverage before the collection period is complete."

    stride(wta) > 1 && @warn "WindowedTimeAverage can be erroneous when stride > 1 and either the timestep is variable or there are floating point rounding errors in times, both of which result in a decoupling of the model clock times (used in the OutputWriters) and iteration numbers (used for stride)."

    return wta.result
end

function accumulate_result!(wta, model)
    integrand = wta.fetch_operand ? fetch_output(wta.operand, model) : wta.operand
    return accumulate_result!(wta, model.clock, integrand)
end

period_to_number(p::Period) = value(p)
period_to_number(n::Number) = n

function accumulate_result!(wta, clock::Clock, integrand=wta.operand)
    # Time increment:
    Δt = period_to_number(clock.time - wta.previous_collection_time)
    # Time intervals:
    T_current = period_to_number(clock.time - wta.window_start_time)
    T_previous = period_to_number(wta.previous_collection_time - wta.window_start_time)

    # Accumulate left Riemann sum
    @. wta.result = (wta.result * T_previous + integrand * Δt) / T_current

    # Save time of integrand collection
    wta.previous_collection_time = clock.time

    return nothing
end

function advance_time_average!(wta::WindowedTimeAverage, model)

    unscheduled = model.clock.iteration == 0 || outside_window(wta.schedule, model.clock)
    if !(unscheduled)
        if !(wta.schedule.collecting)
            # Zero out result to begin new accumulation window
            wta.result .= 0

            # Begin collecting window-averaged increments
            wta.schedule.collecting = true

            wta.window_start_time = next_actuation_time(wta.schedule) - wta.schedule.window
            wta.previous_collection_time = wta.window_start_time
            wta.window_start_iteration = model.clock.iteration - 1
        end

        if end_of_window(wta.schedule, model.clock)
            accumulate_result!(wta, model)
            # Save averaging start time and the initial data collection time
            wta.schedule.collecting = false
            wta.schedule.actuations += 1

        elseif mod(model.clock.iteration - wta.window_start_iteration, stride(wta)) == 0
            accumulate_result!(wta, model)
        else
            # Off stride, so do nothing.
        end

    end
    return nothing
end

function advance_time_average!(wta::SpecifiedWindowedTimeAverage, model)

    unscheduled = model.clock.iteration == 0 || outside_window(wta.schedule, model.clock)
    if !(unscheduled)
        if !(wta.schedule.collecting)
            # Zero out result to begin new accumulation window
            wta.result .= 0

            # Begin collecting window-averaged increments
            wta.schedule.collecting = true

            wta.window_start_time = next_actuation_time(wta.schedule) - get_next_window(wta.schedule)
            wta.previous_collection_time = wta.window_start_time
            wta.window_start_iteration = model.clock.iteration - 1
        end

        if end_of_window(wta.schedule, model.clock)
            accumulate_result!(wta, model)
            # Save averaging start time and the initial data collection time
            wta.schedule.collecting = false
            wta.schedule.specified_times.previous_actuation += 1

        elseif mod(model.clock.iteration - wta.window_start_iteration, stride(wta)) == 0
            accumulate_result!(wta, model)
        else
            # Off stride, so do nothing.
        end

    end
    return nothing
end


# So it can be used as a Diagnostic
run_diagnostic!(wta::WindowedTimeAverage, model) = advance_time_average!(wta, model)

function prognostic_state(wta::WindowedTimeAverage)
    return (result = prognostic_state(wta.result),
            window_start_time = wta.window_start_time,
            window_start_iteration = wta.window_start_iteration,
            previous_collection_time = wta.previous_collection_time,
            schedule = prognostic_state(wta.schedule))
end

function restore_prognostic_state!(restored::WindowedTimeAverage, from)
    restore_prognostic_state!(restored.result, from.result)
    restored.window_start_time = from.window_start_time
    restored.window_start_iteration = from.window_start_iteration
    restored.previous_collection_time = from.previous_collection_time
    restore_prognostic_state!(restored.schedule, from.schedule)
    return restored
end

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

Base.show(io::IO, schedule::AveragedTimeInterval) = print(io, summary(schedule))

Base.summary(schedule::AveragedTimeInterval) = string("AveragedTimeInterval(",
                                                      "window=", prettytime(schedule.window), ", ",
                                                      "stride=", schedule.stride, ", ",
                                                      "interval=", prettytime(schedule.interval),  ")")

show_averaging_schedule(schedule) = ""
show_averaging_schedule(schedule::AveragedTimeInterval) = string(" averaged on ", summary(schedule))

output_averaging_schedule(output::WindowedTimeAverage) = output.schedule

#####
##### Utils for OutputWriters
#####

time_average_outputs(schedule, outputs, model) = schedule, outputs # fallback

const AveragedTimeSchedule = Union{AveragedTimeInterval, AveragedSpecifiedTimes}

"""
$(TYPEDSIGNATURES)

Wrap each `output` in a `WindowedTimeAverage` on the time-averaged `schedule` and with `field_slicer`.

Returns the `TimeInterval` associated with `schedule` and a `NamedTuple` or `Dict` of the wrapped
outputs.
"""
function time_average_outputs(schedule::AveragedTimeSchedule, outputs::Dict, model)
    averaged_outputs = Dict(name => WindowedTimeAverage(output, model; schedule=copy(schedule))
                            for (name, output) in outputs)

    return TimeInterval(schedule), averaged_outputs
end

function time_average_outputs(schedule::AveragedTimeSchedule, outputs::NamedTuple, model)
    averaged_outputs = NamedTuple(name => WindowedTimeAverage(outputs[name], model; schedule=copy(schedule))
                                  for name in keys(outputs))

    return TimeInterval(schedule), averaged_outputs
end
