#include <GL/glew.h>
#include <GLFW/glfw3.h>
#include <cmath>
#include <iostream>
#include <vector>
#include "shader.h"

// Constants
const float g = 9.81f; // Gravitational acceleration
const float r = 0.6f; // Length of the string
const float initialDt = 0.01f; // Initial time step

const float initTheta = 0.785f; // Initial angle
const float initOmega = 2.0f; // Initial angular velocity
const float initPrevTheta = initTheta - initialDt * initOmega; // Initial previous angle


float dt = initialDt;
float currentT = 0.0f;

// Pendulum state structure
struct Pendulum {
    float theta; // Angle
    float omega; // Angular velocity
    float prev_theta; // Previous angle (for Verlet method)
};

// Pendulums
Pendulum pendulums[3] = {
        {initTheta, initOmega, initPrevTheta}, // Pendulum 1 (Euler)
        {initTheta, initOmega, initPrevTheta}, // Pendulum 2 (Verlet)
        {initTheta, initOmega, initPrevTheta}, // Pendulum 3 (Runge-Kutta)
};

GLuint VBO, VAO, shaderProgram;

void compileShaders(){
    Shader shader("pendulum_vs.glsl", "pendulum_fs.glsl");
    shaderProgram = shader.programID();
}

void initOpenGL()
{
    compileShaders();

    // Setup Vertex Array Object and Vertex Buffer Object
    glGenVertexArrays(1, &VAO);
    glGenBuffers(1, &VBO);

    glBindVertexArray(VAO);

    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(float) * 40, nullptr, GL_DYNAMIC_DRAW); // 4 pendulums * 2 vertices * (2 positions + 3 colors)

    // Position attribute
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 5 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);

    // Color attribute
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 5 * sizeof(float), (void*)(2 * sizeof(float)));
    glEnableVertexAttribArray(1);

    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindVertexArray(0);
}

float f(float theta){
    return -(g / r) * sin(theta);
}

// Euler method
void updateEuler(Pendulum& p)
{
    p.theta = p.theta + dt * p.omega;
    p.omega = p.omega + f(p.theta) * dt;
}

// Verlet method
void updateVerlet(Pendulum& p)
{
    float new_theta = 2 * p.theta - p.prev_theta + f(p.theta) * dt * dt;
    float old_theta = p.theta;
    p.theta = new_theta;
    p.prev_theta = old_theta;

    // Aktualizacja omega
    p.omega = (p.theta - p.prev_theta) / (2 * dt);
}

// Runge-Kutta 4th order method
void updateRungeKutta(Pendulum& p)
{

    float k1_theta = p.omega;
    float k1_omega = f(p.theta);

    float k2_theta = p.omega + 0.5f * dt * k1_omega;
    float k2_omega = f(p.theta + 0.5f * dt * k1_theta);

    float k3_theta = p.omega + 0.5f * dt * k2_omega;
    float k3_omega = f(p.theta + 0.5f * dt * k2_theta);

    float k4_theta = p.omega + dt * k3_omega;
    float k4_omega = f(p.theta + dt * k3_theta);

    p.theta += (dt / 6.0f) * (k1_theta + 2.0f * k2_theta + 2.0f * k3_theta + k4_theta);
    p.omega += (dt / 6.0f) * (k1_omega + 2.0f * k2_omega + 2.0f * k3_omega + k4_omega);
}

void drawPendulums()
{
    glClear(GL_COLOR_BUFFER_BIT);
    glUseProgram(shaderProgram);
    glBindVertexArray(VAO);

    std::vector<std::vector<float>> colors = {
            {1.0f, 0.0f, 0.0f}, // EULER
            {0.0f, 1.0f, 0.0f}, // VERLET
            {0.0f, 0.0f, 1.0f}, // RK
            {1.0f, 1.0f, 0.0f}  // SIMPLE
    };

    std::vector<float> vertices;

    for (int i = 0; i < 3; ++i) {
        float x = r * sin(pendulums[i].theta);
        float y = -r * cos(pendulums[i].theta);

        // Append vertices for the current pendulum
        vertices.insert(vertices.end(), {
                0.0f, 0.0f, colors[i][0], colors[i][1], colors[i][2], // Attachment point
                x, y, colors[i][0], colors[i][1], colors[i][2]        // Mass point
        });

        // Debugging output for each pendulum's vertices
//        std::cout << "Pendulum " << i + 1 << ": (" << x << ", " << y << "), Color: ("
//                  << colors[i][0] << ", " << colors[i][1] << ", " << colors[i][2] << ")" << std::endl;
    }

    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferSubData(GL_ARRAY_BUFFER, 0, sizeof(float) * vertices.size(), vertices.data());

    for (int i = 0; i < 3; ++i) {
        glDrawArrays(GL_LINES, i * 2, 2);
        glDrawArrays(GL_POINTS, i * 2 + 1, 1);
    }

    glBindVertexArray(0);
}

int main()
{
    if (!glfwInit())
    {
        std::cerr << "Failed to initialize GLFW" << std::endl;
        return -1;
    }

    glfwWindowHint(GLFW_SAMPLES, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    GLFWwindow* window = glfwCreateWindow(800, 600, "Mathematical Pendulums", NULL, NULL);
    if (!window)
    {
        std::cerr << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return -1;
    }

    glfwMakeContextCurrent(window);
    glewExperimental = GL_TRUE;
    if (glewInit() != GLEW_OK)
    {
        std::cerr << "Failed to initialize GLEW" << std::endl;
        return -1;
    }

    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    glOrtho(-2, 2, -2, 2, -1, 1);
    glMatrixMode(GL_MODELVIEW);

    glPointSize(35.0f);

    initOpenGL();


    while (!glfwWindowShouldClose(window))
    {
        updateEuler(pendulums[0]);
        updateVerlet(pendulums[1]);
        updateRungeKutta(pendulums[2]);

        drawPendulums();

        glfwSwapBuffers(window);
        glfwPollEvents();

        currentT += dt;

    }

    glDeleteVertexArrays(1, &VAO);
    glDeleteBuffers(1, &VBO);

    glfwDestroyWindow(window);
    glfwTerminate();

    return 0;
}