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, Statevector, Operator

def getDensityMatrix(circuit):
    return DensityMatrix(circuit).data

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

from functools import reduce

Dag = lambda matrix: matrix.conj().T
Kron = lambda *matrices: reduce(np.kron, matrices)

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

# H R H, where R = 2|0><0| - I 
# A R A^(-1)

def initAACircuit(n):
    circuit = QuantumCircuit(n, n)
    
    circuit.x(0)
    circuit.h(1)
    circuit.z(1)
    circuit.h(0)
    circuit.barrier()
    
    return circuit

inputAACircuit = initAACircuit(2)
inputAACircuit.draw(output='mpl')

A = Operator(initAACircuit(2))

np.allclose(A.data @ Dag(A.data), np.eye(4))

inputAACircuit.measure([0, 1], [0, 1])

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

plot_histogram(counts, figsize=(7, 5), color="#CCCC66", title="measurements result of A")

def createOracle_3():
    circuit = QuantumCircuit(2, 2)

    # Oracle for find 3

    # U_f 
    circuit.cz(0, 1)
    circuit.barrier()
    
    return circuit

oracleCircuit_3 = createOracle_3()
oracleCircuit_3.draw(output='mpl')

def createR_2q():
    circuit = QuantumCircuit(2, 2)
    
    circuit.z(0)
    circuit.z(1)
    circuit.cz(0, 1)
    
    return circuit

def createAADiffuser():
    circuit = QuantumCircuit(2, 2)
    
    circuit.append(A, [0, 1])
    circuit = circuit.compose(createR_2q())
    circuit.append(A.conjugate(), [0, 1])
    circuit.barrier()
    
    return circuit

diffuserAACircuit = createAADiffuser()
diffuserAACircuit.draw(output='mpl')

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

aaGroverIteration = createGroverIteration(createOracle_3(), createAADiffuser())
aaGroverIteration.draw(output='mpl')

def createDiffuser_2q():
    circuit = QuantumCircuit(2, 2)
    
    circuit.h(0)
    circuit.h(1)
    circuit = circuit.compose(createR_2q())
    circuit.h(0)
    circuit.h(1)
    circuit.barrier()
    
    return circuit

incorrectGroverIteration = createGroverIteration(createOracle_3(), createDiffuser_2q())

grover_incorrect_diffuser = initAACircuit(2).compose(incorrectGroverIteration)
grover_incorrect_diffuser.draw(output='mpl')

grover_incorrect_diffuser.measure([0, 1], [0, 1])

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

plot_histogram(counts, figsize=(1, 5), color="#99CC33", title="using incorrect diffuser - find 3")

### 3.5.2 correct diffuser