/**
 * TODO:
 * - implement cudaMemset and cudaMemsetAsync
 * - flesh out basic structure of cuda kernel:
 *   - spawn gridDim blocks
 *   - spawn blockDim threads
 *     Alternatively, spawn blockDim warps and then spawn warp's threads
 *   - wait for blocks to finish
 * - Add in block-level barriers and a __syncthreads() nested function that uses the barrier.
 * - Add in data race checking support and implement atomicAdd for integers
 * - Handle dependencies at atomic blocks (replace some with local blocks if possible)
 * - Create device scope for cudaMallocs
 */

#include <concurrency.cvh>
#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <string.h>

enum cudaError {
  cudaSuccess
};
typedef enum cudaError cudaError_t;

typedef enum cudaMemcpyKind {
	cudaMemcpyHostToHost,
	cudaMemcpyHostToDevice,
	cudaMemcpyDeviceToHost,
	cudaMemcpyDeviceToDevice,
	cudaMemcpyDefault
} cudaMemcpyKind;

typedef struct {
  unsigned int x, y, z;
} dim3;

/* used to represent a location in a three dimensional grid
 */
typedef struct {
  unsigned int x, y, z;
} uint3;

typedef struct $cuda_op_state* $cuda_op_state_t;
struct $cuda_op_state {
	_Bool start;
	$proc op;
};

typedef struct $cuda_op_state_node* $cuda_op_state_node_t;
struct $cuda_op_state_node {
	$cuda_op_state_t opState;
	$cuda_op_state_node_t next;
};

typedef struct cudaStream* cudaStream_t;
typedef struct $cuda_stream_node* $cuda_stream_node_t;
struct cudaStream {
	$cuda_op_state_node_t head;
	$cuda_op_state_node_t tail;
	int numOps;
	$cuda_stream_node_t containingNode;
	_Bool alive;
};
cudaStream_t $cuda_default_stream;

struct $cuda_stream_node {
	cudaStream_t stream;
	$cuda_stream_node_t prev;
	$cuda_stream_node_t next;
};

typedef struct $cuda_context {
	$cuda_stream_node_t head; //list of streams
	int numStreams;
} $cuda_context;
$cuda_context $cuda_global_context;

// Helper function to get the default stream if passed NULL, and just returns stream otherwise
cudaStream_t $default_stream_if_null(cudaStream_t stream) {
	return stream == NULL ? $cuda_default_stream : stream;
}

$cuda_stream_node_t $create_new_stream_node() {
	cudaStream_t newStream = (cudaStream_t) malloc(sizeof(struct cudaStream));
	newStream->head = NULL;
	newStream->tail = NULL;
	newStream->numOps = 0;
	newStream->alive = true;
	
	$cuda_stream_node_t newHead = ($cuda_stream_node_t) malloc(sizeof(struct $cuda_stream_node));
	newHead->stream = newStream;
	newStream->containingNode = newHead;
	newHead->prev = NULL;
	newHead->next = NULL;
	
	return newHead;
}

/**
 * TODO:
 *   - test
 *   - atomic?
 */
cudaError_t cudaStreamCreate(cudaStream_t * pStream) {
	// Create new stream node in linked list
	$cuda_stream_node_t newHead = $create_new_stream_node();
	newHead->next = $cuda_global_context.head;
	$cuda_global_context.head->prev = newHead;
	
	// Update cuda context's head to be the new node we created
	$cuda_global_context.head = newHead;
	$cuda_global_context.numStreams++;
	
	return cudaSuccess;
}

/**
 * TODO:
 *   - test
 *   - atomic?
 */
cudaError_t cudaStreamSynchronize(cudaStream_t stream) {
	stream = $default_stream_if_null(stream);
	$assert(stream->alive, "Attempt to synchronize with a destroyed stream");
	$when(stream->head == NULL) return cudaSuccess;
}

$proc $destroy_stream_node($cuda_stream_node_t node) {
    $proc lastOpProc = $proc_null;
    cudaStream_t stream = node->stream;
  
    $local_start();
    if (node->prev != NULL) {
      node->prev->next = node->next;
    }
    if (node->next != NULL) {
      node->next->prev = node->prev;
    }
    free(node);
    
    stream->alive = false;
    if(stream->tail != NULL)
      lastOpProc = stream->tail->opState->op;
    $local_end();
	
	void $destroy_stream_when_complete($proc lastOpProc, cudaStream_t stream) {
		$wait(lastOpProc);
		free(stream);
	}
  
	return $spawn $destroy_stream_when_complete(lastOpProc, stream);
}

// TODO: atomic
cudaError_t cudaStreamDestroy(cudaStream_t stream) {
	$assert(stream != NULL && stream != $cuda_default_stream, "Attempt to destroy default stream");
	$assert(stream->alive, "Attempt to destroy an already destroyed stream");
	$destroy_stream_node(stream->containingNode);
	return cudaSuccess;
}

/**
 * Enqueues the calling $proc as a new cuda operation onto stream. Then blocks until the cuda operation is allowed to execute.
 */
$cuda_op_state_t $stream_enqueue(_Bool* enqueuedFlag, cudaStream_t stream) {
  $cuda_op_state_t newOpState = ($cuda_op_state_t) malloc(sizeof(struct $cuda_op_state));
  newOpState->start = false;
  newOpState->op = $self;

  $cuda_op_state_node_t newOpStateNode = ($cuda_op_state_node_t) malloc(sizeof(struct $cuda_op_state_node));
  newOpStateNode->opState = newOpState;
  newOpStateNode->next = NULL;
    
  $local_start();
  stream = $default_stream_if_null(stream);
  $assert(stream->alive, "Attempt to enqueue a CUDA operation onto a destroyed stream");
	
  if (stream->tail == NULL) {
    stream->head = newOpStateNode;
    stream->tail = newOpStateNode;
    newOpState->start = true;
  } else {
    stream->tail->next = newOpStateNode;
    stream->tail = newOpStateNode;
  }
  stream->numOps++;
  *enqueuedFlag = true;
  $local_end();

  return newOpState;
}

void $stream_dequeue(cudaStream_t stream) {
	stream = $default_stream_if_null(stream);
	$assert(stream->head != NULL, "Attempt to dequeue an empty stream");

  	$local_start();
    if (stream->head == stream->tail) {
      stream->tail = NULL;
    }
	
    $cuda_op_state_node_t oldHead = stream->head;
    stream->head = oldHead->next;
    if (stream->head != NULL) {
      stream->head->opState->start = true;
    }
	
    stream->numOps--;
    free(oldHead->opState);
    free(oldHead);
    $local_end();
}

void $cuda_memcpy_proc(void* dst, const void* src, size_t count, _Bool* enqueuedFlag, cudaStream_t stream) {
  $cuda_op_state_t opState = $stream_enqueue(enqueuedFlag, stream);
  $when(opState->start);
	memcpy(dst, src, count);
	$stream_dequeue(stream);
}

cudaError_t cudaMemcpy(void* dst, const void* src, size_t count, cudaMemcpyKind kind) {
	if (kind == cudaMemcpyHostToHost) {
		memcpy(dst, src, count);
	} else {
    _Bool enqueuedFlag = false;
		$proc memcpyProc = $spawn $cuda_memcpy_proc(dst, src, count, &enqueuedFlag, $cuda_default_stream);
    $when(enqueuedFlag);
		if (kind != cudaMemcpyDeviceToDevice) {
			$wait(memcpyProc);
		}
	}
	
	return cudaSuccess;
}

cudaError_t cudaMemcpyAsync(void* dst, const void* src, size_t count,
                            cudaMemcpyKind kind, cudaStream_t stream) {
	if (kind == cudaMemcpyHostToHost) {
		memcpy(dst, src, count);
	} else {
    _Bool enqueuedFlag = false;
		$spawn $cuda_memcpy_proc(dst, src, count, &enqueuedFlag, stream);
    $when(enqueuedFlag);
	}
	
	return cudaSuccess;
}

cudaError_t cudaDeviceSynchronize() {
	$proc* opsToWaitOn;
	int numOps = 0;
	
	$atomic {
		opsToWaitOn = ($proc*) malloc(sizeof($proc) * $cuda_global_context.numStreams);
		
		for ($cuda_stream_node_t node = $cuda_global_context.head;
         node != NULL;
         node = node->next) {
      if (node->stream->tail != NULL) {
        opsToWaitOn[numOps] = node->stream->tail->opState->op;
        numOps++;
      }
    }
  }
  $waitall(opsToWaitOn, numOps);
	
  return cudaSuccess;
}

/**
 * Only called at start of program
 */
void $cuda_setup() {
  $cuda_stream_node_t defaultStreamNode = $create_new_stream_node();
  $cuda_default_stream = defaultStreamNode->stream;

  $cuda_global_context.head = defaultStreamNode;
  $cuda_global_context.numStreams = 1;
}

/**
 * Only called at end of program
 */
void $cuda_teardown() {
  $proc destructor = $destroy_stream_node($cuda_default_stream->containingNode);
  $wait(destructor);
}

// Helper function
int $dim3_index(dim3 size, uint3 location) {
  return location.x + size.x * (location.y + size.y * location.z);
}

// Helper function
int $cuda_kernel_index (dim3 gDim, dim3 bDim, uint3 bIdx, uint3 tIdx) {
  return $dim3_index(gDim, bIdx) * (bDim.x * bDim.y * bDim.z) + $dim3_index(bDim, tIdx);
}

void $cuda_run_and_wait_on_procs(dim3 dim, void spawningFunction(uint3)) {
  //TODO: calculate length and index, replace this function in the kernel
  //$local_start();
  int length = dim.x * dim.y * dim.z;
  $proc proc_array[length];
  $range rx = 0 .. dim.x - 1;
  $range ry = 0 .. dim.y - 1;
  $range rz = 0 .. dim.z - 1;
  $domain(3) dom = ($domain){rx, ry, rz};
  #For some reason there is depends on all here
  $for(int x,y,z : dom){
    uint3 id = { x, y, z };
    int index = $dim3_index(dim, id);
    proc_array[index] = $spawn spawningFunction(id);
  }
  //$local_end();
  $waitall(proc_array,length);
}

void _cuda_kernel_1(dim3 gridDim, dim3 blockDim, size_t _cuda_mem_size,
                    const float *A, const float *B, float *C, int numElements) {
  void _cuda_block(uint3 blockIdx) {
    int numThreads = (blockDim.x * blockDim.y) * blockDim.z;
    $scope _block_root = $here;
    $gbarrier _cuda_block_barrier = $gbarrier_create($here, blockDim.x * blockDim.y * blockDim.z);
    void _cuda_thread(uint3 threadIdx) {
      int _cuda_tid = $dim3_index(blockDim, threadIdx);
      int _cuda_kid = $cuda_kernel_index(gridDim, blockDim, blockIdx, threadIdx);
      $barrier _cuda_thread_barrier = $barrier_create($here, _cuda_block_barrier, _cuda_tid);
      $local_start();
      // Kernel definition start
      
      int i = blockDim.x * blockIdx.x + threadIdx.x;

      if (i < numElements)
      {
        C[i] = A[i] + B[i];
      }
      
      // Kernel definition end
      $local_end();
      $barrier_destroy(_cuda_thread_barrier);
    }
    $cuda_run_and_wait_on_procs(blockDim, _cuda_thread);
    $gbarrier_destroy(_cuda_block_barrier);
  }
  $cuda_run_and_wait_on_procs(gridDim, _cuda_block);
}

void $proc_kernel_1(dim3 gridDim, dim3 blockDim, size_t _cudaMemSize,
                    _Bool* enqueuedFlag, cudaStream_t _cudaStream,
                    const float *A, const float *B, float *C, int numElements) {
  $cuda_op_state_t opState = $stream_enqueue(enqueuedFlag, _cudaStream);
  $when(opState->start);
  _cuda_kernel_1(gridDim, blockDim, _cudaMemSize, A, B, C, numElements);
  $stream_dequeue(_cudaStream);
}

$input int N;
$assume (N > 0);
$input float A[N];
$input float B[N];

void _civl_main() {
  int size = N * sizeof(float);
  int numBlocks = 2;
  int numThreads = N%2 == 0? N/2 : (N+1)/2;

  float* cuda_A;
  // cudaMalloc((void **)&cuda_A, size);
  {
    cuda_A = (float *) malloc(size);
  }
  cudaMemcpy(cuda_A, A, size, cudaMemcpyHostToDevice);

  float* cuda_B;
  // cudaMalloc((void **)&cuda_B, size);
  {
    cuda_B = (float *) malloc(size);
  }
  cudaMemcpy(cuda_B, B, size, cudaMemcpyHostToDevice);

  float* cuda_C;
  // cudaMalloc((void **)&cuda_C, size);
  {
    cuda_C = (float *) malloc(size);
  }

  { // kernel_1<<<gridDim, blockDim>>>(cuda_A, cuda_B, cuda_C, N);
    dim3 gridDim = {numBlocks, 1, 1};
    dim3 blockDim = {numThreads, 1, 1};
    _Bool enqueuedFlag = false;
    $spawn $proc_kernel_1(gridDim, blockDim, 0, &enqueuedFlag, NULL, cuda_A, cuda_B, cuda_C, N);
    $when(enqueuedFlag);
  }
  
  //Checking correctness
  float* C = (float *)malloc(size);
  
  cudaMemcpy(C, cuda_C, size, cudaMemcpyDeviceToHost);
  
  for(int i = 0; i < N; i++)
  	$assert(C[i] == A[i] + B[i]);
  
  free(C);
  
  //cudaFree(cuda_A);...
  free(cuda_A);
  free(cuda_B);
  free(cuda_C);
}

int main() {
  $cuda_setup();
  _civl_main();
  $cuda_teardown();
}