#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <mpi.h>

using namespace std;

void readArrayFromFile(string filename, int* array) {
    ifstream fin(filename);
    int n;
    fin >> n;

    for (int i = 0; i < n; ++i) {
        fin >> array[i];
    }
}

void createArrayFile(string filename, int n) {
    ofstream fout;
    fout.open(filename);
    fout << n << endl;

    for (int i = 0; i < n; i++) {
        fout << rand() % 100000 + 1 << " ";
    }
}

void writeArrayInFile(string filename, int* array, int n) {
    ofstream fout;
    fout.open(filename);
    fout << n << endl;

    for (int i = 0; i < n; i++ ) {
        fout << array[i] << " ";
    }
}

int compare(const void* x1, const void* x2) {
    return ( *(int*)x1 - *(int*)x2 );
}

int getPivot(int* array, int size) {
    if (size == 0) {
        return 0;
    } else if (size < 3) {
        return array[0];
    }
    return (array[0] + array[size - 1] + array[(size - 1) / 2]) / 3;
}

int getPairRank(int rank, int size) {
    int bitToCheck = 1 << ((int) log2(size) - 1);
    return rank ^ bitToCheck;
}

void swapPairs(int pivot, int*& array, int& size, int pairRank, int rank, MPI_Comm MPI_Local_Comm) {
    bool overPivotFlag = rank < pairRank;
    int lowerIter = 0;
    int overIter = size - 1;
    int temp;

    while (lowerIter <= overIter) {
        if (array[lowerIter] >= pivot) {
            temp = array[lowerIter];
            array[lowerIter] = array[overIter];
            array[overIter] = temp;
            overIter--;
        } else {
            lowerIter++;
        }
    }

    int sendSize = overPivotFlag ? size - lowerIter : lowerIter;
    int sendIndex = overPivotFlag ? lowerIter : 0;
    int recvSize = 0;

    MPI_Sendrecv(&sendSize, 1, MPI_INT, pairRank, 100, &recvSize, 1, MPI_INT, pairRank, 100, MPI_Local_Comm, MPI_STATUS_IGNORE);

    int* recvSubArray = new int[recvSize];

    MPI_Sendrecv(&(array[sendIndex]), sendSize, MPI_INT, pairRank, 100, recvSubArray, recvSize, MPI_INT, pairRank, 100, MPI_Local_Comm, MPI_STATUS_IGNORE);
    size -= sendSize;

    int* finalArray = new int[size + recvSize];
    int j = 0;
    for (int i = 0; i < size + recvSize; i++) {
        if (i < size) {
            finalArray[i] = overPivotFlag ? array[i] : array[lowerIter + i];
        } else {
            finalArray[i] = recvSubArray[j];
            j++;
        }
    }

    array = finalArray;
    size += recvSize;
}

int main(int argc, char** argv) {
    string inputFile = argv[1];
    string outputFile = argv[2];
    int n;

    if (argc > 3) {
        n = atoi(argv[3]);
        createArrayFile(inputFile, n);
    }

    MPI_Init(&argc, &argv);
    int* array = new int[0];
    int size = 0;
    double time = 0.0; // счетчик времени

    int rank, worldSize;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &worldSize);

    if (rank == 0) {
        ifstream fin(inputFile);
        fin >> size;
        array = new int[size];
        readArrayFromFile(inputFile, array);
        time = MPI_Wtime();
    }
    MPI_Bcast(&size, 1, MPI_INT, 0, MPI_COMM_WORLD);

    auto* sizes = new int[worldSize];
    auto* displs = new int[worldSize];
    int remainder = size % worldSize;
    int sum = 0;

    for (int i = 0; i < worldSize; i++) {
        sizes[i] = size / worldSize;
        if (remainder > 0) {
            sizes[i]++;
            remainder--;
        }
        displs[i] = sum;
        sum += sizes[i];
    }

//    if (size == 1) {
//        qsort(array, size, sizeof(int), compare);
//    }

    int cyclesCount = (int) log2(worldSize);
    int subArraySize = sizes[rank];
    int* subArray = new int[subArraySize];
    int* subArrayDispls = new int[worldSize];
    int* subArraySizes = new int[worldSize];

    MPI_Scatterv(array, sizes, displs, MPI_INT, subArray, subArraySize, MPI_INT, 0, MPI_COMM_WORLD);

    MPI_Comm MPI_Local_Comm = MPI_COMM_WORLD;
    int localRank = rank;
    int localSize = worldSize;
    int localPivot = 0;

    for(int iteration = 0; iteration < cyclesCount; iteration++) {
        if (localRank == 0) {
            localPivot = getPivot(subArray, subArraySize);
        }
        MPI_Bcast(&localPivot, 1, MPI_INT, 0, MPI_Local_Comm);

        int pairRank = getPairRank(localRank, localSize);

        swapPairs(localPivot, subArray, subArraySize, pairRank, localRank, MPI_Local_Comm);

        int color = localRank < pairRank ? 0 : 1;
        MPI_Comm_split(MPI_Local_Comm, color , localRank, &MPI_Local_Comm);
        MPI_Comm_size(MPI_Local_Comm, &localSize);
        MPI_Comm_rank(MPI_Local_Comm, &localRank);
    }

    qsort(subArray, subArraySize, sizeof(int), compare);
    MPI_Allgather(&subArraySize, 1, MPI_INT, subArraySizes, 1, MPI_INT, MPI_COMM_WORLD);
    subArrayDispls[0] = 0;
    for (int i = 1; i < worldSize; i++) {
        subArrayDispls[i] = subArrayDispls[i - 1] + subArraySizes[i - 1];
    }
    cout << rank << " size: " << subArraySizes[rank] << " disp: " << subArrayDispls[rank] << endl;
    MPI_Gatherv(subArray, subArraySize, MPI_INT, array, subArraySizes, subArrayDispls, MPI_INT, 0, MPI_COMM_WORLD);

    cout << rank << " gath" << endl;
    if (rank == 0) {
        time = MPI_Wtime() - time;
        cout << time << endl;
        writeArrayInFile(outputFile, array, size);
    }

    MPI_Finalize();
    return 0;
}