path: root/wizard.cpp

#include "wizard.h"

#include <Magick++.h>
#include <allheaders.h>
#include <hkutil/string.h>
#include <tesseract/baseapi.h>

#include <iostream>
#include <list>
#include <memory>
#include <random>
#include <set>
#include <tuple>
#include <vector>

#include "designer.h"
#include "prefix_search.h"

std::string stripSpaces(std::string in) {
  in.erase(std::remove_if(std::begin(in), std::end(in), ::isspace),
           std::end(in));

  return in;
}

class tesseract_deleter {
 public:
  void operator()(tesseract::TessBaseAPI* ptr) const { ptr->End(); }
};

using tesseract_ptr =
    std::unique_ptr<tesseract::TessBaseAPI, tesseract_deleter>;

class pix_deleter {
 public:
  void operator()(Pix* ptr) const { pixDestroy(&ptr); }
};

using pix_ptr = std::unique_ptr<Pix, pix_deleter>;

wizard::wizard(const cardset& cards, const imagestore& images, std::string text,
               std::mt19937& rng)
    : cards_(cards), images_(images), text_(text), rng_(rng()) {
  std::cout << "Characters: ";

  for (char ch : cards_.getCharacters()) {
    std::cout << ch;
  }

  std::cout << std::endl;
}

void wizard::set_status_callback(status_callback_type callback) {
  status_callback_ = callback;
}

Magick::Image wizard::run() {
  tesseract_ptr api{new tesseract::TessBaseAPI()};

  if (api->Init(nullptr, "eng", tesseract::OEM_TESSERACT_LSTM_COMBINED)) {
    throw std::runtime_error("Could not initialize tesseract");
  }

  std::string allChars;
  for (char ch : cards_.getCharacters()) {
    allChars.push_back(ch);
    if (std::isalpha(ch)) {
      allChars.push_back(std::toupper(ch));
    }
  }

  if (!api->SetVariable("tessedit_char_whitelist", allChars.c_str())) {
    throw std::runtime_error("Could not set tesseract allowlist");
  }

  std::string text = text_;  //"what the heck, it's just some gay guy";
  // getline(std::cin, text);
  if (text.back() == '\r') {
    text.pop_back();
  }

  std::cout << "Calculating card list..." << std::endl;
  if (status_callback_) {
    status_callback_("Calculating card list...");
  }

  std::string canonText = hatkirby::lowercase(text);
  designer des(canonText, cards_.getTitles());

  Magick::Image endImage;
  bool success = false;
  for (int attempt = 0; attempt < 10; attempt++) {
    endImage = {};

    std::list<usage> res = des.generate(rng_);

    bool firstSlice = false;
    int ui = 0;

    bool failure = false;
    for (const usage& u : res) {
      const card& theCard = cards_.getCard(u.cardId);
      const std::string& cardName = theCard.name;

      std::cout << cardName.substr(0, u.strIndex) << "["
                << cardName.substr(u.strIndex, u.strLen) << "]"
                << cardName.substr(u.strIndex + u.strLen) << std::endl;

      std::cout << "Downloading image..." << std::endl;
      if (status_callback_) {
        std::ostringstream status_stream;
        status_stream << "Processing card " << (ui + 1) << "/" << res.size()
                      << " (attempt " << (attempt + 1) << ")...";
        status_callback_(status_stream.str());
      }

      Magick::Image cardImg = images_.get(theCard.uuid, theCard.imageUri);

      std::cout << "Reading text..." << std::endl;

      Magick::Image titleImg = cardImg;
      titleImg.magick("TIFF");

      // titleImg.threshold(MaxRGB / 2);
      titleImg.write("pre.tif");

      Magick::Geometry margin;

      if (theCard.frame == card_frame::m2015) {
        margin = Magick::Geometry{595, 50, 57, 54};
      } else if (theCard.frame == card_frame::modern) {
        margin = Magick::Geometry{581, 50, 63, 57};
      }

      titleImg.crop(margin);
      titleImg.zoom({margin.width() * 5, margin.height() * 5});

      // titleImg.quantizeColorSpace(Magick::GRAYColorspace);
      // titleImg.quantizeColors(2);
      // titleImg.quantize();
      titleImg.backgroundColor("white");
      titleImg.matte(false);
      titleImg.resolutionUnits(Magick::PixelsPerInchResolution);
      titleImg.density({300, 300});
      titleImg.type(Magick::GrayscaleType);
      titleImg.gamma(3);
      // titleImg.level(0.2, 0.8);

      titleImg.write("title.tif");

      Magick::Blob titleBlob;
      titleImg.write(&titleBlob);

      pix_ptr titlePix{pixReadMemTiff(
          reinterpret_cast<const unsigned char*>(titleBlob.data()),
          titleBlob.length(), 0)};

      api->SetImage(titlePix.get());
      api->Recognize(nullptr);

      tesseract::ResultIterator* ri = api->GetIterator();
      tesseract::PageIteratorLevel level = tesseract::RIL_TEXTLINE;
      bool foundName = false;
      size_t extraChars = 0;

      if (ri) {
        do {
          const char* line = ri->GetUTF8Text(level);

          if (line) {
            std::string lineStr(line);

            size_t leadin = hatkirby::lowercase(lineStr).find(
                hatkirby::lowercase((cardName)));
            if (leadin != std::string::npos) {
              foundName = true;
              extraChars = leadin;

              break;
            } /* else {
                std::cout << "WRONG: " << lineStr << std::endl;
              }*/
          }
        } while (ri->Next(level));
      }

      if (!foundName) {
        failure = true;
        break;
      }

      level = tesseract::RIL_SYMBOL;

      for (int i = 0; i < extraChars; i++) {
        ri->Next(level);
      }

      std::vector<std::tuple<unsigned int, unsigned int>> characters;
      size_t cur = 0;

      do {
        int x1, y1, x2, y2;
        ri->BoundingBox(level, &x1, &y1, &x2, &y2);

        x1 /= 5;
        x2 /= 5;

        if (cardName.at(cur) == ' ') {
          if (cur == 0) {
            characters.emplace_back(0, x1);
          } else {
            const auto& prev = characters.back();
            characters.emplace_back(std::get<1>(characters.back()), x1);
          }

          cur++;
        }

        characters.emplace_back(x1, x2);
        std::cout << characters.size() << " " << x1 << "," << x2 << " ("
                  << ri->GetUTF8Text(level) << ")" << std::endl;

        cur++;
      } while (ri->Next(level) && (cur < cardName.length()));

      if (cur != cardName.length()) {
        throw std::runtime_error("Error detecting character bounds");
      }

      Magick::Image bakImg = cardImg;

      unsigned int left = margin.xOff() + std::get<0>(characters[u.strIndex]);
      unsigned int right =
          margin.xOff() + std::get<1>(characters[u.strIndex + u.strLen - 1]);

      cardImg.crop({right - left, cardImg.rows(), left, 0});

      if (ui == 0) {
        bakImg.crop({margin.xOff(), bakImg.rows(), 0, 0});

        cardImg.extent({cardImg.columns() + bakImg.columns(), cardImg.rows()},
                       "black", Magick::GravityType::EastGravity);

        cardImg.composite(bakImg, 0, 0);
      } else if (ui == (res.size() - 1)) {
        bakImg.crop(
            {bakImg.columns() - margin.xOff() - std::get<1>(characters.back()),
             bakImg.rows(), margin.xOff() + std::get<1>(characters.back()), 0});

        int xoff = cardImg.columns();

        cardImg.extent({cardImg.columns() + bakImg.columns(), cardImg.rows()});

        cardImg.composite(bakImg, xoff, 0);
      }

      cardImg.magick("PNG");
      cardImg.write("slice.png");

      if (!firstSlice) {
        firstSlice = true;
        endImage = cardImg;

        if (theCard.frame == card_frame::m2015) {
          // endImage.extent({endImage.columns(), endImage.rows() + 6},
          // Magick::SouthGravity);
        }
      } else {
        int xoff = endImage.columns();

        endImage.backgroundColor("black");

        endImage.extent(
            {endImage.columns() + cardImg.columns(), cardImg.rows()},
            Magick::WestGravity);

        endImage.composite(cardImg, xoff,
                           (theCard.frame == card_frame::m2015) ? 6 : 0);
      }

      // break;
      ui++;
    }

    if (!failure) {
      success = true;
      break;
    }
  }

  if (!success) {
    throw std::domain_error("Could not generate card");
  }

  endImage.magick("PNG");

  return endImage;
}