/* File: a2-1.c
   
   Sample sort parallel program
*/


/* MPI header */
#include <mpi.h>

#include <stdio.h>
#include <unistd.h>
#include <string.h>

//to get a set of randomly generated data
#include <sys/types.h>
#include <time.h>

//a2-1.h = header file containing the MPI type and proto-type declarations
#include "a2-1.h"

int my_rank, p;
int i;
Element *records;

//list_size must be divisible by p (total processors running the code)
int list_size=1024;
clock_t start,stop;

int main(int argc, char* argv[]) {
	
    LOCAL_LIST_T local_keys;
	//MPI_Status status;
   
    if(argc>0) list_size=atoi(argv[1]);
   
    //Generate Data
    init_data(&records,list_size);
    
	/**** MPI stuff ****/

	/* Start up MPI */ 
    	MPI_Init(&argc, &argv);
        
    	/* Find out process rank */  
    	MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    	/* Find out number of processes */ 
        MPI_Comm_size(MPI_COMM_WORLD, &p);

        /* set up the MPI description of MPI_Record */

        MPI_Address(&record, &baseaddress);
        lena[0] = 2; MPI_Address(&record.key,&loca[0]); 
        loca[0] -= baseaddress; typa[0] = MPI_INT;
        lena[1] = 1; MPI_Address(&record.data   ,&loca[1]); 
        loca[1] -= baseaddress; typa[1] = MPI_FLOAT;
        MPI_Type_struct(2, lena, loca, typa, &MPI_Record);
        MPI_Type_commit(&MPI_Record);

	/**** End of MPI stuff ****/
    
        
    init_local_keys(&local_keys);
    Redistribute_keys(&local_keys);

    //done...... Simple Sample Parallel Sort
    if(my_rank==0)
    {
        //sorting complete
        //data available in *(records+i) for i=0 --> list_size-1
        
        //Uncomment to play data
        //for(i=0;i<list_size;i++) printf("Key[%d]: %d Data: %f\n",i,(records+i)->key,(records+i)->data);
        
        //print stats;
        printf("Time to sort %d in parallel = %d\n",list_size,(long)stop-start);
    //}
    //else if(my_rank==1)
    //{
        //start timer
        start = clock();
        //run qsort on data
        qsort(records, list_size, sizeof(Element),
                (int(*)(const void*, const void*))(element_cmp));

        //stop timer
        stop = clock();

        //print stats;
        printf("Time to sort %d sequentially = %d\n",list_size,(long)(stop-start));
    }
        
    	/* Shut down MPI */ 
    	MPI_Finalize();
  
	return 0;
	
}

/*
 *
 * Redistribute Keys
 * 
 *
 */
void Redistribute_keys(
         LOCAL_LIST_T* local_keys) {

    int new_list_size, i;//, error = 0;
    int* send_counts;
    int* send_displacements;
    int* recv_counts;
    int* recv_displacements;
    Element* new_keys;
    
    // Allocate space for the counts and displacements 
    send_counts = (int*) malloc(p*sizeof(int));
    send_displacements = (int*) malloc(p*sizeof(int));
    recv_counts = (int*) malloc(p*sizeof(int));
    recv_displacements = (int*) malloc(p*sizeof(int));

    //start timer;
    if(my_rank==0)start=clock();
    
    Local_sort(local_keys);
    Find_alltoall_send_params(local_keys, 
        send_counts, send_displacements);
    
    // Distribute the counts
    MPI_Alltoall(send_counts, 1, MPI_INT, recv_counts,
        1, MPI_INT, MPI_COMM_WORLD);

    // Allocate space for new list
    new_list_size = recv_counts[0];
    for (i = 1; i < p; i++)
        new_list_size += recv_counts[i];
    new_keys = (Element*) 
        malloc(new_list_size*sizeof(Element));
 
    //printf("P%d: New List Size =%d\n",my_rank,new_list_size);
    //exit(0);

   
    Find_recv_displacements(recv_counts, recv_displacements);

    // Exchange the keys
    MPI_Alltoallv(List(local_keys), send_counts, 
        send_displacements, MPI_Record, new_keys, 
        recv_counts, recv_displacements, MPI_Record,
        MPI_COMM_WORLD);

    // Replace old list with new list 
    List_free(local_keys);
    List_size(local_keys) = new_list_size;
    List(local_keys) = new_keys;

    Local_sort(local_keys);
    //sort complete
    //stop timer;
    if(my_rank==0)stop=clock();

    //Gather answer
    
    //all processor send_counts
    send_counts[0]=new_list_size;
    for(i=1;i<p;i++)send_counts[i]=0;
    
    // Find the counts
    MPI_Alltoall(send_counts, 1, MPI_INT, recv_counts, 1, 
                 MPI_INT, MPI_COMM_WORLD);
    
    //matching disp
    recv_displacements[0]=0;
    for(i=1;i<p;i++)
        recv_displacements[i]=recv_displacements[i-1]+recv_counts[i-1];

    //Gather list by P0 into P0's records list  
    MPI_Gatherv(List(local_keys), List_size(local_keys), MPI_Record, records,
                recv_counts, recv_displacements, MPI_Record, 0, MPI_COMM_WORLD);
    

    // Free temporary storage 
    free(send_counts);
    free(send_displacements);
    free(recv_counts);
    free(recv_displacements);

} /* Redistribute_keys */

void init_data(Element **records,int count)
{
    int i=0;
    time_t now;
    (void) time(&now);
    srand48((long) now);
    
    *records = (Element *)malloc(list_size*sizeof(Element));
    for(i=0;i<list_size;i++)
    {
        ((*records)+i)->key = lrand48()%list_size;
        ((*records)+i)->data = lrand48()/lrand48(); 
    }
      
}

void Local_sort(LOCAL_LIST_T* local_keys)
{ 
    qsort(List(local_keys), List_size(local_keys), sizeof(Element), 
        (int(*)(const void*, const void*))(element_cmp));
}

void Splitter_sort(int* splits, int size)
{
    qsort(splits, size, sizeof(int), 
        (int(*)(const void*, const void*))(int_cmp));
}

int element_cmp(const Element *a , const Element *b)
{
    return (a->key - b->key) ;
}

int int_cmp(const int *a , const int *b)
{
    //printf("P%d: A=%d, B=%d, A-B=%d\n",my_rank,*a,*b,*a-*b);
    return (*a-*b);
}


void Find_alltoall_send_params(LOCAL_LIST_T* local_keys, int* send_counts,
    int* send_displacements) 
{
    int i, j, *mysplits, *allsplits, *allcounts, *alldisp, *final_splits;
    int split_size = list_size/(p*p);
    
    mysplits = (int*)malloc(split_size*sizeof(int));
    allsplits = (int*)malloc(split_size*p*sizeof(int));
    allcounts = (int*)malloc(p*sizeof(int));
    alldisp = (int*)malloc(p*sizeof(int));
    final_splits = (int*)malloc(p*sizeof(int));
    
    
    //pick_splitters
    for(i=p-1,j=0;i<List_size(local_keys) && j<split_size;i+=p,j++)
    {
        mysplits[j]=(List_key(local_keys,i)).key;
    }

    //Allgatherv splitters
        
        //setting up recieve buffers
        
        allcounts[0]=split_size;
        alldisp[0]=0;
        for(i=1; i<p; i++)
        {
            allcounts[i]=split_size;
            alldisp[i]= alldisp[i-1] + allcounts[i-1];
        }

        //MPI_Collective op
        MPI_Allgatherv ( mysplits, split_size, MPI_INT, allsplits, 
                      allcounts, alldisp, MPI_INT, MPI_COMM_WORLD);

    //local sort splitter set
    Splitter_sort(allsplits,split_size*p);
    
    //ALL pick_global_splitters
    for(j=0;j<p;j++)
    {
        final_splits[j]=allsplits[j*p];
        //if(my_rank==0)printf("P%d Final[%d]=%d\n",my_rank,j,final_splits[j]);
    }

    //find_counts & displacements
    j=0;
    send_counts[0]=0;
    for(i=0;i<List_size(local_keys);i++)
    {
        if(List_key(local_keys,i).key<=final_splits[j]) send_counts[j]++;
        else
        {
            send_counts[++j]=0;
            if(j == p-1)
            {
                send_counts[j]=List_size(local_keys)-i;
                break;
            }
            i--;
        }
    }
    send_displacements[0]=0;
    for(j=1;j<p;j++) 
        send_displacements[j]= send_counts[j-1] + send_displacements[j-1];

    free(mysplits);
    free(allsplits);
    free(allcounts);
    free(alldisp); 
    free(final_splits);
                     
    //for(i=0;i<p;i++) 
    //    printf("P%d Counts[%d]=%d, Disp[%d]=%d\n",my_rank,i,
    //    send_counts[i],i,send_displacements[i]);
    //exit(0);
    
}

void init_local_keys(LOCAL_LIST_T *local_keys)
{
    int i,j;

    List_size(local_keys)= list_size/p;
    local_keys->keys = (Element *)malloc((list_size/p)*sizeof(Element));
            
    for(i=my_rank*(list_size/p),j=0;j<(list_size/p);i++,j++)
    {    
        local_keys->keys[j].key = records[i].key;
        local_keys->keys[j].data = records[i].data;
    }
}

void Find_recv_displacements(int recv_counts[], int recv_disp[])
{
    int i;

    recv_disp[0] = 0;
    for (i = 1; i < p; i++)
        recv_disp[i] = recv_disp[i-1]+recv_counts[i-1];
}