 #include <boost/timer/timer.hpp>
 #include <boost/chrono.hpp>
 #include <boost/program_options.hpp>
+#include <thread>
+#include <chrono>
+#include <mutex>
 
 #include "marian.h"
 #include "optimizers/optimizers.h"

  int dimRnn = options.get<int>("dim-rnn");
  int dimBatch = options.get<int>("mini-batch");
  int dimMaxiBatch = options.get<int>("maxi-batch");

  std::vector<int> dims = {
    dimVocabs[0], dimEmb, dimRnn,
    dimVocabs[1], dimEmb, dimRnn,
    dimBatch
  };

  auto devices = options.get<std::vector<int>>("device");
  int device = devices[0];

  auto graph = New<ExpressionGraph>();
  graph->setDevice(device);

  auto nematus = New<Nematus>(dims);

  if(options.has("init")) {
    LOG(info) << "Loading parameters from " << options.get<std::string>("init");
    nematus->load(graph, options.get<std::string>("init"));
  }

  graph->reserveWorkspaceMB(options.get<size_t>("workspace"));

  Ptr<ClipperBase> clipper = nullptr;

  float clipNorm = options.get<double>("clip-norm");
  float lrate = options.get<double>("lrate");
  if(clipNorm > 0)
    clipper = Clipper<Norm>(clipNorm);
  auto opt = Optimizer<Adam>(lrate, clip=clipper);

  auto update = [graph, opt, nematus](Ptr<CorpusBatch> batch) -> float {
    auto costNode = nematus->construct(graph, batch);
    opt->update(graph);
    return costNode->scalar();
  };

  auto save = [graph, nematus](const std::string& name) {
    LOG(info) << "Saving parameters to " << name;
    nematus->save(graph, name);
  };

  auto trainSets = options.get<std::vector<std::string>>("trainsets");
  auto vocabs = options.get<std::vector<std::string>>("vocabs");
  size_t maxSentenceLength = options.get<size_t>("max-length");

  auto corpus = New<Corpus>(trainSets, vocabs, dimVocabs, maxSentenceLength);
  auto bg = New<BatchGenerator<Corpus>>(corpus, dimBatch, dimMaxiBatch);

  TrainingLoop(options, bg, update, save);

  return 0;
}

namespace marian {

  void TrainingLoop(Config& options,
                    Ptr<data::BatchGenerator<data::Corpus>> batchGenerator) {
 
 
    auto dimVocabs = options.get<std::vector<int>>("dim-vocabs");
    int dimEmb = options.get<int>("dim-emb");
    int dimRnn = options.get<int>("dim-rnn");
    int dimBatch = options.get<int>("mini-batch");
    int dimMaxiBatch = options.get<int>("maxi-batch");

    std::vector<int> dims = {
      dimVocabs[0], dimEmb, dimRnn,
      dimVocabs[1], dimEmb, dimRnn,
      dimBatch
    };

    auto devices = options.get<std::vector<int>>("device");
    
    std::mutex dispMutex;
    
    auto save = [](const std::string& name) {
      LOG(info) << "Saving parameters to " << name;
      //nematus->save(graph, name);
    };

    ThreadPool pool(devices.size(), devices.size());
 
    boost::timer::cpu_timer timer;

    size_t epochs = 1;
