#include <CL/cl.h>
#include <CL/cl_platform.h>

#include "cl_simple.h"
#include "cl_util.h"

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


int main(int argc, char ** argv)
{
   cl_int error;
   struct cl_simple_context context;
   int * out;
   int * new_out;
   size_t global_work_size = 10;
   unsigned prestore = 12345;
   unsigned out_size;
   int iterations, i, j;
   unsigned pass = 1;

   if (argc !=2 ) {
      fprintf(stderr, "Usage: use-host-ptr iterations\n");
      return EXIT_FAILURE;
   }

   iterations = atoi(argv[1]);
   out_size = (global_work_size * iterations + 1) * sizeof(int);
   out = malloc(out_size);
   out[global_work_size * iterations] = prestore;

   if (!clSimpleSimpleInit(&context, "loop_lt")) {
      return EXIT_FAILURE;
   }

   context.out_buffer = clCreateBuffer(context.cl_ctx,
                                       CL_MEM_WRITE_ONLY | CL_MEM_USE_HOST_PTR,
                                       out_size,
                                       out,
                                       &error);
   if (error != CL_SUCCESS) {
      fprintf(stderr, "clCreateBuffer() failed: %s\n",
                       clUtilErrorString(error));
      return EXIT_FAILURE;
   }

   if (   !clSimpleKernelSetArg(context.kernel, 0, sizeof(cl_mem),
                                &context.out_buffer)
       || !clSimpleKernelSetArg(context.kernel, 1, sizeof(int),
                                &iterations)) {
      return EXIT_FAILURE;
   }

   if (!clSimpleEnqueueNDRangeKernel(context.command_queue,
                                     context.kernel,
                                     1, &global_work_size, &global_work_size)) {
      return EXIT_FAILURE;
   }

   new_out = clEnqueueMapBuffer(context.command_queue,
                            context.out_buffer,
                            CL_TRUE, /* Blocking */
                            CL_MAP_READ,
                            0, /* Offset */
                            out_size, /* cb */
                            0, NULL, NULL, &error);
   if (error != CL_SUCCESS) {
      fprintf(stderr, "clEnqueueMapBuffer() failed: %s\n",
                      clUtilErrorString(error));
      return EXIT_FAILURE;
   }


   error = clEnqueueUnmapMemObject(context.command_queue, context.out_buffer,
                                   new_out, 0, NULL, NULL);
   if (error != CL_SUCCESS) {
      fprintf(stderr, "clEnqueueUnmapMemObject() failed: %s\n",
                      clUtilErrorString(error));
   }

   for (i = 0; i < global_work_size; i++) {
      for (j = 0; j < iterations; j++) {
         unsigned id = (i * iterations) + j;
         if (new_out[id] != id) {
            fprintf(stderr, "Expected %u, but got %u\n", id, new_out[id]);
            pass = 0;
         }
         fprintf(stderr, "%2d ", new_out[id]);
      }
      fprintf(stderr, "\n");
   }
   if (out[global_work_size * iterations] != prestore) {
      fprintf(stderr, "Prestore expected %d, got %d\n", prestore,
              out[global_work_size * iterations]);
      pass = 0;
   }

   if (pass) {
      fprintf(stderr, "Pass\n");
      return EXIT_SUCCESS;
   } else {
      fprintf(stderr, "Fail\n");
      return EXIT_FAILURE;
   }
}