import qiskit
from qiskit import QuantumCircuit
from qiskit.visualization import plot_histogram

import numpy as np

from qiskit import execute, Aer

simulator = Aer.get_backend('qasm_simulator')

import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

from qiskit.quantum_info import DensityMatrix, Operator

# Quantum circuit to unitary matrix
def getDensityMatrix(circuit):
    return DensityMatrix(circuit).data

state_0 = np.array([1, 0]) # |0>
state_1 = np.array([0, 1]) # |1>

from functools import reduce

Dag = lambda matrix: matrix.conj().T # Dag(M) = M†
Kron = lambda *matrices: reduce(np.kron, matrices) # Kron(state_0, state_0) is |00>

def pm(matrix):
    for row in range(len(matrix)):
        for col in range (len(matrix[row])):
            print("{:.3f}".format(matrix[row][col]), end = " ")
        print()

# https://qiskit.org/textbook/ch-algorithms/grover.html

def initCircuit(n):
    circuit = QuantumCircuit(n, n)
    
    for i in range(n):
        circuit.h(i)    
    circuit.barrier()
    
    return circuit

inputCircuit_3q = initCircuit(3)
inputCircuit_3q.draw(output='mpl')

def createOracle_6():
    circuit = QuantumCircuit(3, 3)

    # Oracle for find 6

    # U_f 
    circuit.x(0)
    circuit.h(2)
    circuit.ccx(0, 1, 2)
    circuit.h(2)
    circuit.x(0)
    
    circuit.barrier()
    
    return circuit

oracleCircuit_6 = createOracle_6()
oracleCircuit_6.draw(output='mpl')

# where the entry that correspond to the marked item will have a negative phase
pm(Operator(oracleCircuit_6).data)

Operator(oracleCircuit_6).data

# H R H, where R = 2|0><0| - I 

def createR_3q():
    circuit = QuantumCircuit(3, 3)
    
    circuit.x(2)
    circuit.x(0)
    circuit.x(1)
    circuit.h(2)
    circuit.ccx(0, 1, 2)
    circuit.barrier(0)
    circuit.barrier(1)
    circuit.h(2)
    circuit.x(2)
    circuit.x(0)
    circuit.x(1)
    
    return circuit

R_3q = createR_3q()
R_3q.draw(output='mpl')

pm(Operator(R_3q).data)

print('|000>', Kron(state_0, state_0, state_0))
print('R|000>',Operator(R_3q).data @ Kron(state_0, state_0, state_0) )

print('|010>', Kron(state_0, state_1, state_0))
print('R|010>',Operator(R_3q).data @ Kron(state_0, state_1, state_0) )

print('|110>', Kron(state_1, state_1, state_0))
print('R|110>',Operator(R_3q).data @ Kron(state_1, state_1, state_0) )

# H R H

def createDiffuser_3q():
    circuit = QuantumCircuit(3, 3)
    
    circuit.h(0)
    circuit.h(1)
    circuit.h(2)
    circuit = circuit.compose(createR_3q())
    circuit.h(0)
    circuit.h(1)
    circuit.h(2)
    circuit.barrier()
    
    return circuit

diffuserCircuit_3q = createDiffuser_3q()
diffuserCircuit_3q.draw(output='mpl')

pm(Operator(diffuserCircuit_3q).data)

def createGroverIteration(oracle, diffuser):
    return oracle.compose(diffuser)

groverIteration_3q = createGroverIteration(createOracle_6(), createDiffuser_3q())
groverIteration_3q.draw(output='mpl')

groverIteration_3q = createGroverIteration(createOracle_6(), createDiffuser_3q())

inputCircuit_3q = initCircuit(3)
inputCircuit_3q.draw(output='mpl')

inputCircuit_3q.measure([0, 1, 2], [0, 1, 2])

job = execute(inputCircuit_3q, simulator, shots = 10000)
results = job.result()
counts = results.get_counts(inputCircuit_3q)
print(counts)

plot_histogram(counts, figsize=(10, 5), color="#FFFFCC", title="superposition state - 3 qubits")

grover_3q_1 = initCircuit(3).compose(groverIteration_3q.copy())
grover_3q_1.draw(output='mpl')

grover_3q_1.measure([0, 1, 2], [0, 1, 2])

job = execute(grover_3q_1, simulator, shots = 10000)
results = job.result()
counts = results.get_counts(grover_3q_1)
print(counts)

plot_histogram(counts, figsize=(10, 5), color="#FFBB00", title="one iteration - find 6")

grover_3q_2 = initCircuit(3).compose(groverIteration_3q.copy()).compose(groverIteration_3q.copy())
# grover_3q_2.draw(output='mpl')

grover_3q_2.measure([0, 1, 2], [0, 1, 2])

job = execute(grover_3q_2, simulator, shots = 10000)
results = job.result()
counts = results.get_counts(grover_3q_2)
print(counts)

plot_histogram(counts, figsize=(10, 5), color="#FF8822", title="two iteration - find 6")

grover_3q_3 = initCircuit(3).compose(groverIteration_3q.copy()).compose(groverIteration_3q.copy()).compose(groverIteration_3q.copy())
# grover_3q_3.draw(output='mpl')

grover_3q_3.measure([0, 1, 2], [0, 1, 2])

job = execute(grover_3q_3, simulator, shots = 10000)
results = job.result()
counts = results.get_counts(grover_3q_3)
print(counts)

plot_histogram(counts, figsize=(10, 5), color="#FFFF3F", title="three iteration - find 6")

