import numpy as np
from qiskit import QuantumCircuit, execute, Aer
from qiskit.quantum_info import Statevector

import kaleidoscope.qiskit
from kaleidoscope import qsphere, probability_distribution

# import basic plot tools
from qiskit.visualization import plot_histogram, plot_state_qsphere


# A Circuit is something as simple as a Qubit with a 0 or a 1.

qc = QuantumCircuit(1)
state = Statevector.from_instruction(qc)
qsphere(state)

# Or a 1
qc.x(0)
state = Statevector.from_instruction(qc)
qsphere(state)

# Now we simulate 3 Qubits
qc = QuantumCircuit(3)
qc.h(range(3))
qc.cz(0,1)
state = Statevector.from_instruction(qc)
qsphere(state)

qc.draw()

# Representing 5  Qubits

qc2 = QuantumCircuit(5)
qc2.h(range(4))
qc2.x(4)
qc2.h(4)
#qc2.barrier()

# "Mark" |11011>
qc2.x(2)
qc2.mct([0,1,2,3], 4)
qc2.x(2)

# "Mark" |11110>
#qc2.x(0)
#qc2.mct([0,1,2,3], 4)
#qc2.x(0)
qc2.h(4)
#qc2.barrier()
qc2.draw('mpl')

statevector = Statevector.from_instruction(qc2)
qsphere(statevector)

qc2.measure_all()
backend = Aer.get_backend("qasm_simulator")
result = execute(qc2, backend=backend, shots=1000).result()
probability_distribution(result.get_counts(qc2))

# Very simple circuit explanation

qc3 = QuantumCircuit(3)

# 1: No Gates
qc3.i(range(3))

# 2: Some Gates
#qc3.x(0)
#qc3.x(1)

# 3: X on all gates
#qc3.x(range(3))

# 4: Superposition
#qc3.x(range(3))

# 5: With some X as well
#qc3.h(range(3))

qc3.draw('mpl')

statevector = Statevector.from_instruction(qc3)
qsphere(statevector)

qc3.measure_all()
backend = Aer.get_backend("qasm_simulator")
result = execute(qc3, backend=backend, shots=1000).result()

probability_distribution(result.get_counts(qc3))

### NEW CIRCUIT WITH 3 QUBITS 
### SOLUTIONS: 101 AND 110

def phase_oracle(circuit):
    circuit.cz(0, 2)
    circuit.cz(1, 2)
    
    
def diffuser(circuit):
    """Apply inversion about the average step of Grover's algorithm."""
    qubits = circuit.qubits
    nqubits = len(qubits)
    
    for q in range(nqubits):
        circuit.h(q)
        circuit.x(q)
    
    # Do controlled-Z
    circuit.h(2)
    circuit.ccx(0,1,2)
    circuit.h(2)
    
    for q in range(nqubits):
        circuit.x(q)
        circuit.h(q)

## First we create superposition

n = 3
grover_circuit = QuantumCircuit(n)
grover_circuit.h(range(n))
grover_circuit.barrier()

grover_circuit.draw(output="mpl")

statevector = Statevector.from_instruction(grover_circuit)
qsphere(statevector)

# The Quantum Computer "marks" the solutions
phase_oracle(grover_circuit)

grover_circuit.draw(output="mpl")

statevector = Statevector.from_instruction(grover_circuit)
qsphere(statevector)

# User interference to keep only the solution states

grover_circuit.barrier()
diffuser(grover_circuit)

grover_circuit.draw(output="mpl")

statevector = Statevector.from_instruction(grover_circuit)
qsphere(statevector)

grover_circuit.measure_all()
backend = Aer.get_backend('qasm_simulator')
results = execute(grover_circuit, backend=backend, shots=1000).result()
probability_distribution(results.get_counts(grover_circuit))
