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

#include <future>
#include <thread>
#include <chrono>

#include <dune/solvers/norms/energynorm.hh>
#include <dune/solvers/iterationsteps/multigridstep.hh>
#include <dune/solvers/iterationsteps/cgstep.hh>
#include <dune/solvers/solvers/loopsolver.hh>

#include "../spatial-solving/preconditioners/multilevelpatchpreconditioner.hh"

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

#include "adaptivetimestepper.hh"
#include "step.hh"

const unsigned int N_THREADS = std::thread::hardware_concurrency();

void IterationRegister::registerCount(FixedPointIterationCounter count) {
  totalCount += count;
}

void IterationRegister::registerFinalCount(FixedPointIterationCounter count) {
  finalCount = count;
}

void IterationRegister::reset() {
  totalCount = FixedPointIterationCounter();
  finalCount = FixedPointIterationCounter();
}

/*
 * Implementation: AdaptiveTimeStepper
 */
template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::AdaptiveTimeStepper(
        const StepBase& stepBase,
        ContactNetwork& contactNetwork,
        Updaters &current,
        double relativeTime,
        double relativeTau,
        std::function<bool(Updaters &, Updaters &)> mustRefine)
    : relativeTime_(relativeTime),
      relativeTau_(relativeTau),
      stepBase_(stepBase),
      contactNetwork_(contactNetwork),
      current_(current),
      R1_(),
      mustRefine_(mustRefine) {

    std::cout << N_THREADS << " concurrent threads are supported." << std::endl;
}

template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
bool AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::reachedEnd() {
  return relativeTime_ >= 1.0;
}

// compute C and R2 in parallel
// returns number of coarsenings done
template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
int AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::coarsen() {
    using Step = Step<Factory, ContactNetwork, Updaters, ErrorNorms>;

    int coarseCount = 1; // one coarsening step was already done in determineStrategy()

    UpdatersWithCount C;
    UpdatersWithCount R2;

    const auto& currentNBodyAssembler = contactNetwork_.nBodyAssembler();

    while (relativeTime_ + relativeTau_ <= 1.0) {
      //std::cout << "tau: " << relativeTau_ << std::endl;

      setDeformation(current_);
      auto C_Step = Step(stepBase_, current_, currentNBodyAssembler, relativeTime_, 2 * relativeTau_, iterationRegister_);
      C_Step.run(Step::Mode::newThread); //newThread

      //updateReductionFactors(reductionFactors);
      //std::cout << "AdaptiveTimeStepper C computed!" << std::endl << std::endl;

      /*using ScalarVector = typename Updaters::StateUpdater::ScalarVector;
      std::vector<ScalarVector> cAlpha(contactNetwork_.nBodies());
      C.updaters.state_->extractAlpha(cAlpha);
      print(cAlpha, "cAlpha: ");*/

      setDeformation(R1_.updaters);
      //auto R2_linearSolver = makeLinearSolver();
      auto&& nBodyAssembler = step(currentNBodyAssembler);
      auto R2_Step = Step(stepBase_, R1_.updaters, nBodyAssembler, relativeTime_ + relativeTau_, relativeTau_, iterationRegister_);
      R2_Step.run(Step::Mode::newThread); //newThread

      //updateReductionFactors(reductionFactors);
      //std::cout << "AdaptiveTimeStepper R2 computed!" << std::endl << std::endl;

      C = C_Step.get();
      R2 = R2_Step.get();

      /*std::vector<ScalarVector> rAlpha(contactNetwork_.nBodies());
      R2.updaters.state_->extractAlpha(rAlpha);
      print(rAlpha, "rAlpha: ");*/

      if (mustRefine_(C.updaters, R2.updaters))
        break;

      relativeTau_ *= 2;
      R1_ = C;

      coarseCount++;
    }

    current_ = R1_.updaters;
    R1_ = R2;

    return coarseCount;
}

// compute F1 and F2 sequentially
// returns number of refinements done
template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
int AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::refine(UpdatersWithCount& R2) {
    using Step = Step<Factory, ContactNetwork, Updaters, ErrorNorms>;

    int refineCount = 1; // one refinement step was already done in determineStrategy()

    UpdatersWithCount F1;
    UpdatersWithCount F2;

    const auto& currentNBodyAssembler = contactNetwork_.nBodyAssembler();

    while (true) {
        setDeformation(current_);
        //auto F1_linearSolver = makeLinearSolver();
        auto F1_Step = Step(stepBase_, current_, currentNBodyAssembler, relativeTime_, relativeTau_ / 2.0, iterationRegister_);
        F1_Step.run(Step::Mode::sameThread);
        F1 = F1_Step.get();

        //updateReductionFactors(reductionFactors);
        //std::cout << "AdaptiveTimeStepper F1 computed!" << std::endl << std::endl;

        setDeformation(F1.updaters);
        //auto F2_linearSolver = makeLinearSolver();
        auto&& nBodyAssembler = step(currentNBodyAssembler);
        auto F2_Step = Step(stepBase_, F1.updaters, nBodyAssembler, relativeTime_ + relativeTau_ / 2.0,
                  relativeTau_ / 2.0, iterationRegister_);
        F2_Step.run(Step::Mode::sameThread);
        F2 = F2_Step.get();
        //updateReductionFactors(reductionFactors);

        if (!mustRefine_(R1_.updaters, F2.updaters)) {
          std::cout << "Sufficiently refined!" << std::endl;
          break;
        }

        relativeTau_ /= 2.0;
        R1_ = F1;
        R2 = F2;

        refineCount++;
    }

    current_ = R1_.updaters;
    R1_ = R2;

    return refineCount;
}


/*
 * determines how to adapt time step, returns
 * -1: coarsen
 *  0: keep
 *  1: refine
 */
template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
int AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::determineStrategy(UpdatersWithCount& R2) {
    using Step = Step<Factory, ContactNetwork, Updaters, ErrorNorms>;

    int strategy = 0;

    const auto& currentNBodyAssembler = contactNetwork_.nBodyAssembler();

    UpdatersWithCount C;
    UpdatersWithCount F1;
    UpdatersWithCount F2;

    //if (relativeTime_ + relativeTau_ > 1.0) {
    //  return false;
    //}

    setDeformation(current_);
    //auto C_linearSolver = makeLinearSolver();
    auto C_Step = Step(stepBase_, current_, currentNBodyAssembler, relativeTime_, 2 * relativeTau_, iterationRegister_);
    C_Step.run(Step::Mode::newThread); // newThread
    //updateReductionFactors(reductionFactors);
    //std::cout << "AdaptiveTimeStepper C computed!" << std::endl << std::endl;

    setDeformation(R1_.updaters);
    //auto R2_linearSolver = makeLinearSolver();
    auto&& nBodyAssembler = step(currentNBodyAssembler);
    auto R2_Step = Step(stepBase_, R1_.updaters, nBodyAssembler, relativeTime_ + relativeTau_, relativeTau_, iterationRegister_);
    R2_Step.run(Step::Mode::newThread); //newThread

    //updateReductionFactors(reductionFactors);
    //std::cout << "AdaptiveTimeStepper R2 computed!" << std::endl << std::endl;

    if (N_THREADS < 3) {
      C = C_Step.get();
      R2 = R2_Step.get();

      if (!mustRefine_(C.updaters, R2.updaters)) {
          strategy = -1;
      }
    }

    if (strategy>-1) {
        setDeformation(current_);
        //auto F1_linearSolver = makeLinearSolver();
        auto F1_Step = Step(stepBase_, current_, currentNBodyAssembler, relativeTime_, relativeTau_ / 2.0, iterationRegister_);
        F1_Step.run(Step::Mode::newThread); //newThread
        //updateReductionFactors(reductionFactors);
        //std::cout << "AdaptiveTimeStepper F1 computed!" << std::endl << std::endl;

        if (N_THREADS > 2) {
          C = C_Step.get();
          R2 = R2_Step.get();

          if (!mustRefine_(C.updaters, R2.updaters)) {
              strategy = -1;
          }
        }

        F1 = F1_Step.get();

        if (strategy>-1) {
            setDeformation(F1.updaters);
            //auto F2_linearSolver = makeLinearSolver();
            auto&& nBodyAssembler = step(currentNBodyAssembler);
            auto F2_Step = Step(stepBase_, F1.updaters, nBodyAssembler, relativeTime_ + relativeTau_ / 2.0,
                                relativeTau_ / 2.0, iterationRegister_);
            F2_Step.run(Step::Mode::sameThread);
            F2 = F2_Step.get();
            //updateReductionFactors(reductionFactors);

            if (mustRefine_(R1_.updaters, F2.updaters)) {
                strategy = 1;
            }
        }
    }

    switch (strategy) {
        case -1:
            relativeTau_ *= 2;
            R1_ = C;
            break;
        case 0:
            current_ = R1_.updaters;
            R1_ = R2;
            break;
        case 1:
            relativeTau_ /= 2.0;
            R1_ = F1;
            R2 = F2;
            break;
    }

    return strategy;
}

template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
IterationRegister AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::advance() {
  /*
    |     C     | We check here if making the step R1 of size tau is a
    |  R1 | R2  | good idea. To check if we can coarsen, we compare
    |F1|F2|     | the result of (R1+R2) with C, i.e. two steps of size
                  tau with one of size 2*tau. To check if we need to
    refine, we compare the result of (F1+F2) with R1, i.e. two steps
    of size tau/2 with one of size tau. The method makes multiple
    coarsening/refining attempts, with coarsening coming first. */

  //std::cout << "AdaptiveTimeStepper::advance()" << std::endl;

  using Step = Step<Factory, ContactNetwork, Updaters, ErrorNorms>;

  if (R1_.updaters == Updaters()) {
      //setDeformation(current_);
      //auto R1_linearSolver = makeLinearSolver();
      auto R1_Step = Step(stepBase_, current_, contactNetwork_.nBodyAssembler(), relativeTime_, relativeTau_, iterationRegister_);
      R1_Step.run(Step::Mode::sameThread);
      R1_ = R1_Step.get();
   }

  iterationRegister_.reset();
  UpdatersWithCount R2;
  int strat = determineStrategy(R2);

  // coarsen
  if (strat<0) {
      int coarseningCount = coarsen();
      std::cout << " done with coarseningCount: " << coarseningCount << std::endl;
  }

  // refine
  if (strat>0) {
      int refineCount = refine(R2);
      std::cout << " done with refineCount: " << refineCount << std::endl;
  }

  iterationRegister_.registerFinalCount(R1_.count);
  relativeTime_ += relativeTau_;

  return iterationRegister_;
}

template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
void AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::setDeformation(const Updaters& updaters) {
    std::vector<Vector> u;
    updaters.rate_->extractDisplacement(u);

    for (size_t i=0; i<contactNetwork_.nBodies(); i++) {
        contactNetwork_.body(i)->setDeformation(u[i]);
    }

    // note: levelContactNetworks are not up-to-date; build() has to be called in order to update;
    // unnecessary for standard multigrid as linear solver, might have to be called for patch preconditioner
}

template <class Factory, class ContactNetwork, class Updaters, class ErrorNorms>
typename AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::NBodyAssembler
AdaptiveTimeStepper<Factory, ContactNetwork, Updaters, ErrorNorms>::step(const NBodyAssembler& oldNBodyAssembler) const {
    NBodyAssembler nBodyAssembler = oldNBodyAssembler;

    nBodyAssembler.assembleTransferOperator();
    nBodyAssembler.assembleObstacle();

    return nBodyAssembler;
}

#include "adaptivetimestepper_tmpl.cc"