/*
 * Kruskal's MPI
 * Copyright (C) 2015 George Piskas
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Contact: geopiskas@gmail.com
 */

#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <inttypes.h>
#include <time.h>
#include "edge.h"
#include "disjointset.h"

FILE *input = NULL;

// MPI related variables.
MPI_Datatype mpiEdge; // MPI datatype of edge.
int mpiNp; // Total # of processors.
int mpiRank; // Rank of a processor.

uint64_t nVerts; // Total # of vertices.
uint64_t nEdges; // Total # of edges.

edge *edges; // Contains the edges of a processor.
edge *edges_sort_buffer; // Used for parallel quicksorting of 'edges' for performance.
uint64_t edgeCount;

edge *mst; // Contains the MST edges of a processor. Also used as send/recv MPI buffer.
uint64_t mstEdgeCount;
double mstLength;

// Timing variables.
double commTime;
double procTime;
double parseTime;


// Checks if (l  <=  x  <  r)
/*inline int in(uint32_t x, uint32_t l, uint32_t r) {
    return ((x >= l) && (x < r));
}*/

// Finalizes MPI and frees buffers.
inline void finalize() {
    MPI_Type_free(&mpiEdge);
    MPI_Finalize();
}

// Force kill due to initialization error.
inline void die(char *msg) {
    printf("%d: %s\n",mpiRank, msg);
    if (input != NULL) fclose(input);
    MPI_Type_free(&mpiEdge);
    MPI_Finalize();
    exit(0);
}

// Initializes MPI and mpiEdge datatype.
inline void initMPI(int argc, char **argv) {
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
    MPI_Comm_size(MPI_COMM_WORLD, &mpiNp);
    MPI_Type_contiguous(3, MPI_UNSIGNED, &mpiEdge);
    MPI_Type_commit(&mpiEdge);
}

// Constructs a random graph.
inline void initInput() {

    if ((mpiNp & (mpiNp - 1)) != 0)
        die("* Please make sure that the number of processors is a power of 2\n");
    nEdges = nVerts * (nVerts - 1) / 2;

    uint64_t ePerProc = nEdges/mpiNp; //Edges per processor
    uint64_t firstEdge = mpiRank * ePerProc;
    uint64_t lastEdge = (mpiRank + 1) * ePerProc;
    if (mpiRank == mpiNp - 1) {
        ePerProc += nEdges%mpiNp;
        lastEdge += nEdges%mpiNp;
    }

    // Mem allocation.
    edges = malloc(ePerProc * sizeof(edge));

    // Allocating x2 because of merging. The first part of mst buffer contains the
    // MST edges and the second part is used as a receive MPI buffer.
    mst = malloc(2 * (nVerts - 1) * sizeof(edge));

    edgeCount = 0;
    uint64_t count = 0;
    uint64_t count2 = 0;
    uint32_t i, j, minj, maxj;
    edge e;
    for (i = nVerts-1; i > 0; i--){
        count2 += nVerts - i;
        if (count2 < firstEdge){
            count = count2;
            continue;
        }

        maxj = nVerts;
        if (count2 > lastEdge)
            maxj -= count2 - lastEdge;

        minj = i + 1;
        if (count < firstEdge)
            minj += firstEdge - count;

        for (j = minj; j <= maxj; j++) {
            e.u = i-1;
            e.v = j-1;
            e.weight = drand48();
            edges[edgeCount++] = e;
        }
        count = count2;
        if (count >= lastEdge)
            break;
    }

    parseTime = MPI_Wtime() - parseTime;
    MPI_Barrier(MPI_COMM_WORLD);
}

// Calculates the MST.
inline void calculateMst() {
    // Quicksort.
    qsort(edges, edgeCount, sizeof(edge), compareEdges);
    free(dsSet);
    dsMakeSet(nVerts);
    // Looping trhough all edges in increasing weight order.
    mstEdgeCount = 0;
    mstLength = 0;
    uint32_t i = 0;
    for(; i < edgeCount; i++) {
        // Find parent sets of the nodes.
        dsNode *vParent = dsFind(&dsSet[edges[i].v]);
        dsNode *uParent = dsFind(&dsSet[edges[i].u]);
        // If they are from two different sets, merge them.
        if(vParent != uParent) {
            mst[mstEdgeCount++] = edges[i];
            mstLength += edges[i].weight;
            dsUnion(vParent, uParent);
        }
    }
}

inline void processTree(int n) {

    parseTime = MPI_Wtime();
    nVerts = n;
    initInput();

    // Processing.
    double tmpTime;
    procTime = MPI_Wtime();
    // In case of 1 processor, act as serial application.
    if (mpiNp == 1) {
        calculateMst();
    } else {
        int processors = mpiNp;
        int pow2 = 1; // Used to find out where to send/recv.
        MPI_Status mpiStatus;
        int recvEdgeCount;
        while(processors > 1) {
            // Calculate the local MST using 'edges' buffer.
            calculateMst();

            // Communication part.
            if((mpiRank/pow2)%2 != 0) {
                // Send from the first half of 'mst'.
                MPI_Send(mst, mstEdgeCount, mpiEdge, (mpiRank-pow2), 0, MPI_COMM_WORLD);
                break; // Processor did his job and can now exit.
            } else {
                // Receive into the second half of 'mst'.
                tmpTime = MPI_Wtime();
                MPI_Recv(mst+mstEdgeCount, nVerts-1, mpiEdge, (mpiRank+pow2), 0, MPI_COMM_WORLD, &mpiStatus);
                MPI_Get_count(&mpiStatus, mpiEdge, &recvEdgeCount);
                commTime += MPI_Wtime() - tmpTime; // Communication time is equal to the waiting-to-recv time.
            }
            // Pointer swap between 'edges' and 'mst'.
            // 'edges' will be reused to calculate a new local MST.
            // 'mst' will contain that new local MST.
            edge *tmp = edges;
            edges = mst;
            mst = tmp;
            edgeCount = mstEdgeCount + recvEdgeCount; // New edge count after merging.

            processors /= 2;
            pow2 *= 2;
        }

        // Only the root will execute this last calculation.
        if(mpiRank == 0) calculateMst();
    }
    procTime = (MPI_Wtime() - procTime) - commTime;

    if(mpiRank == 0)
        printf("%d nodes, %" PRIu64 " vertices: %.3fs\n", mpiNp, nVerts, parseTime + commTime + procTime);

}

int main (int argc, char **argv) {
    initMPI(argc, argv);
        int numv = 256;
        if (argc > 1)
                numv = atoi(argv[1]);
        srand48(time(NULL));
        processTree(numv);
        MPI_Barrier(MPI_COMM_WORLD);
    finalize();
        return 0;
}