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

// calibration.hpp — カメラキャリブレーション (Zhang 法)
//
// 新規実装:
//   - ホモグラフィ推定 (DLT + SVD)
//   - 内部パラメータ抽出 (Zhang 法)
//   - LM 最適化で歪み係数推定
//   - チェッカーボードコーナー検出

#ifndef CALX_IMAGE_CV_CALIBRATION_HPP
#define CALX_IMAGE_CV_CALIBRATION_HPP

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

namespace calx::cv {

    /// 2D 対応点ペア
    struct PointCorrespondence {
        double srcX, srcY;  // ワールド (モデル) 座標
        double dstX, dstY;  // 画像座標
    };

    // ── ホモグラフィ推定 ──

    /// DLT (Direct Linear Transform) でホモグラフィ行列 (3x3) を推定
    /// @param points  4 点以上の 2D-2D 対応
    /// @return        3x3 ホモグラフィ行列 H (dst = H * src)
    inline calx::Matrix<double> estimateHomography(const std::vector<PointCorrespondence>& points) {
        std::size_t n = points.size();
        if (n < 4) throw std::runtime_error("estimateHomography: need at least 4 points");

        // 正規化 (DLT の精度向上)
        double srcMeanX = 0, srcMeanY = 0, dstMeanX = 0, dstMeanY = 0;
        for (const auto& p : points) {
            srcMeanX += p.srcX; srcMeanY += p.srcY;
            dstMeanX += p.dstX; dstMeanY += p.dstY;
        }
        srcMeanX /= n; srcMeanY /= n;
        dstMeanX /= n; dstMeanY /= n;

        double srcScale = 0, dstScale = 0;
        for (const auto& p : points) {
            srcScale += std::sqrt((p.srcX - srcMeanX) * (p.srcX - srcMeanX) +
                                  (p.srcY - srcMeanY) * (p.srcY - srcMeanY));
            dstScale += std::sqrt((p.dstX - dstMeanX) * (p.dstX - dstMeanX) +
                                  (p.dstY - dstMeanY) * (p.dstY - dstMeanY));
        }
        srcScale = n * std::sqrt(2.0) / srcScale;
        dstScale = n * std::sqrt(2.0) / dstScale;

        // 正規化行列
        calx::Matrix<double> Ts(3, 3, 0.0);
        Ts(0, 0) = srcScale; Ts(0, 2) = -srcScale * srcMeanX;
        Ts(1, 1) = srcScale; Ts(1, 2) = -srcScale * srcMeanY;
        Ts(2, 2) = 1.0;

        calx::Matrix<double> Td(3, 3, 0.0);
        Td(0, 0) = dstScale; Td(0, 2) = -dstScale * dstMeanX;
        Td(1, 1) = dstScale; Td(1, 2) = -dstScale * dstMeanY;
        Td(2, 2) = 1.0;

        // DLT 行列構築 (2n × 9)
        calx::Matrix<double> A(2 * n, 9, 0.0);
        for (std::size_t i = 0; i < n; ++i) {
            double sx = srcScale * (points[i].srcX - srcMeanX);
            double sy = srcScale * (points[i].srcY - srcMeanY);
            double dx = dstScale * (points[i].dstX - dstMeanX);
            double dy = dstScale * (points[i].dstY - dstMeanY);

            // [-sx, -sy, -1, 0, 0, 0, dx*sx, dx*sy, dx]
            A(2 * i, 0) = -sx; A(2 * i, 1) = -sy; A(2 * i, 2) = -1;
            A(2 * i, 6) = dx * sx; A(2 * i, 7) = dx * sy; A(2 * i, 8) = dx;

            // [0, 0, 0, -sx, -sy, -1, dy*sx, dy*sy, dy]
            A(2 * i + 1, 3) = -sx; A(2 * i + 1, 4) = -sy; A(2 * i + 1, 5) = -1;
            A(2 * i + 1, 6) = dy * sx; A(2 * i + 1, 7) = dy * sy; A(2 * i + 1, 8) = dy;
        }

        // SVD で最小特異値に対応するベクトルを取得
        auto [U, S, Vt] = calx::algorithms::svd_decomposition(A);

        // V の最後の列 = Vt の最後の行
        calx::Matrix<double> Hn(3, 3);
        for (int i = 0; i < 9; ++i) {
            Hn(i / 3, i % 3) = Vt(8, i);
        }

        // 正規化を元に戻す: H = Td^{-1} * Hn * Ts
        calx::Matrix<double> TdInv(3, 3, 0.0);
        TdInv(0, 0) = 1.0 / dstScale; TdInv(0, 2) = dstMeanX;
        TdInv(1, 1) = 1.0 / dstScale; TdInv(1, 2) = dstMeanY;
        TdInv(2, 2) = 1.0;

        calx::Matrix<double> H = TdInv * Hn * Ts;

        // スケール正規化 (H(2,2) = 1)
        double scale = H(2, 2);
        if (std::abs(scale) > 1e-15) {
            for (int i = 0; i < 3; ++i)
                for (int j = 0; j < 3; ++j)
                    H(i, j) /= scale;
        }

        return H;
    }

    // ── Zhang 法による内部パラメータ推定 ──

    /// 複数のホモグラフィから内部パラメータを推定 (Zhang 法)
    /// @param homographies  3 枚以上のホモグラフィ行列
    /// @return              PinholeCamera (歪み係数はゼロ初期化)
    inline PinholeCamera calibrateCamera(const std::vector<calx::Matrix<double>>& homographies) {
        std::size_t nViews = homographies.size();
        if (nViews < 3)
            throw std::runtime_error("calibrateCamera: need at least 3 views");

        // Zhang 法: B = K^{-T} K^{-1} は対称正定値
        // b = (B11, B12, B22, B13, B23, B33)^T
        // 各ホモグラフィから 2 つの制約 (直交制約 + 等スケール制約)

        auto vij = [](const calx::Matrix<double>& H, int i, int j) -> calx::Vector<double> {
            calx::Vector<double> v(6);
            v[0] = H(0, i) * H(0, j);
            v[1] = H(0, i) * H(1, j) + H(1, i) * H(0, j);
            v[2] = H(1, i) * H(1, j);
            v[3] = H(2, i) * H(0, j) + H(0, i) * H(2, j);
            v[4] = H(2, i) * H(1, j) + H(1, i) * H(2, j);
            v[5] = H(2, i) * H(2, j);
            return v;
        };

        calx::Matrix<double> V(2 * nViews, 6, 0.0);
        for (std::size_t k = 0; k < nViews; ++k) {
            const auto& H = homographies[k];
            auto v12 = vij(H, 0, 1);
            auto v11 = vij(H, 0, 0);
            auto v22 = vij(H, 1, 1);

            for (int j = 0; j < 6; ++j) {
                V(2 * k, j) = v12[j];
                V(2 * k + 1, j) = v11[j] - v22[j];
            }
        }

        // SVD で b を求める
        auto [U2, S2, Vt2] = calx::algorithms::svd_decomposition(V);
        calx::Vector<double> b(6);
        for (int i = 0; i < 6; ++i) b[i] = Vt2(5, i);

        // B 行列から内部パラメータを復元
        double B11 = b[0], B12 = b[1], B22 = b[2];
        double B13 = b[3], B23 = b[4], B33 = b[5];

        double cy = (B12 * B13 - B11 * B23) / (B11 * B22 - B12 * B12);
        double lambda = B33 - (B13 * B13 + cy * (B12 * B13 - B11 * B23)) / B11;
        double fx = std::sqrt(std::abs(lambda / B11));
        double fy = std::sqrt(std::abs(lambda * B11 / (B11 * B22 - B12 * B12)));
        double cx = -B13 * fx * fx / lambda;

        return PinholeCamera(fx, fy, cx, cy);
    }

    /// 外部パラメータ (R, t) をホモグラフィから取得
    /// @param H     ホモグラフィ行列 (3x3)
    /// @param cam   キャリブレーション済みカメラ
    /// @param R     出力: 回転行列 (3x3)
    /// @param t     出力: 並進ベクトル (3x1)
    inline void decomposeHomography(const calx::Matrix<double>& H,
                                    const PinholeCamera& cam,
                                    calx::Matrix<double>& R,
                                    calx::Vector<double>& t) {
        calx::Matrix<double> Kinv = cam.Kinv();
        calx::Matrix<double> M = Kinv * H;

        // r1, r2, t を列ベクトルとして取得
        double lambda1 = 0;
        for (int i = 0; i < 3; ++i) lambda1 += M(i, 0) * M(i, 0);
        lambda1 = 1.0 / std::sqrt(lambda1);

        R = calx::Matrix<double>(3, 3);
        t = calx::Vector<double>(3);

        // r1
        for (int i = 0; i < 3; ++i) R(i, 0) = M(i, 0) * lambda1;
        // r2
        for (int i = 0; i < 3; ++i) R(i, 1) = M(i, 1) * lambda1;
        // t
        for (int i = 0; i < 3; ++i) t[i] = M(i, 2) * lambda1;
        // r3 = r1 × r2
        R(0, 2) = R(1, 0) * R(2, 1) - R(2, 0) * R(1, 1);
        R(1, 2) = R(2, 0) * R(0, 1) - R(0, 0) * R(2, 1);
        R(2, 2) = R(0, 0) * R(1, 1) - R(1, 0) * R(0, 1);

        // R を最寄りの回転行列に射影 (SVD)
        auto [Ur, Sr, Vtr] = calx::algorithms::svd_decomposition(R);
        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));
        if (det < 0) {
            for (int i = 0; i < 3; ++i)
                for (int j = 0; j < 3; ++j)
                    R(i, j) = -R(i, j);
            for (int i = 0; i < 3; ++i) t[i] = -t[i];
        }
    }

    // ── チェッカーボードコーナー検出 ──

    /// チェッカーボード画像を生成 (キャリブレーション用テスト画像)
    /// @param rows       チェッカーボードの行数 (交差点)
    /// @param cols       チェッカーボードの列数 (交差点)
    /// @param squareSize 各マスのサイズ (ピクセル)
    /// @param margin     画像の余白
    /// @return           グレースケール画像
    template<typename T>
    image::Image<T> generateCheckerboard(int rows, int cols, int squareSize, int margin = 20) {
        int imgH = (rows + 1) * squareSize + 2 * margin;
        int imgW = (cols + 1) * squareSize + 2 * margin;
        image::Image<T> img(imgH, imgW);

        // 白で初期化
        T white = std::is_floating_point_v<T> ? T{1} : T{255};
        for (std::size_t y = 0; y < img.height(); ++y)
            for (std::size_t x = 0; x < img.width(); ++x)
                img(y, x) = white;

        // チェッカーパターンを描画
        for (int r = 0; r <= rows; ++r) {
            for (int c = 0; c <= cols; ++c) {
                if ((r + c) % 2 == 0) continue;
                int y0 = margin + r * squareSize;
                int x0 = margin + c * squareSize;
                for (int dy = 0; dy < squareSize; ++dy) {
                    for (int dx = 0; dx < squareSize; ++dx) {
                        img(y0 + dy, x0 + dx) = T{0};
                    }
                }
            }
        }
        return img;
    }

    /// チェッカーボードのコーナー座標 (ワールド座標) を生成
    /// @param rows       交差点の行数
    /// @param cols       交差点の列数
    /// @param squareSize ワールド座標系でのマスのサイズ
    inline std::vector<std::array<double, 2>> checkerboardWorldPoints(
        int rows, int cols, double squareSize = 1.0) {
        std::vector<std::array<double, 2>> pts;
        pts.reserve(rows * cols);
        for (int r = 0; r < rows; ++r)
            for (int c = 0; c < cols; ++c)
                pts.push_back({c * squareSize, r * squareSize});
        return pts;
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_CALIBRATION_HPP
