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

// visual_odometry.hpp — Visual Odometry パイプライン
//
// 新規実装:
//   - フレーム間の SE(3) 推定
//   - 特徴点追跡 → Essential → R,t → 累積姿勢
//   - 軌跡の可視化

#ifndef CALX_IMAGE_CV_VISUAL_ODOMETRY_HPP
#define CALX_IMAGE_CV_VISUAL_ODOMETRY_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 <math/core/matrix.hpp>
#include <array>
#include <cmath>
#include <vector>

namespace calx::cv {

    /// カメラ姿勢 (回転 + 並進)
    struct Pose {
        calx::Matrix<double> R;   // 3x3 回転行列
        calx::Vector<double> t;   // 3x1 並進ベクトル

        Pose() : R(calx::Matrix<double>::identity(3)), t(calx::Vector<double>(3, 0.0)) {}
        Pose(const calx::Matrix<double>& R, const calx::Vector<double>& t) : R(R), t(t) {}

        /// 位置 (ワールド座標) を取得: -R^T * t
        std::array<double, 3> position() const {
            double x = -(R(0,0)*t[0] + R(1,0)*t[1] + R(2,0)*t[2]);
            double y = -(R(0,1)*t[0] + R(1,1)*t[1] + R(2,1)*t[2]);
            double z = -(R(0,2)*t[0] + R(1,2)*t[1] + R(2,2)*t[2]);
            return {x, y, z};
        }
    };

    /// Visual Odometry のステート
    class VisualOdometry {
    public:
        explicit VisualOdometry(const PinholeCamera& cam)
            : camera_(cam), currentPose_() {}

        /// フレームを処理して姿勢を更新
        /// @param frame  入力画像 (1ch グレースケール)
        /// @return       現在の姿勢
        template<typename T>
        Pose processFrame(const image::Image<T>& frame) {
            // ORB 特徴点検出
            auto [kps, descs] = orbDetectAndCompute(frame, 500, 4,
                std::is_floating_point_v<T> ? T{0.05} : T{15});

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

                if (good.size() >= 8) {
                    // 正規化座標に変換
                    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});
                    }

                    // Essential 行列推定
                    auto [E, inliers] = findEssentialRansac(normMatches, 300, 0.005);

                    if (inliers.size() >= 8) {
                        // R, t を分解
                        auto candidates = decomposeEssential(E);

                        // インライアのみで正しいポーズを選択
                        std::vector<NormalizedMatch> inlierMatches;
                        for (int idx : inliers) inlierMatches.push_back(normMatches[idx]);

                        auto [R_rel, t_rel] = selectCorrectPose(candidates, inlierMatches);

                        // 累積姿勢の更新: T_new = T_rel * T_prev
                        calx::Vector<double> newT(3);
                        for (int i = 0; i < 3; ++i) {
                            newT[i] = 0;
                            for (int j = 0; j < 3; ++j)
                                newT[i] += R_rel(i, j) * currentPose_.t[j];
                            newT[i] += t_rel[i];
                        }

                        calx::Matrix<double> newR = R_rel * currentPose_.R;
                        currentPose_.R = newR;
                        currentPose_.t = newT;
                    }
                }
            }

            // 次フレーム用に保存
            prevKps_ = kps;
            prevDescs_ = descs;
            trajectory_.push_back(currentPose_.position());

            return currentPose_;
        }

        /// 軌跡を取得
        const std::vector<std::array<double, 3>>& trajectory() const { return trajectory_; }

        /// 軌跡を BMP 画像に描画 (XZ 平面の俯瞰)
        image::Image<float> drawTrajectory(int imgSize = 600, double scale = 50.0) const {
            image::Image<float> canvas(imgSize, imgSize, 3, 0.1f);

            int cx = imgSize / 2;
            int 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}, 2);
            }

            // 開始点 (赤) と現在位置 (青)
            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, 5, {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, 5, {0.0f, 0.0f, 1.0f}, 2);
            }

            return canvas;
        }

    private:
        PinholeCamera camera_;
        Pose currentPose_;
        std::vector<KeyPoint> prevKps_;
        std::vector<OrbDescriptor> prevDescs_;
        std::vector<std::array<double, 3>> trajectory_;
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_VISUAL_ODOMETRY_HPP
