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

import re
import subprocess

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

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

        f.write('using Distributed\n')
        f.write('@everywhere empty!(DEPOT_PATH)\n')
        f.write('@everywhere push!(DEPOT_PATH, \"/data/user/gangap_a/jlib_new/\")\n')
        f.write('@everywhere using Yao, YaoBlocks, YaoAPI, CuYao\n')
        f.write('@everywhere using Base.Threads\n')
        f.write('@everywhere using NPZ\n')

        f.write('mutable struct U2{T<:Number} <: PrimitiveBlock{2}\n')
        f.write('    phi::T\n')
        f.write('    lam::T\n')
        f.write('end\n')

        f.write('YaoBlocks.nqudits(fs::U2) = 1\n')
        f.write('YaoBlocks.print_block(io::IO, block::U2) = print(io, \"U2(phi=$(block.phi), lam=$(block.lam))\")\n')

        f.write('function YaoAPI.mat(::Type{T}, fs::U2) where T\n')
        f.write('    p, l = fs.phi, fs.lam\n')
        f.write('    ((cos(0.5*(p + l)) - 1im*sin(0.5*(p + l)))/sqrt(2))*T[1 -cos(l)-1im*sin(l); cos(p) + 1im*sin(p) cos(l + p) + 1im*sin(l + p)]\n')
        f.write('end\n')
        
        f.write('mutable struct U3{T<:Number} <: PrimitiveBlock{2}\n')
        f.write('    theta::T\n')
        f.write('    phi::T\n')
        f.write('    lam::T\n')
        f.write('end\n')

       	f.write('YaoBlocks.nqudits(fs::U3) = 1\n')
        f.write('YaoBlocks.print_block(io::IO, block::U3) = print(io, \"U3(theta=$(block.theta), phi=$(block.phi), lam=$(block.lam))\")\n')

        f.write('function YaoAPI.mat(::Type{T}, fs::U3) where T\n')
        f.write('    t, p, l = fs.theta, fs.phi, fs.lam\n')
        f.write('    (cos(0.5*(p + l)) - 1im*sin(0.5*(p + l)))*T[cos(t/2.) -sin(t/2.)*(cos(l) + 1im*sin(l)); sin(t/2.)*(cos(p) + 1im*sin(p)) cos(t/2.)*(cos(l + p) + 1im*sin(l + p))]\n')   
        f.write('end\n')

        f.write('@const_gate Sx = 0.5*[1. + 1im 1-1im; 1-1im 1+1im]\n')

        f.write('@const_gate Sxdg = 0.5*[1. - 1im 1+1im; 1+1im 1-1im]\n')

        f.write('@const_gate S = [1. 0; 0 1im]\n')

        f.write('function main(n)\n')
        f.write('cir = chain(n)\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('push!(cir, put({}=>{}))\n'.format(t_qbit+1, s.group().upper()))
                    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('push!(cir, put({}=>Sx))\n'.format(t_qbit+1))

                elif s.group() == 'sxdg':
                    lc = lc + 1
                    m_i = re.findall(r'\[\d\d*\]', line)
                    t_qbit = int(m_i[0].strip('[]'))
                    f.write('push!(cir, put({}=>Sxdg))\n'.format(t_qbit+1))

                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('push!(cir, put({}=>R{}(float(pi*{}))))\n'.format(t_qbit+1, s.group()[1], 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('push!(cir, cnot({}, {}))\n'.format(c_qbit+1, t_qbit+1))
                    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('push!(cir, put({}=>U2(float(pi*{}), float(pi*{}))))\n'.format(t_qbit+1, 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('push!(cir, put({}=>U3(float(pi*{}), float(pi*{}), float(pi*{}))))\n'.format(t_qbit+1, float(m_r[0][0]), float(m_r[1][0]), float(m_r[2][0])))
                    continue


        if precision == 'single':
            f.write('zs = cuzero_state(ComplexF32, n)\n')
        elif precision == 'double':
            f.write('zs = cuzero_state(ComplexF64, n)\n')


        # f.write('zs = cuzero_state(n)\n')
        f.write('apply!(zs, cir)\n')
        f.write('return zs\n')
        f.write('end\n')

        f.write('main({})\n'.format(N))
        f.write('_time = @elapsed final_state = main({})\n'.format(N))
        f.write('println(_time)\n')

        f.write('cd(\"{}\")\n'.format(save_data_path))
        f.write('npzwrite(\"time_n{}.npy\", [_time])\n'.format(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)

# parse_qasm_to_package_gen(12, "../qasm_test/qasm_rqc_test.qasm", "test_yao.jl")
# parse_qasm_to_package_gen(12, "/data/user/gangap_a/rqc/data_files/circuit_n12_m14_s0_e0_pEFGH.qasm", "test_cirq.py")

task = 'qft'
sim_pack = 'yao'
com_cap = 'gpu'
prec = 'sp'

for N in range(6, 40, 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, com_cap, prec)

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

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


r"""
for N in range(12, 42, 2):
    input_file_path = '/data/user/gangap_a/rqc/data_files/'
    input_file = input_file_path + 'circuit_n{}_m14_s0_e0_pEFGH.qasm'.format(N)

    output_file_path = '/data/user/gangap_a/rqc/yao/yao_gpu_run_files/'
    output_file = output_file_path + 'yao_rqc_n{}.jl'.format(N)
    parse_qasm_to_package_gen(N, input_file, output_file)

for N in range(41, 51):
    input_file_path = '/data/user/gangap_a/rqc/data_files/'
    input_file = input_file_path + 'circuit_n{}_m14_s0_e0_pEFGH.qasm'.format(N)

    output_file_path = '/data/user/gangap_a/rqc/yao/yao_gpu_run_files/'
    output_file = output_file_path + 'yao_rqc_n{}.jl'.format(N)
    parse_qasm_to_package_gen(N, input_file, output_file)

"""