#include <stdio.h>
#include <stdlib.h>
#include <limits.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>

// Настройки эксперимента
#define USE_SHARED_MEMORY true
#define BLOCKS_COUNT 10000 // 1, 10, 100, 1000, 10000
// Используется, если USE_SHARED_MEMORY == false
#define THREADS_COUNT 100 // 1, 9, 100, 1024
// Используется, если USE_SHARED_MEMORY == true
#define BLOCK_SIZE 1024 // 1, 3, 10, 32
#define MATRIX_SIZE 100 // 500, 1000, 1500

#define OBSTACLE_PROB 10 // Процент препятствий на полигоне
#define START_X 2 // Начальная точка
#define START_Y 2
#define FINISH_X (MATRIX_SIZE - 3) // Конечная точка
#define FINISH_Y (MATRIX_SIZE - 3)

#define INF UINT_MAX // Используем беззнаковый максимум

// Вспомогательные функции
void generate_polygon(int* P, int n) {
    srand(42);
    for (int i = 0; i < n*n; i++) {
        P[i] = (rand() % 100 < OBSTACLE_PROB) ? -1 : 0;
    }
    P[START_X + START_Y * MATRIX_SIZE] = 0;   // Гарантируем, что старт свободен
    P[FINISH_X + FINISH_Y * MATRIX_SIZE] = 0; // Гарантируем, что финиш свободен
}

void print_distance_map(int* P, unsigned int* dist, int n) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            int idx = i * n + j;
            if (P[idx] == -1) 
                printf("████"); // Препятствие
            else if (i == START_X && j == START_Y) 
                printf("S   "); // Старт
            else if (i == FINISH_X && j == FINISH_Y) 
                printf("F   "); // Финиш
            else if (dist[idx] == INF) 
                printf(".   "); // Недостижимая область
            else 
                printf("%-4u", dist[idx]); // Вывод расстояния
        }
        printf("\n");
    }
}

#if USE_SHARED_MEMORY
// Ядро, использующее разделяемую память
__global__ void wave_step_shared(int* P, unsigned int* dist, int n, bool* global_changed) {
    int local_row = threadIdx.x;
    int local_col = threadIdx.y;

    __shared__ bool local_changed;
    __shared__ int shared_p[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ int shared_dist[BLOCK_SIZE][BLOCK_SIZE];

    bool changed_at_least_once = false;

    int block_index = blockIdx.x;
    const int matrix_blocks_size = (MATRIX_SIZE + BLOCK_SIZE - 1) / BLOCK_SIZE;

    while (block_index < matrix_blocks_size * matrix_blocks_size) {
        int block_row = block_index / matrix_blocks_size;
        int block_col = block_index % matrix_blocks_size;
        int i = local_row + block_row * BLOCK_SIZE;
        int j = local_col + block_col * BLOCK_SIZE;

        // Заполняем shared_p и shared_dist в разделяемой памяти
        shared_p[local_row][local_col] = 0;
        shared_dist[local_row][local_col] = INF;
        if (i < MATRIX_SIZE && j < MATRIX_SIZE && P[i * MATRIX_SIZE + j] != -1) {
            shared_dist[local_row][local_col] = dist[i * MATRIX_SIZE + j];
        } else {
            shared_p[local_row][local_col] = -1;
        }
    
        __syncthreads(); // Ждём, пока все потоки скопируют данные

        // Тот же волновой алгоритм, но в масштабах одного блока
        do {
            // Обновляем флаг перед началом новой итерации
            if (threadIdx.x == 0 && threadIdx.y == 0) {
                local_changed = false;
            }

            // Избавляемся от race conditions 
            // (если нулевой поток отстанет, то он может невовремя сбросить флаг)
            __syncthreads();

            if (shared_p[local_row][local_col] != -1) {
                unsigned int current_dist = shared_dist[local_row][local_col];
                unsigned int min_dist = current_dist;

                // Берём соседей из shared memory (shared_dist), 
                // если они там есть, иначе берём из global memory (dist)

                // Сверху
                if (local_row > 0 && shared_dist[local_row - 1][local_col] != INF) 
                    min_dist = min(min_dist, shared_dist[local_row - 1][local_col] + 1);
                else if (i > 0 && dist[(i-1)*n + j] != INF) 
                    min_dist = min(min_dist, dist[(i-1)*n + j] + 1);
            
                // Снизу
                if (local_row < BLOCK_SIZE - 1 && shared_dist[local_row + 1][local_col] != INF) 
                    min_dist = min(min_dist, shared_dist[local_row + 1][local_col] + 1);
                else if (i < n-1 && dist[(i+1)*n + j] != INF) 
                    min_dist = min(min_dist, dist[(i+1)*n + j] + 1);
                
                // Слева
                if (local_col > 0 && shared_dist[local_row][local_col - 1] != INF) 
                    min_dist = min(min_dist, shared_dist[local_row][local_col - 1] + 1);
                else if (j > 0 && dist[i*n + (j-1)] != INF) 
                    min_dist = min(min_dist, dist[i*n + (j-1)] + 1);
                
                // Справа
                if (local_col < BLOCK_SIZE - 1 && shared_dist[local_row][local_col + 1] != INF) 
                    min_dist = min(min_dist, shared_dist[local_row][local_col + 1] + 1);
                else if (j < n-1 && dist[i*n + (j+1)] != INF) 
                    min_dist = min(min_dist, dist[i*n + (j+1)] + 1);

                // Обновляем дистанцию в разделяемой памяти
                if (min_dist < current_dist) {
                    shared_dist[local_row][local_col] = min_dist;
                    local_changed = true;
                    changed_at_least_once = true;
                }
            }

            // Ждём, пока все потоки попробуют обновить ячейку
            // и флаг будет в актуальном состоянии
            __syncthreads();
        } while (local_changed);

        // Копируем результат из разделяемой памяти (shared_dist) в глобальную (dist)
        if (i < MATRIX_SIZE && j < MATRIX_SIZE && shared_p[local_row][local_col] != -1) {
            dist[i * MATRIX_SIZE + j] = shared_dist[local_row][local_col];
        }
        
        // Каждый блок потоков обрабатывает каждый BLOCKS_COUNT блок полигона
        block_index += BLOCKS_COUNT;
    }

    // Если блок потоков обновил хотя бы одно значение, то, возможно, 
    // нужны ещё глобальные итерации
    if (changed_at_least_once) {
        *global_changed = true;
    }
}
#else
// Ядро, не использующее разделяемую память
__global__ void wave_step(int* P, unsigned int* dist, int n, bool* changed) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    // printf("Hello from CUDA kernel! I'm thread #%d\n", tid);

    while (tid < n * n) {
        int i = tid / n;
        int j = tid % n;
        // printf("TID = %d (real %d); i = %d; j = %d\n", tid, threadIdx.x + blockIdx.x * blockDim.x, i, j);

        if (i >= n || j >= n) return;
        if (P[tid] != -1) {
            unsigned int current_dist = dist[tid];
            unsigned int min_dist = current_dist;

            // Проверка соседей с защитой от переполнения
            if (i > 0 && dist[(i-1)*n + j] != INF) 
                min_dist = min(min_dist, dist[(i-1)*n + j] + 1);
            
            if (i < n-1 && dist[(i+1)*n + j] != INF) 
                min_dist = min(min_dist, dist[(i+1)*n + j] + 1);
            
            if (j > 0 && dist[i*n + (j-1)] != INF) 
                min_dist = min(min_dist, dist[i*n + (j-1)] + 1);
            
            if (j < n-1 && dist[i*n + (j+1)] != INF) 
                min_dist = min(min_dist, dist[i*n + (j+1)] + 1);

            if (min_dist < current_dist) {
                dist[tid] = min_dist;
                *changed = true;
            }
        }

        // Каждый поток обрабатывает каждую blockDim.x * gridDim.x клетку
        // printf("Increment will be: %d\n", blockDim.x * gridDim.x);
        tid += blockDim.x * gridDim.x;
    }
}
#endif

int main() {
    const int n = MATRIX_SIZE;

    // Инициализация полигона
    int* P = (int*)malloc(n * n * sizeof(int));
    generate_polygon(P, n);

    // Выделение памяти на GPU
    int* d_P;
    unsigned int* d_dist;
    bool* d_changed;
    cudaMalloc(&d_P, n*n*sizeof(int));
    cudaMalloc(&d_dist, n*n*sizeof(unsigned int));
    cudaMalloc(&d_changed, sizeof(bool));

    // Инициализация расстояний
    unsigned int* dist = (unsigned int*)malloc(n*n*sizeof(unsigned int));
    for (int i = 0; i < n*n; i++) dist[i] = INF;
    dist[START_X + START_Y * MATRIX_SIZE] = 0; // Стартовая точка

    // Копирование данных на GPU
    cudaMemcpy(d_P, P, n*n*sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_dist, dist, n*n*sizeof(unsigned int), cudaMemcpyHostToDevice);

    // Замер времени
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);

    // Основной цикл волны
    int iterations = 0;
    bool changed;
    do {
        // printf("Wave step #%d\n", iterations);
        changed = false;
        cudaMemcpy(d_changed, &changed, sizeof(bool), cudaMemcpyHostToDevice);

        #if USE_SHARED_MEMORY
            dim3 block(BLOCK_SIZE, BLOCK_SIZE);
            wave_step_shared<<<BLOCKS_COUNT, block>>>(d_P, d_dist, n, d_changed);
        #else
            wave_step<<<BLOCKS_COUNT, THREADS_COUNT>>>(d_P, d_dist, n, d_changed);
        #endif

        cudaDeviceSynchronize(); // Синхронизация после ядра
        cudaMemcpy(&changed, d_changed, sizeof(bool), cudaMemcpyDeviceToHost);
        iterations++;
    } while (changed && iterations < 2*n); // Защита от бесконечного цикла

    // Финализация времени
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    float milliseconds = 0;
    cudaEventElapsedTime(&milliseconds, start, stop);

    // Проверка результата
    cudaMemcpy(dist, d_dist, n*n*sizeof(unsigned int), cudaMemcpyDeviceToHost);
    if (dist[FINISH_X + FINISH_Y * MATRIX_SIZE] == INF) {
        printf("Path not found!\n");
    } else {
        printf("Success! Path length: %u\n", dist[FINISH_X + FINISH_Y * MATRIX_SIZE]);
    }

    // Вывод результатов 
    printf("Time: %.2f ms\n", milliseconds);
    #if USE_SHARED_MEMORY
        printf("Results of the algorithm using shared memory.\n");
        printf("Matrix: %dx%d | BlocksXThreadXThread: %dx%dx%d | Obstacles: %d%%\n\n", 
           n, n, BLOCKS_COUNT, BLOCK_SIZE, BLOCK_SIZE, OBSTACLE_PROB);
    #else
        printf("Matrix: %dx%d | BlocksXThreads: %dx%d | Obstacles: %d%%\n\n", 
           n, n, BLOCKS_COUNT, THREADS_COUNT, OBSTACLE_PROB);
    #endif

    if (MATRIX_SIZE <= 100)
        print_distance_map(P, dist, MATRIX_SIZE);

    // Освобождение памяти
    free(P);
    free(dist);
    cudaFree(d_P);
    cudaFree(d_dist);
    cudaFree(d_changed);
    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    return 0;
}