#include <iostream>
#include <fstream>
#include <cstring>
#include <random>
#include <map>
#include <Magick++.h>
#include <fcgio.h>

std::ifstream t10k_images("./t10k-images-idx3-ubyte", std::ios_base::binary | std::ios_base::in);
std::ifstream t10k_labels("./t10k-labels-idx1-ubyte", std::ios_base::binary | std::ios_base::in);
size_t count = 0;

enum {
  IMAGE_IN = 16,
  IMAGE_NUM = 5000,
  IMAGE_SIZE = 28,
  IMAGE_BYTE = IMAGE_SIZE * IMAGE_SIZE,
  LABEL_IN = 8,
};

std::multimap<int, size_t> categories;
std::random_device rd;
std::mt19937 mt(rd());

void init() {
  t10k_labels.seekg(LABEL_IN);
  for (size_t i = 0; i < IMAGE_NUM; i++) {
    unsigned char c;
    t10k_labels.read(reinterpret_cast<char*>(&c), 1);
    categories.insert({c, i * IMAGE_BYTE + IMAGE_IN});
  }
  std::ifstream fcounter("./fcounter.db", std::ios_base::binary | std::ios_base::in);
  fcounter.read(reinterpret_cast<char*>(&count), sizeof(count));
  fcounter.close();
}

void select(std::array<unsigned char, IMAGE_BYTE>& img, unsigned char c) {
  auto range = categories.equal_range(c);
  auto first = range.first; auto last = range.second;
  auto n = std::distance(first, last);
  std::uniform_int_distribution<> dist(0, n - 1);
  auto sk = std::next(first, dist(mt))->second;
  t10k_images.seekg(sk);
  t10k_images.read(reinterpret_cast<char*>(img.data()), IMAGE_BYTE);
}

void hit(std::ostream& os) {
  count++;
  std::ofstream fcounter("./fcounter.db", std::ios_base::binary | std::ios_base::out);
  fcounter.write(reinterpret_cast<char*>(&count), sizeof(count));
  fcounter.close();
  std::string str = std::to_string(count);
  if (str.length() < 6)
    str = std::string(6 - str.length(), '0') + str;
  size_t w = IMAGE_SIZE * str.length(), h = IMAGE_SIZE;
  std::vector<unsigned char> canvas(w*h, 0);
  size_t i = 0;
  for (auto&& c : str) {
    std::array<unsigned char, IMAGE_BYTE> img;
    select(img, c - '0');
    for (int y = 0; y < IMAGE_SIZE; y++) {
      std::memcpy(&canvas[y * w + i * IMAGE_SIZE], &img[y * IMAGE_SIZE], IMAGE_SIZE);
    }
    i++;
  }
  Magick::Image image(IMAGE_SIZE*str.length(), IMAGE_SIZE, "I", Magick::CharPixel, canvas.data());
  Magick::Blob blob;
  image.type(Magick::GrayscaleType);
  image.magick("PNG");
  image.write(&blob);
  os << "Content-Type: image/png\r\n";
  os << "Content-length: " << blob.length() << "\r\n\r\n";
  os.write(reinterpret_cast<const char*>(blob.data()), blob.length()) << std::flush;
}

int main() {
  FCGX_Request request;
  init();
  FCGX_Init();
  FCGX_InitRequest(&request, 0, 0);
  while (FCGX_Accept_r(&request) == 0) {
    fcgi_streambuf osbuf(request.out);
    std::ostream os(&osbuf);
    hit(os);
  }
  return 0;
}