// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// matching.hpp — 特徴点マッチング
//
// 新規実装:
//   - ブルートフォース (ハミング距離)
//   - Lowe's ratio test
//   - KNN マッチング

#ifndef CALX_IMAGE_CV_MATCHING_HPP
#define CALX_IMAGE_CV_MATCHING_HPP

#include <image/cv/orb.hpp>
#include <algorithm>
#include <limits>
#include <vector>

namespace calx::cv {

    /// マッチ結果
    struct DMatch {
        int queryIdx;   // クエリ側のインデックス
        int trainIdx;   // トレイン側のインデックス
        int distance;   // ハミング距離

        DMatch() : queryIdx(-1), trainIdx(-1), distance(std::numeric_limits<int>::max()) {}
        DMatch(int q, int t, int d) : queryIdx(q), trainIdx(t), distance(d) {}
    };

    /// ブルートフォースマッチング (ハミング距離)
    /// @param desc1     クエリ記述子
    /// @param desc2     トレイン記述子
    /// @return          最良マッチのリスト
    inline std::vector<DMatch> bruteForceMatch(
        const std::vector<OrbDescriptor>& desc1,
        const std::vector<OrbDescriptor>& desc2) {

        std::vector<DMatch> matches;
        matches.reserve(desc1.size());

        for (std::size_t i = 0; i < desc1.size(); ++i) {
            int bestDist = std::numeric_limits<int>::max();
            int bestIdx = -1;

            for (std::size_t j = 0; j < desc2.size(); ++j) {
                int d = desc1[i].hammingDistance(desc2[j]);
                if (d < bestDist) {
                    bestDist = d;
                    bestIdx = static_cast<int>(j);
                }
            }

            if (bestIdx >= 0) {
                matches.emplace_back(static_cast<int>(i), bestIdx, bestDist);
            }
        }

        return matches;
    }

    /// KNN マッチング (k=2, ratio test 用)
    inline std::vector<std::array<DMatch, 2>> knnMatch(
        const std::vector<OrbDescriptor>& desc1,
        const std::vector<OrbDescriptor>& desc2) {

        std::vector<std::array<DMatch, 2>> matches;
        matches.reserve(desc1.size());

        for (std::size_t i = 0; i < desc1.size(); ++i) {
            DMatch best1, best2;
            best1.queryIdx = static_cast<int>(i);
            best2.queryIdx = static_cast<int>(i);

            for (std::size_t j = 0; j < desc2.size(); ++j) {
                int d = desc1[i].hammingDistance(desc2[j]);
                if (d < best1.distance) {
                    best2 = best1;
                    best1.trainIdx = static_cast<int>(j);
                    best1.distance = d;
                } else if (d < best2.distance) {
                    best2.trainIdx = static_cast<int>(j);
                    best2.distance = d;
                }
            }

            matches.push_back({best1, best2});
        }

        return matches;
    }

    /// Lowe's ratio test でマッチをフィルタ
    /// @param knnMatches  knnMatch の結果
    /// @param ratio       比率閾値 (通常 0.75)
    /// @return            フィルタ済みマッチ
    inline std::vector<DMatch> ratioTest(
        const std::vector<std::array<DMatch, 2>>& knnMatches,
        double ratio = 0.75) {

        std::vector<DMatch> good;
        good.reserve(knnMatches.size());

        for (const auto& knn : knnMatches) {
            if (knn[1].distance > 0 &&
                static_cast<double>(knn[0].distance) / knn[1].distance < ratio) {
                good.push_back(knn[0]);
            }
        }

        return good;
    }

    /// マッチ結果を画像に描画
    /// @param img1, img2  入力画像
    /// @param kp1, kp2    キーポイント
    /// @param matches     マッチ結果
    /// @return            横に並べた画像にマッチ線を描画
    template<typename T>
    image::Image<T> drawMatches(const image::Image<T>& img1, const std::vector<KeyPoint>& kp1,
                                const image::Image<T>& img2, const std::vector<KeyPoint>& kp2,
                                const std::vector<DMatch>& matches) {
        std::size_t h = std::max(img1.height(), img2.height());
        std::size_t w = img1.width() + img2.width();
        int nc = std::max(img1.numChannels(), img2.numChannels());
        if (nc < 3) nc = 3;

        image::Image<T> canvas(h, w, nc);

        // img1 をコピー (必要に応じてグレー→RGB)
        auto copyImg = [&](const image::Image<T>& src, int offsetX) {
            for (std::size_t y = 0; y < src.height(); ++y) {
                for (std::size_t x = 0; x < src.width(); ++x) {
                    if (src.numChannels() >= 3) {
                        for (int c = 0; c < 3; ++c)
                            canvas.at(y, x + offsetX, c) = src.at(y, x, c);
                    } else {
                        T v = src.at(y, x);
                        for (int c = 0; c < 3; ++c)
                            canvas.at(y, x + offsetX, c) = v;
                    }
                }
            }
        };

        copyImg(img1, 0);
        copyImg(img2, static_cast<int>(img1.width()));

        // マッチ線を描画
        T maxVal = std::is_floating_point_v<T> ? T{1} : T{255};
        std::vector<T> green = {T{0}, maxVal, T{0}};
        std::vector<T> red = {maxVal, T{0}, T{0}};

        for (const auto& m : matches) {
            if (m.queryIdx < 0 || m.queryIdx >= static_cast<int>(kp1.size())) continue;
            if (m.trainIdx < 0 || m.trainIdx >= static_cast<int>(kp2.size())) continue;

            int x1 = static_cast<int>(kp1[m.queryIdx].x);
            int y1 = static_cast<int>(kp1[m.queryIdx].y);
            int x2 = static_cast<int>(kp2[m.trainIdx].x) + static_cast<int>(img1.width());
            int y2 = static_cast<int>(kp2[m.trainIdx].y);

            canvas.drawLine(x1, y1, x2, y2, green, 1);
            canvas.drawCircle(x1, y1, 3, red, 1);
            canvas.drawCircle(x2, y2, 3, red, 1);
        }

        return canvas;
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_MATCHING_HPP
