/*
** Indexed type example
** Copyright Rolf Riesen 2008
**
*/
#include <stdio.h>
#include <stdlib.h>     /* For malloc() */
#include <string.h>     /* For memset() */
#include <mpi.h>


/* Local Functions */
char **malloc_2D(int n);
void print_line(char *str);
#define MAX_PRINT_LINE_LEN      (256)



int
main(int argc, char *argv[])
{

int my_rank, nproc;
int tag, source;
int dest1, dest2;
char str1[MAX_PRINT_LINE_LEN];
char str2[MAX_PRINT_LINE_LEN];
char **array;
char *ptr;
int x, y;
int n, i;
MPI_Datatype low_trig_t;
MPI_Datatype upper_trig_t;
int *blocklens;
int *displacements;
MPI_Status status;
int size;
MPI_Aint lb, extent;



    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nproc);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    if (nproc < 3)   {
	if (my_rank == 0)   {
	    printf("Need at least 3 processes\n");
	}
	MPI_Finalize();
	exit(-1);
    }


    /* Create a 2-D array */
    n= 3 * nproc;
    n= 8;
    array= malloc_2D(n);


    /* On the root node, fill it with some data */
    if (my_rank == 0)   {
        i= 0;
        for (y= 0; y < n; y++)   {
            for (x= 0; x < n; x++)   {
                array[y][x]= 'A' + (i % 26);
                i++;
            }
        }
    }


    /*
    ** Create two data types. One for the strictly lower triangular part of the array
    ** and a second for the strictly upper triangular part of the array.
    */
    blocklens= (int *)malloc(n * sizeof(int));
    if (blocklens == NULL)   {
        fprintf(stderr, "Out of Memory!\n");
        MPI_Abort(MPI_COMM_WORLD, -1);
    }

    displacements= (int *)malloc(n * sizeof(int));
    if (displacements == NULL)   {
        fprintf(stderr, "Out of Memory!\n");
        MPI_Abort(MPI_COMM_WORLD, -1);
    }

    /* Displacements and Offsets for the lower */
    for (i= 0; i < n; i++)   {
	displacements[i]= i * n;
	blocklens[i]= i;
    }

    MPI_Type_indexed(n, blocklens, displacements, MPI_CHAR, &low_trig_t);
    MPI_Type_commit(&low_trig_t);



    /* Displacements and Offsets for the upper */
    for (i= 0; i < n; i++)   {
	displacements[i]= i * n + i + 1;
	blocklens[i]= n - i - 1;
    }

    MPI_Type_indexed(n, blocklens, displacements, MPI_CHAR, &upper_trig_t);
    MPI_Type_commit(&upper_trig_t);


    if (my_rank == 0)   {
	printf("Type            LB   Extent   Size\n");
	printf("--------------+----+--------+-----\n");
	MPI_Type_get_extent(low_trig_t, &lb, &extent);
	MPI_Type_size(low_trig_t, &size);
	printf("low_trig_t      %2lld     %3lld    %3lld\n",
		(long long)lb, (long long)extent, (long long)size);
	MPI_Type_get_extent(upper_trig_t, &lb, &extent);
	MPI_Type_size(upper_trig_t, &size);
	printf("upper_trig_t    %2lld     %3lld    %3lld\n\n",
		(long long)lb, (long long)extent, (long long)size);
    }



    /* Send the lower triangular matrix to rank 1 */
    tag= 200;
    source= 0;
    dest1= 1;
    dest2= 2;

    if (my_rank == source)   {
	MPI_Send(&(array[0][0]), 1, low_trig_t, dest1, tag,
		MPI_COMM_WORLD);
	MPI_Send(&(array[0][0]), 1, upper_trig_t, dest2, tag,
		MPI_COMM_WORLD);
    } else if (my_rank == dest1)   {
        MPI_Recv(&(array[0][0]), 1, low_trig_t, source, tag,
		MPI_COMM_WORLD, &status);
    } else if (my_rank == dest2)   {
        MPI_Recv(&(array[0][0]), 1, upper_trig_t, source, tag,
		MPI_COMM_WORLD, &status);
    } else   {
	/* Do nothing */
    }


    /* Print a Title Line */
    sprintf(str1, "Rank %3d", my_rank);
    for (x= 0; x < n; x++)   {
        strcat(str1, " ");
    }
    print_line(str1);


    /* Build each row and display */
    ptr= array[0];
    for (y= 0; y < n; y++)   {
        sprintf(str1, "[%3d]   ", my_rank);
        for (x= 0; x < n; x++)   {
            sprintf(str2, "%1c", *ptr);
            strcat(str1, str2);
            ptr++;
        }
        print_line(str1);
    }

    MPI_Finalize();
    return 0;

}  /* end of main() */



/*
** Allocate a 2-D array of chars of size n
*/
char **
malloc_2D(int n)
{

int array_size, ptr_size;
char **array;
char *ptr;
int i;


    /* Size for the data: n^2 char */
    array_size= n * n * sizeof(char);

    /* Size for the pointers: n pointers to char */
    ptr_size= n * sizeof(char *);

    /* Allocate all memory in one chunk */
    array= (char **)malloc(array_size + ptr_size);
    if (array == NULL)   {
        fprintf(stderr, "Memory allocation failed!\n");
        MPI_Abort(MPI_COMM_WORLD, -1);
    }

    /* Clear it */
    memset(array, '_', array_size + ptr_size);

    /* Fill in the pointer array */
    ptr= (char *)(array + n);
    for (i= 0; i < n; i++)   {
        array[i]= ptr;
        ptr= ptr + n;
    }

    return array;

}  /* end of malloc_2D() */



/*
** Each node calls print_line() with a string. Rank 0 collects
** all strings and prints them on a single line in rank order.
** The strings should not contain formating characters like \n.
*/
void
print_line(char *str)
{

char line[MAX_PRINT_LINE_LEN];
int nproc, my_rank;
int i, len;
int source, tag;
MPI_Status status;


    len= strlen(str) + 1;
    if (len > MAX_PRINT_LINE_LEN)   {
        MPI_Abort(MPI_COMM_WORLD, -1);
    }

    MPI_Comm_size(MPI_COMM_WORLD, &nproc);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    tag= 953;

    if (my_rank == 0)   {
        printf("%s", str);
        for (i= 1; i < nproc; i++)   {
            source= i;
            MPI_Recv(line, MAX_PRINT_LINE_LEN, MPI_CHAR, source, tag, MPI_COMM_WORLD, &status);
            printf("     %s", line);
            MPI_Barrier(MPI_COMM_WORLD);
        }
        printf("\n");
    } else   {
        for (i= 1; i < nproc; i++)   {
            if (i == my_rank)   {
                MPI_Send(str, len, MPI_CHAR, 0, tag, MPI_COMM_WORLD);
            }
            MPI_Barrier(MPI_COMM_WORLD);
        }
    }

}  /* end of [rint_line() */