#pragma once

#include <map>
#include <string>
#include <utility>

#include <dune/common/hybridutilities.hh>
#include <dune/common/timer.hh>
#include <dune/typetree/childextraction.hh>

#include <amdis/AdaptInfo.hpp>
#include <amdis/FileWriter.hpp>
#include <amdis/Assembler.hpp>
#include <amdis/GridFunctionOperator.hpp>
#include <amdis/GridTransferManager.hpp>

namespace AMDiS {

template <class Traits>
void ProblemStat<Traits>::initialize(
    Flag initFlag,
    Self* adoptProblem,
    Flag adoptFlag)
{
  // create grids
  if (grid_) {
    warning("grid already created");
  }
  else {
    if (initFlag.isSet(CREATE_MESH) ||
        (!adoptFlag.isSet(INIT_MESH) &&
        (initFlag.isSet(INIT_SYSTEM) || initFlag.isSet(INIT_FE_SPACE)))) {
      createGrid();
    }

    if (adoptProblem &&
        (adoptFlag.isSet(INIT_MESH) ||
        adoptFlag.isSet(INIT_SYSTEM) ||
        adoptFlag.isSet(INIT_FE_SPACE))) {
      adoptGrid(adoptProblem->grid_, adoptProblem->boundaryManager_);
    }
  }

  if (!grid_)
    warning("no grid created");

  if (initFlag.isSet(INIT_MESH)) {
    int globalRefinements = 0;
    Parameters::get(gridName_ + "->global refinements", globalRefinements);
    if (globalRefinements > 0) {
      grid_->globalRefine(globalRefinements);
    }
  }

  // create fespace
  if (globalBasis_) {
    warning("globalBasis already created");
  }
  else {
    if (initFlag.isSet(INIT_FE_SPACE) ||
        (initFlag.isSet(INIT_SYSTEM) && !adoptFlag.isSet(INIT_FE_SPACE))) {
      createGlobalBasis();
    }

    if (adoptProblem &&
        (adoptFlag.isSet(INIT_FE_SPACE) || adoptFlag.isSet(INIT_SYSTEM))) {
      adoptGlobalBasis(adoptProblem->globalBasis_);
    }
  }

  if (!globalBasis_)
    warning("no globalBasis created\n");


  // create system
  if (initFlag.isSet(INIT_SYSTEM))
    createMatricesAndVectors();

  if (adoptProblem && adoptFlag.isSet(INIT_SYSTEM)) {
    systemMatrix_ = adoptProblem->systemMatrix_;
    solution_ = adoptProblem->solution_;
    rhs_ = adoptProblem->rhs_;
    estimates_ = adoptProblem->estimates_;
  }


  // create solver
  if (linearSolver_) {
    warning("solver already created\n");
  }
  else {
    if (initFlag.isSet(INIT_SOLVER))
      createSolver();

    if (adoptProblem && adoptFlag.isSet(INIT_SOLVER)) {
      test_exit(!linearSolver_, "solver already created\n");
      linearSolver_ = adoptProblem->linearSolver_;
    }
  }

  if (!linearSolver_) {
    warning("no solver created\n");
  }

  // create marker
    if (initFlag.isSet(INIT_MARKER))
      createMarker();

    if (adoptProblem && adoptFlag.isSet(INIT_MARKER))
      marker_ = adoptProblem->marker_;


  // create file writer
  if (initFlag.isSet(INIT_FILEWRITER))
    createFileWriter();

  solution_->compress();
}


template <class Traits>
void ProblemStat<Traits>::createGrid()
{
  Parameters::get(name_ + "->mesh", gridName_);
  grid_ = MeshCreator<Grid>::create(gridName_);
  boundaryManager_ = std::make_shared<BoundaryManager<Grid>>(grid_);

  msg("Create grid:");
  msg("#elements = {}"   , grid_->size(0));
  msg("#faces/edges = {}", grid_->size(1));
  msg("#vertices = {}"   , grid_->size(dim));
  msg("");
}


template <class T, class GV>
using HasCreate = decltype(T::create(std::declval<GV>()));


template <class Traits>
void ProblemStat<Traits>::createGlobalBasis()
{
  createGlobalBasisImpl(Dune::Std::is_detected<HasCreate,Traits,GridView>{});
  initGlobalBasis(*globalBasis_);
}


template <class Traits>
void ProblemStat<Traits>::createGlobalBasisImpl(std::true_type)
{
  assert( bool(grid_) );
  static_assert(std::is_same<GridView, typename Grid::LeafGridView>::value, "");
  globalBasis_ = std::make_shared<GlobalBasis>(Traits::create(grid_->leafGridView()));
}


template <class Traits>
void ProblemStat<Traits>::createGlobalBasisImpl(std::false_type)
{
  error_exit("Cannot create GlobalBasis from type. Pass a BasisCreator instead!");
}


template <class Traits>
void ProblemStat<Traits>::initGlobalBasis(GlobalBasis const& globalBasis)
{
  dirichletBCs_.init(*globalBasis_, *globalBasis_);
  periodicBCs_.init(*globalBasis_, *globalBasis_);
}


template <class Traits>
void ProblemStat<Traits>::createMatricesAndVectors()
{
  systemMatrix_ = std::make_shared<SystemMatrix>(*globalBasis_, *globalBasis_);
  solution_ = std::make_shared<SystemVector>(*globalBasis_, DataTransferOperation::INTERPOLATE);
  rhs_ = std::make_shared<SystemVector>(*globalBasis_, DataTransferOperation::NO_OPERATION);

  auto localView = globalBasis_->localView();
  for_each_node(localView.tree(), [&,this](auto const& node, auto treePath)
  {
    std::string i = to_string(treePath);
    estimates_[i].resize(globalBasis_->gridView().indexSet().size(0));
    for (std::size_t j = 0; j < estimates_[i].size(); j++)
      estimates_[i][j] = 0.0; // TODO: Remove when estimate() is implemented
  });
}


template <class Traits>
void ProblemStat<Traits>::createSolver()
{
  std::string solverName = "default";
  Parameters::get(name_ + "->solver->name", solverName);

  auto solverCreator
    = named(CreatorMap<LinearSolverType>::getCreator(solverName, name_ + "->solver->name"));

  linearSolver_ = solverCreator->createWithString(name_ + "->solver");
}


template <class Traits>
void ProblemStat<Traits>::createMarker()
{
  marker_.clear();
  auto localView = globalBasis_->localView();
  for_each_node(localView.tree(), [&,this](auto const& node, auto treePath)
  {
    std::string componentName = name_ + "->marker[" + to_string(treePath) + "]";

    if (!Parameters::get<std::string>(componentName + "->strategy"))
      return;

    std::string tp = to_string(treePath);
    auto newMarker
      = EstimatorMarker<Grid>::createMarker(componentName, tp, estimates_[tp], grid_);
    if (newMarker)
      marker_.push_back(std::move(newMarker));

    // If there is more than one marker, and all components are defined
    // on the same grid, the maximum marking has to be enabled.
    // TODO: needs to be checked!
    if (marker_.size() > 1)
      marker_.back()->setMaximumMarking(true);
  });
}


template <class Traits>
void ProblemStat<Traits>::createFileWriter()
{
  filewriter_.clear();
  auto localView = globalBasis_->localView();
  for_each_node(localView.tree(), [&,this](auto const& node, auto treePath)
  {
    std::string componentName = name_ + "->output[" + to_string(treePath) + "]";

    if (!Parameters::get<std::string>(componentName + "->filename"))
      return;

    auto writer = makeFileWriterPtr(componentName, this->solution(treePath));
    filewriter_.push_back(std::move(writer));
  });
}


// Adds a Dirichlet boundary condition
template <class Traits>
  template <class Predicate, class RowTreePath, class ColTreePath, class Values>
void ProblemStat<Traits>::
addDirichletBC(Predicate const& predicate, RowTreePath row, ColTreePath col, Values const& values)
{
  static_assert( Concepts::Functor<Predicate, bool(WorldVector)>,
    "Function passed to addDirichletBC for `predicate` does not model the Functor<bool(WorldVector)> concept");

  auto localView = globalBasis_->localView();
  auto i = child(localView.tree(), makeTreePath(row));
  auto j = child(localView.tree(), makeTreePath(col));

  auto valueGridFct = makeGridFunction(values, globalBasis_->gridView());

  using Range = RangeType_t<decltype(i)>;
  using BcType = DirichletBC<WorldVector,Range>;
  auto bc = std::make_unique<BcType>(predicate, valueGridFct);
  dirichletBCs_[i][j].push_back(std::move(bc));
}


// Adds a Dirichlet boundary condition
template <class Traits>
  template <class RowTreePath, class ColTreePath, class Values>
void ProblemStat<Traits>::
addDirichletBC(BoundaryType id, RowTreePath row, ColTreePath col, Values const& values)
{
  auto localView = globalBasis_->localView();
  auto i = child(localView.tree(), makeTreePath(row));
  auto j = child(localView.tree(), makeTreePath(col));

  auto valueGridFct = makeGridFunction(values, globalBasis_->gridView());

  using Range = RangeType_t<decltype(i)>;
  using BcType = DirichletBC<WorldVector,Range>;
  auto bc = std::make_unique<BcType>(boundaryManager_, id, valueGridFct);
  dirichletBCs_[i][j].push_back(std::move(bc));
}


template <class Traits>
void ProblemStat<Traits>::
addPeriodicBC(BoundaryType id, WorldMatrix const& matrix, WorldVector const& vector)
{
  auto localView = globalBasis_->localView();
  using BcType = PeriodicBC<WorldVector, typename GlobalBasis::MultiIndex>;
  using FaceTrafo = FaceTransformation<typename WorldVector::field_type, WorldVector::dimension>;
  auto bc = std::make_unique<BcType>(boundaryManager_, id, FaceTrafo{matrix, vector});
  periodicBCs_[localView.tree()][localView.tree()].push_back(std::move(bc));
}


template <class Traits>
void ProblemStat<Traits>::
solve(AdaptInfo& adaptInfo, bool createMatrixData, bool storeMatrixData)
{
  Dune::Timer t;

  SolverInfo solverInfo(name_ + "->solver");
  solverInfo.setCreateMatrixData(createMatrixData);
  solverInfo.setStoreMatrixData(storeMatrixData);

  solution_->compress();

  linearSolver_->solve(systemMatrix_->matrix(), solution_->vector(), rhs_->vector(),
                      solverInfo);

  if (solverInfo.info() > 0) {
    msg("solution of discrete system needed {} seconds", t.elapsed());

    if (solverInfo.absResidual() >= 0.0) {
      if (solverInfo.relResidual() >= 0.0)
        msg("Residual norm: ||b-Ax|| = {}, ||b-Ax||/||b|| = {}",
          solverInfo.absResidual(), solverInfo.relResidual());
      else
        msg("Residual norm: ||b-Ax|| = {}", solverInfo.absResidual());
    }
  }

  if (solverInfo.doBreak()) {
    std::stringstream tol_str;
    if (solverInfo.absTolerance() > 0
        && solverInfo.absResidual() > solverInfo.absTolerance())
      tol_str << "absTol = " << solverInfo.absTolerance() << " ";
    if (solverInfo.relTolerance() > 0
        && solverInfo.relResidual() > solverInfo.relTolerance())
      tol_str << "relTol = " << solverInfo.relTolerance() << " ";
    error_exit("Tolerance {} could not be reached!", tol_str.str());
  }
}


template <class Traits>
Flag ProblemStat<Traits>::
markElements(AdaptInfo& adaptInfo)
{
  Dune::Timer t;

  Flag markFlag = 0;
  for (auto& currentMarker : marker_)
    markFlag |= currentMarker->markGrid(adaptInfo);

  msg("markElements needed {} seconds", t.elapsed());

  return markFlag;
}


template <class Traits>
Flag ProblemStat<Traits>::
globalCoarsen(int n)
{
  Dune::Timer t;
  bool adapted = false;
  for (int i = 0; i < n; ++i) {
    // mark all entities for grid refinement
    for (const auto& element : elements(grid_->leafGridView()))
      grid_->mark(-1, element);

    adapted |= GridTransferManager::adapt(*grid_, *globalBasis_);
  }

  msg("globalCoarsen needed {} seconds", t.elapsed());
  return adapted ? MESH_ADAPTED : Flag(0);
}


// grid has globalRefine(int, AdaptDataHandleInterface&)
template <class G>
using HasGlobalRefineADHI = decltype(std::declval<G>().globalRefine(1,std::declval<GridTransfer<G>&>()));

template <class Traits>
Flag ProblemStat<Traits>::
globalRefine(int refCount)
{
  Dune::Timer t;
  bool adapted = false;
  Dune::Hybrid::ifElse(Dune::Std::is_detected<HasGlobalRefineADHI, Grid>{},
  /*then*/ [&](auto id) {
    // TODO(FM): Add a way to pass a GridTransfer as ADH with *globalBasis_ argument
    id(grid_)->globalRefine(refCount, GridTransferManager::gridTransfer(*grid_));
    globalBasis_->update(this->gridView());
  },
  /*else*/ [&](auto id) {
    for (int i = 0; i < refCount; ++i) {
      // mark all entities for grid refinement
      for (const auto& element : elements(grid_->leafGridView()))
        grid_->mark(1, element);

      adapted |= GridTransferManager::adapt(*id(grid_), *id(globalBasis_));
    }
  });

  msg("globalRefine needed {} seconds", t.elapsed());
  return adapted ? MESH_ADAPTED : Flag(0);
}


template <class Traits>
Flag ProblemStat<Traits>::
adaptGrid(AdaptInfo& adaptInfo)
{
  Dune::Timer t;

  bool adapted = GridTransferManager::adapt(*grid_, *globalBasis_);

  msg("adaptGrid needed {} seconds", t.elapsed());
  return adapted ? MESH_ADAPTED : Flag(0);
}


template <class Traits>
void ProblemStat<Traits>::
buildAfterAdapt(AdaptInfo& /*adaptInfo*/, Flag /*flag*/, bool asmMatrix, bool asmVector)
{
  Dune::Timer t;

  // 1. init matrix and rhs vector and initialize dirichlet boundary conditions
  systemMatrix_->init(asmMatrix);
  rhs_->init(asmVector);

  auto localView = globalBasis_->localView();
  for_each_node(localView.tree(), [&,this](auto const& rowNode, auto rowTp) {
    auto rowBasis = Dune::Functions::subspaceBasis(*globalBasis_, rowTp);
    for_each_node(localView.tree(), [&,this](auto const& colNode, auto colTp) {
      auto colBasis = Dune::Functions::subspaceBasis(*globalBasis_, colTp);
      for (auto bc : dirichletBCs_[rowNode][colNode])
        bc->init(rowBasis, colBasis);
      for (auto bc : periodicBCs_[rowNode][colNode])
        bc->init(rowBasis, colBasis);
    });
  });
  msg("{} total DOFs", globalBasis_->dimension());

  // 2. traverse grid and assemble operators on the elements
  for (auto const& element : elements(gridView())) {
    localView.bind(element);

    if (asmMatrix)
      systemMatrix_->assemble(localView, localView);
    if (asmVector)
      rhs_->assemble(localView);

    localView.unbind();
  }

  // 3. finish matrix insertion and apply dirichlet boundary conditions
  systemMatrix_->finish(asmMatrix);
  rhs_->finish(asmVector);

  for_each_node(localView.tree(), [&,this](auto const& rowNode, auto row_tp) {
    for_each_node(localView.tree(), [&,this](auto const& colNode, auto col_tp) {
      // finish boundary condition
      for (auto bc : dirichletBCs_[rowNode][colNode])
        bc->fillBoundaryCondition(*systemMatrix_, *solution_, *rhs_, rowNode, row_tp, colNode, col_tp);
      for (auto bc : periodicBCs_[rowNode][colNode])
        bc->fillBoundaryCondition(*systemMatrix_, *solution_, *rhs_, rowNode, row_tp, colNode, col_tp);
    });
  });

  msg("fill-in of assembled matrix: {}", systemMatrix_->nnz());
  msg("buildAfterAdapt needed {} seconds", t.elapsed());
}


template <class Traits>
void ProblemStat<Traits>::
writeFiles(AdaptInfo& adaptInfo, bool force)
{
  Dune::Timer t;
  for (auto writer : filewriter_)
    writer->writeFiles(adaptInfo, force);
  msg("writeFiles needed {} seconds", t.elapsed());
}


} // end namespace AMDiS