#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <math.h>
#include <string.h>
#include <omp.h>
#include <gmp.h>
#include <ncurses.h>

#include "lodepng.h"

typedef struct
{
    //position
    mpf_t real;
    mpf_t imag;

    //zoom, smaller is more zoomed in
    mpf_t pixelSize;

    //max iters so we dont freeze on something in the set
    uint64_t maxIters;
} MandelbrotRenderer;

uint64_t getIter(uint64_t maxIters, const mpf_t creal, const mpf_t cimag)
{

    //this can be done faster, as this allocates data and then frees it, which is not necessary
    mpf_t real;
    mpf_t imag;
    mpf_init_set(real, creal);
    mpf_init_set(imag, cimag);

    //scratchpad
    mpf_t temp;
    mpf_init(temp);

    //deal with real^2
    mpf_t real2;
    mpf_init(real2);

    //deal with imag^2
    mpf_t imag2;
    mpf_init(imag2);

    uint64_t iter = 0;
    while(1)
    {
        //if we have seen too deep
        if(maxIters < iter) break;

        //square real and imag
        mpf_pow_ui(real2, real, 2);
        mpf_pow_ui(imag2, imag, 2);

        //calculate squared distance from 0,0
        mpf_add(temp, real2, imag2);

        //check if escaped the set
        if(mpf_cmp_ui(temp, 4) > 0) break;

        //update position
        //temp = real*real - imag*imag + creal;
        mpf_sub(temp, real2, imag2);
        mpf_add(temp, temp, creal);

        //imag = 2 * real * imag + cimag;
        mpf_mul(imag, real, imag);
        mpf_mul_ui(imag, imag, 2);
        mpf_add(imag, imag, cimag);

        //real = temp;
        mpf_set(real, temp);
        iter++;
    }

    //free stuff
    mpf_clear(imag2);
    mpf_clear(real2);
    mpf_clear(temp);
    mpf_clear(imag);
    mpf_clear(real);
    //printf("uh oh\n");

    return iter;
}

//take a set of inputs and print an image to the screen.
//alternatively, print an image to a buffer when it inevitably comes time to do that

//this function takes a zoom and image and point and calculates the number of iterations.
//this is only necessary to use stack space for the scratchpad variables
//and not the heap, but i really dont want to stack overflow.
//...gmp variables are already stored on the heap. ignore me

uint64_t getIterOnImage(MandelbrotRenderer* mr, int32_t x, int32_t z)
{
    //x and z are coords across the image we are currently rendering
    mpf_t tempReal;
    mpf_t tempImag;
    mpf_init(tempReal);
    mpf_init(tempImag);

    //multiply our x and z by pizelsize(zoom)
    mpf_mul_ui(tempReal, mr->pixelSize, abs(x));
    mpf_mul_ui(tempImag, mr->pixelSize, abs(z));
    if(x < 0) mpf_neg(tempReal, tempReal);
    if(z < 0) mpf_neg(tempImag, tempImag);

    //offset values to get to the actual real and imaginary spot
    mpf_add(tempReal, tempReal, mr->real);
    mpf_add(tempImag, tempImag, mr->imag);

    //only now we get iteration count and return
    uint32_t iterCount = getIter(mr->maxIters, tempReal, tempImag);

    mpf_clear(tempImag);
    mpf_clear(tempReal);
    return iterCount;
}

//trust that the user puts in the right numbers, have a buffer that can be written to
void renderImage(MandelbrotRenderer* mr, uint64_t* buffer, int32_t w, int32_t h, int useScreen)
{
    //zero out buffer. not necessary because we are already overwriting all
    //pixels
    memset((void*)buffer, 0x00, w * h * sizeof(uint64_t));

    //parallalize this
    #pragma omp parallel for
    for(int32_t z = 0; z < h; z++)
    {
        for(int32_t x = 0; x < w; x++)
        {
            uint32_t iterCount = getIterOnImage(mr, x-(w/2), z-(h/2));
            #pragma omp critical
            {
                if(useScreen)
                {
                    int color = (iterCount - (mr->maxIters%256) + COLOR_WHITE + 8) % 256;
                    attron(COLOR_PAIR(color));
                    mvaddch(z, 2*x, iterCount % 94 + 32);
                    mvaddch(z, 2*x + 1, iterCount % 94 + 32);
                    attroff(COLOR_PAIR(color));
                    refresh();
                }
            }
            buffer[w*z + x] = iterCount;
        }
    }
}

void loadRenderer(MandelbrotRenderer* mr, FILE* fd)
{
    //just do the thing
    char inputBuffer[64];

    memset(inputBuffer, 0, 64); //clear
    fgets(inputBuffer, 64, fd); //prec
    mpf_set_default_prec(atoi(inputBuffer));

    memset(inputBuffer, 0, 64);
    fgets(inputBuffer, 64, fd); //iterCount
    mr->maxIters = atoi(inputBuffer);

    mpf_init(mr->real);
    mpf_init(mr->imag);
    mpf_init(mr->pixelSize);

    mpf_inp_str(mr->real, fd, 10);
    mpf_inp_str(mr->imag, fd, 10);
    mpf_inp_str(mr->pixelSize, fd, 10);
}

void saveRenderer(MandelbrotRenderer* mr, FILE* fd)
{
    fprintf(fd, "%ld\n", mpf_get_default_prec());
    fprintf(fd, "%ld\n", mr->maxIters);
    mpf_out_str(fd, 10, 0, mr->real);
    fprintf(fd, "\n");
    mpf_out_str(fd, 10, 0, mr->imag);
    fprintf(fd, "\n");
    mpf_out_str(fd, 10, 0, mr->pixelSize);
    fprintf(fd, "\n");
}

void saveImage(uint64_t* imageBuffer, char* imageName, int w, int h)
{
    uint32_t* pngBuffer = (uint32_t*)malloc(w * h * sizeof(uint32_t));

    //its already in the right format, just modify and copy it 1:1
    for(uint64_t i = 0; i < (w*h); i++)
    {
        uint32_t color = (10000 * imageBuffer[i]) % (1 << 24);

        //the first byte here is the alpha, because memory stuff. little endian vs big endian
        //abgr
        pngBuffer[i] = 0xff000000 | color;
    }

    unsigned error = lodepng_encode32_file(imageName, (unsigned char*)pngBuffer, w, h);
    if(error) printf("error %u: %s\n", error, lodepng_error_text(error));
    free(pngBuffer);
}

void doNoGui(MandelbrotRenderer* mr, int w, int h, char* fileName)
{
    uint64_t* imageBuffer = (uint64_t*)malloc(w * h * sizeof(uint64_t));
    //if there is no gui, set what to do here
    int i = 0;
    char imageName[64];

    //for zoomOuts we will need to do # frames per halve
    //this is achieved by multiplying pixelSize by the nth root of 2
    mpf_t scalingFactor;
    mpf_init_set_d(scalingFactor, pow(2, 1.0/60));
    while(mpf_cmp_d(mr->pixelSize, 0.1) < 0)
    {
        renderImage(mr, imageBuffer, w, h, 0);
        mpf_mul(mr->pixelSize, mr->pixelSize, scalingFactor);

        memset(imageName, 0, 64);
        sprintf(imageName, "zoomOuts/%.5d.png", i);
        saveImage(imageBuffer, imageName, w, h);
        printf("did layer %.5d\n", i);
        i++;
    }
    
    mpf_clear(scalingFactor);
    free(imageBuffer);
}

void doGui(MandelbrotRenderer* mr, int w, int h, char* fileName)
{
    initscr();
    cbreak(); //disable line buffering
    noecho(); //dont echo chars back
    curs_set(0); //hide cursor
    start_color();
    use_default_colors();
    //init color map
    for (int i = 0; i < 256; i++) {
        init_pair(i + 1, i, 0);  // pair 1..8 = color 0..7 on black
    }

    uint64_t* imageBuffer = (uint64_t*)malloc(w * h * sizeof(uint64_t));

    //make first frame
    clear();
    renderImage(mr, imageBuffer, w, h, 1);
    mvprintw(8, 8, "Welcome to meoxi's mandelbrot set viewer");
    mvprintw(9, 8, "WASD to move");
    mvprintw(10, 8, "Z to zoom in, X to zoom out");
    mvprintw(11, 8, "This program eats all your cores. be warned!");
    refresh();


    int ch;
    char strBuffer[64];
    while((ch = getch()) != 'q')
    {
        switch(ch)
        {
            case 'w':
                mpf_sub(mr->imag, mr->imag, mr->pixelSize);
                break;
            case 'a':
                mpf_sub(mr->real, mr->real, mr->pixelSize);
                break;
            case 's':
                mpf_add(mr->imag, mr->imag, mr->pixelSize);
                break;
            case 'd':
                mpf_add(mr->real, mr->real, mr->pixelSize);
                break;
            case 'z':
                mpf_div_ui(mr->pixelSize, mr->pixelSize, 2);
                break;
            case 'x':
                mpf_mul_ui(mr->pixelSize, mr->pixelSize, 2);
                break;
            case 'c':
                mr->maxIters <<= 1;
                break;
            case 'v':
                mr->maxIters >>= 1;
                break;
            case 'f':
                //save
                FILE* fd = fopen(fileName, "w");
                saveRenderer(mr, fd);
                fclose(fd);
                break;

        }

        //put an image on the screen
        renderImage(mr, imageBuffer, w, h, 1);
        gmp_snprintf(strBuffer, 64, "%.63Ff", mr->real);
        mvprintw(0, 0, "real: %s", strBuffer);
        gmp_snprintf(strBuffer, 64, "%.63Ff", mr->imag);
        mvprintw(1, 0, "imag: %s", strBuffer);
        mvprintw(2, 0, "iter: %ld", mr->maxIters);
        refresh();
    }
    free(imageBuffer);
    endwin();
}

int main(int argc, char** argv)
{
    // ./mandelbrot <fileName> <w> <h> <gui>

    //build renderer from save file
    MandelbrotRenderer mr;

    //read save file and close
    FILE* saveFile = fopen(argv[1], "r");
    if(saveFile == NULL)
    {
        //erm.. akward
        mpf_init(mr.real);
        mpf_init(mr.imag);
        mpf_init(mr.pixelSize);

        mpf_set_default_prec(64);
        mpf_set_d(mr.real, -0.5);
        mpf_set_d(mr.imag, 0.0);
        mpf_set_d(mr.pixelSize, 0.1);
        mr.maxIters = 128;
    }
    else
    {
        loadRenderer(&mr, saveFile);
        fclose(saveFile);
    }


    int w = atoi(argv[2]);
    int h = atoi(argv[3]);
    int gui = atoi(argv[4]);

    //init ncurses
    if(gui)
    {
        doGui(&mr, w, h, argv[1]);
    }
    else
    {
        doNoGui(&mr, w, h, argv[1]);
    }

    mpf_clear(mr.pixelSize);
    mpf_clear(mr.imag);
    mpf_clear(mr.real);
}

//blah blah