#include <stdio.h>
#include <stdlib.h>

#define MAXT 20
#define CT 0.335

#define END 3
#define N 20 
//#define END 20
//#define N  5000
//float mt_1[N][N]; // matrix with the plate's state at time t-1
//float mt[N][N];   // idem for time t

void debug_chp(float *chp)
{
	int i, j;
	printf("\n");
	for(i=0; i< N; i++) {
	  printf("i:%d ", i);
	  for(j=0; j<N; j++)
	     printf("%.1f ", *(chp+i*N+j));
	   printf("\n");
	}

}


__global__ void initChp(float * pt_1)
{
	int x = threadIdx.x + blockIdx.x * blockDim.x;
	int y = threadIdx.y + blockIdx.y * blockDim.y;
	int offset = x + y * blockDim.x * gridDim.x;

	if ( (x== 0) || (y== 0) || ( x==N-1) || (y==N-1) ) 
		pt_1[offset] = 200;
	else 
		pt_1[offset] = 20;

	__syncthreads();
}

__global__ void stepChp(float *pt_1, float *pt)
{
	int x = threadIdx.x + blockIdx.x * blockDim.x;
	int y = threadIdx.y + blockIdx.y * blockDim.y;
	int offset = x + y * blockDim.x * gridDim.x;

	int border = (x== 0) || (y== 0) || ( x==N-1) || (y==N-1);
	
	if(!border) {
		pt[offset]= (CT/4)*(pt_1[offset-1] + pt_1[offset+1] + 
			pt_1[offset-blockDim.x*gridDim.x] + 
			pt_1[offset+blockDim.x*gridDim.x]);
	}
/*
	else
		pt[offset] = 200;			
*/
	
	__syncthreads();
}


int main (int argc, char *argv[]) 
{
	int numthreads, step;
	FILE *f; float *t;
	float *pt_1;
	float *pt;

	if(argc < 2 ) {
		printf("Usage: %s n_threads\n", argv[0]);
		exit(0);
	}
 
	numthreads = strtol(argv[1], NULL, 10);

	t = (float *)malloc(N*N*sizeof(float));

	cudaMalloc((void **)&pt_1, N*N*sizeof(float));
	cudaMalloc((void **)&pt, N*N*sizeof(float));

	dim3 blocks(N/numthreads,N/numthreads);
	dim3 bthreads(numthreads, numthreads);

	initChp<<<blocks, bthreads>>>(pt_1); 
	//cudaMemcpy(t, pt_1, N*N*sizeof(float), cudaMemcpyDeviceToHost);
	//debug_chp(t);
	
	for(step=0; step < END; step++) {
		stepChp<<<blocks, bthreads>>>(pt_1,pt);
		float *tmp;
		tmp = pt_1;
  		pt_1 = pt;
  		pt = tmp;
  		
  	}

	cudaMemcpy(t, pt_1, N*N*sizeof(float), cudaMemcpyDeviceToHost);
	//debug_chp(t);

	cudaFree(pt_1);
	cudaFree(pt);
  	
	f= fopen("simulacao.bin", "w");
	fwrite(t, sizeof(float), N*N, f);
	fclose(f);
	free(t);


	return 0;
 }