/*
!===============================================================================
!
! cublas_wrap.c - cuda version for some lapack functions
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2011 Yoshifumi Nakamura
!
! This file is part of BQCD -- Berlin Quantum ChromoDynamics program
!
! BQCD is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! BQCD is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with BQCD.  If not, see <http://www.gnu.org/licenses/>.
!
!-------------------------------------------------------------------------------
*/
#ifdef NamesToLower
# define  ZMINV_NV  zminv_nv
# define ZGETF2_NV zgetf2_nv
# define  ZGEMM_NV  zgemm_nv
# define INIT_CUBLAS init_cublas
# define FIN_CUBLAS fin_cublas
#endif
#ifdef NamesToLower_
# define  ZMINV_NV  zminv_nv_
# define ZGETF2_NV zgetf2_nv_
# define  ZGEMM_NV  zgemm_nv_
# define INIT_CUBLAS init_cublas_
# define FIN_CUBLAS fin_cublas_
#endif
#ifdef NamesToLower__
# define  ZMINV_NV  zminv_nv__
# define ZGETF2_NV zgetf2_nv__
# define  ZGEMM_NV  zgemm_nv__
# define INIT_CUBLAS init_cublas__
# define FIN_CUBLAS fin_cublas__
#endif

#include <cublas.h>
#include <stdio.h>
#define min(a, b) ((a) < (b) ? (a) : (b))

int init_cublas_alloc=1;
cuDoubleComplex *devA, *devB, *devC, *zero;


void  INIT_CUBLAS(const int* m, const int* n, int* info){
  int stat, i;
  int size=(*m) * (*n);

  *info=0;

  if (init_cublas_alloc == 0)return;

  stat = cublasInit();
  if (stat != CUBLAS_STATUS_SUCCESS) {
    printf ("cublasInit failed");
    *info =1;
    return;
  }

  stat = cublasAlloc (size, 16, (void**)&devA);
  if (stat != CUBLAS_STATUS_SUCCESS) {
    printf ("device memory allocation failed %d %d\n",stat,size);
    *info =1;
    return;
  }
  stat = cublasAlloc (size, 16, (void**)&devB);
  if (stat != CUBLAS_STATUS_SUCCESS) {
    printf ("device memory allocation failed %d %d\n",stat,size);
    *info =1;
    return;
  }
  stat = cublasAlloc (size, 16, (void**)&devC);
  if (stat != CUBLAS_STATUS_SUCCESS) {
    printf ("device memory allocation failed %d %d\n",stat,size);
    *info =1;
    return;
  }

  zero=( cuDoubleComplex * )malloc(*n *16);
  for (i=0; i< *n; i++){ zero[i]=make_cuDoubleComplex( 0, 0);}

  //  printf ("cublasAlloc done\n");
  init_cublas_alloc=0;
}


void  ZMINV_NV(const int* m, const int* n, const double* A, const int* lda, 
	       int* ipiv, int* info, double* det){
  //  printf ("now zminv_nv_ started\n");
  int stat, i, count;
  cuDoubleComplex mone, one, tmp, ttmp, det_tmp;
  count=0;
   one=make_cuDoubleComplex( 1, 0);
  mone=make_cuDoubleComplex(-1, 0);
  det_tmp=make_cuDoubleComplex( 0, 0);

  *info=0;
  int mm=(*m);
  int nn=(*n);

  cublasSetMatrix(*m, *n, 16, A, *n, devA, *m);

//ZGETF2 start
  int j, j0n, j1n, jp;
  for (j=0; j<min(mm,nn); j++) {
    j0n= j     * nn;
    j1n= (j+1) * nn;
    jp = j - 1 + cublasIzamax( mm-j, &devA[ j+j0n ], 1 ) ;

    ipiv[ j ] = jp + 1 ;
    //    if (j0n+j  >= nn*mm || j0n+j  <0){printf ("zgetf2_nv_ j0n+j  =%d \n",j0n+j  );exit;}
    //    if (j0n+jp >= nn*mm || j0n+jp <0){printf ("zgetf2_nv_ j0n+jp =%d \n",j0n+jp );exit;}
    //    if (j      >= nn*mm || j      <0){printf ("zgetf2_nv_ j      =%d \n",j      );exit;}
    //    if (jp     >= nn*mm || jp     <0){printf ("zgetf2_nv_ jp     =%d \n",jp     );exit;}
    
    cublasGetMatrix(1, 1, 16,  &devA[ j0n+jp ], 1, (void*)&tmp, 1);

    if(!( tmp.x == 0.0d && tmp.y == 0.0d) ) {
      if( jp != j) cublasZswap(nn,  &devA[ j ], *lda, &devA[ jp ], *lda);
      cublasGetMatrix(1, 1, 16,  &devA[j0n+j], 1, (void*)&tmp, 1);  
      ttmp = cuCdiv(one , tmp);
      if( j  < mm-1) cublasZscal(mm-j-1, ttmp, &devA[ j+1+j0n], 1);
    } else if( *info == 0 ) {
      printf ("zgetf2_nv_ tmp=0 at j=%d\n",j);  
      *info = j;
    }
    if ( j < min(mm,nn)-1)
      cublasZgeru( mm-j-1, nn-j-1, mone, &devA[ j+1+j0n ], 1, &devA[ j+j1n ], *lda, &devA[ j+1+j1n ], *lda);
    //ZGERU( M-J, N-J, -ONE, A( J+1, J ), 1, A( J, J+1 ), LDA, A( J+1, J+1 ), LDA )
  }
//ZGETF2 end

//ZTRTI2 for Upper & No-unit start
  for (j=0;j<nn;j++) {
    j0n= j     * nn;
    cublasGetMatrix(1, 1, 16,  &devA[j0n+j], 1, (void*)&tmp, 1);
    //det_tmp = cuCmul(det_tmp, tmp); // for determinant
    det_tmp.x = det_tmp.x + log(sqrt(tmp.x*tmp.x + tmp.y*tmp.y));
    det_tmp.y = det_tmp.y + atan2(tmp.y, tmp.x);
    if (ipiv[j]-1 != j)count++;
    ttmp = cuCdiv(one , tmp);
    cublasSetMatrix(1, 1, 16,  (void*)&ttmp, 1, &devA[j0n+j], 1);
    ttmp.x = -ttmp.x;
    ttmp.y = -ttmp.y;
    cublasZtrmv('U','N', 'N', j, devA, *lda, &devA[j0n],1);
    cublasZscal( j, ttmp, &devA[j0n],1);
  }
  det[0]=det_tmp.x;
  det[1]=det_tmp.y;
  if (count%2 == 1)det[1]+=3.1415926535897932384626433832795028841971;
//ZTRTI2 for Upper & No-unit end

//===============================================================
//ZGETRI unblocked start
  for (j=nn-1;j>=0;j--) {
    j0n= j     * nn;
    j1n= (j+1) * nn;
    if (j!=nn-1)cublasZcopy(nn-j-1, &devA[j0n+j+1], 1, devB, 1);
    if (j!=nn-1)cublasSetVector(nn-j-1,   16, zero, 1, &devA[j0n+j+1], 1);
    if( j < nn-1 )
      cublasZgemv('N', nn, nn-j-1, mone, &devA[j1n], *lda, devB, 1, one, &devA[ j0n],1);   
  }
//ZGETRI unblocked end

// interchange
  for (j=nn-2;j>=0;j--) {
    jp=ipiv[j]-1;
    if (jp != j) cublasZswap(nn, &devA[j*nn],1, &devA[jp*nn],1);
  }

  cublasGetMatrix(*m, *n, 16, devA, *n, (void*)A, *m);
}

void ZGETF2_NV(const int* m, const int* n, const double* A, const int* lda, 
	       int* ipiv, int* info){
  int stat, i;
  cuDoubleComplex mone, one, tmp, ttmp;
   one=make_cuDoubleComplex( 1, 0);
  mone=make_cuDoubleComplex(-1, 0);

  *info=0;

  int mm=(*m);
  int nn=(*n);

  cublasSetMatrix(*m, *n, 16, A, *n, devA, *m);

//ZGETF2 start
  int j, j0n, j1n, jp;
  for (j=0; j<min(mm,nn); j++) {
    j0n= j     * nn;
    j1n= (j+1) * nn;
    jp = j - 1 + cublasIzamax( mm-j, &devA[ j+j0n ], 1 ) ;

    ipiv[ j ] = jp + 1 ;
    //    if (j0n+j  >= nn*mm || j0n+j  <0){printf ("zgetf2_nv_ j0n+j  =%d \n",j0n+j  );exit;}
    //    if (j0n+jp >= nn*mm || j0n+jp <0){printf ("zgetf2_nv_ j0n+jp =%d \n",j0n+jp );exit;}
    //    if (j      >= nn*mm || j      <0){printf ("zgetf2_nv_ j      =%d \n",j      );exit;}
    //    if (jp     >= nn*mm || jp     <0){printf ("zgetf2_nv_ jp     =%d \n",jp     );exit;}
    
    cublasGetMatrix(1, 1, 16,  &devA[ j0n+jp ], 1, (void*)&tmp, 1);

    if(!( tmp.x == 0.0d && tmp.y == 0.0d) ) {
      if( jp != j) cublasZswap(nn,  &devA[ j ], *lda, &devA[ jp ], *lda);
      cublasGetMatrix(1, 1, 16,  &devA[j0n+j], 1, (void*)&tmp, 1);  
      ttmp = cuCdiv(one , tmp);
      if( j  < mm-1) cublasZscal(mm-j-1, ttmp, &devA[ j+1+j0n], 1);
    } else if( *info == 0 ) {
      printf ("zgetf2_nv_ tmp=0 at j=%d\n",j);  
      *info = j;
    }
    if ( j < min(mm,nn)-1)
      cublasZgeru( mm-j-1, nn-j-1, mone, &devA[ j+1+j0n ], 1, &devA[ j+j1n ], *lda, &devA[ j+1+j1n ], *lda);
    //ZGERU( M-J, N-J, -ONE, A( J+1, J ), 1, A( J, J+1 ), LDA, A( J+1, J+1 ), LDA )
  }
//ZGETF2 end

  cublasGetMatrix(*m, *n, 16, devA, *n, (void*)A, *m);
}

//================================================================================
    void ZGEMM_NV(const char *transa, const char *transb, const int *m,
		  const int *n, const int *k, const cuDoubleComplex  *alpha,
		  const cuDoubleComplex  *A, const int *lda,
		  const cuDoubleComplex  *B, const int *ldb, 
		  const cuDoubleComplex  *beta, const cuDoubleComplex *C,
		  const int *ldc)
{
  cublasSetMatrix(*m, *n, 16, A, *n, devA, *m);
  cublasSetMatrix(*m, *n, 16, B, *n, devB, *m);
  if ((*beta).x !=0 || (*beta).y !=0) cublasSetMatrix(*m, *n, 16, C, *n, devC, *m);
  cublasZgemm (transa[0], transb[0], *m, *n, *k, *alpha, devA, *lda, devB, *ldb, *beta, devC, *ldc);
  cublasGetMatrix(*m, *n, 16, devC, *n, (void*)C, *m);
}

void FIN_CUBLAS(){
  free(zero);
  cublasFree(devA);
  cublasFree(devB);
  cublasFree(devC);
  cublasShutdown();
  int init_cublas_alloc=1;
}
