using QuantumInformation
using LinearAlgebra
using Statistics

steps = 100
haar = HaarKet(2)
Ïˆ = (ket(1, 4) + ket(4, 4))/sqrt(2)
Î³s = 0.0:0.01:1.0
Î¦(Î³) = KrausOperators([[1 0; 0 sqrt(1-Î³)], [0 sqrt(Î³); 0 0]])
post = [PostSelectionMeasurement(proj(ket(i, 4)) âŠ— ğ�•€(2)) for i=1:4]
rots = [UnitaryChannel(ğ�•€(2)), UnitaryChannel(sx), UnitaryChannel(sz), UnitaryChannel(sx*sz)]
had = UnitaryChannel{Matrix{ComplexF64}}(hadamard(2))
cnot = UnitaryChannel{Matrix{ComplexF64}}([1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0])
r = zeros(steps, length(Î³s), 4)
for (k, Î³) in enumerate(Î³s)
    for i = 1:steps
        Ï• = rand(haar)
        Î¾ = Ï• âŠ— Ïˆ
        Ï� = ((had âŠ— IdentityChannel(4))âˆ˜(cnot âŠ— IdentityChannel(2))âˆ˜(IdentityChannel(4) âŠ— Î¦(Î³)))(Î¾)
        for j=1:4
            Ïƒ = rots[j](ptrace(post[j](Ï�), [2, 2, 2], [1, 2]))
            r[i, k, j] = fidelity(Ï•, Ïƒ/tr(Ïƒ))
        end
    end
end
mean(r, dims=1)
