from qcompile import comp_utils as utils
from qiskit import *
import numpy as np
from pprint import pprint
import random 
import matplotlib.pyplot as plt

random.seed(1)  # set random seed 1

# Intro to some helper functions (which can be found in comp_utils.py): 

circ = qiskit.QuantumCircuit(3)   # construct a simple 3 qbit circuit 
circ.h(0)
circ.cx(0, 2)
circ.i(1)
circ.z(0)

print(circ)    ## visualize the circuit 

## read_circ .......................................................................................................

gate_lst, num_qbits = utils.read_circ(circ)  # here we use the read_circ helper function to extract the meta data 
print('gate_lst: {}\n'.format(gate_lst))     # the gate_lst is an ordered list of tuples containing info about 
                                             # the gate, the qbits being applied to, and the parameters of the gate 
    
## general_replace .................................................................................................

utils.general_replace(gate_lst, 'Z', [('X', [], [])])  # the general_replace function allows us to manipulate the gate_lst
print('replaced Z: {}\n'.format(gate_lst))             # here I replaced each instance of Z gate with an X gate. 

utils.general_replace(gate_lst, 'I', [])               # we can remove gates by providing an empty list of replacement gates
print('removed I: {}\n'.format(gate_lst))              # here I just removes every instance of 'I'

## write_circ ......................................................................................................

new_circ = utils.write_circ(gate_lst, num_qbits)  # this function takes the gate_lst and creates a qiskit circuit object
print(new_circ)


# Here we will construct the gates defined above and numerically check if they are equal :

def Rx(theta):
    """Produces the Rx matrix given float theta (in radians) """
    Rx_mat = np.array([[np.cos(theta/2), complex(0, -np.sin(theta/2))],
                       [complex(0, -np.sin(theta/2)), np.cos(theta/2)]])
    return Rx_mat 

def Ry(theta):
    """Produces the Ry matrix given float theta (in radians) """
    Ry_mat = np.array([[np.cos(theta/2), -np.sin(theta/2)],
                       [np.sin(theta/2),  np.cos(theta/2)]])
    return Ry_mat 

def Rz(phi): 
    """Produces the Rz matrix given float phi (in radians) """
    Rz_mat = np.array([[np.exp(-complex(0, phi/2)), 0],
                       [0, np.exp(complex(0, phi/2))]])
    return Rz_mat 

I = np.eye(2)

H = (1 / np.sqrt(2)) * np.array([[1,  1],
                                 [1, -1]])

X = np.array([[0, 1],
              [1, 0]])

Y =  np.array([[0, -complex(0, 1)],
               [complex(0, 1), 0]])

Z =  np.array([[1,  0],
               [0, -1]])


print('I matrix: --------------------- ') 
pprint(I)  
pprint(np.round(Rz(0), decimals=3))
pprint(np.round(Rx(0), decimals=3))

print('\n')

print('H matrix: --------------------- ') 
pprint(H)  
pprint(np.round(Rz(np.pi/2) @ Rx(np.pi/2) @ Rz(np.pi/2), decimals=3))
pprint(np.round(Rx(np.pi/2) @ Rz(np.pi/2) @ Rx(np.pi/2), decimals=3))

print('\n')

print('X matrix: --------------------- ') 
pprint(X)  
pprint(np.round(Rx(np.pi), decimals=3))

print('\n')

print('Z matrix: --------------------- ') 
pprint(Z)  
pprint(np.round(Rz(np.pi), decimals=3))

print('\n')

print('Y matrix: --------------------- ') 
pprint(Y)  
pprint(np.round(Rz(-np.pi/2) @ Rx(np.pi) @ Rz(np.pi/2), decimals=3))
pprint(np.round(Rx(np.pi/2) @ Rz(np.pi) @ Rx(-np.pi/2), decimals=3))

print('I matrix: --------------------- ') 
global_phase = np.exp(0)
pprint(I)  
pprint(np.round( global_phase * Rz(0) , decimals=3))
pprint(np.round( global_phase * Rx(0) , decimals=3))

print('\n')

print('H matrix: --------------------- ') 
global_phase = np.exp(complex(0, np.pi/2))
pprint(H)  
pprint(np.round( global_phase * (Rz(np.pi/2) @ Rx(np.pi/2) @ Rz(np.pi/2)) , decimals=3))
pprint(np.round( global_phase * (Rx(np.pi/2) @ Rz(np.pi/2) @ Rx(np.pi/2)) , decimals=3))

print('\n')

print('X matrix: --------------------- ') 
global_phase = np.exp(complex(0, np.pi/2))
pprint(X)  
pprint(np.round( global_phase * Rx(np.pi) , decimals=3))

print('\n')

print('Z matrix: --------------------- ') 
global_phase = np.exp(complex(0, np.pi/2))
pprint(Z)  
pprint(np.round( global_phase * Rz(np.pi) , decimals=3))

print('\n')

print('Y matrix: --------------------- ')
global_phase = np.exp(complex(0, 3 * np.pi/2))
pprint(Y)  
pprint(np.round( global_phase * (Rz(-np.pi/2) @ Rx(np.pi) @ Rz(np.pi/2)) , decimals=3))
pprint(np.round( global_phase * (Rx(np.pi/2) @ Rz(np.pi) @ Rx(-np.pi/2)) , decimals=3))

def simple_compiler(circ): 
    '''A quantum compiler that produces a new quantum circuit from the 
     restricted subset of available gates. '''
    
    gate_lst, num_qbits = utils.read_circ(circ)  
    
    # replace CNOT: 
    replacement_gates = [('H', utils.get_second, []), ('Cz', [], []), ('H', utils.get_second, [])] 
    utils.general_replace(gate_lst, 'Cx', replacement_gates)
    
    # replace Identity: 
    replacement_gates = [('Rz', [], [0])]
    utils.general_replace(gate_lst, 'I', replacement_gates)
    
    # replace Hadamard: 
    replacement_gates = [('Rz', [], [np.pi/2]), ('Rx', [], [np.pi/2]), ('Rz', [], [np.pi/2])]
    utils.general_replace(gate_lst, 'H', replacement_gates)
    
    # replace X: 
    replacement_gates = [('Rx', [], [np.pi])]
    utils.general_replace(gate_lst, 'X', replacement_gates)
    
    # replace Z: 
    replacement_gates = [('Rz', [], [np.pi])]
    utils.general_replace(gate_lst, 'Z', replacement_gates)
    
    # replace y: 
    replacement_gates = [('Rz', [], [-np.pi/2]), ('Rx', [], [np.pi]), ('Rz', [], [np.pi/2])]
    utils.general_replace(gate_lst, 'Y', replacement_gates)
    
    # replace Ry(theta):
    replacement_gates = [('Rz', [], [-np.pi/2]), ('Rx', [], utils.get_first), ('Rz', [], [np.pi/2])] 
    utils.general_replace(gate_lst, 'Ry', replacement_gates)
    
    compiled_circ = utils.write_circ(gate_lst, num_qbits)
    
    return compiled_circ



# Testing ground: 
circ = utils.random_circ_generator(num_qbits=3, num_gates=5)   # randomly generate a circuit, 
print(circ)                            

compiled_circ = simple_compiler(circ)  # compile it 
print(compiled_circ)
     
equal = utils.circ_equal(circ, compiled_circ) # this helper function compares the magnitudes of each state_vector (element wise)
print(equal)                                  # to determine if they are identical (up to a global phase)


# Brute force test: 
for i in range(1000): 
    circ = utils.random_circ_generator()
    compiled_circ = simple_compiler(circ)
    equal = utils.circ_equal(circ, compiled_circ)
    
    if not equal.all():
        print('FAILED at circuit {}'.format(i))
        break
    
    else: 
        print('passed circuit {}'.format(i))
        
print('Passed all tests!')

# Here we analyze the overhead 
depth_array = []
compiled_depth_array = []
ratio = []

for i in range(100):
    circ = utils.random_circ_generator(num_qbits=5, num_gates=15)   # randomly generate a circuit with 5 qbits and 15 gates 
    compiled_circ = simple_compiler(circ)
    
    equal = utils.circ_equal(circ, compiled_circ)
    if not equal.all():                                # make sure we are compiling properly! 
        print("FAIL @ circuit {}".format(i))
        break
    
    depth_circ = circ.depth()
    depth_comp = compiled_circ.depth()
    
    depth_array.append(depth_circ)                    # store depth 
    compiled_depth_array.append(depth_comp)                  # store new circuit depth 
    
    ratio.append(((depth_comp - depth_circ) / depth_circ) * 100)
    

print('average initial circuit depth = {} +/- {}'.format(np.mean(depth_array), np.std(depth_array)))
print('average compiled circuit depth = {} +/- {}'.format(np.mean(compiled_depth_array), np.std(compiled_depth_array)))
print('average increase in depth = {}%'.format(np.round(np.mean(ratio), decimals=2)))


plt.plot(depth_array, label='circuit')
plt.plot(compiled_depth_array, label='compiled')
plt.legend()
plt.show()

## Here we implement the optimized compiler: 

def compiler(circ):
    '''A quantum compiler that produces a new quantum circuit from the 
     restricted subset of available gates. '''
    
    gate_lst, num_qbits = utils.read_circ(circ)  
    
    # Preprocessing (Step1): 
    
    utils.general_replace(gate_lst, 'I', [])  # remove Identity 
    
    length = len(gate_lst)
    for index in range(length - 1):         # iterate over the lst and remove redundant Cx, Cz gates
        
        if index >= (len(gate_lst) - 1):    # by removing the repetitive Cz and Cx gates 
            break                           # we reduce the size of the list, so we need to check this edge case
    
        curr_gate_str = gate_lst[index][0]
        curr_qbit_lst = gate_lst[index][1]
        
        if curr_gate_str in ['Cx', 'Cz']:          # Check if this gate is a Cz or Cx gate
            nxt_gate_str = gate_lst[index+1][0]
            nxt_qbit_lst = gate_lst[index+1][1]
            
            if ((nxt_gate_str == curr_gate_str) and  # check that we are applying a Cz or Cx gate twice 
                (nxt_qbit_lst == curr_qbit_lst)):    # consecutively on the same control and target qbits
                
                del gate_lst[index + 1]             # remove both gates
                del gate_lst[index]
    
    
    # Compile (similar to the simple compiler): 
    
    # replace CNOT: 
    replacement_gates = [('H', utils.get_second, []), ('Cz', [], []), ('H', utils.get_second, [])] 
    utils.general_replace(gate_lst, 'Cx', replacement_gates)
    
    # replace Hadamard: 
    replacement_gates = [('Rz', [], [np.pi/2]), ('Rx', [], [np.pi/2]), ('Rz', [], [np.pi/2])]
    utils.general_replace(gate_lst, 'H', replacement_gates)
    
    # replace X: 
    replacement_gates = [('Rx', [], [np.pi])]
    utils.general_replace(gate_lst, 'X', replacement_gates)
    
    # replace Z: 
    replacement_gates = [('Rz', [], [np.pi])]
    utils.general_replace(gate_lst, 'Z', replacement_gates)
    
    # replace y: 
    replacement_gates = [('Rz', [], [-np.pi/2]), ('Rx', [], [np.pi]), ('Rz', [], [np.pi/2])]
    utils.general_replace(gate_lst, 'Y', replacement_gates)
    
    # replace Ry(theta):
    replacement_gates = [('Rz', [], [-np.pi/2]), ('Rx', [], utils.get_first), ('Rz', [], [np.pi/2])] 
    utils.general_replace(gate_lst, 'Ry', replacement_gates)
                
            
    # simplification (Step2): 
    
    index = 0 
    while(index < len(gate_lst) - 1): 
#         print('index: {}'.format(index))
        curr_gate_str = gate_lst[index][0]
        curr_qbit_lst = gate_lst[index][1]
        curr_qbit_params = gate_lst[index][2]
        
        if curr_gate_str in ['Rx', 'Rz']:          # Check if this gate is a Rz or Rx gate
            i = 1                                  # another dummy index to look at gates ahead
            while(index + i < len(gate_lst)):
#                 print('dummy: {}'.format(i))
                nxt_gate_str = gate_lst[index+i][0]
                nxt_qbit_lst = gate_lst[index+i][1]
                nxt_qbit_params = gate_lst[index+i][2]

                if ((nxt_gate_str == curr_gate_str) and  # check that we are applying a Rz or Rx gate twice 
                    (nxt_qbit_lst == curr_qbit_lst)):    # consecutively on the same control and target qbits

                    del gate_lst[index + i]             # remove both gates
                    del gate_lst[index]

                    new_gate = (curr_gate_str, curr_qbit_lst, [curr_qbit_params[0] + nxt_qbit_params[0]])
                    gate_lst.insert(index, new_gate)    # add the combined gate and
                    break                               # break current while loop 
                
                elif ((nxt_gate_str != curr_gate_str) and  # if the next gate applied to the same qbit is different 
                      (nxt_qbit_lst == curr_qbit_lst or    # i.e instead of another Rx gate we apply a Rz or a Cz to 
                       curr_qbit_lst[0] in nxt_qbit_lst)): # the same qbit then 
                    
                    index += 1  # move forward nothing left here to simplify 
                    break
                    
                else:         # the next gate is being applied to a different set of qbits 
                    i += 1    # so we can safely check the next gate in the list 
                
            index += 1
        
        else: 
            index += 1
    
            
    compiled_circ = utils.write_circ(gate_lst, num_qbits)
    
    return compiled_circ

# Here we analyze the overhead 
depth_array = []
compiled_depth_array = []
optimized_depth_array = []
ratio_simple = []
ratio_opt = []

for i in range(100):
    circ = utils.random_circ_generator(num_qbits=5, num_gates=15)   # randomly generate a circuit with 5 qbits and 15 gates 
    compiled_circ = simple_compiler(circ)
    optimized_circ = compiler(circ)
    
    equal1 = utils.circ_equal(circ, compiled_circ)
    equal2 = utils.circ_equal(circ, optimized_circ)
    
    if not equal1.all():                                # make sure we are compiling properly! 
        print("simple compiler FAIL @ circuit {}".format(i))
        break
    
    if not equal2.all():                               
        print("optimized compiler FAIL @ circuit {}".format(i))
        break
       
    depth_circ = circ.depth()
    depth_comp = compiled_circ.depth()
    depth_opt = optimized_circ.depth()
    
    depth_array.append(depth_circ)                    # store depth 
    compiled_depth_array.append(depth_comp)           # store new circuit depth 
    optimized_depth_array.append(depth_opt)           # store optimized circuit depth
    
    ratio_simple.append(((depth_comp - depth_circ) / depth_circ) * 100)
    ratio_opt.append(((depth_opt - depth_circ) / depth_circ) * 100)
    
    

print('average initial circuit depth = {} +/- {}'.format(np.mean(depth_array), np.std(depth_array)))
print('average compiled circuit depth = {} +/- {}'.format(np.mean(compiled_depth_array), np.std(compiled_depth_array)))
print('average optimized circuit depth = {} +/- {}'.format(np.mean(optimized_depth_array), np.std(optimized_depth_array)))
print('average increase in depth from simple compiler = {}%'.format(np.round(np.mean(ratio_simple), decimals=2)))
print('average increase in depth from optimized compiler = {}%'.format(np.round(np.mean(ratio_opt), decimals=2)))

plt.plot(depth_array, label='circuit')
plt.plot(compiled_depth_array, label='compiled')
plt.plot(optimized_depth_array, label='optimized')
plt.legend()
plt.show()


## Here we develop the router: 

topology = {0:[4, 1], 1:[0, 2], 2:[1, 3], 3:[2, 4], 4:[3, 0]} # the key is the qbit and the value is the set of connected qbits
                                                              # the order of the qbit indicies in the value lst correspond to 
                                                              # a natural orientation for traversing the ring 

def get_path(topology, start, end):   
    ''' Takes a dict (topology) representing the geometry of the 
    connections, an int (start) representing the starting index 
    and an int (end) representing the ending index and returns 
    a list corresponding to the shortest path from end --> start
    (assuming a ring topology)'''
    
    path_cw = []              # initialize the clockwise traversed path 
    path_ccw = []             # initialize the counter clockwise traversed path 
    
    path_cw.append(end)
    path_ccw.append(end)      # add the first point 
    
    current = end
    while start not in topology[current]:    # traverse clockwise while adding each intermediate qbit index
        current = topology[current][1]           
        path_cw.append(current)                    
    
    path_cw.append(start)                   
    
    current = end
    while start not in topology[current]:    # traverse counter clockwise while adding each intermediate qbit index
        current = topology[current][0]           
        path_ccw.append(current)                    
    
    path_ccw.append(start)
    
    if len(path_cw) <= len(path_ccw):        # return the shorter among the two paths 
        return path_cw
    else:
        return path_ccw

    
def get_swaps(path): 
    '''Take a list (path) between an end qbit index and 
    a start qbit index. Return a list of tuples (replacement_gates)
    which correspond to the set of swap gates required to swap the 
    end qbit with the start qbit'''
    
    replacement_gates = []
    
    for i in range(len(path) - 1):                 # iterate over the path
        swap = ('S', [path[i], path[i+1]], [])     # swap gate between consecutive qbits along the path
        replacement_gates.append(swap)            
                                              # at this point, we have shifted each qbit along the path
    fix_offset = replacement_gates[:-1]       # by 1, this may be a problem if we swapped the control qbit
    fix_offset.reverse()                      # along the path, so we need to fix the shift before we go ahead
    replacement_gates += fix_offset           # this simply involves performing the same swaps in the reverse order
                                              # excluding the very final swap in the first case
    return replacement_gates                                                                                  
    

def circ_router(circ, topology): 
    '''Takes a compiled circuit, and a topology to produce a 
    properly routed circuit. '''
    
    gate_lst, num_qbits = utils.read_circ(circ)
    
    for index, gate in enumerate(gate_lst):      # iterate through the circuit
        curr_gate_str = gate_lst[index][0]
        curr_qbit_lst = gate_lst[index][1]
        curr_parms = gate_lst[index][2]
        
        if curr_gate_str == 'Cz':                # check if this gate is a cz gate 
            cntrl_qbit = curr_qbit_lst[0]
            trgt_qbit = curr_qbit_lst[1]
            
            if not trgt_qbit in topology[cntrl_qbit]:              # check if the control and target qbits are 'connected'
                new_target = topology[cntrl_qbit][1]               # if not, choose a qbit that is connected to the control
                                                                   # to be the new target qbit 
                path = get_path(topology, new_target, trgt_qbit)   # find the path between the new target and the old target
                first_swaps = get_swaps(path)                      # the swap gates required to swap new_target w/ old target
                path.reverse()                                    
                swap_backs = get_swaps(path)                       # the swap gates required to swap them back to original 
                
                replacement_lst = first_swaps + [(curr_gate_str, [cntrl_qbit, new_target], curr_parms)] + swap_backs
                
                del gate_lst[index]                                 # we delete the old Cz gate 
                for j, replacement in enumerate(replacement_lst):   # and add the swap + cz gate + swap back gates
                    gate_lst.insert(index + j, replacement)            
    
    compiled_circ = utils.write_circ(gate_lst, num_qbits)           # 
    
    return compiled_circ     
        
    

## Testing the qbit router: 
topology = {0:[4, 1], 1:[0, 2], 2:[1, 3], 3:[2, 4], 4:[3, 0]}  # the key is the qbit and the value is the set of connected qbits

for i in range(1000):
    circ = utils.random_circ_generator(num_qbits=5, num_gates=15)   # generate random circuit 
    # print(circ)
    # print('\n\n')

    compiled_circ = compiler(circ)       # compile it 
    # print(compiled_circ)
    # print('\n\n')

    routed_circ = circ_router(compiled_circ, topology)   # route it 
    # print(routed_circ)

    equal1 = utils.circ_equal(circ, routed_circ)           # make sure the routed circuit matches the original circuit 
    equal2 = utils.circ_equal(compiled_circ, routed_circ)  # make sure the routed circuit matches the compiled circuit 

    if not equal1.all():
        print(equal1)
        break

    if not equal2.all():
        print(equal2)
        break
    
    print('passed test #{}'.format(i+1))


# Here we analyze the overhead (again, again)
topology = {0:[4, 1], 1:[0, 2], 2:[1, 3], 3:[2, 4], 4:[3, 0]} 

depth_array = []
compiled_depth_array = []
optimized_depth_array = []
routed_depth_array = []
ratio_simple = []
ratio_opt = []
ratio_routed = []

for i in range(100):
    circ = utils.random_circ_generator(num_qbits=5, num_gates=15)   # randomly generate a circuit with 5 qbits and 15 gates 
    compiled_circ = simple_compiler(circ)
    optimized_circ = compiler(circ)
    routed_circ = circ_router(optimized_circ, topology)
    
    equal1 = utils.circ_equal(circ, compiled_circ)
    equal2 = utils.circ_equal(circ, optimized_circ)
    equal3 = utils.circ_equal(circ, routed_circ)
    
    if not equal1.all():                                       # make sure we are compiling and routing properly! 
        print("simple compiler FAIL @ circuit {}".format(i))
        break
    
    if not equal2.all():                               
        print("optimized compiler FAIL @ circuit {}".format(i))
        break
    
    if not equal3.all():                               
        print("circuit router FAIL @ circuit {}".format(i))
        break
       
    depth_circ = circ.depth()
    depth_comp = compiled_circ.depth()
    depth_opt = optimized_circ.depth()
    depth_rout = routed_circ.depth()
    
    depth_array.append(depth_circ)                    # store depth 
    compiled_depth_array.append(depth_comp)           # store new circuit depth 
    optimized_depth_array.append(depth_opt)           # store optimized circuit depth
    routed_depth_array.append(depth_rout)
    
    ratio_simple.append(((depth_comp - depth_circ) / depth_circ) * 100)
    ratio_opt.append(((depth_opt - depth_circ) / depth_circ) * 100)
    ratio_routed.append(((depth_rout - depth_circ) / depth_circ) * 100)
    

print('average initial circuit depth = {} +/- {}'.format(np.mean(depth_array), np.std(depth_array)))
print('average compiled circuit depth = {} +/- {}'.format(np.mean(compiled_depth_array), np.std(compiled_depth_array)))
print('average optimized circuit depth = {} +/- {}'.format(np.mean(optimized_depth_array), np.std(optimized_depth_array)))
print('average routed circuit depth = {} +/- {}'.format(np.mean(routed_depth_array), np.std(routed_depth_array)))
print('average increase in depth from simple compiler = {}%'.format(np.round(np.mean(ratio_simple), decimals=2)))
print('average increase in depth from optimized compiler = {}%'.format(np.round(np.mean(ratio_opt), decimals=2)))
print('average increase in depth from circuit rout = {}%'.format(np.round(np.mean(ratio_routed), decimals=2)))

plt.plot(depth_array, label='circuit')
plt.plot(compiled_depth_array, label='compiled')
plt.plot(optimized_depth_array, label='optimized')
plt.plot(routed_depth_array, label='routed', linestyle='--')
plt.legend()
plt.show()

