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 *

# 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']


def abs_time_pack(task, package, cc, N_end):

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

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

    fig, ax = plt.subplots()

    dir = os.getcwd()

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

    dat_fsp = dir + '/data/{}/{}_{}_sp.h5'.format(task, package, cc)
    dat_fdp = dir + '/data/{}/{}_{}_dp.h5'.format(task, package, cc)

    if not os.path.isfile(dat_fsp) and not os.path.isfile(dat_fdp):
        return mr.Md(f"Compute capability {cc} possibly not supported")

    mr.Md(f"TtS performance of simulation packages with different precision")

    if os.path.isfile(dat_fsp):
        h5f_sp = h5py.File(dat_fsp, 'r')
        dat_sp = h5f_sp[storage_dict[package]][:]
        h5f_sp.close()
        plot_abs_data_n_arr(N_arr, dat_sp, package+'_'+task+'_{}_'.format(cc)+'sp')

    if os.path.isfile(dat_fdp):
        h5f_dp = h5py.File(dat_fdp, 'r')
        dat_dp = h5f_dp[storage_dict[package]][:]
        h5f_dp.close()
        plot_abs_data_n_arr(N_arr, dat_dp, package+'_'+task+'_{}_'.format(cc)+'dp')

    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)

    mr.Md("___")
    mr.Md(f"Relative performance to single precision")

    fig, ax = plt.subplots()

    if os.path.isfile(dat_fdp) and os.path.isfile(dat_fsp):
        plot_comp_data_n_arr(N_arr, dat_dp, dat_sp, package+'_'+task+'_{}_'.format(cc)+'dp')
        plot_comp_data_n_arr(N_arr, dat_sp, dat_sp, package+'_'+task+'_{}_'.format(cc)+'sp')

    gen_settings(fig, ax, r"N (system size)", r"Relative to double precision", False, True, True, N_arr[0]-2, N_arr[-1], True, 10**-1, 10**3, "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_pack("Heisenberg dynamics", 'qsimcirq', 'Double', 36)

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