#include <SDL2/SDL.h>

#define CL_HPP_TARGET_OPENCL_VERSION 200
#include <CL/opencl.hpp>

#include <iostream>
#include <vector>
#include <chrono>
#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <sstream>
#include <fstream>

struct Agent
{
    float x,y,direction;
};

bool get_source(std::string* source);

unsigned long long sys_time_us();

int main(int argc, char* argv[])
{
    srand(time(0));
    uint32_t u32_width, u32_height;
    uint32_t u32_dimensions[2];

    if (argc == 1)
    {
        u32_width = u32_height = 300;
    }
    else
    {
        u32_width = atoi(argv[1]);
        u32_height = atoi(argv[2]);
    }

    int i_size = u32_width*u32_height;

    u32_dimensions[0] = u32_width;
    u32_dimensions[1] = u32_height;

    if (SDL_Init(SDL_INIT_VIDEO) < 0)
    {
        std::cout << "Video Init Error." << std::endl;
    }

    SDL_Window* window = SDL_CreateWindow("Slime Mold Simulator", SDL_WINDOWPOS_CENTERED-u32_width/2, SDL_WINDOWPOS_CENTERED-u32_height/2, u32_width, u32_height, SDL_WINDOW_SHOWN);

    SDL_Renderer* renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED | SDL_RENDERER_TARGETTEXTURE);

    SDL_RenderClear(renderer);
    SDL_RenderPresent(renderer);

    SDL_Texture* field_texture = SDL_CreateTexture(renderer, SDL_GetWindowPixelFormat(window), SDL_TEXTUREACCESS_STREAMING, u32_width, u32_height);

    SDL_Texture* start_texture = SDL_CreateTexture(renderer, SDL_GetWindowPixelFormat(window), SDL_TEXTUREACCESS_STREAMING, u32_width, u32_height);

    float* ground = new float[i_size];
    for(int i = 0; i < i_size; i++)
    {
        ground[i] = 0.0f;
    }

    int i_agent_count = (int)sqrt(u32_width*u32_height)*1000;
    Agent* agents = new Agent[i_agent_count];

    {
        uint8_t* u8_pixels;
        int i_pitch;
        SDL_LockTexture(start_texture, NULL, (void**)&u8_pixels,&i_pitch);

        int i_created = 0;
        uint32_t u32_radius = 1*((u32_width<u32_height)?u32_width:u32_height)/10;
        while(i_created < i_agent_count)
        {
            float r = u32_radius*sqrt((float)rand()/(float)RAND_MAX);
            float theta = 2.0f*M_PI*(float)rand()/(float)RAND_MAX;
            int x = r*cos(theta)+(int)u32_width/2;
            int y = r*sin(theta)+(int)u32_height/2;
            agents[i_created].x = x;
            agents[i_created].y = y;
            agents[i_created].direction = theta+M_PI;

            u8_pixels[4*(x+y*u32_width)] = 255;
            u8_pixels[4*(x+y*u32_width)+1] = 255;
            u8_pixels[4*(x+y*u32_width)+2] = 255;
            u8_pixels[4*(x+y*u32_width)+3] = 255;

            i_created++;
        }

        SDL_UnlockTexture(start_texture);
    }

    /*
    bool* bl_start_positions = new bool[u32_width*u32_height];
    for(int i = 0; i < u32_width*u32_height; i++)
    {
        bl_start_positions[i] = false;
    }

    int i_created = 0;
    while(i_created < i_agent_count)
    {
        int rng_cur = rand()%(u32_width*u32_height/i_agent_count);
        int i_x = rand()%u32_width;
        int i_y = rand()%u32_height;
        if (rng_cur == 0 && !bl_start_positions[i_x+i_y*u32_width])
        {
            bl_start_positions[i_x+i_y*u32_width] = true;

            agents[i_created].x = (float)i_x;
            agents[i_created].y = (float)i_y;
            agents[i_created].direction = 2.0f*M_PI*((float)(rand()%100)/100.0f);
            i_created++;
            //ground[i_x+i_y*u32_width] = 1.0f;
        }
    }

    {
        uint8_t* u8_pixels;
        int i_pitch;
        SDL_LockTexture(start_texture, NULL, (void**)&u8_pixels,&i_pitch);
        for(uint32_t i = 0; i < u32_width; i++)
        {
            for(uint32_t j = 0; j < u32_height; j++)
            {
                u8_pixels[4*(i+j*u32_width)] = (bl_start_positions[i+j*u32_width])?255:0;
                u8_pixels[4*(i+j*u32_width)+1] = (bl_start_positions[i+j*u32_width])?255:0;
                u8_pixels[4*(i+j*u32_width)+2] = (bl_start_positions[i+j*u32_width])?255:0;
                u8_pixels[4*(i+j*u32_width)+3] = 255;
            }
        }
        SDL_UnlockTexture(start_texture);
    }

    delete[] bl_start_positions;//*/

    std::vector<cl::Platform> all_platforms;
    cl::Platform::get(&all_platforms);

    cl::Platform default_platform=all_platforms[0];

    std::cout << "Using platform :" << default_platform.getInfo<CL_PLATFORM_NAME>() << std::endl;

    std::vector<cl::Device> all_devices;
    default_platform.getDevices(CL_DEVICE_TYPE_ALL, &all_devices);

    if (all_devices.size() == 0)
    {
        std::cout << "Could not find OpenCL device." << std::endl;
        exit(1);
    }

    cl::Device default_device = all_devices[0];

    std::cout << "Using device :" << default_device.getInfo<CL_DEVICE_NAME>() << std::endl;

    cl::Context context({default_device});

    std::string kernel_code;
    if (!get_source(&kernel_code))
    {
        std::cout << "Error reading kernel source." << std::endl;
        exit(1);
    }
      
    cl::Program::Sources sources;
    sources.push_back({kernel_code.c_str(), kernel_code.length()});
    cl::Program program(context, sources);

    if (program.build({default_device}, "-cl-std=CL2.0") != CL_SUCCESS)
    {
        std::cout << "Error building: " << program.getBuildInfo<CL_PROGRAM_BUILD_LOG>(default_device) << std::endl;
        exit(1);
    }

    cl::Buffer buffer_agents(context, CL_MEM_READ_WRITE, sizeof(Agent)*i_agent_count);
    cl::Buffer buffer_ground(context, CL_MEM_READ_WRITE, sizeof(float)*i_size);
    cl::Buffer buffer_ground_swap(context, CL_MEM_READ_WRITE, sizeof(float)*i_size);
    cl::Buffer buffer_ground_dimensions(context, CL_MEM_READ_WRITE, sizeof(uint32_t)*2);
    cl::Buffer buffer_draw_data(context, CL_MEM_READ_WRITE, sizeof(uint8_t)*4*i_size);

    cl::CommandQueue queue(context,default_device);

    queue.enqueueWriteBuffer(buffer_agents, CL_TRUE, 0, sizeof(Agent)*i_agent_count, agents);
    queue.enqueueWriteBuffer(buffer_ground, CL_TRUE, 0, sizeof(float)*i_size, ground);
    queue.enqueueWriteBuffer(buffer_ground_swap, CL_TRUE, 0, sizeof(float)*i_size, ground);
    queue.enqueueWriteBuffer(buffer_ground_dimensions, CL_TRUE, 0, sizeof(uint32_t)*2, &u32_dimensions[0]);

    cl::Kernel kernel_initialize_globals(program, "initialize_globals");
    kernel_initialize_globals.setArg(0, buffer_agents);
    kernel_initialize_globals.setArg(1, buffer_ground);
    kernel_initialize_globals.setArg(2, buffer_ground_swap);
    kernel_initialize_globals.setArg(3, buffer_ground_dimensions);

    queue.enqueueNDRangeKernel(kernel_initialize_globals, cl::NullRange, cl::NDRange(1), cl::NullRange);
    queue.finish();

    cl::Kernel kernel_update_ground(program, "update_ground");
    cl::Kernel kernel_iterate_ground(program, "iterate_ground");
    cl::Kernel kernel_move_agents(program, "move_agents");
    cl::Kernel kernel_emit_agent_pheromones(program, "emit_agent_pheromones");
    cl::Kernel kernel_fill_draw_data(program, "fill_draw_data");

    queue.enqueueNDRangeKernel(kernel_emit_agent_pheromones, cl::NullRange, cl::NDRange(i_agent_count), cl::NullRange);
    queue.enqueueReadBuffer(buffer_ground, CL_TRUE, 0, sizeof(float)*i_size, ground);
    queue.finish();//*/

    SDL_Rect screen_rect = {0,0,(int)u32_width,(int)u32_height};

    unsigned long long ull_time_old, ull_time_now;
    ull_time_now = ull_time_old = sys_time_us();

    int iteration_duration = 16666/10;
    bool pause = true;

    while(true)
    {
        SDL_Event event;
        if (1 == SDL_PollEvent(&event))
        {
            if (event.type == SDL_KEYDOWN)
            {
                if (event.key.keysym.scancode == SDL_SCANCODE_Q)
                {
                    break;
                }
                else if (event.key.keysym.scancode == SDL_SCANCODE_P)
                {
                    pause = !pause;
                }
                else if (event.key.keysym.scancode == SDL_SCANCODE_T)
                {
                    SDL_SetRenderDrawColor(renderer, 0,0,0,0);
                    SDL_RenderClear(renderer);

                    SDL_RenderCopy(renderer, start_texture, &screen_rect, &screen_rect);

                    SDL_RenderPresent(renderer);
                }
            }
        }

        ull_time_now = sys_time_us();
        if (ull_time_now >= ull_time_old + iteration_duration)
        {
            ull_time_old = ull_time_now;
            if (pause)
            {
                continue;
            }

            
            queue.enqueueNDRangeKernel(kernel_update_ground, cl::NullRange, cl::NDRange(i_size), cl::NullRange);
            queue.enqueueNDRangeKernel(kernel_iterate_ground, cl::NullRange, cl::NDRange(i_size), cl::NullRange);
            
            queue.enqueueNDRangeKernel(kernel_move_agents, cl::NullRange, cl::NDRange(i_agent_count), cl::NullRange);
            queue.enqueueNDRangeKernel(kernel_emit_agent_pheromones, cl::NullRange, cl::NDRange(i_agent_count), cl::NullRange);

            kernel_fill_draw_data.setArg(0, buffer_draw_data);
            queue.enqueueNDRangeKernel(kernel_fill_draw_data, cl::NullRange, cl::NDRange(i_size), cl::NullRange);

            uint8_t* u8_pixels;
            int i_pitch;
            SDL_LockTexture(field_texture, NULL, (void**)&u8_pixels, &i_pitch);

            queue.enqueueReadBuffer(buffer_draw_data, CL_TRUE, 0, sizeof(uint8_t)*4*i_size, u8_pixels);
            queue.finish();

            SDL_UnlockTexture(field_texture);

            SDL_SetRenderDrawColor(renderer, 0,0,0,0);
            SDL_RenderClear(renderer);

            SDL_RenderCopy(renderer, field_texture, &screen_rect, &screen_rect);

            SDL_RenderPresent(renderer);
        }
    }

    SDL_DestroyRenderer(renderer);
    SDL_DestroyWindow(window);

    SDL_Quit();

    delete[] agents;
    delete[] ground;

    return 0;
}

bool get_source(std::string* source)
{
    if (source == nullptr)
    {
        return false;
    }

    std::ifstream in_file("gpu_kernel.clcpp", std::ios::in);
    std::stringstream buffer;

    buffer << in_file.rdbuf();

    *source = buffer.str();

    return true;
}

unsigned long long sys_time_us()
{
    return (unsigned long long)(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count());
}