#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <complex>

#include "petsc.h"
#include "petscmat.h"

#include "fileinterface.h"

int sys_read(char *fn, Mat *msm, Vec *mlv,
						 unsigned int *n_eq, unsigned int *n_nd) {

	unsigned int nnz_msm, nnz_mlv, nnz_lcl, *rs, n_rs, k, kr, i_mlv;
	int i_strt, i_end;
  PetscErrorCode ec;
	FILE *fd;
	PetscInt nrow_lcl, *i, *j, *d_nnz, *o_nnz, rnk;
	PetscScalar *v_msm, v_mlv;
	double vr, vi;
	char hn[256];
  #if defined (PETSC_USE_LOG)
    PetscLogEvent  event1;
	PetscLogEvent  event2;
	PetscLogEvent  event3;
  #endif

	MPI_Comm_rank(MPI_COMM_WORLD, &rnk); /* determine node rank */
	ec = PetscGetHostName(hn, 256); CHKERRQ(ec);

	/* open sys-file */
	fd=fopen(fn, "r");
  if (fd==NULL) {
    PetscPrintf(PETSC_COMM_SELF,
								"%s (%d): cannot read sys-file: (%s): %s\n",
                hn, rnk, fn, strerror(errno));
		ec = PetscFinalize(); CHKERRQ(ec);
    exit(EXIT_FAILURE);
  }

	/* read header */
  fread(n_nd,     sizeof(int), 1, fd); /* number of finite element nodes */
  fread(n_eq,     sizeof(int), 1, fd); /* total number of equations */
  fread(&nnz_msm, sizeof(int), 1, fd); /* non-0 elements in system matrix */
  fread(&nnz_mlv, sizeof(int), 1, fd); /* non-0 elements in RHS vector */
	n_rs = *n_eq+1; /* length of "row-start" vector */

  /* create MPI AIJ (=CRS) matrix; dim. of local matrix decided by PETSc */
  printf("Now create MPI AIJ (=CRS) matrix; dim. of local matrix decided by PETSc... \n");
  ec = MatCreate(PETSC_COMM_WORLD,msm); CHKERRQ(ec);
  ec = MatSetSizes(*msm,PETSC_DECIDE,PETSC_DECIDE,*n_eq,*n_eq); CHKERRQ(ec);
  ec = MatSetType(*msm,MATAIJCUSP); CHKERRQ(ec);/*CUDA*/
  ec = PetscObjectSetName((PetscObject)(*msm), "msm"); CHKERRQ(ec);
	ec =  MatSetUp(*msm);
	ec = MatGetOwnershipRange(*msm, &i_strt, &i_end); CHKERRQ(ec);
	nrow_lcl = (PetscInt)(i_end-i_strt); /* number of local rows */

	PetscPrintf(PETSC_COMM_SELF, 
							"sys_read (%d: %s): %d nodes, %d equations, "
							"%d nnz (msm), %d nnz (mlv), %d local rows (%d-%d)\n",
							rnk, hn, *n_nd, *n_eq, nnz_msm, nnz_mlv, nrow_lcl, i_strt, i_end-1);


	/* read data for local submatrix from file: */
	printf("Now read data for local submatrix from file... \n");
	ec = PetscMalloc(n_rs*sizeof(int),&rs); CHKERRQ(ec); /* row-start vec. */
	fread(rs, sizeof(int), n_rs, fd);
	nnz_lcl = rs[i_end]-rs[i_strt]; /* number of non-0s in local submatrix */

 	ec = PetscMalloc(nnz_lcl*sizeof(PetscInt),&i); CHKERRQ(ec);
 	ec = PetscMalloc(nnz_lcl*sizeof(PetscInt),&j); CHKERRQ(ec);
 	ec = PetscMalloc(nnz_lcl*sizeof(PetscScalar),&v_msm); CHKERRQ(ec);
 	ec = PetscMalloc(nrow_lcl*sizeof(PetscInt),&o_nnz); CHKERRQ(ec);
 	ec = PetscMalloc(nrow_lcl*sizeof(PetscInt),&d_nnz); CHKERRQ(ec);

	/* initialize to 0 */
	/* PetscMemzero */
	printf("Now initialize to 0... \n");
 	for (k=0; k<(unsigned)nrow_lcl; k++)
 		d_nnz[k] = o_nnz[k] = 0;

	i[0] = kr = i_strt; /* row counter */

	/* set file position to start of local data */
	fseek(fd, rs[i_strt]*(sizeof(int)+2*sizeof(double)), SEEK_CUR);

	/* read values & determine matrix preallocation size */
	printf("Now read values & determine matrix preallocation size... \n");
	for (k=0; k<nnz_lcl; k++) {

		fread(&(j[k]), sizeof(int), 1, fd);
		fread(&vr,     sizeof(double), 1, fd);
		fread(&vi,     sizeof(double), 1, fd);

		//v_msm[k] = std::complex<double>(vr,vi);
		v_msm[k] = (PetscScalar)(vr + vi);//Hongwang
		
		if (k>=(rs[kr+1]-rs[i_strt])) /* new row */
			kr++;

		i[k] = kr;

		/* count non-0s in diag. & off-diag. submatrix */
		if ((j[k]>=i_strt)&&(j[k]<i_end)) /* diagonal submatrix */
			d_nnz[kr-i_strt]++;
		else /* off-diagonal submatix */
			o_nnz[kr-i_strt]++;

	} /* k */
	
	PetscLogEventRegister("Preallocation",0,&event1);
	PetscLogEventBegin(event1,0,0,0,0);
	
	ec = MatMPIAIJSetPreallocation(*msm, PETSC_NULL, d_nnz, PETSC_NULL, o_nnz);
	CHKERRQ(ec);
	
	PetscLogEventEnd(event1,0,0,0,0);
	
	/* insert values into master stiffness matrix (msm) */
	
	PetscLogEventRegister("Insert values",0,&event2);
	PetscLogEventBegin(event2,0,0,0,0);
	
	printf("Now insert values into master stiffness matrix (msm)... \n");
	for (k=0; k<nnz_lcl; k++) {
		if ((i[k]>=(signed)(*n_eq))||(i[k]>=(signed)(*n_eq)))
			PetscPrintf(PETSC_COMM_SELF, 
									"%s (%d): i=%d or j=%d out of range (%d equations)\n",
									hn, rnk, i[k], j[k], *n_eq);
		ec = MatSetValue(*msm, i[k], j[k], v_msm[k], INSERT_VALUES);
 		CHKERRQ(ec);
	}
	PetscLogEventEnd(event2,0,0,0,0);
	
	PetscLogEventRegister("Mat assembly",0,&event3);
	PetscLogEventBegin(event3,0,0,0,0);
	
    printf("Now mat assembly begin... \n");
    ec = MatAssemblyBegin(*msm,MAT_FINAL_ASSEMBLY); CHKERRQ(ec);
    ec = MatAssemblyEnd(*msm,MAT_FINAL_ASSEMBLY); CHKERRQ(ec);
	
	PetscLogEventEnd(event3,0,0,0,0);

	ec = PetscFree(i); CHKERRQ(ec);
	ec = PetscFree(j); CHKERRQ(ec);
	ec = PetscFree(v_msm); CHKERRQ(ec);
	ec = PetscFree(o_nnz); CHKERRQ(ec);
	ec = PetscFree(d_nnz); CHKERRQ(ec);

	/* master load vector */
	printf("Now same work for the mlv... \n");
	ec = VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE,*n_eq,mlv); CHKERRQ(ec);
    ec = VecSetType(*mlv,VECCUSP); CHKERRQ(ec);/*CUDA*/
    ec = PetscObjectSetName((PetscObject)(*mlv), "mlv"); CHKERRQ(ec);
	ec = VecGetOwnershipRange(*mlv,&i_strt,&i_end); CHKERRQ(ec);

	/* set file position to start of master load vector data */
	fseek(fd, (SYS_HDR_NEL+n_rs)*sizeof(int)+nnz_msm*(sizeof(int)+2*sizeof(double)),
				SEEK_SET);

	/* read all vector data, but assemble only local data */
	for (k=0; k<nnz_mlv; k++) {
		fread(&i_mlv, sizeof(int), 1, fd);
		fread(&vr, sizeof(double), 1, fd);
		if ((i_mlv>=(unsigned int)i_strt)&&(i_mlv<(unsigned int)i_end)) {
			//v_mlv = std::complex<double>(vr,0);
			v_mlv = (PetscScalar)(vr + 0);//Hongwang
			ec = VecSetValue(*mlv,i_mlv,v_mlv,INSERT_VALUES);
			CHKERRQ(ec);
		}
	}

	ec = VecAssemblyBegin(*mlv); CHKERRQ(ec);
	ec = VecAssemblyEnd(*mlv); CHKERRQ(ec);

	fclose(fd);

  return(EXIT_SUCCESS);
} /* sys_read */


/* write first n_nd values of solution to sol-file */
int	sol_write(char *fn, Vec sol, unsigned int n_el, unsigned int n_nd) {
	unsigned int k;
	PetscScalar *v_sol;
  PetscErrorCode ec;
	float vr, vi;
	FILE *fd;
  PetscInt rnk;
	VecScatter sc;
	Vec ss; /* sequential solution */
	char hn[256];

	MPI_Comm_rank(PETSC_COMM_WORLD, &rnk); /* determine node rank */
	ec = PetscGetHostName(hn, 256); CHKERRQ(ec);

	if (rnk==0)
		PetscPrintf(PETSC_COMM_SELF, "%s (%d): writing solution to %s\n",
								hn, rnk, fn);


	ec = VecScatterCreateToZero(sol,&sc,&ss); CHKERRQ(ec);

	/* do the actual communication */
 	ec = VecScatterBegin(sc,sol,ss,INSERT_VALUES,SCATTER_FORWARD); CHKERRQ(ec);
	ec = VecScatterEnd(sc,sol,ss,INSERT_VALUES,SCATTER_FORWARD); CHKERRQ(ec);
	ec = VecScatterDestroy(&sc); CHKERRQ(ec);//Hongwang

	if (rnk==0) { /* writing is only done on processor 0 */
		ec = VecGetArray(ss,&v_sol); CHKERRQ(ec);

		/* open sys-file */
		if ((fd=fopen(fn, "w"))==NULL) {
			PetscPrintf(PETSC_COMM_SELF,
									"%s (%d): cannot open sol-file for write: (%s): %s\n",
									hn, rnk, fn, strerror(errno));
			ec = PetscFinalize(); CHKERRQ(ec);
			exit(EXIT_FAILURE);
		}

		/* write sys-file header */
		if ((fwrite(&n_nd, sizeof(int), 1, fd))!=1) {
			PetscPrintf(PETSC_COMM_SELF,
									"%s (%d): failed to write sol-file header\n",
									hn, rnk);
			ec = PetscFinalize(); CHKERRQ(ec);
			fclose(fd);
			exit(EXIT_FAILURE);
		}


		/* write 1st n_nd solution values to sys-file */
		for (k=0; k<n_nd; k++) {
			vr = (float)(PetscRealPart(v_sol[k]));
			vi = (float)(PetscImaginaryPart(v_sol[k]));
			fwrite(&vr, sizeof(float), 1, fd);
			fwrite(&vi, sizeof(float), 1, fd);
		}

		fclose(fd);

		ec = VecRestoreArray(ss,&v_sol); CHKERRQ(ec);
	} /* rnk==0, i.e., only on processor 0 */
	ec = VecDestroy(&ss); CHKERRQ(ec);//Hongwang


  return(EXIT_SUCCESS);
} /* sol_write */