#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <mpi.h>
#include <math.h>

#define MAX_TASKS 5
#define ITERATIONS_COUNT 100
#define L 2
#define REQUEST_TASK 228
#define STOP_CODE 229
#define TO_RECEIVER 1
#define SEND_WORK_TAG 2

int *tasks;
int size;
int rank;
int offset;
pthread_mutex_t mutex;

int do_job(int length) {
    int res = 0;
    while (1) {
//        if(rank == 0){
//            printf("ZZZZZ %d %d\n", offset, length);
//        }

        pthread_mutex_lock(&mutex);
        if (offset >= length) {
            pthread_mutex_unlock(&mutex);
            break;
        }
        int current_offset = offset++;
        pthread_mutex_unlock(&mutex);

        int weight = tasks[current_offset];
//        if(rank == 0){
//            printf("W %d\n", weight);
//        }
        for (int j = 0; j < weight; j++) {
            res += (int)sqrt(j);
        }
    }
    return res;
}

void fill_tasks(int iter_count) {
    for (int i = 0; i < size * MAX_TASKS; i++) {
        tasks[i] = abs(50 - i % 100) * abs(rank - iter_count % size) * L;
    }
}

int do_request_tasks() {
    for (int i = 0; i < size; i++) {
        if (i == rank) continue;
        int req_code = REQUEST_TASK;
        int help_length;
        //printf("VLASENKO DURAK ..vsdf.sd.fsd.f\n");
        MPI_Send(&req_code, 1, MPI_INT, i, TO_RECEIVER, MPI_COMM_WORLD);
        MPI_Recv(&help_length, 1, MPI_INT, i, SEND_WORK_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        printf("!! resv length tasks VLASENKO TOP\n");
        if (help_length > 0) {
            MPI_Recv(tasks, help_length, MPI_INT, i, SEND_WORK_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            printf("!!! recv tasks\n");
            pthread_mutex_lock(&mutex);
            offset = 0;
            pthread_mutex_unlock(&mutex);
        }

        return do_job(help_length);
    }
}



int do_tasks() {
    int local_res = 0;
    local_res = do_job(size * MAX_TASKS);

    local_res += do_request_tasks();
    return local_res;
}



void* worker_func(void* arg) {
    for (int i = 0; i < ITERATIONS_COUNT; i++) {
        pthread_mutex_lock(&mutex);
        offset = 0;
        fill_tasks(i);
        pthread_mutex_unlock(&mutex);
        int res = do_tasks();
        if (rank == 0) {
            printf("%d\n", res);
        }
    }
    int stop_code = STOP_CODE;
    MPI_Send(&stop_code, 1, MPI_INT, 0, TO_RECEIVER, MPI_COMM_WORLD);
    return NULL;
}

void* receiver_func(void* arg) {
    while (1) {
        int req_code_buf;
        MPI_Status status_worker_requester;
        MPI_Recv(&req_code_buf, 1, MPI_INT, MPI_ANY_SOURCE, TO_RECEIVER, MPI_COMM_WORLD, &status_worker_requester);

        if (req_code_buf == STOP_CODE) break;

        size_t length = size * MAX_TASKS;
        int new_offset = offset + (int)((int)(length - offset) * 0.3);
        int tasks_length = new_offset - offset;
        printf("ASDDSA 1 %d\n", tasks_length);
        MPI_Send(&tasks_length, 1, MPI_INT, status_worker_requester.MPI_SOURCE, SEND_WORK_TAG, MPI_COMM_WORLD);
        printf("ASDDSA 2\n");
        if (tasks_length > 0 ) {
            pthread_mutex_lock(&mutex);
            int old_offset = offset;
            offset = new_offset;
            pthread_mutex_unlock(&mutex);
            printf("ASDDSA 3\n");
            MPI_Send(&tasks[old_offset], tasks_length, MPI_INT, status_worker_requester.MPI_SOURCE, SEND_WORK_TAG, MPI_COMM_WORLD);
            printf("ASDDSA 4\n");
        }
    }
    return NULL;
}

int main(int argc, char* argv[]) {
    MPI_Init_thread(NULL, NULL, MPI_THREAD_MULTIPLE, NULL);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    tasks = malloc(MAX_TASKS * sizeof(int));
    pthread_t worker, receiver;
    pthread_mutex_init(&mutex, NULL);
    pthread_create(&worker, NULL, worker_func, NULL);
    pthread_create(&receiver, NULL, receiver_func, NULL);

    pthread_join(worker, NULL);
    pthread_join(receiver, NULL);

    free(tasks);
    MPI_Finalize();
    return 0;
}