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

// pnp.hpp — Perspective-n-Point (PnP)
//
// 新規実装:
//   - DLT ベースの PnP (6 点以上)
//   - RANSAC との組み合わせ

#ifndef CALX_IMAGE_CV_PNP_HPP
#define CALX_IMAGE_CV_PNP_HPP

#include <math/core/matrix.hpp>
#include <math/linalg/decomposition.hpp>
#include <image/cv/ransac.hpp>
#include <image/cv/camera.hpp>
#include <array>
#include <cmath>
#include <vector>

namespace calx::cv {

    /// 3D-2D 対応
    struct Point3D2D {
        double X, Y, Z;    // 3D ワールド座標
        double u, v;        // 2D 画像正規化座標
    };

    /// DLT ベースの PnP (6 点以上)
    /// @param points  3D-2D 対応点
    /// @return        (R, t) — カメラ姿勢
    inline std::pair<calx::Matrix<double>, calx::Vector<double>>
    solvePnPDLT(const std::vector<Point3D2D>& points) {
        std::size_t n = points.size();

        // 2n × 12 の DLT 行列
        calx::Matrix<double> A(2 * n, 12, 0.0);
        for (std::size_t i = 0; i < n; ++i) {
            double X = points[i].X, Y = points[i].Y, Z = points[i].Z;
            double u = points[i].u, v = points[i].v;

            // [X Y Z 1 0 0 0 0 -uX -uY -uZ -u]
            A(2*i, 0) = X; A(2*i, 1) = Y; A(2*i, 2) = Z; A(2*i, 3) = 1.0;
            A(2*i, 8) = -u*X; A(2*i, 9) = -u*Y; A(2*i, 10) = -u*Z; A(2*i, 11) = -u;

            // [0 0 0 0 X Y Z 1 -vX -vY -vZ -v]
            A(2*i+1, 4) = X; A(2*i+1, 5) = Y; A(2*i+1, 6) = Z; A(2*i+1, 7) = 1.0;
            A(2*i+1, 8) = -v*X; A(2*i+1, 9) = -v*Y; A(2*i+1, 10) = -v*Z; A(2*i+1, 11) = -v;
        }

        auto [U, S, Vt] = calx::algorithms::svd_decomposition(A);

        // 最後の行から P (3x4) を復元
        calx::Matrix<double> P(3, 4);
        for (int i = 0; i < 12; ++i) {
            P(i / 4, i % 4) = Vt(11, i);
        }

        // P = [R | t] を分離
        calx::Matrix<double> M(3, 3);
        calx::Vector<double> t(3);
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) M(i, j) = P(i, j);
            t[i] = P(i, 3);
        }

        // M を最寄りの回転行列に射影
        auto [Ur, Sr, Vtr] = calx::algorithms::svd_decomposition(M);
        calx::Matrix<double> R = Ur * Vtr;

        // det(R) = +1 を保証
        double det = R(0,0)*(R(1,1)*R(2,2)-R(1,2)*R(2,1))
                   - R(0,1)*(R(1,0)*R(2,2)-R(1,2)*R(2,0))
                   + R(0,2)*(R(1,0)*R(2,1)-R(1,1)*R(2,0));

        double scale = (Sr[0] + Sr[1] + Sr[2]) / 3.0;
        if (scale > 1e-15) {
            t[0] /= scale; t[1] /= scale; t[2] /= scale;
        }

        if (det < 0) {
            for (int i = 0; i < 3; ++i) {
                for (int j = 0; j < 3; ++j) R(i, j) = -R(i, j);
                t[i] = -t[i];
            }
        }

        return {R, t};
    }

    /// RANSAC ベースの PnP
    inline std::pair<calx::Matrix<double>, calx::Vector<double>>
    solvePnPRansac(const std::vector<Point3D2D>& points,
                   int maxIter = 200,
                   double threshold = 5.0) {
        using Pose = std::pair<calx::Matrix<double>, calx::Vector<double>>;

        auto fitFn = [](const std::vector<Point3D2D>& sample) -> Pose {
            return solvePnPDLT(sample);
        };

        auto errorFn = [](const Pose& pose, const Point3D2D& pt) -> double {
            const auto& [R, t] = pose;
            // 3D → カメラ座標
            double cx = R(0,0)*pt.X + R(0,1)*pt.Y + R(0,2)*pt.Z + t[0];
            double cy = R(1,0)*pt.X + R(1,1)*pt.Y + R(1,2)*pt.Z + t[1];
            double cz = R(2,0)*pt.X + R(2,1)*pt.Y + R(2,2)*pt.Z + t[2];
            if (cz < 1e-10) return 1e10;
            double proj_u = cx / cz;
            double proj_v = cy / cz;
            double dx = proj_u - pt.u;
            double dy = proj_v - pt.v;
            return dx*dx + dy*dy;
        };

        auto result = ransac<Pose, Point3D2D>(
            points, 6, maxIter, threshold * threshold, fitFn, errorFn);

        return result.bestModel;
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_PNP_HPP
