#include <iostream>
#include <cstdlib>
#include <sys/time.h>
#include "DKSImageReconstruction.h"

struct voxelPosition {
  float x;
  float y;
  float z;
};

void initImage(float *image, int size) {
  for (int i = 0; i < size; i++)
    image[i] = (float)rand() / RAND_MAX;
}

void initPosition(voxelPosition *voxel, int N) {
  for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j++) {
      for (int k = 0; k < N; k++) {
	int idx = i * N * N + j * N + k;
	if (k == 0)
	  voxel[idx].x = 0.0;
	else
	  voxel[idx].x = voxel[idx - 1].x + 0.1;

	if (j == 0)
	  voxel[idx].y = 0.0;
	else
	  voxel[idx].y = voxel[idx - N].y + 0.1;

	if (i == 0)
	  voxel[idx].z = 0.0;
	else
	  voxel[idx].z = voxel[idx - N * N].z + 0.1;
      }
    }
  }
}

void printPosition(voxelPosition *voxel, int size) {
  for (int i = 0; i < size; i++)
    std::cout << voxel[i].x << "\t";
  std::cout << std::endl;
  for (int i = 0; i < size; i++)
    std::cout << voxel[i].y << "\t";
  std::cout << std::endl;
  for (int i = 0; i < size; i++)
    std::cout << voxel[i].z << "\t";
  std::cout << std::endl;
}

#define DIAMETER 2.0
bool select_source(voxelPosition *image_tmp, voxelPosition source_temp, int id)
{
  float distance_x = pow(image_tmp[id].x-source_temp.x,2);
  float distance_y = pow(image_tmp[id].y-source_temp.y,2);
  float distance_z = pow(image_tmp[id].z-source_temp.z,2);
  float distance = sqrt(distance_x + distance_y + distance_z);

  if ( distance < DIAMETER*0.5 ) {
      return true;
  }
  else
    return false;
}

void calculate_source(float *image_space , voxelPosition *image_geometry, 
		      voxelPosition source, int total_voxels, 
		      float *average, float *std)
{
  
  int number_selected_maximum = 10000;
  float *select;
  select = new float[number_selected_maximum];
  for (int j=0;j<number_selected_maximum;j++)
    select[j] = 0.0;
  int number_selected=0;

  for (int voxel_id = 0; voxel_id < total_voxels; voxel_id++) {
    if ( select_source( image_geometry, source, voxel_id ) ) {
      select[number_selected] = image_space[voxel_id];
      number_selected += 1;
    }
  }

  *average = 0.0;
  *std = 0.0;

  for (int j=0;j<number_selected;j++)
    *average += select[j];
  *average /= float(number_selected);

  for (int j=0;j<number_selected;j++)
    *std += pow(*average-select[j],2);
  *std = sqrt(*std/number_selected/(number_selected-1));

  delete[] select;
}

int main(int argc, char *argv[]) {

  int N = 8;
  if (argc == 2)
    N = atoi(argv[1]);

  double ttotal;
  struct timeval timeStart, timeEnd;

  int total = N*N*N;
  float *image = new float[total];
  voxelPosition *geometry = new voxelPosition[total];

  initImage(image, total);
  initPosition(geometry, N);

  voxelPosition source;
  float avg[total], stdev[total];

  gettimeofday(&timeStart, NULL);
  for (int i = 0; i < total; i++) {
    source.x = geometry[i].x;
    source.y = geometry[i].y;
    source.z = geometry[i].z;
    calculate_source(image , geometry, source, total, &avg[i], &stdev[i]);
  }    
  gettimeofday(&timeEnd, NULL);
  ttotal = ( (timeEnd.tv_sec - timeStart.tv_sec) * 1000000 + 
	     (timeEnd.tv_usec - timeStart.tv_usec)) * 1e-6;

  float avgavg = 0;
  float avgstdev = 0;
  for (int i = 0; i < total; i++) {
    avgavg += avg[i] / total;
    avgstdev += stdev[i] / total;
  }

  std::cout << "Total voxels: " << N*N*N << std::endl;
  std::cout << "Dimensions [" << geometry[0].x << ":" << geometry[N-1].x << "]"
	    << "[" << geometry[0].y << ":" << geometry[N*N-1].x << "]"
	    << "[" << geometry[0].z << ":" << geometry[N*N*N-1].x << "]" << std::endl;
  std::cout << "Average: " << avgavg << ", stddev: " << avgstdev << ", time : " << ttotal<< std::endl;


  void *image_space, *image_position, *source_position, *davg, *dstd;

  int ierr;
  DKSImageRecon base;
  base.setAPI("Cuda", 4);
  base.setDevice("-gpu", 4);
  base.initDevice();

  image_space = base.allocateMemory<float>(total, ierr);
  image_position = base.allocateMemory<voxelPosition>(total, ierr);
  source_position = base.allocateMemory<voxelPosition>(total, ierr);
  davg = base.allocateMemory<float>(total, ierr);
  dstd = base.allocateMemory<float>(total, ierr);

  base.writeData<float>(image_space, image, total);
  base.writeData<voxelPosition>(image_position, geometry, total);
  base.writeData<voxelPosition>(source_position, geometry, total);


  gettimeofday(&timeStart, NULL);
  base.callCalculateSource(image_space, image_position, source_position, 
			   davg, dstd, DIAMETER, total, total);


  base.readData<float>(davg, avg, total);
  base.readData<float>(dstd, stdev, total);
  gettimeofday(&timeEnd, NULL);
  ttotal = ( (timeEnd.tv_sec - timeStart.tv_sec) * 1000000 + 
	     (timeEnd.tv_usec - timeStart.tv_usec)) * 1e-6;

  base.freeMemory<float>(image_space, total);
  base.freeMemory<voxelPosition>(image_position, total);
  base.freeMemory<voxelPosition>(source_position, total);
  base.freeMemory<float>(dstd, total);
  base.freeMemory<float>(davg, total);
  
  avgavg = 0;
  avgstdev = 0;
  for (int i = 0; i < total; i++) {
    avgavg += avg[i] / total;
    avgstdev += stdev[i] / total;
  }
  std::cout << "Average: " << avgavg << ", stddev: " << avgstdev << ", time : " << ttotal<< std::endl;

  return N;

}