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

// tracker.hpp — 物体追跡 (Kalman + Hungarian)
//
// 新規実装:
//   - カルマンフィルタ予測
//   - Hungarian アルゴリズムでデータ関連付け
//   - 追跡オブジェクトの管理

#ifndef CALX_IMAGE_CV_TRACKER_HPP
#define CALX_IMAGE_CV_TRACKER_HPP

#include <math/core/matrix.hpp>
#include <math/core/vector.hpp>
#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>

namespace calx::cv {

    /// 2D バウンディングボックス
    struct BBox {
        double x, y, w, h;  // 中心座標、幅、高さ

        BBox() : x(0), y(0), w(0), h(0) {}
        BBox(double x, double y, double w, double h) : x(x), y(y), w(w), h(h) {}

        /// IoU (Intersection over Union) を計算
        double iou(const BBox& other) const {
            double x1 = std::max(x - w/2, other.x - other.w/2);
            double y1 = std::max(y - h/2, other.y - other.h/2);
            double x2 = std::min(x + w/2, other.x + other.w/2);
            double y2 = std::min(y + h/2, other.y + other.h/2);

            double interArea = std::max(0.0, x2 - x1) * std::max(0.0, y2 - y1);
            double unionArea = w * h + other.w * other.h - interArea;
            return (unionArea > 0) ? interArea / unionArea : 0;
        }
    };

    /// 簡易カルマンフィルタ (等速モデル)
    class KalmanTracker {
    public:
        KalmanTracker() = default;

        explicit KalmanTracker(const BBox& initial, int id = 0)
            : id_(id), age_(0), hitStreak_(0), timeSinceUpdate_(0) {
            // 状態: [x, y, w, h, vx, vy]
            state_ = calx::Vector<double>(6, 0.0);
            state_[0] = initial.x; state_[1] = initial.y;
            state_[2] = initial.w; state_[3] = initial.h;
        }

        /// 予測
        BBox predict() {
            // 等速モデル: x' = x + vx, y' = y + vy
            state_[0] += state_[4];
            state_[1] += state_[5];
            ++age_;
            ++timeSinceUpdate_;
            return {state_[0], state_[1], state_[2], state_[3]};
        }

        /// 観測で更新
        void update(const BBox& measurement) {
            // 簡易カルマンゲイン (固定)
            double alpha = 0.5;
            state_[4] = alpha * (measurement.x - state_[0]) + (1 - alpha) * state_[4];
            state_[5] = alpha * (measurement.y - state_[1]) + (1 - alpha) * state_[5];
            state_[0] = measurement.x;
            state_[1] = measurement.y;
            state_[2] = alpha * measurement.w + (1 - alpha) * state_[2];
            state_[3] = alpha * measurement.h + (1 - alpha) * state_[3];
            timeSinceUpdate_ = 0;
            ++hitStreak_;
        }

        BBox bbox() const { return {state_[0], state_[1], state_[2], state_[3]}; }
        int id() const { return id_; }
        int age() const { return age_; }
        int timeSinceUpdate() const { return timeSinceUpdate_; }
        int hitStreak() const { return hitStreak_; }

    private:
        int id_ = 0;
        int age_ = 0;
        int hitStreak_ = 0;
        int timeSinceUpdate_ = 0;
        calx::Vector<double> state_;
    };

    /// 貪欲マッチング (Hungarian の簡易版)
    /// @param costMatrix  コスト行列 (M×N)
    /// @param threshold   最大コスト閾値
    /// @return            (検出インデックス, トラッカーインデックス) のペア
    inline std::vector<std::pair<int, int>> greedyMatch(
        const calx::Matrix<double>& costMatrix, double threshold = 0.5) {
        int M = static_cast<int>(costMatrix.rows());
        int N = static_cast<int>(costMatrix.cols());

        std::vector<bool> usedRow(M, false), usedCol(N, false);
        std::vector<std::pair<int, int>> matches;

        // コストの小さい順にマッチング
        while (true) {
            double minCost = threshold;
            int bestI = -1, bestJ = -1;

            for (int i = 0; i < M; ++i) {
                if (usedRow[i]) continue;
                for (int j = 0; j < N; ++j) {
                    if (usedCol[j]) continue;
                    if (costMatrix(i, j) < minCost) {
                        minCost = costMatrix(i, j);
                        bestI = i; bestJ = j;
                    }
                }
            }

            if (bestI < 0) break;
            matches.push_back({bestI, bestJ});
            usedRow[bestI] = true;
            usedCol[bestJ] = true;
        }

        return matches;
    }

    /// マルチオブジェクトトラッカー (SORT 風)
    class MultiObjectTracker {
    public:
        /// @param maxAge      追跡を維持する最大フレーム数 (更新なし)
        /// @param minHits     追跡を出力する最小ヒット数
        /// @param iouThresh   IoU マッチング閾値
        MultiObjectTracker(int maxAge = 5, int minHits = 3, double iouThresh = 0.3)
            : maxAge_(maxAge), minHits_(minHits), iouThresh_(iouThresh), nextId_(0) {}

        /// 検出結果でトラッカーを更新
        /// @param detections  検出されたバウンディングボックス
        /// @return            追跡結果 (ID 付き BBox)
        std::vector<std::pair<int, BBox>> update(const std::vector<BBox>& detections) {
            // 予測
            for (auto& tracker : trackers_) {
                tracker.predict();
            }

            // IoU コスト行列
            int M = static_cast<int>(detections.size());
            int N = static_cast<int>(trackers_.size());

            if (M > 0 && N > 0) {
                calx::Matrix<double> costMatrix(M, N, 1.0);
                for (int i = 0; i < M; ++i) {
                    for (int j = 0; j < N; ++j) {
                        costMatrix(i, j) = 1.0 - detections[i].iou(trackers_[j].bbox());
                    }
                }

                auto matches = greedyMatch(costMatrix, 1.0 - iouThresh_);

                std::vector<bool> matchedDet(M, false), matchedTrack(N, false);
                for (const auto& [di, ti] : matches) {
                    trackers_[ti].update(detections[di]);
                    matchedDet[di] = true;
                    matchedTrack[ti] = true;
                }

                // 未マッチの検出 → 新規トラッカー
                for (int i = 0; i < M; ++i) {
                    if (!matchedDet[i]) {
                        trackers_.emplace_back(detections[i], nextId_++);
                    }
                }
            } else if (M > 0) {
                // トラッカーが空
                for (const auto& det : detections) {
                    trackers_.emplace_back(det, nextId_++);
                }
            }

            // 古いトラッカーを削除
            trackers_.erase(
                std::remove_if(trackers_.begin(), trackers_.end(),
                    [this](const KalmanTracker& t) { return t.timeSinceUpdate() > maxAge_; }),
                trackers_.end());

            // 出力 (十分なヒットがあるもの)
            std::vector<std::pair<int, BBox>> results;
            for (const auto& t : trackers_) {
                if (t.hitStreak() >= minHits_ || t.age() <= minHits_) {
                    results.push_back({t.id(), t.bbox()});
                }
            }

            return results;
        }

    private:
        int maxAge_, minHits_;
        double iouThresh_;
        int nextId_;
        std::vector<KalmanTracker> trackers_;
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_TRACKER_HPP
