# Credit to: https://github.com/qulacs/cirq-qulacs/blob/master/benchmark/benchmark_state_vector_qulacs.py

import re
import subprocess
import os

def parse_qasm_to_package_gen(N, com_cap, input_filename, output_filename, save_data_path):

    with open(output_filename, 'w') as f:

        f.write('import os\n')
        f.write('from time import process_time, perf_counter\n')
        f.write('import numpy as np\n\n')
        f.write('from pyqpanda import *\n')

        f.write('def SX_gate(qp, qb):\n')
        f.write('    qp << U4(qb, np.pi/4., np.pi/2., -np.pi/2., -np.pi/2.)\n')

        f.write('def SXDG_gate(qp, qb):\n')
        f.write('    qp << U4(qb, -np.pi/4., np.pi/2., np.pi/2., -np.pi/2.)\n')

        f.write('def U2_gate(qp, qb, p, l):\n')
        f.write('    qp << U4(qb, 0., p, np.pi/2., l)\n')

        f.write('def U3_gate(qp, qb, t, p, l):\n')
        f.write('    qp << U4(qb, 0., p, t, l)\n')

        f.write('t_sp = process_time()\n')
        f.write('t_s = perf_counter()\n')

        if com_cap == 'st':
            # f.write('init(QMachineType.CPU_SINGLE_THREAD)\n')
            f.write('qvm = init_quantum_machine(QMachineType.CPU_SINGLE_THREAD)\n')
        elif com_cap == 'mt':
            f.write('qvm = CPUQVM()\n')
            f.write('qvm.init_qvm()\n')
        elif com_cap == 'gpu':
            f.write('qvm = init_quantum_machine(QMachineType.GPU)\n')

        f.write('qvm.set_configure({}, {})\n'.format(N, N))
        f.write('q_bits = qvm.qAlloc_many({})\n'.format(N))
        f.write('qprog = QProg()\n')

        with open(input_filename, "r") as ifile:
            lines = ifile.readlines()

            lc = 0
            cirq_c = 0
            for line in lines:

                s = re.search(r"Generated|QASM|include|qreg|Qubits|cirq|x|y|z|h|rx|ry|rz|cx|u2|u3|\bsx\b|\bsxdg\b|s", line)

                if s is None:
                    continue

                elif s.group() in ['Generated', 'QASM', 'include', 'qreg', 'Qubits']:#, 'cirq']:
                    lc = lc + 1
                    f.write('# {}\n'.format(s.group()))
                    continue

                elif s.group() == 'cirq':
                    lc = lc + 1
                    f.write('# {} {}\n'.format(s.group(), cirq_c))
                    cirq_c = cirq_c + 1
                    continue

                elif s.group() == 'x' or s.group() == 'y' or s.group() == 'z' or s.group() == 's' or s.group() == 'h':
                    lc = lc + 1
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('qprog << {}(q_bits[{}])\n'.format(s.group().upper(), t_qbit))
                    continue

                elif s.group() == 'sx':
                    lc = lc + 1
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('SX_gate(qprog, q_bits[{}])\n'.format(t_qbit))

                elif s.group() == 'sxdg':
                    lc = lc + 1
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('SXDG_gate(qprog, q_bits[{}])\n'.format(t_qbit))

                elif s.group() == 'rx' or s.group() == 'ry' or s.group() == 'rz':
                    lc = lc + 1
                    m_r = re.findall(r'\((.*?)\)', line)
                    # print(m_r)
                    if 'pi' in m_r[0] and 'e' in m_r[0]:
                        sp_str = m_r[0].split('e-')
                        m_r[0] = [re.findall(r'[-]?\d*\.\d\d*', sp_str[0])[0] + 'E-' + sp_str[1]]
                    elif 'pi' in m_r[0]:
                        m_r[0] = re.findall(r'[-]?\d*\.\d\d*', m_r[0])
                    # print(m_r)
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('qprog << {}(q_bits[{}], np.pi*{})\n'.format(s.group().upper(), t_qbit, float(m_r[0][0])))
                    continue

                elif s.group() == 'cx':
                    lc = lc + 1
                    match = re.findall(r'\[\d\d*\]', line)
                    c_qbit = int(match[0].strip('[]'))
                    t_qbit = int(match[1].strip('[]'))
                    # f.write('cir.append(cirq.X(q[{}]).controlled_by(q[{}]))\n'.format(t_qbit,  c_qbit))
                    f.write('qprog << CNOT(q_bits[{}], q_bits[{}])\n'.format(c_qbit,  t_qbit))
                    continue

                elif s.group() == 'u2':
                    lc = lc + 1
                    m_r = re.findall(r'\((.*?,.*?)\)', line)
                    m_r = m_r[0].split(',')
                    for i, m in enumerate(m_r):
                        if 'pi' in m and 'e' in m:
                            sp_str = m_r[i].split('e-')
                            m_r[i] = [re.findall(r'[-]?\d*\.\d\d*', sp_str[0])[0] + 'E-' + sp_str[1]]
                        elif 'pi' in m:
                            m_r[i] = re.findall(r'[-]?\d*\.\d\d*', m)
                        else:
                            m_r[i] = '0'
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('U2_gate(qprog, q_bits[{}], np.pi*{}, np.pi*{})\n'.format(t_qbit, float(m_r[0][0]), float(m_r[1][0])))
                    continue

                elif s.group() == 'u3':
                    lc = lc + 1
                    m_r = re.findall(r'\((.*?,.*?,.*?)\)', line)
                    m_r = m_r[0].split(',')
                    for i, m in enumerate(m_r):
                        if 'pi' in m and 'e' in m:
                            sp_str = m_r[i].split('e-')
                            m_r[i] = [re.findall(r'[-]?\d*\.\d\d*', sp_str[0])[0] + 'E-' + sp_str[1]]
                        elif 'pi' in m:
                            m_r[i] = re.findall(r'[-]?\d*\.\d\d*', m)
                        else:
                            m_r[i] = '0'
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('U3_gate(qprog, q_bits[{}], np.pi*{}, np.pi*{}, np.pi*{})\n'.format(t_qbit, float(m_r[0][0]), float(m_r[1][0]),float(m_r[2][0])))
                    continue

        f.write('result = qvm.directly_run(qprog)\n')
        f.write('t_e = perf_counter()\n')
        f.write('t_ep = process_time()\n')
        f.write('os.chdir(\'{}\')\n'.format(save_data_path))
        f.write('np.save(\'time_n{}.npy\', [t_e - t_s, t_ep - t_sp])\n'.format(N))
        f.write('print(t_e - t_s)\n')
        f.write('print(t_ep - t_sp)\n')
        # f.write('print(qvm.get_qstate())\n')

        # result = subprocess.run(["grep", ".", input_filename, "|", "wc", "-l"], text=True)#stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        _cmd = "grep . " + input_filename + "| wc -l"
        lc_res = subprocess.check_output(_cmd, shell=True)
        print(int(lc_res) == lc)
        # print(lc)

task = 'hdyn'
sim_pack = 'qpanda'
c_cap = 'gpu'
prec = 'dp'

for N in range(6, 38, 2):
    input_file_path = '/data/user/gangap_a/{}_singular/qasm_files/'.format(task)
    input_file = input_file_path + 'qasm_N_{}.qasm'.format(N)

    save_file_path = '/data/user/gangap_a/{}_singular/{}/data_{}_{}'.format(task, sim_pack, c_cap, prec)

    output_file_path = '/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/'.format(task, sim_pack, c_cap, prec)

    output_file = output_file_path + '{}_n{}.py'.format(task, N)
    parse_qasm_to_package_gen(N, c_cap, input_file, output_file, save_file_path)