#include <filesystem>
#include "fmt/core.h"
#include "util.hpp"
#include "BaseConfig.hpp"
#include "SubConfig.hpp"
#include "TransitionSet.hpp"
#include "ReadingMachine.hpp"
#include "Trainer.hpp"
#include "Decoder.hpp"

int main(int argc, char * argv[])
{
  if (argc != 5)
  {
    fmt::print(stderr, "needs 4 arguments.\n");
    exit(1);
  }

  std::string model = argv[1];
  std::string mcdFile = argv[2];
  std::string tsvFile = argv[3];
  //std::string rawFile = argv[4];
  std::string rawFile = "";

  std::filesystem::path modelPath(model);
  auto machinePath = modelPath / "machine.rm";

  ReadingMachine machine(machinePath.string());

  BaseConfig goldConfig(mcdFile, tsvFile, rawFile);
  SubConfig config(goldConfig);

  Trainer trainer(machine);
  trainer.createDataset(config);

  Decoder decoder(machine);

  int nbEpoch = 10;

  for (int i = 0; i < nbEpoch; i++)
  {
    float loss = trainer.epoch();
    auto devConfig = goldConfig;
    decoder.decode(devConfig, 1);
    decoder.evaluate(devConfig, modelPath, tsvFile);
    fmt::print(stderr, "\r{:80}\rEpoch {}/{} loss = {} dev = {}\n", " ", i+1, nbEpoch, loss, decoder.getF1Score("UPOS"));
  }

  return 0;
}