#include "tricorn_V2.h"
#include <immintrin.h>
#include "color_table.h"

static unsigned converge(float c_rea, float c_imag, unsigned n);


/**
 * @brief Generates a Tricorn fractal image using highly optimized SSE/SIMD instructions, however the algoriithm is not fully optimized.
 * @param start The starting complex coordinate (bottom-left or top-left depending on height sign).
 * @param width The width of the target image in pixels.
 * @param height The height of the target image. A negative height reverses the Y-axis calculation.
 * @param res The step size (resolution) per pixel in the complex plane.
 * @param n The maximum number of mathematical iterations per pixel.
 * @param color If true, generates a 24-bit RGB image. If false, writes a grayscale image.
 * @param img Pointer to the pre-allocated memory buffer where the bitmap data will be stored.
 */
void tricorn_V2(float complex start, size_t width, ssize_t height, float res, unsigned n, _Bool color, unsigned char* img){
    int bytesPerColor = color ? 3 : 1;

    // Absolute height for iteration purposes
    ssize_t actual_height = height < 0 ?  (ssize_t)(-height) : height;
    size_t simd_width = (width / 4) * 4;

    // Depending on the sign of height we generate the image top-down or bottom up
    // As we iterate the image buffer from the top, if height is negative we would have to start at the end of the imaginary part and go in the oposite direction when generating
    ssize_t imag_direction = height < 0 ? -1 : 1;
    float c_start_real = creal(start);
    float c_start_imag = height > 0 ? cimag(start) : cimag(start) + (actual_height - 1) * res;


    // Prepared vectors for later
    __m128 four = _mm_set1_ps(4.0f);
    __m128 minus_two = _mm_set1_ps(-2.0f);

    for(ssize_t y = 0; y < actual_height;y++) {
        // Calculate the imaginary part once per row and create a vector from it
        float c_imag = c_start_imag + (y * imag_direction) * res;
        __m128 c_imag_vector = _mm_set1_ps(c_imag);

        size_t x = 0;
        for(; x < simd_width;x += 4) { //SIMD Loop
            // Calculate real part for the first pixel and then create vector for the group of for pixels
            float c_real = c_start_real + (x*res);
            __m128 c_real_vector = _mm_setr_ps(c_real,c_real+1.0f*res,c_real+2.0f*res,c_real+3.0f*res);


            // Initialize start values for the calculation
            __m128 real = _mm_setzero_ps();
            __m128 imag = _mm_setzero_ps();
            __m128i iterations = _mm_setzero_si128();
            __m128i active_mask = _mm_set1_epi32(-1); // Active mask will track which pixels have digerved - 0 for diverged, 1 - for contained, at start 0xFFFF

            for(unsigned i = 0; i < n;i++) {
                // Calculate squared values of real and imaginary
                __m128 real2 = _mm_mul_ps(real,real);
                __m128 imag2 = _mm_mul_ps(imag,imag);
                

                // Detect which have diverged
                __m128 unbounded = _mm_cmple_ps(_mm_add_ps(real2,imag2),four);
                __m128i unbounded_i = _mm_castps_si128(unbounded); // Cast to __m128i

                // Update the active mask
                active_mask = _mm_and_si128(active_mask,unbounded_i);

                // If all have diverged break
                if(_mm_test_all_zeros(active_mask,active_mask)) {
                    break;
                }

                // Add the iterations for each pixel only if it hasn't diverged
                iterations = _mm_sub_epi32(iterations,active_mask);

                imag = _mm_add_ps(c_imag_vector,_mm_mul_ps(minus_two, _mm_mul_ps(real,imag)));
                real = _mm_add_ps(_mm_sub_ps(real2,imag2),c_real_vector);
            }
            
            // Calculate the value for the pixel
            // value = 255 - 4 * (i % 64)
            __m128i mod64 = _mm_and_si128(iterations, _mm_set1_epi32(63)); // mod is realised with a bitmask
            __m128i times4 = _mm_slli_epi32(mod64, 2); // then multiplication by 4 is shift by 2
            __m128i values = _mm_sub_epi32(_mm_set1_epi32(255), times4); // subtraction from 255

            // pixels which have not escaped need to be colored black so value = 0
            __m128i max_iter = _mm_set1_epi32(n);
            __m128i is_max = _mm_cmpeq_epi32(iterations, max_iter);
            values = _mm_andnot_si128(is_max, values);

            // Convert vector to an array of ints
            unsigned int value_array[4];
            _mm_storeu_si128((__m128i*)value_array,values);

            // Write the 4 values in the image buffer
            for(unsigned i = 0 ; i < 4;i++) {
                unsigned int value = value_array[i];
                size_t position = (y*width + (x+i)) * bytesPerColor;

                // If grayscale just take value
                if(!color) {
                    img[position] = value;
                }
                else {
                    // If colored use the lookup table to get the RGB values
                    img[position] = color_lookup_table[value][0];
                    img[position + 1] = color_lookup_table[value][1];
                    img[position + 2] = color_lookup_table[value][2];
                }
            }
        }

        // Follow up loop for the remaining pixels
        // Calculates them scalarly
        for(;x < width; x++) {
            float c_real = c_start_real + (x*res);
            unsigned int value = converge(c_real,c_imag,n);

            size_t position = (y*width + x) * bytesPerColor;

            if(!color) {
                img[position] = value;
            }
            else {
                img[position] = ((value*3) % 255);
                img[position + 1] = ((value*2) % 255);
                img[position + 2] = ((value*5) % 255);
            }
        }

    }
}

static unsigned converge(float c_real, float c_imag, unsigned n) {
    float real = 0.0f;
    float imag = 0.0f;

    for (unsigned i = 0; i < n; i++) {
        float real2 = real * real;
        float imag2 = imag * imag;

        if (real2 + imag2 > 4.0f) {
            return 255 - 4 * (i % 64);
        }

        imag = -2.0f * real * imag + c_imag;
        real = real2 - imag2 + c_real;
    }
    return 0; 
}