/*
** (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 <hgproj.H>
#include <globals.H>
#include <HGPROJ_F.H>
#include <BCTypes.H>
#include <ParmParse.H>
#include <names.H>

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

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

hg_projector::hg_projector(const BoxArray& Grids, const Geometry& Geom,
                           MultiFab* Area, MultiFab& Vol) :
       grids(Grids), geom(Geom), area(Area), vol(Vol)

{

  ParmParse pp("hg");
  pp.get("tol",tol);

  pp.query("numSmoothCoarsen",numSmoothCoarsen);
  pp.query("numSmoothRefine",numSmoothRefine);
 
  ng = 2*std::max(numSmoothCoarsen,numSmoothRefine);
}

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

hg_projector::~hg_projector() { }

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

void hg_projector::project(MultiFab * state, 
                           MultiFab * pressure,
                           MultiFab * rho, 
                           MultiFab * divu_src, 
                           Real time, 
                           Real dt)

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

  BoxArray nodal_grids(grids);
  nodal_grids.surroundingNodes();

  MultiFab *sigma = new MultiFab(grids,1,ng);
  sigma->setVal(0.);

#if (BL_SPACEDIM == 3)
  MultiFab *coeff = new MultiFab(nodal_grids,3,ng);
  coeff->setVal(0.);
#endif

  MultiFab *phi    = new MultiFab(nodal_grids,1,ng);
  MultiFab *source = new MultiFab(nodal_grids,1,ng);
  MultiFab *resid  = new MultiFab(nodal_grids,1,ng);

  phi->setVal(0.0);
  source->setVal(0.);
  resid->setVal(0.);

  MultiFab::Copy(*sigma,vol,0,0,1,0);
  for (MFIter mfi(*sigma); mfi.isValid(); ++mfi)
  {
     int i = mfi.index();
     (*sigma)[mfi].divide((*rho)[mfi],grids[i],0,0,1);
  }
  sigma->FillBoundary(0,1);
  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*sigma,0,1,true);

  Real rhsnorm = -1.e20;
  Real    norm = -1.e20;

  MultiFab gradp(grids,BL_SPACEDIM,1);
  gradient(&gradp,pressure);

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

     for (int n = 0; n < BL_SPACEDIM; n++)
     {
       gradp[mfi].divide((*rho)[mfi],grids[i],0,n,1);
       (*state)[mfi].plus(gradp[mfi],grids[i],n,n,1); 
     }
  }

  state->FillBoundary(0,BL_SPACEDIM);
  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*state,0,BL_SPACEDIM,true);
 
  int is_singular = 1;
  for (int n = 0; n < 2*BL_SPACEDIM; n++)
    if (domain_bc[n] == OUTLET) is_singular = 0;

  Real sum_src = 0.0;
  Real sum_fac = 0.0;

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

     Real ss = 0.;
     Real sf = 0.;
     FORT_RHSHG((*source)[mfi].dataPtr(),(*state)[mfi].dataPtr(),
                (*divu_src)[mfi].dataPtr(),vol[mfi].dataPtr(),
                ARLIM(lo),ARLIM(hi),dx,bc[i].dataPtr(),&norm,&ng,
                &is_singular,&ss,&sf);
     sum_src += ss;
     sum_fac += sf;

     rhsnorm = std::max(rhsnorm,norm);
  }
  ParallelDescriptor::ReduceRealSum(sum_src);
  ParallelDescriptor::ReduceRealSum(sum_fac);
  ParallelDescriptor::ReduceRealMax(rhsnorm);

  if (is_singular == 1)
  {
    sum_src /= sum_fac;
    if (ParallelDescriptor::IOProcessor())
      std::cout << "Singular adjustment for nodal projection is " << sum_src << std::endl;
    for (MFIter mfi(*source); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();
      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();
      FORT_ADJUSTRHS((*source)[mfi].dataPtr(),ARLIM(lo),ARLIM(hi),
                     bc[i].dataPtr(), &ng,&sum_src);
    }
  }

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

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

  if (rhsnorm > 1.e-16) {

    phi->setVal(0.0);

    Real goal = tol*rhsnorm;

    hgprojection_mg *solver =
      new hgprojection_mg(grids,geom,phi,source,resid,sigma,
#if (BL_SPACEDIM == 3)
                          coeff,
#endif
                          dx);

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

    MultiFab gradphi(grids,BL_SPACEDIM,1);
    gradient(&gradphi,phi);

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

       FORT_PROJUHG((*state)[mfi].dataPtr(), (*pressure)[mfi].dataPtr(),
                    (*phi)[mfi].dataPtr(), gradphi[mfi].dataPtr(),
                    (*rho)[mfi].dataPtr(), ARLIM(lo),ARLIM(hi),&ng);
     }
   }

  delete sigma;
  delete phi;
  delete source;
  delete resid;
#if (BL_SPACEDIM == 3)
  delete coeff;
#endif
}

// ************************************************************************
// ** gradient **
// ************************************************************************

void hg_projector::gradient(MultiFab * gradp, MultiFab * pressure)
{
   const Real* dx = geom.CellSize();
   for (MFIter pmfi(*pressure); pmfi.isValid(); ++pmfi)
   {
     int i = pmfi.index();
     const int* lo = grids[i].loVect();
     const int* hi = grids[i].hiVect();

     const int* p_lo = (*pressure)[pmfi].loVect();
     const int* p_hi = (*pressure)[pmfi].hiVect();

     const int* g_lo = (*gradp)[pmfi].loVect();
     const int* g_hi = (*gradp)[pmfi].hiVect(); 

     FORT_GRADHG((*gradp)[pmfi].dataPtr(),ARLIM(g_lo),ARLIM(g_hi),
                 (*pressure)[pmfi].dataPtr(),ARLIM(p_lo),ARLIM(p_hi),
   	         ARLIM(lo),ARLIM(hi),dx);
   }
   if (gradp->nGrow() > 0) 
   {
     gradp->FillBoundary(0,BL_SPACEDIM);
     if (geom.isAnyPeriodic())
       geom.FillPeriodicBoundary(*gradp,0,BL_SPACEDIM,true);
   }
}

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

hgprojection_mg::hgprojection_mg(const BoxArray& Grids,
                                 const Geometry& Geom,
                                 MultiFab* Phi,
                                 MultiFab* Source,
                                 MultiFab* Resid,
                                 MultiFab* Sigma,
#if (BL_SPACEDIM == 3)
                                 MultiFab* Coeff,
#endif
                                 const Real* Dx) : 
       multigrid(Grids, Geom, Phi, Source, Resid)
{
  int i,n;
  int ng = Phi->nGrow();
  sigma = Sigma;

#if (BL_SPACEDIM == 3)
  coeff = Coeff;
  for (MFIter mfi(*coeff); mfi.isValid(); ++mfi)
  {
    i = mfi.index();
    const int* lo = grids[i].loVect();
    const int* hi = grids[i].hiVect();

    FORT_MKCOEFF((*coeff)[mfi].dataPtr(), (*sigma)[mfi].dataPtr(),
                 ARLIM(lo),ARLIM(hi),&ng);
   }
   coeff->FillBoundary(0,3);
   if (geom.isAnyPeriodic())
     geom.FillPeriodicBoundary(*coeff,0,3,true);
#endif

// 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.

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

  bool is_odd = false;
  for (i = 0; i < grids.size(); i++)
    for (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 (n = 1; n < BL_SPACEDIM; n++)
    domain_len = std::min(domain_len,prob_domain.length()[n]);

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

  if (domain_is_odd || domain_len < 4)
  {
    Next = NULL;

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

    BoxArray nodal_grids_mf(grids_mf);
    nodal_grids_mf.surroundingNodes();

    int ng = phi->nGrow();

    cgwork = new MultiFab(nodal_grids_mf,5,ng);
    cgwork->setVal(0.);

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

      FORT_MAKESUM((*cgwork)[mfi].dataPtr(0), (*coeff)[mfi].dataPtr(),
                   ARLIM(lo),ARLIM(hi),dx,domain_bc,&ng);
    }
#endif

  } 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);

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

    BoxArray new_nodal_grids(newgrids);
    new_nodal_grids.surroundingNodes();

    MultiFab* newsigma  = new MultiFab(newgrids,1,ng);
    newsigma->setVal(0.);

#if (BL_SPACEDIM == 3)
    MultiFab* newcoeff  = new MultiFab(new_nodal_grids,3,ng);
    newcoeff->setVal(0.);
#endif

    MultiFab* newphi    = new MultiFab(new_nodal_grids,1,ng);
    MultiFab* newsource = new MultiFab(new_nodal_grids,1,ng);
    MultiFab* newresid  = new MultiFab(new_nodal_grids,1,ng);

    newphi->setVal(0.0);
    newresid->setVal(0.);
    newsource->setVal(0.);

    if (grids.size() == newgrids.size()) 
    {
      for (MFIter sigmfi(*sigma); sigmfi.isValid(); ++sigmfi)
      {
        i = sigmfi.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();
  
        FORT_COARSIG((*sigma)[sigmfi].dataPtr(), (*newsigma)[sigmfi].dataPtr(),
                     ARLIM(lo), ARLIM(hi),
    		     ARLIM(newlo), ARLIM(newhi),&ng);
      }
      newsigma->FillBoundary(0,1);

    } else {

      BoxList pd;
      pd.push_back(prob_domain); 
      BoxArray grids_mf(pd);
   
      MultiFab sigma_mf(grids_mf,1,sigma->nGrow());
      sigma_mf.copy(*sigma,0,0,1);

      for (MFIter sigmfi(sigma_mf); sigmfi.isValid(); ++sigmfi)
      {
        i = sigmfi.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();

        FORT_COARSIG(sigma_mf[sigmfi].dataPtr(), (*newsigma)[sigmfi].dataPtr(),
                     ARLIM(lo), ARLIM(hi),
 		     ARLIM(newlo), ARLIM(newhi),&ng);
      }
    }

    if (newgeom.isAnyPeriodic())
      newgeom.FillPeriodicBoundary(*newsigma,0,1,true);

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

    Next = new hgprojection_mg(newgrids, newgeom, newphi,
                               newsource, newresid, newsigma, 
#if (BL_SPACEDIM == 3)
			       newcoeff,
#endif
			       newdx);
  }
  next = Next;
}

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

hgprojection_mg::~hgprojection_mg()
{
  if (Next != NULL) {
    delete Next->phi;
    delete Next->source;
    delete Next->resid;
    delete Next->sigma;
#if (BL_SPACEDIM ==3 )
    delete Next->coeff;
#endif
    delete Next;
  }
    delete cgwork;
}

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

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

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

  int ng = phi->nGrow();

  MultiFab dgphi(phi->boxArray(),1,ng);
  for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
  {
    int i = mfi.index();
    const int* lo = grids[i].loVect();
    const int* hi = grids[i].hiVect();
  
#if (BL_SPACEDIM == 2)
    FArrayBox& fab = (*sigma)[mfi];
#elif (BL_SPACEDIM == 3)
    FArrayBox& fab = (*coeff)[mfi];
#endif

    int * bcptr = bc[i].dataPtr();
    if (grids.size() == 1) bcptr = domain_bc;
  
    FORT_RESIDUAL((*resid)[mfi].dataPtr(),(*phi)[mfi].dataPtr(),
                  (*source)[mfi].dataPtr(),fab.dataPtr(),dgphi[mfi].dataPtr(),
                  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 hgprojection_mg::step(int nngsrb)
{
  int maxiter = 100;

  int ng = phi->nGrow();

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

  if (cgwork != NULL) 
  {
    BL_ASSERT(grids.size() == 1);

    Real resnorm  = 0.0;
    MultiFab dest0(phi->boxArray(),1,ng);

    for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
    {
      int i = mfi.index();

      const int* lo = grids[i].loVect();
      const int* hi = grids[i].hiVect();

#if (BL_SPACEDIM == 2)
      FArrayBox& fab = (*sigma)[mfi];
#elif (BL_SPACEDIM == 3)
      FArrayBox& fab = (*coeff)[mfi];
#endif
  
      FORT_SOLVEHG((*phi)[mfi].dataPtr(), dest0[mfi].dataPtr(), 
                   (*source)[mfi].dataPtr(), fab.dataPtr(),
                   (*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, &maxiter, &resnorm, &prob_norm, &ng);
    }
 }
 else 
 {
      MultiFab temp(phi->boxArray(),1,ng);
      for (MFIter mfi(*phi); mfi.isValid(); ++mfi)
      {
        int i = mfi.index();

        const int* lo = grids[i].loVect();
        const int* hi = grids[i].hiVect();
  
#if (BL_SPACEDIM == 2)
        FArrayBox& fab = (*sigma)[mfi];
#elif (BL_SPACEDIM == 3)
        FArrayBox& fab = (*coeff)[mfi];
#endif

        int * bcptr = bc[i].dataPtr();
        if (grids.size() == 1) bcptr = domain_bc;
  
        FORT_RELAX((*phi)[mfi].dataPtr(),(*source)[mfi].dataPtr(),
                   fab.dataPtr(),temp[mfi].dataPtr(),
                   ARLIM(lo),ARLIM(hi),dx,bcptr,&nngsrb,&ng);
      }
  }
  phi->FillBoundary(0,1);
  if (geom.isAnyPeriodic())
    geom.FillPeriodicBoundary(*phi,0,1,true);
}

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

void hgprojection_mg::Restrict()
{
  int ng = phi->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();
  
      int * bcptr = bc[i].dataPtr();
      if (grids.size() == 1) bcptr = domain_bc;
  
      FORT_RESTRICT((*resid)[mfi].dataPtr(), (*next->source)[mfi].dataPtr(),
                    ARLIM(lo),ARLIM(hi),ARLIM(loc),ARLIM(hic),
                    bcptr, &ng);
    }
    next->source->FillBoundary(0,1);

  } else {

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

    BoxArray nodal_grids_mf(grids_mf);
    nodal_grids_mf.surroundingNodes();
   
    MultiFab res_mf(nodal_grids_mf,1,resid->nGrow());

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

    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),
                    domain_bc, &ng);
    }
  }

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

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

void hgprojection_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_INTERP((*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);

    BoxArray nodal_grids_mf(grids_mf);
    nodal_grids_mf.surroundingNodes();
   
    MultiFab phi_interp(nodal_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_INTERP(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);
}
