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

// slam.hpp — SLAM パイプライン
//
// 新規実装:
//   - Visual Odometry + マップ管理 + ループ閉合 + Pose Graph 最適化
//   - 軌跡とマップの可視化

#ifndef CALX_IMAGE_CV_SLAM_HPP
#define CALX_IMAGE_CV_SLAM_HPP

#include <image/Image.hpp>
#include <image/io/bmp.hpp>
#include <image/cv/camera.hpp>
#include <image/cv/orb.hpp>
#include <image/cv/matching.hpp>
#include <image/cv/essential.hpp>
#include <image/cv/triangulation.hpp>
#include <image/cv/pose_graph.hpp>
#include <image/cv/bundle_adjustment.hpp>
#include <image/cv/loop_closure.hpp>
#include <array>
#include <cmath>
#include <vector>

namespace calx::cv {

    /// SLAM システム
    class SlamSystem {
    public:
        explicit SlamSystem(const PinholeCamera& cam, int vocabSize = 50)
            : camera_(cam), loopDetector_(vocabSize), frameCount_(0) {
            // 初期ポーズ (原点)
            poseGraph_.addNode(SE3Pose());
        }

        /// フレームを処理
        template<typename T>
        void processFrame(const image::Image<T>& frame) {
            // ORB 特徴点
            auto [kps, descs] = orbDetectAndCompute(frame, 300, 4,
                std::is_floating_point_v<T> ? T{0.05} : T{15});

            if (frameCount_ == 0) {
                // 辞書の初期化用にデータを蓄積
                allDescs_.insert(allDescs_.end(), descs.begin(), descs.end());
            }

            if (frameCount_ > 0 && !prevDescs_.empty() && descs.size() >= 8) {
                // マッチング
                auto knnResults = knnMatch(descs, prevDescs_);
                auto good = ratioTest(knnResults, 0.8);

                SE3Pose relativePose;

                if (good.size() >= 8) {
                    // Essential 行列推定
                    std::vector<NormalizedMatch> normMatches;
                    for (const auto& m : good) {
                        auto [x1, y1] = camera_.unproject(kps[m.queryIdx].x, kps[m.queryIdx].y);
                        auto [x2, y2] = camera_.unproject(prevKps_[m.trainIdx].x, prevKps_[m.trainIdx].y);
                        normMatches.push_back({x2, y2, x1, y1});
                    }

                    auto [E, inliers] = findEssentialRansac(normMatches, 200, 0.01);

                    if (inliers.size() >= 8) {
                        auto candidates = decomposeEssential(E);
                        std::vector<NormalizedMatch> inlierMatches;
                        for (int idx : inliers) inlierMatches.push_back(normMatches[idx]);
                        auto [R, t] = selectCorrectPose(candidates, inlierMatches);
                        relativePose = SE3Pose(R, t);
                    }
                }

                // ポーズグラフにノードとエッジを追加
                SE3Pose prevPose = poseGraph_.node(frameCount_ - 1);
                SE3Pose newPose = prevPose.compose(relativePose);
                poseGraph_.addNode(newPose);
                poseGraph_.addEdge(frameCount_ - 1, frameCount_, relativePose);
            }

            // ループ閉合検出
            if (frameCount_ > 0 && frameCount_ % 5 == 0 && !allDescs_.empty()) {
                // 辞書が未構築なら構築
                if (frameCount_ == 5) {
                    allDescs_.insert(allDescs_.end(), descs.begin(), descs.end());
                    loopDetector_.trainVocabulary(allDescs_);
                    allDescs_.clear();
                }

                loopDetector_.addFrame(descs);
                auto loops = loopDetector_.detectLoops(5, 0.7);

                for (const auto& [loopFrame, sim] : loops) {
                    // ループエッジを追加
                    SE3Pose loopPose = poseGraph_.node(loopFrame).inverse().compose(
                        poseGraph_.node(frameCount_));
                    poseGraph_.addEdge(loopFrame, frameCount_, loopPose, sim);
                }
            }

            // 軌跡を記録
            if (frameCount_ < poseGraph_.numNodes()) {
                trajectory_.push_back(poseGraph_.node(frameCount_).position());
            }

            prevKps_ = kps;
            prevDescs_ = descs;
            ++frameCount_;
        }

        /// ポーズグラフ最適化を実行
        double optimize(int maxIter = 20) {
            return poseGraph_.optimize(maxIter);
        }

        /// 軌跡の可視化
        image::Image<float> drawMap(int imgSize = 600, double scale = 50.0) const {
            image::Image<float> canvas(imgSize, imgSize, 3, 0.05f);
            int cx = imgSize / 2, cy = imgSize / 2;

            // 軌跡
            for (std::size_t i = 1; i < trajectory_.size(); ++i) {
                int x0 = cx + static_cast<int>(trajectory_[i-1][0] * scale);
                int y0 = cy - static_cast<int>(trajectory_[i-1][2] * scale);
                int x1 = cx + static_cast<int>(trajectory_[i][0] * scale);
                int y1 = cy - static_cast<int>(trajectory_[i][2] * scale);
                canvas.drawLine(x0, y0, x1, y1, {0.0f, 1.0f, 0.0f}, 1);
            }

            // 開始点・現在位置
            if (!trajectory_.empty()) {
                int sx = cx + static_cast<int>(trajectory_.front()[0] * scale);
                int sy = cy - static_cast<int>(trajectory_.front()[2] * scale);
                canvas.drawCircle(sx, sy, 4, {1.0f, 0.0f, 0.0f}, 2);

                int ex = cx + static_cast<int>(trajectory_.back()[0] * scale);
                int ey = cy - static_cast<int>(trajectory_.back()[2] * scale);
                canvas.drawCircle(ex, ey, 4, {0.0f, 0.0f, 1.0f}, 2);
            }

            return canvas;
        }

        int frameCount() const { return frameCount_; }
        const std::vector<std::array<double, 3>>& trajectory() const { return trajectory_; }

    private:
        PinholeCamera camera_;
        PoseGraph poseGraph_;
        LoopClosureDetector loopDetector_;
        int frameCount_;

        std::vector<KeyPoint> prevKps_;
        std::vector<OrbDescriptor> prevDescs_;
        std::vector<OrbDescriptor> allDescs_;
        std::vector<std::array<double, 3>> trajectory_;
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_SLAM_HPP
