#include <mpi.h>
    #include <stdio.h>
    #include <windows.h>
    #include <intrin.h>
    #include <processthreadsapi.h>
    #include <barrier>
    #include <memory>

    #define MSG_ROWS 8
    #define MSG_COLS 158
    #define DISP_ROWS 8
    #define DISP_COLS 34
    #define NUM_PROCS DISP_ROWS * DISP_COLS
    #define MILLISECONDS_PER_TICK 5000

    // 8 rows x 158 columns
    // each hex value is an 8-pixel column

    const int message[MSG_COLS] = {
        0x00, 0x00, 0xe7, 0xe7, 0xe7, 0x00, 0x00, 0xff, 0xfb, 0x1b, 0x0b, 0x6b, 0x03, 0x03, 0xff, 0x00,
        0x00, 0xde, 0xcc, 0xc0, 0xe1, 0xff, 0x00, 0x00, 0xde, 0xcc, 0xc0, 0xe1, 0xff, 0xfc, 0x30, 0x03,
        0x83, 0xe0, 0xfc, 0xff, 0xff, 0xff, 0xfe, 0xfe, 0x00, 0x00, 0xf0, 0xc3, 0x0f, 0x00, 0x00, 0xff,
        0x87, 0x03, 0x0b, 0x4b, 0x43, 0xe7, 0xff, 0xf3, 0x83, 0x1f, 0x07, 0xe3, 0x03, 0x1f, 0x83, 0xf3,
        0xff, 0xff, 0xff, 0xfe, 0xf8, 0x01, 0x07, 0xf1, 0xf8, 0xfe, 0x87, 0x03, 0x0b, 0x4b, 0x43, 0xe7,
        0xff, 0xfb, 0x1b, 0x0b, 0x6b, 0x03, 0x03, 0xff, 0x03, 0x03, 0xe3, 0xf3, 0xff, 0xff, 0xff, 0xff,
        0x03, 0x03, 0xe3, 0xf3, 0x3f, 0x1f, 0x8f, 0xc7, 0xe3, 0xf3, 0xff, 0x00, 0x00, 0xf7, 0xe7, 0x07,
        0x0f, 0xff, 0x87, 0x03, 0x7b, 0x33, 0x03, 0x87, 0xff, 0x01, 0x01, 0xfd, 0xf9, 0x01, 0x01, 0xfd,
        0x01, 0x01, 0xff, 0x87, 0x03, 0x0b, 0x4b, 0x43, 0xe7, 0xff, 0x00, 0x00, 0xff, 0xfb, 0x1b, 0x0b,
        0x6b, 0x03, 0x03, 0xff, 0x00, 0x00, 0x33, 0x7b, 0x33, 0x03, 0xcf, 0xff, 0x20, 0x20 };

    typedef struct ThreadArgs {
        int thread_rank;
        std::shared_ptr<std::barrier<std::_No_completion_function>> barrier;
    } MYDATA, *PMYDATA;

    DWORD WINAPI run_thread(LPVOID arg_ptr);

    int main(int argc, char** argv) {

        PMYDATA pDataArray[NUM_PROCS];
        DWORD   dwThreadIdArray[NUM_PROCS - 1];
        HANDLE  hThreadArray[NUM_PROCS - 1];

        // Set own affinity to 0
        auto newAffinity = new _GROUP_AFFINITY;
        newAffinity->Group = 0;
        newAffinity->Mask = static_cast<KAFFINITY>(1);
        newAffinity->Reserved[0] = 0;
        newAffinity->Reserved[1] = 0;
        newAffinity->Reserved[2] = 0;
        SetThreadGroupAffinity(GetCurrentThread(), newAffinity, NULL);

        // Create barrier
        auto barrier = std::make_shared<std::barrier<std::_No_completion_function>>(NUM_PROCS);

        // Populate own data
        pDataArray[0] = (PMYDATA)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(MYDATA));
        pDataArray[0]->barrier = std::shared_ptr(barrier);
        pDataArray[0]->thread_rank = 0;

        // Create MAX_THREADS - 1 worker threads.

        for (int i = 0; i < NUM_PROCS - 1; i++)
        {
            // Create the thread to begin execution on its own.
            pDataArray[i + 1] = (PMYDATA)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(MYDATA));
            pDataArray[i + 1]->barrier = std::shared_ptr(barrier);
            pDataArray[i + 1]->thread_rank = i + 1;
            hThreadArray[i] = CreateThread(
                NULL,                   // default security attributes
                0,                      // use default stack size
                run_thread,             // thread function name
                pDataArray[i + 1],      // argument to thread function
                0,                      // use default creation flags
                &dwThreadIdArray[i]);   // returns the thread identifier

            // Set affinity.
            auto newAffinity = new _GROUP_AFFINITY;
            newAffinity->Group = (i+1) / 64;
            newAffinity->Mask = static_cast<KAFFINITY>(1) << ((i+1) % 64);
            newAffinity->Reserved[0] = 0;
            newAffinity->Reserved[1] = 0;
            newAffinity->Reserved[2] = 0;
            SetThreadGroupAffinity(hThreadArray[i], newAffinity, NULL);
        }

        run_thread(pDataArray[0]);

        // Wait until all threads have terminated.

        WaitForMultipleObjects(NUM_PROCS - 1, hThreadArray, TRUE, INFINITE);

        // Close all thread handles.

        for (int i = 0; i < NUM_PROCS - 1; i++)
        {
            CloseHandle(hThreadArray[i]);
        }
    }

    DWORD WINAPI run_thread(LPVOID arg_ptr) {
        PMYDATA args = (PMYDATA)arg_ptr;
        int rank = args->thread_rank;
        args->barrier->arrive_and_wait();
        auto procnum = new PROCESSOR_NUMBER;
        GetCurrentProcessorNumberEx(procnum);

        int my_row = rank / DISP_COLS;
        int my_col = rank % DISP_COLS;

        // count cols from -disp_cols to msg_cols + disp_cols
        // add my_col. If count < 0 or count > msg_cols - 1, sleep/barrier
        // otherwise, look up in array and bit-shift to find out whether to spin
        for (int i = -DISP_COLS; i <= MSG_COLS + DISP_COLS; i++)
        {
            int cur_pixel_val;
            if (i + my_col < 0 || i + my_col > MSG_COLS - 1) {
                cur_pixel_val = 1;        }
            else {
                cur_pixel_val = message[i + my_col] & (1 << my_row);
            }
            if (cur_pixel_val) {
                args->barrier->arrive_and_wait();
            }
            else {
                LARGE_INTEGER frequency;
                LARGE_INTEGER ticks;
                QueryPerformanceFrequency(&frequency);
                QueryPerformanceCounter(&ticks);
                int64_t start_ticks = ticks.QuadPart;
                int64_t cur_ticks = start_ticks;
                while (cur_ticks - start_ticks < frequency.QuadPart * MILLISECONDS_PER_TICK / 1000.0) {
                    QueryPerformanceCounter(&ticks);
                    cur_ticks = ticks.QuadPart;
                    for (int j = 0; j < 10000; j++) {
                        __nop();
                    }
                }
                args->barrier->arrive_and_wait();
            }
        }

        return 0;
    }