/*@@
   @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$
 @@*/


#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"

static char *rcsid = "$Id$";
CCTK_FILEVERSION(CactusPUGHIO_IOFlexIO_Write2D_c)


/* 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
   @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;
  void *locdat = NULL, *alldat = NULL;
  int hi = 0, lo = 0;
  int locbnd [4];
  int timelevel, variable_type, element_size, ioflex_type;
  IOFile *IEEEfile_2D = NULL;
#ifdef CCTK_MPI
  void *recd;
  MPI_Status ms;
  MPI_Datatype mpi_type;
#endif
  pGA *GA;

  /* first, check if variable has storage assigned */
  if (! CCTK_QueryGroupStorageI (GH, CCTK_GroupIndexFromVarI (index))) 
  {
    char *fullname = CCTK_FullName (index);

    CCTK_VWarn (2, __LINE__, __FILE__, CCTK_THORNSTRING,
                "No IOFlexIO 2D output for '%s' (no storage)", fullname);
    free (fullname);
    return;
  }

  /* Get the handles for PUGH, IOUtil, and IOFlexIO extensions */
  pughGH = PUGH_pGH (GH);
  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);
  }

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

    case CCTK_VARIABLE_INT:
      ioflex_type = FLEXIO_INT;
      element_size = sizeof (CCTK_INT);
#ifdef CCTK_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 CCTK_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;
  }

  /* 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) 
    {
      char *fname;

      IEEEfile_2D = (IOFile *) malloc (3 * sizeof (IOFile));
      fname = (char *) malloc (strlen (myGH->outdir2D) + strlen (alias) +
                               strlen (pughGH->identity_string) + 20);
      assert (IEEEfile_2D && fname);
 
      /* Open/Create files */
      for (i = 0; i < 3; i++)
      {
        const char *extensions [3] = {"yz", "xz", "xy"};

        sprintf (fname, "%s/%s_2d_%s%s.ieee", myGH->outdir2D, alias,
                 extensions [i], pughGH->identity_string);

        /* if restart from recovery, try to open an exisiting file ... */
        IEEEfile_2D [i] = NULL;
        if (ioUtilGH->recovered)
          IEEEfile_2D [i] = IEEEopen (fname, "a");

        /* otherwise or if that failed, create a new one */
        if (! IOisValid (IEEEfile_2D [i]))
          IEEEfile_2D [i] = IEEEopen (fname, "w");
        if (! IOisValid (IEEEfile_2D [i])) 
        {
          CCTK_VWarn (1, __LINE__, __FILE__, CCTK_THORNSTRING,
                      "Cannot open 2D output file '%s'", fname);
          return;
        }
      }
      
      /* store file desriptors in database */
      StoreNamedData (&myGH->filenameList2D, alias, IEEEfile_2D);

      free (fname);
    }

    if (reuse_filehandles)
    {
      /* resume the files (reopen descriptors)
         This allows descriptor reuse without requiring expensive destruction
         and reallocation of the associated data structures */
      if (verbose)
        CCTK_VInfo (CCTK_THORNSTRING, "Resuming 2D output files for variable "
                    "'%s'", CCTK_VarName (index));

      for (dir = 0; dir < 3; dir++)
        CACTUS_IEEEIO_ERROR (IOresume (IEEEfile_2D [dir]));
    }
  }

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

  timelevel = CCTK_NumTimeLevelsFromVarI (index) - 1;
  if (timelevel > 0)
    timelevel--;

  /* Get the GA for this index */
  GA = ((pGA ***)pughGH->variables)[index][timelevel];

  data = CCTK_VarDataPtrI (GH, timelevel, index);

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

    /* 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 CCTK_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->GFExtras[2]->ownership[0][0][(dir+1)%3];
      bnd [lo+1] = GH->cctk_lbnd[(dir+1)%3] + 
                   pughGH->GFExtras[2]->ownership[0][1][(dir+1)%3];

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

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

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

    } 
    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 (GA->extras, myGH->sp2xyz [dir], j, k);
            if (dir==1) pt = DATINDEX (GA->extras, j, myGH->sp2xyz [dir], k);
            if (dir==2) pt = DATINDEX (GA->extras, j, k, myGH->sp2xyz [dir]);
            assert (l <= npoints && pt <= GA->extras->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 (GA->extras, myGH->sp2xyz [dir], j, k);
            if (dir==1) pt = DATINDEX (GA->extras, j, myGH->sp2xyz [dir], k);
            if (dir==2) pt = DATINDEX (GA->extras, j, k, myGH->sp2xyz [dir]);
            assert (l <= npoints && pt <= GA->extras->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 (GA->extras, myGH->sp2xyz [dir], j, k);
	    }
            if (dir==1) 
	    {
	      pt = DATINDEX (GA->extras, j, myGH->sp2xyz [dir], k);
	    }
            if (dir==2) 
	    {
	      pt = DATINDEX (GA->extras, j, k, myGH->sp2xyz [dir]);
	    }
            assert (l <= npoints && pt <= GA->extras->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 CCTK_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) 
    {
      int dims [2];
      CCTK_REAL min_ext [2], max_ext [2];

      dims [0] = ni;
      dims [1] = nj;

      /* 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 CCTK_MPI
    if (myproc == 0)
    {
      free (alldat);
    }
#endif /* If not MPI then alldat and locdat are not distinct */

  }

  free (nrempoints);

  if (reuse_filehandles && myproc == 0)
  {
    /* pause the files (close system descriptors) */
    if (verbose)
      CCTK_VInfo (CCTK_THORNSTRING, "Pausing 2D output files for variable '%s'",
                  CCTK_VarName (index));

    for (dir = 0; dir < 3; dir++)
      CACTUS_IEEEIO_ERROR (IOpause (IEEEfile_2D [dir]));
  }

}