using Flux

using ..PLRNNs
using ..Utilities
using ..ObservationModels

abstract type AbstractTFRecur end
Flux.trainable(tfrec::AbstractTFRecur) = (tfrec.model, tfrec.O)

(tfrec::AbstractTFRecur)(X::AbstractArray{T, 3}, ZÌ‚::AbstractArray{T, 3}) where {T} =
    forward(tfrec, X, ZÌ‚)

(tfrec::AbstractTFRecur)(
    X::AbstractArray{T, 3},
    ZÌ‚::AbstractArray{T, 3},
    S::AbstractArray{T, 3},
) where {T} = forward(tfrec, X, ZÌ‚, S)

# compute forcing signals using observation model
estimate_forcing_signals(tfrec::AbstractTFRecur, X::AbstractArray{T, 3}) where {T} =
    apply_inverse(tfrec.O, X)

"""
Inspired by `Flux.Recur` struct, which by default has no way
of incorporating teacher forcing.

This is just a convenience wrapper around stateful models,
to be used during training.
"""
mutable struct TFRecur{â„³ <: AbstractPLRNN, M <: AbstractMatrix, ð�’ª <: ObservationModel} <:
               AbstractTFRecur
    # stateful model, e.g. PLRNN
    model::â„³
    # observation model
    O::ð�’ª
    # state of the model
    z::M
    # forcing interval
    const Ï„::Int
end
Flux.@functor TFRecur

function forward(tfrec::TFRecur, X::AbstractArray{T, 3}, ZÌ‚::AbstractArray{T, 3}) where {T}
    T_ = size(X, 3)

    # initialize latent state
    tfrec.z = init_state(tfrec.O, @view(X[:, :, 1]))

    # process sequence of forcing signals ZÌ‚
    Z = [tfrec(zÌ‚, t) for (zÌ‚, t) âˆˆ zip(Flux.eachlastdim(ZÌ‚), 2:T_)]

    # reshape to 3d array and return
    return reshape(reduce(hcat, Z), size(tfrec.z)..., :)
end

function forward(
    tfrec::TFRecur,
    X::AbstractArray{T, 3},
    ZÌ‚::AbstractArray{T, 3},
    S::AbstractArray{T, 3},
) where {T}
    T_ = size(X, 3)

    # initialize latent state
    tfrec.z = init_state(tfrec.O, @view(X[:, :, 1]))

    # process sequence of forcing signals ZÌ‚
    Z = [
        tfrec(zÌ‚, s, t) for (zÌ‚, s, t) âˆˆ
        zip(Flux.eachlastdim(ZÌ‚), Flux.eachlastdim(@view(S[:, :, 2:end])), 2:T_)
    ]

    # reshape to 3d array and return
    return reshape(reduce(hcat, Z), size(tfrec.z)..., :)
end

function (tfrec::TFRecur)(x::AbstractMatrix, t::Int)
    # determine if it is time to force the model
    z = tfrec.z

    # perform one step using the model, update model state
    z = tfrec.model(z)

    # force
    zÌƒ = (t - 1) % tfrec.Ï„ == 0 ? force(z, x) : z
    tfrec.z = zÌƒ
    return z
end

function (tfrec::TFRecur)(x::AbstractMatrix, s::AbstractMatrix, t::Int)
    # determine if it is time to force the model
    z = tfrec.z

    # perform one step using the model, update model state
    z = tfrec.model(z, s)

    # force
    zÌƒ = (t - 1) % tfrec.Ï„ == 0 ? force(z, x) : z
    tfrec.z = zÌƒ
    return z
end

"""
    GTFRecur(model, O, z, Î±)

Generalized teacher forcing wrapper.
"""
mutable struct GTFRecur{
    â„³ <: AbstractPLRNN,
    ð�’ª <: ObservationModel,
    M <: AbstractMatrix,
    T <: AbstractFloat,
} <: AbstractTFRecur
    model::â„³
    # ObservationModel
    O::ð�’ª
    # state of the model
    z::M
    # forcing Î±
    Î±::T
end
Flux.@functor GTFRecur

"""
    forward(GTFRecur, X, ZÌ‚, [S,])

Forward pass using Generalized Teacher Forcing (GTF). If the latent dimension `M` of
the RNN is larger than the dimension of the provided forcing signals `ZÌ‚`, 
partial teacher forcing of the first `size(ZÌ‚, 1)` neurons is
used. Initializing latent state `zâ‚�` is performed by inversion of the observation model applied to data `X`.
Optionally, the forward pass incorporates external inputs `S`.

Returns a `M Ã— S Ã— T` array of latent states (predictions).
"""
function forward(tfrec::GTFRecur, X::AbstractArray{T, 3}, ZÌ‚::AbstractArray{T, 3}) where {T}
    # initialize latent state
    tfrec.z = init_state(tfrec.O, @view(X[:, :, 1]))

    # process sequence of forcing signals ZÌ‚
    Z = [tfrec(z) for z âˆˆ Flux.eachlastdim(ZÌ‚)]

    # reshape to 3d array and return
    return reshape(reduce(hcat, Z), size(tfrec.z)..., :)
end

function forward(
    tfrec::GTFRecur,
    X::AbstractArray{T, 3},
    ZÌ‚::AbstractArray{T, 3},
    S::AbstractArray{T, 3},
) where {T}
    # initialize latent state
    tfrec.z = init_state(tfrec.O, @view(X[:, :, 1]))

    # process sequence of forcing signals ZÌ‚
    Z = [
        tfrec(zÌ‚, s) for
        (zÌ‚, s) âˆˆ zip(Flux.eachlastdim(ZÌ‚), Flux.eachlastdim(@view(S[:, :, 2:end])))
    ]

    # reshape to 3d array and return
    return reshape(reduce(hcat, Z), size(tfrec.z)..., :)
end

function (tfrec::GTFRecur)(zÌ‚::AbstractMatrix)
    z = tfrec.z
    D, M = size(zÌ‚, 1), size(z, 1)
    z = tfrec.model(z)
    # gtf
    zÌƒ = force(@view(z[1:D, :]), zÌ‚, tfrec.Î±)
    zÌƒ = (D == M) ? zÌƒ : force(z, zÌƒ)

    tfrec.z = zÌƒ
    return z
end

function (tfrec::GTFRecur)(zÌ‚::AbstractMatrix, s::AbstractMatrix)
    z = tfrec.z
    D, M = size(zÌ‚, 1), size(z, 1)
    z = tfrec.model(z, s)
    # gtf
    zÌƒ = force(@view(z[1:D, :]), zÌ‚, tfrec.Î±)
    zÌƒ = (D == M) ? zÌƒ : force(z, zÌƒ)

    tfrec.z = zÌƒ
    return z
end