#include <iostream>		// has std::cout, ...
#include <stdlib.h>		// has EXIT_SUCCESS
#include "Pooma/Arrays.h"	// has Pooma's Array

// Doof2d: Pooma Arrays, element-wise implementation

int main(int argc, char *argv[])
{
  // Prepare the Pooma library for execution.
  Pooma::initialize(argc,argv);
  
  // Ask the user for the number of averagings.
  long nuAveragings, nuIterations;
  std::cout << "Please enter the number of averagings: ";
  std::cin >> nuAveragings;
  nuIterations = (nuAveragings+1)/2; // Each iteration performs two averagings.

  // Ask the user for the number n of elements along one dimension of
  // the grid.
  long n;
  std::cout << "Please enter the array size: ";
  std::cin >> n;

  // Specify the arrays' domains [0,n) x [0,n).
  Interval<1> N(0, n-1);
  Interval<2> vertDomain(N, N);

  // Create the arrays.
  // The template parameters indicate 2 dimensions, a 'double' element
  // type, and ordinary 'Brick' storage.
  Array<2, double, Brick> a(vertDomain);
  Array<2, double, Brick> b(vertDomain);

  // Set up the initial conditions.
  // All grid values should be zero except for the central value.
  a = b = 0.0;
  b(n/2,n/2) = 1000.0;

  // In the average, weight element with this value.
  const double weight = 1.0/9.0;

  // Perform the simulation.
  for (int k = 0; k < nuIterations; ++k) {
    // Read from b.  Write to a.
    for (int j = 1; j < n-1; j++)
      for (int i = 1; i < n-1; i++)
        a(i,j) = weight *
          (b(i+1,j+1) + b(i+1,j  ) + b(i+1,j-1) +
           b(i  ,j+1) + b(i  ,j  ) + b(i  ,j-1) +
           b(i-1,j+1) + b(i-1,j  ) + b(i-1,j-1));

    // Read from a.  Write to b.
    for (int j = 1; j < n-1; j++)
      for (int i = 1; i < n-1; i++)
        b(i,j) = weight *
          (a(i+1,j+1) + a(i+1,j  ) + a(i+1,j-1) +
           a(i  ,j+1) + a(i  ,j  ) + a(i  ,j-1) +
           a(i-1,j+1) + a(i-1,j  ) + a(i-1,j-1));
  }

  // Print out the final central value.
  std::cout << "before: " << (nuAveragings % 2 ? a(n/2,n/2) : b(n/2,n/2)) << std::endl; // TMP
  Pooma::blockAndEvaluate();	// Ensure all computation has finished.
  std::cout << (nuAveragings % 2 ? a(n/2,n/2) : b(n/2,n/2)) << std::endl;

  // The arrays are automatically deallocated.

  // Tell the Pooma library execution has finished.
  Pooma::finalize();
  return EXIT_SUCCESS;
}