abstract type AbstractGPLayer <: AbstractGP end;

const jitt = 1e-5

struct SVGPLayer{GPType,TV<:AbstractVector{<:GPType}} <: AbstractGPLayer
    dim::Int
    gps::TV
end

function SVGPLayer(
    dim::Int, μ::AbstractVector, L::AbstractMatrix, Z::AbstractVector, kernel::Kernel, μ₀
)
    gps = [
        SVGP_Base(length(μ), copy(μ), copy(L), GP(deepcopy(μ₀), deepcopy(kernel)), deepcopy(Z)) for _ in 1:dim
    ]
    return SVGPLayer(dim, gps)
end

Base.length(l::SVGPLayer) = l.dim
Base.getindex(l::SVGPLayer, i::Int) = getindex(l.gps, i)
Base.iterate(l::SVGPLayer) = iterate(l.gps)
Base.iterate(l::SVGPLayer, state::Any) = iterate(l.gps, state)

bsqrt(v::AbstractVector) = sqrt.(v)

## Return the parametrization m + eps * sqrt(S)
function propagate(l::AbstractGPLayer, samples::AbstractVector)
    m_v_s = [mean_and_var(gp, samples) for gp in l]
    out =
        mapreduce(first, hcat, m_v_s) +
        randn(length(samples), l.dim) .* mapreduce(bsqrt ∘ last, hcat, m_v_s)
    return RowVecs(out)
end

function propagate(l::AbstractGPLayer, sample::AbstractArray{<:Real})
    m_v_s = [mean_and_var(gp, sample) for gp in l]
    return first(m_v_s) .+ randn(l.dim) .* sqrt.(last(m_v_s))
end

# mu(l::AbstractGPLayer,x) = reduce(hcat,mu(gp,x) for gp in l.gps)
function Statistics.mean(l::AbstractGPLayer, x::AbstractVector)
    return mapreduce(Base.Fix2(mean, x), hcat, l)
end
Statistics.var(l::AbstractGPLayer, x) = mapreduce(Base.Fix2(var, x), hcat, l)
Statistics.cov(l::AbstractGPLayer, x) = cov.(l, Ref(x))

kl_divergence(l::AbstractGPLayer) = sum(kl_divergence, l)

abstract type AbstractGPBase end;

"""
    SVGP_Base

Corresponds to one node in the Deep GP
"""
struct SVGP_Base{
    T,Tμ<:AbstractVector{T},TL<:AbstractMatrix{T},Tgp<:AbstractGP,TZ<:AbstractVector
} <: AbstractGPs.AbstractGP
    dim::Int
    μ::Tμ
    L::TL
    prior::Tgp
    Z::TZ
end

function SVGP_Base{T}(
    dim, Z::AbstractVector, kernel=SqExponentialKernel(), μ₀=AbstractGPs.ZeroMean()
) where {T}
    return SVGP_Base(dim, zeros(T, dim), diagm(ones(T, dim)), GP(μ₀, kernel), Z)
end

function ParameterHandling.flatten(x::SVGP_Base{T}) where {T}
    θprior, build_prior = flatten(x.prior)
    θZ, build_Z = flatten(x.Z)
    nZ = length(θZ)
    θ = vcat(x.μ, vec(x.L), θZ, θp)
    function unflatten_svgp_base(θ)
        μ = θ[1:x.dim]
        
        SVGP_Base{T}(x.dim, μ)
    end
end


"""
    kappa(gp::SVGP_Base, x::AbstractVector)

"""
function kappa(gp::SVGP_Base, x::AbstractVector, K::AbstractMatrix=cov(gp.prior, gp.Z))
    ## Return the K_XZ K_ZZ^{-1}
    return cov(gp.prior, x, gp.Z) / K
end

function Statistics.mean(
    gp::SVGP_Base,
    x::AbstractVector,
    K::AbstractMatrix=cov(gp.prior, gp.Z),
    κ::AbstractMatrix=kappa(gp, x, K),
)
    return κ * gp.μ - mean(gp.prior, x)
end

function Statistics.var(
    gp::SVGP_Base,
    x::AbstractVector,
    K::AbstractMatrix=cov(gp.prior, gp.Z),
    κ::AbstractMatrix=kappa(gp, x, K),
)
    return var(gp.prior, x) - opt_diag(κ * (K - gp.L * gp.L'), κ)
end
function Statistics.cov(
    gp::SVGP_Base,
    x::AbstractVector,
    K::AbstractMatrix=cov(gp.prior, gp.Z),
    κ::AbstractMatrix=kappa(gp, x, K),
)
    return cov(gp.prior, x) - κ * (K - gp.L * gp.L') * κ'
end
function StatsBase.mean_and_var(gp::SVGP_Base, x::AbstractVector)
    K = cov(gp.prior, gp.Z)
    κ = kappa(gp, x, K)
    return mean(gp, x, K, κ), var(gp, x, K, κ)
end

function kl_divergence(gp::SVGP_Base)
    K = cov(gp.prior, gp.Z)
    return 0.5 * (
        -2logdet(gp.L) +
        logdet(K) +
        opt_trace(inv(K), gp.L * gp.L') +
        dot(gp.μ, K \ (gp.μ)) - length(gp.μ)
    ) #WARNING removed μ₀
end #WARNING removed μ₀
