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

// orb.hpp — ORB 特徴点 (Oriented FAST + Rotated BRIEF)
//
// 新規実装:
//   - Oriented FAST (intensity centroid で方向推定)
//   - rBRIEF 記述子 (256 bit, 回転不変)
//   - 特許フリー

#ifndef CALX_IMAGE_CV_ORB_HPP
#define CALX_IMAGE_CV_ORB_HPP

#include <image/Image.hpp>
#include <image/cv/fast.hpp>
#include <image/cv/harris.hpp>
#include <image/filters/ndimage.hpp>
#include <image/pyramid.hpp>
#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <numeric>
#include <random>
#include <vector>

namespace calx::cv {

    /// ORB 記述子 (256 bit = 32 bytes)
    struct OrbDescriptor {
        std::array<uint8_t, 32> data{};

        /// ハミング距離を計算
        int hammingDistance(const OrbDescriptor& other) const {
            int dist = 0;
            for (int i = 0; i < 32; ++i) {
                uint8_t xorVal = data[i] ^ other.data[i];
                // popcount
                while (xorVal) { ++dist; xorVal &= xorVal - 1; }
            }
            return dist;
        }
    };

    /// ORB で使う BRIEF サンプリングパターン (事前生成)
    /// 256 ペア × 2 点 × 2 座標 (dx, dy) — パッチ中心からの相対座標
    struct BriefPattern {
        static constexpr int NUM_PAIRS = 256;
        static constexpr int PATCH_SIZE = 31;

        struct Point { int dx, dy; };
        std::array<Point, NUM_PAIRS> p1, p2;

        /// ガウス分布でサンプリングパターンを生成
        BriefPattern() {
            std::mt19937 rng(42);  // 固定シード (再現性)
            std::normal_distribution<double> dist(0.0, PATCH_SIZE / 5.0);
            int half = PATCH_SIZE / 2;

            for (int i = 0; i < NUM_PAIRS; ++i) {
                p1[i].dx = std::clamp(static_cast<int>(dist(rng)), -half, half);
                p1[i].dy = std::clamp(static_cast<int>(dist(rng)), -half, half);
                p2[i].dx = std::clamp(static_cast<int>(dist(rng)), -half, half);
                p2[i].dy = std::clamp(static_cast<int>(dist(rng)), -half, half);
            }
        }
    };

    /// Intensity Centroid で特徴点の方向を計算
    template<typename T>
    double computeOrientation(const calx::Matrix<T>& img, int cx, int cy, int radius = 15) {
        int H = static_cast<int>(img.rows());
        int W = static_cast<int>(img.cols());

        double m01 = 0, m10 = 0;
        for (int dy = -radius; dy <= radius; ++dy) {
            for (int dx = -radius; dx <= radius; ++dx) {
                if (dx * dx + dy * dy > radius * radius) continue;
                int y = cy + dy, x = cx + dx;
                if (y < 0 || y >= H || x < 0 || x >= W) continue;
                double val = static_cast<double>(img(y, x));
                m10 += dx * val;
                m01 += dy * val;
            }
        }
        return std::atan2(m01, m10);
    }

    /// rBRIEF 記述子を計算 (回転不変)
    template<typename T>
    OrbDescriptor computeBrief(const calx::Matrix<T>& img, int cx, int cy,
                               double angle, const BriefPattern& pattern) {
        int H = static_cast<int>(img.rows());
        int W = static_cast<int>(img.cols());
        double cosA = std::cos(angle);
        double sinA = std::sin(angle);

        OrbDescriptor desc;
        for (int i = 0; i < BriefPattern::NUM_PAIRS; ++i) {
            // 回転を適用
            int x1 = cx + static_cast<int>(std::round(cosA * pattern.p1[i].dx - sinA * pattern.p1[i].dy));
            int y1 = cy + static_cast<int>(std::round(sinA * pattern.p1[i].dx + cosA * pattern.p1[i].dy));
            int x2 = cx + static_cast<int>(std::round(cosA * pattern.p2[i].dx - sinA * pattern.p2[i].dy));
            int y2 = cy + static_cast<int>(std::round(sinA * pattern.p2[i].dx + cosA * pattern.p2[i].dy));

            // 境界チェック
            x1 = std::clamp(x1, 0, W - 1); y1 = std::clamp(y1, 0, H - 1);
            x2 = std::clamp(x2, 0, W - 1); y2 = std::clamp(y2, 0, H - 1);

            if (img(y1, x1) < img(y2, x2)) {
                desc.data[i / 8] |= (1 << (i % 8));
            }
        }
        return desc;
    }

    /// ORB 特徴点検出 + 記述子計算
    /// @param img         入力画像 (1ch, グレースケール)
    /// @param maxFeatures 検出する最大特徴点数
    /// @param nLevels     ピラミッドレベル数
    /// @param threshold   FAST 閾値
    /// @return            (キーポイント, 記述子) のペア
    template<typename T>
    std::pair<std::vector<KeyPoint>, std::vector<OrbDescriptor>>
    orbDetectAndCompute(const image::Image<T>& img,
                        int maxFeatures = 500,
                        int nLevels = 8,
                        T threshold = T{20}) {
        // ガウシアンぼかし
        int kSize = 5;
        auto kernel = image::gaussianKernel<T>(kSize, T{1.0});
        image::Image<T> blurred(image::convolve(img.channel(0), kernel));

        // ピラミッド構築
        auto pyramid = image::buildPyramid(blurred, nLevels);

        // 各レベルの特徴点数を分配
        std::vector<int> featuresPerLevel(pyramid.size());
        {
            double factor = 1.0 / 1.2;  // スケールファクタの逆数
            double nDesired = maxFeatures * (1 - factor) /
                             (1 - std::pow(factor, static_cast<double>(pyramid.size())));
            int sum = 0;
            for (std::size_t lv = 0; lv < pyramid.size() - 1; ++lv) {
                featuresPerLevel[lv] = static_cast<int>(nDesired + 0.5);
                sum += featuresPerLevel[lv];
                nDesired *= factor;
            }
            featuresPerLevel.back() = std::max(maxFeatures - sum, 0);
        }

        BriefPattern pattern;
        std::vector<KeyPoint> allKeypoints;
        std::vector<OrbDescriptor> allDescriptors;

        for (std::size_t lv = 0; lv < pyramid.size(); ++lv) {
            const auto& levelImg = pyramid[lv];
            double scale = std::pow(2.0, static_cast<double>(lv));

            // FAST 検出
            auto kps = fastDetect(levelImg, threshold);

            // 上位 N 個に制限
            if (static_cast<int>(kps.size()) > featuresPerLevel[lv]) {
                std::partial_sort(kps.begin(), kps.begin() + featuresPerLevel[lv], kps.end(),
                    [](const KeyPoint& a, const KeyPoint& b) { return a.response > b.response; });
                kps.resize(featuresPerLevel[lv]);
            }

            const auto& ch = levelImg.channel(0);
            int half = BriefPattern::PATCH_SIZE / 2;
            int H = static_cast<int>(levelImg.height());
            int W = static_cast<int>(levelImg.width());

            for (auto& kp : kps) {
                int px = static_cast<int>(kp.x);
                int py = static_cast<int>(kp.y);

                // パッチが画像内に収まるか
                if (px < half || px >= W - half || py < half || py >= H - half) continue;

                // 方向計算
                double angle = computeOrientation(ch, px, py);

                // rBRIEF 記述子
                OrbDescriptor desc = computeBrief(ch, px, py, angle, pattern);

                // 元画像座標に変換
                KeyPoint finalKp(kp.x * scale, kp.y * scale, kp.response, angle,
                                 static_cast<int>(lv), kp.size * scale);
                allKeypoints.push_back(finalKp);
                allDescriptors.push_back(desc);
            }
        }

        return {allKeypoints, allDescriptors};
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_ORB_HPP
