begin
	using Pkg
	Pkg.activate(mktempdir())
	Pkg.Registry.update()
	Pkg.add("Yao")
	Pkg.add("YaoPlots")
	Pkg.add("StatsBase")
	Pkg.add("Plots")
	Pkg.add("BitBasis")
end

using Yao, YaoPlots

# make the bell circuit
bellcircuit = chain(2, put(1 => H), control(1, 2 => X))

# make the reverse bell circuit
reversebellcircuit = chain(2, control(1, 2 => X), put(1 => H))

# circuit that takes two qubits and passes it through 
# the bell circuit then the reverse bell circuit
circuit = chain(2, bellcircuit, reversebellcircuit)
# plot(circuit)

# make the qubits |00>, |01>, |10>, |11>
qubit00 = ArrayReg(bit"00")
qubit10 = ArrayReg(bit"10") # circuit reads in reverse order
qubit01 = ArrayReg(bit"01")
qubit11 = ArrayReg(bit"11")

using StatsBase: Histogram, fit
using Plots: bar, scatter!, gr;
gr();
using BitBasis
function plotmeasure(x::Array{BitStr{n, Int}, 1}) where n
	hist = fit(Histogram, Int.(x), 0:2^n)
	x = 0
	if (n <= 3)
		s = 8
	elseif (n > 3 && n <= 6)
		s = 5
	elseif (n > 6 && n <= 10)
		s = 3.2
	elseif (n > 10 && n <= 15)
		s = 2
	elseif (n > 15)
		s = 1
	end
	bar(hist.edges[1] .- 0.5, hist.weights, legend = :none, size = (600 * (2^n) / s, 400), ylims = (0:maximum(hist.weights)), xlims = (0:2^n), grid = :false, ticks = false, border = :none, color = :lightblue, lc = :lightblue)
	scatter!(0:2^n-1, ones(2^n, 1), markersize = 0,
		series_annotations = "|" .* string.(hist.edges[1]; base = 2, pad = n) .* "⟩")
	scatter!(0:2^n-1, zeros(2^n, 1) .+ maximum(hist.weights), markersize = 0,
		series_annotations = string.(hist.weights))
end

# pass them through the circuit 1024 times, taking measurements
println("Qubit 00")
measured_qubits00 = qubit00 |> circuit -> measure(circuit, nshots = 1024)
println(measured_qubits00)
plotmeasure(measured_qubits00)

println("Qubit 01")
measured_qubits01 = qubit01 |> circuit -> measure(circuit, nshots = 1024)
println(measured_qubits01)
plotmeasure(measured_qubits01)

println("Qubit 10")
measured_qubits10 = qubit10 |> circuit -> measure(circuit, nshots = 1024)
println(measured_qubits10)
plotmeasure(measured_qubits10)

println("Qubit 11")
measured_qubits11 = qubit11 |> circuit -> measure(circuit, nshots = 1024)
println(measured_qubits11)
plotmeasure(measured_qubits11)