/***************************************
  $Header$

  This is a neat 3-D graphing application of Illuminator.  Just put whatever
  function you like down in line 110 or so (or graph the examples provided), or
  run with -twodee and use PETSc's native 2-D graphics (though that would be
  BORING!).  You might want to run it as:
  +latex+\begin{quote}{\tt
  +html+ <blockquote><tt>
  ./3dgf -da_grid_x 50 -da_grid_y 50 -da_grid_z 50
  +latex+}\end{quote}
  +html+ </tt></blockquote>
  and hit return to end.
***************************************/


static char help[] = "A neat 3-D graphing application of Illuminator.\n\n\
Use -da_grid_x etc. to set the discretization size.\n\
Other options:\n\
  -twodee : (obvious)\n\
So you would typically run this using:\n\
  3dgf -da_grid_x 20 -da_grid_y 20 -da_grid_z 20\n";


#include "illuminator.h"


/* Set #if to 1 to debug, 0 otherwise */
#undef DPRINTF
#if 0
#define DPRINTF(fmt, args...) \
  ierr = PetscSynchronizedPrintf (PETSC_COMM_WORLD, "[%d] %s: " fmt, rank, \
  __FUNCT__, args); CHKERRQ (ierr); \
  ierr = PetscSynchronizedFlush (PETSC_COMM_WORLD); CHKERRQ(ierr)
#else
#define DPRINTF(fmt, args...)
#endif


#undef __FUNCT__
#define __FUNCT__ "function_3d"

/*++++++++++++++++++++++++++++++++++++++
  This is where you put the 3-D function you'd like to graph, or the 2-D
  function you'd like to graph in 3-D using the zero contour of
  +latex+$f(x,y)-z$.
  +html+ <i>f</i>(<i>x</i>,<i>y</i>)-<i>z</i>.

  PetscScalar function_3d It returns the function value.

  PetscScalar x The
  +latex+$x$-coordinate
  +html+ <i>x</i>-coordinate
  at which to calculate the function value.

  PetscScalar y The
  +latex+$y$-coordinate
  +html+ <i>y</i>-coordinate
  at which to calculate the function value.

  PetscScalar z The
  +latex+$z$-coordinate
  +html+ <i>z</i>-coordinate
  at which to calculate the function value.
  ++++++++++++++++++++++++++++++++++++++*/

static inline PetscScalar function_3d
(PetscScalar x, PetscScalar y, PetscScalar z)
{
  /* A real simple function for testing */
  /* return x+y+z; */

  /* The potential Green's function in 3-D: -1/(4 pi r), with one
     octant sliced out */
  if (x<0 || y<0 || z<0)
    return -.25/sqrt(x*x+y*y+z*z)/M_PI;
  else
    return 1000000000;

  /* The x-derivative of that Green's function: x/(4 pi r^3) */
  /* return .25*x/sqrt(x*x+y*y+z*z)/(x*x+y*y+z*z)/M_PI; */

  /* Demo of graphing z=f(x,y): the x-derivative of the 2-D Green's
     function; you need to modify the DATriangulate call below
     to plot one contour at f=0 */
  /* return x/(x*x+y*y)/2./M_PI - z; */
}


#undef __FUNCT__
#define __FUNCT__ "function_2d"

/*++++++++++++++++++++++++++++++++++++++
  This is where you put the 2-D function you'd like to graph using PETSc's
  native 2-D "contour" color graphics.

  PetscScalar function_2d It returns the function value.

  PetscScalar x The
  +latex+$x$-coordinate
  +html+ <i>x</i>-coordinate
  at which to calculate the function value.

  PetscScalar y The
  +latex+$y$-coordinate
  +html+ <i>y</i>-coordinate
  at which to calculate the function value.
  ++++++++++++++++++++++++++++++++++++++*/

static inline PetscScalar function_2d (PetscScalar x, PetscScalar y)
{
  /* The potential Green's function in 2-D: ln(r)/(2 pi) */
  return -log(1./sqrt(x*x+y*y))/2./M_PI;

  /* And its x-derivative: x/(2 pi r^2) */
  /* return x/(x*x+y*y)/2./M_PI; */
}


#undef __FUNCT__
#define __FUNCT__ "main"

/*++++++++++++++++++++++++++++++++++++++
  The usual main function.

  int main Returns 0 or error.

  int argc Number of args.

  char **argv The args.
  ++++++++++++++++++++++++++++++++++++++*/

int main (int argc, char **argv)
{
  DA            da;
  Vec           vector; /* solution vector */
  int           xstart,ystart,zstart, xwidth,ywidth,zwidth, xglobal,yglobal,
    zglobal, i,j,k, rank, ierr;
  PetscScalar   xmin=-.8,xmax=.8, ymin=-.8,ymax=.8, zmin=-.8,zmax=.8;
  PetscTruth    threedee;
  PetscViewer   theviewer;
  ISurface      Surf;
  IDisplay      Disp;

  /*+The program first calls
    +latex+{\tt PetscInitialize()}
    +html+ <tt>PetscInitialize()</tt>
    and creates the distributed arrays.  Note that even though this program
    doesn't do any communication between the CPUs, illuminator must do so in
    order to make the isoquants at CPU boundaries, so the stencil width must be
    at least one. +*/
  ierr = PetscInitialize (&argc, &argv, (char *)0, help); CHKERRQ (ierr);
  ierr = MPI_Comm_rank (PETSC_COMM_WORLD, &rank); CHKERRQ (ierr);
  DPRINTF ("Petsc initialized, moving forward\n", 0);

  /* Create DA */
  ierr = PetscOptionsHasName (PETSC_NULL, "-twodee", &threedee);
  threedee = !threedee;
  CHKERRQ (ierr);

  if (threedee)
    {
      ierr = DACreate3d
	(PETSC_COMM_WORLD, DA_NONPERIODIC, DA_STENCIL_BOX, PETSC_DECIDE,
	 PETSC_DECIDE,PETSC_DECIDE, PETSC_DECIDE,PETSC_DECIDE,PETSC_DECIDE,
	 1, 1, PETSC_NULL,PETSC_NULL,PETSC_NULL, &da); CHKERRQ (ierr);
    }
  else
    {
      ierr = DACreate2d
	(PETSC_COMM_WORLD, DA_NONPERIODIC, DA_STENCIL_BOX, PETSC_DECIDE,
	 PETSC_DECIDE, PETSC_DECIDE,PETSC_DECIDE, 1, 1, PETSC_NULL,PETSC_NULL,
	 &da); CHKERRQ (ierr);
    }

  /*+Next it gets the distributed array's local corner and global size
    information.  It gets the global vector, and loops over the part stored on
    this CPU to set all of the function values, using
    +latex+{\tt function\_3d()} or {\tt function\_2d()}
    +html+ <tt>function_3d()</tt> or <tt>function_2d()</tt>
    depending on whether the
    +latex+{\tt -twodee}
    +html+ <tt>-twodee</tt>
    command line switch was used at runtime. +*/
  ierr = DAGetCorners (da, &xstart, &ystart, &zstart, &xwidth, &ywidth,
		       &zwidth); CHKERRQ (ierr);
  ierr = DAGetInfo (da, PETSC_NULL, &xglobal, &yglobal, &zglobal, PETSC_NULL,
		    PETSC_NULL, PETSC_NULL, PETSC_NULL, PETSC_NULL, PETSC_NULL,
		    PETSC_NULL); CHKERRQ (ierr);
  if (xglobal == 1 && yglobal == 1 && zglobal == 1)
    {
      ierr = PetscPrintf (PETSC_COMM_WORLD, "Grid dimensions 1x1x1 indicate you probably want to use -da_grid_x etc.\n");
      CHKERRQ (ierr);
    }
  ierr = DAGetGlobalVector (da, &vector); CHKERRQ (ierr);

  DPRINTF ("About to calculate function values\n", 0);

  if (threedee)
    {
      PetscScalar ***array3d;

      ierr = VecGetArray3d (vector, zwidth, ywidth, xwidth, zstart, ystart,
			    xstart, &array3d); CHKERRQ (ierr);
      for (k=zstart; k<zstart+zwidth; k++)
	for (j=ystart; j<ystart+ywidth; j++)
	  for (i=xstart; i<xstart+xwidth; i++)
	    {
	      PetscScalar x,y,z;

	      x = xmin + (xmax-xmin) * (double) i/(xglobal-1);
	      y = ymin + (ymax-ymin) * (double) j/(yglobal-1);
	      z = zmin + (zmax-zmin) * (double) k/(zglobal-1);

	      array3d [k][j][i] = function_3d (x, y, z);
	    }
      ierr = VecRestoreArray3d (vector, zwidth, ywidth, xwidth, zstart, ystart,
				xstart, &array3d); CHKERRQ (ierr);
    }
  else
    {
      PetscScalar **array2d;

      ierr = VecGetArray2d (vector, ywidth, xwidth, ystart, xstart, &array2d);
      CHKERRQ (ierr);
      DASetFieldName (da, 0, "2-D Potential Green's Function");
      for (j=ystart; j<ystart+ywidth; j++)
	for (i=xstart; i<xstart+xwidth; i++)
	  {
	    double x,y;

	    x = xmin + (xmax-xmin) * (double) i/(xglobal-1);
	    y = ymin + (ymax-ymin) * (double) j/(yglobal-1);

	    array2d [j][i] = function_2d (x,y);
	  }
      ierr = VecRestoreArray2d (vector, ywidth, xwidth, ystart, xstart,
				&array2d); CHKERRQ (ierr);
    }

  /*+It then uses
    +latex+{\tt GeomviewBegin()} or {\tt PetscViewerDrawOpen()}
    +html+ <tt>GeomviewBegin()</tt> or <tt>PetscViewerDrawOpen()</tt>
    to start the viewer, and either
    +latex+{\tt DATriangulate()} and {\tt GeomviewDisplayTriangulation()} or
    +latex+{\tt VecView()}
    +html+ <tt>DATriangulate()</tt> and <tt>GeomviewDisplayTriangulation()</tt>
    +html+ or <tt>VecView()</tt>
    to display the solution. +*/

  DPRINTF ("About to open display\n", 0);

  if (threedee)
    {
      ierr = SurfCreate (&Surf); CHKERRQ (ierr);
      ierr = GeomviewBegin (PETSC_COMM_WORLD, &Disp); CHKERRQ (ierr);
    }
  else
    {
      PetscDraw draw;
      ierr = PetscViewerDrawOpen (PETSC_COMM_WORLD, 0, "", PETSC_DECIDE,
				  PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE,
				  &theviewer); CHKERRQ (ierr);
      ierr = PetscViewerDrawGetDraw (theviewer, 0, &draw);
      CHKERRQ (ierr);
      ierr = PetscDrawSetDoubleBuffer (draw); CHKERRQ (ierr);
    }

  DPRINTF ("About to do the graphing\n", 0);

  if (threedee)
    {
      PetscScalar minmax[6] = { xmin,xmax, ymin,ymax, zmin,zmax };
      PetscScalar isovals[4] = { -.1, -.2, -.3, -.4 };
      PetscScalar colors[16] = { 1,0,0,.4, 1,1,0,.4, 0,1,0,.4, 0,0,1,.4 };

      DPRINTF ("Starting triangulation\n", 0);
      ierr = DATriangulate (Surf, da, vector, 0, minmax, 4, isovals, colors,
			    PETSC_FALSE, PETSC_FALSE, PETSC_FALSE);
      CHKERRQ (ierr);

      DPRINTF ("Sending to Geomview\n", 0);
      ierr = GeomviewDisplayTriangulation
	(PETSC_COMM_WORLD, Surf, Disp, minmax, "Function-Contours",PETSC_TRUE);
      CHKERRQ (ierr);

      ierr = SurfClear (Surf); CHKERRQ (ierr);
    }
  else
    {
      ierr = VecView (vector, theviewer); CHKERRQ (ierr);
    }

  /*+ Finally, it prompts the user to hit <return> before wrapping up. +*/

  {
    char instring[100];

    ierr = PetscPrintf (PETSC_COMM_WORLD, "Press <return> to close up... ");
    CHKERRQ (ierr);
    ierr = PetscSynchronizedFGets (PETSC_COMM_WORLD, stdin, 99, instring);
    CHKERRQ (ierr);
  }

  DPRINTF ("Cleaning up\n", 0);
  if (threedee)
    {
      ierr = GeomviewEnd (PETSC_COMM_WORLD, Disp); CHKERRQ (ierr);
      ierr = ISurfDestroy (Surf); CHKERRQ (ierr);
    }
  else
    {
      ierr = PetscViewerDestroy (theviewer); CHKERRQ (ierr);
    }
  ierr = DARestoreGlobalVector (da, &vector); CHKERRQ (ierr);
  ierr = DADestroy (da); CHKERRQ (ierr);

  PetscFinalize ();
  return 0;
}