module SequenceGenerator

import Flux

# This import is only for documentation reference purposes.
import ..LevelFormatter
using ..InputStatistics
using ..ModelUtils: LearningModel, togpu, tocpu
using ..LevelStatistics

export generatesequence

# TODO more `view`s
# TODO Take `reverse_rows` into account!
# TODO and `each_tile`

"""
    generatesequence(model::LearningModel, initialinput::AbstractArray)

Return a sequence generated by applying the given model to the given first screen or input.
An initial screen is given as an `AbstractVector{AbstractVector}` of model inputs while a
single first input is any other `AbstractArray`.

The input will automatically be converted to the correct form in the 1D case.
"""
function generatesequence(model::LearningModel, initialinput::AbstractArray)
    isempty(initialinput) && return (), []
    constantinput = initialinput[eachindex(initialinput)[1:constantinputsize - 1]]
    # Strip constant input part.
    sequence = [initialinput[eachindex(initialinput)[constantinputsize:end]]]
    Flux.reset!(model)
    Flux.testmode!(model)
    # While the "sequence has not ended yet" bit is not set...
    while sequence[end][1] != 0 && length(sequence) < LevelStatistics.maxcolshori
        push!(sequence, tocpu(Flux.data(
            model(togpu(vcat(constantinput, sequence[end]))))))
    end
    if sequence[end][1] != 0
        println("Force stopped generation due to maximum level length.")
    end
    Flux.testmode!(model, false)
    return constantinput, postprocess(sequence, model.hyperparams[:dimensionality])
end

function generatesequence(model::LearningModel, initialinput::AbstractMatrix)
    isempty(initialinput) && return (), empty(initialinput)
    constantinputs = initialinput[1:constantinputsize - 1, :]
    constantinput = constantinputs[:, 1]
    # TODO preallocate large sequence instead of growing
    sequence = convert(Matrix, initialinput)
    prediction = sequence[constantinputsize:end, :]
    Flux.reset!(model)
    Flux.testmode!(model)
    # While the "sequence has not ended yet" bit is not set...
    while (prediction[1, end] != 0
           && size(prediction, 2) < LevelStatistics.maxcolshori)
        prediction = tocpu(Flux.data(model(togpu(sequence))))
        sequence = hcat(sequence, vcat(constantinput, prediction[:, end]))
    end
    if prediction[1, end] != 0
        println("Force stopped generation due to maximum level length.")
        sequence = sequence[:, 1:end - 1]
    end
    Flux.testmode!(model, false)
    return constantinput, postprocess(sequence[constantinputsize:end, :],
                                      model.hyperparams[:dimensionality])
end

"""
    generatesequence(model::LearningModel,
                     initialscreen::AbstractVector{<:AbstractVector{<:Number})

Return a sequence generated by applying the given model to the given first screen.
"""
function generatesequence(model::LearningModel, initialscreen::AbstractVector{T}
                          ) where {T <: AbstractVector{<:Number}}
    isempty(initialscreen) && return [], []
    initialinput = first(initialscreen)
    constantinput = initialinput[eachindex(initialinput)[1:constantinputsize - 1]]
    sequence = [inputpart[eachindex(inputpart)[constantinputsize:end]]
                for inputpart in initialscreen]
    Flux.reset!(model)
    Flux.testmode!(model)
    # Give the model the initial inputs, pre-tuning it (and disregard the predictions.)
    model.(initialscreen[1:end - 1])
    while sequence[end][1] != 0 && length(sequence) < LevelStatistics.maxcolshori
        push!(sequence, tocpu(Flux.data(
            model(togpu(vcat(constantinput, sequence[end]))))))
    end
    if sequence[end][1] != 0
        println("Force stopped generation due to maximum level length.")
    end
    Flux.testmode!(model, false)
    return constantinput, postprocess(sequence, model.hyperparams[:dimensionality])
end

"""
Return the given sequence's elements as one concatenated `Array`.
The dimensionality of the returned `Array` is determined by the argument of the same name
which allows any `Symbol` in [`LevelFormatter.dimensionality_defaultflags`](@ref).
"""
function postprocess(sequence::AbstractVector{T},
                     dimensionality::Symbol) where {T <: AbstractVector{Float32}}
    # TODO We could just do `x -> reshape(x, screenrowshori, 1, :)` for most cases.
    # TODO Refactor code so the `fromXd` functions are applied here (need more args then).
    # Remove the "has-not-ended"-bit.
    map!(x -> x[2:end], sequence, sequence)
    if     dimensionality === Symbol("1d")
        return reduce(vcat, sequence)
    elseif dimensionality === Symbol("2d")
        return reduce(hcat, sequence)
    # Explicit code for safety:
    # elseif dimensionality === Symbol("3dtiles")
    #     return reduce(hcat, map(x -> reshape(x, (convert(Int, screenrowshori), 1,
    #             convert(Int, uniquevanillatiles))), sequence))
    # elseif dimensionality === Symbol("3d")
    #     return reduce(hcat, map(x -> reshape(x, (convert(Int, screenrowshori), 1,
    #                                              convert(Int, layers3d))), sequence))
    # Implicit code for extensibility:
    elseif String(dimensionality)[1:2] == "3d"
        return reduce(hcat, map(x -> reshape(x, (convert(Int, screenrowshori), 1,
                                                 :)), sequence))
    else
        throw(ValueError("unknown dimensionality"))
    end
end

"""
Return the given sequence's elements as one concatenated `Array`.
The dimensionality of the returned `Array` is determined by the argument of the same name
which allows any `Symbol` in [`LevelFormatter.dimensionality_defaultflags`](@ref).
"""
function postprocess(sequence::AbstractMatrix{Float32}, dimensionality::Symbol)
    # Remove the "has-not-ended"-bit.
    sequence = sequence[2:end, :]
    if     dimensionality === Symbol("1d")
        return vec(sequence)
    elseif dimensionality === Symbol("2d")
        return sequence
    # Explicit code for safety:
    # elseif dimensionality === Symbol("3dtiles")
    #     return reduce(hcat, map(x -> reshape(x, (convert(Int, screenrowshori), 1,
    #             convert(Int, uniquevanillatiles))), sequence))
    # elseif dimensionality === Symbol("3d")
    #     return reduce(hcat, map(x -> reshape(x, (convert(Int, screenrowshori), 1,
    #                                              convert(Int, layers3d))), sequence))
    # Implicit code for extensibility:
    elseif String(dimensionality)[1:2] == "3d"
        return reduce(hcat, map(x -> reshape(x, (convert(Int, screenrowshori), 1,
                                                 :)), eachcol(sequence)))
    else
        throw(ValueError("unknown dimensionality"))
    end
end

end # module

