// Oldham, Jeffrey D.
// 2001Mar27
// Pooma

// Hydrodynamics Kernel Written Using POOMA

#include <iostream>
#include <stdlib.h>
#include <cmath>
#include "Pooma/NewFields.h"
#include "ComputeVolumes.h"
#include "Product.h"
#include "ConstantVectorComponentBC.h"

// This hydrodynamics program implements "The Construction of
// Compatible Hydrodynamics Algorithms Utilizing Conservation of Total
// Energy," by E.J. Caramana, D.E. Burton, M.J. Shashkov, and
// P.P. Whalen, \emph{Journal of Computational Physics}, vol. 146
// (1998), pp. 227-262.

// The code uses fields with two different granularities: coarse and
// fine.  Both grid conceptually cover the same domain, but the
// fine-grain one has twice the number of values for each dimension.
// Currently, Pooma does not nicely support operations between fields
// with different granularities.  The code occurring between
// `CORNER_FIELD_TEMPORARY' preprocessor values substitutes for this
// lack of support.  See the separately distributed proposal for
// details.

// Preprocessor symbols:
// CORNER_FIELD_TEMPORARY: Define this symbol.  Include code to deal
// 			   with different granularity fields.
// PSEUDOCODE: Do not define this symbol.  Surrounds desired code to
//	       deal with different granularity fields.
// DEBUG: If defined, print some information about internal program
//	  values.
// DEBUG2: If defined, print much more information about internal
//	   program values.

// Sample compilation command:
// g++ -g -DCORNER_FIELD_TEMPORARY -I/nfs/oz/home/oldham/pooma/r2/src/ -I/nfs/oz/home/oldham/pooma/r2/lib/LINUXgcc/ -L. hydrodynamics.cc -lpooma-gcc -o hydrodynamics

#ifdef CORNER_FIELD_TEMPORARY

/** FORWARD DECLARATIONS **/

template<class InputGeometry, class OutputGeometry, class T, class EngineTag>
inline
void
sumAroundCell (const Field<InputGeometry, T, EngineTag> & input,
	       Field<OutputGeometry, T, EngineTag> & output);

template<class InputGeometry, class OutputGeometry, class T, class EngineTag>
inline
void
sumAroundVertex(const Field<InputGeometry, T, EngineTag> & input,
		Field<OutputGeometry, T, EngineTag> & output);

template<class InputGeometry1,
         class InputT1, class InputT2,
         class InputEngineTag1, class InputEngineTag2,
         class OutputGeometry, class OutputT, class OutputEngineTag>
inline
void
computeCornerMasses(const Field<InputGeometry1, InputT1, InputEngineTag1> & pressureDensity,
		    const Field<OutputGeometry, InputT2, InputEngineTag2> & fineCoordinates,
		    Field<OutputGeometry, OutputT, OutputEngineTag> & output);

template <class InputGeometry1,
          class InputT1, class InputT2,
          class InputEngineTag1, class InputEngineTag2,
          class OutputGeometry, class OutputT, class OutputEngineTag>
inline
void
computeCornerForces(const Field<InputGeometry1, InputT1, InputEngineTag1> & pressure,
		    const Field<OutputGeometry, InputT2, InputEngineTag2> & coordinates,
		    Field<OutputGeometry, OutputT, OutputEngineTag> & output);

template <class InputGeometry1,
          class InputT1,
          class InputEngineTag1,
          class OutputGeometry, class OutputT, class OutputEngineTag>
inline
void
copyVelocities(const Field<InputGeometry1, InputT1, InputEngineTag1> & input1,
	       const Field<InputGeometry1, InputT1, InputEngineTag1> & input2,
	       Field<OutputGeometry, OutputT, OutputEngineTag> & output);
#endif // CORNER_FIELD_TEMPORARY



/** THE HYDRODYNAMICS PROGRAM **/

int main(int argc, char *argv[])
{
  // Set up the Pooma library.
  Pooma::initialize(argc,argv);
#ifdef DEBUG
  std::cout << "Hello, world." << std::endl;
#endif // DEBUG

  /* DECLARATIONS */

  // Values
  const double gamma = 4.0/3;		// gas pressure constant $\gamma$
  const double timeStep = 0.01;		// $\Delta t$
  unsigned nuIterations = 1;		// number of iterations

  // Create a simple layout.
  const unsigned Dim = 2;		// Work in a 2D world.
  const unsigned nHorizVerts = 11;	// number of horizontal vertices
  const unsigned nAngles = 5;		// number of angles
  Interval<Dim> vertexDomain;
  vertexDomain[0] = Interval<1>(nHorizVerts);
  vertexDomain[1] = Interval<1>(nAngles);
  DomainLayout<Dim> layout(vertexDomain, GuardLayers<2>(1));

  // Preparation for Field creation.
  Vector<Dim> origin(0.0);
  Vector<Dim> spacings(1.0,1.0);  // TODO: What does this do?
  typedef UniformRectilinear<Dim, double, Cartesian<Dim> > Geometry_t;
  typedef Field<Geometry_t, double,      Brick> Fields_t;
  typedef Field<Geometry_t, double,      Brick> ConstFields_t; // TODO: Change to ConstField when ConstField is available.
  typedef Field<Geometry_t, Vector<Dim>, Brick> Fieldv_t;
  typedef Field<Geometry_t, Vector<Dim>, Brick> ConstFieldv_t; // TODO: Change to ConstField when ConstField is available.

  // Cell-centered Fields.
  Cell cell;
  Fields_t internalEnergy  (cell, layout, origin, spacings);
  ConstFields_t zoneMass   (cell, layout, origin, spacings);
  Fields_t pressure	   (cell, layout, origin, spacings);
  Fields_t pressureDensity (cell, layout, origin, spacings);
  Fields_t volume	   (cell, layout, origin, spacings);

  // Vertex-centered Fields.
  Vert vert;
  ConstFields_t pointMass  (vert, layout, origin, spacings);
  Fieldv_t coordinates	   (vert, layout, origin, spacings);
  Fieldv_t velocity	   (vert, layout, origin, spacings);
  Fieldv_t velocityChange  (vert, layout, origin, spacings);

  // Corner Fields.
  // TODO: How should I implement corner Fields?
#ifdef CORNER_FIELD_TEMPORARY
  // For now, implement corner Fields using a Field with twice as many
  // entries in each dimension.
  Interval<Dim> cornerVertexDomain;
  for (unsigned d = 0; d < Dim; ++d)
    cornerVertexDomain[d] = Interval<1>(2*vertexDomain[d].min(),2*vertexDomain[d].max());
  DomainLayout<Dim> cornerLayout(cornerVertexDomain, GuardLayers<2>(1));
  Fieldv_t cornerForces	   (cell, cornerLayout, origin, spacings);
  Fields_t cornerMasses    (cell, cornerLayout, origin, spacings);
  Fieldv_t cornerCoordinates(vert, cornerLayout, origin, spacings);
  Fieldv_t cornerVelocities(vert, cornerLayout, origin, spacings);
  Fields_t tmp1		   (cell, cornerLayout, origin, spacings);
  Fields_t tmp2		   (cell, cornerLayout, origin, spacings);
#endif // CORNER_FIELD_TEMPORARY

  /* INITIALIZATION */

  // TODO: Is coordinates initialization necessary?  What is the best way to do this?
  for (unsigned xIndex = 0; xIndex <= vertexDomain[0].last (); ++xIndex)
    for (unsigned yIndex = 0; yIndex <= vertexDomain[1].last (); ++yIndex)
      coordinates(xIndex,yIndex) =
	Vector<2>(static_cast<double>(xIndex)/(nHorizVerts-1) * cos(yIndex*M_PI/(2*(nAngles-1))),
		  static_cast<double>(xIndex)/(nHorizVerts-1) * sin(yIndex*M_PI/(2*(nAngles-1))));
  for (unsigned xIndex = 0; xIndex <= cornerVertexDomain[0].last (); ++xIndex)
    for (unsigned yIndex = 0; yIndex <= cornerVertexDomain[1].last (); ++yIndex)
      cornerCoordinates(xIndex,yIndex) =
	Vector<2>(0.5*xIndex/(nHorizVerts-1) * cos(0.5*yIndex*M_PI/(2*(nAngles-1))),
		  0.5*xIndex/(nHorizVerts-1) * sin(0.5*yIndex*M_PI/(2*(nAngles-1))));
#ifdef DEBUG
  std::cout << "initial coordinates:\n" << coordinates << std::endl;
#endif // DEBUG
  internalEnergy = 1.0;
  pressureDensity = 1.0;
#ifdef PSEUDOCODE
  // This is the desired code:
  cornerMasses = replicate<2>(pressureDensity) *
    computeVolumes(interpolate<2>(coordinates));
  pointMass = total<2>(cornerMasses);
  zoneMass = total<2>(cornerMasses);
  // This is the code that Pooma currently supports:
#endif // PSEUDOCODE
#ifdef CORNER_FIELD_TEMPORARY
  computeCornerMasses(pressureDensity, cornerCoordinates, cornerMasses);
  sumAroundVertex(cornerMasses, pointMass);
  sumAroundCell(cornerMasses, zoneMass);
#endif // CORNER_FIELD_TEMPORARY
  PInsist(min(pointMass) > 0.0, "Zero masses are not supported.");
#ifdef DEBUG2
  std::cout << "pointMass:\n" << pointMass << std::endl;
  std::cout << "zoneMass:\n" << zoneMass << std::endl;
#endif // DEBUG2
  velocity = Vector<Dim>(0.0);
  velocityChange.addUpdaters(AllConstantFaceBC<Vector<Dim> >(Vector<Dim>(0.0), false));
  for (int d = 0; d < Dim; ++d)
    velocityChange.addUpdater(ConstantVectorComponentBC<Vector<Dim>::Element_t>(2*d, d, 0.0, true));
  velocityChange.addUpdater(ConstantVectorComponentBC<Vector<Dim>::Element_t>(3, 0, 0.0, true));

  /* ITERATIONS */

  for (; nuIterations > 0; --nuIterations) {
#ifdef DEBUG
    std::cout << "Begin an iteration." << std::endl;
#endif // DEBUG
    pressure = (gamma - 1.0) * pressureDensity * internalEnergy;
#ifdef DEBUG2
    std::cout << "pressure:\n" << pressure << std::endl;
#endif // DEBUG2
#ifdef PSEUDOCODE
    // This is the desired code.
    forces = replicate<2>(pressure) * addNormals(coordinates);
    velocityChange = (timeStep / pointMass) * total<2>(cornerForces);
    // This is the code that Pooma currently supports:
#endif // PSEUDOCODE
#ifdef CORNER_FIELD_TEMPORARY
    computeCornerForces(pressure, coordinates, cornerForces);
    sumAroundVertex(cornerForces, velocityChange);
    velocityChange *= timeStep / pointMass;
#endif // CORNER_FIELD_TEMPORARY
#ifdef DEBUG2
    std::cout << "corner forces:\n" << cornerForces << std::endl;
    std::cout << "velocity changes:\n" << velocityChange << std::endl;
#endif // DEBUG2
    velocityChange.update();
#ifdef PSEUDOCODE
    // This is the desired code.
    internalEnergy +=
      -(timeStep / pointMass) *
      total<2>(dot(cornerForces, replicate<2>(velocity + 0.5*velocityChange)));
    // This is the code that Pooma currently supports:
#endif // PSEUDOCODE
#ifdef CORNER_FIELD_TEMPORARY
    copyVelocities(velocity, velocityChange, cornerVelocities);
    tmp1 = dot(cornerForces, cornerVelocities);
    sumAroundCell(tmp1, tmp2);
    internalEnergy += -(timeStep / pointMass) * tmp2;
#endif // CORNER_FIELD_TEMPORARY
    coordinates += (velocity + 0.5*velocityChange) * timeStep;
    velocity += velocityChange;
    volume = computeVolumes(coordinates);
    pressureDensity = zoneMass / volume;
#ifdef DEBUG2
    std::cout << "pressure density:\n" << pressureDensity << std::endl;
#endif // DEBUG2
  }

  /* TERMINATION */

  std::cout << "final coordinates:\n" << coordinates << std::endl;
#ifdef DEBUG
  std::cout << "Goodbye, world." << std::endl;
#endif // DEBUG
  Pooma::finalize();
  return EXIT_SUCCESS;
}


/** HELPER FUNCTIONS **/

#ifdef CORNER_FIELD_TEMPORARY
// Following is temporary code to permit compilation using the current
// Pooma code.  Adding new Pooma abstractions may eliminate the need
// for this code.
static Loc<2> offset[] = { Loc<2>(0,0), Loc<2>(0,1),
			     Loc<2>(1,0), Loc<2>(1,1) };

// Given a Field with twice the refinement, form a cell-centered Field
// by summing the 1<<Dim corresponding corners.
template<class InputGeometry, class OutputGeometry, class T, class EngineTag>
inline
void
sumAroundCell(const Field<InputGeometry, T, EngineTag> & input,
	      Field<OutputGeometry, T, EngineTag> & output)
{
  const int dim = Field<OutputGeometry, T, EngineTag>::dimensions;
  Field<OutputGeometry,T,EngineTag>::Domain_t range = output.physicalDomain();
  output(range) = 0;
  for (unsigned counter = 0; counter < (1<<dim); ++counter)
    output(range) += input(2*range+offset[counter]);
  return;
}


// Given a Field with twice the refinement, form a vertex-centered Field
// by summing the 1<<Dim corresponding corners.
template<class InputGeometry, class OutputGeometry, class T, class EngineTag>
inline
void
sumAroundVertex(const Field<InputGeometry, T, EngineTag> & input,
		Field<OutputGeometry, T, EngineTag> & output)
{
  const int dim = Field<OutputGeometry, T, EngineTag>::dimensions;
  Field<OutputGeometry,T,EngineTag>::Domain_t range = output.physicalDomain();
  output(range) = 0;
  for (unsigned counter = 0; counter < (1<<dim); ++counter)
    output(range) += input(2*range-offset[counter]);
  return;
}


// Given a quadrilateral's two diagonals as vectors, return the
// quadrilateral's volume.
template <int Dim, class T, class E>
inline
double
quadrilateralVolume (const Vector<Dim,T,E> &v, const Vector<Dim,T,E> &w)
{
  return 0.5 * std::abs (sum (product (v, w)));
}


// Given a pressure density Field and a fine-mesh coordinate Field, compute
// corner masses, placing in "output".
// TODO: Extend to multiple dimensions.
template<class InputGeometry1,
         class InputT1, class InputT2,
         class InputEngineTag1, class InputEngineTag2,
         class OutputGeometry, class OutputT, class OutputEngineTag>
inline
void
computeCornerMasses(const Field<InputGeometry1, InputT1, InputEngineTag1> & pressureDensity,
		    const Field<OutputGeometry, InputT2, InputEngineTag2> & fineCoordinates,
		    Field<OutputGeometry, OutputT, OutputEngineTag> & output)
{
  // Assume the temporary finer coordinates Field fineCoordinates
  // already has values.

  // Compute the volumes in fineCoordinates.
  output = computeVolumes(fineCoordinates);

  // Multiply each volume by the pressure density.
  Field<InputGeometry1,InputT1,InputEngineTag1>::Domain_t range = pressureDensity.physicalDomain();
  const unsigned corners =
    (1 << (Field<InputGeometry1, InputT1, InputEngineTag1>::dimensions));;
  for (unsigned c = 0; c < corners; ++c)
    output(2*range+offset[c]) *= pressureDensity(range);

  return;
}


// Given a pressure Field and a coordinate Field, compute the corner
// forces in a fine-mesh Field.
template <class InputGeometry1,
          class InputT1, class InputT2,
          class InputEngineTag1, class InputEngineTag2,
          class OutputGeometry, class OutputT, class OutputEngineTag>
inline
void
computeCornerForces(const Field<InputGeometry1, InputT1, InputEngineTag1> & pressure,
		    const Field<OutputGeometry, InputT2, InputEngineTag2> & coordinates,
		    Field<OutputGeometry, OutputT, OutputEngineTag> & output)
{
  Field<InputGeometry1,InputT1,InputEngineTag1>::Domain_t range = pressure.physicalDomain();
  for (unsigned xIndex = 0; xIndex <= range.unwrap()[0].last(); ++xIndex)
    for (unsigned yIndex = 0; yIndex <= range.unwrap()[1].last(); ++yIndex) {
      Loc<2> range (xIndex, yIndex);
      output(2*range + offset[0]) =
	pressure(range) *
	0.5 * product(InputT2(1.0),
		      coordinates(range+offset[2])
		      - coordinates(range+offset[1]));
      output(2*range+offset[1]) =
	pressure(range) *
	0.5 * product(InputT2(1.0),
		      coordinates(range+offset[0])
		      - coordinates(range+offset[3]));
      output(2*range+offset[2]) =
	pressure(range) *
	0.5 * product(InputT2(1.0),
		      coordinates(range+offset[3])
		      - coordinates(range+offset[0]));
      output(2*range+offset[3]) =
	pressure(range) *
	0.5 * product(InputT2(1.0),
		      coordinates(range+offset[1])
		      - coordinates(range+offset[2]));
    }
  return;

#ifdef PSEUDOCODE
  // TODO: I really wanted to use "range" in a data-parallel
  // statement, but I do not know how to extend "product" to work
  // on Fields for such a statement.
  Field<InputGeometry1,InputT1,InputEngineTag1>::Domain_t range = pressure.physicalDomain();
  output(2*range+offset[0]) =
    pressure(range) *
    0.5 * product(InputT2(1.0),
		  coordinates(range+offset[2])
		  - coordinates(range+offset[1]));
  output(2*range+offset[1]) =
    pressure(range) *
    0.5 * product(InputT2(1.0),
		  coordinates(range+offset[0])
		  - coordinates(range+offset[3]));
  output(2*range+offset[2]) =
    pressure(range) *
    0.5 * product(InputT2(1.0),
		  coordinates(range+offset[3])
		  - coordinates(range+offset[0]));
  output(2*range+offset[3]) =
    pressure(range) *
    0.5 * product(InputT2(1.0),
		  coordinates(range+offset[1])
		  - coordinates(range+offset[2]));
#endif // PSEUDOCODE
}


template <class InputGeometry, class InputT, class InputEngineTag,
          class OutputGeometry, class OutputT, class OutputEngineTag>
inline
void
copyVelocities(const Field<InputGeometry, InputT, InputEngineTag> & input1,
	       const Field<InputGeometry, InputT, InputEngineTag> & input2,
	       Field<OutputGeometry, OutputT, OutputEngineTag> & output)
{
  Field<InputGeometry,InputT,InputEngineTag>::Domain_t range =
    input1.physicalDomain();
  output(2*range-offset[0]) =
  output(2*range-offset[1]) =
  output(2*range-offset[2]) =
  output(2*range-offset[3]) = input1(range) + 0.5*input2(range);
  return;
}

#endif // CORNER_FIELD_TEMPORARY