/*@@
   @file      Write2D.c
   @date      Tue Feb  3 19:06:12 1998
   @author    Paul Walker
   @desc 
   Two dimensional slices into IEEEIO files on MPI distributed cactus data.
   @enddesc 
   @history
   @hauthor Gabrielle Allen @hdate 26 Sep 1998
   @hdesc Changed subroutine name (SliceTwoD -> Write2D), renamed IO parameters
   @hauthor Gabrielle Allen @hdate Oct 17 1998
   @hdesc Added the IO_verbose parameter instead of an ifdef
   @hauthor Gabrielle Allen @hdate 19 Oct 1998      
   @hdesc Changed names ready for thorn_IO
   @hauthor Thomas Radke @hdate 16 Mar 1999
   @hdesc Converted to Cactus 4.0
   @hauthor Thomas Radke @hdate 17 Mar 1999
   @hdesc Changed naming: IEEEIO -> IOFlexIO
   @hendhistory
   @version $Id$
 @@*/

static char *rcsid = "$Id$";

#include <stdio.h>
#include <assert.h>
#include <stdlib.h>

#include "cctk.h"
#include "cctk_parameters.h"
#include "CactusPUGH/PUGH/src/include/pugh.h"
#include "CactusBase/IOUtil/src/ioGH.h"
#include "ioFlexGH.h"

/* MPI message tag */
#define BASE 1234


/*@@
   @routine    IOFlexIO_Write2D
   @date       Sat March 6 1999
   @author     Gabrielle Allen
   @desc       Writes the 2D data of a variable into separate IEEEIO files
   @enddesc
   @calls    CCTK_GHExtensionHandle
             CCTK_MyProc
             CCTK_nProcs
             CCTK_WARN
   @calledby IOFlexIO_Output2DVarAs
             IOFlexIO_TriggerOutput2D
   @history
   @endhistory
   @var     GH
   @vdesc   Pointer to CCTK GH
   @vtype   cGH
   @vio     in
   @vcomment
   @endvar
   @var     index
   @vdesc   index of variable to output
   @vtype   int
   @vio     in
   @vcomment
   @endvar
   @var     alias
   @vdesc   alias name of variable to output
   @vtype   const char *
   @vio     in
   @vcomment
   @endvar
@@*/

void IOFlexIO_Write2D (cGH *GH, int index, const char *alias)
{
  DECLARE_CCTK_PARAMETERS
  int i, j, k, l;
  int myproc, nprocs;
  pGH *pughGH;
  ioGH *ioUtilGH;
  flexioGH *myGH;
  int dir, ngpoints;
  CCTK_INT npoints, *nrempoints, bnd [4];
  void *data;
  int ni, nj;                        /* "x" and "y" direction ns */
  /* Attributes */
  CCTK_REAL origin [2], delta [2];
  int maxrempt;
  int dims [2];
  void *locdat = NULL, *alldat = NULL, *recd;
  int hi = 0, lo = 0;
  int locbnd [4];
  int timelevel, variable_type, element_size, ioflex_type;
  IOFile *IEEEfile_2D = NULL;
#ifdef MPI
  MPI_Status ms;
  int mpi_type;
#endif


  /* Get the handles for pugh, IOFlexIO and IOUtil extensions */
  pughGH = (pGH *) GH->extensions [CCTK_GHExtensionHandle ("PUGH")];
  ioUtilGH = (ioGH *) GH->extensions [CCTK_GHExtensionHandle ("IO")];
  myGH = (flexioGH *) GH->extensions [CCTK_GHExtensionHandle ("IOFlexIO")];

  /* What processor are we on? */
  myproc = CCTK_MyProc (GH);
  nprocs = CCTK_nProcs (GH);

  if (verbose) {
    printf ("<><><><><><><><><><><><><><><><><>\n");
    printf ("2D Slice for [%s]\n", alias);
#if        0
    cactus_memtrace ();
#endif
  }

  /* Open the files on the first trip through if we are proc. 0 */
  if (myproc == 0) {

    /* see if output file for this alias name was already created */
    IEEEfile_2D = (IOFile *) GetNamedData (myGH->filenameList2D, alias);
    if (IEEEfile_2D == NULL) {
      int len;
      char *fname;
      char *msg;

      IEEEfile_2D = (IOFile *) malloc (3 * sizeof (IOFile));
      len = strlen (myGH->outpfx_2D) + strlen (alias);
      fname = (char *) malloc (len + 20);
      msg = (char *) malloc (len + 20 + 40);
      assert (IEEEfile_2D && fname && msg);
 
      /* Create files (open mode "w") */
      sprintf (fname, "%s/%s_2d_yz.ieee", myGH->outpfx_2D, alias);
      IEEEfile_2D [0] = IEEEopen (fname, "w");
      if (! IOisValid (IEEEfile_2D [0])) {
        sprintf (msg, "Cannot open output file %s", fname);
        CCTK_WARN (1, msg);
        return;
      }
      
      sprintf (fname, "%s/%s_2d_xz.ieee", myGH->outpfx_2D, alias);
      IEEEfile_2D [1] = IEEEopen (fname, "w");
      if (! IOisValid (IEEEfile_2D [1])) {
        sprintf (msg, "Cannot open output file %s", fname);
        CCTK_WARN (1, msg);
        return;
      }
      
      sprintf (fname, "%s/%s_2d_xy.ieee", myGH->outpfx_2D, alias);
      IEEEfile_2D [2] = IEEEopen (fname, "w");
      if (! IOisValid (IEEEfile_2D [2])) {
        sprintf (msg, "Cannot open output file %s", fname);
        CCTK_WARN (1, msg);
        return;
      }

      /* store file desriptors in database */
      StoreNamedData (&myGH->filenameList2D, alias, IEEEfile_2D);

      free (msg);
      free (fname);
    }
  }

  nrempoints = (CCTK_INT *) malloc (nprocs * sizeof (CCTK_INT));
  assert (nrempoints);

  timelevel = CCTK_NumTimeLevelsFromVarI (index) - 1;
  if (timelevel > 0)
    timelevel--;
  data = CCTK_VarDataPtrI (GH, timelevel, index);

  variable_type = CCTK_VarTypeI (index);
  switch (variable_type) {
    case CCTK_VARIABLE_CHAR:
      ioflex_type = FLEXIO_CHAR;
      element_size = sizeof (CCTK_CHAR);
#ifdef MPI
      mpi_type = PUGH_MPI_CHAR;
#endif
      break;

    case CCTK_VARIABLE_INT:
      ioflex_type = FLEXIO_INT;
      element_size = sizeof (CCTK_INT);
#ifdef MPI
      mpi_type = PUGH_MPI_INT;
#endif
      break;

    case CCTK_VARIABLE_REAL:
      ioflex_type = FLEXIO_REAL;
      element_size = ioUtilGH->out_single ?
                     sizeof (CCTK_REAL4) : sizeof (CCTK_REAL);
#ifdef MPI
      mpi_type = ioUtilGH->out_single ? PUGH_MPI_REAL4 : PUGH_MPI_REAL;
#endif
      break;

    default:
      CCTK_WARN (1, "Unsupported variable type in IOFlexIO_Write2D");
      return;
  }

  for (dir = 0; dir < 3; dir++) {
    npoints = 0;
    /* Figure out the number of points */
    if ((pughGH->ownership[pughGH->stagger][dir][0] <= pughGH->sp2dxyz[dir]) &&
        (pughGH->sp2dxyz[dir] < pughGH->ownership[pughGH->stagger][dir][1]))
      npoints = (pughGH->ownership[pughGH->stagger][(dir+1)%3][1] -
                 pughGH->ownership[pughGH->stagger][(dir+1)%3][0]) * 
                (pughGH->ownership[pughGH->stagger][(dir+2)%3][1] -
                 pughGH->ownership[pughGH->stagger][(dir+2)%3][0]);

    /* Set up the number of global points and the points in the "i"
       and "j" directions */
    if (dir == 0) {
      ngpoints = GH->cctk_gsh[1] * GH->cctk_gsh[2];
      ni = GH->cctk_gsh[1];
      origin[0] = CCTK_CoordOrigin ("y");
      delta[0] = pughGH->dy0;
      nj = GH->cctk_gsh[2];
      origin[1] = CCTK_CoordOrigin ("z");
      delta[1] = pughGH->dz0;
    } else if (dir == 1) {
      ngpoints = GH->cctk_gsh[0] * GH->cctk_gsh[2];
      ni = GH->cctk_gsh[0];
      origin[0] = CCTK_CoordOrigin ("x");
      delta[0] = pughGH->dx0;
      nj = GH->cctk_gsh[2];
      origin[1] = CCTK_CoordOrigin ("z");
      delta[1] = pughGH->dz0;
    } else {
      ngpoints = GH->cctk_gsh[0] * GH->cctk_gsh[1];
      ni = GH->cctk_gsh[0];
      origin[0] = CCTK_CoordOrigin ("x");
      delta[0] = pughGH->dx0;
      nj = GH->cctk_gsh[1];
      origin[1] = CCTK_CoordOrigin ("y");
      delta[1] = pughGH->dy0;
    }

    if (verbose) {
      printf ("Npoints in dir %d is %d -> ",dir,npoints);
      printf ("Global npoints is %d\n",ngpoints);
    }

#ifdef MPI
    /* Send the number of points */
    CACTUS_MPI_ERROR (MPI_Gather (&npoints, 1, PUGH_MPI_INT,
                      nrempoints, 1, PUGH_MPI_INT, 0, pughGH->PUGH_COMM_WORLD));

    if (verbose && myproc == 0)
      for (i = 0; i < nprocs; i++)
        printf ("From proc %d: %d\n", i, (int) nrempoints [i]);

#else
    nrempoints [0] = npoints;
#endif

    maxrempt = 0;
    if (myproc == 0)
      for (i = 0; i < nprocs; i++)
        maxrempt = (nrempoints [i] > maxrempt ? nrempoints [i] : maxrempt);

    if (verbose)
      printf ("Maximum remote points: %d\n", maxrempt);

    /* OK so we want the ordering in bnd to be yz xz or xy but the
       mod naturally gives is yz zx xy so we need to flip the order.
       thus put the dir+1 % 3 in at lo, the dir+2 % 3 in at hi, and
       set lo and hi directionally as follows ... */
    if (dir == 0) {lo = 0; hi = 2;}
    if (dir == 1) {lo = 2; hi = 0;}
    if (dir == 2) {lo = 0; hi = 2;}
    if (npoints > 0) {
      /* See above comment ... */
      bnd [lo]   = GH->cctk_lbnd[(dir+1)%3] + 
                   pughGH->ownership[pughGH->stagger][(dir+1)%3][0];
      bnd [lo+1] = GH->cctk_lbnd[(dir+1)%3] + 
                   pughGH->ownership[pughGH->stagger][(dir+1)%3][1];

      bnd [hi]   = GH->cctk_lbnd[(dir+2)%3] + 
                   pughGH->ownership[pughGH->stagger][(dir+2)%3][0];
      bnd [hi+1] = GH->cctk_lbnd[(dir+2)%3] + 
                   pughGH->ownership[pughGH->stagger][(dir+2)%3][1];

      locbnd [lo]   = pughGH->ownership[pughGH->stagger][(dir+1)%3][0];
      locbnd [lo+1] = pughGH->ownership[pughGH->stagger][(dir+1)%3][1];

      locbnd [hi]   = pughGH->ownership[pughGH->stagger][(dir+2)%3][0];
      locbnd [hi+1] = pughGH->ownership[pughGH->stagger][(dir+2)%3][1];

    } else {
      bnd [lo] = 0;
      bnd [lo+1] = 0;
      bnd [hi] = 0;
      bnd [hi+1] = 0;
    }

    if (verbose)
      printf ("Local bounds %d: j -> %d %d  k -> %d %d\n",
               dir, locbnd [0], locbnd [1], locbnd [2], locbnd [3]);

    /* Subsample the data */
    if (npoints > 0) {
      locdat = malloc (npoints * element_size);
      assert (locdat);
      l = 0;
      if (variable_type == CCTK_VARIABLE_CHAR) {
        CCTK_CHAR *char_locdat = (CCTK_CHAR *) locdat;

        for (k = locbnd [2]; k < locbnd [3]; k++)
          for (j = locbnd [0]; j < locbnd [1]; j++) {
            int pt = 0;
            if (dir==0) pt = DATINDEX (pughGH, pughGH->sp2dxyz [dir], j, k);
            if (dir==1) pt = DATINDEX (pughGH, j, pughGH->sp2dxyz [dir], k);
            if (dir==2) pt = DATINDEX (pughGH, j, k, pughGH->sp2dxyz [dir]);
            assert (l <= npoints && pt <= pughGH->npoints);
            char_locdat [l++] = ((CCTK_CHAR *) data) [pt];
          }
      } else if (variable_type == CCTK_VARIABLE_INT) {
        CCTK_INT *int_locdat = (CCTK_INT *) locdat;

        for (k = locbnd [2]; k < locbnd [3]; k++)
          for (j = locbnd [0]; j < locbnd [1]; j++) {
            int pt = 0;
            if (dir==0) pt = DATINDEX (pughGH, pughGH->sp2dxyz [dir], j, k);
            if (dir==1) pt = DATINDEX (pughGH, j, pughGH->sp2dxyz [dir], k);
            if (dir==2) pt = DATINDEX (pughGH, j, k, pughGH->sp2dxyz [dir]);
            assert (l <= npoints && pt <= pughGH->npoints);
            int_locdat [l++] = ((CCTK_INT *) data) [pt];
          }
      } else if (variable_type == CCTK_VARIABLE_REAL) {
        CCTK_REAL4 *real4_locdat = (CCTK_REAL4 *) locdat;
        CCTK_REAL  *real_locdat = (CCTK_REAL *) locdat;

        for (k = locbnd [2]; k < locbnd [3]; k++)
          for (j = locbnd [0]; j < locbnd [1]; j++) {
            int pt = 0;
            if (dir==0) pt = DATINDEX (pughGH, pughGH->sp2dxyz [dir], j, k);
            if (dir==1) pt = DATINDEX (pughGH, j, pughGH->sp2dxyz [dir], k);
            if (dir==2) pt = DATINDEX (pughGH, j, k, pughGH->sp2dxyz [dir]);
            assert (l <= npoints && pt <= pughGH->npoints);
            if (ioUtilGH->out_single)
              real4_locdat [l++] = (CCTK_REAL4) (((CCTK_REAL *) data) [pt]);
            else
              real_locdat [l++] = ((CCTK_REAL *) data) [pt];
          }
      }
    } 

    /* Send the data */
#ifdef MPI
    /* Only send if we have points and we are not on proc 0 */
    if (npoints > 0 && myproc != 0) {
      CACTUS_MPI_ERROR (MPI_Send (bnd, 4, PUGH_MPI_INT, 0, 2*myproc+1+BASE,
                        pughGH->PUGH_COMM_WORLD));
      CACTUS_MPI_ERROR (MPI_Send (locdat, (int) npoints, mpi_type, 0,
                        2*myproc+BASE, pughGH->PUGH_COMM_WORLD));
    }

    /* Receive the data on processor 0 */
    if (myproc == 0) {
      alldat = malloc (ngpoints * element_size);
      recd   = malloc (maxrempt * element_size);
      assert (alldat && recd);

      for (i = 0; i < nprocs; i++) {
        void *ud;

        if (verbose)
          printf ("Proceesing proc %d\n",i);fflush(stdout);

        if (nrempoints [i] > 0) {
          if (i == 0)
            ud = locdat;
          else {

            if (verbose) {
              printf ("Receiving %d from %d\n", nrempoints [i], i);
              fflush (stdout);
            }

            CACTUS_MPI_ERROR (MPI_Recv (bnd, 4, PUGH_MPI_INT,
                              i, 2*i+1+BASE, pughGH->PUGH_COMM_WORLD, &ms));
            CACTUS_MPI_ERROR (MPI_Recv (recd, (int) nrempoints [i], mpi_type,
                              i, 2*i+BASE, pughGH->PUGH_COMM_WORLD, &ms));
            ud = recd;

            if (verbose) {
              printf ("   Received %d points from %d\n",
                      (int) nrempoints [i], i);
              printf ("   Data lives at [%d %d] -> [%d %d]\n",
                      (int) bnd [0], (int) bnd [2], (int) bnd [1], (int)bnd[3]);
            }

          } /* End proc != 0 */

          if (verbose)
            printf ("Remote Bounds %d: j -> %d %d  k -> %d %d\n",
                    dir, (int) bnd [0], (int) bnd[1], (int)bnd[2], (int)bnd[3]);

          l = 0;
          if (variable_type == CCTK_VARIABLE_CHAR)
            for (k = bnd [2]; k < bnd [3]; k++)
              for (j = bnd [0]; j < bnd [1]; j++)
                ((CCTK_CHAR *) alldat) [j + k*ni] = ((CCTK_CHAR *) ud) [l++];
          else if (variable_type == CCTK_VARIABLE_INT)
            for (k = bnd [2]; k < bnd [3]; k++)
              for (j = bnd [0]; j < bnd [1]; j++)
                ((CCTK_INT *) alldat) [j + k*ni] = ((CCTK_INT *) ud) [l++];
          else if (variable_type == CCTK_VARIABLE_REAL && ioUtilGH->out_single)
            for (k = bnd [2]; k < bnd [3]; k++)
              for (j = bnd [0]; j < bnd [1]; j++)
                ((CCTK_REAL4 *) alldat) [j + k*ni] = ((CCTK_REAL4 *) ud) [l++];
          else if (variable_type == CCTK_VARIABLE_REAL && !ioUtilGH->out_single)
            for (k = bnd [2]; k < bnd [3]; k++)
              for (j = bnd [0]; j < bnd [1]; j++)
                ((CCTK_REAL *) alldat) [j + k*ni] = ((CCTK_REAL *) ud) [l++];
        } /* End if > 0 */
      } /* End processor loop */
      free (recd);
    }
#else
    alldat = locdat;
#endif

    /* Proc 0 writes */
    if (myproc == 0) {
      dims [0] = ni;
      dims [1] = nj;

      /*** FIXME: since slice center is not yet set up, npoints might be 0 ***/
      if (npoints <= 0)
        CCTK_WARN (1, "2D output doesn't work since slice center is not set up");
      else {
        CCTK_REAL min_ext [2], max_ext [2];

        /* the data itself */
        CACTUS_IEEEIO_ERROR (IOwrite (IEEEfile_2D [dir],
                             ioflex_type, 2, dims, alldat));
      
        /* Time and space attributes */
        CACTUS_IEEEIO_ERROR (IOwriteAttribute (IEEEfile_2D [dir],
                             "time", FLEXIO_REAL, 1, &GH->cctk_time));
        CACTUS_IEEEIO_ERROR (IOwriteAttribute (IEEEfile_2D [dir],
                             "origin", FLEXIO_REAL, 2, origin));
        CACTUS_IEEEIO_ERROR (IOwriteAttribute (IEEEfile_2D [dir],
                             "delta", FLEXIO_REAL, 2, delta));

        /* and also the coordinate ranges */
        CCTK_CoordRange (GH, &min_ext [0], &max_ext [0], "x");
        CCTK_CoordRange (GH, &min_ext [1], &max_ext [1], "x");
        CACTUS_IEEEIO_ERROR (IOwriteAttribute (IEEEfile_2D [dir],
                             "min_ext",  FLEXIO_REAL, 2, min_ext));
        CACTUS_IEEEIO_ERROR (IOwriteAttribute (IEEEfile_2D [dir],
                             "max_ext",  FLEXIO_REAL, 2, max_ext));
      }
    }

    /* Free memory for this time through */
    if (npoints > 0) 
      free (locdat);
#ifdef MPI
    if (myproc == 0)
      free (alldat);
#endif /* If not MPI then alldat and locdat are not distinct */

  }

  free (nrempoints);

#if        0
  if (verbose)
    cactus_memtrace();
#endif
}
 


#if        0
void FMODIFIER FORTRAN_NAME(io_write2d_,
                            IO_WRITE2D,
                            io_write2d)
     (CCTK_REAL *x)
{
  pGF *GF;

  GF = locateGFbyDataPtr(x);
  IO_Write2D(GF->parentGH,GF);
}

#endif