#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <dune/common/exceptions.hh>

#include <dune/matrix-vector/axpy.hh>
#include <dune/solvers/norms/sumnorm.hh>

#include <dune/tectonic/utils/reductionfactors.hh>

#include "functionalfactory.hh"
#include "nonlinearsolver.hh"

#include "../utils/debugutils.hh"

#include "fixedpointiterator.hh"

void FixedPointIterationCounter::operator+=(
    FixedPointIterationCounter const &other) {
  iterations += other.iterations;
  multigridIterations += other.multigridIterations;
}

template <class Factory, class NBodyAssembler, class Updaters, class ErrorNorms>
FixedPointIterator<Factory, NBodyAssembler, Updaters, ErrorNorms>::FixedPointIterator(
    Dune::ParameterTree const &parset,
    const NBodyAssembler& nBodyAssembler,
    const IgnoreVector& ignoreNodes,
    GlobalFriction& globalFriction,
    const std::vector<const BitVector*>& bodywiseNonmortarBoundaries,
    const ErrorNorms& errorNorms)
    : parset_(parset),
      nBodyAssembler_(nBodyAssembler),
      ignoreNodes_(ignoreNodes),
      globalFriction_(globalFriction),
      bodywiseNonmortarBoundaries_(bodywiseNonmortarBoundaries),
      fixedPointMaxIterations_(parset.get<size_t>("v.fpi.maximumIterations")),
      fixedPointTolerance_(parset.get<double>("v.fpi.tolerance")),
      lambda_(parset.get<double>("v.fpi.lambda")),
      solverParset_(parset.sub("v.solver")),
      errorNorms_(errorNorms) {}

template <class Factory, class NBodyAssembler, class Updaters, class ErrorNorms>
bool FixedPointIterator<Factory, NBodyAssembler, Updaters, ErrorNorms>::displacementCriterion(const Updaters& updaters, std::vector<Vector>& last_u) const {
    bool criterion = true;

    std::vector<Vector> u;
    updaters.rate_->extractDisplacement(u);

    const auto& mats = updaters.rate_->getMatrices();

    for (int i=0; i<u.size(); i++) {
        const EnergyNorm<Matrix, Vector> ANorm(*mats.elasticity[i]), MNorm(*mats.mass[i]);
        const SumNorm<Vector> AMNorm(1.0, ANorm, 1.0, MNorm);

        if (u[i].size()==0 || last_u[i].size()==0)
            continue;

        if (AMNorm.diff(u[i], last_u[i]) >= fixedPointTolerance_) {
            criterion = false;
            break;
        }
    }

    last_u = u;

    return criterion;
}

template <class Factory, class NBodyAssembler, class Updaters, class ErrorNorms>
bool FixedPointIterator<Factory, NBodyAssembler, Updaters, ErrorNorms>::stateCriterion(const std::vector<ScalarVector>& alpha, const std::vector<ScalarVector>& newAlpha) const {
    bool criterion = true;
    for (int i=0; i<alpha.size(); i++) {
        if (alpha[i].size()==0 || newAlpha[i].size()==0)
            continue;

        if (errorNorms_[i]->diff(alpha[i], newAlpha[i]) >= fixedPointTolerance_) {
            criterion = false;
            break;
        }
    }

    return criterion;
}

template <class Factory, class NBodyAssembler, class Updaters, class ErrorNorms>
template <class LinearSolver>
FixedPointIterationCounter
FixedPointIterator<Factory, NBodyAssembler, Updaters, ErrorNorms>::run(
    Updaters updaters, std::shared_ptr<LinearSolver>& linearSolver,
    const std::vector<Matrix>& velocityMatrices, const std::vector<Vector>& velocityRHSs,
    std::vector<Vector>& velocityIterates) {

    // debugging
    /*const auto& contactCouplings = nBodyAssembler_.getContactCouplings();
    for (size_t i=0; i<contactCouplings.size(); i++) {
      print(*contactCouplings[i]->nonmortarBoundary().getVertices(), "nonmortarBoundaries:");
    }*/

  const auto nBodies = nBodyAssembler_.nGrids();

  FunctionalFactory<std::decay_t<decltype(nBodyAssembler_)>, decltype(globalFriction_), Matrix, Vector> functionalFactory(nBodyAssembler_, globalFriction_, velocityMatrices, velocityRHSs);
  functionalFactory.build();
  auto functional = functionalFactory.functional();

  NonlinearSolver<std::remove_reference_t<decltype(*functional)>, IgnoreVector> solver(parset_.sub("solver.tnnmg"), linearSolver, functional, ignoreNodes_);

  size_t fixedPointIteration;
  size_t multigridIterations = 0;

  std::vector<Vector> last_u;
  updaters.rate_->extractOldDisplacement(last_u);

  std::vector<ScalarVector> alpha(nBodies);
  updaters.state_->extractAlpha(alpha);

  std::vector<Vector> v_o;
  updaters.rate_->extractOldVelocity(v_o);

  Vector total_v_o;
  nBodyAssembler_.nodalToTransformed(v_o, total_v_o);

  Vector total_v;
  nBodyAssembler_.nodalToTransformed(velocityIterates, total_v);

  //print(velocityIterates, "velocityIterates start:");

  for (fixedPointIteration = 0; fixedPointIteration < fixedPointMaxIterations_;
       ++fixedPointIteration) {

    // update friction
    globalFriction_.updateAlpha(alpha);

    // solve rate problem
    auto res = solver.solve(solverParset_, total_v);
    multigridIterations += res.iterations;

    Vector v_m = total_v_o;
    v_m *= 1.0 - lambda_;
    Dune::MatrixVector::addProduct(v_m, lambda_, total_v);

    // extract relative velocities in mortar basis
    std::vector<Vector> v_rel;
    split(v_m, v_rel);

    // solve state problem
    updaters.state_->solve(v_rel);
    std::vector<ScalarVector> newAlpha(nBodies);
    updaters.state_->extractAlpha(newAlpha);

    nBodyAssembler_.postprocess(total_v, velocityIterates);
    //Rprint(velocityIterates, "velocityIterates loop:");
    updaters.rate_->postProcess(velocityIterates);
    bool breakCriterion = true; //displacementCriterion(updaters, last_u); //stateCriterion(alpha, newAlpha);

    printRegularityTruncation(globalFriction_, total_v);

    if (lambda_ < 1e-12 or breakCriterion) {
      fixedPointIteration++;
      break;
    }

    if (res.iterations>200) {
        std::cout << "FixedPointIterator: TNNMG did not converge, iterations: " << res.iterations << std::endl;
        //DUNE_THROW(Dune::Exception, "FixedPointIterator: TNNMG did not converge");
        break;
    }

    alpha = newAlpha;
  }

  if (fixedPointIteration == fixedPointMaxIterations_)
    DUNE_THROW(Dune::Exception, "FPI failed to converge");

  nBodyAssembler_.postprocess(total_v, velocityIterates);
  updaters.rate_->postProcess(velocityIterates);

  //print(velocityIterates, "velocityIterates end:");

  // Cannot use return { fixedPointIteration, multigridIterations };
  // with gcc 4.9.2, see also http://stackoverflow.com/a/37777814/179927
  FixedPointIterationCounter ret;
  ret.iterations = fixedPointIteration;
  ret.multigridIterations = multigridIterations;
  return ret;
}

/*std::ostream &operator<<(std::ostream &stream,
                         FixedPointIterationCounter const &fpic) {
  return stream << "(" << fpic.iterations << "," << fpic.multigridIterations
                << ")";
}*/

template <class Factory, class NBodyAssembler, class Updaters, class ErrorNorms>
void FixedPointIterator<Factory, NBodyAssembler, Updaters, ErrorNorms>::split(
    const Vector& v,
    std::vector<Vector>& splitV) const {

    const size_t nBodies = nBodyAssembler_.nGrids();

    size_t globalIdx = 0;
    splitV.resize(nBodies);
    for (size_t bodyIdx=0; bodyIdx<nBodies; bodyIdx++) {
        auto& splitV_ = splitV[bodyIdx];

        splitV_.resize(nBodyAssembler_.grids_[bodyIdx]->size(dims));

        for (size_t i=0; i<splitV_.size(); i++) {
            splitV_[i] = v[globalIdx];
            globalIdx++;
        }
    }


   /*
        boundaryNodes =

        const auto gridView = contactCouplings[couplingIndices[0]]->nonmortarBoundary().gridView();

        Dune::MultipleCodimMultipleGeomTypeMapper<
            decltype(gridView), Dune::MCMGVertexLayout> const vertexMapper(gridView, Dune::mcmgVertexLayout());

        for (auto it = gridView.template begin<block_size>(); it != gridView.template end<block_size>(); ++it) {
            const auto i = vertexMapper.index(*it);

            for (size_t j=0; j<couplingIndices.size(); j++) {
                const auto couplingIdx = couplingIndices[j];

                if (not contactCouplings[couplingIdx]->nonmortarBoundary().containsVertex(i))
                  continue;

                localToGlobal_.emplace_back(i);
                restrictions_.emplace_back(weights[bodyIdx][i], weightedNormalStress[bodyIdx][i],
                                          couplings[i]->frictionData()(geoToPoint(it->geometry())));
                break;
            }

            globalIdx++;
        }
        maxIndex_[bodyIdx] = globalIdx;
    }*/
}

#include "fixedpointiterator_tmpl.cc"