#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <immintrin.h>
#include <stdint.h>
#include <math.h>
#include <thread>
#include <vector>

#pragma comment(lib, "user32.lib")
#pragma comment(lib, "gdi32.lib")

#define WIDTH  1280
#define HEIGHT 800
#define MAX_ITER 1000

static uint32_t framebuffer[WIDTH * HEIGHT];

static double centerX = -0.5;
static double centerY = 0.0;
static double scale = 3.0 / WIDTH;

static int dragging = 0;
static int lastX, lastY;

static int numThreads = std::thread::hardware_concurrency();

static void render_tile(int yStart, int yEnd)
{
    __m256d two = _mm256_set1_pd(2.0);
    __m256d four = _mm256_set1_pd(4.0);

    for (int y = yStart; y < yEnd; y++)
    {
        double ci_scalar = centerY + (y - HEIGHT / 2.0) * scale;
        __m256d ci = _mm256_set1_pd(ci_scalar);

        for (int x = 0; x < WIDTH; x += 4)
        {
            __m256d cr = _mm256_set_pd(
                centerX + (x + 3 - WIDTH / 2.0) * scale,
                centerX + (x + 2 - WIDTH / 2.0) * scale,
                centerX + (x + 1 - WIDTH / 2.0) * scale,
                centerX + (x + 0 - WIDTH / 2.0) * scale
            );

            __m256d zr = _mm256_setzero_pd();
            __m256d zi = _mm256_setzero_pd();

            int iter[4] = { 0,0,0,0 };

            for (int i = 0; i < MAX_ITER; i++)
            {
                __m256d zr2 = _mm256_mul_pd(zr, zr);
                __m256d zi2 = _mm256_mul_pd(zi, zi);
                __m256d mag = _mm256_add_pd(zr2, zi2);

                __m256d mask = _mm256_cmp_pd(mag, four, _CMP_LT_OS);
                if (_mm256_movemask_pd(mask) == 0)
                    break;

                __m256d zrzi = _mm256_mul_pd(zr, zi);

                zi = _mm256_add_pd(_mm256_mul_pd(two, zrzi), ci);
                zr = _mm256_add_pd(_mm256_sub_pd(zr2, zi2), cr);

                for (int k = 0; k < 4; k++)
                    if (iter[k] < MAX_ITER)
                        iter[k]++;
            }

            for (int k = 0; k < 4; k++)
            {
                double smooth = iter[k];
                if (iter[k] < MAX_ITER)
                {
                    double zr_s = ((double*)&zr)[k];
                    double zi_s = ((double*)&zi)[k];
                    double mag = sqrt(zr_s * zr_s + zi_s * zi_s);
                    smooth = iter[k] + 1 - log2(log2(mag));
                }

                uint8_t c = (uint8_t)(255.0 * smooth / MAX_ITER);
                framebuffer[y * WIDTH + x + k] =
                    (c << 16) | (c << 8) | c;
            }
        }
    }
}

static void render_mandelbrot()
{
    std::vector<std::thread> threads;

    int tile = HEIGHT / numThreads;

    for (int i = 0; i < numThreads; i++)
    {
        int yStart = i * tile;
        int yEnd = (i == numThreads - 1) ? HEIGHT : yStart + tile;
        threads.emplace_back(render_tile, yStart, yEnd);
    }

    for (auto& t : threads)
        t.join();
}

LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    switch (msg)
    {
    case WM_MOUSEWHEEL:
    {
        int delta = GET_WHEEL_DELTA_WPARAM(wParam);
        scale *= (delta > 0) ? 0.8 : 1.25;
        render_mandelbrot();
        InvalidateRect(hwnd, NULL, FALSE);
        break;
    }

    case WM_LBUTTONDOWN:
        dragging = 1;
        lastX = LOWORD(lParam);
        lastY = HIWORD(lParam);
        break;

    case WM_MOUSEMOVE:
        if (dragging)
        {
            int x = LOWORD(lParam);
            int y = HIWORD(lParam);

            centerX -= (x - lastX) * scale;
            centerY -= (y - lastY) * scale;

            lastX = x;
            lastY = y;

            render_mandelbrot();
            InvalidateRect(hwnd, NULL, FALSE);
        }
        break;

    case WM_LBUTTONUP:
        dragging = 0;
        break;

    case WM_PAINT:
    {
        PAINTSTRUCT ps;
        HDC hdc = BeginPaint(hwnd, &ps);

        BITMAPINFO bmi = {};
        bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
        bmi.bmiHeader.biWidth = WIDTH;
        bmi.bmiHeader.biHeight = -HEIGHT;
        bmi.bmiHeader.biPlanes = 1;
        bmi.bmiHeader.biBitCount = 32;
        bmi.bmiHeader.biCompression = BI_RGB;

        StretchDIBits(hdc,
            0, 0, WIDTH, HEIGHT,
            0, 0, WIDTH, HEIGHT,
            framebuffer,
            &bmi,
            DIB_RGB_COLORS,
            SRCCOPY);

        EndPaint(hwnd, &ps);
        break;
    }

    case WM_DESTROY:
        PostQuitMessage(0);
        break;

    default:
        return DefWindowProc(hwnd, msg, wParam, lParam);
    }
    return 0;
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE, LPSTR, int nCmdShow)
{
    render_mandelbrot();

    WNDCLASS wc = {};
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = "MandelbrotSIMD";

    RegisterClass(&wc);

    HWND hwnd = CreateWindowEx(
        0,
        "MandelbrotSIMD",
        "Mandelbrot Explorer (SIMD + Threads)",
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT,
        WIDTH, HEIGHT,
        NULL, NULL, hInstance, NULL);

    ShowWindow(hwnd, nCmdShow);

    MSG msg;
    while (GetMessage(&msg, NULL, 0, 0))
    {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    return 0;
}