//===========================================================================
// SINTEF LSMG library - version 1.1
//
// Copyright (C) 2000-2005 SINTEF ICT, Applied Mathematics, Norway.
//
// This program is free software; you can redistribute it and/or          
// modify it under the terms of the GNU General Public License            
// as published by the Free Software Foundation version 2 of the License. 
//
// This program is distributed in the hope that it will be useful,        
// but WITHOUT ANY WARRANTY; without even the implied warranty of         
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the          
// GNU General Public License for more details.                           
//
// You should have received a copy of the GNU General Public License      
// along with this program; if not, write to the Free Software            
// Foundation, Inc.,                                                      
// 59 Temple Place - Suite 330,                                           
// Boston, MA  02111-1307, USA.                                           
//
// Contact information: e-mail: tor.dokken@sintef.no                      
// SINTEF ICT, Department of Applied Mathematics,                         
// P.O. Box 124 Blindern,                                                 
// 0314 Oslo, Norway.                                                     
//
// Other licenses are also available for this software, notably licenses
// for:
// - Building commercial software.                                        
// - Building software whose source code you wish to keep private.        
//===========================================================================
#include <MGsystem.h>

#include <UCBsplines.h>
#include <MBA.h>

#include <math.h>
#include <algorithm>
#include <valarray>
using namespace std;


static void errorAndExit(char* message) {
  cout << "ERROR: ";
  cout << message << endl;
  exit(-1);
}
#define DEBUG_0  1    

#include <UCBsplineSurface.h>
#include <UCButils.h>
static void printGLgrid(const char filename[], boost::shared_ptr<GenMatrixType> pmat, double scale) {
  

  double umin = 0.0;
  double vmin = 0.0;
  int m = pmat->noX()-3;
  int n = pmat->noY()-3;
  double umax = (double)m;
  double vmax = (double)n;

  
  UCBspl::SplineSurface surf(pmat, umin, vmin, umax, vmax);
  UCBspl::printGLgrid(filename, surf, m+1, n+1);
}

static void printVTKgrid(const char filename[], const GenMatrixType& mat, double scale) {
  
  
  cout << "Printing grid to vtk structured points file: " << filename << "...... ";
  ofstream os(filename);

  int noU = mat.noX();
  int noV = mat.noY();

  double umin = 0.0;
  double vmin = 0.0;
  double umax = (double)(noU-1);;
  double vmax = (double)(noV-1);;
  double du = (umax - umin)/(double)(noU-1);
  double dv = (vmax - vmin)/(double)(noV-1);  

  os << "# vtk DataFile Version 2.0" << endl;
  os << "This file was generated by class UCButils (structured points)" << endl;
  os << "ASCII" << endl;
  os << "DATASET STRUCTURED_POINTS" << endl;
  os << "DIMENSIONS " << noU << " " << noV << " " << 1 << endl;
  os << "ORIGIN " << umin << " " << vmin << " " << 999 << endl;
  os << "SPACING " << du << " " << dv << " " << 999 << endl;
  os << "POINT_DATA " << noU*noV*1 << endl;
  os << "SCALARS volume_scalars float" << endl;
  os << "LOOKUP_TABLE default" << endl;

  int noPoints = noU*noV;

  double maxVal = mat(-1,-1);
  double minVal = mat(-1,-1);
  double sum=0;

  int i,j;
  for (j = 0; j < noV; j++) {
    double v = vmin + j*dv;
    for (i = 0; i < noU; i++) {
      double u = umin + i*du;
      double val = mat(i-1,j-1);
      sum += val;
      if (val > maxVal)
        maxVal = val;
      if (val < minVal)
        minVal = val;
      os << val*scale << '\n';
    }
  }

  cout << "minVal = " << minVal << "  maxVal = " << maxVal << " Mean = " << sum/(double)(noU*noV)<< endl; 
}



void MGsystem::residual(const LSsystem& lssys, GenMatrix<UCBspl_real>& r) const {




  r.resize(lssys.n1_, lssys.n2_);
  int n = lssys.n1_*lssys.n2_; 

  for (int i = 0; i < n; i++) { 
    
    double sum = LSsystem::rowMatVecMult(i, lssys.A_, *lssys.x_);
    
    int k = i%lssys.n1_ - 1; 
    int l = i/lssys.n1_ - 1;
    
    r(k,l) = lssys.b_(k,l) - sum; 
  }
}


static double residual(const std::vector<LSsystem*>& systems, int normType, bool scaled=false) {
  

  int noLevels = systems.size();

  if (noLevels == 0 || (normType != 0 && normType != 2)) {
#ifdef DEBUG_0
    cout << "ERROR: No equation system present" << endl;
#endif
    return 0.0;
  }

  if (normType == 0)
    return systems[noLevels-1]->residual_linf();
  else if (normType == 2)
    return systems[noLevels-1]->residual_l2(scaled);
  else
    return 0.;
}


double MGsystem::residual_l2(bool scaled) const {
  
  int normType = 2;
  return ::residual(systems_,normType,scaled);
}

double MGsystem::residual_linf() const {

  int normType = 0;
  return ::residual(systems_,normType);
}


void MGsystem::residual_restrictionInjection(const LSsystem& lssys, GenMatrix<UCBspl_real>& r) const {

  cout << "residual_restrictionInjection..." << endl;



  int old_noX = lssys.n1_;
  int old_noY = lssys.n2_;

  if ((old_noX-3)%2 != 0 || (old_noY-3)%2 != 0) {
    cout << "ERROR, invalid grid for projection; see comments in the code above: " << old_noX << " " << old_noY << endl;
    exit(-1);
  }

  int noX = (old_noX-3)/2 + 3;
  int noY = (old_noY-3)/2 + 3;

  r.resize(noX, noY);

  int n = lssys.n1_*lssys.n2_; 

  int kk = -1;
  int ll = -1;

  int i = 0;
  while (i < n) {
    
    double sum = LSsystem::rowMatVecMult(i, lssys.A_, *lssys.x_);



    int k = i%lssys.n1_ - 1; 
    int l = i/lssys.n1_ - 1;        
    r(kk,ll) = lssys.b_(k,l) - sum; 

    if (kk == -1 || kk == noX-3) { 
      kk++;
      i += 1;
    }
    else if (kk == noX-2) {        
      if (ll == -1 || ll == noY-3) 
        i += 1;
      else                         
        i += (lssys.n1_+1);           
      kk = -1;
      ll++;
    }
    else {                         
      kk++;
      i += 2;
    }
  }

#ifdef DEBUG_1
  cout << "\n(new) noX,noY = " << noX << "," << noY << endl; 
  cout << "n1_ = " << lssys.n1_ << " n2_ = " << lssys.n2_ << endl;
  cout << "n = n1_*n2_: " << n << endl; 
  cout << "DONE residual_restrictionInjection" << endl;
#endif
}


static void restrictionFullWeightingLaplace(const GenMatrix<UCBspl_real>& rr, GenMatrix<UCBspl_real>& r) {

  cout << "restrictionFullWeightingLaplace..." << endl;
  
  int old_noX = rr.noX();
  int old_noY = rr.noY();

  if ((old_noX-3)%2 != 0 || (old_noY-3)%2 != 0) {
    cout << "ERROR, invalid grid for projection; see comments in the code above: " << old_noX << " " << old_noY << endl;
    exit(-1);
  }

  int noX = (old_noX-3)/2 + 3;
  int noY = (old_noY-3)/2 + 3;

  r.resize(noX, noY);


  double op_corners;
  double op_sides; 
  double op_center;
  double op_div;

 
  op_corners =  1.0;
  op_sides  =  2.0; 
  op_center =  4.0;
  op_div    = 16.0; 
  

  bool halfWeighting=false;
  if (halfWeighting) {
    op_corners =  0.0;
    op_sides  =  1.0; 
    op_center =  4.0;
    op_div    = 8.0;
  }






  int n = old_noX*old_noY; 

  int kk = -1;
  int ll = -1;

  int i = 0;
  while (i < n) {



    
    int k = i%old_noX - 1; 
    int l = i/old_noX - 1;    

    


    if (kk == -1 || kk == noX-2 || ll == -1 || ll == noY-2)
      r(kk,ll) = rr(k,l);
    else {
      double val = op_corners * (rr(k-1,l-1) + rr(k-1,l+1) + rr(k+1,l-1) + rr(k+1,l+1)) 
         +  op_sides  * (rr(k,l-1) + rr(k,l+1) + rr(k-1,l) + rr(k+1,l))
         +  op_center * rr(k,l);
      r(kk,ll) = val/op_div;
    }


    if (kk == -1 || kk == noX-3) { 
      kk++;
      i += 1;
    }
    else if (kk == noX-2) {        
      if (ll == -1 || ll == noY-3) 
        i += 1;
      else                         
        i += (old_noX+1); 
      kk = -1;
      ll++;
    }
    else {                         
      kk++;
      i += 2;
    }
  }


#ifdef DEBUG_1
  cout << "\n(new) noX,noY = " << noX << "," << noY << endl; 
  cout << "n1_ = " << lssys.n1_ << " n2_ = " << lssys.n2_ << endl;
  cout << "n = n1_*n2_: " << n << endl; 
  cout << "DONE restrictionFullWeightingLaplace" << endl;
#endif
}


void MGsystem::residual_restrictionFullWeightingLaplace(const LSsystem& lssys, GenMatrix<UCBspl_real>& r) const {

  int old_noX = lssys.n1_;
  int old_noY = lssys.n2_;

  if ((old_noX-3)%2 != 0 || (old_noY-3)%2 != 0) {
    cout << "ERROR, invalid grid for projection; see comments in the code above: " << old_noX << " " << old_noY << endl;
    exit(-1);
  }

  GenMatrix<UCBspl_real> rr; 
  residual(lssys, rr);    
  restrictionFullWeightingLaplace(rr, r);
}

void MGsystem::residual_restrictionSplines(const LSsystem& lssys, GenMatrix<UCBspl_real>& r) const {

  int old_noX = lssys.n1_;
  int old_noY = lssys.n2_;

  if ((old_noX-3)%2 != 0 || (old_noY-3)%2 != 0) {
    cout << "ERROR, invalid grid for projection; see comments in the code above: " << old_noX << " " << old_noY << endl;
    exit(-1);
  }

  GenMatrix<UCBspl_real> rr;
  residual(lssys, rr);

  UCBspl::restrictCoeffsC2(rr,r);
  

  
  return;
}

void MGsystem::relaxAndCorrectVcycle(int noLevelsInVcycle) {
  
  
  std::vector<int> noIterArr;
  noIterArr.resize(noLevelsInVcycle,noIter_);
  noIterArr[0] = noIterCoarsest_; 
  
#ifdef DEBUG_0
  cout << "\nrelaxAndCorrectVcycle..." << endl;
  cout << "\nLevel " << noLevelsInVcycle-1 << ":\n";
#endif

  systems_[noLevelsInVcycle-1]->relaxGaussSeidel(noIterArr[noLevelsInVcycle-1]);
#ifdef DEBUG_0
#endif
  
  
  const int restrictionType = 1; 
  for (int i = noLevelsInVcycle-2; i >= 0; i--) {
    if (restrictionType == 1)
      residual_restrictionSplines(*systems_[i+1],r_);  
    else if (restrictionType == 2) {
      cout << "!!!!!!!!!!!!!!!!!!!!!!!!!!!  Note: Direct injection..." << endl;
      residual_restrictionInjection(*systems_[i+1],r_);
    }
    else if (restrictionType == 3)
      residual_restrictionFullWeightingLaplace(*systems_[i+1],r_);  
    
     
    systems_[i]->x_->fill(0.0); 
    systems_[i]->b_.swap(r_);
#ifdef DEBUG_0
    cout << "\nLevel " << i << ":\n";
#endif
    systems_[i]->relaxGaussSeidel(noIterArr[i]); 
#ifdef DEBUG_0
#endif
  }
  
  for (int i = 1; i < noLevelsInVcycle; i++) {
    UCBspl::refineCoeffsC2(*systems_[i-1]->x_,e_refined_);
    (*systems_[i]->x_) += e_refined_;
#ifdef DEBUG_0
    cout << "\nLevel " << i << ":\n";
#endif
    systems_[i]->relaxGaussSeidel(noIterArr[i]);
#ifdef DEBUG_0
#endif
  }
#ifdef DEBUG_0  
  cout << ".................................done relaxAndCorrectVcycle.................." << endl;
#endif
}

UCBspl::SplineSurface MGsystem::getSplineSurface() const {
  if (systems_.size() == 0)
    return UCBspl::SplineSurface();
  else {
    return UCBspl::SplineSurface(systems_[systems_.size()-1]->x_,umin(), vmin(), umax(), vmax());
  }
}


static void restrictionSplines(GenMatrix<UCBspl_real>& PHI, int noLevels) {
  GenMatrix<UCBspl_real> PHIrestricted;
  for (int kk = 0; kk < noLevels; kk++) {
    UCBspl::restrictCoeffsC2(PHI, PHIrestricted);
    PHI.swap(PHIrestricted);
  }
}

static void calculateDomain(const std::vector<double>& U,
                            const std::vector<double>& V,
                            std::vector<double>& domain) {
  
  if (domain.size() != 4)
    domain.resize(4);

  domain[0] = *std::min_element(U.begin(), U.end());
  domain[1] = *std::min_element(V.begin(), V.end());
  domain[2] = *std::max_element(U.begin(), U.end());
  domain[3] = *std::max_element(V.begin(), V.end());
}


void MGsystem::setDomain(double umin, double vmin, double umax, double vmax) {
  if (domain_.size() != 4)
    domain_.resize(4);

  domain_[0] = umin;
  domain_[1] = vmin;
  domain_[2] = umax;
  domain_[3] = vmax;
}

void MGsystem::getDomain(double& umin, double& vmin, double& umax, double& vmax) {
  umin = domain_[0];
  vmin = domain_[1];
  umax = domain_[2];
  vmax = domain_[3];
}

void MGsystem::reserveResidual() {
   
  int noLevels = systems_.size();
  int m_nextToFinest = (systems_[noLevels-1]->n1_-3)/2;
  int n_nextToFinest = (systems_[noLevels-1]->n2_-3)/2;
  r_.reserve(m_nextToFinest+3, n_nextToFinest+3); 
}

void MGsystem::initMultilevel(int m0, int n0, int h,
                              boost::shared_ptr<std::vector<double> > U,
                              boost::shared_ptr<std::vector<double> > V,
                              boost::shared_ptr<std::vector<double> > Zorig,
                              boost::shared_ptr<GenMatrix<UCBspl_real> > x, 
                              boost::shared_ptr<std::map<int, double> > weights) {
  
  

  
  if (domain_.size() != 4) 
    calculateDomain(*U,*V,domain_);

  m0_ = m0;
  n0_ = n0;

  m0 *= 8;
  n0 *= 8;
  h  -= 3;

  int noLevels = h+1;
  
  if (h < 1) {
    cout << "ERROR in input: At least two layers are required (h >= 4 in MBA) " << endl;
    exit(-1);
  }

  int m_finest = m0 * (2<<(h-1));
  int n_finest = n0 * (2<<(h-1));
  
  int noX = m_finest+3;
  int noY = n_finest+3;

#ifdef DEBUG_LSMG
  cout << "Number of unknowns coefficients: " << noX << 'x' << noY << '=' << noX*noY << endl;
#endif
  
  if (!x.get())
    x.reset(new GenMatrix<UCBspl_real> );


  if (x->noX() != noX || x->noY() != noY) 
    x->resize(noX,noY);
  

  systems_.resize(noLevels);
  

  systems_[noLevels-1] = new LSsystem(U,V,Zorig,m_finest+3,n_finest+3,x,lambdaFac_);
  const bool buildRHS=false; 
  systems_[noLevels-1]->setDomain(umin(),vmin(),umax(),vmax());
  systems_[noLevels-1]->setWeights(weights);
  systems_[noLevels-1]->buildEqSystem(buildRHS);
  
  int prev_noX = systems_[noLevels-1]->n1_;
  int prev_noY = systems_[noLevels-1]->n2_;

  e_arr_.resize(noLevels-1); 

  for (int i = noLevels-2; i >= 0; i--) {
    noX = (prev_noX-3)/2 + 3;
    noY = (prev_noY-3)/2 + 3;
    e_arr_[i].reset(new GenMatrix<UCBspl_real> );
    e_arr_[i]->resize(noX,noY); 

    systems_[i] = new LSsystem(U,V,Zorig,noX,noY,e_arr_[i],lambdaFac_);

    systems_[i]->setDomain(umin(),vmin(),umax(),vmax());


    systems_[i]->setWeights(weights);
    systems_[i]->buildEqSystem(buildRHS); 


    

    prev_noX = noX;
    prev_noY = noY;    
  }

  
  e_refined_.reserve(m_finest+3,n_finest+3);


#ifdef DEMUG_LSMG
  cout << "initializing done" << endl;
#endif
}


void MGsystem::cleanup() {
  int noLevels = systems_.size();
  for (int i = 0; i < noLevels; i++)
    delete systems_[i];
  systems_.clear();

  r_.clear();
  e_refined_.clear();

  domain_.clear();

  rhsFinestLevel_ = false;
}


static double average(const std::vector<double>& arr) {

  double sum=0.;
  std::vector<double>::const_iterator it;
  for(it = arr.begin(); it != arr.end(); ++it)
    sum += *it;

  return sum/arr.size();

}

void MGsystem::calculateStartVector() const {

#ifdef DEBUG_LSMG
  cout << "Calculating start vector with MBA" << endl;
#endif

  MBA mba(systems_[0]->U_, systems_[0]->V_, systems_[0]->Zvals_);
  mba.setDomain(umin(),vmin(),umax(),vmax());

  int h = 3;
  mba.MBAalg(m0_,n0_,h);

  boost::shared_ptr<GenMatrix<UCBspl_real> > phi =  mba.getSplineSurface().getCoefficients();
  
  systems_[0]->x_->init(*phi);
}

void MGsystem::solveFMG() {
  
  
  
  


  

  if (systems_.size() == 0)
    errorAndExit("the multilevel system has not been initialized");

  reserveResidual();

  systems_[0]->buildRHS();
  int noLevels = systems_.size();

  calculateStartVector(); 

  systems_[0]->relaxCG(noIterCoarsest_);
#ifdef DEBUG_LSMG
  cout << "Residual at coarsest level with CG,  r_l2= " << systems_[0]->residual_l2() << "   l_linf=" << systems_[0]->residual_linf() << endl;
#endif

  for (int ii = 1; ii < noLevels; ii++) {
    UCBspl::refineCoeffsC2(*systems_[ii-1]->x_,e_refined_);
    systems_[ii]->x_->swap(e_refined_);

    systems_[ii]->buildRHS();
    int noLevelsInVcycle = ii+1;
    relaxAndCorrectVcycle(noLevelsInVcycle);
  }
 

  rhsFinestLevel_ = true;
}

void MGsystem::solveAscend() {

  if (systems_.size() == 0)
    errorAndExit("the multilevel system has not been initialized");


  systems_[0]->buildRHS();
  int noLevels = systems_.size();

  calculateStartVector();
  
  systems_[0]->relaxCG(noIterCoarsest_);
  for (int i = 1; i < noLevels; i++) {
    UCBspl::refineCoeffsC2(*systems_[i-1]->x_,e_refined_);
    systems_[i]->x_->swap(e_refined_);
    systems_[i]->buildRHS();
    systems_[i]->relaxGaussSeidel(noIter_);
  }
  rhsFinestLevel_ = true;
}




  


















	  
	  






void MGsystem::solveVcycle() {
  cout << "MGsystem::solveVcycle..." << endl;

  if (systems_.size() == 0)
    errorAndExit("the multilevel system has not been initialized");

  reserveResidual();


  int noLevels = systems_.size();
  if (!rhsFinestLevel_)
    systems_[noLevels-1]->buildRHS();  

  relaxAndCorrectVcycle(noLevels);
  rhsFinestLevel_ = true;
}


void MGsystem::relax(int numberOfIterations, int relaxType) {

  if (systems_.size() == 0)
    errorAndExit("the multilevel system has not been initialized");
  
  int noLevels = systems_.size();
  
  if (!rhsFinestLevel_)
    systems_[noLevels-1]->buildRHS();
  rhsFinestLevel_ = true;

  if (relaxType == 1)
    systems_[noLevels-1]->relaxGaussSeidel(numberOfIterations);
  else if (relaxType == 2)
    systems_[noLevels-1]->relaxJacobi(numberOfIterations);
  else if (relaxType == 3)
    systems_[noLevels-1]->relaxCG(numberOfIterations);
  else if (relaxType == 4)
    systems_[noLevels-1]->relaxCGPrecond(numberOfIterations);
  else {
    cout << "No relaxion of that type: " << relaxType << endl;
    exit(-1);
  }
}

void MGsystem::setSmoothingFactor(double smoothingFac) {

  lambdaFac_ = smoothingFac;
  int noLevels = systems_.size();
  for (int i = 0; i < noLevels; i++) {
    systems_[i]->setSmoothingFactor(lambdaFac_);
  }
}

bool MGsystem::addPoint(double u, double v, double z, double weight) {
  
  int noLevels = systems_.size();
  const bool addToRHS=true; 
  for (int i = 0; i < noLevels; i++) {
    if (!systems_[i]->addPoint(u,v,z,weight,addToRHS))
      return false;
  }

  systems_[0]->U_->push_back(u);
  systems_[0]->V_->push_back(v);
  systems_[0]->Zvals_->push_back(z);
  
  if (weight != 1.0)
    (*systems_[0]->weights_.get())[systems_[0]->U_->size()-1] = weight;

  return true;
}


void MGsystem::setWeight(int pointIndex, double new_weight) {
  

  double old_weight = 1.0;
  bool wls = false;

  std::map<int, double>* weights = systems_[0]->weights_.get();
  if (weights && weights->size() != 0) {
    map<int, double, less<int> >::iterator mit = weights->find(pointIndex);
    if(mit != weights->end()) {
       old_weight = (*mit).second;
       wls = true;
    }
  }

  double u = (*systems_[0]->U_)[pointIndex];
  double v = (*systems_[0]->V_)[pointIndex];
  double z = (*systems_[0]->Zvals_)[pointIndex];
  int noLevels = systems_.size();

  const bool addToRHS = false;

  for (int i = 0; i < noLevels; i++)
    systems_[i]->addPoint(u,v,z,-old_weight,addToRHS);

  for (int i = 0; i < noLevels; i++)
    systems_[i]->addPoint(u,v,z,new_weight,addToRHS);

  (*systems_[0]->weights_.get())[pointIndex] = new_weight;

  rhsFinestLevel_ = false; 
}





 
