import numpy as np
import h5py
import os

import mercury as mr

import sys
sys.path.append('/plot_scripts/')
from map_packages_colors_mgpu import *
from plot_scripts_mgpu 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(t, pr, n_gpu, compare_pack, N_end):

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

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

    fig, ax = plt.subplots()

    dir = os.getcwd()

    pack_list = []

    mr.Md(f"TtS performance of the different packages")

    for package in package_str:
        data_file = dir + '/data/{}/{}_gpu_{}_{}.h5'.format(t, package, n_gpu, pr)

        if os.path.isfile(data_file):
            h5f = h5py.File(data_file, 'r')
            dat = h5f[storage_dict[package]][:]
            h5f.close()
            pack_list.append(package)

            if t == 'hdyn' or t == 'qft':
                N_arr = range(6, N_end, 2)
            else:
                N_arr = range(12, N_end, 2)
            plot_abs_data_n_arr(N_arr, dat, package+'_'+t+'_'+'gpu_'+str(n_gpu)+'_'+pr)

    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**4, "out", None)
    # else:
    #     print(" Re-select the options as the requested option data is not available.")

    p_main = dir + '/data/{}/{}_gpu_{}_{}.h5'.format(t, compare_pack, n_gpu, pr)
    if not os.path.isfile(p_main):
        return "Please select other package"
    fig, ax = plt.subplots()
    h5f = h5py.File(p_main, 'r')
    main_dat = h5f[storage_dict[compare_pack]][:]
    h5f.close()

    mr.Md("___")
    mr.Md(f"Relative performance to the package {compare_pack}")

    for package in package_str:
        dat_file = dir + '/data/{}/{}_gpu_{}_{}.h5'.format(t, package, n_gpu, pr)
        if os.path.isfile(dat_file):
            h5f = h5py.File(dat_file, 'r')
            dat = h5f[storage_dict[package]][:]
            h5f.close()
            plot_comp_data_n_arr(N_arr, dat, main_dat, package+'_'+t+'_'+'gpu_'+str(n_gpu)+'_'+pr)

    gen_settings(fig, ax, r"N (system size)", r"Relative time", False, False, True, N_arr[0]-2, N_arr[-1], True, -1, 10, "out", None)

# abs_time("Heisenberg dynamics", "Single", 1, "cuquantum_qsimcirq", 36)

def abs_time_ngpus(t, pr, pack, compare_ngpu, N_end):

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

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

    fig, ax = plt.subplots()

    dir = os.getcwd()

    if t == 'hdyn' or t == 'qft':
        N_arr = range(6, N_end, 2)
    else:
        N_arr = range(12, N_end, 2)

    mr.Md(f"TtS performance of the package with respect to different number of GPU's")

    for ngpu in [1, 2, 4, 8]:

        data_file = dir + '/data/{}/{}_gpu_{}_{}.h5'.format(t, pack, ngpu, pr)
        if not os.path.isfile(data_file):
            return "Please select other package"

        if os.path.isfile(data_file):
            h5f = h5py.File(data_file, 'r')
            dat = h5f[storage_dict[pack]][:]
            h5f.close()

            plot_abs_data_n_arr(N_arr, dat, pack+'_'+t+'_'+'gpu_'+str(ngpu)+'_'+pr)

    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**4, "out", None)
    # else:
    #     print(" Re-select the options as the requested option data is not available.")

    mr.Md("___")
    mr.Md(f"Relative performance to the {compare_ngpu} GPUs")

    fig, ax = plt.subplots()

    p_main = dir + '/data/{}/{}_gpu_{}_{}.h5'.format(t, pack, compare_ngpu, pr)
    h5f = h5py.File(p_main, 'r')
    main_dat = h5f[storage_dict[pack]][:]
    h5f.close()

    for n_gpu in [1, 2, 4, 8]:
        dat_file = dir + '/data/{}/{}_gpu_{}_{}.h5'.format(t, pack, n_gpu, pr)
        if os.path.isfile(dat_file):
            h5f = h5py.File(dat_file, 'r')
            dat = h5f[storage_dict[pack]][:]
            h5f.close()
            plot_comp_data_n_arr(N_arr, dat, main_dat, pack+'_'+t+'_'+'gpu_'+str(n_gpu)+'_'+pr)

    gen_settings(fig, ax, r"N (system size)", r"Relative time", False, False, True, N_arr[0]-2, N_arr[-1], True, -1, 10, "out", None)

# abs_time_ngpus("Heisenberg dynamics", "Single", "cuquantum_qsimcirq", 8, 36)