'''
(C) Copyright Renata Wong 2023.

This code is licensed under the Apache License, Version 2.0. You may obtain a copy of this license 
at http://www.apache.org/licenses/LICENSE-2.0.

Any modifications or derivative works of this code must retain this copyright notice, and modified files 
need to carry a notice indicating that they have been altered from the originals.

This is the accompanying Qiskit code for the paper 
'Fast Quantum Algorithm for Protein Structure Prediction in Hydrophobic-Hydrophilic Model' (DOI: 10.1016/j.jpdc.2022.03.011)
The code works for amino acid sequences of length > 4.

NOTE: This code assumes the sequence 01001 for testing pursposes. Other sequences may require adjustments of variables
num_solutions and j. 
'''

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
import numpy as np
import math

'''
Taking in user input of the amino acid sequence.
Bit 0 stands for a hydrophilic amino acid, while bit 1 stands for a hydrophobic amino acid
'''

sequence_input = list(input('Enter a sequence of length at least 5 encoded as a string of 0s and 1s:' ))
length = len(sequence_input)

'''
Quantum register holding the sequence encoding
'''
sequence = QuantumRegister(length,'s')


'''
Quantum registers for each of the three coordinates x, y, and z. 
Coordinate order within each register: x0, x1, x2, etc., where x0 is the 
x-coordinate of the 1st amino acid, x1 is the x-coordinate of the 2nd amino acid, etc.
Each register holds 2(length - 1) + 1 values.
Each value can therefore be repserented with a minimum of ceil(log2(2*length - 1)) bits.
There are a 'length' number of amino acids in a sequence.
'''
register_length = int(math.ceil(np.log2(2*length-1)))
coord_reg_length = length * register_length

x_coord = QuantumRegister(coord_reg_length, 'x')
y_coord = QuantumRegister(coord_reg_length, 'y')

'''
Quantum register holding the controls w0, w1, etc.
Each w is of length 2. 
Amino acid at position (0, 0, 0) is fixed. 
Hence, length - 1 amino acids in a sequence require directional information encoded in w. 
'''
w = QuantumRegister(2*(length-1), 'w')

'''
Quantum register holding the binary 1, and the two's complement of 1 (= -1).
The two's complement of 1 can be obtained by negating all the qubits in binary one but the last.
NOTE: Encoding coordinates in two's complement notation allows one to treat subtraction as addition
with the same ripple carry circuit. 
'''
one_reg_length = int(math.ceil(np.log2(2*length-1)))
binary_one = QuantumRegister(one_reg_length,'one')

'''
Quantum register holding the carry bit for ripple-carry adder
'''
carry = QuantumRegister(1, 'c')
  
    
'''
Register holding the Grover auxialliary qubit
'''
g = QuantumRegister(1, 'g')
   
    
'''
Classical register holding the output of the quantum computation
'''
output = ClassicalRegister(length+2*(length-1), 'out')#(length + 2*(length-1) + 2*coord_reg_length,'out')   


'''
The ancilla register of 3 qubits in needed for the multiplexer (controlled Toffoli gates) in the ripple-carry quantum adder.
'''
ancilla = QuantumRegister(3, 'anc')   


'''
Registers for the section FINDING ENERGY VALUES
'''
k_range = length-3
j_range = int((k_range * (k_range+1))/2)
    
x_star = QuantumRegister(4*j_range*register_length, 'x_st')
y_star = QuantumRegister(4*j_range*register_length, 'y_st')
    
r_reg = QuantumRegister(4*j_range*(2*register_length + 1), 'r')
    
delta = QuantumRegister(4*2*j_range, 'delta')
    
psi = QuantumRegister(4*j_range, 'psi')


'''
Registers for section SUMMING UP ENERGY VALUES. 
We create 4 index matrices that hold registers for each of 0 <= p <= 3 in the z_reg to obtain the value for kj in register z.
'''

p0_matrix = [ [ 0 for m in range(4 * (length-4) + 2) ] for n in range(length-2) ]

kj = 0
for k in range(1, length-2):
    for j in range(4 * (k-1) + 2):
        p0_matrix[k][j] = kj
        kj += 1

p1_matrix = [ [ 0 for m in range(4 * (length-4) + 3) ] for n in range(length-2) ]

for k in range(1, length-2):
    for j in range(4 * (k-1) + 3):
        p1_matrix[k][j] = kj
        kj += 1

p2_matrix = [ [ 0 for m in range(4 * (length-4) + 4) ] for n in range(length-2) ]

for k in range(1, length-2):
    for j in range(4 * (k-1) + 4):
        p2_matrix[k][j] = kj
        kj += 1

p3_matrix = [ [ 0 for m in range(4 * (length-4) + 5) ] for n in range(length-2) ]

for k in range(1, length-2):
    for j in range(4 * (k-1) + 5):
        p3_matrix[k][j] = kj
        kj += 1

        
        
z = QuantumRegister(kj, 'z')      #((4 * length - 12) * (4 * length - 11) / 2, 'z') this value is too low  
z_or = QuantumRegister(1, 'z_or')       # ancilla qubit for the multiplexer in Fig. 4   
z_and = QuantumRegister(1, 'z_and')    # ancilla to hold the results of the AND operation




'''
Defining a quantum circuit consisting of the above quantum and classical registers
'''
qc = QuantumCircuit(g, sequence, x_coord, y_coord, w, binary_one, carry, ancilla, x_star, y_star, r_reg, delta, psi, 
                    z, z_or, z_and, output)




'''
SECTION: PREPARING A UNIFORM SUPERPOSITION STATE

Initialisation of sequence values based on user input
'''
for index, bit in enumerate(sequence_input):
    if bit == '1':
        qc.x(sequence[index])

'''
Initialising binary one and the two's complement of one
''' 
qc.x(binary_one[one_reg_length-1])

'''
Setting the quantum state in a uniform superposition wrt. w
'''
qc.h(w[0:2*(length-1)])


'''
Initialization of r_reg
'''
multiplier = int(len(r_reg)/(4*j_range))
for i in range(4*j_range):
    qc.x(r_reg[i*multiplier])
  

'''
Initialize z_or to |1> for cascading OR gates implemented with mcx gate.
'''
qc.x(z_or)



'''
Setting the Grover qubit into the |-> state
'''
qc.x(g)
qc.h(g)


'''
SECTION: CALCULATING COORDINATES
Requires a doubly-controlled ripple-carry adder as a subroutine sc
The adder below is for amino acid sequences of length > 4.
'''
summand_register_length = int(math.ceil(np.log2(2*length-1)))

sw = QuantumRegister(2,'sw')     # control qubits (w)
sa = QuantumRegister(3,'sa')     # ancilla qubits
ss = QuantumRegister(1,'ss')     # carry 
sx = QuantumRegister(summand_register_length,'sx')   # summand x 
sy = QuantumRegister(summand_register_length,'sy')   # summand y
sc = QuantumCircuit(sw, sa, ss, sx, sy, name='ccrca')

sc.ccx(sw[0],sw[1],sa[0])

for i in range(1, summand_register_length):
    sc.ccx(sa[0],sx[i],sa[1])
    sc.cx(sa[1],sy[i])
    sc.ccx(sa[0],sx[i],sa[1])

sc.ccx(sa[0],sx[1],sa[1])
sc.cx(sa[1],ss[0])
sc.ccx(sa[0],sx[1],sa[1])

sc.ccx(sa[0],sx[0],sa[1])
sc.ccx(sa[1],sy[0],sa[2])
sc.cx(sa[2],ss[0])
sc.ccx(sa[1],sy[0],sa[2])
sc.ccx(sa[0],sx[0],sa[1])

sc.ccx(sa[0],sx[2],sa[1])
sc.cx(sa[1],sx[1])
sc.ccx(sa[0],sx[2],sa[1])

sc.ccx(sa[0],ss[0],sa[1])
sc.ccx(sa[1],sy[1],sa[2])
sc.cx(sa[2],sx[1])
sc.ccx(sa[1],sy[1],sa[2])
sc.ccx(sa[0],ss[0],sa[1])

sc.ccx(sa[0],sx[3],sa[1])
sc.cx(sa[1],sx[2])
sc.ccx(sa[0],sx[3],sa[1])

for i in range(2, summand_register_length-2):
    sc.ccx(sa[0],sy[i],sa[1])
    sc.ccx(sa[1],sx[i-1],sa[2])
    sc.cx(sa[2],sx[i])
    sc.ccx(sa[1],sx[i-1],sa[2])
    sc.ccx(sa[0],sy[i],sa[1])
    sc.ccx(sa[0],sx[i+2],sa[1])
    sc.cx(sa[1],sx[i+1])
    sc.ccx(sa[0],sx[i+2],sa[1])
    
sc.ccx(sa[0],sy[summand_register_length-2],sa[1])
sc.ccx(sa[1],sx[summand_register_length-3],sa[2])
sc.cx(sa[2],sx[summand_register_length-2])
sc.ccx(sa[1],sx[summand_register_length-3],sa[2])
sc.ccx(sa[0],sy[summand_register_length-2],sa[1])

for i in range(1, summand_register_length-1):
    sc.cx(sa[0],sy[i])
    
sc.ccx(sa[0],ss[0],sa[1])
sc.cx(sa[1],sy[1])
sc.ccx(sa[0],ss[0],sa[1])

for i in range(2, summand_register_length):
    sc.ccx(sa[0],sx[i-1],sa[1])
    sc.cx(sa[1],sy[i])
    sc.ccx(sa[0],sx[i-1],sa[1])
    
sc.ccx(sa[0],sx[summand_register_length-3],sa[1])
sc.ccx(sa[1],sy[summand_register_length-2],sa[2])
sc.cx(sa[2],sx[summand_register_length-2])
sc.ccx(sa[1],sy[summand_register_length-2],sa[2])
sc.ccx(sa[0],sx[summand_register_length-3],sa[1])

for i in range(summand_register_length-3, 1, -1):
    sc.ccx(sa[0],sx[i-1],sa[1])
    sc.ccx(sa[1],sy[i],sa[2])
    sc.cx(sa[2],sx[i])
    sc.ccx(sa[1],sy[i],sa[2])
    sc.ccx(sa[0],sx[i-1],sa[1])
    sc.ccx(sa[0],sx[i+2],sa[1])
    sc.cx(sa[1],sx[i+1])
    sc.ccx(sa[0],sx[i+2],sa[1])
    sc.cx(sa[0],sy[i+1])

sc.ccx(sa[0],ss[0],sa[1])
sc.ccx(sa[1],sy[1],sa[2])
sc.cx(sa[2],sx[1])
sc.ccx(sa[1],sy[1],sa[2])
sc.ccx(sa[0],ss[0],sa[1])

sc.ccx(sa[0],sx[3],sa[1])
sc.cx(sa[1],sx[2])
sc.ccx(sa[0],sx[3],sa[1])

sc.cx(sa[0],sy[2])

sc.ccx(sa[0],sx[0],sa[1])
sc.ccx(sa[1],sy[0],sa[2])
sc.cx(sa[2],ss[0])
sc.ccx(sa[1],sy[0],sa[2])
sc.ccx(sa[0],sx[0],sa[1])

sc.ccx(sa[0],sx[2],sa[1])
sc.cx(sa[1],sx[1])
sc.ccx(sa[0],sx[2],sa[1])

sc.cx(sa[0],sy[1])

sc.ccx(sa[0],sx[1],sa[1])
sc.cx(sa[1],ss[0])
sc.ccx(sa[0],sx[1],sa[1])

for i in range(summand_register_length):
    sc.ccx(sa[0],sx[i],sa[1])
    sc.cx(sa[1],sy[i])
    sc.ccx(sa[0],sx[i],sa[1])
    
sc.ccx(sw[0],sw[1],sa[0])

subinst = sc.to_instruction()





'''
Ripple-carry adder without control qubits |w>
'''
ssr = QuantumRegister(1,'ssr')     # carry 
sxr = QuantumRegister(summand_register_length,'sxr')   # summand x 
syr = QuantumRegister(summand_register_length,'syr')   # summand y
rca = QuantumCircuit(ssr, sxr, syr, name='rca')

for i in range(1, summand_register_length):
    rca.cx(sxr[1],syr[i])

rca.cx(sxr[1],ssr[0])
rca.ccx(sxr[0],syr[0],ssr[0])
rca.cx(sxr[2],sxr[1])
rca.ccx(ssr[0],syr[1],sxr[1])
rca.cx(sxr[3],sxr[2])

for i in range(2, summand_register_length-2):
    rca.ccx(syr[i],sxr[i-1],sxr[i])
    rca.cx(sxr[i+2],sxr[i+1])
    
rca.ccx(syr[summand_register_length-2],sxr[summand_register_length-3],sxr[summand_register_length-2])

for i in range(1, summand_register_length-1):
    rca.x(syr[i])
    
rca.cx(ssr[0],syr[1])

for i in range(2, summand_register_length):
    rca.cx(sxr[i-1],syr[i])
    
rca.ccx(sxr[summand_register_length-3],syr[summand_register_length-2],sxr[summand_register_length-2])

for i in range(summand_register_length-3, 1, -1):
    rca.ccx(sxr[i-1],syr[i],sxr[i])
    rca.cx(sxr[i+2],sxr[i+1])
    rca.cx(sxr[i+2],syr[i+1])

rca.ccx(ssr[0],syr[1],sxr[1])
rca.cx(sxr[3],sxr[2])
rca.x(syr[2])
rca.ccx(sxr[0],syr[0],ssr[0])
rca.cx(sxr[2],sxr[1])
rca.x(syr[1])
rca.cx(sxr[1],ssr[0])

for i in range(summand_register_length):
    rca.cx(sxr[i],syr[i])
    
subinst_rca = rca.to_instruction()






'''
AMPLITUDE AMPLIFICATION PROCEDURE:

Calculation of coordinates using the ripple-carry adder
b is a global variable used to navigate among the values of vector w.
d stands for the position of an amino acid in a sequence.

MODIFY NUMBER OF SOLUTIONS num_solutions FOR DIFFERENT SEQUENCES.
'''
num_solutions = 8
num_iter = int(math.ceil(np.pi*(np.sqrt(2**(2*(length-1)/num_solutions)))/4))    # Number of iterations 


def calculate_neighbour(register, b, d):
    
    arglist = []
    
    arglist.append(w[b])
    arglist.append(w[b+1])
    
    for i in range(len(ancilla)):
        arglist.append(ancilla[i])
        
    arglist.append(carry[0])

    for i in range(summand_register_length-1, -1, -1): # digit order reversed for adder
        arglist.append(binary_one[i])
        
    for i in range(d*summand_register_length-1, (d-1)*summand_register_length-1, -1): # digit order reversed for adder 
        arglist.append(register[i])
    
    return arglist
    
    
    
    
def update_coordinates(register_1, register_2):
    
    arglist = []
        
    arglist.append(carry[0])

    for i in range(summand_register_length-1, -1, -1): # digit order reversed for adder
        arglist.append(register_1[i])
        
    reg = len(register_2)
    for i in reversed(range(reg)): # digit order reversed for adder 
        arglist.append(register_2[i])
    
    return arglist




'''
THIS IS THE MAIN PART OF THE CODE. 
'''
for grover_iteration in range(num_iter): 
    
    b = 0

    # Subroutine 1: Generating conformational space by calculating coordinates
    # This part has been tested and performs correctly

    for d in range(2, length+1):
        for q in range(summand_register_length):
            qc.cx(x_coord[(d-2)*summand_register_length+q],x_coord[(d-1)*summand_register_length+q])
            qc.cx(y_coord[(d-2)*summand_register_length+q],y_coord[(d-1)*summand_register_length+q])
          
            
        '''
        Calculating the western neighbour of site d-1 (this is the case when |w> = |11>)
        '''
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
        neighbour_west = calculate_neighbour(x_coord, b, d)
        qc.append(subinst, neighbour_west)
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])

        '''
        Calculating the eastern neighbour of site d-1 (this is the case when |w> = |01>)
        '''
        qc.x(w[b])
        neighbour_east = calculate_neighbour(x_coord, b, d)
        qc.append(subinst, neighbour_east)
        qc.x(w[b])
        
        '''
        Calculating the northern neighbour of site d-1 (this is the case when |w> = |00>)
        '''
        qc.x(w[b])
        qc.x(w[b+1])
        neighbour_north = calculate_neighbour(y_coord, b, d)
        qc.append(subinst, neighbour_north)
        qc.x(w[b])
        qc.x(w[b+1])
        
        '''
        Calculating the southern neighbour of site d-1 (this is the case when |w> = |10>)
        '''
        qc.x(w[b+1])
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
        neighbour_south = calculate_neighbour(x_coord, b, d)
        qc.append(subinst, neighbour_south)
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
        qc.x(w[b+1])

        b = b+2
        
    
    
    

    
    


    '''
    SECTION: FINDING ENERGY VALUES
    For the conformation that corresponds to the ground energy of the protein, the energy value will be e = 1, 
    while e = 0 for all other confornations. 
    The energy is calculated based on the number of hydrophobic-hydrophobic contacts in the lattice. 
    For each such contact, a 1 is added to the energy value for the given conformation.
    
    New registers for calculating energy values based on the coordinates x and y: |x_star> and |y_star>.
    Register r_reg stores the result of coordinate comparison between sites k and j. 
    Register delta stores the information as to whether the k-th and j-th lattice sites are both hydrophobic and form a 
    loose contact (i.e., are adjacent in the lattice but not in the sequence). 
    
    At the end of the computation, vector psi contains all the information about hydrophobic-hydrophobic contacts:
    psi(k, j, p) = |1> iff amino acid k and amino acid j, where j is p-neighbour of k, are both hydrophobic and have 
    a contact in the lattice but not in the sequence.
    '''
    
    
    
    
    
    
    
    # star_matrix: 3D array to hold the |x_star> and y_star indices
    # r_matrix: 3D array to hold the indices for r_reg
    # delta_matrix: 3D array to hold pairs of delta indices
    # psi_matrix: 3D array to hold psi indices. 
    # All indexed in the coding convention, i.e. from 0 to n-1 instead of from 1 to n. 
     
    star_matrix = [[[0 for k in range(length)] for j in range(length)] for p in range(4)]
    r_matrix = [[[0 for k in range(length)] for j in range(length)] for p in range(4)]
    delta_matrix = [[[0 for k in range(length)] for j in range(length)] for p in range(4)]
    psi_matrix = [[[0 for k in range(length)] for j in range(length)] for p in range(4)]
    
    star_index, r_index, delta_index, psi_index = 0, 0, 0, 0
    
    for k in range(length-3):
        for j in range(k+3, length):
            for p in range(4):
                
                star_matrix[k][j][p] = [m for m in range(star_index, star_index + register_length)]
                star_index += register_length
                
                r_matrix[k][j][p] = [m for m in range(r_index, r_index + 2*register_length + 1)]
                r_index += 2 * register_length + 1
                
                delta_matrix[k][j][p] = (delta_index, delta_index+1)
                delta_index += 2
                
                psi_matrix[k][j][p] = psi_index
                psi_index += 1
    
    

    
    #The code below has not been tested within the algorithm but the indices are assigned correctly and so it should work
   
    
    for k in range(length-3):
        for j in range(k+3, length):
            for p in range(4):
                
                for r in range(register_length):
                    qc.cx(x_coord[k+r], x_star[star_matrix[k][j][p][r]])
                    qc.cx(y_coord[k+r], y_star[star_matrix[k][j][p][r]])
                    
                if p == 0:
                    neighbour_north = update_coordinates(binary_one, [y_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_north)
                
                                                                      
                if p == 1:
                    neighbour_east = update_coordinates(binary_one, [x_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_east)
                                                                      
                if p == 2:
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                    neighbour_south = update_coordinates(binary_one, [y_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_south)
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                     
                if p == 3:
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                    neighbour_west = update_coordinates(binary_one, [x_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_west)
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                
                for r in range(register_length):
                    qc.cx(x_coord[j+r], x_star[star_matrix[k][j][p][r]])
                    qc.cx(y_coord[j+r], y_star[star_matrix[k][j][p][r]])
                  
                
                # rows 13 to 19
                for r in range(register_length):
                    qc.x(y_star[star_matrix[k][j][p][r]])
                    qc.x(x_star[star_matrix[k][j][p][r]])
                
                qc.ccx(y_star[star_matrix[k][j][p][0]], r_reg[r_matrix[k][j][p][0]], r_reg[r_matrix[k][j][p][1]])
                    
                for r in reversed(range(1, register_length)):
                    qc.ccx(y_star[star_matrix[k][j][p][r]], r_reg[r_matrix[k][j][p][r]], r_reg[r_matrix[k][j][p][r+1]])
                    
                for r in reversed(range(register_length, 2*register_length)):
                    qc.ccx(x_star[star_matrix[k][j][p][r-register_length]], r_reg[r_matrix[k][j][p][r]], 
                           r_reg[r_matrix[k][j][p][r+1]])
                    
                for r in range(register_length):
                    qc.x(y_star[star_matrix[k][j][p][r]])
                    qc.x(x_star[star_matrix[k][j][p][r]]) 
                     
                
                # row 20 onwards
                qc.ccx(sequence[k], r_reg[r_matrix[k][j][p][2*register_length]], delta[delta_matrix[k][j][p][0]])
                qc.ccx(sequence[j], delta[delta_matrix[k][j][p][0]], delta[delta_matrix[k][j][p][1]])
                
                qc.cx(delta[delta_matrix[k][j][p][1]], psi[psi_matrix[k][j][p]])
                

    

 


            
            
    '''
    Uncomputing calculation of H-H contacts (stored in vector psi) to free up qubits
    '''
    #This part has not been tested due to running out of memory. But should be correct. 
    
    
    for k in reversed(range(length-3)):
        for j in reversed(range(k+3, length)):
            for p in reversed(range(3)):
                
                qc.ccx(sequence[j], delta[delta_matrix[k][j][p][0]], delta[delta_matrix[k][j][p][1]])
                qc.ccx(sequence[k], r_reg[r_matrix[k][j][p][2*register_length]], delta[delta_matrix[k][j][p][0]])
                
                for r in range(register_length):
                    qc.x(y_star[star_matrix[k][j][p][r]])
                    qc.x(x_star[star_matrix[k][j][p][r]])  
                
                for r in range(register_length, 2*register_length):
                    qc.ccx(x_star[star_matrix[k][j][p][r-register_length]], r_reg[r_matrix[k][j][p][r]], 
                           r_reg[r_matrix[k][j][p][r+1]])
                  
                for r in range(1, register_length):
                    qc.ccx(y_star[star_matrix[k][j][p][r]], r_reg[r_matrix[k][j][p][r]], r_reg[r_matrix[k][j][p][r+1]]) 
                
                qc.ccx(y_star[star_matrix[k][j][p][0]], r_reg[r_matrix[k][j][p][0]], r_reg[r_matrix[k][j][p][1]]) 
                
                for r in range(register_length):
                    qc.x(y_star[star_matrix[k][j][p][r]])
                    qc.x(x_star[star_matrix[k][j][p][r]])
                    
                for r in reversed(range(register_length)):
                    qc.cx(x_coord[j+r], x_star[star_matrix[k][j][p][r]])
                    qc.cx(y_coord[j+r], y_star[star_matrix[k][j][p][r]])    
                
                if p == 0:
                    neighbour_north = update_coordinates(binary_one, [y_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_north)
                
                                                                      
                if p == 1:
                    neighbour_east = update_coordinates(binary_one, [x_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_east)
                                                                      
                if p == 2:
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                    neighbour_south = update_coordinates(binary_one, [y_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_south)
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                     
                if p == 3:
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                    neighbour_west = update_coordinates(binary_one, [x_star[star_matrix[k][j][p][r]] 
                                                                      for r in range(register_length)])
                    qc.append(subinst_rca, neighbour_west)
                    for i in range(summand_register_length-1):
                        qc.x(binary_one[i])
                
                for r in reversed(range(register_length)):
                    qc.cx(x_coord[k+r], x_star[star_matrix[k][j][p][r]])
                    qc.cx(y_coord[k+r], y_star[star_matrix[k][j][p][r]])

                
    
    
    
    
    
    

    
    
    
    

    
    
    
    
    
    # This part has been tested and performs correctly
    '''
    # Reversing coordinate calculation 
    '''
    b = b-2

    for d in range (length, 1, -1):

        
        # Uncomputing the western neighbour of site d-1 (|w> = |11>)
        
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
        neighbour_west = calculate_neighbour(x_coord, b, d)
        qc.append(subinst.inverse(), neighbour_west)
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
              

        
        # Uncomputing the eastern neighbour of site d-1 (|w> = |01>)
        
        qc.x(w[b])
        neighbour_east = calculate_neighbour(x_coord, b, d)
        qc.append(subinst.inverse(), neighbour_east)
        qc.x(w[b])

        
        # Uncomputing the northern neighbour of site d-1 (|w> = |00>)
        
        qc.x(w[b])
        qc.x(w[b+1])
        neighbour_north = calculate_neighbour(y_coord, b, d)
        qc.append(subinst.inverse(), neighbour_north)
        qc.x(w[b])
        qc.x(w[b+1])
        
        
        
        # Uncomputing the southern neighbour of site d-1 (|w> = |10>)
        
        qc.x(w[b+1])
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
        neighbour_south = calculate_neighbour(x_coord, b, d)
        qc.append(subinst.inverse(), neighbour_south)
        for i in range(summand_register_length-1):
            qc.x(binary_one[i])
        qc.x(w[b+1])

        b = b-2

        for q in range(summand_register_length-1, -1, -1):
            qc.cx(x_coord[(d-2)*summand_register_length+q],x_coord[(d-1)*summand_register_length+q])
            qc.cx(y_coord[(d-2)*summand_register_length+q],y_coord[(d-1)*summand_register_length+q])

     



            
            
            
            
            
    '''
    SUMMING UP ENERGY VALUES IN PSI = THIS IS THE CRUCIAL ORACLE PART
    STORING THE SUM OF ENERGIES IN VECTOR Z_{length-3, 3, j}
    FLIPPING THE BIT ON GROVER QUBIT g FOR THE LARGEST ENERGY VALUE 
    
    
    p0_matrix = [ [ 0 for m in range(4 * (length-4) + 2) ] for n in range(length-2) ]

    kj = 0
    for k in range(1, length-2):
        for j in range(4 * (k-1) + 2):
            p0_matrix[k][j] = kj
            kj += 1
            
    p1_matrix = [ [ 0 for m in range(4 * (length-4) + 3) ] for n in range(length-2) ]

    for k in range(1, length-2):
        for j in range(4 * (k-1) + 3):
            p1_matrix[k][j] = kj
            kj += 1
    
    p2_matrix = [ [ 0 for m in range(4 * (length-4) + 4) ] for n in range(length-2) ]

    for k in range(1, length-2):
        for j in range(4 * (k-1) + 4):
            p2_matrix[k][j] = kj
            kj += 1
    
    p3_matrix = [ [ 0 for m in range(4 * (length-4) + 5) ] for n in range(length-2) ]

    for k in range(1, length-2):
        for j in range(4 * (k-1) + 5):
            p3_matrix[k][j] = kj
            kj += 1
    
    ''' 
    
    

    
    
    
    # Adding up the energies. 
    # The energies for each conformation are stored in qubits z_{length-3, 3, j}.
    # If j = 0, there are no HH contacts and the energy is 0.
    # If j = 1, there is 1 HH contact and, hence, the energy is 1. And so on. 
    
    
    
    for k in range(length-3):
        for p in range(4):
            
            multiplexer = [psi[psi_matrix[k][i][p]] for i in range(k+3, length)]
                
            for j in reversed(range(4*k + p + 1)):
                
                
                if p == 0:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.cx(z_and, z[p0_matrix[k+1][j+1]])
                    # uncomputing to reuse the z_or and z_and ancillas
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.cx(z_and, z[p0_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                elif p == 1:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p1_matrix[k+1][j+1]])
                    
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p1_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    

                elif p == 2:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p2_matrix[k+1][j+1]])
                    
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p2_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                elif p == 3:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p3_matrix[k+1][j+1]])
                    
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p3_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
    
    
    
    
    
    # This is the crucial part of the oracle
    # WE ASSUME j = 1 SINCE WE ASSUME THAT OUR SEQUENCE IS 01001, 
    # i.e. second and fifth amino acids are hydrophobic, first, third and fourth are hydrophilic.
    # PARAMETER j IS ADJUSTABLE.
    
    j = 1
    qc.cx(z[p3_matrix[length-3][j]], g)
    
    
    
    
    
    # Uncomputing the summation of energy
    
    for k in reversed(range(length-3)):
        for p in reversed(range(4)):
            
            multiplexer = [psi[psi_matrix[k][i][p]] for i in range(k+3, length)]
                
            for j in range(4*k + p + 1):
                
                
                if p == 0:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.cx(z_and, z[p0_matrix[k+1][j+1]])
                    # uncomputing to reuse the z_or and z_and ancillas
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.cx(z_and, z[p0_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p3_matrix[k][3]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                elif p == 1:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p1_matrix[k+1][j+1]])
                    
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p1_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p0_matrix[k+1][j]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    

                elif p == 2:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p2_matrix[k+1][j+1]])
                    
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p2_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p1_matrix[k+1][j]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                elif p == 3:
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p3_matrix[k+1][j+1]])
                    
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])
                    qc.mcx(multiplexer, z_or)
                    qc.x(z_or)
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.cx(z_and, z[p3_matrix[k+1][j]])
                    
                    qc.ccx(z_or, z[p2_matrix[k+1][j]], z_and)
                    qc.x(z_or)
                    qc.mcx(multiplexer, z_or)
                    for index, qubit in enumerate(multiplexer):
                        qc.x(multiplexer[index])    
    
    
    
    
    
    
    

    
    '''
    SECTION: IDENTIFYING THE CONFORMATION WITH THE MINIMAL ENERGY
    
    Using Grover's diffusion operation on register w
    '''
    
    for i in range(len(w)):
        qc.h(w[i])
        qc.x(w[i])
        
    
    multiplexer = [w[i] for i in range(len(w) - 1)]
    
    qc.h(w[len(w) - 1]) 
    qc.mcx(multiplexer, w[len(w) - 1])
    qc.h(w[len(w) - 1])


    for i in range(len(w)):
        qc.x(w[i])
        qc.h(w[i])  
    

    
    







'''
MEASUREMENT.
WE ASSUME THAT THE AMINO ACID SEQUENCE IS PHPPH = 01001. 
HENCE, WE ARE EXPECTING TO SEE THE FOLLOWING VALUES IN THE CHART:

01011011
01101111
00000110
00001110
11111001
11110001
10100100
10101100
'''

conformations = []

for i in range(length):
    conformations.append(sequence[i])
for i in range(2*(length-1)):
    conformations.append(w[i])

# Reverse the order in which the output is shown so that it can be read from left to right.
conformations.reverse()

qc.measure(conformations, output)

from qiskit import Aer, execute
#from qiskit.providers.aer import QasmSimulator, StatevectorSimulator
simulator = Aer.get_backend('qasm_simulator')
result = execute(qc, backend = simulator, shots = 10).result()
counts = result.get_counts()
#statevector = result.get_statevector(qc)
#print(statevector)
#from qiskit.tools.visualization import plot_histogram
#plot_histogram(counts)
print(counts, sep='\n')

qc.draw('mpl')