#include <algorithm>
#include <png.h>

#include "DOFMatrix.h"
#include "QPsiPhi.h"
#include "BasisFunction.h"
#include "Boundary.h"
#include "DOFAdmin.h"
#include "ElInfo.h"
#include "FiniteElemSpace.h"
#include "Mesh.h"
#include "DOFVector.h"
#include "Operator.h"
#include "BoundaryCondition.h"
#include "BoundaryManager.h"
#include "ElementMatrix.h"
#include "Assembler.h"
#include "Utilities.h"

#include <boost/numeric/mtl/mtl.hpp>

namespace AMDiS {

  using namespace mtl;

  DOFMatrix *DOFMatrix::traversePtr = NULL;

  DOFMatrix::DOFMatrix()
    : rowFESpace(NULL),
      colFESpace(NULL),
      elementMatrix(NULL),
      inserter(NULL)
  {}

  DOFMatrix::DOFMatrix(const FiniteElemSpace* rowFESpace_,
		       const FiniteElemSpace* colFESpace_,
		       std::string name_)
    : rowFESpace(rowFESpace_),
      colFESpace(colFESpace_),
      name(name_), 
      coupleMatrix(false),
      inserter(NULL)
  {
    TEST_EXIT(rowFESpace)("no rowFESpace\n");
  
    if (!colFESpace)
      colFESpace = rowFESpace;

    if (rowFESpace && rowFESpace->getAdmin())
      (const_cast<DOFAdmin*>(rowFESpace->getAdmin()))->addDOFIndexed(this);

    boundaryManager = NEW BoundaryManager(rowFESpace_);
    elementMatrix = NEW ElementMatrix(rowFESpace->getBasisFcts()->getNumber(),
				      colFESpace->getBasisFcts()->getNumber());

    applyDBCs.clear();

#ifdef HAVE_PARALLEL_AMDIS
    applicationOrdering = NULL;
#endif
  }

  DOFMatrix::DOFMatrix(const DOFMatrix& rhs)
    : name(rhs.name + "copy")
  {
    *this = rhs;
    if (rowFESpace && rowFESpace->getAdmin())
      (const_cast<DOFAdmin*>( rowFESpace->getAdmin()))->addDOFIndexed(this);

    TEST_EXIT(rhs.inserter == 0)("Cannot copy during insertion");
    inserter= 0;
  }

  DOFMatrix::~DOFMatrix()
  {
    FUNCNAME("DOFMatrix::~DOFMatrix()");
    if (rowFESpace && rowFESpace->getAdmin()) {
      (const_cast<DOFAdmin*>(rowFESpace->getAdmin()))->removeDOFIndexed(this);
    }  

    if (boundaryManager)
      DELETE boundaryManager;

    if (elementMatrix)
      DELETE elementMatrix;
    if (inserter) delete inserter;
  }

  void DOFMatrix::print() const
  {
    FUNCNAME("DOFMatrix::print()");

    using mtl::tag::major; using mtl::tag::nz; using mtl::begin; using mtl::end;
    namespace traits= mtl::traits;
    typedef base_matrix_type Matrix;

    traits::row<Matrix>::type                                 row(matrix);
    traits::col<Matrix>::type                                 col(matrix);
    traits::const_value<Matrix>::type                         value(matrix);

    typedef traits::range_generator<major, Matrix>::type      cursor_type;
    typedef traits::range_generator<nz, cursor_type>::type    icursor_type;
    
    for (cursor_type cursor = begin<major>(matrix), cend = end<major>(matrix); cursor != cend; ++cursor) {
	for (icursor_type icursor = begin<nz>(cursor), icend = end<nz>(cursor); icursor != icend; ++icursor)
	    Msg::print(" (%3d,%3d,%20.17lf)", row(*icursor), col(*icursor), value(*icursor));
	Msg::print("\n");
    }
  }

  bool DOFMatrix::symmetric()
  {
    FUNCNAME("DOFMatrix::symmetric()");

    double tol = 1e-5;

    using mtl::tag::major; using mtl::tag::nz; using mtl::begin; using mtl::end;
    namespace traits= mtl::traits;
    typedef base_matrix_type   Matrix;

    traits::row<Matrix>::type                                 row(matrix);
    traits::col<Matrix>::type                                 col(matrix);
    traits::const_value<Matrix>::type                         value(matrix);

    typedef traits::range_generator<major, Matrix>::type      cursor_type;
    typedef traits::range_generator<nz, cursor_type>::type    icursor_type;
    
    for (cursor_type cursor = begin<major>(matrix), cend = end<major>(matrix); cursor != cend; ++cursor)
      for (icursor_type icursor = begin<nz>(cursor), icend = end<nz>(cursor); icursor != icend; ++icursor)
	// Compare each non-zero entry with its transposed
	if (abs(value(*icursor) - matrix[col(*icursor)][row(*icursor)]) > tol)
	  return false;
    return true;
  }

  void DOFMatrix::test()
  {
    FUNCNAME("DOFMatrix::test()");

    int non_symmetric = !symmetric();

    if (non_symmetric) {
      MSG("matrix `%s' not symmetric.\n", name.data());
    } else {
      MSG("matrix `%s' is symmetric.\n", name.data());
    }
  }

  DOFMatrix& DOFMatrix::operator=(const DOFMatrix& rhs)
  {
    rowFESpace = rhs.rowFESpace;
    colFESpace = rhs.colFESpace;
    operators = rhs.operators;
    operatorFactor = rhs.operatorFactor;
    coupleMatrix = rhs.coupleMatrix;

    /// The matrix values may only be copyed, if there is no active insertion.
    if (rhs.inserter == 0 && inserter == 0)
      matrix = rhs.matrix;

    if (rhs.boundaryManager) {
      boundaryManager = NEW BoundaryManager(*rhs.boundaryManager);
    } else {
      boundaryManager = NULL;
    }

    if (rhs.elementMatrix) {
      elementMatrix = NEW ElementMatrix(rowFESpace->getBasisFcts()->getNumber(),
					colFESpace->getBasisFcts()->getNumber());
    } else {
      elementMatrix = NULL;
    }

    return *this;
  }

  void DOFMatrix::addElementMatrix(double sign, 
				   const ElementMatrix &elMat, 
				   const BoundaryType *bound,
				   bool add)
  {
    FUNCNAME("DOFMatrix::addElementMatrix()");

    TEST_EXIT_DBG(inserter)("DOFMatrix is not in insertion mode");
    inserter_type &ins= *inserter;

    DegreeOfFreedom row, col;
    double entry;

    int n_row = elMat.rowIndices.getSize();
    int n_col = elMat.colIndices.getSize();

    for (int i = 0; i < n_row; i++)  {   // for all rows of element matrix
      row = elMat.rowIndices[i];

      BoundaryCondition *condition = 
	bound ? boundaryManager->getBoundaryCondition(bound[i]) : NULL;

      if (condition && condition->isDirichlet()) {
	if (!coupleMatrix)
	  applyDBCs.insert(static_cast<int>(row));
      } else
	for (int j = 0; j < n_col; j++) {  // for all columns
	  col = elMat.colIndices[j];
	  entry = elMat[i][j];
	  if (add)
	    ins[row][col]+= sign * entry;
	  else
	    ins[row][col]= sign * entry;
	}   
    }
  }

  double DOFMatrix::logAcc(DegreeOfFreedom a, DegreeOfFreedom b) const
  {
    return matrix[a][b];
  }

  void DOFMatrix::freeDOFContent(int index)
  {
      if (matrix.nnz() == 0) return;

    using mtl::tag::major; using mtl::tag::nz; using mtl::begin; using mtl::end;
    namespace traits= mtl::traits;
    typedef base_matrix_type Matrix;

    traits::row<Matrix>::type                                 row(matrix);
    traits::col<Matrix>::type                                 col(matrix);

    typedef traits::range_generator<major, Matrix>::type      cursor_type;
    typedef traits::range_generator<nz, cursor_type>::type    icursor_type;
    
    cursor_type cursor = begin<major>(matrix);
    // Jump directly to corresponding row or column
    cursor+= index;

    // Requires structural symmetry !!!
    for (icursor_type icursor = begin<nz>(cursor), icend = end<nz>(cursor); icursor != icend; ++icursor) {
      int my_row= row(*icursor), my_col= col(*icursor);
      // Not very efficient (but general)
      matrix.lvalue(my_row, my_col) = 0.0; // Need to call crop somewhere !!!! Peter
      matrix.lvalue(my_col, my_row) = 0.0;
    }
  }

  // Should work as before
  void DOFMatrix::assemble(double factor, ElInfo *elInfo, const BoundaryType *bound)
  {
    FUNCNAME("DOFMatrix::assemble()");

    operators[0]->getAssembler(omp_get_thread_num())->
      initElementMatrix(elementMatrix, elInfo);

    std::vector<Operator*>::iterator it = operators.begin();
    std::vector<double*>::iterator factorIt = operatorFactor.begin();
    for (; it != operators.end(); ++it, ++factorIt) {
      if ((*it)->getNeedDualTraverse() == false) {
	  (*it)->getElementMatrix(elInfo,
				  elementMatrix, 
				  *factorIt ? **factorIt : 1.0);
      }
    }        

    addElementMatrix(factor, *elementMatrix, bound);   
  }

  void DOFMatrix::assemble(double factor, ElInfo *elInfo, const BoundaryType *bound,
			   Operator *op)
  {
      FUNCNAME("DOFMatrix::assemble()");

      TEST_EXIT_DBG(op)("No operator!\n");

      op->getAssembler(omp_get_thread_num())->initElementMatrix(elementMatrix, elInfo);
      op->getElementMatrix(elInfo, elementMatrix, factor);
      addElementMatrix(factor, *elementMatrix, bound);
  }

  void DOFMatrix::assemble(double factor, 
			   ElInfo *rowElInfo, ElInfo *colElInfo,
			   ElInfo *smallElInfo, ElInfo *largeElInfo,
			   const BoundaryType *bound, Operator *op)
  {
    FUNCNAME("DOFMatrix::assemble()");

    if(!op && operators.size() == 0) {
      return;
    }

    Operator *operat = op ? op : operators[0];
    operat->getAssembler(omp_get_thread_num())->
      initElementMatrix(elementMatrix, rowElInfo, colElInfo);
    
    if (op) {
      op->getElementMatrix(rowElInfo, colElInfo, smallElInfo, largeElInfo,
			   elementMatrix);
    } else {
      std::vector<Operator*>::iterator it = operators.begin();
      std::vector<double*>::iterator factorIt = operatorFactor.begin();
      for (; it != operators.end(); ++it, ++factorIt) {
	(*it)->getElementMatrix(rowElInfo, colElInfo,
				smallElInfo, largeElInfo,
				elementMatrix, 
				*factorIt ? **factorIt : 1.0);	
      }      
    }

    addElementMatrix(factor, *elementMatrix, bound);       
  }

  void DOFMatrix::assemble2(double factor, 
			    ElInfo *mainElInfo, ElInfo *auxElInfo,
			    ElInfo *smallElInfo, ElInfo *largeElInfo,			    
			   const BoundaryType *bound, Operator *op)
  {
    FUNCNAME("DOFMatrix::assemble2()");

    if (!op && operators.size() == 0) {
      return;
    }

    Operator *operat = op ? op : operators[0];
    operat->getAssembler(omp_get_thread_num())->
      initElementMatrix(elementMatrix, mainElInfo);
    
    if(op) {
      ERROR_EXIT("TODO");
//       op->getElementMatrix(rowElInfo, colElInfo, 
// 			   smallElInfo, largeElInfo,
// 			   elementMatrix);
    } else {
      std::vector<Operator*>::iterator it;
      std::vector<double*>::iterator factorIt;
      for(it = operators.begin(), factorIt = operatorFactor.begin();	
	   it != operators.end(); 
	   ++it, ++factorIt) {
	if ((*it)->getNeedDualTraverse()) {
	  (*it)->getElementMatrix(mainElInfo, auxElInfo,
				  smallElInfo, largeElInfo,
				  elementMatrix, 
				  *factorIt ? **factorIt : 1.0);
	}
      }      
    }

    addElementMatrix(factor, *elementMatrix, bound);       
  }

  void DOFMatrix::finishAssembling()
  {
    // call the operatos cleanup procedures
    for (std::vector<Operator*>::iterator it = operators.begin();
	 it != operators.end();
	 ++it) {     
      (*it)->finishAssembling();
    }
  }

  // Should work as before
  Flag DOFMatrix::getAssembleFlag()
  {
    Flag fillFlag(0);
    std::vector<Operator*>::iterator op;
    for(op = operators.begin(); op != operators.end(); ++op) {
      fillFlag |= (*op)->getFillFlag();
    }

    return fillFlag;
  }

  void DOFMatrix::axpy(double a, const DOFMatrix& x, const DOFMatrix& y)
  {
    matrix+= a * x.matrix + y.matrix;
  }


  void DOFMatrix::scal(double b) 
  {
    matrix*= b;
  }

  void DOFMatrix::addOperator(Operator *op, double* factor, double* estFactor) 
  { 
    operators.push_back(op);
    operatorFactor.push_back(factor);
    operatorEstFactor.push_back(estFactor);
  }

  void DOFMatrix::copy(const DOFMatrix& rhs) 
  {
    matrix= rhs.matrix;
  }

  void DOFMatrix::removeRowsWithDBC(std::set<int> *rows)
  {      
    inserter_type &ins= *inserter;
   
    for (std::set<int>::iterator it = rows->begin(); it != rows->end(); ++it)
      ins[*it][*it] = 1.0;

    rows->clear();
  }

  void DOFMatrix::createPictureFile(const char* filename, int dim)
  {
    TEST_EXIT(0)("Not yet re-implemented.");

#if 0
    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, 
						  NULL, NULL, NULL);

    if (!png_ptr)
       return;

    png_bytep rowPointers[dim];
    for (int i = 0; i < dim; i++) {
      rowPointers[i] = (png_byte*)png_malloc(png_ptr, dim);

      for (int j = 0; j < dim; j++)
	rowPointers[i][j] = 255;
    }

    double scalFactor = static_cast<double>(dim) / static_cast<double>(matrix.size());

    for (int i = 0; i < static_cast<int>(matrix.size()); i++) {    
      int pi = static_cast<int>(static_cast<double>(i) * scalFactor);

      TEST_EXIT_DBG((pi >= 0) && (pi < dim))("PI");

      for (int j = 0; j < static_cast<int>(matrix[i].size()); j++) {
	
	int pj = static_cast<int>(static_cast<double>(matrix[i][j].col) * scalFactor);

	TEST_EXIT_DBG((pj >= 0) && (pj < dim))("PJ");

	rowPointers[pi][pj] = 0;
      }
    }
   
    FILE *fp = fopen(filename, "wb");
    TEST_EXIT(fp)("Cannot open file for writing matrix picture file!\n");

    png_infop info_ptr = png_create_info_struct(png_ptr);

    if (!info_ptr) {
       png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
       return;
    }

    png_init_io(png_ptr, fp);

    png_set_IHDR(png_ptr, info_ptr, dim, dim, 8,
		 PNG_COLOR_TYPE_GRAY, 
		 PNG_INTERLACE_NONE,
		 PNG_COMPRESSION_TYPE_DEFAULT,
		 PNG_FILTER_TYPE_DEFAULT);

    png_set_rows(png_ptr, info_ptr, rowPointers);

    png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

    png_destroy_write_struct(&png_ptr, &info_ptr);

    fclose(fp);
#endif
  }


  int DOFMatrix::memsize() 
  {   
    return (num_rows(matrix) + matrix.nnz()) * sizeof(base_matrix_type::size_type)
      + matrix.nnz() * sizeof(base_matrix_type::value_type);
  }

}