using MDPs
using BSON
using ImageShow
using StatsBase

struct GCHook <: AbstractHook
end

function MDPs.postepisode(::GCHook; kwargs...)
    GC.gc()
end

struct ModelsSaveHook <: AbstractHook
    models
    dirname::String
    interval::Int
end
function MDPs.postepisode(msh::ModelsSaveHook; returns, steps, kwargs...)
    if length(returns) == 1
        rm(msh.dirname, recursive=true, force=true)
    end
    if length(returns) % msh.interval == 0
        mkpath(msh.dirname)
        models = msh.models
        models = cpu(models)
        BSON.@save "$(msh.dirname)/ep-$(length(returns))-steps-$(steps).bson" models
    end
end
function MDPs.postexperiment(msh::ModelsSaveHook; returns, steps, kwargs...)
    mkpath(msh.dirname)
    models = msh.models
    models = cpu(models)
    BSON.@save "$(msh.dirname)/ep-$(length(returns))-steps-$(steps).bson" models
end

function loadmodels(filename)
    BSON.@load filename models
    return models
end


function smoothen(x::Vector{T}, window_size::Int) where {T <: AbstractFloat}
    y = similar(x)
    for i in 1:length(x)
        if i <= window_size
            y[i] = mean(x[1:i])
        else
            y[i] = mean(x[i-window_size+1:i])
        end
    end
    return y
end





struct SleepHook <: AbstractHook
    Δt::Float64  # seconds
end

function MDPs.postepisode(hook::SleepHook; kwargs...)
    sleep(hook.Δt)
    nothing
end

