#include <stdio.h>
#include <stdlib.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 compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo_num,
                double* t1,
                double* tau,
                double* cc_space_v_vo_chol,
                double* cc_space_v_vv_chol,
                double* r2)
{
    int m,n,k, lda, ldb, ldc;
    double alpha, beta;
    double* A;
    double* B;
    double* C;

    double* tmp_cc = malloc(cholesky_mo_num*nV*nV*sizeof(double));

    cublasHandle_t handle;
    cublasCreate(&handle);

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

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

    double* d_cc_space_v_vv_chol;
    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);

    double* d_cc_space_v_vo_chol;
    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);

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

    double* d_tmp_cc;
    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);
    cublasGetMatrix(m, n, sizeof(double), d_tmp_cc, lda, tmp_cc, lda);

// ---

/*
    m=cholesky_mo_num*nV;  n=nV;  k=nO;
    alpha=1.0; beta=0.0;
    lda=m ; ldb=k ; ldc=m;
    A=cc_space_v_vo_chol; B=t1; C=tmp_cc;
    dgemm_("N","N", &m, &n, &k, &alpha, A, &lda, B, &ldb, &beta, C, &ldc);
*/


    #pragma omp parallel
    {
        double* tmp_cc2 = malloc(cholesky_mo_num*nV*sizeof(double));
        double* d_tmp_cc2;
        cudaMalloc((void **)&d_tmp_cc2, cholesky_mo_num*nV*sizeof(double));

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

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

        #pragma omp for
        for (size_t 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]);
//           double* cc_space_v_vv_chol_ = &(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 (size_t l=0 ; l<nV*cholesky_mo_num ; ++l) {
             tmp_cc2[l] = cc_space_v_vv_chol_[l] - tmp_cc_[l];
           }
           cublasSetMatrix(cholesky_mo_num, nV, sizeof(double), tmp_cc2, cholesky_mo_num, d_tmp_cc2, cholesky_mo_num);
*/

           for (size_t iblock=0 ; iblock<nV ; iblock += BLOCK_SIZE)
           {
                const size_t 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);

/*
                A=&(tmp_cc[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
                B=cc_space_v_vv_chol_; ldb=cholesky_mo_num;
                C=tmpB1 ; ldc=nV*BLOCK_SIZE;
                dgemm_("T","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);
                cublasGetMatrix(m, n, sizeof(double), d_tmpB1, ldc, tmpB1, ldc);
/*
                A=&(cc_space_v_vv_chol[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
                B=tmp_cc2; ldb=cholesky_mo_num;
                C=tmpB1 ; ldc=nV*BLOCK_SIZE;
                dgemm_("T","N", &m, &n, &k, &alpha, A, &lda, B, &ldb, &beta, C, &ldc);
*/



                for (size_t beta=iblock ; beta<(nV < iblock+BLOCK_SIZE ? nV : iblock+BLOCK_SIZE) ; ++beta) {
                  for (size_t b=0 ; b<nV ; ++b) {
                    for (size_t a=0 ; a<nV ; ++a) {
                       B1[a + nV*(b + nV*(beta-iblock))] =
                           tmpB1[a + nV*(beta-iblock + BLOCK_SIZE*b)];
                    }
                  }
                }
                cublasSetMatrix(nV*nV, mbs, sizeof(double), B1, nV*nV, d_B1, nV*nV);

                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);
/*
                A=tau; lda=nO*nO;
                B=B1 ; ldb=nV*nV;
                C=&(r2[nO*nO*(iblock + nV*gam)]); ldc=nO*nO;
                dgemm_("N","N", &m, &n, &k, &alpha, A, &lda, B, &ldb, &beta, C, &ldc);
*/

           }
        }

        free(tmp_cc2);
        free(B1);
        free(tmpB1);
    }
    lda=nO*nO;
    cublasGetMatrix(nO*nO, nV*nV, sizeof(double), d_r2, lda, r2, lda);
    free(tmp_cc);

    cudaFree(d_cc_space_v_vo_chol);
    cudaFree(d_cc_space_v_vv_chol);
    cudaFree(d_tau);
    cudaFree(d_t1);
    cudaFree(d_tmp_cc);
    cudaFree(d_r2);
    cublasDestroy(handle);

}


/*

    cublasHandle_t handle;
    cublasCreate(&handle);

    double *d_A, *d_B, *d_C;
    cudaMalloc((void **)&d_A, m * k * sizeof(double));
    cudaMalloc((void **)&d_B, k * n * sizeof(double));
    cudaMalloc((void **)&d_C, m * n * sizeof(double));

    cublasSetMatrix(m, k, sizeof(double), A, m, d_A, m);
    cublasSetMatrix(k, n, sizeof(double), B, k, d_B, k);
    cublasSetMatrix(m, n, sizeof(double), C, m, d_C, m);

    cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, d_A, m, d_B, k, &beta, d_C, m);

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

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

*/