using ..PLRNNs
using ..ObservationModels

"""
    generate(model, zâ‚�, T, [S,])

Generate a trajectory of length `T` using model `â„³` given initial condition `zâ‚�`.

Returns a `T Ã— M` matrix of generated orbits in latent (model) space. If `S` is provided,
then `S` must be a `T' Ã— N` matrix of external inputs with `T' â‰¥ T`.
"""
function generate(â„³, zâ‚�::AbstractVector, T::Int)
    # trajectory placeholder
    Z = similar(zâ‚�, T, length(zâ‚�))

    # initial condition for model
    Z[1, :] .= zâ‚�

    # evolve initial condition in time
    @views for t = 2:T
        Z[t, :] .= â„³(Z[t-1, :])
    end
    return Z
end

function generate(â„³, zâ‚�::AbstractVector, T::Int, S::AbstractMatrix)
    # trajectory placeholder
    Z = similar(zâ‚�, T, length(zâ‚�))

    # initial condition for model
    Z[1, :] .= zâ‚�

    # evolve initial condition in time
    @views for t = 2:T
        Z[t, :] .= â„³(Z[t-1, :], S[t, :])
    end
    return Z
end

"""
    generate(â„³, ğ�’ª, xâ‚�, T, [S,])

Generate a trajectory of length `T` using latent model `â„³` and observation model
`ğ�’ª` given initial condition `xâ‚�` in observation space. Estimates latent state by inversion
of `ğ�’ª` and evolves latent state using `â„³`.

Returns a `T Ã— M` matrix of generated orbits in observation space. If `S` is provided,
then `S` must be a `T' Ã— N` matrix of external inputs with `T' â‰¥ T`.
"""
function generate(â„³, ğ�’ª::ObservationModel, xâ‚�::AbstractVector, T::Int)
    zâ‚� = init_state(ğ�’ª, xâ‚�)
    Z = generate(â„³, zâ‚�, T)
    return permutedims(ğ�’ª(Z'), (2, 1))
end

function generate(â„³, ğ�’ª::ObservationModel, xâ‚�::AbstractVector, T::Int, S::AbstractMatrix)
    zâ‚� = init_state(ğ�’ª, xâ‚�)
    Z = generate(â„³, zâ‚�, T, S)
    return permutedims(ğ�’ª(Z'), (2, 1))
end

keep_connectivity_offdiagonal!(m, g) = nothing
keep_connectivity_offdiagonal!(m::Union{AbstractVanillaPLRNN, AbstractDendriticPLRNN}, g) =
    offdiagonal!(g[m.W])