import matplotlib.pyplot as plt
import numpy as np       
import copy

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, assemble
from collections.abc import Iterable


# Construct A operator that takes |0>_{n+1} to sqrt(1-a) |psi_0>|0> + sqrt(a) |psi_1>|1>
def A_gen(num_state_qubits, a, psi_zero=None, psi_one=None):

    if psi_zero==None:
        psi_zero = '0'*num_state_qubits
    if psi_one==None:
        psi_one = '1'*num_state_qubits
        
    theta = 2 * np.arcsin(np.sqrt(a))
    # Let the objective be qubit index n; state is on qubits 0 through n-1
    qc_A = QuantumCircuit(num_state_qubits+1, name=f"A")
    
    # qc_A.ry(theta, num_state_qubits)       #ry is not present in the list of basis_gates (qiskit-aakash/qiskit/providers/basicaer/unitary_simulator.py)
    
    qc_A.x(num_state_qubits)

    for i in range(num_state_qubits):
        if psi_zero[i]=='1':
            qc_A.cnot(num_state_qubits,i)
            
    qc_A.x(num_state_qubits)
    
    for i in range(num_state_qubits):
        if psi_one[i]=='1':
            qc_A.cnot(num_state_qubits,i)

    return qc_A

# Construct the grover-like operator and a controlled version of it
def Ctrl_Q(num_state_qubits, A_circ):

    # index n is the objective qubit, and indexes 0 through n-1 are state qubits
    qc = QuantumCircuit(num_state_qubits+1, name=f"Q")
    
    temp_A = copy.copy(A_circ)
    A_gate = temp_A.to_gate()
    A_gate_inv = temp_A.inverse().to_gate()
    
    ### Each cycle in Q applies in order: S_chi, A_circ_inverse, S_0, A_circ 
    # S_chi
    qc.z(num_state_qubits)
        
    # A_circ_inverse
    qc.append(A_gate_inv, [i for i in range(num_state_qubits+1)])
        
    # S_0
    for i in range(num_state_qubits+1):
        qc.x(i)
    qc.h(num_state_qubits)
    
    qc.mcx([x for x in range(num_state_qubits)], num_state_qubits)
    
    qc.h(num_state_qubits)
    for i in range(num_state_qubits+1):
        qc.x(i)
        
    # A_circ
    qc.append(A_gate, [i for i in range(num_state_qubits+1)])

    # add "global" phase
    qc.x(num_state_qubits)
    qc.z(num_state_qubits)
    qc.x(num_state_qubits)
    qc.z(num_state_qubits)
    
    # Create a gate out of the Q operator
    qc.to_gate(label='Q')
    
    # and also a controlled version of it
    Ctrl_Q_ = qc.control(1)
    
    # and return both
    return Ctrl_Q_, qc

num_q = 3

qr_ = QuantumRegister(num_q+1)
cr_ = ClassicalRegister(num_q+1)
qc_ = QuantumCircuit(qr_, cr_)

a = 1/8

A = A_gen(num_q, a)
_, Q = Ctrl_Q(num_q, A)

qc_.append(A, qr_)
# qc_.append(Q, qr_)
qc_.measure(qr_, cr_)


from qiskit import execute, BasicAer
backend = BasicAer.get_backend("dm_simulator")
job = execute(qc_, backend, shots=1000)
result = job.result()
counts = result.get_counts()
print(counts)

num_q = 1

qr_ = QuantumRegister(num_q+1)
cr_ = ClassicalRegister(num_q+1)
qc_ = QuantumCircuit(qr_, cr_)

a = 1/5

A = A_gen(num_q, a)
_, Q = Ctrl_Q(num_q, A)

qc_.append(Q, qr_)
display(qc_.draw())
qc_ = qc_.decompose().decompose()

usim = BasicAer.get_backend('unitary_simulator')
qobj = assemble(qc_)
unitary = usim.run(qobj).result().get_unitary()
print(np.round(unitary, 6))



