### Quantum indexed Bidirectional Associative Memory by A. Sarkar, Z. Al-Ars, C. G. Almudever, K. Bertels
### Repository reference: https://gitlab.com/prince-ph0en1x/QaGs (by A. Sarkar)
## Importing libraries

%matplotlib inline
import qiskit
from qiskit import IBMQ
from qiskit import Aer
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister, QiskitError
from qiskit.quantum_info.operators import Operator
from qiskit.tools.visualization import circuit_drawer
from qiskit.tools.visualization import plot_histogram
from qiskit.providers.aer import noise
from qiskit.providers.aer.noise import NoiseModel, errors

import random
from math import *
import os
import re
import sys
import math
import matplotlib.pyplot as plt
import numpy as np


## Defining some auxiliary functions

def hamming_distance(str1, str2):
    count = sum(c1 != c2 for c1, c2 in zip(str1, str2))
    return count


def convertToNumEncoding(genome_str):
    bin_str = ""
    for i in range(0, len(genome_str)):
        if genome_str[i] == 'A':
            bin_str = bin_str + "0"
        elif genome_str[i] == 'C':
            bin_str = bin_str + "1"
        elif genome_str[i] == 'G':
            bin_str = bin_str + "2"
        elif genome_str[i] == 'T':
            bin_str = bin_str + "3"
    return bin_str

def convertReadToBin(read):
    bin_read = ""
    for i in range(0, len(read)):
        if read[i] == '0':
            bin_read = bin_read + "00"
        elif read[i] == '1':
            bin_read = bin_read + "01"
        elif read[i] == '2':
            bin_read = bin_read + "10"
        elif read[i] == '3':
            bin_read = bin_read + "11"
    return bin_read


## Initializing global variables
# Alphabet set {0, 1, 2, 3} := {A, C, G, T} for DNA Nucleotide bases
AS = {'00', '01', '10', '11'}
A = len(AS)

################# Default Test #################
# Reference Genome: "AATTGTCTAGGCGACC"
#w = "0033231302212011"
#N = len(w)
# Short Read: "CA"
#p = "10"
# Distributed query around zero Hamming distance
#pb = "0000"
################################################


############## Test on HIV genome ##############
genome_file = open("HIVgenome.txt", "r") #select the data source
w = genome_file.read()
genome_file.close()

w = convertToNumEncoding(w)

w = w[0:128] #changes the number of characters you want to read from the file

N = len(w)

# Short read: "TCT" #Try with chains of three
p = "313"

# Distributed query around zero Hamming distance
pb = "000000"
################################################


# Short Read size
M = len(p)

# Number of qubits to encode one character
Q_A = ceil(log2(A))

# Number of data qubits
Q_D = Q_A * M

# Tag Qubits
Q_T = ceil(log2(N-M))
#print('valor Q_T:',Q_T)

# Ancilla Qubits
Q_anc = 8

# Ancilla qubits ids
anc = []

for qi in range(0, Q_anc):
    anc.append(Q_D + Q_T + qi)

total_qubits = Q_D + Q_T + Q_anc

## Inizializing Oracle O Matrix

gamma = 0.25
SS = 2**Q_D
bp = np.empty([1, SS])

for i in range(0, SS):
    i_binary = format(i, '0'+str(Q_D)+'b')
    hd = hamming_distance(i_binary, pb)
    bp[0][i] = sqrt((gamma**hd)*((1 - gamma)**(Q_D - hd)))
    

BO = np.identity(SS) - 2*np.dot(np.conjugate(np.transpose(bp)), bp)
orc = Operator(BO)

qbsp = []
for qi in range (0, Q_D):
    qbsp.append(Q_T+qi)
qbsp.reverse()

## Initialization of IBM QX

IBMQ.enable_account('My Token')
provider = IBMQ.get_provider()

# Pick an available backend
# If this isn't available pick a backend whose name containes '_qasm_simulator' from the output above
backend = provider.get_backend('ibmq_qasm_simulator')

# Uncomment if you want to use local simulator
#backend= Aer.get_backend('qasm_simulator')

## Defining functions
def QIBAM():
    print('Reference genome:', w)
    print('Chosen pattern for testing:', p)
    print('Total number of qubits:', total_qubits)
    print('Number of ancilla qubits:', Q_anc)
    qr = qiskit.QuantumRegister(total_qubits)
    cr = qiskit.ClassicalRegister(Q_T)
    qc = qiskit.QuantumCircuit(qr, cr)
    
    # Initialise
    generateInitialState(qc, qr)
    
    # Transform to Hamming distance
    evolveToHammingDistances(qc, qr)
    
    # Oracle call
    oracle(qc)
    
    # Inversion about mean
    inversionAboutMean(qc, qr)

    it = 0
    for r in range(0, 2):
        # Oracle call
        oracle(qc)

        # Inversion about mean
        inversionAboutMean(qc, qr)

        it = it + 1
        
    print("Grover's algorithm had {} iterations.".format(int(it)))

    # Measurement
    finalGroverMeasurement(qc, qr, cr)

    return qc

# Initialization as suggested by L. C. L. Hollenberg
def generateInitialState(qc, qr):
    for qi in range(0, Q_T):
        qc.h(qr[qi])
    
    control_qubits = []
    for ci in range(0,Q_T):
        control_qubits.append(qr[ci])
    
    ancilla_qubits = []
    for qi in anc:
        ancilla_qubits.append(qr[qi])
    
    for qi in range(0, N - M + 1):
        qis = format(qi, '0' + str(Q_T) + 'b')        
        for qisi in range(0, Q_T):
            if qis[qisi] == '0':
                qc.x(qr[qisi])
    
        wMi = w[qi:qi + M]
        print("Tag: {} - Data: {} - Hamming distance: {}".format(qis, wMi, hamming_distance(convertReadToBin(wMi), convertReadToBin(p))))

        for wisi in range(0, M):
            wisia = format(int(wMi[wisi]), '0' + str(Q_A) + 'b')
            for wisiai in range(0, Q_A):
                if wisia[wisiai] == '1':
                    qc.mct(control_qubits, qr[Q_T + wisi * Q_A + wisiai], ancilla_qubits)

        for qisi in range(0, Q_T):
            if qis[qisi] == '0':
                qc.x(qr[qisi])
                
    wMi = p
    
    for qi in range(N - M + 1, 2**Q_T):
        qis = format(qi, '0' + str(Q_T) + 'b')
        for qisi in range(0, Q_T):
            if qis[qisi] == '0':
                qc.x(qr[qisi])

        for wisi in range(0, M):
            wisia = format(int(wMi[wisi]), '0' + str(Q_A) + 'b')
            for wisiai in range(0, Q_A):
                if wisia[wisiai] == '0':
                    qc.mct(control_qubits, qr[Q_T + wisi * Q_A + wisiai], ancilla_qubits)

        for qisi in range(0, Q_T):
            if qis[qisi] == '0':
                qc.x(qr[qisi])
    
    return


# Calculate Hamming Distance
def evolveToHammingDistances(qc, qr):
    for pi in range(0, M):
        ppi = format(int(p[pi]), '0' + str(Q_A) + 'b')
        for ppii in range(0, Q_A):
            if ppi[ppii] == '1':
                qc.x(qr[Q_T + pi * Q_A + ppii])
    return


# Oracle to mark zero Hamming distance
def oracle(qc):
    qc.unitary(orc, qbsp, label='orc')
    return

# Inversion about mean
def inversionAboutMean(qc, qr):
    for si in range(0, Q_D + Q_T):
        qc.h(qr[si])
        qc.x(qr[si])

    control_qubits = []
    for sj in range(0, Q_D + Q_T - 1):
        control_qubits.append(qr[sj])
    
    ancilla_qubits = []
    for qi in anc:
        ancilla_qubits.append(qr[qi])

    qc.h(qr[Q_D + Q_T - 1])
    qc.mct(control_qubits, qr[Q_D + Q_T - 1], ancilla_qubits)
    qc.h(qr[Q_D + Q_T - 1])
    
    for si in range(0,Q_D + Q_T):
        qc.x(qr[si])
        qc.h(qr[si])
        
    return


# Oracle to mark Memory States
def markStoredStates(qc, qr):
    control_qubits = []
    for qsi in range(0, Q_T + Q_D - 1):
        control_qubits.append(qr[qsi])
    
    ancilla_qubits = []
    for qi in anc:
        ancilla_qubits.append(qr[qi])
    
    for qi in range(0, N - M + 1):
        qis = format(qi, '0' + str(Q_T) + 'b')
        wMi = w[qi:qi + M]
        wt = qis
        for wisi in range(0, M):
            hd = int(format(int(wMi[wisi]), '0'+str(Q_A)+'b'), 2) ^ int(format(int(p[wisi]), '0'+str(Q_A)+'b'), 2)
            wisia = format(hd, '0' + str(Q_A) + 'b')
            wt = wt + wisia
        for qisi in range(0, Q_T + Q_D):
            if wt[qisi] == '0':
                qc.x(qr[qisi])
            qc.h(qr[Q_D + Q_T - 1])
            qc.mct(control_qubits, qr[Q_D + Q_T - 1], ancilla_qubits)
            qc.h(qr[Q_D + Q_T - 1])

            if wt[qisi] == '0':
                qc.x(qr[qisi])
    return


# Final measurement
def finalGroverMeasurement(qc, qr, cr):
    for qi in range(0, Q_T):
        qc.measure(qr[qi], cr[qi])
    return


## Main function
if __name__ == '__main__':
    # Printing some data for testing
    qc = QIBAM()
    print("Circuit depth: {}".format(qc.depth()))
    
    # Total number of gates
    print("Number of gates: {}".format(len(qc.data)))
    gate_num = 1
    for item in qc.data:
        qb_list = ''
        for qb in item[1]:
            qb_list = qb_list + str(qb.index) + ', '
        qb_list = qb_list[:len(qb_list)-2]
        print("#{}: {}, {}".format(gate_num, item[0].name, qb_list))
        gate_num = gate_num + 1
    
    # Drawing circuit
    # qc.draw()
    # Showing histogram
    # BE CAREFUL!
    # Qiskit uses a LSB ordering, meaning the first qubit is all the way to the right!
    # For example, a state of |01> would mean the first qubit is 1 and the second qubit is 0!
    sim = qiskit.execute(qc, backend=backend, shots=8192)
    result = sim.result()
    final=result.get_counts(qc)
    print(final)
    plot_histogram(final)



