# Gaussian Process State-Space Model (GP-SSM)

using LinearAlgebra
using Random
using AbstractGPs
using Plots
using Distributions
using SSMProblems

Parameters = @NamedTuple begin
a::Float64
q::Float64
kernel
end

mutable struct GPSSM <: SSMProblems.AbstractStateSpaceModel
X::Vector{Float64}
observations::Vector{Float64}
θ::Parameters

GPSSM(params::Parameters) = new(Vector{Float64}(), params)
GPSSM(y::Vector{Float64}, params::Parameters) = new(Vector{Float64}(), y, params)
end

seed = 1
T = 100
Nₚ = 20
Nₛ = 250
a = 0.9
q = 0.5

params = Parameters((a, q, SqExponentialKernel()))

f(θ::Parameters, x, t) = Normal(θ.a * x, θ.q)
h(θ::Parameters) = Normal(0, θ.q)
g(θ::Parameters, x, t) = Normal(0, exp(0.5 * x)^2)

rng = Random.MersenneTwister(seed)

x = zeros(T)
y = similar(x)
x[1] = rand(rng, h(params))
for t in 1:T
if t < T
x[t + 1] = rand(rng, f(params, x[t], t))
end
y[t] = rand(rng, g(params, x[t], t))
end

function gp_update(model::GPSSM, state, step)
gp = GP(model.θ.kernel)
prior = gp(1:(step - 1))
post = posterior(prior, model.X[1:(step - 1)])
μ, σ = mean_and_cov(post, [step])
return Normal(μ[1], σ[1])
end

SSMProblems.transition!!(rng::AbstractRNG, model::GPSSM) = rand(rng, h(model.θ))
function SSMProblems.transition!!(rng::AbstractRNG, model::GPSSM, state, step)
return rand(rng, gp_update(model, state, step))
end

function SSMProblems.emission_logdensity(model::GPSSM, state, step)
return logpdf(g(model.θ, state, step), model.observations[step])
end
function SSMProblems.transition_logdensity(model::GPSSM, prev_state, current_state, step)
return logpdf(gp_update(model, prev_state, step), current_state)
end

AdvancedPS.isdone(::GPSSM, step) = step > T

model = GPSSM(y, params)
plot!(mean_trajectory; color=:dodgerblue, label="Mean trajectory", opacity=0.9)