#include <SDL2/SDL_events.h>
#include <SDL2/SDL_pixels.h>
#include <SDL2/SDL_render.h>
#include <SDL2/SDL_scancode.h>
#include <SDL2/SDL_surface.h>
#include <SDL2/SDL_timer.h>
#include <SDL2/SDL_video.h>
#include <stdio.h>
#include <SDL2/SDL.h>
#include <complex.h>
#include <CL/cl.h>
#include <omp.h>

#include "mandelbrot.h"
#include "mouse.h"
#include "ui.h"

/*TODO
 *  1. Store iterations instead of colors ?
 *      1.1 memory usage vs performance loss ??? (cant draw in parallel anyway)
 *  2. Look into SDL2 rendering
 *  3. employ SSE2
 *      3.1 x86intrin.h
 *      3.2 inline assembly
*/

int 
mandelbrot(Complex c) {
    Complex z = {0.0, 0.0};
    int i;
    
    for (i = 0; i < MAX_ITERATIONS; i++) {
        double temp_real = z.real * z.real - z.imag * z.imag + c.real;
        double temp_imag = 2 * z.real * z.imag + c.imag;
        z.real = temp_real;
        z.imag = temp_imag;
        
        if (__builtin_hypot(z.real, z.imag) > 2)
            return i;
    }
    return MAX_ITERATIONS;
} 

typedef struct{
    short int r;
    short int g;
    short int b; 
} Color;

typedef struct {
    Color *pointer;
    int size;
}Array;


void
calculate_set(Array *arr, int height, int width, ViewInfo view)
{
    #pragma omp parallel for schedule(guided)
    for (int x = 0; x < width; x++) {
        for (int y = 0; y < height; y++) {
            double real = view.x_min + (x * (view.x_max - view.x_min)) / width;
            double imag = view.y_min + (y * (view.y_max - view.y_min)) / height;
            Complex c = {real, imag};
            
            int iterations = mandelbrot(c);

            if (iterations == MAX_ITERATIONS){
                arr->pointer[x * height + y].r = 0;
                arr->pointer[x * height + y].g = 0;
                arr->pointer[x * height + y].b = 0;
            } else {
                double t = (double)iterations / MAX_ITERATIONS;
                t = 0.5 + 0.5 * tan(log(t + 0.0001) * 3.0);

                arr->pointer[x * height + y].r = 255 * t * 0.5;
                arr->pointer[x * height + y].g = 255 * t * 0.4;
                arr->pointer[x * height + y].b = 255 * t;
            }
        }
    }
}


void 
update_texture(App *app, Array *arr) {
    void *pixels;
    int pitch;
    
    if (SDL_LockTexture(app->texture, NULL, &pixels, &pitch) != 0) {
        fprintf(stderr, "Failed to lock texture: %s\n", SDL_GetError());
        return;
    }
    
    uint32_t *dst = pixels;
    Color *src = arr->pointer;
    
    for (int y = 0; y < app->win_height; y++) {
        for (int x = 0; x < app->win_width; x++) {
            Color c = src[x * app->win_height + y];
            // Convert RGB to RGBA8888 format
            *dst++ = (c.b << 0) | (c.g << 8) | (c.r << 16) | (255 << 24);
        }
        dst += pitch / sizeof(uint32_t) - app->win_width;
    }
    
    SDL_UnlockTexture(app->texture);
}

void 
render_cl(Array *arr, App *app, ViewInfo view) {
    
    SDL_SetRenderDrawColor(app->renderer, 0, 0, 50, 255);
    SDL_RenderClear(app->renderer);
    
    if (arr->size < app->win_height * app->win_width * (int)sizeof(Color)){
        arr->pointer = realloc(arr->pointer, app->win_height * app->win_width * sizeof(Color));
        arr->size = app->win_height * app->win_width * sizeof(Color);
    }
    calculate_set(arr, app->win_height, app->win_width, view); 
    
    update_texture(app, arr);

    SDL_RenderCopy(app->renderer, app->texture, NULL, NULL);
}


void 
render(App *app, ViewInfo view, int start_x, int end_x, int start_y, int end_y) {
    
    SDL_SetRenderDrawColor(app->renderer, 0, 0, 50, 255);
    SDL_RenderClear(app->renderer);

    //#pragma omp parallel for schedule(guided)
    for (int x = start_x; x < end_x; x++) {
        for (int y = start_y; y < end_y; y++) {
            double real = view.x_min + (x * (view.x_max - view.x_min)) / app->win_width;
            double imag = view.y_min + (y * (view.y_max - view.y_min)) / app->win_height;
            Complex c = {real, imag};
            
            int iterations = mandelbrot(c);

            if (iterations == MAX_ITERATIONS){
                SDL_SetRenderDrawColor(app->renderer, 0, 0, 0, 255);
            } else {
                double t = (double)iterations / MAX_ITERATIONS;
                t = 0.5 + 0.5 * cos(log(t + 0.0001) * 3.0);
                
                int r = (int)(255 * t);  
                int g = (int)(255 * t * 0.4);  
                int b = (int)(255 * t * 0.2);        
                
                SDL_SetRenderDrawColor(app->renderer, r, g, b, 255);
            }
            SDL_RenderDrawPoint(app->renderer, x, y);
        }
    }
}



int
main(void)
{
    if (SDL_Init(SDL_INIT_EVERYTHING) != 0){
        printf("error initializing SDL: %s\n", SDL_GetError());
        return 1;
    }
    
    App app = {
        SDL_CreateWindow("Mandelbrot set exploer", 
            0, 0, 
            SCREEN_WIDTH,  SCREEN_HEIGHT, 
            SDL_WINDOW_SHOWN | SDL_WINDOW_RESIZABLE | SDL_WINDOW_OPENGL), 
        SDL_CreateRenderer(app.window, 
                -1, 
                SDL_RENDERER_ACCELERATED),
        SDL_CreateTexture(app.renderer, 
                SDL_PIXELFORMAT_RGBA8888,
                SDL_TEXTUREACCESS_STREAMING,
                SCREEN_WIDTH, SCREEN_HEIGHT),
        SCREEN_HEIGHT,
        SCREEN_WIDTH
            };

    ViewInfo view = {
        .x_min = -2.0,
        .x_max = 1.0,
        .y_min = -1.5,
        .y_max = 1.5,
        .zoom = 1.0
    };
    
    Mouse mouse = {
        .dragging = 0,
        .start_x = 0,
        .start_y = 0,
        .original_view = view
    };

    Uint32 frame_start;

    UI ui;
    init_ui(&ui);
    
    Array screen = {
    malloc(app.win_height * app.win_width * sizeof(Color)),
    app.win_height * app.win_width * sizeof(Color)
    };


    int running = 1;
    while (running)
    {
        frame_start = SDL_GetTicks();
        SDL_Event event;
        while (SDL_PollEvent(&event))
        {
            switch (event.type)
            {
                case SDL_QUIT:
                    running = 0;
                    break;
                case SDL_WINDOWEVENT:
                    if(event.window.event == SDL_WINDOWEVENT_RESIZED){
                        SDL_GetWindowSize(app.window, &app.win_width, &app.win_height);
                    }
                    break;
                case SDL_KEYDOWN:
                        switch (event.key.keysym.scancode)
                        {
                            case SDL_SCANCODE_Q:
                            case SDL_SCANCODE_ESCAPE:
                                running = 0;
                                break;
                            case SDL_SCANCODE_MINUS:
                                view.x_min = view.x_min * 1.1;
                                view.x_max = view.x_max * 1.1;
                                view.y_min= view.y_min * 1.1;
                                view.y_max = view.y_max * 1.1;
                                view.zoom = view.zoom * 1.1;
                                break;
                            case SDL_SCANCODE_EQUALS:
                                view.x_min = view.x_min * 0.9;
                                view.x_max = view.x_max * 0.9;
                                view.y_min= view.y_min * 0.9;
                                view.y_max = view.y_max * 0.9;
                                view.zoom = view.zoom * 0.9;
                                break;
                            case SDL_SCANCODE_R:
                                view.x_min = -2.0;
                                view.x_max = 1.0;
                                view.y_min = -1.5;
                                view.y_max = 1.5;
                                view.zoom = 1.0;
                                break;
                        }
                default:
                    handle_mouse(event, &mouse, &view);
                    break;
            }
        }
        
        //drawing happens here
        SDL_SetRenderDrawColor(app.renderer, 0xff, 0xff, 0xff, 0xff);
        SDL_RenderClear(app.renderer);
        render_cl(&screen, &app, view);

       // render(&app, 
       //         view, 
       //         0, app.win_width, 
       //         0,app.win_height);

        int msec = SDL_GetTicks() - frame_start;
        float fps = 0;
        if(msec > 0)
            fps = 1000.0 / (double) msec;

        render_ui(&ui, app.renderer, view, fps);

        SDL_RenderPresent(app.renderer);
    }
        //SDL_FreeSurface(app.surfaceA);
        free(screen.pointer);

        SDL_DestroyRenderer(app.renderer);
        SDL_DestroyWindow(app.window);
        SDL_Quit();
}
