import numpy as np
import h5py
import os

import mercury as mr

import sys
sys.path.append('/plot_scripts/')
from map_packages_colors_1v1 import *
from plot_scripts_1v1 import *

# print(" Tasks to choose from : ")
# print(" Heisenberg dynamics (hdyn), Random Quantum Circuits (rqc), Quantum Fourier Transform (qft)")
# print("###############################")
# print(" Package list to choose from :")
# print(" cirq,  hybridq, intel_qs_cpp, pennylane_l, projectq, qcgpu, qibojit, qrack_sch, qsimcirq, quest, svsim, yao, hiq, pennylane, qibo, qiskit, qulacs")
# 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 for the first package : ")
# package_1 = input(" Enter the choice of the package, package 1 : ")
# p1_com_cap = input(" Enter the choice of the compute capability for package 1 : ")
# p1_prec = input(" Enter the choice of the precision for package 1 : ")

# task_2 = input("Enter the task for the second package : ")
# package_2 = input(" Enter the choice of the package, package 2 : ")
# p2_com_cap = input(" Enter the choice of the compute capability for package 2 : ")
# p2_prec = input(" Enter the choice of the precision for package 2 : ")

def abs_time(t1, p1, p1_cc, p1_pr, t2, p2, p2_cc, p2_pr, N_end):

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

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

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

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

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

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

    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)

    dir = os.getcwd()
    data_file_p1 = dir + '/data/{}/{}_{}_{}.h5'.format(t1, p1, p1_cc, p1_pr)
    data_file_p2 = dir + '/data/{}/{}_{}_{}.h5'.format(t2, p2, p2_cc, p2_pr)

    fig, ax = plt.subplots()

    mr.Md(f"TtS performance of the selected options")

    if os.path.isfile(data_file_p1) and os.path.isfile(data_file_p2):
            h5f_1 = h5py.File(data_file_p1, 'r')
            dat_1 = h5f_1[storage_dict[p1]][:]
            h5f_1.close()

            h5f_2 = h5py.File(data_file_p2, 'r')
            dat_2 = h5f_2[storage_dict[p2]][:]
            h5f_2.close()

            plot_abs_data_n_arr(N_arr_t1, dat_1, p1+'_'+t1+'_'+p1_cc+'_'+p1_pr)
            plot_abs_data_n_arr(N_arr_t2, dat_2, p2+'_'+t2+'_'+p2_cc+'_'+p2_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_t2
            else:
                N_arr = N_arr_t1

            gen_settings(fig, ax, r"N (system size)", r"Time ($t_{package}$)", False, True, True, N_arr[0]-2, N_arr[-1], True, 10**-1, 10**5, "out", None)
    else:
        mr.Md(f" Re-select the options as the requested configuration is not supported (check the table in the index page for supported configurations)")

# abs_time(task_1, package_1, p1_com_cap, p1_prec, task_2, package_2, p2_com_cap, p2_prec)

def relative_time_wrt_pack(t1, p1, p1_cc, p1_pr, t2, p2, p2_cc, p2_pr, N_end):

    mr.Md("___")
    mr.Md(f"Relative performance")

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

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

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

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

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

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

    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)

    fig, ax = plt.subplots()

    dir = os.getcwd()
    data_file_p1 = dir + '/data/{}/{}_{}_{}.h5'.format(t1, p1, p1_cc, p1_pr)
    data_file_p2 = dir + '/data/{}/{}_{}_{}.h5'.format(t2, p2, p2_cc, p2_pr)

    if os.path.isfile(data_file_p1) and os.path.isfile(data_file_p2):

        h5f_1 = h5py.File(data_file_p1, 'r')
        dat_1 = h5f_1[storage_dict[p1]][:]
        h5f_1.close()

        h5f_2 = h5py.File(data_file_p2, 'r')
        dat_2 = h5f_2[storage_dict[p2]][:]
        h5f_2.close()

        if np.sum(dat_1) > np.sum(dat_2):
            if N_arr_t1[0] > N_arr_t2[0]:
                dat_2 = dat_2[3:]
                N_arr = N_arr_t1
            elif N_arr_t1[0] < N_arr_t2[0]:
                dat_1 = dat_1[3:]
                N_arr = N_arr_t2
            else:
                N_arr = N_arr_t1
            plot_comp_data_n_arr(N_arr, dat_1, dat_2, p1+'_'+t1+'_'+p1_cc+'_'+p1_pr)
            plot_comp_data_n_arr(N_arr, dat_2, dat_2, p2+'_'+t2+'_'+p2_cc+'_'+p2_pr)
            # save_flag = input("Do you want to save the plot?")
            # if save_flag == "Y":
            #     gen_settings(fig, ax, r"N (system size)", r"Relative time - " + p2, False, True, True, 10**-1, 10**3, "out", "relative_perf_{}_{}_{}_{}_{}_{}_{}_{}.pdf".format(t1, p1, p1_cc, p1_pr, t2, p2, p2_cc, p2_pr))
            # else:
            gen_settings(fig, ax, r"N (system size)", r"Relative time", False, True, True, N_arr[0]-2, N_arr[-1], True, 10**-1, 10**3, "out", None)
        else:
            if N_arr_t1[0] > N_arr_t2[0]:
                dat_2 = dat_2[3:]
                N_arr = N_arr_t1
            elif N_arr_t1[0] < N_arr_t2[0]:
                dat_1 = dat_1[3:]
                N_arr = N_arr_t2
            else:
                N_arr = N_arr_t1
            plot_comp_data_n_arr(N_arr, dat_2, dat_1, p2+'_'+t2+'_'+p2_cc+'_'+p2_pr)
            plot_comp_data_n_arr(N_arr, dat_1, dat_1, p1+'_'+t1+'_'+p1_cc+'_'+p1_pr)
            # save_flag = input("Do you want to save the plot?")
            # if save_flag == "Y":
            # gen_settings(fig, ax, r"N (system size)", r"Relative time", False, True, True, 10**-1, 10**3, "out", "relative_perf_{}_{}_{}_{}_{}_{}_{}_{}.pdf".format(t1, p1, p1_cc, p1_pr, t2, p2, p2_cc, p2_pr))
            # else:
            gen_settings(fig, ax, r"N (system size)", r"Relative time", False, True, True, N_arr[0]-2, N_arr[-1], True, 10**-1, 10**3, "out", None)

# relative_time_wrt_pack(task_1, package_1, p1_com_cap, p1_prec, task_2, package_2, p2_com_cap, p2_prec)