#include <yandex/maps/mrc/tensorflow/rotationlib.h>

#include <maps/libs/cmdline/include/cmdline.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/log8/include/log8.h>

#include <opencv2/opencv.hpp>


namespace {

const size_t CHANNELS_NUM = 3;
const std::string INPUT_LAYER_NAME = "inference_input";
const std::string OUTPUT_SOFTMAX_LAYER_NAME = "inference_softmax";
const std::string OUTPUT_CLASS_NAMES_LAYER_NAME = "class_names";

std::vector<uint8_t> toBGR(const cv::Mat& image)
{
    std::vector<uint8_t> result;
    result.reserve(image.size().area() * 3);
    if (image.isContinuous()) {
        result.assign(image.datastart, image.dataend);
    }
    else {
        for (int i = 0; i < image.rows; ++i) {
            result.insert(result.end(), image.ptr<uint8_t>(i),
                          image.ptr<uint8_t>(i)
                              + image.cols * image.channels());
        }
    }
    return result;
}


} // namespace

int main(int argc, char** argv)
 try {
    maps::cmdline::Parser parser;
    auto modelPath
        = parser.string("model").required().help("path to trained model");

    auto imagePath
        = parser.string("image").required().help("path to trained model");

    parser.parse(argc, argv);
    INFO() << "loading file " << imagePath;
    auto image = cv::imread(imagePath);

    INFO() << "loading model " << modelPath;
    maps::mrc::rotationdetect::RotationClassifier classifier(modelPath);

    INFO() << "supported classes:";
    for (const auto& name : classifier.classNames()) {
        INFO() << "\t" << name;
    }

    INFO() << "classifying image";
    auto result = classifier.classify(
        toBGR(image),
        image.cols,
        image.rows
    );

    INFO() << result.first << "\t" << result.second;

    return EXIT_SUCCESS;
}
catch (const maps::Exception& e) {
    INFO() << e;
    return EXIT_FAILURE;
}
catch (const std::exception& e) {
    INFO() << e.what();
    return EXIT_FAILURE;
}
catch (...) {
    INFO() << "Caught unknown exception";
    return EXIT_FAILURE;
}
