//Try to solve the Schrödinger Equation for the spherical Harmonic Oszi

#include <deal.II/base/logstream.h>
#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/function.h>
#include <deal.II/base/utilities.h>
#include <deal.II/grid/tria.h>
#include <deal.II/grid/grid_generator.h>
#include <deal.II/grid/grid_out.h>
#include <deal.II/grid/tria_accessor.h>
#include <deal.II/grid/tria_iterator.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/dofs/dof_accessor.h>
#include <deal.II/dofs/dof_tools.h>
#include <deal.II/dofs/dof_renumbering.h>
#include <deal.II/fe/fe_nedelec.h>
#include <deal.II/fe/fe_raviart_thomas.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/numerics/vector_tools.h>
#include <deal.II/numerics/matrix_tools.h>
#include <deal.II/numerics/data_out.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/base/index_set.h>
#include <deal.II/lac/petsc_sparse_matrix.h>
#include <deal.II/lac/petsc_parallel_vector.h>
#include <deal.II/lac/slepc_solver.h>
#include <deal.II/grid/manifold_lib.h>
#include <set>
#include <fstream>
#include <iostream>

namespace MR_fesc{
  using namespace dealii;
  template<int dim>
  class EigenvalueProblem{
  public:
    EigenvalueProblem();
    void run();
  private:
    //unsigned int num_eigs;
    
    void make_grid_and_dofs();
    void assemble_system();
    unsigned int solve();
    void output_results() const;

    Triangulation<dim> triang;
    //for TE modes:
    //FE_Nedelec<dim> fe;
    //for TM modes:
    FE_RaviartThomas<dim> fe;
    DoFHandler<dim> dof_h;

    PETScWrappers::SparseMatrix H,S;
    std::vector<PETScWrappers::MPI::Vector> eigenfunctions;
    std::vector<double> eigenvals;
    ConstraintMatrix constraints;
  };

  template<int dim>
  EigenvalueProblem<dim>::EigenvalueProblem():
    fe(1),dof_h(triang){}

  template<int dim>
  void EigenvalueProblem<dim>::make_grid_and_dofs(){
    GridGenerator::hyper_ball(triang); //default radius 1 around 0,0
    static const SphericalManifold<dim> bound;
    triang.set_all_manifold_ids_on_boundary(0);
    triang.set_manifold(0,bound);
    triang.refine_global(3);

    GridOut gridout;
    std::ofstream gridofstream("grid.eps");
    gridout.write_eps(triang,gridofstream);

    dof_h.distribute_dofs(fe);
    //renumbering, should benchmark this at some point
    //DoFRenumbering::Cuthill_McKee(dof_h);
    //DoFRenumbering::block_wise(dof_h);
    //BOUNDARY VALUES:
    //for TE Waves, we actually compute the electric field and
    //demand n x E = 0
    //VectorTools::project_boundary_values_curl_conforming_l2
    //  (dof_h,0,ZeroFunction<dim>(dim),0,constraints); //Mapping?
    //for TM Waves, we compute the magnetic field B and demand
    //dot(n,B) = 0
    VectorTools::project_boundary_values_div_conforming
      (dof_h,0,ZeroFunction<dim>(dim),0,constraints);
						    
    constraints.close();

    H.reinit(dof_h.n_dofs(),
	     dof_h.n_dofs(),
	     dof_h.max_couplings_between_dofs());
    S.reinit(dof_h.n_dofs(),
	     dof_h.n_dofs(),
	     dof_h.max_couplings_between_dofs());
    IndexSet eigs_ind_set = dof_h.locally_owned_dofs();
    eigenfunctions.resize(13);
    for(unsigned int i=0; i<eigenfunctions.size();++i){
      eigenfunctions[i].reinit(eigs_ind_set,MPI_COMM_WORLD);
    }
    eigenvals.resize(eigenfunctions.size());
  }

  // template<int dim>
  // std::vector<double> pot_func(const std::vector<Point<dim> > &pts,double scale){
  //   std::vector<double> retval(pts.size());
  //   for(int p=0; p<pts.size();++p){
  //     for(int i=0; i<dim; ++i){
  // 	retval[p]+=scale*pts[p][i]*pts[p][i];
  //     }
  //   }
  //   return retval;
  // }

  template<int dim>
  void EigenvalueProblem<dim>::assemble_system(){
    QGauss<dim> quadrature_formula(3);
    FEValues<dim> fe_values(fe,quadrature_formula,
			    update_values | update_gradients |
			    update_quadrature_points | update_JxW_values);

    const unsigned int dpc = fe.dofs_per_cell;
    const unsigned int nqp = quadrature_formula.size();
    FullMatrix<double> cell_H(dpc,dpc);
    FullMatrix<double> cell_S(dpc,dpc);

    std::vector<types::global_dof_index> local_dof_indices (dpc);

    //std::vector<double> potential_values(nqp);
    std::vector<typename internal::CurlType<dim>::type> curl_E(dpc); //curl_type is probs not defined here
    std::vector<Tensor<1,dim> > val_E(dpc); //FIXME: is that the right type?

    const FEValuesExtractors::Vector Efield(0);

    //weak form of Maxwell's Wave-Equation:
    //(rot(E),rot(v)) = omega_h (eps*E,v) for all v

    for(auto cell = dof_h.begin_active(); cell!=dof_h.end(); ++cell){
      fe_values.reinit(cell);
      cell_H=0;
      cell_S=0;
      //potential_values = pot_func(fe_values.get_quadrature_points(),500);
      for(unsigned int q_point=0; q_point<nqp; ++q_point){
	//just for debugging, output cell_H and cell_S
	//std::cout << "Cell " << q_point << "\n";
	for(unsigned int k=0; k<dpc;++k){
	  curl_E[k] = fe_values[Efield].curl(k,q_point);
	  val_E[k]  = fe_values[Efield].value(k,q_point);
	}
	for(unsigned int i=0; i<dpc; ++i){
	  for(unsigned int j=0; j<dpc; ++j){
	    cell_H(i,j)+=(curl_E[i]*curl_E[j]*
			  fe_values.JxW(q_point));
	    cell_S(i,j)+=(val_E[i]*val_E[j]*
			  fe_values.JxW(q_point));
	    //std::cout << 
	  }
	}
      }
      
      cell->get_dof_indices(local_dof_indices);
      constraints.distribute_local_to_global(cell_H,local_dof_indices,H);
      constraints.distribute_local_to_global(cell_S,local_dof_indices,S);
    }
    H.compress(VectorOperation::add);
    S.compress(VectorOperation::add);
    double min_spurious_eigenvalue = std::numeric_limits<double>::max(),
           max_spurious_eigenvalue = -std::numeric_limits<double>::max();

    for (unsigned int i = 0; i < dof_h.n_dofs(); ++i)
      if (constraints.is_constrained(i))
        {
          const double ev = H(i,i)/S(i,i);
          min_spurious_eigenvalue = std::min (min_spurious_eigenvalue, ev);
          max_spurious_eigenvalue = std::max (max_spurious_eigenvalue, ev);
        }

    std::cout << "   Spurious eigenvalues are all in the interval "
              << "[" << min_spurious_eigenvalue << "," << max_spurious_eigenvalue << "]"
              << std::endl;
  }

  template<int dim>
  unsigned int EigenvalueProblem<dim>::solve(){
    SolverControl solver_control(dof_h.n_dofs(),1e-9);
    SLEPcWrappers::SolverKrylovSchur eigensolver(solver_control);

    eigensolver.set_which_eigenpairs(EPS_SMALLEST_REAL);
    eigensolver.set_problem_type (EPS_GHEP);
    eigensolver.solve(H,S,eigenvals,eigenfunctions,eigenvals.size());
    for (unsigned int i=0; i<eigenfunctions.size(); ++i){
      eigenfunctions[i] /= eigenfunctions[i].linfty_norm ();
    }
    return solver_control.last_step ();
  }

  template <int dim>
  void EigenvalueProblem<dim>::output_results () const
  {
    DataOut<dim> data_out;

    data_out.attach_dof_handler (dof_h);

    for (unsigned int i=0; i<eigenfunctions.size(); ++i)
      data_out.add_data_vector (eigenfunctions[i],
                                std::string("eigenfunction_") +
                                Utilities::int_to_string(i));

    data_out.build_patches ();

    std::ofstream output ("eigenvectors.vtk");
    data_out.write_vtk (output);
  }

  template<int dim>
  void EigenvalueProblem<dim>::run(){
    make_grid_and_dofs();
    assemble_system();
    const unsigned int n_iterations = solve ();
    std::cout << "   Solver converged in " << n_iterations
              << " iterations." << std::endl;

    output_results ();

    std::cout << std::endl;
    for (unsigned int i=0; i<eigenvals.size(); ++i)
      std::cout << "      Eigenvalue " << i
                << " : " << eigenvals[i]
                << std::endl;
  }
}

int main(int argc, char **argv){
  using namespace MR_fesc;
  using namespace dealii;

  Utilities::MPI::MPI_InitFinalize mpi_initialization(argc, argv, 1);
  
  EigenvalueProblem<2> prob;
  prob.run();
  return 0;
}