#include "coupledtimestepper.hh"

template <typename T1, typename T2>
std::pair<T1, T2> clonePair(std::pair<T1, T2> in) {
  return { in.first->clone(), in.second->clone() };
}

template <class Factory, class UpdaterPair> class AdaptiveTimeStepper {
  using StateUpdater = typename UpdaterPair::first_type::element_type;
  using VelocityUpdater = typename UpdaterPair::second_type::element_type;
  using Vector = typename Factory::Vector;
  using ConvexProblem = typename Factory::ConvexProblem;
  using Nonlinearity = typename ConvexProblem::NonlinearityType;

  using MyCoupledTimeStepper =
      CoupledTimeStepper<Factory, StateUpdater, VelocityUpdater>;

public:
  AdaptiveTimeStepper(
      Factory &factory, Dune::ParameterTree const &parset,
      std::shared_ptr<Nonlinearity> globalFriction, UpdaterPair &current,
      std::function<void(double, Vector &)> externalForces,
      std::function<bool(UpdaterPair &, UpdaterPair &)> mustRefine)
      : finalTime_(parset.get<double>("problem.finalTime")),
        relativeTime_(0.0),
        relativeTau_(1e-6), // FIXME (not really important, though)
        factory_(factory),
        parset_(parset),
        globalFriction_(globalFriction),
        current_(current),
        R1_(clonePair(current_)),
        externalForces_(externalForces),
        mustRefine_(mustRefine),
        iterationWriter_("iterations", std::fstream::out) {
    MyCoupledTimeStepper coupledTimeStepper(finalTime_, factory_, parset_,
                                            globalFriction_, R1_.first,
                                            R1_.second, externalForces_);
    stepAndReport("R1", coupledTimeStepper, relativeTime_, relativeTau_);
    iterationWriter_ << std::endl;
  }

  bool reachedEnd() { return relativeTime_ >= 1.0; }

  bool coarsen() {
    bool didCoarsen = false;

    while (relativeTime_ + relativeTau_ < 1.0) {
      R2_ = clonePair(R1_);
      {
        MyCoupledTimeStepper coupledTimeStepper(finalTime_, factory_, parset_,
                                                globalFriction_, R2_.first,
                                                R2_.second, externalForces_);
        stepAndReport("R2", coupledTimeStepper, relativeTime_ + relativeTau_,
                      relativeTau_);
      }

      UpdaterPair C = clonePair(current_);
      {
        MyCoupledTimeStepper coupledTimeStepper(finalTime_, factory_, parset_,
                                                globalFriction_, C.first,
                                                C.second, externalForces_);

        stepAndReport("C", coupledTimeStepper, relativeTime_,
                      2.0 * relativeTau_);
      }

      if (!mustRefine_(C, R2_)) {
        R2_ = { nullptr, nullptr };
        R1_ = C;
        relativeTau_ *= 2.0;
        didCoarsen = true;
      } else {
        break;
      }
    }
    return didCoarsen;
  }

  void refine() {
    while (true) {
      UpdaterPair F2 = clonePair(current_);
      UpdaterPair F1;
      {
        MyCoupledTimeStepper coupledTimeStepper(finalTime_, factory_, parset_,
                                                globalFriction_, F2.first,
                                                F2.second, externalForces_);
        stepAndReport("F1", coupledTimeStepper, relativeTime_,
                      relativeTau_ / 2.0);

        F1 = clonePair(F2);
        stepAndReport("F2", coupledTimeStepper,
                      relativeTime_ + relativeTau_ / 2.0, relativeTau_ / 2.0);
      }

      if (!mustRefine_(R1_, F2)) {
        break;
      } else {
        R1_ = F1;
        R2_ = F2;
        relativeTau_ /= 2.0;
      }
    }
  }

  void advance() {
    if (!coarsen())
      refine();

    iterationWriter_ << std::endl;

    current_ = R1_;
    R1_ = R2_;
    relativeTime_ += relativeTau_;
  }

  double getRelativeTime() { return relativeTime_; }
  double getRelativeTau() { return relativeTau_; }

private:
  void stepAndReport(std::string type, MyCoupledTimeStepper &stepper,
                     double rTime, double rTau) {
    iterationWriter_ << type << " " << stepper.step(rTime, rTau) << " "
                     << std::flush;
  }

  double finalTime_;
  double relativeTime_;
  double relativeTau_;
  Factory &factory_;
  Dune::ParameterTree const &parset_;
  std::shared_ptr<Nonlinearity> globalFriction_;
  UpdaterPair &current_;
  UpdaterPair R1_;
  UpdaterPair R2_;
  std::function<void(double, Vector &)> externalForces_;
  std::function<bool(UpdaterPair &, UpdaterPair &)> mustRefine_;
  std::fstream iterationWriter_;
};