#ifndef MNIST_HPP
#define MNIST_HPP

#include <iostream>
#include <string>
#include <fstream>
#include <cstdint>

#include "../dataset.hpp"

using namespace std;

class Mnist:public Dataset{
private:
  unsigned char* train_labels;
  unsigned char* test_labels;
  unsigned char* train_images;
  unsigned char* test_images;
  mutable Vector x;
  mutable Vector y;
  size_t load_labels(string filename,unsigned char** dst);
  size_t load_images(string filename,unsigned char** dst);
  int reverse_int32(unsigned char* buffer);
  pair<Vector,Vector> get(const size_t i,const unsigned char* const* labels,const unsigned char* const * images) const;
public:
  Mnist();
  pair<Vector,Vector> get_train(const size_t i) const;
  pair<Vector,Vector> get_test(const size_t i) const;
};

inline pair<Vector,Vector>
Mnist::get_train(const size_t i) const{
  assert(i<train_size);
  return get(i,&train_labels,&train_images);
}

inline pair<Vector,Vector>
Mnist::get_test(const size_t i) const{
  assert(i<test_size);
  return get(i,&test_labels,&test_images);
}
#endif
