import warnings; warnings.simplefilter('ignore')

# Importing the necessary libraries
import numpy as np
from numpy import pi
import matplotlib.pyplot as plt

# Importing standard Qiskit libraries
from qiskit import *
from qiskit.quantum_info import Statevector
from qiskit.tools.jupyter import *
from qiskit.visualization import *

#we have 2 system qubits
system_qubits = 2 
#following the argument above we need 2 additional qubits
nqubits = system_qubits + 2
q = QuantumRegister(nqubits, 'q')
circuit = QuantumCircuit(q)  

x = np.linspace(-1, 1)
test = np.polynomial.Chebyshev((0, 0, 0, 1)) # The 3rd Chebyshev polynomial!
y = test(x)
plt.plot(x, y)

# Define projectors.
# Note that p_left is actually p_right from above because the order
# of operations are reversed when equations are turned into circuits
# due to how time-flow is defined in circuit structures

# In the Qiskit implementation qubit indexes start from 0
# and the most significant qubit is the highest index
# keeping this in mind e.g for 4 nqubits = {q0,q1,q2,q3}
# q0 and q1 are the system qubits
# q2 is the signal qubit
# q3 is the ancillary rotation qubit
# nqubits-1 = 4-1 = 3 below, then corresponds to q3

def p_left(q, phi): #right projector
    qc = QuantumCircuit(q)
    n = q
    ctrl_range = list(range(0,n-1))
    
    for qubit in range(n-1): # Implement a simple multi 0-controlled
        qc.x(qubit)
    qc.mcx(ctrl_range , n-1) # 0-Controlled on all but the last qubits, acts on the last qubit
    for qubit in range(n-1):
        qc.x(qubit)
        
    #qc.barrier(0, 1, 2, 3)
    qc.rz(phi, n-1) # RZ(phi) on the last qubit
    #qc.barrier(0, 1, 2, 3)
    
    for qubit in range(n-1): # Reverse the effect of the first multi-control
        qc.x(qubit)
    qc.mcx(ctrl_range ,n-1) 
    for qubit in range(n-1):
        qc.x(qubit)
    
    p_left_gate = qc.to_gate() # Compiles all this into a gate
    p_left_gate.name = "P$_l$(Φ)"
    return p_left_gate

def p_right(phi): # Left projector acts just on the signal and the ancillary qubit
    qc = QuantumCircuit(2)
    
    qc.cx(0, 1)
    qc.rz(phi, 1)
    qc.cx(0 ,1)
    
    p_right_gate = qc.to_gate()
    p_right_gate.name = "P$_r$(Φ)"
    return p_right_gate

#Define Oracle and the reverse-gate for 
#constructing the dagger later

def U(q): 
    qc = QuantumCircuit(q)
    n = q + 1
    
    for qubit in range(n-2):
        qc.h(qubit)
    
    qc.mcx(list(range(0,n-2)), n-2)
    
    U_gate = qc.to_gate()
    U_gate.name = "U"
    return U_gate

def reverse_gate(gate):
    gate_rev = gate.reverse_ops()
    gate_rev.name = gate.name + "$^†$"
    return gate_rev

#we have 2 system qubits
system_qubits = 2 
#following the argument above we need 2 additional qubits
nqubits = system_qubits + 2
q = QuantumRegister(nqubits, 'q')
circuit = QuantumCircuit(q)

d = 3

u = U(nqubits-1)
u_dag = reverse_gate(u) #construct U_dagger

p_right_range = [nqubits-2, nqubits-1] # build the ranges of qubits for gates to act upon
u_range = list(range(0, nqubits-1))
p_left_range = list(range(0, nqubits))

circuit.append(p_left(nqubits, (1-d)*pi), p_left_range) # in general, starting from this line, 
circuit.append(u, u_range)                              # the circuit would iterate over the phases,
                                                        # but the phases for pure Cheby are just trivial pi/2
                                                        # so we chose to directly putting pi's in here for simplification
for i in range((d-1)//2):                               # we put pi instead of pi/2 because of how RZ is defined in Qiskit
    circuit.append(p_right(pi), p_right_range) 
    circuit.append(u_dag, u_range)
    circuit.append(p_left(nqubits, pi), p_left_range)
    circuit.append(u, u_range)

circuit.measure_all()
circuit.draw('mpl')

qasm_sim = Aer.get_backend('qasm_simulator')
transpiled_circuit = transpile(circuit, qasm_sim)
qobj = assemble(transpiled_circuit)
results = qasm_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

# We have 2 system qubits
system_qubits = 2 
# Following the argument above we need 2 additional qubits
nqubits = system_qubits + 2
q = QuantumRegister(nqubits, 'q')
circuit = QuantumCircuit(q)
############################## d in terms of n!
d = (2*system_qubits) - 1

if system_qubits > 6 and system_qubits < 10:
    for i in range(1, system_qubits - 6 + 1):
        d += 2 * i
###############################

u = U(nqubits-1)
u_dag = reverse_gate(u)

p_right_range = [nqubits-2, nqubits-1]
u_range = list(range(0, nqubits-1))
p_left_range = list(range(0, nqubits))

circuit.append(p_left(nqubits, (1-d)*pi), p_left_range)
circuit.append(U(nqubits-1), u_range)

for i in range((d-1)//2):
    circuit.append(p_right(pi), p_right_range) #debug this, doesnt work just as a 2 qubit gate
    circuit.append(u_dag, u_range)
    circuit.append(p_left(nqubits,pi), p_left_range)
    circuit.append(u, u_range)

circuit.measure_all()
circuit.draw('mpl')

def qsvt_search(target): # target = marked element, is a bit-string!
    
    systemqubits = len(target)
    nqubits = systemqubits + 2
    q = QuantumRegister(nqubits, 'q')
    circuit = QuantumCircuit(q)  

    d = (2*systemqubits) - 1

    if systemqubits > 6 and systemqubits < 10:
        for i in range(1, systemqubits - 6 + 1):
            d += 2 * i
    
    u = U(nqubits-1)
    u_dag = reverse_gate(u)
    
    p_right_range = [nqubits-2, nqubits-1]
    u_range = list(range(0, nqubits-1))
    p_left_range = list(range(0, nqubits))

    circuit.append(p_left(nqubits,(1-d)*pi), p_left_range)
    circuit.append(U(nqubits-1), u_range)

    for i in range((d-1)//2):
        circuit.append(p_right(pi), p_right_range) #debug this, doesnt work just as a 2 qubit gate
        circuit.append(u_dag, u_range)
        circuit.append(p_left(nqubits,pi), p_left_range)
        circuit.append(u, u_range)

    for i in range(len(target)): # The operation for acquiring arbitrary marked element
        bts = target [::-1]      # bitstring is reversed to be compatible with the reverse qubit order in Qiskit
        if bts[i] == '0':
            circuit.x(i)

    circuit.measure_all()
    return circuit

circuit_test = qsvt_search('1101')
circuit_test.draw('mpl')

qasm_sim = Aer.get_backend('qasm_simulator')
transpiled_circuit = transpile(circuit_test, qasm_sim)
qobj = assemble(transpiled_circuit)
results = qasm_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

qc = qsvt_search('001')
qc.draw('mpl')

qasm_sim = Aer.get_backend('qasm_simulator')
transpiled_circuit = transpile(qc, qasm_sim)
qobj = assemble(transpiled_circuit)
results = qasm_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

# We have 4 system qubits
system_qubits = 4 
nqubits = system_qubits + 2
q = QuantumRegister(nqubits, 'q')
circuit = QuantumCircuit(q)

d = 3 # set d = something you chose to test different values!

u = U(nqubits-1)
u_dag = reverse_gate(u)

p_right_range = [nqubits-2, nqubits-1]
u_range = list(range(0, nqubits-1))
p_left_range = list(range(0, nqubits))

circuit.append(p_left(nqubits,(1-d)*pi), p_left_range)
circuit.append(U(nqubits-1), u_range)

for i in range((d-1)//2):
    circuit.append(p_right(pi), p_right_range) #debug this, doesnt work just as a 2 qubit gate
    circuit.append(u_dag, u_range)
    circuit.append(p_left(nqubits,pi), p_left_range)
    circuit.append(u, u_range)

circuit.measure_all()
circuit.draw('mpl')

qasm_sim = Aer.get_backend('qasm_simulator')
transpiled_circuit = transpile(circuit, qasm_sim)
qobj = assemble(transpiled_circuit)
results = qasm_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

qc = qsvt_search('001110')
qc.draw('mpl')

qasm_sim = Aer.get_backend('qasm_simulator')
transpiled_circuit = transpile(qc, qasm_sim)
qobj = assemble(transpiled_circuit)
results = qasm_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

qc = qsvt_search('001110110')
qc.draw('mpl')

qasm_sim = Aer.get_backend('qasm_simulator')
transpiled_circuit = transpile(qc, qasm_sim)
qobj = assemble(transpiled_circuit)
results = qasm_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

