/**
 * @brief CUDA optimizing strategy.
 *
 * @author Miao Zhang
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */
#ifdef WITH_CUDA
#ifndef CUDA_CUDACONFIG_H
#define CUDA_CUDACONFIG_H
#include<cuda_runtime.h>
#include<cuComplex.h>
#include<complex>
#include<iostream>
#include<set>
#include<vector>
#include<eigen3/Eigen/Dense>
#include"../base/utils.h"
#include"base/datatypes.h"
using namespace tbplas::base;
__host__ cudaError_t checkError(cudaError_t err_code, const char* filename, int lineNumber);

template<typename T>
__host__ void act_GPU(
	const T* bessel_series_H, const int num_series,
	const dim3 block,
	const dim3 grid,
	cuDoubleComplex* t0_D, const size_t t0_size,
	cuDoubleComplex* t1_D, const size_t t1_size,
	cuDoubleComplex* wf_in_D, const size_t wf_in_size,
	cuDoubleComplex* wf_out_D, const size_t wf_out_size,
	const uint64_t* indptr_D, const  uint64_t* indices_D, const  cuDoubleComplex* values_D,
	const size_t dim);


#define CHECKERROR(err_code) \
	do{if (err_code != cudaSuccess) {\
		printf("CUDA error:\r\ncode=%d,name=%s,description=%s\r\nfile=%s,line=%d\r\n", \
			err_code, cudaGetErrorName(err_code), cudaGetErrorString(err_code), \
			__FILE__, __LINE__); \
		exit(1);}\
	}while (0);

__host__ void checkGPUInfo(int iDev);


class CudaConfig {
public:
	CudaConfig() {
		int iDeviceCount = 0, iMaxGPU = 1;
		cudaGetDeviceCount(&iDeviceCount);
		if(iDeviceCount==0){
			printf("No GPU device exists..\n");
			exit(0);
		}
		for (size_t i = 0; i < iDeviceCount; i++) {
			checkGPUInfo(i);
		}
		//cudaSetDevice(iMaxGPU);
	};
	~CudaConfig() {
	};
	__host__ void calc_dyn_pol_GPU(
		complex_t* dyn_pol_H, const int dyn_pol_width, const int dyn_pol_height,
		complex_t* product_H, double* delta_eng_H,
		const int delta_eng_product_size, const int omegas_size,
		const double delta, const int iq, double* omegas_H);

	__host__   void tbpm_qe_dos_GPU(
		const bool calc_corr_dos, const size_t num_steps,
		const double time_step, const double* qe_energies_H,
		const size_t  num_energies, const  size_t  wfn_check_time_steps,
		const  double wfn_check_threshold,
		const complex_t* wf0_H, const size_t wf0_size,
		const complex_t* wf_t_pos_H, const size_t wf_t_pos_size,
		const complex_t* wf_t_neg_H, const size_t wf_t_neg_size,
		const complex_t norm_ref,
		const complex_t* bessel_series_T_H,
		const complex_t* bessel_series_F_H,
		const int bessel_series_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const size_t dim_size,
		complex_t* corr_H, const size_t corr_size,
		std::vector<Eigen::VectorXcd>& wfq);

	__host__ void tbpm_dos_GPU(
		complex_t* wf_t_H, const size_t wf_t_size,
		const size_t num_steps, const  size_t  wfn_check_time_steps, const  double wfn_check_threshold,
		const complex_t* wf0_H, const size_t wf0_size,
		const complex_t norm_ref,
		const complex_t* bessel_series_H,
		const int bessel_series_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		complex_t* wf_out_H, const size_t wf_out_size,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const size_t dim_size,
		complex_t* corr_H, const size_t corr_size,
		ProgressBar& prog_bar);

	/*__host__   void tbpm_dos_GPU(
		complex_t* wf_t_H, const int wf_t_size,
		const int num_steps, const  int  wfn_check_time_steps, const  double wfn_check_threshold,
		const complex_t* wf0_H, const int wf0_size,
		const complex_t norm_ref,
		const complex_t* bessel_series_H,
		const int bessel_series_size,
		complex_t* t0_H, const int t0_size,
		complex_t* t1_H, const int t1_size,
		complex_t* wf_out_H, const int wf_out_size,
		const uint64_t* indptr_H, const int indptr_size,
		const uint64_t* indices_H, const int indices_size,
		const complex_t* values_H, const int values_size,
		const size_t dim_size,
		complex_t* corr_H, const int corr_size);*/

	__host__ void tbpm_ac_cond_GPU(
		const size_t num_steps, const  size_t  wfn_check_time_steps, const  double wfn_check_threshold,
		const complex_t norm_ref,
		const complex_t* bessel_series_H, const size_t bessel_series_size,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const uint64_t* curr_x_indptr_H, const size_t curr_x_indptr_size,
		const uint64_t* curr_x_indices_H, const size_t curr_x_indices_size,
		const complex_t* curr_x_values_H, const size_t curr_x_values_size,
		const uint64_t* curr_y_indptr_H, const size_t curr_y_indptr_size,
		const uint64_t* curr_y_indices_H, const size_t curr_y_indices_size,
		const complex_t* curr_y_values_H, const size_t curr_y_values_size,
		const size_t dim_size,
		complex_t* psi2_H, const size_t psi2_size,
		complex_t* psi1_x_H, const size_t psi1_x_size,
		complex_t* psi1_y_H, const size_t psi1_y_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		complex_t* corr_H, const size_t corr_size,
		const bool isRowMajor,
		ProgressBar& prog_bar);

	__host__   void tbpm_dc_cond_GPU(
		const size_t num_steps_dc,
		const complex_t* bessel_series_H, const size_t bessel_series_size,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const uint64_t* curr_x_indptr_H, const size_t curr_x_indptr_size,
		const uint64_t* curr_x_indices_H, const size_t curr_x_indices_size,
		const complex_t* curr_x_values_H, const size_t curr_x_values_size,
		const uint64_t* curr_y_indptr_H, const size_t curr_y_indptr_size,
		const uint64_t* curr_y_indices_H, const size_t curr_y_indices_size,
		const complex_t* curr_y_values_H, const size_t curr_y_values_size,
		const size_t dim_size,
		double* factor_H, const size_t factor_size,
		complex_t* psi0_x_H, const size_t psi0_x_size,
		complex_t* psi0_y_H, const size_t psi0_y_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		complex_t* corr_dc_xx_H, const size_t corr_dc_xx_size,
		complex_t* corr_dc_yy_H, const size_t corr_dc_yy_size,
		std::vector<Eigen::VectorXcd>& wf_qe,
		const bool isRowMajor,
		const size_t num_energies,
		ProgressBar& prog_bar);

	__host__   void tbpm_dyn_pol_GPU(
		const size_t num_steps, const  size_t  wfn_check_time_steps, const  double wfn_check_threshold,
		const complex_t norm_ref,
		const complex_t* bessel_series_H, const size_t bessel_series_size,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const size_t dim_size,
		complex_t* psi1_H, const size_t psi1_size,
		complex_t* psi2_H, const size_t psi2_size,
		const complex_t* den_coeff_q_H, const size_t den_coeff_q_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		double* corr_H, const size_t corr_size);

	__host__   void tbpm_wft_GPU(
		const size_t num_steps, const  size_t  wfn_check_time_steps,
		const  double wfn_check_threshold,
		const std::set<size_t>& wft_time_save,
		const complex_t* wf0_H, const size_t wf0_size,
		const complex_t norm_ref,
		const complex_t* bessel_series_H, const int bessel_series_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const size_t dim_size,
		complex_t* corr_H, const size_t corr_size,
		const size_t corr_colWidth, const bool isRowMajor);

	__host__   void haydock_coef_GPU(
		const double rescale, const int n_depth,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		complex_t* p1_H, const size_t p1_size,
		complex_t* p2_H, const size_t p2_size,
		complex_t* p_buf_H, const size_t p_buf_size,
		complex_t* coefa_H, const size_t coefa_size,
		double* coefb_H, const size_t coefb_size);

	__host__   void kbdc_mu_GPU(
		const complex_t* wf0_H, const int wf0_size,
		const int dckb_component, const int num_kernel,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const uint64_t* curr_x_indptr_H, const size_t curr_x_indptr_size,
		const uint64_t* curr_x_indices_H, const size_t curr_x_indices_size,
		const complex_t* curr_x_values_H, const size_t curr_x_values_size,
		const uint64_t* curr_y_indptr_H, const size_t curr_y_indptr_size,
		const uint64_t* curr_y_indices_H, const size_t curr_y_indices_size,
		const complex_t* curr_y_values_H, const size_t curr_y_values_size,
		std::vector<Eigen::VectorXcd>& buffer,
		std::vector<Eigen::VectorXcd>& wf_vb_tn,
		complex_t* mu_sample_H, const bool isRowMajor,
		ProgressBar& prog_bar);

	__host__   void kbdc_mu_GPU2(
		const complex_t* wf0_H, const int wf0_size,
		const int dckb_component, const int num_kernel,
		const uint64_t* indptr_H, const size_t indptr_size,
		const uint64_t* indices_H, const size_t indices_size,
		const complex_t* values_H, const size_t values_size,
		const uint64_t* curr_x_indptr_H, const size_t curr_x_indptr_size,
		const uint64_t* curr_x_indices_H, const size_t curr_x_indices_size,
		const complex_t* curr_x_values_H, const size_t curr_x_values_size,
		const uint64_t* curr_y_indptr_H, const size_t curr_y_indptr_size,
		const uint64_t* curr_y_indices_H, const size_t curr_y_indices_size,
		const complex_t* curr_y_values_H, const size_t curr_y_values_size,
		complex_t* t0_H, const size_t t0_size,
		complex_t* t1_H, const size_t t1_size,
		complex_t* mu_sample_H, const bool isRowMajor,
		ProgressBar& prog_bar,
		std::vector<Eigen::VectorXcd>& wf_vb_tn);

private:
	__host__ complex_t vec_dot_GPU(
		const dim3 block, const dim3 grid,
		const cuDoubleComplex* x, const size_t x_size,
		const cuDoubleComplex* y, const size_t y_size,
		cuDoubleComplex* tmp_D, cuDoubleComplex* tmp_H);

	__host__ complex_t xmy_GPU(
		const dim3 block, const dim3 grid,
		const cuDoubleComplex* x, const size_t x_size,
		const cuDoubleComplex* y, const size_t y_size, const size_t dim,
		const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
		cuDoubleComplex* tmp_D, cuDoubleComplex* tmp_H);

	__host__ complex_t xmy_diag_GPU(
		const dim3 block, const dim3 grid,
		const cuDoubleComplex* x, const cuDoubleComplex* y,
		const cuDoubleComplex* m, const size_t dim,
		cuDoubleComplex* tmp_D, cuDoubleComplex* tmp_H);

	__host__ double  vec_norm_GPU(
		const dim3 block, const dim3 grid,
		const cuDoubleComplex* x, const size_t x_size,
		cuDoubleComplex* tmp_D, cuDoubleComplex* tmp_H);
};
#endif //CUDA_CUDACONFIG_H
#endif //WITH_CUDA