#ifndef TRAINER__H
#define TRAINER__H

#include "ReadingMachine.hpp"
#include "ConfigDataset.hpp"
#include "SubConfig.hpp"

class Trainer
{
  private :

  static constexpr std::size_t safetyNbExamplesMax = 10*1000*1000;

  struct Examples
  {
    std::vector<torch::Tensor> contexts;
    std::vector<torch::Tensor> classes;

    int currentExampleIndex{0};
    int lastSavedIndex{0};

    void saveIfNeeded(const std::string & state, std::filesystem::path dir, std::size_t threshold);
    void addContext(std::vector<std::vector<long>> & context);
    void addClass(int goldIndex);
  };

  private :

  using Dataset = ConfigDataset;
  using DataLoader = std::unique_ptr<torch::data::StatefulDataLoader<Dataset>>;

  private :

  ReadingMachine & machine;
  std::unique_ptr<Dataset> trainDataset{nullptr};
  std::unique_ptr<Dataset> devDataset{nullptr};
  DataLoader dataLoader{nullptr};
  DataLoader devDataLoader{nullptr};
  std::size_t epochNumber{0};
  int batchSize;

  private :

  void extractExamples(SubConfig & config, bool debug, std::filesystem::path dir, int epoch, int dynamicOracleInterval);
  float processDataset(DataLoader & loader, bool train, bool printAdvancement, int nbExamples);
  void fillDicts(SubConfig & config);

  public :

  Trainer(ReadingMachine & machine, int batchSize);
  void createDataset(BaseConfig & goldConfig, bool debug, std::filesystem::path dir, int epoch, int dynamicOracleInterval);
  void createDevDataset(BaseConfig & goldConfig, bool debug, std::filesystem::path dir, int epoch, int dynamicOracleInterval);
  void fillDicts(BaseConfig & goldConfig);
  float epoch(bool printAdvancement);
  float evalOnDev(bool printAdvancement);
};

#endif