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]) # |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)

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_2q = initCircuit(2)
inputCircuit_2q.draw(output='mpl')

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')

O_3 = Operator(oracleCircuit_3).data
pm(O_3)

# all possible states
state_00 = Kron(state_0, state_0)
state_01 = Kron(state_0, state_1)
state_10 = Kron(state_1, state_0)
state_11 = Kron(state_1, state_1)

print("O|00>", O_3 @ state_00)
print("O|01>", O_3 @ state_01)
print("O|10>", O_3 @ state_10)
print("O|11>", O_3 @ state_11) # flip phase

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

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

R_2q = createR_2q()
R_2q.draw(output='mpl')

R_2q = Operator(R_2q).data
pm(R_2q)

# flip phase(no work on zero states) => Conditional Phase Shift gate
print("RO|00>", R_2q @ O_3 @ state_00) # zero state => would not flip
print("RO|01>", R_2q @ O_3 @ state_01)
print("RO|10>", R_2q @ O_3 @ state_10)
print("RO|11>", R_2q @ O_3 @ state_11)

# H R H

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

diffuserCircuit_2q = createDiffuser_2q()
diffuserCircuit_2q.draw(output='mpl')

diff_2q = Operator(diffuserCircuit_2q).data
pm(diff_2q)

DB = 0.5 * state_00 + 0.5 * state_01 + 0.5 * state_10 + 0.5 * state_11
print("DO|DB>", diff_2q @ O_3 @ DB)

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

groverIteration_2q = createGroverIteration(createOracle_3(), createDiffuser_2q())
groverIteration_2q.draw(output='mpl')

groverIteration_2q = createGroverIteration(createOracle_3(), createDiffuser_2q())

grover_2q_1 = initCircuit(2).compose(groverIteration_2q.copy())
grover_2q_1.draw(output='mpl')

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

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

plot_histogram(counts, figsize=(1, 5), color="#CC3333", title="one iteration - find 3")

grover_2q_2 = initCircuit(2).compose(groverIteration_2q.copy()).compose(groverIteration_2q.copy())
grover_2q_2.draw(output='mpl')

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

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

plot_histogram(counts, figsize=(7, 5), color="#FF9999", title="two iterations - find 3")

