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

// bundle_adjustment.hpp — バンドル調整
//
// 新規実装:
//   - カメラ姿勢 + 3D 点の同時最適化
//   - Levenberg-Marquardt
//   - 再射影誤差の最小化

#ifndef CALX_IMAGE_CV_BUNDLE_ADJUSTMENT_HPP
#define CALX_IMAGE_CV_BUNDLE_ADJUSTMENT_HPP

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

namespace calx::cv {

    /// 観測データ (カメラ i から 3D 点 j への 2D 観測)
    struct Observation {
        int cameraIdx;
        int pointIdx;
        double u, v;        // 画像座標
    };

    /// バンドル調整の結果
    struct BundleAdjustmentResult {
        std::vector<SE3Pose> cameras;
        std::vector<std::array<double, 3>> points;
        double finalError;
        int iterations;
    };

    /// バンドル調整 (Levenberg-Marquardt)
    /// @param cameras       初期カメラ姿勢
    /// @param points3d      初期 3D 点
    /// @param observations  2D 観測
    /// @param fx, fy        焦点距離 (全カメラ共通と仮定)
    /// @param maxIter       最大反復回数
    /// @return              最適化結果
    inline BundleAdjustmentResult bundleAdjust(
        std::vector<SE3Pose> cameras,
        std::vector<std::array<double, 3>> points3d,
        const std::vector<Observation>& observations,
        double fx, double fy,
        int maxIter = 50) {

        double lambda = 1e-3;
        double lastError = 1e20;
        int iter = 0;

        for (; iter < maxIter; ++iter) {
            double totalError = 0;

            // 各カメラと 3D 点を個別に更新 (交互最適化)
            // Step 1: 3D 点の更新 (カメラ固定)
            for (std::size_t j = 0; j < points3d.size(); ++j) {
                double grad[3] = {0, 0, 0};
                double hess[3] = {0, 0, 0};  // 対角近似
                int obsCount = 0;

                for (const auto& obs : observations) {
                    if (obs.pointIdx != static_cast<int>(j)) continue;
                    ++obsCount;

                    const auto& cam = cameras[obs.cameraIdx];
                    double X = points3d[j][0], Y = points3d[j][1], Z = points3d[j][2];

                    // カメラ座標
                    double cx = cam.R(0,0)*X + cam.R(0,1)*Y + cam.R(0,2)*Z + cam.t[0];
                    double cy = cam.R(1,0)*X + cam.R(1,1)*Y + cam.R(1,2)*Z + cam.t[1];
                    double cz = cam.R(2,0)*X + cam.R(2,1)*Y + cam.R(2,2)*Z + cam.t[2];
                    if (cz < 0.01) continue;

                    double proj_u = fx * cx / cz;
                    double proj_v = fy * cy / cz;
                    double eu = proj_u - obs.u;
                    double ev = proj_v - obs.v;
                    totalError += eu*eu + ev*ev;

                    // ∂e/∂X の近似
                    double inv_cz = 1.0 / cz;
                    for (int k = 0; k < 3; ++k) {
                        double dcdX_u = cam.R(0, k) * inv_cz - cx * cam.R(2, k) * inv_cz * inv_cz;
                        double dcdX_v = cam.R(1, k) * inv_cz - cy * cam.R(2, k) * inv_cz * inv_cz;
                        grad[k] += fx * dcdX_u * eu + fy * dcdX_v * ev;
                        hess[k] += fx*fx * dcdX_u*dcdX_u + fy*fy * dcdX_v*dcdX_v;
                    }
                }

                if (obsCount >= 2) {
                    for (int k = 0; k < 3; ++k) {
                        double denom = hess[k] + lambda;
                        if (denom > 1e-15) {
                            points3d[j][k] -= grad[k] / denom;
                        }
                    }
                }
            }

            // Step 2: カメラ並進の更新 (3D 点固定, カメラ 0 固定)
            for (std::size_t i = 1; i < cameras.size(); ++i) {
                double grad[3] = {0, 0, 0};
                double hess[3] = {0, 0, 0};

                for (const auto& obs : observations) {
                    if (obs.cameraIdx != static_cast<int>(i)) continue;

                    const auto& cam = cameras[i];
                    double X = points3d[obs.pointIdx][0];
                    double Y = points3d[obs.pointIdx][1];
                    double Z = points3d[obs.pointIdx][2];

                    double cx = cam.R(0,0)*X + cam.R(0,1)*Y + cam.R(0,2)*Z + cam.t[0];
                    double cy = cam.R(1,0)*X + cam.R(1,1)*Y + cam.R(1,2)*Z + cam.t[1];
                    double cz = cam.R(2,0)*X + cam.R(2,1)*Y + cam.R(2,2)*Z + cam.t[2];
                    if (cz < 0.01) continue;

                    double proj_u = fx * cx / cz;
                    double proj_v = fy * cy / cz;
                    double eu = proj_u - obs.u;
                    double ev = proj_v - obs.v;

                    double inv_cz = 1.0 / cz;

                    // ∂e/∂t
                    grad[0] += fx * inv_cz * eu;
                    grad[1] += fy * inv_cz * ev;
                    grad[2] += -(fx * cx * inv_cz * inv_cz * eu + fy * cy * inv_cz * inv_cz * ev);

                    hess[0] += fx*fx * inv_cz * inv_cz;
                    hess[1] += fy*fy * inv_cz * inv_cz;
                    hess[2] += (fx*fx*cx*cx + fy*fy*cy*cy) * inv_cz*inv_cz*inv_cz*inv_cz;
                }

                for (int k = 0; k < 3; ++k) {
                    double denom = hess[k] + lambda;
                    if (denom > 1e-15) {
                        cameras[i].t[k] -= grad[k] / denom;
                    }
                }
            }

            // 収束判定
            if (totalError > lastError) {
                lambda *= 10.0;
            } else {
                lambda *= 0.3;
            }

            if (std::abs(totalError - lastError) < 1e-8 * lastError) break;
            lastError = totalError;
        }

        return {cameras, points3d, lastError, iter};
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_BUNDLE_ADJUSTMENT_HPP
