//
// Computes the minimum spanning tree (MST) of a weighted graph
// The graph is provided as an adjacency matrix A
//
// Warning: Return values of calls are not checked for error to keep
// the code simple.
//
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <sys/time.h>
#include <time.h> // mod

enum { FALSE, TRUE };

#define MAX_N   	65536

//#define INFINITY   	1.0e300

// mod
struct timespec start, stop;
double total_time;

int numthreads = 4;
pthread_t p_threads[4]; // Threads
pthread_attr_t attr; // Thread attributes
// end

// MST
typedef struct _mst {
    unsigned int n;		// number of nodess in the tree
    unsigned int * node; 	// edge (i,node[i]) is in the MST
    double * weight; 		// weight of edge (i,node[i])
    double mst_weight; 		// weight of MST (sum of weights of all edges
} MST_t;
// Adjacency matrix structure
typedef struct _adj_matrix {
    double ** weight;
    int n;
} ADJ_MATRIX_t;

// Variable Declarations
ADJ_MATRIX_t * A;		// Adjacency Matrix of graph
MST_t * mst;			// MST of graph
MST_t * my_mst;			// mod -> my parellel MST of graph

// -----------------------------------------------------------------------------
// MST Routines
//
// Create and initialize MST
//
MST_t * init_mst ( ADJ_MATRIX_t * A ) {
    unsigned i;
    MST_t * mst = calloc(1, sizeof(MST_t));
    mst->n = A->n;
    mst->node = calloc(mst->n, sizeof(unsigned int));
    mst->weight = calloc(mst->n, sizeof(double));
    for (i = 0; i < mst->n; i++) {
        mst->node[i] = i;
        mst->weight[i] = 0.0;
    }
    mst->mst_weight = 0.0;
    return mst;
}
// -----------------------------------------------------------------------------
// Adjacency Matrix Routines
//
// Create and initialize adjacency matrix A with edge weights of the graph
//   A->n denotes size of matrix
//   A->weight[i][j] denotes weight edge from nodes i and j
//
ADJ_MATRIX_t * init_adj_matrix ( unsigned int num_nodes, unsigned int seed ) {
    unsigned int i, j;
    ADJ_MATRIX_t * A = calloc(1, sizeof(ADJ_MATRIX_t));
    A->n = num_nodes;
    // Allocate storage for matrix elements
    A->weight = calloc(A->n, sizeof(double *));
    for (i = 0; i < A->n; i++) {
        A->weight[i] = calloc(A->n, sizeof(double));
    }
    // Initialize matrix elements
    srand(seed);
    double randwt;
    for (i = 0; i < A->n; i++) {
        for (j = i; j < A->n; j++) {
            if (i == j) {
                A->weight[i][j] = 0.0;
            } else {
                // Insert random edge weights
                A->weight[i][j] = (double)(rand_r(&seed))/(double)(RAND_MAX);
                A->weight[j][i] = A->weight[i][j];
            }
        }
    }
    return A;
}
// Compute minimum spanning tree of graph A
MST_t *  minimum_spanning_tree(ADJ_MATRIX_t * A, unsigned int root) {
    unsigned int i, j, u, nodes;
    double minwt;

    // Initialize MST
    MST_t * mst = init_mst(A);

    // Initialize weight vector d with weight of edges connecting node i to root
    double * d = calloc(A->n, sizeof(double));
    for (i = 0; i < A->n; i++) {
        d[i] = A->weight[root][i];	// d[i]= min wt. edge from node i to MST
        mst->node[i] = root;		// closest MST node to node i
    }
    // Initialize flag that indicates whether a node has been inserted in the tree
    short int * inserted_in_mst = calloc(A->n, sizeof(short int));
    for (i = 0; i < A->n; i++)
        inserted_in_mst[i] = FALSE;

    // Prim's MST algorithm

    // Insert root into MST
    inserted_in_mst[root] = TRUE;
    mst->node[root] = root;
    mst->weight[root] = 0.0;
    mst->mst_weight = 0.0;

    nodes = 1;
    while (nodes < A->n) {
        // Find node u not belonging to MST with minimum weight edge in d
        minwt = INFINITY;
        for (i = 0; i < A->n; i++) {
            if ((!inserted_in_mst[i]) && (d[i] < minwt)) {
                u = i;
                minwt = d[i];
            }
        }
        // Add u to MST
        inserted_in_mst[u] = TRUE;
        mst->weight[u] = d[u];
        mst->mst_weight += d[u];
        nodes ++;

        // Update d vector
        for (i = 0; i < A->n; i++) {
            if ((!inserted_in_mst[i]) && (d[i] > A->weight[u][i])) {
                d[i] = A->weight[u][i];
                mst->node[i] = u;	// set u as closest MST node of i
            }
        }
    }
    return mst;
}


// mod
typedef struct _passme {
    ADJ_MATRIX_t * A;
    short int * inserted_in_mst;
    double * d;
    MST_t * mst;
    unsigned int u;
    int threadid;
} passme;


void *pprim(void *s) {
    passme* data = (passme*)s;
    unsigned int i;
    MST_t * mst = data->mst;
    ADJ_MATRIX_t * A = data->A;
    unsigned int u = data->u;
    double * d = data->d;
    short int * inserted_in_mst = data->inserted_in_mst;
    int threadid = data->threadid;

    int blocksize = A->n / numthreads;
    int index = threadid * blocksize;

    //printf("my threadid is %d\n", threadid);
    //for(;;);

    for(i = index; i < (index + blocksize); i++){
        if ((!inserted_in_mst[i]) && (d[i] > A->weight[u][i])) {
            d[i] = A->weight[u][i];
            mst->node[i] = u;	// set u as closest MST node of i
        }
    }
}


MST_t * parallel_mst(ADJ_MATRIX_t * A, unsigned int root) {
    unsigned int i, j, u, nodes;
    double minwt;

    // Initialize MST
    MST_t * mst = init_mst(A);

    // Initialize weight vector d with weight of edges connecting node i to root
    double * d = calloc(A->n, sizeof(double));
    for(i = 0; i < A->n; i++) {
        d[i] = A->weight[root][i];	// d[i]= min wt. edge from node i to MST
        mst->node[i] = root;		// closest MST node to node i
    }

    // Initialize flag that indicates whether a node has been inserted in the tree
    short int * inserted_in_mst = calloc(A->n, sizeof(short int));
    for(i = 0; i < A->n; i++)
        inserted_in_mst[i] = FALSE;

    // Prim's MST algorithm

    // Insert root into MST
    inserted_in_mst[root] = TRUE;
    mst->node[root] = root;
    mst->weight[root] = 0.0;
    mst->mst_weight = 0.0;

    nodes = 1;
    while (nodes < A->n) {
        // Find node u not belonging to MST with minimum weight edge in d
        minwt = INFINITY;
        for(i = 0; i < A->n; i++) {
            if((!inserted_in_mst[i]) && (d[i] < minwt)) {
                u = i;
                minwt = d[i];
            }
        }
        // Add u to MST
        inserted_in_mst[u] = TRUE;
        mst->weight[u] = d[u];
        mst->mst_weight += d[u];
        nodes ++;

        passme data[4];

        for(i = 0; i < numthreads; i++) {
            data[i].A = A;
            data[i].inserted_in_mst = inserted_in_mst;
            data[i].d = d;
            data[i].mst = mst;
            data[i].u = u;
            data[i].threadid = i;
            //printf("creating thread with value %d\n", data.threadid);
            pthread_create(&p_threads[i], &attr, pprim, (void *)&data[i]);
        }

        for(i = 0; i < numthreads; i++) {
            pthread_join(p_threads[i], NULL);
        }
    }
    return mst;
}
// end


// Main program
// - Initialize adjacency matrix A of a graph
// - Computes minimum spanning tree with a given starting node
int main(int argc, char *argv[]) {

    unsigned int num_nodes, root, seed;

    if (argc != 4) {
        printf("Use: <executable_name> <num_nodes> <root_node> <seed>\n");
        exit(0);
    }
    if ((num_nodes = atoi(argv[argc-3])) > MAX_N) {
        printf("Maximum number of nodes allowed: %d\n", MAX_N);
        exit(0);
    };
    if ((root = atoi(argv[argc-2])) > num_nodes-1) {
        printf("Root node (%d) should be in the range [0,...,%d]\n", root, num_nodes-1);
        exit(0);
    };
    seed = atoi(argv[argc-1]);

    // Initialize adjacency matrix of graph
    A = init_adj_matrix(num_nodes, seed);

    clock_gettime(CLOCK_REALTIME, &start);

    // Compute single source shortest path with root node r
    mst = minimum_spanning_tree(A, root);

    // Compute time taken
    clock_gettime(CLOCK_REALTIME, &stop);
    total_time = (stop.tv_sec-start.tv_sec)	+ 0.000000001 * (stop.tv_nsec-start.tv_nsec);

    // Print results
    unsigned int i, j;
    printf("         MST: Nodes=%d, MST Weight = %e, time (sec) = %8.4f\n", A->n, mst->mst_weight, total_time);

    // mod
    clock_gettime(CLOCK_REALTIME, &start);
    my_mst = parallel_mst(A, root);
    clock_gettime(CLOCK_REALTIME, &stop);
    total_time = (stop.tv_sec-start.tv_sec)	+ 0.000000001 * (stop.tv_nsec-start.tv_nsec);
    printf("parallel MST: Nodes=%d, MST Weight = %e, time (sec) = %8.4f\n", A->n, my_mst->mst_weight, total_time);
    // end
}