/*
   Solves Poisson's equation with Dirichlet BC on a cubical grid of arbitrary size
   Processors: n
*/

#define MAIN_CPP

#include <mpi.h>
#include <iostream>
#include <stack>
#include <queue>

#include "petscsolver.h"

/* constants that would normally be imported from input card */
const    int   g_Ni =  100; //number of nodes per dimension
const double      h = 1E-4; //distance between grid points
const double reltol = 1E-5; //rtol in KSPSetTolerances

int main(int argc, char **argv)
{
  MPI_Init(&argc,&argv);
  
  int size,rank;
  MPI_Comm_size(MPI_COMM_WORLD,&size);
  MPI_Comm_rank(MPI_COMM_WORLD,&rank);

	if (rank==0) {
	  for (int i=0; i<argc; i++)
	    std::cout<<argv[i]<<' ';
	  std::cout<<'\n'<<std::flush;
  }

  /* Acquire the most evenly split factorization for domain decomposition */
  std::stack<int> primes;
	std::priority_queue<int> factors;
	int product = size;

	for (int i = 2; i <= product; i++){
		while (product%i==0) {
			primes.push(i);
			product /= i;
		}
	}

	while (primes.size()<3) {
		primes.push(1);
	}

	for (int i=0; i<3; i++) {
		factors.push(primes.top());
		primes.pop();
	}

	while (!primes.empty()) {
		int i = factors.top();
		factors.pop();
		i *= primes.top();
		primes.pop();
		factors.push(i);
	}

  int dims[3]; //process partition in each direction, as passed to DMDACreate3d
  for (int i = 0; i < 3; i++) {
    dims[i] = factors.top();
	  factors.pop();
  }

  if (rank==0) std::cout<<'\n'<<dims[0]<<'x'<<dims[1]<<'x'<<dims[2]<<" subdomains\n\n"<<std::flush; 


  /* Get subdomain position of own rank in global grid */
  /* Note that this is also the way PETSc partitions the global grid, otherwise this wouldn't work */
	int rankpos[3];
	rankpos[0] = rank % dims[0];
	rankpos[1] = (rank / dims[0]) % dims[1];
	rankpos[2] = (rank / (dims[0]*dims[1])) % dims[2];

	int l_Ni[dims[0]], l_Nj[dims[1]], l_Nk[dims[2]]; //number of local nodes along the x,y and z axis for each cell,
	                                                 //as passed to DMDACreate3d

	for (int i=0; i<dims[0]; i++)	l_Ni[i]=0;
	for (int i=0; i<dims[1]; i++)	l_Nj[i]=0;
	for (int i=0; i<dims[2]; i++)	l_Nk[i]=0;

	int j=0;
	do l_Ni[j%dims[0]]++; while (++j<g_Ni);

	int imin = 0;
	for (int i=0; i<rankpos[0]; i++)
		imin += l_Ni[i];
	int imax = imin+l_Ni[rankpos[0]]-1;

	j=0;
	do l_Nj[j%dims[1]]++; while (++j<g_Ni);

	int jmin = 0;
	for (int i=0; i<rankpos[1]; i++)
		jmin += l_Nj[i];
	int jmax = jmin+l_Nj[rankpos[1]]-1;

	j=0;
	do l_Nk[j%dims[2]]++; while (++j<g_Ni);

	int kmin = 0;
	for (int i=0; i<rankpos[2]; i++)
		kmin += l_Nk[i];
	int kmax = kmin+l_Nk[rankpos[2]]-1;

	int Ni = l_Ni[rankpos[0]];
	int Nj = l_Nj[rankpos[1]];
	int Nk = l_Nk[rankpos[2]];
  int Nall = Ni*Nj*Nk;
  
  MPI_Barrier(MPI_COMM_WORLD);
  std::cout<<"Rank "<<rank<<": Ni="<<Ni<<", Nj = "<<Nj<<", Nk = "<<Nk<<'\n'<<std::flush;
  MPI_Barrier(MPI_COMM_WORLD);
  if (rank==0) std::cout<<std::endl<<std::flush; 
  
  MPI_Barrier(MPI_COMM_WORLD);
  std::cout<<"Rank "<<rank<<": i="<<imin<<".."<<imax<<", j="<<jmin<<".."<<jmax<<", k="<<kmin<<".."<<kmax<<'\n'<<std::flush;
  MPI_Barrier(MPI_COMM_WORLD);
  if (rank==0) std::cout<<std::endl<<std::flush; 

  /*
     Assemble local arrays that need to be passed to the solver.
     Note that these can't be PETSc vectors yet, as they are normally accessed from other part of the program as well.
     The triplet (i,j,k) in the local grid is addressed via element i+j*Ni+k*Ni*Nj.
  */
  int *bound; //used to determine whether a node is a Dirichlet BC (1) or not (0)
  bound = new int[Nall];
  double *phi, *charge; //hold electric potential and charge
  phi = new double[Nall];
  charge = new double[Nall];

  /* 
     Define a hollow cylinder of constant potential
     Charge is set to 0, but can be varied
  */
  int center = (int) ceil(((double) g_Ni)/2. - 1.);
  for (int k=kmin; k<=kmax; k++) {
    for (int j=jmin; j<=jmax; j++) {
      for (int i=imin; i<=imax; i++) {
        int idx = i-imin+(j-jmin)*Ni+(k-kmin)*Ni*Nj;
	      
	      int R = (int) ceil(((double) g_Ni)/2-1.);

	      if ((i-center)*(i-center)+(j-center)*(j-center)>=R*R || k==0 || k==g_Ni-1) {
		      phi[idx] = 100.;
		      charge[idx] = 0.;
		      bound[idx] = 1;
	      } else {
		      phi[idx] = 0.;
		      charge[idx] = 0.;
		      bound[idx] = 0;
	      }
      }
    }
  }
  
  
  /* Create and initialize solver */
	PETScSolver *petscsolver;
	petscsolver = new PETScSolver(argc,argv,g_Ni,g_Ni,g_Ni,dims,l_Ni,l_Nj,l_Nk,Nall,h);
	petscsolver->init(bound); //initializes matrix
	
	/* Solve */
	petscsolver->initsolve(phi,charge,reltol);
	
	/*
	   Solve without setting options.
	   Note that this is normally repeatedly called with altered charge array and phi array of previous timestep as initial guess
	*/
	petscsolver->solve(0,phi,charge,reltol);
	petscsolver->solve(1,phi,charge,reltol);
	
	
	delete petscsolver;

  delete[] bound;
  delete[] phi;
  delete[] charge;


	MPI_Barrier(MPI_COMM_WORLD);

	if (rank == 0) std::cout<<".. done"<<std::endl;

	MPI::Finalize();

	return 0;
}