import numpy as np
import h5py
import os

import mercury as mr

import sys
sys.path.append('/plot_scripts/')
from map_packages_colors_all import *
from plot_scripts_all import *

# print(" Tasks to choose from : ")
# print(" Heisenberg dynamics (hdyn), Random Quantum Circuits (rqc), Quantum Fourier Transform (qft)")
# print("###############################")
# print(" Compute capability choices for packages :")
# print(" singlethread, multithread, gpu")
# print("###############################")
# print(" Precision choices for different compute capabilities :")
# print(" sp (single precision), dp (double precision)")
# print("###############################")

# task_1 = input(" Enter the task to compare : ")
# package_1 = input(" Enter the choice of the package, package 1 : ")
# p_com_cap = input(" Enter the choice of the compute capability : ")
# p_prec = input(" Enter the choice of the precision : ")

package_str = ['qiskit' , 'cirq', 'qsimcirq', 'pennylane', 'pennylane_l', 'qibo', 'qibojit', 'yao', 'quest', 'qulacs', 'intel_qs_cpp', 'projectq', 'svsim',  'hybridq', 'hiq', 'qcgpu', 'qrack_sch', 'cuquantum_qiskit', 'cuquantum_qsimcirq', 'qpanda',  'qpp', 'myqlm', 'myqlm_cpp', 'braket']

# N = 36
ngates_hdyn = [5544, 7722, 9900, 12078, 14256, 16434, 18612, 20790, 22968, 25146, 27324, 29502, 31680, 33858, 36036, 38214]
ngates_rqc = [3300, 3850, 4400, 5158, 5760, 6362, 7120, 7670, 8220, 8978, 9580, 10182, 10940]
ngates_qft = [186, 344, 550, 804, 1106, 1456, 1854, 2300, 2794, 3336, 3926, 4564, 5250, 5984, 6766, 7596]

def abs_time_pack(t1, t2, cc, pr, N_end):

    if t1 == "Heisenberg dynamics":
        t1 = "hdyn"
        ng1 = ngates_hdyn
    elif t1 == "Random Quantum Circuit":
        t1 = "rqc"
        ng1 = ngates_rqc
    elif t1 == "Quantum Fourier Transform":
        t1 = "qft"
        ng1 = ngates_qft

    if t2 == "Heisenberg dynamics":
        t2 = "hdyn"
        ng2 = ngates_hdyn
    elif t2 == "Random Quantum Circuit":
        t2 = "rqc"
        ng2 = ngates_rqc
    elif t2 == "Quantum Fourier Transform":
        t2 = "qft"
        ng2 = ngates_qft

    if cc == "Singlethread":
        cc = "singlethread"
    elif cc == "Multithread":
        cc = "multithread"
    elif cc == "GPU":
        cc = "gpu"

    if pr == "Single":
        pr = "sp"
    elif pr == "Double":
        pr = "dp"

    fig, ax = plt.subplots()

    dir = os.getcwd()

    if t1 == 'hdyn' or t1 == 'qft':
        N_arr_t1 = np.arange(6, N_end, 2)
    elif t1 == 'rqc':
        N_arr_t1 = np.arange(12, N_end, 2)

    if t2 == 'hdyn' or t2 == 'qft':
        N_arr_t2 = np.arange(6, N_end, 2)
    elif t2 == 'rqc':
        N_arr_t2 = np.arange(12, N_end, 2)

    # pack_list = []

    for package in package_str:
        dat_f1 = dir + '/data/{}/{}_{}_{}.h5'.format(t1, package, cc, pr)
        dat_f2 = dir + '/data/{}/{}_{}_{}.h5'.format(t2, package, cc, pr)

        if os.path.isfile(dat_f1) and os.path.isfile(dat_f2):
            h5f1 = h5py.File(dat_f1, 'r')
            dat1 = h5f1[storage_dict[package]][:]
            h5f1.close()
            # pack_list.append(package)

            h5f2 = h5py.File(dat_f2, 'r')
            dat2 = h5f2[storage_dict[package]][:]
            h5f2.close()

            if N_arr_t1[0] > N_arr_t2[0]:
                dat2 = dat2[3:]
                N_arr = N_arr_t1
            elif N_arr_t1[0] < N_arr_t2[0]:
                dat1 = dat1[3:]
                N_arr = N_arr_t2
            else:
                N_arr = N_arr_t1
            plot_comp_data_n_arr(N_arr, dat1, dat2, package+'_'+t1+'_'+cc+'_'+pr)


    # save_flag = input("Do you want to save the plot?")
    # if save_flag == "Y":
    #     gen_settings(fig, ax, r"N (system size)", r"Time ($t_{package}$)", False, True, True, 10**-1, 10**5, "out", "perf_{}_{}_{}_{}_{}_{}_{}_{}.pdf".format(t1, p1, p1_cc, p1_pr, t2, p2, p2_cc, p2_pr))
    # else:
    if N_arr_t1[0] > N_arr_t2[0]:
        N_arr = N_arr_t1
        ng2 = ng2[3:]
    elif N_arr_t1[0] < N_arr_t2[0]:
        N_arr = N_arr_t2
        ng1 = ng1[3:]
    else:
        N_arr = N_arr_t1

    mr.Md(f" Number of gates of Task 1: {ng1}")
    mr.Md(f" Number of gates of Task 2: {ng2}")

    plot_comp_data_n_arr(N_arr, ng1, ng2, 'gate_count')
    gen_settings(fig, ax, r"N (system size)", r"Task I/Task II ($t_{T1}/t_{T2}$)", False, False, True, N_arr[0]-2, N_arr[-1], True, -3, 25, "out", None)
    # else:
    #     print(" Re-select the options as the requested option data is not available.")

# pkg_str = ['qiskit' , 'cirq', 'qsimcirq', 'pennylane', 'pennylane_l', 'qibo', 'qibojit', 'yao', 'quest', 'qulacs', 'intel_qs_cpp', 'projectq', 'svsim',  'hybridq', 'hiq', 'qcgpu', 'qrack_sch']

# abs_time(pkg_str, task_1, p_com_cap, p_prec)
# abs_time("Heisenberg dynamics", "Singlethread", "Single", 'qsimcirq')
# abs_time_pack("Heisenberg dynamics", "Random Quantum Circuit", "Singlethread", "Single", 34)
# abs_time_pack("Heisenberg dynamics", "Quantum Fourier Transform", "GPU", "Single", 38)