// SpaFEDTe, a Template based C++ library for creating 
// Discontinuous Finite Element Spaces,
// Copyright (C) 2012 Lorenzo Alessio Botti

/* This library is free software; you can redistribute it and/or */
/* modify it under the terms of the GNU Lesser General Public */ 
/* License as published by the Free Software Foundation either */ 
/* version 3.0 of the License, or (at your option) any later version. */

/* This software 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 */
/* Lesser General Public License for more details. */

/* You should have received a copy of the GNU Lesser General Public */
/* License along with this software; if not, a copy of the full */
/* GNU Lesser General Public License can be found at */
/* http://www.gnu.org/licenses/. */

#include <iostream>
#include <assert.h>
#include <future>

#include "CommonTemplateTypes.hpp"
#include "InputData.hpp"
#include "Mesh.hpp"
#include "ElementIterator.hpp"
#include "ReadMeshVtu.hpp"
#include "LapAssembler.hpp"
#include "ExactSolution.hpp"
#include "VTKWriterSpaMesh.hpp"
#include "Timer.hpp"

#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Core>
#include <Eigen/Sparse>

//#define VIENNACL_WITH_EIGEN 1
//
//#include "viennacl/linalg/gmres.hpp"
//#include "viennacl/linalg/cg.hpp"
//#include "viennacl/linalg/bicgstab.hpp"
//#include "viennacl/tools/tools.hpp"
//#include "viennacl/meta/result_of.hpp"
//#include "viennacl/traits/size.hpp"
//#include "viennacl/linalg/jacobi_precond.hpp"

typedef CommonTemplateTypes::GlobalMatrixType GlobalMatrixType;
typedef CommonTemplateTypes::GlobalVectorType GlobalVectorType;
typedef CommonTemplateTypes::FEManagerType FEManagerType;
typedef CommonTemplateTypes::LiftingOperatorType LiftingOperatorType;
typedef CommonTemplateTypes::MeshType MeshType;

int main( int argc, char* argv[] )
{
  FELib::FELibInit init(argc, argv);
  Eigen::initParallel();
  Timer tmr;
  std::vector<double> l2ErrorNorms;

  for (std::string meshFile : InputData::meshSequence)
  {
    MeshType mesh;
    ReadMeshVtu vtuReader;
    //vtuReader.load_mesh("../../meshes_vtu/square_01/square_01_512.vtu", mesh);
    vtuReader.load_mesh(meshFile.c_str(), mesh);
    bool computeFaces = false;
    mesh.computeElemNeighbors(computeFaces);
    mesh.printInfo();
    std::cout<<"Polynomial degree is "<<InputData::p_order<<", DOFs per element "<<InputData::n_dofs<<std::endl;
    
    DofMap dofMap;
    FEManagerType FEManager;
    dofMap.setPOrder(InputData::p_order);
    dofMap.setNVars(1);
    dofMap.computeMatrixSparsity(mesh);

    GlobalMatrixType matrix(dofMap.getGlobalElementNumber()*InputData::n_dofs, dofMap.getGlobalElementNumber()*InputData::n_dofs);
    Eigen::VectorXi matrixNonZeroes(dofMap.getGlobalElementNumber()*InputData::n_dofs);
    for (unsigned int i = 0; i < dofMap.getGlobalElementNumber(); i++)
      matrixNonZeroes.segment<InputData::n_dofs>(i*InputData::n_dofs).noalias() = Eigen::VectorXi::Constant(InputData::n_dofs, dofMap.getLocalNonZeroBlocks()[i]*InputData::n_dofs);
    matrix.reserve(matrixNonZeroes);
    GlobalVectorType rhs(dofMap.getGlobalElementNumber()*InputData::n_dofs);
    rhs.setZero();
    GlobalVectorType solution(dofMap.getGlobalElementNumber()*InputData::n_dofs);


    tmr.reset();
    //double elemsPerThread = static_cast<double>(dofMap.getGlobalElementNumber())/static_cast<double>(InputData::nThreads);
    //std::future<void> voidF[InputData::nThreads];
    //for (unsigned int i = 0; i < InputData::nThreads; i++)
    //{
    //  std::pair<int,int> threadElemRange = std::make_pair(std::round(elemsPerThread*i),std::round(elemsPerThread*(i+1)));
    //  std::cout<<"Thread "<<i<<" elem range: "<<threadElemRange.first<<"-"<<threadElemRange.second<<std::endl;
    //  voidF[i] = (std::async(std::launch::async, [&mesh, &matrix, &rhs, FEManager, threadElemRange] () mutable -> void 
    //                                             { LapAssembler::assemble_elem_contribution(mesh, matrix, rhs, FEManager, threadElemRange); }));
    //}
    //for (unsigned int i = 0; i < InputData::nThreads; i++)
    //  voidF[i].get();
    //std::pair<int,int> elemRange = std::make_pair(0, mesh.elemCard());
    //LapAssembler::assemble_elem_contribution(mesh, matrix, rhs, FEManager, elemRange);
    LapAssembler::assemble_volume_contribution(mesh, matrix, rhs, FEManager);
    LapAssembler::assemble_boundary_face_contribution(mesh, matrix, rhs, FEManager);
    LapAssembler::assemble_internal_face_contribution(mesh, matrix, rhs, FEManager);
    double assembly_time = tmr.elapsed();
    if (FELib::processor_id() == 0)
      std::cout<<"rank "<<FELib::processor_id()<<" assembly time "<<assembly_time<<std::endl;

    tmr.reset();
    //Eigen::SparseLU<GlobalMatrixType> solver;
    //// Compute the ordering permutation vector from the structural pattern of A
    //solver.analyzePattern(matrix);
    //// Compute the numerical factorization
    //solver.factorize(matrix);
    ////Use the factors to solve the linear system
    //solution = solver.solve(rhs);

    //Eigen::BiCGSTAB<GlobalMatrixType, Eigen::IncompleteLUT<double> > solver;
    Eigen::ConjugateGradient<GlobalMatrixType> solver;
    solver.setTolerance(1.e-10);
    solver.compute(matrix);
    solution = solver.solve(rhs);
    std::cout << "Eigen Linear solver: #iterations: " << solver.iterations() << ", estimated error: " << solver.error() << std::endl;
    //std::cout << "----- Running CG -----" << std::endl;
    //viennacl::linalg::cg_tag my_cg_tag(1e-1, 2000);
    //solution = viennacl::linalg::solve(matrix, rhs, my_cg_tag);

    double solution_time = tmr.elapsed();
    if (FELib::processor_id() == 0)
    {
      std::cout<<"rank "<<FELib::processor_id()<<", Solution time "<<solution_time<<std::endl;
      std::cout<<"total step time "<<assembly_time+solution_time<<std::endl;
    }
    const Eigen::Map<const Eigen::Matrix<double,Eigen::Dynamic,InputData::n_dofs,InputData::MajorOption(InputData::p_order)> > 
                     solutionBlocked(solution.data(), mesh.elemCard(), InputData::n_dofs); 
    {
      double l2_error = 0;
      const FEManagerType::ShapeMatrixType& shapeMatrix           = FEManager.getElemShapeFunction();
      const FEManagerType::ElemJacobianVectorType& jacXweightVec  = FEManager.getElemJacobianWeightVec();
      const FEManagerType::ElemQpointVectorType& qpXVec           = FEManager.getElemQpointX();
      for (SpaMesh::ElementIterator elemIt(mesh); elemIt; ++elemIt)
      {
        int elem = elemIt.value();
        FEManager.initVolumeFE(mesh, elemIt, InputData::p_order, InputData::q_order, 2);
        double l2_elem_error = 0;
        for (unsigned int j = 0; j < shapeMatrix.cols(); j++)
        {
          l2_elem_error += pow((solutionBlocked.row(elem) * shapeMatrix.col(j))-InputData::ExactSolution::exact_solution(qpXVec.row(j)),2) * jacXweightVec[j];
        }
        l2_error += l2_elem_error;
      }
      double allL2_error = sqrt(l2_error);
      l2ErrorNorms.push_back(allL2_error);
      if (FELib::processor_id() == 0)
      {
        std::cout<<"dofs "<<"L2 error "<<dofMap.getGlobalElementNumber()*InputData::n_dofs<<","<<allL2_error<<std::endl;
        std::cout<<std::endl;
      }
      VTKWriterSpaMesh<InputData::dim, FEManagerType> vtkWriter;
      vtkWriter.writeOnTheFly(mesh,FEManager.getFEDataAccessor(),dofMap,solutionBlocked);
    }
  }
  for (unsigned int i = 1; i < l2ErrorNorms.size(); i++)
    std::cout<<"Convergence rate "<<log(l2ErrorNorms[i-1]/l2ErrorNorms[i])/log(2.)<<std::endl;

  return EXIT_SUCCESS;
}

