#define _USE_MATH_DEFINES
//#define HYDROGEN

#include <math.h>
#include <cmath>
#include <array>
#include <chrono>
#include <string>
#include <vector>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>

#include "DataStructures.h"
#include <ComputeEngine.h>

using std::cout;
using std::endl;

void UpdateKeys(ComputeEngine &renderer, std::vector<SceneObject> &objects, shader_data_t &shader_data, Texture &tex_output)
{
    if (renderer.IsKeyClicked(GLFW_KEY_P))
    {
        renderer.SwitchInput();
    }
    if (renderer.GetInput())
    {
        if (renderer.IsKeyClicked(GLFW_KEY_F11))
            renderer.SwitchFullScreen();
        if (renderer.IsKeyClicked(GLFW_KEY_ESCAPE))
            renderer.CloseWindow();
        if (renderer.IsKeyClicked(GLFW_KEY_L))
            renderer.SaveScreen("screenshots/screenshot");
    }
    int width, height;
    shader_data.GetWindowSize(width, height);
    renderer.GetWindowSize(shader_data.screenSize[0], shader_data.screenSize[1]);
    if (width != shader_data.GetWidth() || height != shader_data.GetHeight())
        tex_output.Resize(shader_data.GetWidth(), shader_data.GetHeight());
}

void UpdateInput(ComputeEngine &renderer, std::vector<SceneObject> &objects, shader_data_t &shader_data, Texture &tex_output)
{
    renderer.PollEvents();
    UpdateKeys(renderer, objects, shader_data, tex_output);

    double dx, dy;
    renderer.GetMouseDelta(dx, dy);
    dx *= (renderer.GetFrametime() / 1000.0);
    dy *= (renderer.GetFrametime() / 1000.0);

    shader_data.mouse_movement[0] = (float)dx;
    shader_data.mouse_movement[1] = (float)dy;
}

int main(void)
{
    srand((unsigned)time(0));

    int width = 1024;
    int height = 1024;

    // init data for shader
    shader_data_t shader_data = {};
    shader_data.screenSize = {width, height};
#ifdef HYDROGEN
    // init renderer
    ComputeEngine renderer(width, height, "Hydrogen wave function", true);
    Shader shader("shaders/hydrogen_wave/pixelCompute.compute");
    Program pixel_program(shader);
#else
    // init renderer
    ComputeEngine renderer(width, height, "Attractors", true);
    Shader shader("shaders/attractors/objectCompute.compute");
    Shader pixelShader("shaders/attractors/pixelCompute.compute");
    Program compute_program(shader);
    Program pixel_program(pixelShader);
#endif

    // init textures
    Texture tex_output("img_output", width, height);

    // init buffers
    UBO data_buffer(2);

    std::vector<SceneObject> objects;
    // create scene
#ifdef HYDROGEN
    shader_data.iterations = 2;
    shader_data.samples = 1;
    shader_data.dt = .05f;
#else
    shader_data.dt = 1.5f;
    shader_data.iterations = 10;
    SSBO objects_buffer(3);

    int points = 20000;
    float mul = 25;

    objects.reserve(points);
    for (int i = 0; i < points; i++)
    {
        objects.emplace_back();
        objects[i].position[0] = (getRand() - 0.5f) * mul;
        objects[i].position[1] = (getRand() - 0.5f) * mul;
        objects[i].position[2] = (getRand() - 0.5f) * mul;
    }

    objects_buffer.Set(&objects);
#endif

    // fill buffers with data
    data_buffer.Set(&shader_data);

    // bind textures
    pixel_program.BindTextureImage(tex_output, 0);

#ifdef HYDROGEN
    pixel_program.SetUniform1i("width", width);
    pixel_program.SetUniform1i("height", height);
    int n, l, m;
    n = 4;
    l = 2;
    m = 1;
    pixel_program.SetUniform1i("quantum_n", n);
    pixel_program.SetUniform1i("quantum_l", l);
    pixel_program.SetUniform1i("quantum_m", m);

    float phi = 0;

    pixel_program.DispatchCompute(shader_data.GetWidth(), shader_data.GetHeight());
#else
    compute_program.BindTextureImage(tex_output, 0);

    pixel_program.SetUniform1i("_pass", 0);
    pixel_program.DispatchCompute(shader_data.GetWidth(), shader_data.GetHeight());
    printf("points: %d\n", (int)objects.size());
#endif

    printf("width: %d\ height: %d\n", shader_data.GetWidth(), shader_data.GetHeight());

    while (!renderer.ShouldClose())
    {
#ifdef HYDROGEN
        // update input
        UpdateInput(renderer, objects, shader_data, tex_output);

        phi += (renderer.GetFrametime() / 1000.0) * M_PI * 0.5;

        phi = std::fmod(phi, M_PI);
        printf("phi: %.2f\n", phi);

        // update variables
        pixel_program.SetUniform1f("u_phi", phi);

        // dispatch compute shader
        pixel_program.DispatchCompute(shader_data.GetWidth(), shader_data.GetHeight());

        renderer.UpdateFrametime();

        // draw rendered image
        renderer.DrawTexture(tex_output);
#else
        // update input
        UpdateInput(renderer, objects, shader_data, tex_output);

        shader_data.UpdateSeed();
        // update buffers
        data_buffer.Update(&shader_data);

        if(shader_data.mouse_movement[0] != 0 || shader_data.mouse_movement[1] != 0)
            pixel_program.SetUniform1i("_pass", 0);
        else
            pixel_program.SetUniform1i("_pass", renderer.GetFramecount());

        // dispatch compute shader
        compute_program.DispatchCompute((int)objects.size());

        // dispatch compute shader
        pixel_program.DispatchCompute(shader_data.GetWidth(), shader_data.GetHeight());

        renderer.UpdateFrametime();
        renderer.DrawTexture(tex_output);
#endif
    }
    // print average frametime
    printf("Avg frametime: %.3f", renderer.GetAverageFrametime());

    return 0;
}