#include <GL/glew.h>
#include <GL/freeglut.h>
#include <iostream>
#include <vector>
#include <cmath>

const char* VERTEX_SHADER = 
"#version 330 core\n"
"layout(location = 0) in vec3 position;\n"
"layout(location = 1) in vec3 color;\n"
"out vec3 fragColor;\n"
"void main() {\n"
"    fragColor = color;\n"
"    gl_Position = vec4(position, 1.0);\n"
"    gl_PointSize = 5.0;\n"
"}\n";

const char* FRAGMENT_SHADER = 
"#version 330 core\n"
"in vec3 fragColor;\n"
"out vec4 outColor;\n"
"void main() {\n"
"    outColor = vec4(fragColor, 1.0);\n"
"}\n";

GLuint shaderProgram;
GLuint vao, vbo_pos, vbo_color;
int numPoints = 10000;

GLuint compileShader(const char* source, GLenum type) {
    GLuint shader = glCreateShader(type);
    glShaderSource(shader, 1, &source, NULL);
    glCompileShader(shader);
    
    GLint success;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
    if (!success) {
        char log[512];
        glGetShaderInfoLog(shader, 512, NULL, log);
        std::cerr << "Shader compilation error: " << log << std::endl;
    }
    return shader;
}

void createShaderProgram() {
    GLuint vertShader = compileShader(VERTEX_SHADER, GL_VERTEX_SHADER);
    GLuint fragShader = compileShader(FRAGMENT_SHADER, GL_FRAGMENT_SHADER);
    
    shaderProgram = glCreateProgram();
    glAttachShader(shaderProgram, vertShader);
    glAttachShader(shaderProgram, fragShader);
    glLinkProgram(shaderProgram);
    
    GLint success;
    glGetProgramiv(shaderProgram, GL_LINK_STATUS, &success);
    if (!success) {
        char log[512];
        glGetProgramInfoLog(shaderProgram, 512, NULL, log);
        std::cerr << "Shader linking error: " << log << std::endl;
    }
    
    glDeleteShader(vertShader);
    glDeleteShader(fragShader);
}

void createPointCloud() {
    std::vector<float> vertices;
    std::vector<float> colors;
    
    // Create a grid of points
    for (int i = 0; i < 100; i++) {
        for (int j = 0; j < 100; j++) {
            // X: -0.8 to 0.8, Y: -0.8 to 0.8, Z: 0
            float x = -0.8f + (i / 100.0f) * 1.6f;
            float y = -0.8f + (j / 100.0f) * 1.6f;
            float z = 0.0f;
            
            vertices.push_back(x);
            vertices.push_back(y);
            vertices.push_back(z);
            
            // Color based on position (rainbow)
            float r = (float)i / 100.0f;
            float g = (float)j / 100.0f;
            float b = 1.0f - (r + g) / 2.0f;
            
            colors.push_back(r);
            colors.push_back(g);
            colors.push_back(b);
        }
    }
    
    glGenVertexArrays(1, &vao);
    glBindVertexArray(vao);
    
    glGenBuffers(1, &vbo_pos);
    glBindBuffer(GL_ARRAY_BUFFER, vbo_pos);
    glBufferData(GL_ARRAY_BUFFER, vertices.size() * sizeof(float), vertices.data(), GL_STATIC_DRAW);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, 0);
    glEnableVertexAttribArray(0);
    
    glGenBuffers(1, &vbo_color);
    glBindBuffer(GL_ARRAY_BUFFER, vbo_color);
    glBufferData(GL_ARRAY_BUFFER, colors.size() * sizeof(float), colors.data(), GL_STATIC_DRAW);
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, 0);
    glEnableVertexAttribArray(1);
    
    glBindVertexArray(0);
    
    std::cout << "Created " << numPoints << " points" << std::endl;
}

void display() {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    
    glUseProgram(shaderProgram);
    glBindVertexArray(vao);
    
    glEnable(GL_PROGRAM_POINT_SIZE);
    glPointSize(5.0f);
    
    glDrawArrays(GL_POINTS, 0, numPoints);
    
    glBindVertexArray(0);
    glUseProgram(0);
    
    glutSwapBuffers();
    
    static int frameCount = 0;
    if (frameCount++ % 60 == 0) {
        GLenum err = glGetError();
        if (err != GL_NO_ERROR) {
            std::cout << "OpenGL Error: " << err << std::endl;
        }
    }
}

void init() {
    std::cout << "\n=== OpenGL Point Cloud Test ===" << std::endl;
    std::cout << "GL Vendor: " << glGetString(GL_VENDOR) << std::endl;
    std::cout << "GL Renderer: " << glGetString(GL_RENDERER) << std::endl;
    std::cout << "GL Version: " << glGetString(GL_VERSION) << std::endl;
    std::cout << "GLSL Version: " << glGetString(GL_SHADING_LANGUAGE_VERSION) << std::endl;
    
    glClearColor(0.1f, 0.1f, 0.15f, 1.0f);
    glEnable(GL_DEPTH_TEST);
    
    createShaderProgram();
    createPointCloud();
    
    std::cout << "\nYou should see a 100x100 grid of colored points" << std::endl;
    std::cout << "Points should have rainbow colors across the grid" << std::endl;
    std::cout << "If you see nothing, OpenGL point rendering may not work via X11" << std::endl;
}

int main(int argc, char** argv) {
    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGBA | GLUT_DEPTH);
    glutInitWindowSize(800, 600);
    glutCreateWindow("OpenGL Point Test");
    
    GLenum err = glewInit();
    if (err != GLEW_OK) {
        std::cerr << "GLEW Init Error: " << glewGetErrorString(err) << std::endl;
        return 1;
    }
    
    init();
    
    glutDisplayFunc(display);
    glutIdleFunc(display);
    
    glutMainLoop();
    
    return 0;
}