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

// camera.hpp — ピンホールカメラモデル + レンズ歪み補正
//
// 新規実装:
//   - 内部パラメータ行列 K (fx, fy, cx, cy)
//   - 放射歪み (k1, k2, k3) + 接線歪み (p1, p2)
//   - undistort (歪み除去), project (3D→2D), unproject (2D→3D ray)

#ifndef CALX_IMAGE_CV_CAMERA_HPP
#define CALX_IMAGE_CV_CAMERA_HPP

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

namespace calx::cv {

    /// レンズ歪み係数
    struct DistortionCoeffs {
        double k1 = 0, k2 = 0, k3 = 0;  // 放射歪み
        double p1 = 0, p2 = 0;           // 接線歪み
    };

    /// ピンホールカメラモデル
    class PinholeCamera {
    public:
        PinholeCamera() = default;

        /// 焦点距離と主点から構築
        PinholeCamera(double fx, double fy, double cx, double cy,
                      const DistortionCoeffs& dist = {})
            : fx_(fx), fy_(fy), cx_(cx), cy_(cy), dist_(dist) {}

        /// 内部パラメータ行列 K (3x3) から構築
        explicit PinholeCamera(const calx::Matrix<double>& K,
                               const DistortionCoeffs& dist = {})
            : fx_(K(0, 0)), fy_(K(1, 1)), cx_(K(0, 2)), cy_(K(1, 2)), dist_(dist) {}

        // ── アクセサ ──

        double fx() const { return fx_; }
        double fy() const { return fy_; }
        double cx() const { return cx_; }
        double cy() const { return cy_; }
        DistortionCoeffs& distortion() { return dist_; }
        const DistortionCoeffs& distortion() const { return dist_; }

        /// 内部パラメータ行列 K (3x3) を取得
        calx::Matrix<double> K() const {
            calx::Matrix<double> m(3, 3, 0.0);
            m(0, 0) = fx_;  m(0, 2) = cx_;
            m(1, 1) = fy_;  m(1, 2) = cy_;
            m(2, 2) = 1.0;
            return m;
        }

        /// K の逆行列 (解析的)
        calx::Matrix<double> Kinv() const {
            calx::Matrix<double> m(3, 3, 0.0);
            m(0, 0) = 1.0 / fx_;  m(0, 2) = -cx_ / fx_;
            m(1, 1) = 1.0 / fy_;  m(1, 2) = -cy_ / fy_;
            m(2, 2) = 1.0;
            return m;
        }

        // ── 射影・逆射影 ──

        /// 3D 点 (カメラ座標系) → 2D 画像座標 (歪みあり)
        /// @param pt  (X, Y, Z)
        /// @return    (u, v)
        std::array<double, 2> project(double X, double Y, double Z) const {
            double xn = X / Z;
            double yn = Y / Z;

            // 歪みを適用
            auto [xd, yd] = applyDistortion(xn, yn);

            return {fx_ * xd + cx_, fy_ * yd + cy_};
        }

        /// 2D 画像座標 → 正規化カメラ座標 (歪み除去済み)
        /// @return (xn, yn) — Z=1 平面上の座標
        std::array<double, 2> unproject(double u, double v) const {
            double xd = (u - cx_) / fx_;
            double yd = (v - cy_) / fy_;
            return removeDistortion(xd, yd);
        }

        // ── 歪み補正 ──

        /// 画像全体の歪み除去
        template<typename T>
        image::Image<T> undistort(const image::Image<T>& img) const {
            image::Image<T> result(img.height(), img.width(), img.numChannels());
            int h = static_cast<int>(img.height());
            int w = static_cast<int>(img.width());

            for (int y = 0; y < h; ++y) {
                for (int x = 0; x < w; ++x) {
                    // 歪み除去後の画像座標 → 歪みありの元画像座標
                    double xn = (x - cx_) / fx_;
                    double yn = (y - cy_) / fy_;
                    auto [xd, yd] = applyDistortion(xn, yn);
                    double srcX = fx_ * xd + cx_;
                    double srcY = fy_ * yd + cy_;

                    // バイリニア補間
                    int x0 = static_cast<int>(std::floor(srcX));
                    int y0 = static_cast<int>(std::floor(srcY));
                    if (x0 < 0 || x0 >= w - 1 || y0 < 0 || y0 >= h - 1) continue;

                    double fx = srcX - x0;
                    double fy = srcY - y0;

                    for (int ch = 0; ch < img.numChannels(); ++ch) {
                        double val =
                            (1 - fy) * (1 - fx) * img.at(y0, x0, ch) +
                            (1 - fy) * fx       * img.at(y0, x0 + 1, ch) +
                            fy       * (1 - fx) * img.at(y0 + 1, x0, ch) +
                            fy       * fx       * img.at(y0 + 1, x0 + 1, ch);
                        result.at(y, x, ch) = static_cast<T>(val);
                    }
                }
            }
            return result;
        }

    private:
        double fx_ = 1, fy_ = 1, cx_ = 0, cy_ = 0;
        DistortionCoeffs dist_;

        /// 正規化座標に歪みを適用
        std::array<double, 2> applyDistortion(double xn, double yn) const {
            double r2 = xn * xn + yn * yn;
            double r4 = r2 * r2;
            double r6 = r4 * r2;

            // 放射歪み
            double radial = 1.0 + dist_.k1 * r2 + dist_.k2 * r4 + dist_.k3 * r6;

            // 接線歪み
            double tx = 2.0 * dist_.p1 * xn * yn + dist_.p2 * (r2 + 2.0 * xn * xn);
            double ty = dist_.p1 * (r2 + 2.0 * yn * yn) + 2.0 * dist_.p2 * xn * yn;

            return {xn * radial + tx, yn * radial + ty};
        }

        /// 歪みあり座標 → 歪みなし座標 (反復法)
        std::array<double, 2> removeDistortion(double xd, double yd) const {
            double xn = xd, yn = yd;

            // Newton 法的な反復で歪みを除去
            for (int i = 0; i < 20; ++i) {
                auto [xd2, yd2] = applyDistortion(xn, yn);
                double dx = xd2 - xd;
                double dy = yd2 - yd;
                xn -= dx;
                yn -= dy;
                if (dx * dx + dy * dy < 1e-14) break;
            }

            return {xn, yn};
        }
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_CAMERA_HPP
