/*
 * Utility routines for use in OpenCL programs.
 */
#ifndef OPENCL_UTILITY_H_
#define OPENCL_UTILITY_H_
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <CL/cl.h>
#include "filesize.h"

/*
 * Data structure for information about "platform" -- ID and list
 * of device IDs.
 */
typedef struct platform_info {
    cl_uint num_devices;
    cl_device_id* devices;
} platform_info_t;

/*
 * Macro to allow casting void* to other pointer type when compiled as C++.
 * Parameters are type, expression.
 */

/* 
 * Print error message:
 * msg is message to print, or NULL
 * rval is OpenCL return code, printed if other than CL_SUCCESS
 *   (text description for some values, otherwise numeric value)
 */
void describe_error(const char * msg, cl_int rval, FILE * outfile) 
{
    fprintf(outfile, "Error:  ");
    if (msg != NULL) fprintf(outfile, "%s  ", msg);
    if (rval != CL_SUCCESS) {
        fprintf(outfile, "(error %d)", (int)rval);
        /* partial list of possible errors */
        switch (rval) {
            case CL_INVALID_VALUE:
                fprintf(outfile, " (CL_INVALID_VALUE)");
                break;
            case CL_INVALID_DEVICE:
                fprintf(outfile, " (CL_INVALID_DEVICE)");
                break;
            case CL_INVALID_CONTEXT:
                fprintf(outfile, " (CL_INVALID_CONTEXT)");
                break;
            case CL_INVALID_QUEUE_PROPERTIES:
                fprintf(outfile, " (CL_INVALID_QUEUE_PROPERTIES)");
                break;
            case CL_OUT_OF_RESOURCES:
                fprintf(outfile, " (CL_OUT_OF_RESOURCES)");
                break;
            case CL_INVALID_BUILD_OPTIONS:
                fprintf(outfile, " (CL_INVALID_BUILD_OPTIONS)"); 
                break;
            case CL_BUILD_PROGRAM_FAILURE:
                fprintf(outfile, " (CL_BUILD_PROGRAM_FAILURE)"); 
                break;
            case CL_INVALID_PROGRAM:
                fprintf(outfile, " (CL_INVALID_PROGRAM)"); 
                break;
            case CL_INVALID_PROGRAM_EXECUTABLE:
                fprintf(outfile, " (CL_INVALID_PROGRAM_EXECUTABLE)"); 
                break;
            case CL_INVALID_KERNEL:
                fprintf(outfile, " (CL_INVALID_KERNEL)"); 
                break;
            case CL_INVALID_KERNEL_NAME:
                fprintf(outfile, " (CL_INVALID_KERNEL_NAME)"); 
                break;
            case CL_INVALID_KERNEL_DEFINITION:
                fprintf(outfile, " (CL_INVALID_KERNEL_DEFINITION)"); 
                break;
            case CL_INVALID_KERNEL_ARGS:
                fprintf(outfile, " (CL_INVALID_KERNEL_ARGS)"); 
                break;
            case CL_INVALID_ARG_VALUE:
                fprintf(outfile, " (CL_INVALID_ARG_VALUE)");
                break;
            case CL_INVALID_ARG_SIZE:
                fprintf(outfile, " (CL_INVALID_ARG_SIZE)");
                break;
            case CL_INVALID_WORK_DIMENSION:
                fprintf(outfile, " (CL_INVALID_WORK_DIMENSION)"); 
                break;
            case CL_INVALID_GLOBAL_OFFSET:
                fprintf(outfile, " (CL_INVALID_GLOBAL_OFFSET)"); 
                break;
            case CL_INVALID_WORK_GROUP_SIZE:
                fprintf(outfile, " (CL_INVALID_WORK_GROUP_SIZE)"); 
                break;
            case CL_INVALID_WORK_ITEM_SIZE:
                fprintf(outfile, " (CL_INVALID_WORK_ITEM_SIZE)"); 
                break;
            case CL_INVALID_IMAGE_SIZE:
                fprintf(outfile, " (CL_INVALID_IMAGE_SIZE)"); 
                break;
            case CL_INVALID_EVENT_WAIT_LIST:
                fprintf(outfile, " (CL_INVALID_EVENT_WAIT_LIST)"); 
                break;
            case CL_MEM_OBJECT_ALLOCATION_FAILURE:
                fprintf(outfile, " (CL_MEM_OBJECT_ALLOCATION_FAILURE)"); 
                break;
            case CL_OUT_OF_HOST_MEMORY:
                fprintf(outfile, " (CL_OUT_OF_HOST_MEMORY)"); 
                break;
            default:
                break;
        }
    }
    fprintf(outfile, "\n");
}

/* 
 * Print error message and exit. 
 * Uses 
 */
void error_exit(const char * error_message, cl_int rval)
{
    describe_error(error_message, rval, stderr);
    exit(EXIT_FAILURE);
}

/*
 * "malloc" that prints error message and exits if it fails.
 */
void * malloc_with_check(size_t sz, const char * message)
{
    void * p = malloc(sz);
    if (p == NULL) {
        fprintf(stderr, "Error:  Could not get space for %s\n", message);
        exit(EXIT_FAILURE);
    }
    return p;
}

/* 
 * Print information about OpenCL device to output stream(s).
 * Returns EXIT_SUCCESS or EXIT_FAILURE.
 */
int output_device_info(FILE* out_stream, cl_device_id device_id) {

    cl_int err;

    cl_device_type device_type;
    char *device_string = "";
    err = clGetDeviceInfo(device_id, CL_DEVICE_TYPE, 
            sizeof(device_type), &device_type, NULL);
    if (err == CL_SUCCESS) {
        if(device_type  == CL_DEVICE_TYPE_GPU)
            device_string = "GPU";
        else if (device_type == CL_DEVICE_TYPE_CPU)
            device_string = "CPU";
        else 
            device_string = "Unknown device type";
    } else {
        return EXIT_FAILURE;
    }

    cl_char name[1024] = {0};
    err = clGetDeviceInfo(device_id, CL_DEVICE_NAME, 
            sizeof(name), &name, NULL);

    cl_char vendor[1024] = {0};
    err |= clGetDeviceInfo(device_id, CL_DEVICE_VENDOR, 
                sizeof(vendor), &vendor, NULL);

    cl_uint units;
    err |= clGetDeviceInfo(device_id, CL_DEVICE_MAX_COMPUTE_UNITS, 
            sizeof(units), &units, NULL);

    size_t max_group_size;
    err |= clGetDeviceInfo(device_id, CL_DEVICE_MAX_WORK_GROUP_SIZE,
            sizeof(max_group_size), &max_group_size, NULL);

    if (err != CL_SUCCESS) {
        return EXIT_FAILURE;
    }

    fprintf(out_stream, "%s (%s)\n", name, vendor);
    fprintf(out_stream, "%s, %d compute units, maximum workgroup size %zd\n", 
            device_string, units, max_group_size);
 
    return EXIT_SUCCESS;
}

/* 
 * Get information about platforms.
 * (Helper function for other functions, but could be used alone?)
 * Prints messages about errors.
 * Returns pointer to array of cl_platform_id "objects" (NULL if none
 *   or fatal error), with count in *num_platforms_p.  Caller must free.
 */
cl_platform_id * getPlatforms(
    cl_uint * num_platforms_p
    )
{
    *num_platforms_p = 0;
    cl_platform_id *platforms = NULL;

    cl_int rval = clGetPlatformIDs(0, NULL, num_platforms_p);
    if (rval != CL_SUCCESS) {
        describe_error("Could not get number of platforms", rval, stderr);
        return NULL;
    }
    if (*num_platforms_p > 0) {
        platforms = malloc_with_check(
                *num_platforms_p * sizeof(*platforms),
                "platform info");
        rval = clGetPlatformIDs(*num_platforms_p, platforms, NULL);
        if (rval != CL_SUCCESS) {
            describe_error("Could not get platform(s)", rval, stderr);
            return NULL;
        }
    } 
    else {
        fprintf(stderr, "No platforms\n");
        return NULL;
    }
    return platforms;
}

/* 
 * Get information about devices of specified type.
 * Prints messages about errors.
 * Returns pointer to array of number of platform_info "objects"
 *   (NULL if none or fatal error).
 * Caller should call freePlatformInfo to free.
 */
platform_info_t * getAllDeviceInfo(
    cl_device_type device_type,
    cl_uint * num_platforms_p,
    cl_bool report_not_found /* FIXME or no? */
    )
{
    *num_platforms_p = 0;
    cl_platform_id *platforms = getPlatforms(num_platforms_p);
    platform_info_t *platform_info = NULL;

    if ((platforms == NULL) || (*num_platforms_p == 0)) {
        fprintf(stderr, "No platforms\n");
        return NULL;
    } 

    platform_info = malloc_with_check(
            *num_platforms_p * sizeof(*platform_info),
            "platform info");
    cl_int rval = CL_SUCCESS;

    for (cl_uint i = 0; (i < *num_platforms_p); ++i) {
        cl_uint temp_count = 0;
        cl_device_id * temp_devices = NULL;
        rval = 
            clGetDeviceIDs(platforms[i], device_type, 0, NULL, &temp_count);
        if (rval != CL_SUCCESS) {
            /* FIXME improve? */
            if (report_not_found) {
                fprintf(stderr, "Error on platform %d\n",i);
                describe_error("Could not get devices", rval, stderr);
            }
        }
        else {
            temp_devices = malloc_with_check(
                    temp_count * sizeof(*temp_devices), 
                    "device info");
            rval = 
                clGetDeviceIDs(platforms[i], device_type, 
                        temp_count, temp_devices, NULL);
            if (rval != CL_SUCCESS) {
                if (report_not_found) {
                    fprintf(stderr, "Error on platform %d\n",i);
                    describe_error("Could not get devices", rval, stderr);
                }
            }
        }
        platform_info[i].num_devices = temp_count;
        platform_info[i].devices = temp_devices;
    }
    free(platforms);
    return platform_info;
}

/*
 * Free array of platform_info objects.
 */
void freePlatformInfo(
    cl_uint count,
    platform_info_t * platform_info
    )
{
    for (cl_uint i = 0; i < count; ++i) {
        free(platform_info[i].devices);
    }
}

/* 
 * Get first device of requested type.
 * Exits program on error.
 */
void getDevice(
    cl_device_type device_type,
    cl_device_id * device_id_ptr
    )
{
    cl_uint num_platforms = 0;
    cl_platform_id *platforms = getPlatforms(&num_platforms);

    if ((platforms == NULL) || (num_platforms == 0)) {
        fprintf(stderr, "No platforms\n");
        free(platforms);
        return;
    } 

    cl_int rval = CL_SUCCESS;
    cl_bool found = CL_FALSE;
    for (cl_uint i = 0; (i < num_platforms) && !found; ++i) {
        rval = 
            clGetDeviceIDs(platforms[i], device_type, 1, device_id_ptr, NULL);
        if (rval == CL_SUCCESS) {
            found = CL_TRUE;
        }
    }
    free(platforms);

    if (!found) {
        fprintf(stderr, "Error:  Could not get device of requested type\n");
        exit(EXIT_FAILURE);
    }
}

/* 
 * Do standard initialization.
 * Device ID is input; obtain with getAllDeviceInfo() or getDevice().
 * Prints error message(s) and exits program on error.
 */
void initialize(
    cl_device_id * device_id_ptr,
    cl_context * context_ptr,
    cl_command_queue * commands_ptr,
    cl_program * program_obj_ptr,
    cl_uint num_source_strings,
    const char* kernel_source[]
    )
{
    cl_int rval;

    /* create a compute context */
    *context_ptr = clCreateContext(0, 1, device_id_ptr, NULL, NULL, &rval);
    if ((rval != CL_SUCCESS) || (*context_ptr == NULL))
        error_exit("Could not create compute context", rval);

    /* create a command queue */
    *commands_ptr = clCreateCommandQueue(*context_ptr, *device_id_ptr, 
            0, &rval);
    if ((rval != CL_SUCCESS) || (*commands_ptr == NULL))
        error_exit("Could not create command queue", rval);

    /* create compute program object from source buffer */
    *program_obj_ptr = clCreateProgramWithSource(*context_ptr, 
            num_source_strings, (const char**) kernel_source, NULL, &rval);
    if ((rval != CL_SUCCESS) || (*program_obj_ptr == NULL))
        error_exit("Could not create program object from source", rval);

    /* build program object */
    rval = clBuildProgram(*program_obj_ptr, 0, NULL, NULL, NULL, NULL);
    if (rval != CL_SUCCESS) {
        size_t len;
        char buffer[4096];

        describe_error("Could not build kernel executable", rval, stderr);
        clGetProgramBuildInfo(*program_obj_ptr, *device_id_ptr, 
                CL_PROGRAM_BUILD_LOG, 
                sizeof(buffer), buffer, &len);
        fprintf(stderr, "%s\n", buffer);
        exit(EXIT_FAILURE);
    }
}

/*
 * Do standard cleanup.
 */
void finalize(
    cl_context context,
    cl_command_queue commands,
    cl_program program_obj)
{
    clReleaseProgram(program_obj);
    clReleaseCommandQueue(commands);
    clReleaseContext(context);
}

/*
 * Load function(s) from file.
 * Returns C-style string (pointer to malloc'd space) or NULL.
 * Caller must free allocated string.
 */
char* source_from_file(const char* filename) {
    long file_size = filesize(filename);
    if (file_size < 0) return NULL;
    FILE* in = fopen(filename, "r");
    if (in == NULL) return NULL;
    char* data = malloc_with_check(file_size+1, "file contents");
    size_t data_pos = 0;
    int input;
    while ((input = fgetc(in)) != EOF) {
        data[data_pos++] = input;
    }
    data[data_pos] = '\0';
    fclose(in);
    return data;
}

#endif /* #ifndef OPENCL_UTILITY_H_ */