/*
 * assemble_coarse_matrix.h
 *
 *  Created on: Mar 10, 2020
 *      Author: mwichro
 */

#ifndef ASSEMBLE_COARSE_MATRIX_H_
#define ASSEMBLE_COARSE_MATRIX_H_


#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/logstream.h>
#include <deal.II/base/function.h>
#include <deal.II/base/utilities.h>

#include <deal.II/base/conditional_ostream.h>

#include <deal.II/base/parameter_handler.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/solver_bicgstab.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/lac/solver_gmres.h>
#include <deal.II/lac/affine_constraints.h>

#include <deal.II/lac/block_sparsity_pattern.h>
#include <deal.II/lac/trilinos_parallel_block_vector.h>
#include <deal.II/lac/trilinos_sparse_matrix.h>
#include <deal.II/lac/trilinos_block_sparse_matrix.h>
#include <deal.II/lac/trilinos_precondition.h>
#include <deal.II/lac/trilinos_solver.h>
#include <deal.II/grid/tria.h>

#include <deal.II/grid/grid_tools.h>
#include <deal.II/dofs/dof_handler.h>

#include <deal.II/fe/fe_q.h>
#include <deal.II/fe/fe_dgq.h>
#include <deal.II/fe/fe_dgp.h>
#include <deal.II/fe/fe_system.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/fe/mapping_q.h>
#include <deal.II/numerics/vector_tools.h>
#include <deal.II/numerics/matrix_tools.h>


#include <fstream>
#include <iostream>
#include <limits>
#include <locale>
#include <string>



using namespace dealii;



template <typename LevelNumber>
class CoarseDirectSolver
:
	public MGCoarseGridBase<LinearAlgebra::distributed::Vector<LevelNumber>>{

public:

	typedef LevelNumber number;
	typedef LinearAlgebra::distributed::Vector<LevelNumber> VectorType;

	CoarseDirectSolver(const IndexSet& owned_dofs_,
			const IndexSet& relevant_dofs_);

	template<class Operator>
	void initialize(const Operator & vmult_operator,
			const unsigned int& lvl=0);

	template<class Operator, int dim>
		void initialize(const Operator & vmult_operator,
				const DoFHandler<dim> &dof_handler,
				const unsigned int& lvl=0);

	void vmult(VectorType &dst,
			const VectorType &src) const{
		inverse.solve(dst, src);
	}

	virtual void
	  operator()(const unsigned int lvl,
	             VectorType &      dst,
	             const VectorType &src) const;
private:

	int level;

	IndexSet owned_dofs;
	IndexSet relevant_dofs;

	std::unique_ptr<TrilinosWrappers::SparseMatrix>  matrix;
	SolverControl solver_control;
	mutable TrilinosWrappers::SolverDirect inverse;

	template<int dim>
	void setup_matrix(const DoFHandler<dim>& dof_handler_u);

	template<class Operator>
	void do_assembly(const Operator & vmult_operator);
};

template < typename LevelNumber>
CoarseDirectSolver<LevelNumber>::CoarseDirectSolver(const IndexSet& owned_dofs_,
		const IndexSet& relevant_dofs_)
		: level(-1)
		  , owned_dofs(owned_dofs_)
		  , relevant_dofs(relevant_dofs_)
		  , solver_control(10, 1e-14)
		  , inverse(solver_control)
{}


template < typename LevelNumber>
template<class Operator>
void
CoarseDirectSolver<LevelNumber>::initialize(const Operator & vmult_operator
		, const unsigned int& lvl){
	matrix=std::make_unique<TrilinosWrappers::SparseMatrix> (owned_dofs,
					owned_dofs,
					MPI_COMM_WORLD,
					100);
	level=lvl;
	do_assembly( vmult_operator);
}



template < typename LevelNumber>
template<class Operator, int dim>
void
CoarseDirectSolver<LevelNumber>::initialize(const Operator & vmult_operator,
		const DoFHandler<dim> &dof_handler,
		const unsigned int& lvl){

	level=lvl;

		matrix=std::make_unique<TrilinosWrappers::SparseMatrix> ();

		TrilinosWrappers::SparsityPattern A_sparsity( owned_dofs,
				owned_dofs,
				relevant_dofs,
				MPI_COMM_WORLD,
				100 );

		typename DoFHandler<dim>::cell_iterator cell =
				dof_handler.begin(level),
				endc = dof_handler.end(level);

		std::vector<types::global_dof_index> local_dof_indices_u(dof_handler.get_fe().dofs_per_cell);

		for (; cell != endc; ++cell)
		{

			if(!cell->is_locally_owned_on_level())
				continue;

			cell->get_mg_dof_indices(local_dof_indices_u);

			for(types::global_dof_index & i : local_dof_indices_u ){
				A_sparsity.add_entries(i,
						local_dof_indices_u.begin(),
						local_dof_indices_u.end() );
			}
		}

		A_sparsity.compress();
		matrix->reinit(A_sparsity);


		do_assembly( vmult_operator);

}



template < typename LevelNumber>
template<class Operator>
void
CoarseDirectSolver< LevelNumber>::do_assembly(const Operator & vmult_operator){

    VectorType  dst;
    VectorType  src;

//    vmult_operator.get_matrix_free()->initialize_dof_vector(
//    			src, component);
    src.reinit(owned_dofs, relevant_dofs, MPI_COMM_WORLD);
    dst.reinit(owned_dofs, relevant_dofs, MPI_COMM_WORLD);


	for(types::global_dof_index i =0; i< owned_dofs.size(); ++i){
		src=0;
		dst=0;
		if(owned_dofs.is_element(i) )
			src(i)=1;

		src.compress(VectorOperation::insert);
		vmult_operator.vmult(dst, src);

		dst.update_ghost_values();

		for(IndexSet::ElementIterator  index_iter =owned_dofs.begin();
				index_iter != owned_dofs.end();
				++index_iter){
			if(dst(*index_iter)!=0 )
				matrix->set(*index_iter,i, dst(*index_iter) );
		}
	}

    matrix->compress(VectorOperation::unknown);


    inverse.initialize(*matrix);


}

template <typename LevelNumber>
void
CoarseDirectSolver< LevelNumber>::operator()(const unsigned int lvl,
	             VectorType &      dst,
	             const VectorType &src) const {
	Assert(lvl==level,ExcMessage("You are trying to use coarse solver on the wrong level"));
	this->vmult(dst,src);
}


#endif /* ASSEMBLE_COARSE_MATRIX_H_ */