#include <SDL2/SDL_events.h>
#include <SDL2/SDL_pixels.h>
#include <SDL2/SDL_render.h>
#include <SDL2/SDL_scancode.h>
#include <SDL2/SDL_stdinc.h>
#include <SDL2/SDL_surface.h>
#include <SDL2/SDL_timer.h>
#include <SDL2/SDL_video.h>
#include <immintrin.h>
#include <stdint.h>
#include <stdio.h>
#include <SDL2/SDL.h>
#include <complex.h>
#include <CL/cl.h>
#include <omp.h>
#include <x86intrin.h>
#include <argp.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
 *      2.1 texture tendering   DONE
 *      2.2 fix colors ??? 
 *  3. SSE2
 *      3.1 x86intrin.h DONE
*/

int
mandelbrot_avx(Complex c)
{
    __m256 cr = _mm256_set1_pd(c.real);
    __m256 ci = _mm256_set1_pd(c.imag);

    __m256 zr = cr;
    __m256 zi = ci;

    __m256 threshold = _mm256_set1_pd(4); 

    int k = 1;

    __m256 mk = _mm256_set1_pd(k);
    __m256 one = _mm256_set1_pd(1);
    while (++k < MAX_ITERATIONS){
        /* Compute z1 from z0 */
        __m256 zr2 = _mm256_mul_pd(zr, zr);
        __m256 zi2 = _mm256_mul_pd(zi, zi);
        __m256 zrzi = _mm256_mul_pd(zr, zi);
        /* zr1 = zr0 * zr0 - zi0 * zi0 + cr */
        /* zi1 = zr0 * zi0 + zr0 * zi0 + ci */
        zr = _mm256_add_pd(_mm256_sub_pd(zr2, zi2), cr);
        zi = _mm256_add_pd(_mm256_add_pd(zrzi, zrzi), ci);

        /* Increment k */
        zr2 = _mm256_mul_pd(zr, zr);
        zi2 = _mm256_mul_pd(zi, zi);
        __m256 mag2 = _mm256_add_pd(zr2, zi2);
        __m256 mask = _mm256_cmp_pd(mag2, threshold, _CMP_LT_OS);
        mk = _mm256_add_pd(_mm256_and_pd(mask, one), mk);

        /* Early bailout? */
        if (_mm256_testz_pd(mask, _mm256_set1_pd(-1)))
            break;
    }
    return k;
}


int 
mandelbrot(Complex c) {
    Complex z = {0.0, 0.0};
    int i;
    
    for (i = 0; i < MAX_ITERATIONS; i++) {
        long double temp_real = z.real * z.real - z.imag * z.imag + c.real;
        long 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{
    uint8_t r;
    uint8_t g;
    uint8_t b;
    uint8_t a;
} 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 y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            long double real = view.x_min + (x * (view.x_max - view.x_min)) / width;
            long double imag = view.y_min + (y * (view.y_max - view.y_min)) / height;
            Complex c = {real, imag};
            
            int iterations = mandelbrot_avx(c);

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

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

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;
    }
    
    memcpy(pixels, arr->pointer, app->win_width * app->win_height * sizeof(Color));
    
    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);
    }
    if (app->needs_recalc){
            calculate_set(arr, app->win_height, app->win_width, view); 
            app->needs_recalc = 0;
    }
    
    update_texture(app, arr);

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

struct arguments
{
    int image;
    char *output_file;
    int width;
    int height;
    int use_avx;
    int redcalc;

    double x_min;
    double x_max;
    double y_min;
    double y_max;

    double zoom;
    char *args[2];
};

static struct argp_option options[] =
{
    {"image", 'i', 0, 0, "Render to image"},
    {"output file", 'o', "FILE", 0, "Filename for image"},
    {"width", 'w', "WIDTH", 0, "Width of window/image"},
    {"height", 'h', "HEIGHT", 0, "Height of window/image"},
    {"dont_use_avx", 'x', 0, 0, "NOT use avx for calculations"},
    {"recalc", 'r', 0, 0, "Recalculate set even if not necessary"},

    {"x_min", 'a', "VALUE", 0, "X min"},
    {"x_max", 'b', "VALUE", 0, "X max"},
    {"y_min", 'c', "VALUE", 0, "Y min"},
    {"y_max", 'd', "VALUE", 0, "Y max"},

    {"zoom", 'z', "VALUE", 0, "Zoom"},
  {0}
};

static error_t
parse_opt (int key, char *arg, struct argp_state *state)
{
    struct arguments *arguments = state->input;
    switch (key)
    {
    case 'i':
        arguments->image = 1; 
        break;
    case 'o':
        arguments->output_file = arg; 
        break;
    case 'w':
        arguments->width = atoi(arg); 
        break;
    case 'h':
        arguments->height = atoi(arg);
        break;
    case 'x':
        arguments->use_avx = 0;
        break;
    case 'r':
        arguments->redcalc = 1;
        break;

    case 'a':
        arguments->x_min = atoi(arg);
        break;
    case 'b':
        arguments->redcalc = atoi(arg);
        break;
    case 'c':
        arguments->redcalc = atoi(arg);
        break;
    case 'd':
        arguments->redcalc = atoi(arg);
        break;
    case 'z':
        arguments->zoom= atof(arg);
        break;
    case ARGP_KEY_ARG:
        if (state->arg_num >= 2)
        /* Too many arguments. */
        argp_usage (state);
        arguments->args[state->arg_num] = arg;
      break;
    default:
        return ARGP_ERR_UNKNOWN;
    }
  return 0;
}

static struct argp argp = {options, parse_opt, NULL, NULL};


int
main(int argc, char **argv)
{
    struct arguments args = {
        0, 
        NULL, 
        640, 480, 
        1, 0,
        -2.0, 1.0, -1.5, 1.5, 1.0 
    };
    argp_parse(&argp, argc, argv, 0, 0, &args);

    if (SDL_Init(SDL_INIT_EVERYTHING) != 0){
        printf("error initializing SDL: %s\n", SDL_GetError());
        return 1;
    }
    
    ViewInfo view = {
        .x_min = args.x_min * args.zoom,
        .x_max = args.x_max * args.zoom,
        .y_min = args.y_min * args.zoom,
        .y_max = args.y_max * args.zoom,
        .zoom = args.zoom 
    };

    App app = {
        SDL_CreateWindow("Mandelbrot set exploer", 
            0, 0, 
            args.width,  args.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,
                args.width, args.height),
        args.height,
        args.height,
        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);
                        SDL_DestroyTexture(app.texture);

                        app.texture = SDL_CreateTexture(app.renderer, 
                        SDL_PIXELFORMAT_RGBA8888,
                        SDL_TEXTUREACCESS_STREAMING,
                        app.win_width, app.win_height);
                        app.needs_recalc = 1;
                    }
                    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;

                                app.needs_recalc = 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;

                                app.needs_recalc = 1;
                                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;


                                app.needs_recalc = 1;
                                break;
                            case SDL_SCANCODE_H:
                                app.draw_ui = 1 - app.draw_ui;
                                break;
                        }
                default:
                    //returns one if recalc needed
                    app.needs_recalc = app.needs_recalc == 0 ? handle_mouse(event, &mouse, &view) : 1;
                    app.needs_recalc = (app.needs_recalc || args.redcalc);
                    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;

        if(app.draw_ui)
            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();
}
