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

// kcf.hpp — KCF トラッカー (Kernelized Correlation Filter)
//
// 新規実装:
//   - ガウスカーネル相関フィルタ
//   - 2D FFT による高速相関計算
//   - コサイン窓 + HOG 特徴量 (簡易版)
//   - J. F. Henriques et al., "High-Speed Tracking with Kernelized
//     Correlation Filters", TPAMI 2015.

#ifndef CALX_IMAGE_CV_KCF_HPP
#define CALX_IMAGE_CV_KCF_HPP

#include <image/Image.hpp>
#include <image/cv/tracker.hpp>    // BBox
#include <math/fft/fft2d.hpp>
#include <math/core/matrix.hpp>
#include <algorithm>
#include <cmath>
#include <complex>
#include <numbers>
#include <vector>

namespace calx::cv {

    /// KCF トラッカー
    class KCFTracker {
    public:
        /// @param sigma        ガウスカーネルの帯域幅
        /// @param lambda       正則化パラメータ
        /// @param interpFactor 学習率 (モデル更新の速さ)
        /// @param padding      パッチの拡張率 (1.0 = 拡張なし, 2.5 が一般的)
        KCFTracker(double sigma = 0.5, double lambda = 1e-4,
                   double interpFactor = 0.075, double padding = 2.5)
            : sigma_(sigma), lambda_(lambda)
            , interpFactor_(interpFactor), padding_(padding)
            , initialized_(false) {}

        /// 初期化 (最初のフレーム + 初期 BBox)
        template<typename T>
        void init(const image::Image<T>& frame, const BBox& bbox) {
            bbox_ = bbox;

            // パッチサイズ (パディング適用 + 2 の累乗に丸め)
            patchW_ = nextPow2(static_cast<int>(bbox.w * padding_));
            patchH_ = nextPow2(static_cast<int>(bbox.h * padding_));

            // コサイン窓を生成
            cosWindow_ = createCosineWindow(patchH_, patchW_);

            // 目標応答 (ガウス形状)
            yf_ = createGaussianTarget(patchH_, patchW_, sigma_);

            // 特徴量を抽出
            auto patch = extractPatch(frame, bbox_);
            auto x = applyWindow(patch, cosWindow_);

            // フィルタ学習
            auto xf = calx::fft2d_real(x);
            auto kxxf = gaussianCorrelation(xf, xf, sigma_);

            // αf = yf / (kxxf + λ)
            alphaf_ = calx::Matrix<std::complex<double>>(patchH_, patchW_);
            for (std::size_t r = 0; r < static_cast<std::size_t>(patchH_); ++r) {
                for (std::size_t c = 0; c < static_cast<std::size_t>(patchW_); ++c) {
                    alphaf_(r, c) = yf_(r, c) / (kxxf(r, c) + lambda_);
                }
            }

            modelXf_ = xf;
            initialized_ = true;
        }

        /// 追跡 (新しいフレームで位置を更新)
        /// @return 更新後の BBox
        template<typename T>
        BBox update(const image::Image<T>& frame) {
            if (!initialized_) return bbox_;

            // 前回位置でパッチ抽出
            auto patch = extractPatch(frame, bbox_);
            auto z = applyWindow(patch, cosWindow_);
            auto zf = calx::fft2d_real(z);

            // 相関応答
            auto kzxf = gaussianCorrelation(zf, modelXf_, sigma_);

            // response = IFFT(alphaf * kzxf)
            calx::Matrix<std::complex<double>> respF(patchH_, patchW_);
            for (std::size_t r = 0; r < static_cast<std::size_t>(patchH_); ++r)
                for (std::size_t c = 0; c < static_cast<std::size_t>(patchW_); ++c)
                    respF(r, c) = alphaf_(r, c) * kzxf(r, c);

            auto respC = calx::ifft2d(respF);

            // 応答のピークを探す
            double maxVal = -1e20;
            int maxR = 0, maxC = 0;
            for (std::size_t r = 0; r < static_cast<std::size_t>(patchH_); ++r) {
                for (std::size_t c = 0; c < static_cast<std::size_t>(patchW_); ++c) {
                    double v = respC(r, c).real();
                    if (v > maxVal) { maxVal = v; maxR = static_cast<int>(r); maxC = static_cast<int>(c); }
                }
            }

            // ピーク位置を変位に変換 (周期的シフト)
            double dy = maxR;
            double dx = maxC;
            if (dy > patchH_ / 2) dy -= patchH_;
            if (dx > patchW_ / 2) dx -= patchW_;

            // BBox 更新
            bbox_.x += dx;
            bbox_.y += dy;

            // モデル更新 (指数移動平均)
            auto newPatch = extractPatch(frame, bbox_);
            auto newX = applyWindow(newPatch, cosWindow_);
            auto newXf = calx::fft2d_real(newX);
            auto kxxf = gaussianCorrelation(newXf, newXf, sigma_);

            calx::Matrix<std::complex<double>> newAlphaf(patchH_, patchW_);
            for (std::size_t r = 0; r < static_cast<std::size_t>(patchH_); ++r) {
                for (std::size_t c = 0; c < static_cast<std::size_t>(patchW_); ++c) {
                    newAlphaf(r, c) = yf_(r, c) / (kxxf(r, c) + lambda_);
                }
            }

            double lr = interpFactor_;
            for (std::size_t r = 0; r < static_cast<std::size_t>(patchH_); ++r) {
                for (std::size_t c = 0; c < static_cast<std::size_t>(patchW_); ++c) {
                    alphaf_(r, c) = (1.0 - lr) * alphaf_(r, c) + lr * newAlphaf(r, c);
                    modelXf_(r, c) = (1.0 - lr) * modelXf_(r, c) + lr * newXf(r, c);
                }
            }

            return bbox_;
        }

        /// 現在の BBox
        BBox bbox() const { return bbox_; }

        /// 初期化済みか
        bool initialized() const { return initialized_; }

    private:
        double sigma_, lambda_, interpFactor_, padding_;
        bool initialized_;
        BBox bbox_;
        int patchW_ = 0, patchH_ = 0;

        calx::Matrix<double> cosWindow_;
        calx::Matrix<std::complex<double>> yf_;      // 目標応答 (周波数)
        calx::Matrix<std::complex<double>> alphaf_;   // フィルタ係数
        calx::Matrix<std::complex<double>> modelXf_;  // モデル特徴量 (周波数)

        /// 2 の累乗に丸め上げ
        static int nextPow2(int n) {
            int p = 1;
            while (p < n) p <<= 1;
            return p;
        }

        /// コサイン窓を生成
        static calx::Matrix<double> createCosineWindow(int h, int w) {
            calx::Matrix<double> win(h, w);
            for (int y = 0; y < h; ++y) {
                double wy = 0.5 * (1.0 - std::cos(2.0 * std::numbers::pi * y / h));
                for (int x = 0; x < w; ++x) {
                    double wx = 0.5 * (1.0 - std::cos(2.0 * std::numbers::pi * x / w));
                    win(y, x) = wy * wx;
                }
            }
            return win;
        }

        /// ガウス形状の目標応答 (周波数領域)
        static calx::Matrix<std::complex<double>> createGaussianTarget(int h, int w, double sigma) {
            calx::Matrix<double> y(h, w);
            double sig2 = 2.0 * sigma * sigma;
            for (int r = 0; r < h; ++r) {
                int dy = r;
                if (dy > h / 2) dy -= h;
                for (int c = 0; c < w; ++c) {
                    int dx = c;
                    if (dx > w / 2) dx -= w;
                    y(r, c) = std::exp(-(dx * dx + dy * dy) / sig2);
                }
            }
            return calx::fft2d_real(y);
        }

        /// 画像からパッチを抽出 (バイリニア補間、パディング)
        template<typename T>
        calx::Matrix<double> extractPatch(const image::Image<T>& frame, const BBox& box) const {
            calx::Matrix<double> patch(patchH_, patchW_, 0.0);
            int imgH = static_cast<int>(frame.height());
            int imgW = static_cast<int>(frame.width());

            double cx = box.x, cy = box.y;
            double scaleX = box.w * padding_ / patchW_;
            double scaleY = box.h * padding_ / patchH_;

            for (int r = 0; r < patchH_; ++r) {
                for (int c = 0; c < patchW_; ++c) {
                    double srcX = cx + (c - patchW_ / 2.0) * scaleX;
                    double srcY = cy + (r - patchH_ / 2.0) * scaleY;

                    int x0 = static_cast<int>(std::floor(srcX));
                    int y0 = static_cast<int>(std::floor(srcY));
                    if (x0 < 0 || x0 >= imgW - 1 || y0 < 0 || y0 >= imgH - 1) continue;

                    double fx = srcX - x0;
                    double fy = srcY - y0;
                    patch(r, c) = static_cast<double>(
                        (1 - fy) * (1 - fx) * frame(y0, x0) +
                        (1 - fy) * fx       * frame(y0, x0 + 1) +
                        fy       * (1 - fx) * frame(y0 + 1, x0) +
                        fy       * fx       * frame(y0 + 1, x0 + 1));
                }
            }
            return patch;
        }

        /// コサイン窓を適用
        static calx::Matrix<double> applyWindow(const calx::Matrix<double>& patch,
                                                  const calx::Matrix<double>& window) {
            std::size_t h = patch.rows(), w = patch.cols();
            calx::Matrix<double> result(h, w);

            // 平均を引いて正規化
            double mean = 0;
            for (std::size_t r = 0; r < h; ++r)
                for (std::size_t c = 0; c < w; ++c)
                    mean += patch(r, c);
            mean /= (h * w);

            for (std::size_t r = 0; r < h; ++r)
                for (std::size_t c = 0; c < w; ++c)
                    result(r, c) = (patch(r, c) - mean) * window(r, c);

            return result;
        }

        /// ガウスカーネル相関 (周波数領域)
        /// k(x, x') = exp(-1/(σ²) * (||x||² + ||x'||² - 2 * corr(x, x')) / n)
        static calx::Matrix<std::complex<double>> gaussianCorrelation(
            const calx::Matrix<std::complex<double>>& xf,
            const calx::Matrix<std::complex<double>>& yf,
            double sigma) {

            std::size_t h = xf.rows(), w = xf.cols();
            double n = static_cast<double>(h * w);

            // ||x||² と ||y||² (Parseval: ||x||² = (1/N)Σ|Xf|²)
            double xxNorm = 0, yyNorm = 0;
            for (std::size_t r = 0; r < h; ++r) {
                for (std::size_t c = 0; c < w; ++c) {
                    xxNorm += std::norm(xf(r, c));
                    yyNorm += std::norm(yf(r, c));
                }
            }
            xxNorm /= n;
            yyNorm /= n;

            // xf* ⊙ yf → IFFT → 相互相関 (空間領域)
            calx::Matrix<std::complex<double>> xyf(h, w);
            for (std::size_t r = 0; r < h; ++r)
                for (std::size_t c = 0; c < w; ++c)
                    xyf(r, c) = std::conj(xf(r, c)) * yf(r, c);

            auto xy = calx::ifft2d(xyf);

            // ガウスカーネル: k = exp(-(||x||² + ||y||² - 2*xy) / (n * σ²))
            // xy(r,c).real() は IFFT 結果で、空間領域の相互相関
            double sig2 = sigma * sigma;
            calx::Matrix<double> k(h, w);
            for (std::size_t r = 0; r < h; ++r) {
                for (std::size_t c = 0; c < w; ++c) {
                    double crossCorr = xy(r, c).real();
                    double dist = std::max(0.0, xxNorm + yyNorm - 2.0 * crossCorr);
                    k(r, c) = std::exp(-dist / (n * sig2));
                }
            }

            return calx::fft2d_real(k);
        }
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_KCF_HPP
