export HaarKet, HilbertSchmidtStates, ChoiJamiolkowskiMatrices,
    HaarPOVM, WishartPOVM, VonNeumannPOVM


struct HaarKet{Î²} <: QIContinuousMatrixDistribution
    d::Int
end

HaarKet(d::Int) = HaarKet{2}(d)

function rand(rng::AbstractRNG, h::HaarKet{1})
    Ïˆ = randn(rng, h.d)
    renormalize!(Ïˆ)
    Ïˆ
 end

 function rand(rng::AbstractRNG, h::HaarKet{2})
     Ïˆ = randn(rng, h.d) + 1im * randn(rng, h.d)
     renormalize!(Ïˆ)
     Ïˆ
  end

# Random mixed states
struct HilbertSchmidtStates{Î², K} <: QIContinuousMatrixDistribution
    w::WishartEnsemble
    d::Int

    function HilbertSchmidtStates{Î², K}(d::Int) where {Î², K}
        w = WishartEnsemble{Î², K}(d)
        new(w, w.d)
    end
end
HilbertSchmidtStates{Î²}(d::Int) where Î² = HilbertSchmidtStates{Î², 1}(d)
HilbertSchmidtStates(d::Int) = HilbertSchmidtStates{2, 1}(d)

function rand(rng::AbstractRNG, hs::HilbertSchmidtStates{Î², K}) where {Î², K}
    Ï� = rand(rng, hs.w)
    renormalize!(Ï�)
    Ï�
end

#Random channels
struct ChoiJamiolkowskiMatrices{Î², K} <: QIContinuousMatrixDistribution
    w::WishartEnsemble
    idim::Int
    odim::Int

    function ChoiJamiolkowskiMatrices{Î², K}(idim::Int, odim::Int)  where {Î², K}
        w = WishartEnsemble{Î², K}(idim * odim)
        new(w, idim, odim)
    end
end

function ChoiJamiolkowskiMatrices{Î²}(idim::Int, odim::Int) where Î²
    ChoiJamiolkowskiMatrices{Î², 1}(idim, odim)
end

function ChoiJamiolkowskiMatrices{Î²}(d::Int) where Î²
    ChoiJamiolkowskiMatrices{Î²}(d, d)
end

function ChoiJamiolkowskiMatrices(idim::Int, odim::Int)
    ChoiJamiolkowskiMatrices{2}(idim, odim)
end

function ChoiJamiolkowskiMatrices(d::Int)
    ChoiJamiolkowskiMatrices(d, d)
end

function rand(rng::AbstractRNG, c::ChoiJamiolkowskiMatrices{Î², K}) where {Î², K}
    z = rand(rng, c.w)
    y = ptrace(z, [c.odim, c.idim], [1])
    sy = funcmh!(x -> 1 / sqrt(x), y)
    onesy = Matrix(I, c.odim, c.odim) âŠ— sy # onesy = eye(c.odim) âŠ— sy
    DynamicalMatrix(onesy * z * onesy, c.idim, c.odim)
end

# Random POVMs implemented according to
# https://arxiv.org/pdf/1902.04751.pdf
abstract type AbstractHaarPOVM <: QIContinuousMatrixDistribution
end

struct HaarPOVM{N} <: AbstractHaarPOVM
    idim::Int
    odim::Int
    c::HaarIsometry

    function HaarPOVM{N}(idim::Int, odim::Int) where N
        c = HaarIsometry(idim::Int, N*odim::Int)
        new(idim, odim, c)
    end
end
# N controls the rank (mixedness) of the effects, N=1 gives rank-one effects
HaarPOVM(idim::Int, odim::Int) = HaarPOVM{1}(idim, odim)

#this should use slicing of V
function rand(rng::AbstractRNG, c::HaarPOVM{N}) where N
    V = rand(rng, c.c)
    POVMMeasurement([V'*(ketbra(i, i, c.odim) âŠ— ð�•€(N))*V for i=1:c.odim])
end

struct VonNeumannPOVM <: AbstractHaarPOVM
    d::Int
    c::CUE

    function VonNeumannPOVM(d::Int)
        c = CUE(d)
        new(d, c)
    end
end

function rand(rng::AbstractRNG, c::VonNeumannPOVM)
    V = rand(rng, c.c)
    POVMMeasurement([proj(V[:, i]) for i=1:c.d])
end

struct WishartPOVM{V} <: QIContinuousMatrixDistribution
    idim::Int
    odim::Int
    c::Vector{WishartEnsemble}

    function WishartPOVM{V}(idim::Int) where V
        odim = length(V)
        c = [WishartEnsemble{2, v}(idim) for v=V]
        new(idim, odim, c)
    end
end

function WishartPOVM(idim::Int, odim::Int, K::Real=1)
    V = Tuple(round.(Int, K .* ones(odim)))
    WishartPOVM{V}(idim)
end

function rand(rng::AbstractRNG, c::WishartPOVM)
    Ws = map(x->rand(rng, x), c.c)
    S = sum(Ws)
    Ssq = funcmh!(x->1/sqrt(x), S)
    POVMMeasurement([Ssq * W * Ssq for W=Ws])
end