/* std */
#include <iostream>
#include <vector>
#include <string>
#include <stdlib.h>
#include <unistd.h>  
#include <cmath>  

/* OpenGL */
#include "glad/glad.h"
#include <GLFW/glfw3.h>

/* Project classes */
#include "Helper.hpp"
#include "Config.hpp"
#include "Cell.hpp"
#include "Particle.hpp"
#include "ComputeShader.hpp"
#include "Application.hpp"
#include "Buffer.hpp"
#include "Canvas.hpp"

constexpr float deg_to_rad = (float) (3.141592653589793 / 180.0);

/* Keys callbacks */
Config config;
ComputeShader cs_draw;
Buffer particles_buffer;
Buffer map_buffer;

void key_callback(GLFWwindow* window, int key, int scancode, int action, int mods){
    Config new_conf = config;
    if(action != GLFW_PRESS)
        return;
    char filename[100];
    if(glfwGetKey(window, GLFW_KEY_LEFT_SHIFT) == GLFW_RELEASE
       && (key >= GLFW_KEY_0 && key <= GLFW_KEY_9)){
        char pattern[] = "./configs/config%d.cfg";
        sprintf(filename, pattern, key - GLFW_KEY_0);
        if(load_config(std::string(filename), &new_conf)){
            config = new_conf;
            std::cout << "Loaded '" << filename << "'" << std::endl;
        }else{
            std::cout << "Failed to load '" << filename << "'" << std::endl;
        }
    }else if(glfwGetKey(window, GLFW_KEY_LEFT_SHIFT) == GLFW_PRESS
             && (key >= GLFW_KEY_0 && key <= GLFW_KEY_9)){
        char pattern[] = "./shaders/draw%d.glsl";
        sprintf(filename, pattern, key - GLFW_KEY_0);
        cs_draw.load(filename);
        cs_draw.compile();
        particles_buffer.bind_base(cs_draw.get_program(), "particles_buffer", 1);
        map_buffer.bind_base(cs_draw.get_program(), "map_buffer", 2);
    }else if (key == GLFW_KEY_C){
        float distance = std::min(config.map_width, config.map_height) / 5;
        std::vector<Particle> particles(config.population);
        for(auto &particle : particles){
            float angle = random_float(0, 360) * deg_to_rad;
            particle.dir_x = cosf(angle);
            particle.dir_y = sinf(angle);

            particle.pos_x = config.map_width / 2 + particle.dir_x * distance;
            particle.pos_y = config.map_height / 2 + particle.dir_y * distance;
        }
        particles_buffer.set(particles.data(), particles.size(), sizeof(particles[0]));

        std::vector<Cell> map(config.map_width * config.map_height);
        for(auto &cell : map){
            cell = {
                .pheromone = 0,
                .result = 0,
                .population = 0,
            };
        }
        map_buffer.set(map.data(), map.size(), sizeof(map[0]));
    }
}

int main(int argc, char *argv[]){
    srand(0);
    std::string filename;
    Application app;
    ComputeShader cs_particles_update;
    ComputeShader cs_map_update;
    ComputeShader cs_map_reset;
    Canvas canvas;

    int map_width =  1500;
    int map_height = 800;
    size_t population = 5000000;
    if(argc > 2){
        map_width  = strtol(argv[1], nullptr, 0);
        map_height = strtol(argv[2], nullptr, 0);
    }
    if(argc > 3){
        population = strtol(argv[3], nullptr, 0);
    }

    // Default value
    config = {
        .map_width=map_width,
        .map_height=map_height,

        .population=population,
 
        .increase=10.0f,
        .decay=0.5f,
        .diffusion=5.0f,
 
        .speed=60.0f,
        .steering=180.0f * deg_to_rad,
        .sensing_distance=3.0f,
        .sensing_angle=45.0f * deg_to_rad,
    };

    Config new_conf = config;
    if(load_config("./configs/config0.cfg", &new_conf)){
        config = new_conf;
    }

    float average_population = (float)config.population / (float)(config.map_width * config.map_height);
    float average_pheromone = (float)config.increase * average_population;

    app = Application(config.map_width, config.map_height);
    app.init();
    app.create_window();

    /* Load and compile particles update shader */
    cs_particles_update.load("./shaders/particles_update.glsl");
    cs_particles_update.compile();

    /* Load and compile map update shader */
    cs_map_update.load("./shaders/map_update.glsl");
    cs_map_update.compile();

    /* Load and compile map reset shader */
    cs_map_reset.load("./shaders/map_reset.glsl");
    cs_map_reset.compile();

    /* Load and compile draw shader */
    cs_draw.load("./shaders/draw0.glsl");
    cs_draw.compile();

    /* Particles init */
    std::vector<Particle> particles(config.population);
    for(auto &particle : particles){
        particle.pos_x = random_float(0.0f, config.map_width);
        particle.pos_y = random_float(0.0f, config.map_height);

        float angle = random_float(0, 360 * deg_to_rad);
        particle.dir_x = cosf(angle);
        particle.dir_y = sinf(angle);
    }

    /* Create and fill particles buffer */
    particles_buffer.create();
    particles_buffer.set(particles.data(), particles.size(), sizeof(particles[0]));

    /* Map init */
    std::vector<Cell> map(config.map_width * config.map_height);
    for(auto &cell : map){
        cell = {
            .pheromone = 0,
            .result = 0,
            .population = 0,
        };
    }

    /* Create and fill map buffer */
    map_buffer.create();
    map_buffer.set(map.data(), map.size(), sizeof(map[0]));

    /* Create Canvas */
    canvas = Canvas(config.map_width, config.map_width);
    canvas.create();

    /* Bind buffers to compute shaders */
    particles_buffer.bind_base(cs_particles_update.get_program(), "particles_buffer", 1);
    map_buffer.bind_base(cs_particles_update.get_program(), "map_buffer", 2);
    
    map_buffer.bind_base(cs_map_update.get_program(), "map_buffer", 2);
    
    particles_buffer.bind_base(cs_draw.get_program(), "particles_buffer", 1);
    map_buffer.bind_base(cs_draw.get_program(), "map_buffer", 2);
    
    map_buffer.bind_base(cs_map_reset.get_program(), "map_buffer", 2);

    /* Register key callbacks */
    app.register_callback(key_callback);

    /* Main loop */
    float pervious_time = app.time();
    while(app.running()){
        float tick = app.time();
        float dt = tick - pervious_time;
        pervious_time = tick;

        /* Update uniforms */
        particle_uniform(config, cs_particles_update);    
        map_update_uniform(config, cs_map_update);    
        map_reset_uniform(config, cs_map_reset);    
        cs_draw.set_uniform("map_width", config.map_width);    
        cs_draw.set_uniform("map_height", config.map_height);    
        cs_draw.set_uniform("display_population_value", average_population);
        cs_draw.set_uniform("display_pheromone_value", average_pheromone);

        /* Update particles */
        cs_particles_update.set_uniform("dt", dt);
        cs_particles_update.execute(config.population, 256);
        ComputeShader::wait();

        /* Update map */
        cs_map_update.set_uniform("dt", dt);
        cs_map_update.execute(config.map_width, config.map_height, 16, 16);
        ComputeShader::wait();

        /* Draw map */
        cs_draw.activate();
        canvas.bind();
        cs_draw.execute(config.map_width, config.map_height, 16, 16);
        ComputeShader::wait();
        canvas.draw();
        canvas.unbind();

        /* Reset map */
        cs_map_reset.execute(config.map_width, config.map_height, 16, 16);
        ComputeShader::wait();

        /* Display, move on */
        app.update();
    }

    particles_buffer.release();
    map_buffer.release();
    canvas.release();
    cs_draw.release();
    cs_map_reset.release();
    cs_map_update.release();
    cs_particles_update.release();
    app.terminate();

    return 0;
}