from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.circuit.library import ZGate
from qiskit.quantum_info import Operator, Statevector
from qiskit.extensions import Initialize
from qiskit.providers.aer import AerSimulator
from qiskit.providers.aer.backends import statevector_simulator, qasm_simulator
from qiskit.visualization import plot_histogram, plot_state_city
from qiskit import Aer, execute, transpile
import numpy as np

def targets(n):
    state1 = ""
    state2 = ""
    bits = ["0", "1"]
    for i in range(n):
        state1 += bits[i%2]
        state2 += bits[1-(i%2)]
    return state1, state2

def getData(input_array, mem_bits):
    state = ""
    for i in input_array:
        state += bin(i)[2:].rjust(mem_bits, '0')
    return state[::-1]

def decoder(addr_bits, aux_bits, n):
    addr = QuantumRegister(addr_bits)
    tau = QuantumRegister(aux_bits)
    
    decoder = QuantumCircuit(addr, tau)
    
    decoder.x(tau[0])
    decoder.cx(addr[addr_bits-1], tau[1])
    decoder.cx(tau[1], tau[0])
    
    c = 1
    
    for i in range(1, addr_bits-1):
        p = 2**i
        for j in range(p):
            decoder.ccx(addr[addr_bits-i-1], tau[j], tau[p+j])
            c += 1
        for j in range(p):
            decoder.cx(tau[p+j], tau[j])
    p = aux_bits - 1 - c 
    for j in range(p):
        decoder.ccx(addr[0], tau[j], tau[p+j])
        c += 1
    for j in range(p):
        decoder.cx(tau[p+j], tau[j])
    
    
    return decoder.to_gate(label = "Decoder")

def groverOracle(aux_bits, mem_bits, n):
    winners = targets(mem_bits)
    
    tau = QuantumRegister(aux_bits)
    mem = QuantumRegister(mem_bits*n)
    oracle = QuantumCircuit(tau, mem)
    
    gate1 = ZGate().control(num_ctrl_qubits=mem_bits, ctrl_state=winners[0])
    gate2 = ZGate().control(num_ctrl_qubits=mem_bits, ctrl_state=winners[1])

    for i in range(aux_bits):
        qubits = [mem[mem_bits*i + j] for j in range(mem_bits)]
        oracle.append(gate1, qubits + [tau[i]])
        oracle.append(gate2, qubits + [tau[i]])
    
    return oracle.to_gate(label="Oracle")

def groverDiffuser(addr_bits):
    d = 2**addr_bits
    s = np.array([[1.0/np.sqrt(d)] for i in range(d)])
    op = np.matmul(s, s.T)
    U_S = Operator((2*op)-np.identity(d))
    return U_S

def qram(input_array):
        
    n = len(input_array)
    if n < 8:
        input_array += [0 for i in range(8-n)]
    n = len(input_array)  
    
    addr_bits = len(bin(n-1)) - 2                       # number of address bits
    aux_bits = n                                        # number of memory locations
    mem_bits = len(bin(max(input_array))) - 2           # size of one memory location
    
    
    addr = QuantumRegister(addr_bits, name="a")         # address register
    tau = QuantumRegister(aux_bits, name="tau")         # auxillary register to store selected address
    mem = QuantumRegister(mem_bits*n, name="m")         # memory register
    cr = ClassicalRegister(addr_bits)                   # to store final output
    
    qram_circ = QuantumCircuit(addr, tau, mem, cr)
    
    for i in range(addr_bits):
        qram_circ.h(addr[i])
    
    # Load data into qRAM
    data = getData(input_array[::-1], mem_bits)
    for i in range(len(data)):
        if data[i] == "1":
            qram_circ.x(mem[i])
    # End of loading
    
        
    # Entangle tau with addr
    decode = decoder(addr_bits, aux_bits, n)
    qram_circ.append(decode, addr[:] + tau[:])

    # Grover oracle
    oracle = groverOracle(aux_bits, mem_bits, n)
    qram_circ.append(oracle, tau[:] + mem[:])
    # End of grover oracle

    # Detangle tau from addr
    encode = decode.inverse()
    encode.label = "Encoder"
    qram_circ.append(encode, addr[:] + tau[:])

    # Diffuser
    U_S = groverDiffuser(addr_bits)
    qram_circ.append(U_S, [addr[i] for i in range(addr_bits)])
    
    # Measure    
    qram_circ.measure([addr[i] for i in range(addr_bits)], [cr[i] for i in range(addr_bits)])
    
    return qram_circ

circ = qram([3, 2, 0, 1])
circ.draw('mpl')

backend = Aer.get_backend("qasm_simulator")
job = execute(circ, backend)
res = job.result()
counts = res.get_counts()
print(counts)
plot_histogram(counts)

circ = qram([1, 0, 3, 2, 3])
circ.draw('mpl')

backend = Aer.get_backend("qasm_simulator")
job = execute(circ, backend)
res = job.result()
counts = res.get_counts()
print(counts)
plot_histogram(counts)

circ = qram([0, 1, 3, 2, 2])
circ.draw('mpl')

backend = Aer.get_backend("qasm_simulator")
job = execute(circ, backend)
res = job.result()
counts = res.get_counts()
print(counts)
plot_histogram(counts)