struct WarpedGauss{T<:Real} <: ContinuousMultivariateDistribution
    "Standard deviation of the first dimension, must be > 0"
    σ1::T
    "Standard deviation of the second dimension, must be > 0"
    σ2::T
    function WarpedGauss{T}(σ1, σ2) where {T<:Real}
        σ1 > 0 || error("σ₁ must be > 0")
        σ2 > 0 || error("σ₂ must be > 0")
        return new{T}(σ1, σ2)
    end
end
WarpedGauss(σ1::T, σ2::T) where {T<:Real} = WarpedGauss{T}(σ1, σ2)
WarpedGauss() = WarpedGauss(1.0, 0.12)

Base.length(p::WarpedGauss) = 2
Base.eltype(::WarpedGauss{T}) where {T<:Real} = T
Distributions.sampler(p::WarpedGauss) = p

# Define the transformation function φ and the inverse ϕ⁻¹ for the warped Gaussian distribution
function ϕ!(::WarpedGauss, z::AbstractVector)
    length(z) == 2 || error("Dimension mismatch")
    x, y = z
    r = norm(z)
    θ = atan(y, x) #in [-π , π]
    θ -= r / 2
    z .= r .* [cos(θ), sin(θ)]
    return z
end

function ϕ⁻¹(::WarpedGauss, z::AbstractVector)
    length(z) == 2 || error("Dimension mismatch")
    x, y = z
    r = norm(z)
    θ = atan(y, x) #in [-π , π]
    # increase θ depending on r to "smear"
    θ += r / 2

    # get the x,y coordinates foαtransformed point
    xn = r * cos(θ)
    yn = r * sin(θ)
    # compute jacobian
    logJ = log(r)
    return [xn, yn], logJ
end

function Distributions._rand!(rng::AbstractRNG, p::WarpedGauss, x::AbstractVecOrMat)
    size(x, 1) == 2 || error("Dimension mismatch")
    σ₁, σ₂ = p.σ1, p.σ2
    randn!(rng, x)
    x .*= [σ₁, σ₂]
    for y in eachcol(x)
        ϕ!(p, y)
    end
    return x
end

function Distributions._logpdf(p::WarpedGauss, x::AbstractVector)
    size(x, 1) == 2 || error("Dimension mismatch")
    σ₁, σ₂ = p.σ1, p.σ2
    S = [σ₁, σ₂] .^ 2
    z, logJ = ϕ⁻¹(p, x)
    return -sum(z .^ 2 ./ S) / 2 - IrrationalConstants.log2π - log(σ₁) - log(σ₂) + logJ
end
