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

// loop_closure.hpp — ループ閉合検出
//
// 新規実装:
//   - Bag of Visual Words (BoVW)
//   - ORB 記述子 → KMeans クラスタリング → BoW ヒストグラム
//   - TF-IDF スコアリング

#ifndef CALX_IMAGE_CV_LOOP_CLOSURE_HPP
#define CALX_IMAGE_CV_LOOP_CLOSURE_HPP

#include <image/cv/orb.hpp>
#include <algorithm>
#include <cmath>
#include <numeric>
#include <random>
#include <vector>

namespace calx::cv {

    /// ビジュアルワード辞書 (KMeans クラスタ中心)
    class VisualVocabulary {
    public:
        /// 辞書を構築 (KMeans)
        /// @param descriptors  学習用記述子
        /// @param k            クラスタ数 (ワード数)
        /// @param maxIter      KMeans の最大反復回数
        void build(const std::vector<OrbDescriptor>& descriptors, int k = 100, int maxIter = 20) {
            if (descriptors.empty() || k <= 0) return;
            k = std::min(k, static_cast<int>(descriptors.size()));

            // ランダム初期化
            words_.resize(k);
            std::mt19937 rng(42);
            std::vector<int> indices(descriptors.size());
            std::iota(indices.begin(), indices.end(), 0);
            for (int i = 0; i < k; ++i) {
                std::uniform_int_distribution<int> dist(i, static_cast<int>(indices.size()) - 1);
                std::swap(indices[i], indices[dist(rng)]);
                words_[i] = descriptors[indices[i]];
            }

            // KMeans 反復
            std::vector<int> assignments(descriptors.size());
            for (int iter = 0; iter < maxIter; ++iter) {
                // 割り当て
                bool changed = false;
                for (std::size_t i = 0; i < descriptors.size(); ++i) {
                    int bestWord = 0;
                    int bestDist = descriptors[i].hammingDistance(words_[0]);
                    for (int w = 1; w < k; ++w) {
                        int d = descriptors[i].hammingDistance(words_[w]);
                        if (d < bestDist) { bestDist = d; bestWord = w; }
                    }
                    if (assignments[i] != bestWord) changed = true;
                    assignments[i] = bestWord;
                }
                if (!changed) break;

                // 中心更新 (ハミング空間での多数決ビット)
                for (int w = 0; w < k; ++w) {
                    int bitCount[256] = {};
                    int count = 0;
                    for (std::size_t i = 0; i < descriptors.size(); ++i) {
                        if (assignments[i] != w) continue;
                        ++count;
                        for (int b = 0; b < 256; ++b) {
                            if (descriptors[i].data[b / 8] & (1 << (b % 8)))
                                ++bitCount[b];
                        }
                    }
                    if (count == 0) continue;
                    words_[w].data.fill(0);
                    for (int b = 0; b < 256; ++b) {
                        if (bitCount[b] > count / 2)
                            words_[w].data[b / 8] |= (1 << (b % 8));
                    }
                }
            }
        }

        /// 記述子を最近傍ワードに割り当て
        int quantize(const OrbDescriptor& desc) const {
            int bestWord = 0;
            int bestDist = desc.hammingDistance(words_[0]);
            for (int w = 1; w < static_cast<int>(words_.size()); ++w) {
                int d = desc.hammingDistance(words_[w]);
                if (d < bestDist) { bestDist = d; bestWord = w; }
            }
            return bestWord;
        }

        int size() const { return static_cast<int>(words_.size()); }

    private:
        std::vector<OrbDescriptor> words_;
    };

    /// BoW ヒストグラム
    struct BoWHistogram {
        std::vector<double> counts;

        /// コサイン類似度
        double similarity(const BoWHistogram& other) const {
            double dot = 0, norm1 = 0, norm2 = 0;
            for (std::size_t i = 0; i < counts.size(); ++i) {
                dot += counts[i] * other.counts[i];
                norm1 += counts[i] * counts[i];
                norm2 += other.counts[i] * other.counts[i];
            }
            double denom = std::sqrt(norm1 * norm2);
            return (denom > 1e-15) ? dot / denom : 0.0;
        }
    };

    /// ループ閉合検出器
    class LoopClosureDetector {
    public:
        explicit LoopClosureDetector(int vocabularySize = 100)
            : vocabSize_(vocabularySize) {}

        /// 辞書を学習
        void trainVocabulary(const std::vector<OrbDescriptor>& descriptors) {
            vocab_.build(descriptors, vocabSize_);
        }

        /// フレームの BoW ヒストグラムを計算・追加
        int addFrame(const std::vector<OrbDescriptor>& descriptors) {
            BoWHistogram hist;
            hist.counts.assign(vocab_.size(), 0.0);

            for (const auto& desc : descriptors) {
                int word = vocab_.quantize(desc);
                hist.counts[word] += 1.0;
            }

            // L1 正規化
            double sum = 0;
            for (double c : hist.counts) sum += c;
            if (sum > 0) {
                for (double& c : hist.counts) c /= sum;
            }

            histograms_.push_back(hist);
            return static_cast<int>(histograms_.size()) - 1;
        }

        /// 最新フレームに対してループ閉合候補を検索
        /// @param minFrameGap  最低フレーム間隔 (近すぎるフレームを除外)
        /// @param threshold    類似度閾値
        /// @return             (フレームインデックス, 類似度) のペア
        std::vector<std::pair<int, double>> detectLoops(int minFrameGap = 30, double threshold = 0.6) const {
            if (histograms_.size() < 2) return {};

            int queryIdx = static_cast<int>(histograms_.size()) - 1;
            const auto& query = histograms_.back();

            std::vector<std::pair<int, double>> candidates;
            for (int i = 0; i < queryIdx - minFrameGap; ++i) {
                double sim = query.similarity(histograms_[i]);
                if (sim > threshold) {
                    candidates.push_back({i, sim});
                }
            }

            // 類似度で降順ソート
            std::sort(candidates.begin(), candidates.end(),
                      [](const auto& a, const auto& b) { return a.second > b.second; });

            return candidates;
        }

    private:
        int vocabSize_;
        VisualVocabulary vocab_;
        std::vector<BoWHistogram> histograms_;
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_LOOP_CLOSURE_HPP
