#include <vector>

#include "UmfPackSolver.h"
#include "DOFMatrix.h"
#include "MatVecMultiplier.h"

#include "umfpack.h"


namespace AMDiS {

  template<typename VectorType>
  UmfPackSolver<VectorType>::UmfPackSolver(::std::string name) 
    : OEMSolver<VectorType>(name) 
  {}

  template<typename VectorType>
  UmfPackSolver<VectorType>::~UmfPackSolver() {}

  template<typename VectorType>
  int UmfPackSolver<VectorType>::solveSystem(MatVecMultiplier<VectorType> *matVec,
					     VectorType *x, VectorType *b)
  {
    FUNCNAME("UmfPackSolver::solveSystem()");

    TEST_EXIT(x->getSize() == b->getSize())("Vectors x and b must have the same size!");

    // Extract the matrix of DOF-matrices.
    StandardMatVec<Matrix<DOFMatrix*>, SystemVector> *stdMatVec = 
      dynamic_cast<StandardMatVec<Matrix<DOFMatrix*>, SystemVector> *>(matVec);
    Matrix<DOFMatrix*> *m = stdMatVec->getMatrix();

    // Number of systems.
    int nComponents = m->getSize();
    // Size of the new composed matrix.
    int newMatrixSize = ((*m)[0][0])->getFESpace()->getAdmin()->getUsedSize() * nComponents;

    // The new matrix has to be stored in compressed col format, therefore
    // the cols are collected.
    ::std::vector< ::std::vector< MatEntry > > cols(newMatrixSize, ::std::vector<MatEntry>(0));

    // Counter for the number of non-zero elements in the new matrix.
    int nElements = 0;

    for (int stencilRow = 0; stencilRow < nComponents; stencilRow++) {
      for (int stencilCol = 0; stencilCol < nComponents; stencilCol++) {

	if (!(*m)[stencilRow][stencilCol]) {
	  continue;
	}
	
	DOFMatrix::Iterator matrixRow((*m)[stencilRow][stencilCol], USED_DOFS);
 	int rowIndex = 0;
 	for (matrixRow.reset(); !matrixRow.end(); matrixRow++, rowIndex++) {
 	  for (int i = 0; i < static_cast<int>((*matrixRow).size()); i++) {	      
 	    if ((*matrixRow)[i].col >= 0) {
 	      MatEntry me;
	      me.entry = (*matrixRow)[i].entry;
	      // The col field is used to store the row number of the new element.
     	      me.col = (rowIndex  * nComponents) + stencilRow;      

	      // And save the new element in the corresponding column.
	      cols[((*matrixRow)[i].col * nComponents) + stencilCol].push_back(me);

	      nElements++;
 	    }
 	  }
 	}

      }
    }

    // Data fields for UMFPack.
    int *Ap = (int*)malloc(sizeof(int) * (newMatrixSize + 1));
    int *Ai = (int*)malloc(sizeof(int) * nElements);
    double *Ax = (double*)malloc(sizeof(double) * nElements);
    double *bvec = (double*)malloc(sizeof(double) * newMatrixSize);
    double *xvec = (double*)malloc(sizeof(double) * newMatrixSize);

    // Resort the right hand side of the linear system.
    for (int i = 0; i < b->getSize(); i++) {
      DOFVector<double>::Iterator it(b->getDOFVector(i), USED_DOFS);

      int counter = 0;
      for (it.reset(); !it.end(); ++it, counter++) {	
	bvec[counter * nComponents + i] = *it;
      }
    }

    // Create fields Ap, Ai and Ax.
    int elCounter = 0;
    Ap[0] = 0;    
    for (int i = 0; i < newMatrixSize; i++) {
      Ap[i + 1] = Ap[i] + cols[i].size();

      // The cols has to be sorted for using them in UMFPack.
      sort(cols[i].begin(), cols[i].end(), CmpMatEntryCol());

      for (int j = 0; j < static_cast<int>(cols[i].size()); j++) {
	Ai[elCounter] = cols[i][j].col;
	Ax[elCounter] = cols[i][j].entry;

	elCounter++;
      }     
    }

	   
    void *Symbolic, *Numeric;
    double Control[UMFPACK_CONTROL];
    double Info[UMFPACK_INFO];
    int status = 0;

    MSG("solving system with UMFPack ...\n");
    
    // Use default setings.
    umfpack_di_defaults(Control);
   
    // Run UMFPack.
    status = umfpack_di_symbolic(newMatrixSize, newMatrixSize, Ap, Ai, Ax, &Symbolic, Control, Info);   
    if (!status == UMFPACK_OK) {
      ERROR_EXIT("UMFPACK Error in function umfpack_di_symbolic");
    }    
 
    status = umfpack_di_numeric(Ap, Ai, Ax, Symbolic, &Numeric, Control, Info);
    if (!status == UMFPACK_OK) {
      ERROR_EXIT("UMFPACK Error in function umfpack_di_numberic");
    }
    
    status = umfpack_di_solve(UMFPACK_A, Ap, Ai, Ax, xvec, bvec, Numeric, Control, Info);
    if (!status == UMFPACK_OK) {
      ERROR_EXIT("UMFPACK Error in function umfpack_di_solve");
    }

    
    umfpack_di_free_symbolic(&Symbolic);
    umfpack_di_free_numeric(&Numeric);

    
    // Copy and resort solution.
    for (int i = 0; i < x->getSize(); i++) {
      DOFVector<double>::Iterator it(x->getDOFVector(i), USED_DOFS);
      
      int counter = 0;
      for (it.reset(); !it.end(); it++, counter++) {
	*it = xvec[counter * nComponents + i];
      }
    }

    free(Ap);
    free(Ai);
    free(Ax);
    free(bvec);
    free(xvec);

    // Calculate and print the residual.
    *p = *x;
    *p *= -1.0;
    matVec->matVec(NoTranspose, *p, *r);
    *r += *b;
    
    this->residual = norm(r);

    MSG("Residual: %e\n", this->residual);
      
    if (this->residual < this->tolerance) {
      ERROR_EXIT("UMFPACK could not solve the system!\n");
    }
    
    return(1);
  }
}