# 导入相应的依赖库
import numpy as np
import math
from math import pi
import re
from qiskit import BasicAer,execute
from qiskit import QuantumCircuit,ClassicalRegister,QuantumRegister
from qiskit.tools.visualization import circuit_drawer
from qiskit import Aer, execute

# 实现目标量子态的制备电路
# 输入经典的数据，以向量的形式给出（注意：这里给出的向量元素为实数）
classic_vec = [0.234,0.567]
# 获得经典向量的二范数，存储至变量x_norm
x_norm = np.linalg.norm(classic_vec, ord=None, axis=None, keepdims=False)
# 获得归一化之后的向量，即对应量子态上的概率幅
norm_vec = classic_vec/x_norm
print(norm_vec)
# 计算制备目标量子态所需要的量子比特
num_qubits = int(math.log(len(classic_vec),2))
# 利用Qiskit中的初始化内置函数，生成对应的制备电路
q_inital = QuantumRegister(num_qubits,'q_inital')
Circuit_inital = QuantumCircuit()
Circuit_inital.add_register(q_inital)
Circuit_inital.initialize(norm_vec,q_inital)
# 对初始化制备电路，生成qsam代码，并截取量子门操作部分的代码
Qasm_Code = Circuit_inital.qasm().split(';')[3:-1]
# 去除list中字符串的‘\n’部分
Qasm_Code = [x.strip() for x in Qasm_Code if x.strip()!='']
# 设置参数的提取规则
p1 = re.compile(r'[(](.*?)[)]', re.S)   # 提取圆括号中的数据，返回数据类型为list
p2 = re.compile(r'[[](.*?)[]]', re.S)   # 提取方括号中的数据，返回数据类型为list

# 设置数据寄存器和地址寄存器，对应各自的量子比特数量为num_qubits
# q0表示address地址寄存器，量子比特与数据制备寄存器相等
q0 = QuantumRegister(num_qubits,'q0')
c0 = ClassicalRegister(num_qubits,'c0')
# q1表示data数据寄存器，用于数据制备
q1 = QuantumRegister(num_qubits,'q1')
c1 = ClassicalRegister(num_qubits,'c1')
# 辅助量子比特寄存器B，有一个量子比特
q2 = QuantumRegister(1,'q2')
c2 = ClassicalRegister(1,'c2')
# 将寄存器添加至量子电路
Circuit = QuantumCircuit()
Circuit.add_register(q0,c0)
Circuit.add_register(q1,c1)
Circuit.add_register(q2,c2)
# Circuit.add_register(q0)
# Circuit.add_register(q1)
# Circuit.add_register(q2)

# 制备地址寄存器的状态
for i in range(num_qubits):
    Circuit.h(q0[i])
# 对目标量子态进行制备（概率幅编码）
for i in range(len(Qasm_Code)):
    code_str = Qasm_Code[i]
    if code_str[0:2]=='ry':
        theta = float(re.findall(p1, code_str)[0])
        target_qubit = int(re.findall(p2, code_str)[0])
        Circuit.u3(theta,0,0,q1[target_qubit])
    elif code_str[0:2]=='cx':
        control_qubit = int(re.findall(p2, code_str)[0])
        target_qubit = int(re.findall(p2, code_str)[1])
        Circuit.cx(q1[control_qubit],q1[target_qubit])
    else:
        print('The desired vector must be real number!')                           
Circuit.barrier()
# 执行Hadamard Test操作
Circuit.h(q2[0])
for i in range(len(Qasm_Code)):
    code_str = Qasm_Code[len(Qasm_Code)-1-i]
    if code_str[0:2]=='ry':
        theta = float(re.findall(p1, code_str)[0])
        target_qubit = int(re.findall(p2, code_str)[0])
        Circuit.cu3(-theta,0,0,q2[0],q1[target_qubit])
    elif code_str[0:2]=='cx':
        control_qubit = int(re.findall(p2, code_str)[0])
        target_qubit = int(re.findall(p2, code_str)[1])
        Circuit.ccx(q2[0],q1[control_qubit],q1[target_qubit])
    else:
        print('The desired vector must be real number!')
for i in range(num_qubits):
    Circuit.ccx(q2[0],q0[i],q1[i])
Circuit.h(q2[0])
# 绘制量子电路图
# Circuit.draw(output='mpl')

# 定义两个函数（分解为通用量子门）
# 定义具有两个控制位的受控旋转门（绕Y轴旋转）
def cu3_2(theta,control_1,control_2,target,circuit):
    circuit.cu3(theta/2,0,0,control_1,target)
    circuit.cu3(theta/2,0,0,control_2,target)
    circuit.cx(control_1,control_2)
    circuit.cu3(-theta/2,0,0,control_2,target)
    circuit.cx(control_1,control_2)
    return circuit
# 定义具有三个控制位的受控非门（使用两个初始状态为0态的辅助量子比特）
def cccx(control_1,control_2,control_3,target,circuit):
    circuit.ccx(control_1,control_2,q_auxiliary[0])
    circuit.ccx(control_3,q_auxiliary[0],q_auxiliary[1])
    circuit.cx(q_auxiliary[1],target)
    circuit.ccx(control_3,q_auxiliary[0],q_auxiliary[1])
    circuit.ccx(control_1,control_2,q_auxiliary[0])
    return circuit

# 执行量子相位估计
# 相位估计的第一寄存器q3，量子比特数量决定了计算精度
qpe_control_qubits = 3
q3 = QuantumRegister(qpe_control_qubits,'q3')
c3 = ClassicalRegister(qpe_control_qubits,'c3')
Circuit.add_register(q3,c3)
# Circuit.add_register(q3)
# 三个控制位的受控非门，需要两个辅助量子比特
q_auxiliary = QuantumRegister(2,'q_auxiliary')
c_auxiliary = ClassicalRegister(2,'c_auxiliary')
Circuit.add_register(q_auxiliary,c_auxiliary)
# Circuit.add_register(q_auxiliary)
# 对相位估计的第一寄存器进行全Hadamard门操作
for i in range(qpe_control_qubits):
    Circuit.h(q3[i])
Circuit.barrier()
# 对相位估计的第二寄存器执行受控U操作
for i in range(qpe_control_qubits):
# 由控制位决定U的阶数，使用for循环进行操作
    for j in range(int(math.pow(2, i))):
# 执行受控Z门以及受控H门的操作，其中控制位为第一寄存器的第i个比特
        Circuit.cz(q3[i],q2[0])
        Circuit.ch(q3[i],q2[0])
        Circuit.barrier()
# 执行多位受控非门的操作
        for k in range(num_qubits):
            cccx(control_1=q3[i],control_2=q2[0],control_3=q0[k],target=q1[k],circuit=Circuit)
        Circuit.barrier()
# 执行两比特控制的U^A操作，用到了两比特控制旋转门以及三比特控制非门的分解电路函数
        for t1 in range(len(Qasm_Code)):
            code_str = Qasm_Code[t1]
            if code_str[0:2]=='ry':
                theta = float(re.findall(p1, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[0])
                cu3_2(theta=theta,control_1=q3[i],control_2=q2[0],target=q1[target_qubit],circuit=Circuit)
            elif code_str[0:2]=='cx':
                control_qubit = int(re.findall(p2, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[1])
                cccx(control_1=q3[i],control_2=q2[0],control_3=q1[control_qubit],target=q1[target_qubit],circuit=Circuit)        
            else:
                print('The desired vector must be real number!')
        Circuit.barrier()
# 执行受控H门操作，其中控制位为第一寄存器的第i个比特
        Circuit.ch(q3[i],q2[0])
        Circuit.barrier()
# 执行单比特控制的U^A+操作，用到了单比特控制的受控旋转门以及TOFFOLI门
        for t2 in range(len(Qasm_Code)):
            code_str = Qasm_Code[len(Qasm_Code)-1-t2]
            if code_str[0:2]=='ry':
                theta = float(re.findall(p1, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[0])
                Circuit.cu3(-theta,0,0,q3[i],q1[target_qubit])
            elif code_str[0:2]=='cx':
                control_qubit = int(re.findall(p2, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[1])
                Circuit.ccx(q3[i],q1[control_qubit],q1[target_qubit])
            else:
                print('The desired vector must be real number!')
        Circuit.barrier()
# 多位受控-Z门的操作
#        for k in range(num_qubits):
#            Circuit.x(q1[k])
#        Circuit.ccx(q3[i],q1[0],q_auxiliary[0])
#        Circuit.ccx(q1[1],q_auxiliary[0],q_auxiliary[1])
#        Circuit.cu3(pi,-pi/2,pi/2,q_auxiliary[1],q2[0])
#        Circuit.cu3(pi,pi/2,pi/2,q_auxiliary[1],q2[0])
#        Circuit.ccx(q1[1],q_auxiliary[0],q_auxiliary[1])
#        Circuit.ccx(q3[i],q1[0],q_auxiliary[0])
#        for k in range(num_qubits):
#            Circuit.x(q1[k])
#        Circuit.barrier()
        Circuit.x(q1[0])
        Circuit.ccx(q1[0],q3[i],q_auxiliary[0])
        Circuit.cu3(pi,-pi/2,pi/2,q_auxiliary[0],q2[0])
        Circuit.cu3(pi,pi/2,pi/2,q_auxiliary[0],q2[0])
        Circuit.ccx(q1[0],q3[i],q_auxiliary[0])
        Circuit.x(q1[0])
        Circuit.barrier()
# 执行单比特控制的U^A操作，用到了单比特控制的受控旋转门以及TOFFOLI门
        for t3 in range(len(Qasm_Code)):
            code_str = Qasm_Code[t3]
            if code_str[0:2]=='ry':
                theta = float(re.findall(p1, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[0])
                Circuit.cu3(theta,0,0,q3[i],q1[target_qubit])
            elif code_str[0:2]=='cx':
                control_qubit = int(re.findall(p2, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[1])
                Circuit.ccx(q3[i],q1[control_qubit],q1[target_qubit])
            else:
                print('The desired vector must be real number!')
        Circuit.barrier()
# 执行受控H门操作，其中控制位为第一寄存器的第i个比特
        Circuit.ch(q3[i],q2[0])
        Circuit.barrier()
# 执行两比特控制的U^A+操作，用到了两比特控制旋转门以及三比特控制非门的分解电路函数
        for t4 in range(len(Qasm_Code)):
            code_str = Qasm_Code[len(Qasm_Code)-1-t4]
            if code_str[0:2]=='ry':
                theta = float(re.findall(p1, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[0])
                cu3_2(theta=-theta,control_1=q3[i],control_2=q2[0],target=q1[target_qubit],circuit=Circuit)
            elif code_str[0:2]=='cx':
                control_qubit = int(re.findall(p2, code_str)[0])
                target_qubit = int(re.findall(p2, code_str)[1])
                cccx(control_1=q3[i],control_2=q2[0],control_3=q1[control_qubit],target=q1[target_qubit],circuit=Circuit)        
            else:
                print('The desired vector must be real number!')
        Circuit.barrier()
# 执行多位受控非门的操作
        for k in range(num_qubits):
            cccx(control_1=q3[i],control_2=q2[0],control_3=q0[k],target=q1[k],circuit=Circuit)
        Circuit.barrier()
# 执行受控H门的操作，其中控制位为第一寄存器的第i个比特
        Circuit.ch(q3[i],q2[0])
        Circuit.barrier()

# 执行量子相位估计的第二阶段（量子傅里叶逆变换）
Circuit.h(q3[0])
Circuit.cu1(-pi/2,q3[0],q3[1])
Circuit.cu1(-pi/4,q3[0],q3[2])
Circuit.h(q3[1])
Circuit.cu1(-pi/2,q3[1],q3[2])
Circuit.h(q3[2])
Circuit.barrier()

# Circuit.measure(q0,c0)
# Circuit.measure(q1,c1)
# Circuit.measure(q2,c2)
# Circuit.measure(q3,c3)
# Circuit.measure(q_auxiliary,c_auxiliary)

Circuit.draw(output='mpl').savefig('QC.pdf',dpi=800)

simulator = Aer.get_backend('statevector_simulator')
result = execute(Circuit, simulator).result()
statevector = result.get_statevector(Circuit)
print(statevector)