#include <iostream>
#include "Pooma/Arrays.h"
#include "MyStencils.h"

#include <vector>

static const int DefaultWorldSize = 9;
static const int DefaultNumIter   = 20;
static const double DefaultStepWidth = 0.1;

using namespace std;

Stencil<StateEquation1> f1;
Stencil<StateEquation2> f2;

void rungeKutta( std::vector< Array<2> > x, int worldSize , double h ) 
{
	std::vector< Array<2> > temp;
	std::vector< Array<2> > k1;
	std::vector< Array<2> > k2;
	std::vector< Array<2> > k3;
	std::vector< Array<2> > k4;

	Interval<1> interior_1(1, worldSize-2);
  	Interval<2> interior_2(interior_1, interior_1);
	
	for (int i = 0; i < x.size(); i++) 
	{
		
		Array<2> tempDummy(worldSize, worldSize);
		temp.push_back(tempDummy);
		temp[i] = x[i]; 
		
		Array<2> k1Dummy(worldSize, worldSize);
		Array<2> k2Dummy(worldSize, worldSize);
		Array<2> k3Dummy(worldSize, worldSize);
		Array<2> k4Dummy(worldSize, worldSize);
		
		k1.push_back(k1Dummy);
		k2.push_back(k2Dummy);
		k3.push_back(k3Dummy);
		k4.push_back(k4Dummy);

	
	}
	
	// k1 
	for (int i = 0; i < x.size(); i++) 
	{
		Pooma::blockAndEvaluate();
		k1[i](interior_2) = f1(x[i]); // f[i](x);
		Pooma::blockAndEvaluate();
		temp[i](interior_2) = x[i](interior_2) + 0.5 * h * k1[i](interior_2);
	}
	
	// k2
	for (int i = 0; i < x.size(); i++) 
	{
		Pooma::blockAndEvaluate();
		k2[i](interior_2) =  f1(temp[i])  ; // f[i](x);
		Pooma::blockAndEvaluate();
		temp[i](interior_2) = x[i](interior_2) + 0.5 * h * k2[i](interior_2);
	}
	
	// k3
	for (int i = 0; i < x.size(); i++) 
	{
		Pooma::blockAndEvaluate();
		k3[i](interior_2) =  f1(temp[i])  ; // f[i](x);
		Pooma::blockAndEvaluate();
		temp[i](interior_2) = x[i](interior_2) + h * k3[i](interior_2);
	}
	
	// k4
	for (int i = 0; i < x.size(); i++) 
	{
		Pooma::blockAndEvaluate();
		k4[i](interior_2) =  f1(temp[i]); // f[i](x);
		Pooma::blockAndEvaluate();
		temp[i](interior_2) = x[i](interior_2) + h / 6.0 * (k1[i](interior_2) + 2 * k2[i](interior_2) + 2* k3[i](interior_2) + k4[i](interior_2));
	}


	for (int i = 0; i < x.size(); i++) 
	{

		Pooma::blockAndEvaluate();
		// Euler: temp[i](interior_2) = f1(x[i]) * h + x[i](interior_2) ; // f[i](x);
		
		x[i] = temp[i];
		std::cout << "x[" << i << "]: " << std::endl;
	    std::cout << x[i] << std::endl << std::endl;

	}
	
	std::cout << "-------------------------------" << std::endl;
	
}

int main(int argc, char* argv[])
{
  Pooma::initialize(argc,argv);

  int    worldSize = DefaultWorldSize;
  int    numIter   = DefaultNumIter;
  double h         = DefaultStepWidth;

  // List of x
  std::vector< Array<2> > x;
  
  // x1
  Array<2> x1(worldSize, worldSize);
  x1 = 0.00;
  Pooma::blockAndEvaluate();
  x1(worldSize/2, worldSize/2) = 2.0 * numIter;
  std::cout << "start x1" << std::endl << x1 << std::endl << std::endl;

  // x2
  Array<2> x2(worldSize, worldSize);
  x2 = 0.00;
  x2(worldSize/2, worldSize/2) = - 2.0 * numIter;
  Pooma::blockAndEvaluate();
  std::cout << "start x2" << std::endl << x2 << std::endl << std::endl;

  x.push_back(x1);
  x.push_back(x2);

  for (int i=0; i<numIter; i+=1) 
  {
		rungeKutta(x, worldSize, h);	
  }
  
  Pooma::finalize();
  return 0;
}