#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <omp.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 16

void dgemm_(char*, char*, int*, int*, int*, double*, double*, int*, double*, int*,
		double*, double*, int*);




void gpu_dgemm(char transa, char transb, int m, int n, int k, double alpha,
                  double* A, int lda, double* B, int ldb, double beta, double* C, int ldc)
{
    cublasHandle_t handle;
    cublasCreate(&handle);

    double * d_A;
    double * d_B;
    double * d_C;
    cublasOperation_t ta, tb;

    if (transa == 'N') {
      cudaMalloc((void**)&d_A, lda*k*sizeof(double));
      cublasSetMatrix(m, k, sizeof(double), A, lda, d_A, lda);
      ta = CUBLAS_OP_N;
    } else {
      cudaMalloc((void**)&d_A, lda*m*sizeof(double));
      cublasSetMatrix(k, m, sizeof(double), A, lda, d_A, lda);
      ta = CUBLAS_OP_T;
    }

    if (transb == 'N') {
      cudaMalloc((void**)&d_B, ldb*n*sizeof(double));
      cublasSetMatrix(k, n, sizeof(double), B, ldb, d_B, ldb);
      tb = CUBLAS_OP_N;
    } else {
      cudaMalloc((void**)&d_B, ldb*k*sizeof(double));
      cublasSetMatrix(n, k, sizeof(double), B, ldb, d_B, ldb);
      tb = CUBLAS_OP_T;
    }

    cudaMalloc((void**)&d_C, ldc*n*sizeof(double));
    if (beta != 0.) {
      cublasSetMatrix(m, n, sizeof(double), C, ldc, d_C, ldc);
    }

    cublasDgemm(handle, ta, tb, m, n, k, &alpha, d_A, lda, d_B, ldb, &beta, d_C, ldc);

    cublasGetMatrix(m, n,  sizeof(double), d_C, ldc, C, ldc);

    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    cublasDestroy(handle);
}

void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo_num,
                double* t1,
                double* t2,
                double* tau,
                double* cc_space_v_oo_chol,
                double* cc_space_v_ov_chol,
                double* cc_space_v_vo_chol,
                double* cc_space_v_vv_chol,
                double* cc_space_v_oooo,
                double* cc_space_v_vooo,
                double* cc_space_v_oovv,
                double* cc_space_v_vvoo,
                double* cc_space_v_oovo,
                double* cc_space_v_ovvo,
                double* cc_space_f_vo,
                double* H_vv,
                double* g_occ,
                double* r2)
{

    int ngpus;
    cudaGetDeviceCount(&ngpus);
    #pragma omp parallel num_threads(ngpus)
    {
      int m,n,k, lda, ldb, ldc;
      double alpha, beta;
      double* A;
      double* B;
      double* C;

      int ithread = omp_get_thread_num();
      int igpu = ithread ;
//igpu=1;
      cudaSetDevice(igpu);
      cublasHandle_t handle;

      cublasCreate(&handle);

      double* d_tau;
      double* d_r2;
      double* d_cc_space_v_oo_chol;
      double* d_cc_space_v_ov_chol;
      double* d_cc_space_v_vo_chol;
      double* d_cc_space_v_vv_chol;
      double* d_t1;
      double* d_t2;
      double* d_tmp_cc;

      lda = nO * nO;
      cudaMalloc((void **)&d_tau, lda * nV * nV * sizeof(double));
      cublasSetMatrix(nO*nO, nV*nV, sizeof(double), tau, lda, d_tau, lda);

      lda = nO * nO;
      cudaMalloc((void **)&d_r2, lda * nV * nV * sizeof(double));
      memset(r2, 0, nO*nO*nV*nV*sizeof(double));
      cublasSetMatrix(nO*nO, nV*nV, sizeof(double), r2, lda, d_r2, lda);

      lda = cholesky_mo_num * nV;
      cudaMalloc((void **)&d_cc_space_v_vv_chol, lda * nV * sizeof(double));
      cublasSetMatrix(cholesky_mo_num*nV, nV, sizeof(double), cc_space_v_vv_chol, lda, d_cc_space_v_vv_chol, lda);

      lda = cholesky_mo_num * nO;
      cudaMalloc((void **)&d_cc_space_v_oo_chol, lda * nO * sizeof(double));
      cublasSetMatrix(cholesky_mo_num*nO, nO, sizeof(double), cc_space_v_oo_chol, lda, d_cc_space_v_oo_chol, lda);

      lda = cholesky_mo_num * nO;
      cudaMalloc((void **)&d_cc_space_v_ov_chol, lda * nV * sizeof(double));
      cublasSetMatrix(cholesky_mo_num*nO, nV, sizeof(double), cc_space_v_ov_chol, lda, d_cc_space_v_ov_chol, lda);

      lda = cholesky_mo_num * nV;
      cudaMalloc((void **)&d_cc_space_v_vo_chol, lda * nO * sizeof(double));
      cublasSetMatrix(cholesky_mo_num*nV, nO, sizeof(double), cc_space_v_vo_chol, lda, d_cc_space_v_vo_chol, lda);

      lda = nO;
      cudaMalloc((void **)&d_t1, nO * nV * sizeof(double));
      cublasSetMatrix(nO, nV, sizeof(double), t1, lda, d_t1, lda);

      lda = nO*nO;
      cudaMalloc((void **)&d_t2, nO*nO*nV*nV * sizeof(double));
      cublasSetMatrix(nO*nO, nV*nV, sizeof(double), t2, lda, d_t2, lda);

      double* d_cc_space_f_vo;
      cudaMalloc((void**)&d_cc_space_f_vo, nV*nO*sizeof(double));
      cublasSetMatrix(nV, nO, sizeof(double), cc_space_f_vo, nV, d_cc_space_f_vo, nV);

      #pragma omp sections
      {

        #pragma omp section
        {
           double* d_cc_space_v_vooo;
           cudaMalloc((void**)&d_cc_space_v_vooo, nV*nO*nO*nO*sizeof(double));
           cublasSetMatrix(nV*nO, nO*nO, sizeof(double), cc_space_v_vooo, nV*nO, d_cc_space_v_vooo, nV*nO);

           double* d_Y_oooo;
           cudaMalloc((void**)&d_Y_oooo, nO*nO*nO*nO*sizeof(double));

           alpha = 1.0;
           beta =  0.0;
           m=nO ; n=nO*nO*nO; k=nV;
           A = d_t1 ; lda = nO;
           B = d_cc_space_v_vooo ; ldb = nV;
           C = d_Y_oooo; ldc = nO;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_cc_space_v_vooo);

           double* d_A1;
           cudaMalloc((void**)&d_A1, nO*nO*nO*nO*sizeof(double));

           double* d_cc_space_v_oooo;
           cudaMalloc((void**)&d_cc_space_v_oooo, nO*nO*nO*nO*sizeof(double));
           cublasSetMatrix(nO*nO, nO*nO, sizeof(double), cc_space_v_oooo, nO*nO, d_cc_space_v_oooo, nO*nO);

           alpha = 1.0;
           beta =  1.0;
           A = d_cc_space_v_oooo; lda = nO*nO;
           B = d_Y_oooo; ldb = nO*nO;
           C = d_A1; ldc = nO*nO;
           cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO*nO, nO*nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
           for (int j=0 ; j<nO ; ++j) {
             for (int i=0 ; i<nO ; ++i) {
                alpha = 1.0;
                beta =  1.0;
                A = &(d_A1[nO*nO*(i+nO*j)]); lda = nO;
                B = &(d_Y_oooo[nO*nO*(j+nO*i)]); ldb = nO;
                C = &(d_A1[nO*nO*(i+nO*j)]); ldc = nO;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);

             }
           }
           cudaFree(d_Y_oooo);

           double* d_cc_space_v_vvoo;
           cudaMalloc((void**)&d_cc_space_v_vvoo, nV*nV*nO*nO*sizeof(double));
           cublasSetMatrix(nV*nV, nO*nO, sizeof(double), cc_space_v_vvoo, nV*nV, d_cc_space_v_vvoo, nV*nV);

           alpha = 1.0;
           beta =  1.0;
           m=nO*nO ; n=nO*nO; k=nV*nV;
           A = d_tau ; lda = nO*nO;
           B = d_cc_space_v_vvoo ; ldb = nV*nV;
           C = d_A1; ldc = nO*nO;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_cc_space_v_vvoo);

           alpha = 1.0;
           beta =  0.0;
           m=nO*nO ; n=nV*nV; k=nO*nO;
           A = d_A1  ; lda = nO*nO;
           B = d_tau ; ldb = nO*nO;
           C = d_r2; ldc = nO*nO;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_A1);
        }

        // g_vir
        #pragma omp section
        {
           double* d_g_vir;
           cudaMalloc((void**)&d_g_vir, nV*nV*sizeof(double));
           cublasSetMatrix(nV, nV, sizeof(double), H_vv, nV, d_g_vir, nV);

           alpha = -1.0;
           beta =  1.0;
           m=nV ; n=nV; k=nO;
           A = d_cc_space_f_vo  ; lda = nV;
           B = d_t1 ; ldb = nO;
           C = d_g_vir; ldc = nV;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

           double* d_tmp_k;
           cudaMalloc((void**)&d_tmp_k, cholesky_mo_num*sizeof(double));
           alpha = 1.0;
           beta =  0.0;
           m=cholesky_mo_num ; n=1; k=nO*nV;
           A = d_cc_space_v_ov_chol; lda = cholesky_mo_num;
           B = d_t1 ; ldb = nO*nV;
           C = d_tmp_k; ldc = cholesky_mo_num;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

           alpha = 2.0;
           beta =  1.0;
           m=nV*nV; n=1; k=cholesky_mo_num;
           A = d_cc_space_v_vv_chol; lda = cholesky_mo_num;
           B = d_tmp_k ; ldb = cholesky_mo_num;
           C = d_g_vir; ldc = nV*nV;
           cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_tmp_k);

           double* d_tmp_vo;
           cudaMalloc((void**)&d_tmp_vo, cholesky_mo_num*nV*nO*sizeof(double));
           alpha = 1.0;
           beta =  0.0;
           m=cholesky_mo_num*nV ; n=nO; k=nV;
           A = d_cc_space_v_vv_chol; lda = cholesky_mo_num*nV;
           B = d_t1 ; ldb = nO;
           C = d_tmp_vo; ldc = cholesky_mo_num*nV;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

           double* d_tmp_vo2;
           cudaMalloc((void**)&d_tmp_vo2, cholesky_mo_num*nV*nO*sizeof(double));
           for (int i=0 ; i<nO ; ++i) {
             alpha = -1.0;
             beta =  0.0;
             A = &(d_tmp_vo[cholesky_mo_num*nV*i]); lda = cholesky_mo_num;
             B = &(d_tmp_vo[cholesky_mo_num*nV*i]); ldb = cholesky_mo_num;
             C = &(d_tmp_vo2[cholesky_mo_num*i]); ldc = cholesky_mo_num*nO;
             cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, cholesky_mo_num, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
           }
           cudaFree(d_tmp_vo);

           alpha = 1.0;
           beta =  1.0;
           m=nV ; n=nV; k=nO*cholesky_mo_num;
           A = d_cc_space_v_ov_chol; lda = cholesky_mo_num*nO;
           B = d_tmp_vo2 ; ldb = cholesky_mo_num*nO;
           C = d_g_vir; ldc = nV;
           cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_tmp_vo2);

           double* d_Y_oovv;
           cudaMalloc((void**)&d_Y_oovv, nO*nO*nV*nV*sizeof(double));
           alpha = 1.0;
           beta =  0.0;
           m=nO*nO*nV ; n=nV; k=nV;
           A = d_t2; lda = nO*nO*nV;
           B = d_g_vir; ldb = nV;
           C = d_Y_oovv; ldc = nO*nO*nV;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_g_vir);

           alpha = 1.0;
           beta =  1.0;
           A = d_r2; lda = nO*nO;
           B = d_Y_oovv; ldb = nO*nO;
           C = d_r2; ldc = nO*nO;
           cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO*nO, nV*nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
           for (int j=0 ; j<nV ; ++j) {
             for (int i=0 ; i<nV ; ++i) {
                alpha = 1.0;
                beta =  1.0;
                A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
                B = &(d_Y_oovv[nO*nO*(j+nV*i)]); ldb = nO;
                C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);

             }
           }
           cudaFree(d_Y_oovv);
        }

        // g_occ
        #pragma omp section
        {
          double* d_g_occ;
          lda = nO;
          cudaMalloc((void **)&d_g_occ, nO*nO * sizeof(double));
          cublasSetMatrix(lda, nO, sizeof(double), g_occ, lda, d_g_occ, lda);

          alpha = 1.0;
          beta  = 1.0;
          m=nO;  n=nO;  k=nV;
          A=d_t1;  lda=nO;
          B=d_cc_space_f_vo;  ldb=nV;
          C=d_g_occ;  ldc=nO;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

          double* d_X_oovv;
          lda = nO*nO;
          cudaMalloc((void **)&d_X_oovv, nO*nO*nV*nV * sizeof(double));

          alpha = 1.0;
          beta  = 0.0;
          m=nO;  n=nO*nV*nV;  k=nO;
          A=d_g_occ;  lda=nO;
          B=d_t2;  ldb=nO;
          C=d_X_oovv;  ldc=nO;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
          cudaFree(d_g_occ);

           alpha = 1.0;
           beta = -1.0;
           A = d_r2; lda = nO*nO;
           B = d_X_oovv; ldb = nO*nO;
           C = d_r2; ldc = nO*nO;
           cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO*nO, nV*nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
           for (int j=0 ; j<nV ; ++j) {
             for (int i=0 ; i<nV ; ++i) {
                alpha = 1.0;
                beta = -1.0;
                A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
                B = &(d_X_oovv[nO*nO*(j+nV*i)]); ldb = nO;
                C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
             }
           }
           cudaFree(d_X_oovv);
        }

        #pragma omp section
        {
           double* d_X_vovv;
           lda = nV*nO;
           cudaMalloc((void **)&d_X_vovv, nV*nO*nV*BLOCK_SIZE * sizeof(double));

           double* d_Y_oovv;
           lda = nO*nO;
           cudaMalloc((void **)&d_Y_oovv, nO*nO*nV*nV * sizeof(double));

           for (int iblock=0 ; iblock<nV ; iblock += BLOCK_SIZE) {
             int mbs = nV < iblock+BLOCK_SIZE ? nV : iblock+BLOCK_SIZE;
             for (int gam=iblock ; gam<mbs ; ++gam) {
               alpha = 1.0;
               beta  = 0.0;
               m=nV;  n=nO*nV;  k=cholesky_mo_num;
               A=&(d_cc_space_v_vv_chol[cholesky_mo_num*nV*gam]);  lda=cholesky_mo_num;
               B=d_cc_space_v_ov_chol;  ldb=cholesky_mo_num;
               C=&(d_X_vovv[nV*nO*nV*(gam-iblock)]);  ldc=nV;
               cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
             }
             mbs = BLOCK_SIZE < nV-iblock ? BLOCK_SIZE : nV-iblock;
             alpha = 1.0;
             beta  = 0.0;
             m=nO;  n=nO*nV*mbs;  k=nV;
             A=d_t1;  lda=nO;
             B=d_X_vovv;  ldb=nV;
             C=&(d_Y_oovv[nO*nO*nV*iblock]);  ldc=nO;
             cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           }
           cudaFree(d_X_vovv);
           for (int j=0 ; j<nV ; ++j) {
             for (int i=0 ; i<nV ; ++i) {
                alpha = 1.0;
                beta =  1.0;
                A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
                B = &(d_Y_oovv[nO*nO*(i+nV*j)]); ldb = nO;
                C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
                A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
                B = &(d_Y_oovv[nO*nO*(j+nV*i)]); ldb = nO;
                C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
             }
           }
           cudaFree(d_Y_oovv);
        }

        #pragma omp section
        {
          double* d_tcc2;
          lda = cholesky_mo_num;
          cudaMalloc((void **)&d_tcc2, cholesky_mo_num*nV*nO * sizeof(double));

          alpha = 1.0;
          beta  = 0.0;
          m=cholesky_mo_num*nV;  n=nO;  k=nV;
          A=d_cc_space_v_vv_chol;  lda=cholesky_mo_num*nV;
          B=d_t1;  ldb=nO;
          C=d_tcc2;  ldc=cholesky_mo_num*nV;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

          double* d_tcc;
          lda = cholesky_mo_num;
          cudaMalloc((void **)&d_tcc, cholesky_mo_num*nO*nV * sizeof(double));

          alpha = 1.0;
          beta  = 0.0;
          m=cholesky_mo_num*nO;  n=nV;  k=nO;
          A=d_cc_space_v_oo_chol;  lda=cholesky_mo_num*nO;
          B=d_t1;  ldb=nO;
          C=d_tcc;  ldc=cholesky_mo_num*nO;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

          double* d_X_ovvo;
          lda = nO*nV;
          cudaMalloc((void **)&d_X_ovvo, nO*nV*nV*nO * sizeof(double));

          alpha = 1.0;
          beta  = 0.0;
          m=nO*nV;  n=nV*nO;  k=cholesky_mo_num;
          A=d_tcc;  lda=cholesky_mo_num;
          B=d_tcc2;  ldb=cholesky_mo_num;
          C=d_X_ovvo;  ldc=nO*nV;
          cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
          cudaFree(d_tcc);
          cudaFree(d_tcc2);

          for(int gam = 0; gam < nV; gam++){
            for(int bet = 0; bet < nV; bet++){
               alpha = 1.0;
               beta = -1.0;
               A = &(d_r2[nO*nO*(bet+nV*gam)]); lda = nO;
               B = &(d_X_ovvo[nO*(bet+nV*gam)]); ldb = nO*nV*nV;
               C = &(d_r2[nO*nO*(bet+nV*gam)]); ldc = nO;
               cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
               B = &(d_X_ovvo[nO*(gam+nV*bet)]); ldb = nO*nV*nV;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
            }
          }

          cudaFree(d_X_ovvo);
        }

        #pragma omp section
        {
          double* d_cc_space_v_oovo;
          lda = nO*nO;
          cudaMalloc((void **)&d_cc_space_v_oovo, nO*nO*nV*nO * sizeof(double));
          cublasSetMatrix(lda, nV*nO, sizeof(double), cc_space_v_oovo, lda, d_cc_space_v_oovo, lda);

          double* d_X_oovv;
          lda = nO*nO;
          cudaMalloc((void **)&d_X_oovv, nO*nO*nV*nV * sizeof(double));

          alpha = 1.0;
          beta  = 0.0;
          m=nO*nO*nV;  n=nV;  k=nO;
          A=d_cc_space_v_oovo;  lda=nO*nO*nV;
          B=d_t1;  ldb=nO;
          C=d_X_oovv;  ldc=nO*nO*nV;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

          alpha = 1.0;
          beta = -1.0;
          A = d_r2; lda = nO*nO;
          B = d_X_oovv; ldb = nO*nO;
          C = d_r2; ldc = nO*nO;
          cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO*nO, nV*nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
          for (int j=0 ; j<nV ; ++j) {
            for (int i=0 ; i<nV ; ++i) {
               alpha = 1.0;
               beta = -1.0;
               A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
               B = &(d_X_oovv[nO*nO*(j+nV*i)]); ldb = nO;
               C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
               cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
            }
          }
          cudaFree(d_cc_space_v_oovo);

          double* d_cc_space_v_ovvo;
          lda = nO*nV;
          cudaMalloc((void **)&d_cc_space_v_ovvo, nO*nV*nV*nO * sizeof(double));
          cublasSetMatrix(lda, nV*nO, sizeof(double), cc_space_v_ovvo, lda, d_cc_space_v_ovvo, lda);

          double* d_X_vovo;
          lda = nV*nO;
          cudaMalloc((void **)&d_X_vovo, nV*nO*nV*nO * sizeof(double));

          alpha = 0.0;
          beta =  1.0;
          for (int i=0 ; i<nO ; ++i) {
            for (int gam=0 ; gam<nV ; ++gam) {
               A = &(d_X_vovo[nV*nO*(gam+nV*i)]); lda = nV;
               B = &(d_cc_space_v_ovvo[nO*nV*(gam+nV*i)]); ldb = nO;
               C = &(d_X_vovo[nV*nO*(gam+nV*i)]); ldc = nV;
               cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nV, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
            }
          }
          cudaFree(d_cc_space_v_ovvo);

          double* d_Y_oovo;
          lda = nO*nO;
          cudaMalloc((void **)&d_Y_oovo, nO*nO*nV*nO * sizeof(double));

          alpha = 1.0;
          beta  = 0.0;
          m=nO;  n=nO*nV*nO;  k=nV;
          A=d_t1;  lda=nO;
          B=d_X_vovo;  ldb=nV;
          C=d_Y_oovo;  ldc=nO;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

          cudaFree(d_X_vovo);

          alpha = 1.0;
          beta  = 0.0;
          m=nO*nO*nV;  n=nV;  k=nO;
          A=d_Y_oovo;  lda=nO*nO*nV;
          B=d_t1;  ldb=nO;
          C=d_X_oovv;  ldc=nO*nO*nV;
          cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

          cudaFree(d_Y_oovo);

          for (int j=0 ; j<nV ; ++j) {
            for (int i=0 ; i<nV ; ++i) {
               alpha = 1.0;
               beta = -1.0;
               A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
               B = &(d_X_oovv[nO*nO*(i+nV*j)]); ldb = nO;
               C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
               cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
               B = &(d_X_oovv[nO*nO*(j+nV*i)]); ldb = nO;
               cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
            }
          }
          cudaFree(d_X_oovv);
        }
      }


      lda = cholesky_mo_num * nV;
      cudaMalloc((void **)&d_tmp_cc, lda * nV * sizeof(double));

      alpha=1.0; beta=0.0;
      m=cholesky_mo_num*nV;  n=nV;  k=nO;
      A = d_cc_space_v_vo_chol; B = d_t1; C = d_tmp_cc;
      cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, m, B, k, &beta, C, m);

      double* d_tmp_cc2;
      cudaMalloc((void **)&d_tmp_cc2, cholesky_mo_num*nV*sizeof(double));

      double* d_B1;
      cudaMalloc((void**)&d_B1, nV*nV*BLOCK_SIZE*sizeof(double));

      double* d_tmpB1;
      cudaMalloc((void**)&d_tmpB1, nV*BLOCK_SIZE*nV*sizeof(double));

      #pragma omp for
      for (int gam=0 ; gam<nV ; ++gam)
      {
         double* d_tmp_cc_ = &(d_tmp_cc[gam*nV*cholesky_mo_num]);
         double* d_cc_space_v_vv_chol_ = &(d_cc_space_v_vv_chol[gam*nV*cholesky_mo_num]);

         alpha = 1.0;
         beta = -1.0;
         A = d_cc_space_v_vv_chol_; lda = cholesky_mo_num;
         B = d_tmp_cc_; ldb = cholesky_mo_num;
         C = d_tmp_cc2 ; ldc = cholesky_mo_num;
         cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, cholesky_mo_num, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);

         for (int iblock=0 ; iblock<nV ; iblock += BLOCK_SIZE)
         {
              const int mbs = BLOCK_SIZE < nV-iblock ? BLOCK_SIZE : nV-iblock;

              alpha=-1.0; beta=0.0;
              m=nV*mbs;  n=nV;  k=cholesky_mo_num;

              A=&(d_tmp_cc[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
              B=d_cc_space_v_vv_chol_; ldb=cholesky_mo_num;
              C=d_tmpB1 ; ldc=nV*BLOCK_SIZE;
              cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, lda, &beta, C, ldc);

              alpha=1.0; beta=1.0;
              m=nV*mbs;  n=nV;  k=cholesky_mo_num;

              A=&(d_cc_space_v_vv_chol[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
              B=d_tmp_cc2; ldb=cholesky_mo_num;
              C=d_tmpB1 ; ldc=nV*BLOCK_SIZE;
              cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, lda, &beta, C, ldc);

              for (size_t bet=iblock ; bet<(nV < iblock+BLOCK_SIZE ? nV : iblock+BLOCK_SIZE) ; ++bet)
              {

                alpha = 1.0;
                beta = 0.0;
                A = &(d_tmpB1[nV*(bet-iblock)]); lda = nV*BLOCK_SIZE;
                B = d_tmpB1; ldb = nV;
                C = &(d_B1[nV*nV*(bet-iblock)]) ; ldc = nV;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nV, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
             }

             alpha=1.0; beta=1.0;
             m=nO*nO;  n=mbs;  k=nV*nV;
             A=d_tau; lda=nO*nO;
             B=d_B1 ; ldb=nV*nV;
             C=&(d_r2[nO*nO*(iblock + nV*gam)]); ldc=nO*nO;
             cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

         }
      }
      cudaFree(d_tmpB1);
      cudaFree(d_B1);
      cudaFree(d_tmp_cc2);

      cudaFree(d_cc_space_v_vo_chol);
      cudaFree(d_cc_space_v_vv_chol);
      cudaFree(d_tau);
      cudaFree(d_t1);
      cudaFree(d_tmp_cc);
      double * r2_tmp = malloc(nO*nO*nV*nV*sizeof(double));
      lda=nO*nO;
      cublasGetMatrix(nO*nO, nV*nV, sizeof(double), d_r2, lda, r2_tmp, lda);
      #pragma omp critical
      {
         for (size_t i=0 ; i<(size_t) nO*nO*nV*nV ; ++i) {
           r2[i] += r2_tmp[i];
         }
      }
      free(r2_tmp);

      cudaFree(d_r2);
      cublasDestroy(handle);
  }

  for (size_t i=0 ; i<(size_t) nO*nO*nV*nV ; ++i)
  {
     r2[i] += cc_space_v_oovv[i];
  }
}






void compute_g_vir_chol_gpu(const int nO, const int nV, const int cholesky_mo_num,
                double* t1,
                double* t2,
                double* tau,
                double* cc_space_v_ov_chol,
                double* cc_space_v_vo_chol,
                double* cc_space_v_vv_chol,
                double* cc_space_v_oooo,
                double* cc_space_v_vooo,
                double* cc_space_v_oovv,
                double* cc_space_v_vvoo,
                double* cc_space_f_vo,
                double* H_vv,
                double* r2)
{
      int m,n,k, lda, ldb, ldc;
      double alpha, beta;
      double* A;
      double* B;
      double* C;

      int ithread = omp_get_thread_num();
      int igpu = ithread ;

      cudaSetDevice(igpu);
      cublasHandle_t handle;

      cublasCreate(&handle);

      double* d_tau;
      double* d_r2;
      double* d_cc_space_v_vv_chol;
      double* d_cc_space_v_vo_chol;
      double* d_cc_space_v_ov_chol;
      double* d_t1;
      double* d_t2;
      double* d_tmp_cc;



      lda = nO * nO;
      cudaMalloc((void **)&d_r2, lda * nV * nV * sizeof(double));
      cublasSetMatrix(nO*nO, nV*nV, sizeof(double), r2, lda, d_r2, lda);

      lda = nO*nO;
      cudaMalloc((void **)&d_t2, nO*nO*nV*nV * sizeof(double));
      cublasSetMatrix(nO*nO, nV*nV, sizeof(double), t2, lda, d_t2, lda);
      lda = nO;
      cudaMalloc((void **)&d_t1, nO * nV * sizeof(double));
      cublasSetMatrix(nO, nV, sizeof(double), t1, lda, d_t1, lda);
      lda = cholesky_mo_num * nV;
      cudaMalloc((void **)&d_cc_space_v_vv_chol, lda * nV * sizeof(double));
      cublasSetMatrix(cholesky_mo_num*nV, nV, sizeof(double), cc_space_v_vv_chol, lda, d_cc_space_v_vv_chol, lda);

      lda = cholesky_mo_num * nO;
      cudaMalloc((void **)&d_cc_space_v_ov_chol, lda * nV * sizeof(double));
      cublasSetMatrix(cholesky_mo_num*nO, nV, sizeof(double), cc_space_v_ov_chol, lda, d_cc_space_v_ov_chol, lda);

 // ---
           double* d_g_vir;
           cudaMalloc((void**)&d_g_vir, nV*nV*sizeof(double));
           cublasSetMatrix(nV, nV, sizeof(double), H_vv, nV, d_g_vir, nV);

           double* d_cc_space_f_vo;
           cudaMalloc((void**)&d_cc_space_f_vo, nV*nO*sizeof(double));
           cublasSetMatrix(nV, nO, sizeof(double), cc_space_f_vo, nV, d_cc_space_f_vo, nV);

           alpha = -1.0;
           beta =  1.0;
           m=nV ; n=nV; k=nO;
           A = d_cc_space_f_vo  ; lda = nV;
           B = d_t1 ; ldb = nO;
           C = d_g_vir; ldc = nV;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_cc_space_f_vo);

           double* d_tmp_k;
           cudaMalloc((void**)&d_tmp_k, cholesky_mo_num*sizeof(double));
           alpha = 1.0;
           beta =  0.0;
           m=cholesky_mo_num ; n=1; k=nO*nV;
           A = d_cc_space_v_ov_chol; lda = cholesky_mo_num;
           B = d_t1 ; ldb = nO*nV;
           C = d_tmp_k; ldc = cholesky_mo_num;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

           alpha = 2.0;
           beta =  1.0;
           m=nV*nV; n=1; k=cholesky_mo_num;
           A = d_cc_space_v_vv_chol; lda = cholesky_mo_num;
           B = d_tmp_k ; ldb = cholesky_mo_num;
           C = d_g_vir; ldc = nV*nV;
           cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_tmp_k);


           double* d_tmp_vo;
           cudaMalloc((void**)&d_tmp_vo, cholesky_mo_num*nV*nO*sizeof(double));
           alpha = 1.0;
           beta =  0.0;
           m=cholesky_mo_num*nV ; n=nO; k=nV;
           A = d_cc_space_v_vv_chol; lda = cholesky_mo_num*nV;
           B = d_t1 ; ldb = nO;
           C = d_tmp_vo; ldc = cholesky_mo_num*nV;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);

           double* d_tmp_vo2;
           cudaMalloc((void**)&d_tmp_vo2, cholesky_mo_num*nV*nO*sizeof(double));
           for (int i=0 ; i<nO ; ++i) {
             alpha = -1.0;
             beta =  0.0;
             A = &(d_tmp_vo[cholesky_mo_num*nV*i]); lda = cholesky_mo_num;
             B = &(d_tmp_vo[cholesky_mo_num*nV*i]); ldb = cholesky_mo_num;
             C = &(d_tmp_vo2[cholesky_mo_num*i]); ldc = cholesky_mo_num*nO;
             cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, cholesky_mo_num, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
           }
           cudaFree(d_tmp_vo);

           alpha = 1.0;
           beta =  1.0;
           m=nV ; n=nV; k=nO*cholesky_mo_num;
           A = d_cc_space_v_ov_chol; lda = cholesky_mo_num*nO;
           B = d_tmp_vo2 ; ldb = cholesky_mo_num*nO;
           C = d_g_vir; ldc = nV;
           cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_tmp_vo2);

           double* d_Y_oovv;
           cudaMalloc((void**)&d_Y_oovv, nO*nO*nV*nV*sizeof(double));
           alpha = 1.0;
           beta =  0.0;
           m=nO*nO*nV ; n=nV; k=nV;
           A = d_t2; lda = nO*nO*nV;
           B = d_g_vir; ldb = nV;
           C = d_Y_oovv; ldc = nO*nO*nV;
           cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
           cudaFree(d_g_vir);


           alpha = 1.0;
           beta =  1.0;
           A = d_r2; lda = nO*nO;
           B = d_Y_oovv; ldb = nO*nO;
           C = d_r2; ldc = nO*nO;
           cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO*nO, nV*nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
           for (int j=0 ; j<nV ; ++j) {
             for (int i=0 ; i<nV ; ++i) {
                alpha = 1.0;
                beta =  1.0;
                A = &(d_r2[nO*nO*(i+nV*j)]); lda = nO;
                B = &(d_Y_oovv[nO*nO*(j+nV*i)]); ldb = nO;
                C = &(d_r2[nO*nO*(i+nV*j)]); ldc = nO;
                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);

             }
           }
           cudaFree(d_Y_oovv);

           cublasGetMatrix(nO*nO, nV*nV, sizeof(double), d_r2, nO*nO, r2, nO*nO);
}