#include "alvin/vision.hpp"

#include <algorithm>

#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

namespace alvin::vision {

cv::Mat view(const Frame& frame) {
    return cv::Mat(frame.height(), frame.width(), CV_8UC4, const_cast<std::uint8_t*>(frame.data()), frame.stride());
}

cv::Mat toBgr(const Frame& frame) {
    cv::Mat out;
    cv::cvtColor(view(frame), out, cv::COLOR_BGRA2BGR);
    return out;
}

Result<Frame> fromMat(const cv::Mat& mat) {
    if (mat.empty()) return fail(ErrorCode::InvalidArgument, "empty image");
    if (mat.depth() != CV_8U) return fail(ErrorCode::InvalidArgument, "only 8-bit images are supported");
    Frame frame(mat.cols, mat.rows);
    cv::Mat dst = view(frame);
    switch (mat.channels()) {
        case 1: cv::cvtColor(mat, dst, cv::COLOR_GRAY2BGRA); break;
        case 3: cv::cvtColor(mat, dst, cv::COLOR_BGR2BGRA); break;
        case 4: mat.copyTo(dst); break;
        default: return fail(ErrorCode::InvalidArgument, "unsupported channel count");
    }
    return frame;
}

Result<cv::Mat> loadImage(const std::filesystem::path& path) {
    cv::Mat img = cv::imread(path.string(), cv::IMREAD_COLOR);
    if (img.empty()) return fail(ErrorCode::Io, "cannot read image " + path.string());
    return img;
}

Status saveImage(const Frame& frame, const std::filesystem::path& path) {
    try {
        if (!cv::imwrite(path.string(), view(frame))) return fail(ErrorCode::Io, "cannot write image " + path.string());
    } catch (const cv::Exception& e) {
        return fail(ErrorCode::Io, e.what());
    }
    return {};
}

namespace {

// Brings both images to the same channel layout and returns the TM_CCOEFF_NORMED score map.
Result<cv::Mat> scoreMap(const Frame& haystack, const cv::Mat& needle, bool grayscale) {
    if (needle.empty()) return fail(ErrorCode::InvalidArgument, "empty template");
    if (needle.cols > haystack.width() || needle.rows > haystack.height())
        return fail(ErrorCode::InvalidArgument, "template is larger than the frame");

    cv::Mat hay, tpl;
    if (grayscale) {
        cv::cvtColor(view(haystack), hay, cv::COLOR_BGRA2GRAY);
        if (needle.channels() == 1) tpl = needle;
        else cv::cvtColor(needle, tpl, needle.channels() == 4 ? cv::COLOR_BGRA2GRAY : cv::COLOR_BGR2GRAY);
    } else {
        cv::cvtColor(view(haystack), hay, cv::COLOR_BGRA2BGR);
        if (needle.channels() == 3) tpl = needle;
        else cv::cvtColor(needle, tpl, needle.channels() == 4 ? cv::COLOR_BGRA2BGR : cv::COLOR_GRAY2BGR);
    }
    cv::Mat scores;
    cv::matchTemplate(hay, tpl, scores, cv::TM_CCOEFF_NORMED);
    return scores;
}

}  // namespace

Result<std::optional<Match>> find(const Frame& haystack, const cv::Mat& needle, const MatchOptions& opts) {
    auto scores = scoreMap(haystack, needle, opts.grayscale);
    if (!scores) return std::unexpected(scores.error());
    double best = 0;
    cv::Point loc;
    cv::minMaxLoc(*scores, nullptr, &best, nullptr, &loc);
    if (best < opts.threshold) return std::optional<Match>{};
    return std::optional{Match{{loc.x, loc.y, needle.cols, needle.rows}, best}};
}

Result<std::vector<Match>> findAll(const Frame& haystack, const cv::Mat& needle, const MatchOptions& opts) {
    auto scores = scoreMap(haystack, needle, opts.grayscale);
    if (!scores) return std::unexpected(scores.error());
    std::vector<Match> out;
    // Greedy non-maximum suppression: take the best peak, blank out its neighbourhood, repeat.
    while (static_cast<int>(out.size()) < opts.maxResults) {
        double best = 0;
        cv::Point loc;
        cv::minMaxLoc(*scores, nullptr, &best, nullptr, &loc);
        if (best < opts.threshold) break;
        out.push_back({{loc.x, loc.y, needle.cols, needle.rows}, best});
        cv::Rect suppress(loc.x - needle.cols / 2, loc.y - needle.rows / 2, needle.cols, needle.rows);
        suppress &= cv::Rect(0, 0, scores->cols, scores->rows);
        (*scores)(suppress).setTo(-1.0f);
    }
    return out;
}

}  // namespace alvin::vision