import numpy as np
import h5py
# from map_packages_colors import *
import matplotlib.pyplot as plt

import os
import sys

import warnings
warnings.filterwarnings("ignore")


params = {'backend': 'ps',
          'axes.labelsize': 16,
          'axes.titlesize': 14,
          'font.size': 8,
          'legend.fontsize': 14,
          'xtick.labelsize': 16,
          'ytick.labelsize': 16,
          'figure.figsize': (10, 10),}
#          'text.usetex': True}

from matplotlib import rc
# plt.rc('text.latex', preamble=r'\usepackage{braket}')
plt.rcParams.update(params)

N_arr = np.arange(6, 14, 2)

path = os.getcwd()

data_path = path + '/exp_data/'

def compare_two_expect_all(task, prec, pkg_str):
    fig, ax = plt.subplots()

    comp_mat = np.zeros((len(pkg_str), len(pkg_str)))

    for p1 in range(len(pkg_str)):

        dat_file_p1 = data_path + '{}/{}/exp_mat_{}_{}.npy'.format(task, pkg_str[p1], pkg_str[p1], prec)
        # print(dat_file_p1)

        for p2 in range(len(pkg_str)):

            dat_file_p2 = data_path + '{}/{}/exp_mat_{}_{}.npy'.format(task, pkg_str[p2], pkg_str[p2], prec)

            if os.path.isfile(dat_file_p1) and os.path.isfile(dat_file_p2):

                dat_p1 = np.load(dat_file_p1)
                dat_p2 = np.load(dat_file_p2)

                comp_mat[p1, p2] = np.log10(np.sum(np.abs(dat_p1 - dat_p2))/float(16))
                # print(comp_mat[p1, p2])
            else:
                continue

    if prec == 'sp':
        plt.imshow(comp_mat, cmap='Spectral', vmin=-16, vmax=-3)
    elif prec == 'dp':
        plt.imshow(comp_mat, cmap='Spectral', vmin=-17, vmax=-11)

    plt.xticks(range(len(pkg_str)), pkg_str)
    locs, labels = plt.xticks()

    plt.setp(labels, rotation=90)
    plt.yticks(range(len(pkg_str)), pkg_str)
    locs, labels = plt.yticks()
    plt.colorbar()
    plt.tight_layout()
    # plt.savefig(fn)
    plt.show()

package_str_sp = ['cirq', 'hybridq', 'intel_qs_cpp', 'pennylane', 'pennylane_l', 'qibo', 'qibojit', 'qiskit', 'qsimcirq', 'quest', 'yao', 'qiskit_gpu', 'hybridq_gpu', 'qibojit_gpu', 'qsimcirq_gpu', 'yao_gpu']#, 'pennylane_l_gpu', 'quest_gpu', 'projectq', 'qulacs', 'qulacs_gpu']

sorted_package_str_sp = ['cuquantum_qiskit', 'qiskit', 'qiskit_gpu', 'pennylane', 'cirq', 'qsimcirq', 'qsimcirq_gpu', 'qibojit_gpu', 'qibo', 'qibojit', 'quest', 'intel_qs_cpp', 'myqlm_cpp', 'cuquantum_qsimcirq', 'pennylane_l', 'hybridq_gpu', 'hybridq', 'yao', 'yao_gpu']

# compare_two_expect_all('hdyn', 'sp', sorted_package_str_sp)#, 'compare_packages_qft_sp_N16_log.pdf')

package_str_dp = ['cirq', 'intel_qs_cpp', 'pennylane', 'pennylane_l', 'qibo', 'qibojit', 'qiskit', 'quest', 'yao', 'qiskit_gpu', 'qibojit_gpu', 'yao_gpu', 'pennylane_l_gpu', 'quest_gpu', 'projectq', 'qulacs', 'qulacs_gpu']

sorted_package_str_dp = ['qibojit_gpu', 'qulacs', 'qiskit', 'cuquantum_qiskit', 'qiskit_gpu', 'qpanda', 'intel_qs_cpp', 'quest', 'qibo', 'qibojit', 'pennylane_l_gpu', 'pennylane', 'yao_gpu', 'myqlm_cpp', 'myqlm', 'braket', 'yao', 'pennylane_l', 'cirq', 'quest_gpu', 'projectq', 'qulacs_gpu']

# compare_two_expect_all('hdyn', 'dp', sorted_package_str_dp)#, 'compare_packages_qft_dp_N16_log.pdf')