//===========================================================================
// 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 <LSsystem.h>
#include <stdexcept>

#ifdef WIN32
#include <minmax.h>
#endif

#include <SmoothMatrix.h>

#include <UCBsplines.h>
#include <MBAclock.h> 

#define DEBUG_LSMG_0  1    

#include <algorithm>
using namespace std;

#include <MGsystem.h> 

static double average(std::vector<double> v) {
  
    std::vector<double>::const_iterator it;
    double sum = 0.0;
    for (it = v.begin(); it != v.end(); ++it) {
	sum += *it;
    }
    return sum/v.size();
}

LSsystem::LSsystem(boost::shared_ptr<std::vector<double> > U,
                   boost::shared_ptr<std::vector<double> > V,
                   boost::shared_ptr<std::vector<double> > Zvals,
                   int noX, int noY, 
                   boost::shared_ptr<GenMatrix<UCBspl_real> > x,
                   double lambdaFac)
    : U_(U), V_(V), Zvals_(Zvals), x_(x) {

    n1_ = noX ; n2_ = noY; 
    K_=L_=4;
    lambdaFac_ = lambdaFac;

    if (!x_.get())
	x_.reset(new GenMatrix<UCBspl_real> );

    if (x_->noX() != noX || x_->noY() != noY) {
	x_->resize(noX,noY);
	x_->fill(average(*Zvals));
    }
}


void LSsystem::calculateDomain() {
  
    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 LSsystem::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 LSsystem::getDomain(double& umin, double& vmin, double& umax, double& vmax) {

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


static int noNonZeros(int n1_, int n2_, int K_, int L_) {
    int numNonZeros = 0;
    int imin,imax,jmin,jmax;
    for(int j = 0; j < n2_; j++) {
	jmin = max(0,j-L_+1);
	jmax = min(n2_-1,j+L_-1);
    
	for(int i = 0; i < n1_; i++) {
	    imin = max(0,i-K_+1);
	    imax = min(n1_-1,i+K_-1);
	    numNonZeros += (imax - imin + 1) * (jmax - jmin +1);
	}
    }
    return numNonZeros;
}

void LSsystem::buildSparseStructure() {

  
    if (n1_ < 1 || n2_ < 1) {
	throw runtime_error("ERROR in buildSparseStructure, inproper initialization");
    }

  
    int numNonZeros = noNonZeros(n1_, n2_, K_, L_);
  
    A_.resize(n1_*n2_, n1_*n2_, numNonZeros);
  
    int offset = 0; 
    for(int j = 0; j < n2_; j++) {
	int jmin = max(0,j-L_+1);
	int jmax = min(n2_-1,j+L_-1);
	for(int i = 0; i < n1_; i++) {
	    int ii = j * n1_ + i; 
	    A_.irow(ii) = offset; 
	    int imin = max(0,i-K_+1);
	    int imax = min(n1_-1,i+K_-1);
	    for(int s = jmin; s <= jmax; s++) {
		for(int r = imin; r <= imax; r++) {
		    A_.jcol(offset) = s * n1_ + r;
		    A_(offset) = 0.0;
		    offset++;
		}
	    }
	}
    }
    A_.irow(n1_*n2_) = offset;
}

void LSsystem::calcLambda() {
  
  
    double sa_norm = A_.norml2();
  
    int m_ = n1_-3;
    int n_ = n2_-3;
  
    double du = (umax() - umin()) / (double)m_;
    double dv = (vmax() - vmin()) / (double)n_;
  
    SmoothMatrix E(m_+3,n_+3,du,dv);
  
    double e_norm = E.norm_l2();
    lambda_ = 1.0;
    if (e_norm != 0.0) {
	lambda_ = sa_norm/e_norm;
    } else {
	throw runtime_error("LSsystem::calcLambda(): e_norm is zero when "
			    "calculating lambda, exit now in Debug mode");
    }
}

void LSsystem::addSmoothingTerm(double lambdaVal) {
  
  
  
    int m_ = n1_-3;
    int n_ = n2_-3;  
    double du = (umax() - umin()) / (double)m_;
    double dv = (vmax() - vmin()) / (double)n_;
    SmoothMatrix E(m_+3,n_+3,du,dv);
  
    int noRows = n1_*n2_;
    for(int ii = 0; ii < noRows ; ii++) {
	int i = ii % n1_;
	int j = ii/n1_;
	for(int offset = A_.irow(ii); offset < A_.irow(ii+1); offset++) {
	    int jj = A_.jcol(offset); 
	    int r = jj % n1_;
	    int s = jj/n1_;
	    A_(offset) += lambdaVal * E(i,j,r,s);
	}
    }
}

void LSsystem::setSmoothingFactor(double newLambdaFac) {
    
    double lambdaAdjust = lambda_*(newLambdaFac - lambdaFac_);
  
    lambdaFac_ = newLambdaFac;

    addSmoothingTerm(lambdaAdjust);
}



void LSsystem::buildRHS() {
  

    if (domain_.size() != 4)
	calculateDomain();

    b_.resize(n1_, n2_);
    b_.fill(0.0);
  
    int noPoints = U_->size();  
    int m_ = n1_-3;
    int n_ = n2_-3;
    double u_scaling = double(m_) / (umax() - umin());
    double v_scaling = double(n_) / (vmax() - vmin());

    int i, j, k, l;
    double s, t;
    double w_kl[4][4];

    if (weights_.get() && weights_->size() != 0) {
	for (int ip = 0; ip < noPoints; ip++) {
	    double uc = ((*U_)[ip] - umin()) * u_scaling;
	    double vc = ((*V_)[ip] - vmin()) * v_scaling;
	    
	    UCBspl::ijst(m_, n_, uc, vc, i, j, s, t); 
	    
	    UCBspl::WKL(s, t, w_kl); 
	    
	    double zw = (*Zvals_)[ip];
	    
	    map<int, double, less<int> >::iterator mit = weights_->find(ip);
	    if(mit != weights_->end()) 
		zw *= (*mit).second;
	    for (k = 0; k <= 3; k++) {
		for (l = 0; l <=3; l++) {
		    b_(i+k,j+l) += w_kl[k][l] * zw;
		}
	    }
	}
    } else {
	for (int ip = 0; ip < noPoints; ip++) {
	    double uc = ((*U_)[ip] - umin()) * u_scaling;
	    double vc = ((*V_)[ip] - vmin()) * v_scaling;
	    
	    UCBspl::ijst(m_, n_, uc, vc, i, j, s, t); 
	    
	    UCBspl::WKL(s, t, w_kl); 
	    
	    double zw = (*Zvals_)[ip];
	    
	    for (k = 0; k <= 3; k++) {
		for (l = 0; l <=3; l++) {
		    b_(i+k,j+l) += w_kl[k][l] * zw;
		}
	    }
	}
    }
}


bool LSsystem::addPoint(double u, double v, double z, double weight, bool addToRHS) 
{  
  
    int m_ = n1_-3;
    int n_ = n2_-3;
  
    double uc = (u - umin())/(umax()-umin()) * (double)m_;
    double vc = (v - vmin())/(vmax()-vmin()) * (double)n_;
  
    if (uc < 0.0 || vc < 0.0 || uc > m_ || vc > n_) {
	throw runtime_error("ERROR in LSsystem::buildRHS: A point was mapped "
			    "to outside domain");
    }
  
    int i, j;
    double s, t;
    UCBspl::ijst(m_, n_, uc, vc, i, j, s, t); 
  
    double w_kl[4][4];
  
    UCBspl::WKL(s, t, w_kl); 

    int k, l;
    double weighted_z = z * weight;
    if (addToRHS) {
	for (k = 0; k <= 3; k++) {
	    for (l = 0; l <=3; l++) {
		b_(i+k,j+l) += w_kl[k][l] * weighted_z;
	    }
	}
    }
    
    for(k = 0; k < K_; k++) {
	for(l = 0; l < L_; l++) {
      
      
	    int ii = (j + l + 1) * n1_ + (i + k + 1); 
      
	    int imin = max(0,i+k - K_ + 2);
	    int jmin = max(0,j+l - L_ + 2);
	    int imax = min(n1_-1, i + k + K_ - 2 + 2);
      
	    
	    for(int kk = 0; kk < K_; kk++) {
		for(int ll = 0; ll < L_; ll++) {
          
		    int r = i+kk - K_ + 5;
		    int s = j+ll - L_ + 5;
          
		    int offset = A_.irow(ii) + (s - jmin)*(imax - imin + 1) + (r - imin);
		    
		    double contrib = w_kl[k][l]*w_kl[kk][ll] * weight;
		    A_(offset) += contrib;
		}
	    }
	}
    } 
  
    return true;
}


void LSsystem::buildEqSystem(bool buildRHS) {
  

    MBAclock rolex;

    if (domain_.size() != 4)
	calculateDomain();

    n1_ = x_->noX() ; n2_ = x_->noY();

    buildSparseStructure();
  
  
    if (buildRHS) {
	b_.resize(n1_, n2_);
	b_.fill(0.0);
    }
  
    int noPoints = U_->size();  
    double weight = 1.0;
    bool wls = (weights_.get() && weights_->size() != 0);

    map<int,double>::iterator mit;
    for (int ip = 0; ip < noPoints; ip++) {
	if (wls) {

	    mit = weights_->find(ip);
	    if (mit != weights_->end())
		weight = (*mit).second;
	    else
		weight = 1.0;
	}      
	addPoint((*U_)[ip], (*V_)[ip], (*Zvals_)[ip], weight, buildRHS);
    }
  
    calcLambda(); 
    addSmoothingTerm(lambda_*lambdaFac_);
}

double LSsystem::rowMatVecMult(int row_no, const MatSparse& A, const GenMatrix<UCBspl_real>& x) {

  
  
    double sum = 0.0;
  
    int n1 = x.noX();
    for (int offset = A.irow(row_no); offset < A.irow(row_no+1); offset++) { 
    
	int j = A.jcol(offset); 
	int k = j%n1 - 1; 
	int l = j/n1 - 1;
    
	sum += A(offset)*x(k,l);
    }
    return sum;
}

double LSsystem::currentSolutionNorm() const
{
    return sqrt((*x_).norm_2());
}

double LSsystem::residual_l2(bool scaled) const {
  

    double h=1.0;
    if (scaled) {
	int m_ = n1_-3;
	int n_ = n2_-3;

	double du = (umax() - umin()) / (double)m_;
	double dv = (vmax() - vmin()) / (double)n_;    
	h = du*dv;
    }  

    int n = n1_*n2_; 
    double norm2 = 0.0;
  
    for (int i = 0; i < n; i++) { 
    
	double sum = rowMatVecMult(i,A_,*x_);
    
	int k = i%n1_ - 1; 
	int l = i/n1_ - 1;
    
	double val = b_(k,l) - sum; 

	norm2 += val * val;

    }
    
    if (scaled) {
	norm2 *= h;
    }
  
    return sqrt(norm2);
}

double LSsystem::residual_linf() const {
  
  
    int n = n1_*n2_; 
    double norm = 0.0;
  
    for (int i = 0; i < n; i++) { 
    
	double sum = rowMatVecMult(i,A_,*x_);
    
	int k = i%n1_ - 1; 
	int l = i/n1_ - 1;
    
	double val = fabs(b_(k,l) - sum); 
	if (val > norm)
	    norm = val;
    }
  
    return norm;
}


void LSsystem::residual(GenMatrix<UCBspl_real>& r) const {
  
  
  
    if (r.noX() != n1_ || r.noY() != n2_)
	r.resize(n1_, n2_);
    r.fill(0.0); 
  
    int n = n1_*n2_; 
  
    for (int i = 0; i < n; i++) { 
    
	double sum = rowMatVecMult(i,A_,*x_);
    
	int k = i%n1_ - 1; 
	int l = i/n1_ - 1;
    
	r(k,l) = b_(k,l) - sum; 
    }
}


static double norm2(GenMatrix<UCBspl_real>& v) {
  
    int noX = v.noX();
    int noY = v.noY();

    double sp = 0.0;
    for (int j = -1; j <= noY-2; j++) {
	for (int i = -1; i <= noX-2; i++) {
	    double tmp = v(i,j);
	    sp += tmp*tmp;
	}
    }
    return sp;
}



static double scalarProduct(GenMatrix<UCBspl_real>& v1, GenMatrix<UCBspl_real>& v2) {

    if (&v1 == &v2)
	return norm2(v1);

    int noX = v1.noX();
    int noY = v1.noY();

    double sp = 0.0;
    for (int j = -1; j <= noY-2; j++) {
	for (int i = -1; i <= noX-2; i++) {
	    sp += v1(i,j)*v2(i,j);
	}
    }
    return sp;
}


static void matVecMult(const MatSparse& A, const GenMatrix<UCBspl_real>& P,
		       GenMatrix<UCBspl_real>& T) {

    

    int n1 = P.noX();
    int n2 = P.noY();
    if (T.noX() != n1 || T.noY() != n2)
	T.resize(n1, n2);
 
    int j = 0;

    int dimA = A.noRows();
    for (int row_no = 0; row_no < dimA; row_no++) {
	int k = j%n1 - 1; 
	int l = j/n1 - 1;
	T(k,l) = LSsystem::rowMatVecMult(row_no, A, P);
	j++;
    }
}


static void vecPlusScalVec_1(GenMatrix<UCBspl_real>& G, GenMatrix<UCBspl_real>& H,
			     double scalar) {


    int noX = G.noX();
    int noY = G.noY();

    for (int j = -1; j <= noY-2; j++) {
	for (int i = -1; i <= noX-2; i++) {
	    G(i,j) += scalar*H(i,j);
	}
    }
}
static void vecPlusScalVec_2(GenMatrix<UCBspl_real>& G, GenMatrix<UCBspl_real>& H,
			     double scalar) {


    int noX = G.noX();
    int noY = G.noY();

    for (int j = -1; j <= noY-2; j++) {
	for (int i = -1; i <= noX-2; i++) {
	    double tmp = G(i,j) + scalar*H(i,j);
	    H(i,j) = tmp;
	}
    }
}


void LSsystem::relaxGaussSeidel(int noIterations) 
{
    if (noIterations == 0)
	return;

    int j;
    int num_unknowns = numberOfUnknowns();

    GenMatrix<UCBspl_real>& x = *x_;

    for (int n_it = 1; n_it <= noIterations; n_it++) {
    
	for (int i = 0; i < num_unknowns; i++) { 
	    double sumL = 0.0;
      
	    int offset; 
	    for (offset = A_.irow(i); (j=A_.jcol(offset)) < i; offset++) {
        
		int l = j/n1_;
		int k = j - (n1_ * l) - 1;  
		l -= 1;
        
		sumL += A_(offset) * x(k,l);
	    }
      
	    double sumR = 0.0;
      
	    int offset_diag = offset;
	    offset++; 
      
	    for ( ;offset < A_.irow(i+1); offset++) {
		j = A_.jcol(offset);

		int l = j/n1_;
		int k = j - (n1_ * l) - 1;  
		l -= 1;

		sumR += A_(offset) * x(k,l);
	    }
      

	    int l = i/n1_;
	    int k = i - (n1_ * l) - 1; 
	    l -= 1;

	    x(k,l) = (b_(k,l) - sumL - sumR)/A_(offset_diag);
      
      
	} 
    } 
}

void LSsystem::relaxCG(int noIterations) {

  



    GenMatrix<UCBspl_real> r_k(n1_,n2_); 
    residual(r_k);

    GenMatrix<UCBspl_real> p_k;
    p_k.init(r_k);

    double rho_k = scalarProduct(r_k,r_k);

    int k=0;
  
    GenMatrix<UCBspl_real> t_k(n1_,n2_);



    double eps_num = 1.0e-15;
    while (k < noIterations && rho_k > eps_num) {

	::matVecMult(A_, p_k, t_k);
  
	double alpha_k = rho_k/scalarProduct(p_k,t_k);

	::vecPlusScalVec_1(*x_, p_k, alpha_k);

	::vecPlusScalVec_1(r_k, t_k, -1.0*alpha_k);

	double rho_kPlus1 = scalarProduct(r_k,r_k);


	double rho_fac = rho_kPlus1/rho_k;
	::vecPlusScalVec_2(r_k, p_k, rho_fac);

	k += 1;

	rho_k = rho_kPlus1;
    }
}


void LSsystem::relaxCGPrecond(int noIterations) {


  




    GenMatrix<UCBspl_real> r_k(n1_,n2_); 
    residual(r_k);
  
    boost::shared_ptr<GenMatrix<UCBspl_real> > s_k_temp(new GenMatrix<UCBspl_real> (n1_,n2_));


    s_k_temp->fill(0.0);
    MGsystem mgsys;
    boost::shared_ptr<std::vector<double> > dummy1;
    boost::shared_ptr<std::map<int,double> > dummy2;
    mgsys.initMultilevel(1,1,6,U_,V_,dummy1,s_k_temp,dummy2); 
    mgsys.qweSetRHS(r_k); 
    mgsys.solveVcycle();

    GenMatrix<UCBspl_real> s_k(n1_,n2_);
    s_k.swap(*s_k_temp); 
 

    GenMatrix<UCBspl_real> p_k;
    p_k.init(s_k);


    double rho_k = scalarProduct(s_k,r_k);

    int k=0;
  
    GenMatrix<UCBspl_real> t_k(n1_,n2_);
  
    GenMatrix<UCBspl_real> w_k(n1_,n2_);

    while (k < noIterations) {

	::matVecMult(A_, p_k, t_k);


	s_k_temp->fill(0.0);

    
	mgsys.qweSetRHS(t_k); 
	mgsys.solveVcycle();

	w_k.swap(*s_k_temp); 
    
	double alpha_k = rho_k/scalarProduct(p_k,t_k);

	::vecPlusScalVec_1(*x_, p_k, alpha_k);

	::vecPlusScalVec_1(r_k, t_k, -1.0*alpha_k);

	::vecPlusScalVec_1(s_k, w_k, -1.0*alpha_k);


	double rho_kPlus1 = scalarProduct(s_k,r_k);


	double rho_fac = rho_kPlus1/rho_k;
	::vecPlusScalVec_2(s_k, p_k, rho_fac);

	k += 1;

	rho_k = rho_kPlus1;
    }
}


void LSsystem::relaxJacobi(int noIterations, double omega) 
{
    if (noIterations == 0)
	return;

    GenMatrix<UCBspl_real> x_curr(x_->noX(), x_->noY()); 

    int j;
    for (int n_it = 1; n_it <= noIterations; n_it++) { 
	for (int i = 0; i < numberOfUnknowns(); i++) { 
	    double sumL = 0.0;
      
	    int offset;
	    for (offset = A_.irow(i); (j=A_.jcol(offset)) < i; offset++) {
        
		int k = j%n1_ - 1; 
		int l = j/n1_ - 1;
        
		sumL += A_(offset) * (*x_)(k,l);
	    }
      
	    double sumR = 0.0;
      
	    int offset_diag = offset;
	    offset++; 
      
	    for ( ;offset < A_.irow(i+1); offset++) {
		j = A_.jcol(offset);
		int k = j%n1_ - 1; 
		int l = j/n1_ - 1;
		sumR += A_(offset) * (*x_)(k,l);
	    }
      
	    int k = i%n1_ - 1; 
	    int l = i/n1_ - 1;

	    x_curr(k,l) = (b_(k,l) - sumL - sumR)/A_(offset_diag);
	} 



	if (omega != 1.0) {
	    int idx = x_->noX()-2;
	    int idy = x_->noY()-2;
	    for (int k = -1; k <= idx; k++)
		for (int l = -1; l <= idy; l++)
		    (*x_)(k,l) = (1.0-omega) * (*x_)(k,l) + omega*(x_curr(k,l));
	}
	else
	    x_->swap(x_curr);
    } 
}
 
