#include <opencv2/opencv.hpp>

#include <boost/format.hpp>

#include <cstdint>
#include <vector>
// ============================================================================
class stitcher
{
public:
    virtual ~stitcher() {}

    virtual void add_slice(cv::Mat image) = 0;
    virtual cv::Mat stitched_image() = 0;
};
// ============================================================================
class simple_stitcher
    : public stitcher
{
public:
    simple_stitcher(uint32_t min_overlap_width, uint32_t vertical_tolerance);

    virtual void add_slice(cv::Mat image);
    virtual cv::Mat stitched_image();

private:
    void add_slice(cv::Mat const& image, cv::Point2i const& offset);
    cv::Point2i find_overlap(cv::Mat const& haystack, cv::Mat const& needle) const;

    static bool slice_size_matches(cv::Mat const& source, cv::Mat const& target);

private:
    uint32_t min_overlap_width_;
    uint32_t vertical_tolerance_;

    std::vector<cv::Mat> slices_;
    std::vector<cv::Point2i> offsets_;
};
// ============================================================================
simple_stitcher::simple_stitcher(uint32_t min_overlap_width
    , uint32_t vertical_tolerance)
    : min_overlap_width_(min_overlap_width)
    , vertical_tolerance_(vertical_tolerance)
{
}
// ----------------------------------------------------------------------------
void simple_stitcher::add_slice(cv::Mat image)
{
    if (slices_.empty()) {
        add_slice(image, { 0, 0 });
        return;
    }

    if (!slice_size_matches(image, slices_.front())) {
        throw std::runtime_error("Inconsistent slice image size.");
    }

    cv::Point2i offset(find_overlap(slices_.back(), image));

    add_slice(image, offset);
}
// ----------------------------------------------------------------------------
cv::Mat simple_stitcher::stitched_image()
{
    if (slices_.empty()) {
        return cv::Mat();
    }

    // Calculate slice origins, relative to top-left of stitched image.
    std::vector<cv::Point2i> origins;
    origins.push_back({ 0, 0 });
    for (auto it(++offsets_.begin()); it != offsets_.end(); ++it) {
        origins.push_back(origins.back() + *it);
    }

    auto min_max_y_origin = std::minmax_element(origins.begin(), origins.end()
        , [](cv::Point2i const& l, cv::Point2i const& r) { return l.y < r.y; });

    int32_t min_y_offset(min_max_y_origin.first->y);
    int32_t max_y_offset(min_max_y_origin.second->y);

    uint32_t width(origins.back().x + slices_.back().cols);
    uint32_t height(max_y_offset - min_y_offset + slices_.back().rows);

    cv::Mat result(cv::Mat(height, width, slices_.front().type()));
    for (uint32_t i(0); i < slices_.size(); ++i) {
        cv::Mat const& slice(slices_[i]);
        slice.copyTo(result(cv::Rect(origins[i], slice.size())));
    }

    return result;
}
// ----------------------------------------------------------------------------
void simple_stitcher::add_slice(cv::Mat const& image, cv::Point2i const& offset)
{
    std::cout << boost::format("Adding slice #%d, size=%s, offset=%s\n")
        % slices_.size() % image.size() % offset;

    slices_.push_back(image);
    offsets_.push_back(offset);
}
// ----------------------------------------------------------------------------
cv::Point2i simple_stitcher::find_overlap(cv::Mat const& haystack
    , cv::Mat const& needle) const
{
    cv::Mat needle_roi(needle
        .colRange(0, min_overlap_width_)
        .rowRange(vertical_tolerance_, needle.rows - vertical_tolerance_));

    cv::Mat correlation;
    cv::matchTemplate(haystack
        , needle_roi
        , correlation
        , cv::TM_CCORR_NORMED);

    cv::Point min_point, max_point;
    double min_val, max_val;
    cv::minMaxLoc(correlation, &min_val, &max_val, &min_point, &max_point);

    max_point.y -= vertical_tolerance_;
    return max_point;
}
// ----------------------------------------------------------------------------
bool simple_stitcher::slice_size_matches(cv::Mat const& source, cv::Mat const& target)
{
    return (source.rows == target.rows)
        && (source.cols == target.cols)
        && (source.depth() == target.depth());
}
// ============================================================================
int main()
{
    simple_stitcher s(50, 4);

    for (uint32_t i(0); i < 4; ++i) {
        cv::Mat slice(cv::imread(str(boost::format("slice_%d.png") % i), 0));
        s.add_slice(slice);
    }

    cv::imwrite("stitched.png", s.stitched_image());

    return 0;
}
// ============================================================================