/*
** (c) 1996-2000 The Regents of the University of California (through
** E.O. Lawrence Berkeley National Laboratory), subject to approval by
** the U.S. Department of Energy.  Your use of this software is under
** license -- the license agreement is attached and included in the
** directory as license.txt or you may contact Berkeley Lab's Technology
** Transfer Department at TTD@lbl.gov.  NOTICE OF U.S. GOVERNMENT RIGHTS.
** The Software was developed under funding from the U.S. Government
** which consequently retains certain rights as follows: the
** U.S. Government has been granted for itself and others acting on its
** behalf a paid-up, nonexclusive, irrevocable, worldwide license in the
** Software to reproduce, prepare derivative works, and perform publicly
** and display publicly.  Beginning five (5) years after the date
** permission to assert copyright is obtained from the U.S. Department of
** Energy, and subject to any subsequent five (5) year renewals, the
** U.S. Government is granted for itself and others acting on its behalf
** a paid-up, nonexclusive, irrevocable, worldwide license in the
** Software to reproduce, prepare derivative works, distribute copies to
** the public, perform publicly and display publicly, and to permit
** others to do so.
*/


#include <algorithm>
#include <iostream>

#include <IntVect.H>
#include <Box.H>
#include <MultiFab.H>
#include <multigrid.H>
#include <macproj.H>
#include <globals.H>
#include <MACPROJ_F.H>
#include <ParmParse.H>
#include <names.H>

int mac_projector::numSmoothCoarsen = 2;
int mac_projector::numSmoothRefine  = 2;
int mac_projector::ng               = 4;

// ************************************************************************
// ** constructor **
// ************************************************************************

mac_projector::mac_projector(const BoxArray& Grids, const Geometry& Geom,
                             MultiFab* Area, MultiFab& Vol):
       grids(Grids), geom(Geom), area(Area), vol(Vol)
{
  ParmParse pp("mac");
  pp.get("tol",tol);

  pp.query("numSmoothCoarsen",numSmoothCoarsen);
  pp.query("numSmoothRefine",numSmoothRefine);

  ng = 2*std::max(numSmoothCoarsen,numSmoothRefine);
}

// ************************************************************************
// ** destructor **
// ************************************************************************

mac_projector::~mac_projector()
{
}

// ************************************************************************
// ** project **
// ************************************************************************

void mac_projector::project(MultiFab * uadv, MultiFab * vadv, 
#if (BL_SPACEDIM == 3)
                            MultiFab * wadv,
#endif
                            MultiFab * state,  MultiFab * divu_src)

{
  const Real* dx = geom.CellSize();

  MultiFab* source = new MultiFab(grids,1,ng-1);
  MultiFab* resid = new MultiFab(grids,1,ng-1);

  MultiFab* phi = new MultiFab(grids,1,ng);
  phi->setVal(0.0);

  MultiFab* sigma[BL_SPACEDIM];
  for (int n = 0; n < BL_SPACEDIM; n++) 
  {
    BoxArray edge_grids(grids);
    edge_grids.surroundingNodes(n);
    sigma[n] = new MultiFab(edge_grids,1,ng-1);
    sigma[n]->setVal(0.);
  }

  Real rhsnorm = -1.e20;
  Real norm;

  for (MFIter mfi(*state); mfi.isValid(); ++mfi)
  {
     int i = mfi.index();
     const int* lo = grids[i].loVect();
     const int* hi = grids[i].hiVect();

     FORT_INITSIGMA((*sigma[0])[mfi].dataPtr(),(*sigma[1])[mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                    (*sigma[2])[mfi].dataPtr(),
#endif
                    (*state)[mfi].dataPtr(Density),
                    area[0][mfi].dataPtr(),area[1][mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                    area[2][mfi].dataPtr(),
#endif
   		    ARLIM(lo),ARLIM(hi),bc[i].dataPtr(),&ng);

     // NOTE: The advection velocities are multiplied by area in this routine.
     FORT_RHSMAC((*uadv)[mfi].dataPtr(),(*vadv)[mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                 (*wadv)[mfi].dataPtr(),
#endif
                 (*divu_src)[mfi].dataPtr(),(*source)[mfi].dataPtr(),
                 area[0][mfi].dataPtr(),area[1][mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                 area[2][mfi].dataPtr(),
#endif
                 vol[mfi].dataPtr(),
                 ARLIM(lo),ARLIM(hi),&norm,&ng);
     rhsnorm = std::max(rhsnorm, norm);
  }
  ParallelDescriptor::ReduceRealMax(rhsnorm);

  source->FillBoundary(0,1);
  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*source,0,1,true);

  for (int n = 0; n < BL_SPACEDIM; n++)
  {
    sigma[n]->FillBoundary(0,1);
    geom.FillPeriodicBoundary(*sigma[n],0,1,true);
  }

  if (ParallelDescriptor::IOProcessor()) 
    std::cout << std::endl << "MAC Source norm: " << rhsnorm << std::endl;

  if (rhsnorm > 1.e-16) {

    Real goal = tol*rhsnorm;

    macprojection_mg *solver =
      new macprojection_mg(grids,geom,phi,source,resid,sigma,dx);

    solver->solve(goal,rhsnorm,numSmoothCoarsen,numSmoothRefine);
    delete solver;

    phi->FillBoundary(0,1);
    if (geom.isAnyPeriodic())
      geom.FillPeriodicBoundary(*phi,0,1,true);

    const int is_rz = CoordSys::IsRZ();

    for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

     // NOTE: The advection velocities are divided by area in this routine.
#if (BL_SPACEDIM == 2)
      FORT_PROJUMAC((*uadv)[mfi].dataPtr(), (*vadv)[mfi].dataPtr(),
                    (*sigma[0])[mfi].dataPtr(),(*sigma[1])[mfi].dataPtr(),
                    (*phi)[mfi].dataPtr(),
                    area[0][mfi].dataPtr(),area[1][mfi].dataPtr(),
                    dx,ARLIM(lo), ARLIM(hi),bc[i].dataPtr(),&is_rz,&ng);
#elif (BL_SPACEDIM == 3)
      FORT_PROJUMAC((*uadv)[mfi].dataPtr(), (*vadv)[mfi].dataPtr(), (*wadv)[mfi].dataPtr(),
                    (*sigma[0])[mfi].dataPtr(),(*sigma[1])[mfi].dataPtr(), (*sigma[2])[mfi].dataPtr(),
                    (*phi)[mfi].dataPtr(),
                    area[0][mfi].dataPtr(),area[1][mfi].dataPtr(),area[2][mfi].dataPtr(),
                    dx,ARLIM(lo), ARLIM(hi),bc[i].dataPtr(),&ng);
#endif
    }
  }
  delete phi;
  delete source;
  delete resid;
  for (int n = 0; n < BL_SPACEDIM; n++) 
    delete sigma[n];
}

// ************************************************************************
// ** constructor **
// ************************************************************************

macprojection_mg::macprojection_mg(const BoxArray& Grids,
                                   const Geometry& Geom,
                                   MultiFab* Phi,
                                   MultiFab* Source,
                                   MultiFab* Resid,
                                   MultiFab* Sigma[],
                                   const Real* Dx) :
       multigrid(Grids, Geom, Phi, Source, Resid)
{

  for (int n = 0; n < BL_SPACEDIM; n++)
    sigma[n] = Sigma[n];

  // Check whether we can coarsen all the grids separately
  int len = grids[0].length()[0];
  for (int i = 0; i < grids.size(); i++) 
    for (int n = 0; n < BL_SPACEDIM; n++)
      len = std::min(len,grids[i].length()[n]);

  bool is_odd = false;
  for (int i = 0; i < grids.size(); i++) 
    for (int n = 0; n < BL_SPACEDIM; n++)
      if ( (grids[i].length()[n]&1) != 0 ) is_odd = true;

  Box prob_domain(grids.minimalBox());

  // Check whether we can coarsen the domain as a whole
  int domain_len = prob_domain.length()[0];
  for (int n = 1; n < BL_SPACEDIM; n++)
      domain_len = std::min(domain_len,prob_domain.length()[n]);

  bool domain_is_odd = false;
  for (int n = 0; n < BL_SPACEDIM; n++)
      if ( (prob_domain.length()[n]&1) != 0 ) domain_is_odd = true;

// If any grid has an odd number of cells, or has reached its minimum
// length in either direction, then we can not coarsen the grids
// as such any further.  If the domain as a whole can still be coarsened,
// then we proceed with the domain as one grid.  If not, then we
// use a conjugate gradient bottom solver on a single grid.

  int ng = phi->nGrow();

  if (domain_len < 4 || domain_is_odd) 
  {

    Next = NULL;  

    BL_ASSERT(grids.size() == 1);

    cgwork = new MultiFab(grids,5,1);

    for (MFIter mfi(*cgwork); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

      FORT_MKSUMMAC((*cgwork)[mfi].dataPtr(), 
                    (*sigma[0])[mfi].dataPtr(), (*sigma[1])[mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                    (*sigma[2])[mfi].dataPtr(),
#endif
                    ARLIM(lo), ARLIM(hi),dx,&ng);
    }


  } else {

    cgwork = NULL; 

    BoxArray newgrids;
 
    if (len < 8 || is_odd) 
    {
      BoxList pd;
      pd.push_back(prob_domain);
      newgrids.define(pd);
    }
    else 
    {
      newgrids.define(grids);
    }

    newgrids.coarsen(2);
    MultiFab* newphi = new MultiFab(newgrids,1,ng);
    newphi->setVal(0.0);

    Real newdx[BL_SPACEDIM];
    for (int n = 0; n < BL_SPACEDIM; n++)
      newdx[n] = 2.0*dx[n];

    Geometry newgeom;
    newgeom.define(BoxLib::coarsen(geom.Domain(),2));

    MultiFab* newsigma[BL_SPACEDIM];
    for (int n = 0; n < BL_SPACEDIM; n++)
    {
      BoxArray new_edge_grids(newgrids);
      new_edge_grids.surroundingNodes(n);
      newsigma[n] = new MultiFab(new_edge_grids,1,ng-1);
    }

    if (grids.size() == newgrids.size()) 
    {
     // 
     // In this case, the BoxArray newgrids is a coarsened version of grids.
     // 
     for (MFIter sxmfi(*sigma[0]); sxmfi.isValid(); ++sxmfi)
     {
      int i = sxmfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

      const int* newlo = newgrids[i].loVect();
      const int* newhi = newgrids[i].hiVect();

#if (BL_SPACEDIM == 2)
      FORT_COARSIGMA((*sigma[0])[sxmfi].dataPtr(),(*sigma[1])[sxmfi].dataPtr(),
                     (*newsigma[0])[sxmfi].dataPtr(),(*newsigma[1])[sxmfi].dataPtr(),
                     ARLIM(lo), ARLIM(hi),
                     ARLIM(newlo), ARLIM(newhi),&ng);
#elif (BL_SPACEDIM == 3)
      FORT_COARSIGMA((*sigma[0])[sxmfi].dataPtr(),(*sigma[1])[sxmfi].dataPtr(),(*sigma[2])[sxmfi].dataPtr(),
                     (*newsigma[0])[sxmfi].dataPtr(),(*newsigma[1])[sxmfi].dataPtr(),
                     (*newsigma[2])[sxmfi].dataPtr(),
                     ARLIM(lo), ARLIM(hi),
                     ARLIM(newlo), ARLIM(newhi),&ng);
#endif
     }
     for (int n = 0; n < BL_SPACEDIM; n++)
     {
       newsigma[n]->FillBoundary(0,1);
       newgeom.FillPeriodicBoundary(*newsigma[n],0,1,true);
     }

    }
    else 
    {

     // 
     // In this case, there is only one grid (the size of the domain)
     //   in the BoxArray newgrids.
     // 
     BL_ASSERT(newgrids.size() == 1);
  
     BoxList pd;
     pd.push_back(prob_domain);
     BoxArray grids_mf(pd);

     MultiFab* sigma_mf[BL_SPACEDIM];
     for (int n = 0; n < BL_SPACEDIM; n++) 
     {
       BoxArray edge_grids(grids_mf);
       edge_grids.surroundingNodes(n);
       sigma_mf[n] = new MultiFab(edge_grids,1,ng-1);
       sigma_mf[n]->copy(*sigma[n],0,0,1);
     }

     for (MFIter sxmfi(*sigma_mf[0]); sxmfi.isValid(); ++sxmfi)
     {
      int i = sxmfi.index();
      const int* lo = grids_mf[i].loVect();
      const int* hi = grids_mf[i].hiVect();

      const int* newlo = newgrids[i].loVect();
      const int* newhi = newgrids[i].hiVect();

#if (BL_SPACEDIM == 2)
      FORT_COARSIGMA((*sigma_mf[0])[sxmfi].dataPtr(),(*sigma_mf[1])[sxmfi].dataPtr(),
                     (*newsigma[0])[sxmfi].dataPtr(), (*newsigma[1])[sxmfi].dataPtr(),
                     ARLIM(lo), ARLIM(hi),
                     ARLIM(newlo), ARLIM(newhi),&ng);
#elif (BL_SPACEDIM == 3)
      FORT_COARSIGMA((*sigma_mf[0])[sxmfi].dataPtr(),(*sigma_mf[1])[sxmfi].dataPtr(),(*sigma_mf[2])[sxmfi].dataPtr(),
                     (*newsigma[0])[sxmfi].dataPtr(), (*newsigma[1])[sxmfi].dataPtr(),
                     (*newsigma[2])[sxmfi].dataPtr(),
                     ARLIM(lo), ARLIM(hi),
                     ARLIM(newlo), ARLIM(newhi),&ng);
#endif
     }
     for (int n = 0; n < BL_SPACEDIM; n++) 
       delete sigma_mf[n];
     for (int n = 0; n < BL_SPACEDIM; n++)
       newgeom.FillPeriodicBoundary(*newsigma[n],0,1,true);
    }

    Next = new macprojection_mg(newgrids,
                                newgeom,
                                newphi,
                                new MultiFab(newgrids,1,ng-1),
                                new MultiFab(newgrids,1,ng-1),
			        newsigma,
			        newdx);
  }
  next = Next;
}

// ************************************************************************
// ** destructor **
// ************************************************************************

macprojection_mg::~macprojection_mg()

{
  if (Next != NULL) {
    delete Next->phi;
    delete Next->source;
    delete Next->resid;
    for (int n = 0; n < BL_SPACEDIM; n++) 
      delete Next->sigma[n];
    delete Next;
  }
    delete cgwork;
}

// ************************************************************************
// ** residual **
// ************************************************************************

Real macprojection_mg::residual()
{
  Real  norm = -1.e20;
  Real rnorm = -1.e20;

  int ng = phi->nGrow();

  for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
  {
    int i = mfi.index();
    const int* lo = grids[i].loVect();
    const int* hi = grids[i].hiVect();

    int* bcptr = bc[i].dataPtr();
    if (grids.size() == 1) bcptr = domain_bc;

    FORT_RESIDUAL((*resid)[mfi].dataPtr(),(*phi)[mfi].dataPtr(),(*source)[mfi].dataPtr(),
                  (*sigma[0])[mfi].dataPtr(),(*sigma[1])[mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
  	          (*sigma[2])[mfi].dataPtr(),
#endif
                  ARLIM(lo),ARLIM(hi),dx,&norm,bcptr,&ng);
    rnorm = std::max(rnorm,norm);
  }
  resid->FillBoundary(0,1);
  ParallelDescriptor::ReduceRealMax(rnorm);

  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*resid,0,1,true);

  return rnorm;
}

// ************************************************************************
// ** step **
// ************************************************************************

void macprojection_mg::step(int nngsrb)
{
  Real rnorm = 1.e10;
 
  int ng = phi->nGrow();

  if (cgwork != NULL) {

    BL_ASSERT(grids.size() == 1);

    MultiFab dest0(grids,1,1);

    for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

      FORT_SOLVEMAC((*phi)[mfi].dataPtr(), dest0[mfi].dataPtr(),
                    (*source)[mfi].dataPtr(),
                    (*sigma[0])[mfi].dataPtr(), (*sigma[1])[mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                    (*sigma[2])[mfi].dataPtr(),
#endif
                    (*cgwork)[mfi].dataPtr(0), (*cgwork)[mfi].dataPtr(1), (*cgwork)[mfi].dataPtr(2), 
                    (*cgwork)[mfi].dataPtr(3), (*cgwork)[mfi].dataPtr(4), 
                    ARLIM(lo),ARLIM(hi),dx,domain_bc,&rnorm,&prob_norm,&ng);
    }

  } else {

    for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

      int* bcptr = bc[i].dataPtr();
      if (grids.size() == 1) bcptr = domain_bc;

        FORT_RELAX((*phi)[mfi].dataPtr(),(*source)[mfi].dataPtr(),
                   (*sigma[0])[mfi].dataPtr(),(*sigma[1])[mfi].dataPtr(),
#if (BL_SPACEDIM == 3)
                   (*sigma[2])[mfi].dataPtr(),
#endif
                   ARLIM(lo),ARLIM(hi),dx,bcptr,&nngsrb,&ng);
    }
  }

  //  
  //  Note that we choose to impose bc's after each time we modify phi,
  //    as opposed to before each time we use it.
  //  
  phi->FillBoundary(0,1);
  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*phi,0,1,true);
}

// ************************************************************************
// ** Restrict **
// ************************************************************************

void macprojection_mg::Restrict()
{
  int nextra = resid->nGrow();
  if (grids.size() == next->grids.size())
  {
    for (MFIter mfi(*resid); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();
  
      const int * loc = next->grids[i].loVect();
      const int * hic = next->grids[i].hiVect();

      FORT_RESTRICT((*resid)[mfi].dataPtr(), (*next->source)[mfi].dataPtr(),
    		  ARLIM(lo),ARLIM(hi),ARLIM(loc),ARLIM(hic),&nextra);
    }
  } else {
  
    BoxList pd;
    Box prob_domain(grids.minimalBox());
    pd.push_back(prob_domain);
    BoxArray grids_mf(pd);

    MultiFab res_mf(grids_mf,1,resid->nGrow());
    res_mf.copy(*resid,0,0,1);

    for (MFIter mfi(res_mf); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids_mf[i].loVect();
      const int* hi = grids_mf[i].hiVect();

      const int * loc = next->grids[i].loVect();
      const int * hic = next->grids[i].hiVect();

      FORT_RESTRICT(res_mf[mfi].dataPtr(), (*next->source)[mfi].dataPtr(),
    		    ARLIM(lo),ARLIM(hi),ARLIM(loc),ARLIM(hic),&nextra);
    }
  }
  next->source->FillBoundary(0,1);
  if (next->geom.isAnyPeriodic())
    next->geom.FillPeriodicBoundary(*next->source,0,1,true);
}

// ************************************************************************
// ** interpolate **
// ************************************************************************

void macprojection_mg::interpolate()
{
  int ng = phi->nGrow();
  if (grids.size() == next->grids.size()) 
  {
    for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

      const int * loc = next->grids[i].loVect();
      const int * hic = next->grids[i].hiVect();
  
      FORT_INTERPOLATE((*phi)[mfi].dataPtr(), (*next->phi)[mfi].dataPtr(),
  		     ARLIM(lo),ARLIM(hi),ARLIM(loc),ARLIM(hic),&ng);
    }
  } else {

    BoxList pd;
    Box prob_domain(grids.minimalBox());
    pd.push_back(prob_domain);
    BoxArray grids_mf(pd);

    MultiFab phi_interp(grids_mf,1,phi->nGrow());
    phi_interp.copy(*phi,0,0,1);

    for (MFIter mfi(phi_interp); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids_mf[i].loVect();
      const int* hi = grids_mf[i].hiVect();

      const int * loc = next->grids[i].loVect();
      const int * hic = next->grids[i].hiVect();
  
      FORT_INTERPOLATE(phi_interp[mfi].dataPtr(), (*next->phi)[mfi].dataPtr(),
  		     ARLIM(lo),ARLIM(hi),ARLIM(loc),ARLIM(hic),&ng);
    }
    phi->copy(phi_interp,0,0,1);
  }

  phi->FillBoundary(0,1);
  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*phi,0,1,true);
}
