import Bijectors: inverse, VecCorrBijector, VecCholeskyBijector

"""
    softplus_inv(x::Real)

Returns the inverse of the softplus function such that: \\
`y = softplus(x)` \\
`x = softplus_inv(y)`
"""
softplus_inv(x::T) where {T<:Real} = log(exp(x) - one(T))

constrain_error(p) = (sigma = softplus.(p.error.sigma), )
function constrain_omega(p)
    Ï‰ = softplus.(p.omega.var) # TODO: rename this to sigma or similar, e.g. (prior = (omega = ..., corr = ...), )
    C = inverse(VecCorrBijector())(p.omega.corr)
    return Symmetric(Ï‰ .* C .* Ï‰')
end

"""
    constrain(p::NamedTuple)

Transforms the unconstrained parameter vector to constrained space.

# Examples
`Ïƒ* âˆˆ â„� â†’ softplus(Ïƒ*) âˆˆ â„�â�º` \\
`Ï‰, C  â†’ Ï‰ â‹… C â‹… Ï‰'`
"""
constrain(::O, p::NamedTuple) where O<:SSE = p
constrain(::O, p_::NamedTuple) where O<:LogLikelihood = (weights = p_.weights, error = constrain_error(p_))
constrain(::O, p_::NamedTuple) where O<:MixedObjective = (weights = p_.weights, error = constrain_error(p_), omega = constrain_omega(p_))

"""
    constrain_phi(::MeanField, ð�œ™::NamedTuple)

Transforms unconstrained `ð�œ™` to constrained space. For a MeanField approximation
this function returns `Î¼` and standard deviations `Ïƒ`.
"""
constrain_phi(::Type{MeanField}, ð�œ™::NamedTuple) = (mean = ð�œ™.mean, sigma = softplus.(ð�œ™.sigma))

sigma_corr_to_L(sigma, corr) = sigma .* inverse(VecCholeskyBijector(:L))(corr).L

"""
    constrain_phi(::FullRank, ð�œ™::NamedTuple)

Transforms unconstrained `ð�œ™` to constrained space. For a FullRank approximation
this function returns `Î¼` and the lower cholesky factor `L`.
"""
function constrain_phi(::Type{FullRank}, ð�œ™::NamedTuple)
    Ïƒ = softplus.(ð�œ™.sigma)
    L = sigma_corr_to_L.(eachcol(Ïƒ), eachcol(ð�œ™.corr))
    return (mean = ð�œ™.mean, L = L)
end
