#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 != 8)
  {
    fmt::print(stderr, "needs 7 arguments.\n");
    exit(1);
  }

  std::string model = argv[1];
  std::string mcdFile = argv[2];
  std::string trainTsvFile = argv[3];
  std::string trainRawFile = "";
  std::string devTsvFile = argv[5];
  std::string devRawFile = "";
  int nbEpoch = std::stoi(argv[7]);

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

  ReadingMachine machine(machinePath.string());

  BaseConfig goldConfig(mcdFile, trainTsvFile, trainRawFile);
  SubConfig config(goldConfig);

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

  Decoder decoder(machine);
  BaseConfig devGoldConfig(mcdFile, devTsvFile, devRawFile);

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

  return 0;
}