#include <glad/glad.h>
#include <glfw3.h>
#include <math.h>
#include <vector>
#include <stack>
#include <iostream>
#include <fstream>
#include <chrono>
#include <random>
#include <sstream>
typedef float HeatmapType;
using namespace std;

vector<float> vertices;
// settings
const int IMAGE_HEIGHT = 200;
const int IMAGE_WIDTH = 200;
const int RED_ITERS = 200;
const int BLUE_ITERS = 800;
const int GREEN_ITERS = 200;
const long long int SAMPLE_COUNT = IMAGE_WIDTH * IMAGE_HEIGHT * 100;

class Complex
{
  public:
    Complex(double r = 0.0, double i = 0.0)
        : _r(r), _i(i)
    {
    }

    Complex(const Complex &) = default;

    double r() const { return _r; }
    double i() const { return _i; }

    Complex operator*(const Complex &other)
    {
        // (a + bi) (c + di)
        return Complex(_r * other._r - _i * other._i, _r * other._i + _i * other._r);
    }

    Complex operator+(const Complex &other)
    {
        return Complex(_r + other._r, _i + other._i);
    }

    double sqmagnitude() const
    {
        return _r * _r + _i * _i;
    }

  private:
    double _r, _i;
};

//
// Utility
//
void AllocHeatmap(HeatmapType **&o_heatmap, int width, int height)
{
    o_heatmap = new HeatmapType *[height];
    for (int i = 0; i < height; ++i)
    {
        o_heatmap[i] = new HeatmapType[width];
        for (int j = 0; j < width; ++j)
        {
            o_heatmap[i][j] = 0;
        }
    }
}

void FreeHeatmap(HeatmapType **&o_heatmap, int height)
{
    for (int i = 0; i < height; ++i)
    {
        delete[] o_heatmap[i];
        o_heatmap[i] = nullptr;
    }
    delete o_heatmap;
    o_heatmap = nullptr;
}

vector<Complex> buddhabrotPoints(const Complex &c, int nIterations)
{
    int n = 0;
    Complex z;

    vector<Complex> toReturn;
    toReturn.reserve(nIterations);

    while (n < nIterations && z.sqmagnitude() <= 2.0)
    {
        z = z * z + c;
        ++n;

        toReturn.push_back(z);
    }

    // If point remains bounded through nIterations iterations, the point
    //  is bounded, therefore in the Mandelbrot set, therefore of no interest to us
    if (n == nIterations)
    {
        return vector<Complex>();
    }
    else
    {
        return toReturn;
    }
}

int rowFromReal(double real, double minR, double maxR, int imageHeight)
{
    // [minR, maxR]
    // [0, maxR - minR] // subtract minR from n
    // [0, imageHeight] // multiply by (imageHeight / (maxR - minR))
    return (int)((real - minR) * (imageHeight / (maxR - minR)));
}

int colFromImaginary(double imag, double minI, double maxI, int imageWidth)
{
    return (int)((imag - minI) * (imageWidth / (maxI - minI)));
}

void GenerateHeatmap(HeatmapType **o_heatmap, int imageWidth, int imageHeight,
                     const Complex &minimum, const Complex &maximum, int nIterations, long long nSamples,
                     HeatmapType &o_maxHeatmapValue, string consoleMessagePrefix)
{
    mt19937 rng;
    uniform_real_distribution<double> realDistribution(minimum.r(), maximum.r());
    uniform_real_distribution<double> imagDistribution(minimum.i(), maximum.i());

    rng.seed(chrono::high_resolution_clock::now().time_since_epoch().count());
    // Collect nSamples samples... (sample is just a random number c)
    for (long long sampleIdx = 0; sampleIdx < nSamples; ++sampleIdx)
    {
        //  Each sample, get the list of points as the function
        //    escapes to infinity (if it does at all)

        Complex sample(realDistribution(rng), imagDistribution(rng));
        vector<Complex> pointsList = buddhabrotPoints(sample, nIterations);

        for (Complex &point : pointsList)
        {
            if (point.r() <= maximum.r() && point.r() >= minimum.r() && point.i() <= maximum.i() && point.i() >= minimum.i())
            {
                int row = rowFromReal(point.r(), minimum.r(), maximum.r(), imageHeight);
                int col = colFromImaginary(point.i(), minimum.i(), maximum.i(), imageWidth);
                ++o_heatmap[row][col];

                if (o_heatmap[row][col] > o_maxHeatmapValue)
                {
                    o_maxHeatmapValue = o_heatmap[row][col];
                }
            }
        }
    }
}

float colorFromHeatmap(HeatmapType inputValue, HeatmapType maxHeatmapValue, float maxColor)
{
    double scale = (1.0f * (maxColor)) / (1.0f * maxHeatmapValue);
    return inputValue * scale;
}

void framebuffer_size_callback(GLFWwindow *window, int width, int height);
void mouse_button_callback(GLFWwindow *window, int button, int action, int mods);
void processInput(GLFWwindow *window);

const char *vertexShaderSource = "#version 330 core\n"
                                 "layout (location = 0) in vec3 aPos;\n"
                                 "layout (location = 1) in vec3 aColor;\n"
                                 "out vec3 vertexColor;\n"
                                 "void main()\n"
                                 "{\n"
                                 "   gl_Position = vec4(aPos,1.0);\n"
                                 "vertexColor = aColor;\n"
                                 "}\0";
const char *fragmentShaderSource = "#version 330 core\n"
                                   "out vec4 FragColor;\n"
                                   "in vec3 vertexColor;\n"
                                   "void main()\n"
                                   "{\n"
                                   "   FragColor =  vec4(vertexColor, 1.0);\n"
                                   "}\n\0";

int points;

float mapX(float x)
{
    return x * 2.0 / IMAGE_WIDTH - 1.0;
}

float mapY(float y)
{
    return y * 2.0 / IMAGE_HEIGHT - 1.0;
}

int main()
{
    const Complex MINIMUM(-2.0, -2.0);
    const Complex MAXIMUM(2.0, 2.0);
    vector<float> vertices;
    vertices.reserve(IMAGE_HEIGHT * IMAGE_WIDTH);

    // Allocate a heatmap of the size of our image
    HeatmapType maxHeatmapValue = 0;
    HeatmapType **red;
    HeatmapType **green;
    HeatmapType **blue;
    AllocHeatmap(red, IMAGE_WIDTH, IMAGE_HEIGHT);
    AllocHeatmap(green, IMAGE_WIDTH, IMAGE_HEIGHT);
    AllocHeatmap(blue, IMAGE_WIDTH, IMAGE_HEIGHT);

    GenerateHeatmap(red, IMAGE_WIDTH, IMAGE_HEIGHT, MINIMUM, MAXIMUM, RED_ITERS,
                    SAMPLE_COUNT, maxHeatmapValue, "Red Channel: ");
    GenerateHeatmap(blue, IMAGE_WIDTH, IMAGE_HEIGHT, MINIMUM, MAXIMUM, BLUE_ITERS,
                    SAMPLE_COUNT, maxHeatmapValue, "Blue Channel: ");

    GenerateHeatmap(green, IMAGE_WIDTH, IMAGE_HEIGHT, MINIMUM, MAXIMUM, GREEN_ITERS,
                    SAMPLE_COUNT, maxHeatmapValue, "green Channel: ");

    // Scale the heatmap down
    for (int row = 0; row < IMAGE_HEIGHT; ++row)
    {
        for (int col = 0; col < IMAGE_WIDTH; ++col)
        {

            red[row][col] = colorFromHeatmap(red[row][col], maxHeatmapValue, 1);
            blue[row][col] = colorFromHeatmap(blue[row][col], maxHeatmapValue, 1);
            green[row][col] = colorFromHeatmap(green[row][col], maxHeatmapValue, 1);

            //cout<<red[row][col]<<endl;
            vertices.push_back(-mapX(col));
            vertices.push_back(-mapY(row));
            vertices.push_back(0.0f);
            vertices.push_back(red[row][col]);
            vertices.push_back(green[row][col]);
            vertices.push_back(blue[row][col]);

            //  cout<<red[row][col]<<endl;
            //vertices.push_back(0.0f);
            //vertices.push_back(0.0f);
        }
    }

    // glfw: initialize and configure
    // ------------------------------
    glfwInit();
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

#ifdef __APPLE__
    glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE); // uncomment this statement to fix compilation on OS X
#endif

    // glfw window creation
    // --------------------
    GLFWwindow *window = glfwCreateWindow(IMAGE_WIDTH, IMAGE_HEIGHT, "Algorithmic Modeling Fractal Buddhabrot Set", NULL, NULL);
    if (window == NULL)
    {
        std::cout << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return -1;
    }
    glfwMakeContextCurrent(window);
    glfwSetFramebufferSizeCallback(window, framebuffer_size_callback);
    //    glfwSetMouseButtonCallback(window, mouse_button_callback);

    // glad: load all OpenGL function pointers
    // ---------------------------------------
    if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress))
    {
        std::cout << "Failed to initialize GLAD" << std::endl;
        return -1;
    }

    // build and compile our shader program
    // ------------------------------------
    // vertex shader
    int vertexShader = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertexShader, 1, &vertexShaderSource, NULL);
    glCompileShader(vertexShader);
    // check for shader compile errors
    int success;
    char infoLog[512];
    glGetShaderiv(vertexShader, GL_COMPILE_STATUS, &success);
    if (!success)
    {
        glGetShaderInfoLog(vertexShader, 512, NULL, infoLog);
        std::cout << "ERROR::SHADER::VERTEX::COMPILATION_FAILED\n"
                  << infoLog << std::endl;
    }
    // fragment shader
    int fragmentShader = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragmentShader, 1, &fragmentShaderSource, NULL);
    glCompileShader(fragmentShader);
    // check for shader compile errors
    glGetShaderiv(fragmentShader, GL_COMPILE_STATUS, &success);
    if (!success)
    {
        glGetShaderInfoLog(fragmentShader, 512, NULL, infoLog);
        std::cout << "ERROR::SHADER::FRAGMENT::COMPILATION_FAILED\n"
                  << infoLog << std::endl;
    }
    // link shaders
    int shaderProgram = glCreateProgram();
    glAttachShader(shaderProgram, vertexShader);
    glAttachShader(shaderProgram, fragmentShader);
    glLinkProgram(shaderProgram);
    // check for linking errors
    glGetProgramiv(shaderProgram, GL_LINK_STATUS, &success);
    if (!success)
    {
        glGetProgramInfoLog(shaderProgram, 512, NULL, infoLog);
        std::cout << "ERROR::SHADER::PROGRAM::LINKING_FAILED\n"
                  << infoLog << std::endl;
    }
    glDeleteShader(vertexShader);
    glDeleteShader(fragmentShader);

    // set up vertex data (and buffer(s)) and configure vertex attributes
    // ------------------------------------------------------------------

    points = vertices.size();

    // uncomment this call to draw in wireframe polygons.
    //glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);

    // render loop
    cout << "done" << endl;
    unsigned int VBO, VAO;

    glGenVertexArrays(1, &VAO);
    glGenBuffers(1, &VBO);
    // bind the Vertex Array Object first, then bind and set vertex buffer(s), and then configure vertex attributes(s).
    glBindVertexArray(VAO);

    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(float) * vertices.size(), &vertices[0], GL_STATIC_DRAW);

    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float), (void *)0);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float), (void *)(3 * sizeof(float)));
    glEnableVertexAttribArray(1);
    // note that this is allowed, the call to glVertexAttribPointer registered VBO as the vertex attribute's bound vertex buffer object so afterwards we can safely unbind
    glBindBuffer(GL_ARRAY_BUFFER, 0);

    // You can unbind the VAO afterwards so other VAO calls won't accidentally modify this VAO, but this rarely happens. Modifying other
    // VAOs requires a call to glBindVertexArray anyways so we generally don't unbind VAOs (nor VBOs) when it's not directly necessary.
    glBindVertexArray(0);
    // -----------
    while (!glfwWindowShouldClose(window))
    {

        processInput(window);

        //glClearColor(1,1, 1, 1.0f);
        //glClear(GL_COLOR_BUFFER_BIT);

        // draw our first triangle
        glUseProgram(shaderProgram);
        glBindVertexArray(VAO); // seeing as we only have a single VAO there's no need to bind it every time, but we'll do so to keep things a bit more organized
        glDrawArrays(GL_POINTS, 0, points);
        glfwSwapBuffers(window);
        glfwPollEvents();
        // ------------------------------------------------------------------
    }

    //glDeleteVertexArrays(1, &VAO);
    //glDeleteBuffers(1, &VBO);
    // optional: de-allocate all resources once they've outlived their purpose:
    // ------------------------------------------------------------------------

    // glfw: terminate, clearing all previously allocated GLFW resources.

    glfwTerminate();
    return 0;
}

// process all input: query GLFW whether relevant keys are pressed/released this frame and react accordingly
// ---------------------------------------------------------------------------------------------------------
void processInput(GLFWwindow *window)
{
    if (glfwGetKey(window, GLFW_KEY_ESCAPE) == GLFW_PRESS)
        glfwSetWindowShouldClose(window, true);
}

// glfw: whenever the window size changed (by OS or user resize) this callback function executes
// ---------------------------------------------------------------------------------------------

void framebuffer_size_callback(GLFWwindow *window, int width, int height)
{
    // make sure the viewport matches the new window dimensions; note that width and
    // height will be significantly larger than specified on retina displays.
    glViewport(0, 0, width, height);
}