/* psimeq1.c  using broadcast and gather, unoptimized  */
#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>

#undef  abs
#define abs(x) ((x)<0.0?(-(x)):(x))

static int n=1024;  /* number of equations */

int main(int argc, char *argv[])
{
  int proc, myid, numprocs;
  int count = 1;
  int master = 0;
  double esec = 0.0;
  double * gath;    /* only master receives here from Gather */
  double * Y;       /* only master has solution              */
  int * ROW;        /* only master keeps row interchange     */
  int i, j, jj;
  int knout[2];     /* column being reduced, number of equations */
  int ipivrow;      /* pivot row */
  int * upivrow;    /* local used pivot rows */
  double *rbuf;     /* gather buffer row, ipiv, abspiv */
  double *B;        /* local equations nn by n+1 */
  int nn;           /* nn=(n+numprocs-1)/numprocs local equations */
  int ii;           /* ii=nn*myid  index of first local equation */
  int k;            /* pivot column  global */
  int ipivot;       /* pivot row of max pivot */
  double abspivot;  /* absolute value of largest pivot */

  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &myid);
  knout[0] = 0; /* initialize for Bcast */
  knout[1] = n;
  ipivrow = -1;

  if(myid == master)
  {
    printf("psimeq1.c  numprocs=%d, n=%d \n", numprocs, n);
    printf("Bcast is blocking and must be executed by all.\n");
    printf("Gather automatically indexes through gath.\n");
    printf("only master receives from Gather.\n");
    printf("no order guaranteed for printf output.\n");
  } /* end master */

  /* all */
  count = 2;
  MPI_Bcast(knout, count, MPI_INT, master, MPI_COMM_WORLD);
  k = knout[0];
  n = knout[1];
  nn=(n+numprocs-1)/numprocs; /* number of local equations */
  ii=nn*myid;                 /* index of first local equation */
  B = (double *)malloc((n+1)*nn*sizeof(double));
  rbuf = (double *)malloc((n+3)*sizeof(double));
  upivrow = (int *)malloc(nn*sizeof(int));
  for(i=0; i<nn; i++) upivrow[i] = 0;

  if(myid == master)
  {
    gath = (double *)malloc((n+3)*numprocs*sizeof(double));
    Y = (double *)malloc(nn*numprocs*sizeof(double));
    ROW  = (int *)malloc(nn*numprocs*sizeof(int));
    for(i=0; i<nn*numprocs; i++) ROW[i] = i;
  } /* end master unique setup */

  /* build test case, in general distribute or compute equations */
  /* all */
  for(i=0; i<nn; i++)
  {
    for(j=0; j<n; j++)
    {
      if(ii+i>j) B[i*(n+1)+j] = (double)(ii+i+1); /* test data */
      else       B[i*(n+1)+j] = (double)(j+1);
    }
    B[i*(n+1)+n] = 0.0;
    for(j=0; j<n; j++)
    {                       /* solution n, n-1, n-2, ... 1 */
      B[i*(n+1)+n] = B[i*(n+1)+n]+ (double)(n-j)*B[i*(n+1)+j];
    }
  }
  /* end building test case */

  while(k<n)
  {
    /* find local max pivot */
    ipivot = -1;
    abspivot = 0.0;
    for(i=0; i<nn; i++)
    {
      if(upivrow[i]==0 && abs(B[i*(n+1)+k])>abspivot)
      {
        abspivot = abs(B[i*(n+1)+k]);
        ipivot = i;
      }
    }
    if(ipivot>=0)
    {
      for(j=k+1; j<=n; j++)
        rbuf[j] = B[ipivot*(n+1)+j]/B[ipivot*(n+1)+k];
    }
    else
    {
      for(j=k+1; j<=n; j++)
        rbuf[j] = 0.0;
    }
    rbuf[n+1] = (double)(ipivot); /* local row */
    rbuf[n+2] = abspivot;
    /* unused locations */
    for(j=0; j<k; j++)
    {
      rbuf[j] = 0.0;
    }
    rbuf[k] = 1.0;

    count = n+3;
    MPI_Gather(rbuf, count, MPI_DOUBLE,
               gath, count, MPI_DOUBLE,
               master, MPI_COMM_WORLD);


    if(myid == master)
    {
      if(k%100==0)
      {
        esec = MPI_Wtime();
        printf("master gath vector k=%d at %g sec \n", k, esec);
      }
      ipivrow = -1;
      abspivot = 0.0;
      jj = -1;
      for(i=0; i<numprocs; i++)
      {
        /* printf("gath ipivot[%d]=%g, abspivot=%g \n",
	          i, gath[i*(n+3)+(n+1)], gath[i*(n+3)+(n+2)]); */
        if((int)gath[i*(n+3)+(n+1)]>=0 && gath[i*(n+3)+(n+2)]>abspivot)
        {
          abspivot = gath[i*(n+3)+(n+2)];
          ipivrow = (int)gath[i*(n+3)+(n+1)];
          jj = i;
        }
      }

      if(jj<0)
      {
        printf("Big trouble, non independent equations! \n");
        for(j=0; j<n+3; j++) rbuf[j] = 0.0;
      }
      else
      {
        ROW[k] = ipivrow+jj*nn; /* not interchange, just set */
        for(j=0; j<=n; j++) rbuf[j] = gath[jj*(n+3)+j];
        rbuf[n+1] = (double)ipivrow;
        rbuf[n+2] = (double)jj;
        /* for(j=0; j<=n+2; j++) printf("rbuf[%d]=%g \n", j, rbuf[j]); */
      }
    } /* end master finding maximum pivot row */


    /* all */
    count = n+3;
    MPI_Bcast(rbuf, count, MPI_DOUBLE, master, MPI_COMM_WORLD);
    ipivrow = (int)rbuf[n+1];
    jj = (int)rbuf[n+2];

    /* one */
    if(myid == jj && ipivrow>=0)
    {
      upivrow[ipivrow] = 1;
    }

    /* reduce */
    for(i=0; i<nn; i++)
    {
      if(myid==jj && i==ipivrow) /* only one node will have this true */
      {
        for(j=0; j<=n; j++) B[i*(n+1)+j] = rbuf[j];
      }
      else
      {
        for(j=k+1; j<=n; j++)
        {
          B[i*(n+1)+j] = B[i*(n+1)+j] - B[i*(n+1)+k]*rbuf[j];
        }
        /* unused */
        B[i*(n+1)+k] = 0.0;
      }
    }
    /* end reduce */
    k++;                /* ready for next column */
  } /* end while */

  for(i=0; i<nn; i++)
  {
    rbuf[i] = B[i*(n+1)+n]; /* return solution */
  }
  /* return solution to master */
  count = nn;
  MPI_Gather(rbuf, count, MPI_DOUBLE,
             gath, count, MPI_DOUBLE,
             master, MPI_COMM_WORLD);

  if(myid == master)
  {
    for(i=0; i<n; i++)
    {
      Y[i] = gath[ROW[i]];
      if(i%100==0) printf("Y[%d]=%g \n", i, Y[i]);
    }
    printf("Y[%d]=%g \n", n-1, Y[n-1]);
    esec = MPI_Wtime();
    printf("end psimeq1.c at %g sec\n", esec);
  } /* end master */ 

  /* all */
  MPI_Finalize();
  return 0; 
} /* end psimeq1.c */