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

// Kinematics.hpp — Forward/Inverse kinematics via DH convention
#ifndef CALX_ROBOTICS_KINEMATICS_HPP
#define CALX_ROBOTICS_KINEMATICS_HPP

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

namespace calx::robotics {

    // Denavit-Hartenberg パラメータ
    template<typename T>
    struct DHParam {
        T a;      // リンク長
        T alpha;  // リンクひねり
        T d;      // リンクオフセット
        T theta;  // 関節角度オフセット (回転関節の場合、q に加算)
    };

    // 関節角度構成
    template<typename T>
    struct JointConfig {
        std::vector<T> angles;
    };

    // 単一 DH 変換行列 T_i = Rz(θ) · Tz(d) · Tx(a) · Rx(α)
    template<typename T>
    Matrix<T> dhTransform(const DHParam<T>& p, T q) {
        const T theta = p.theta + q;
        const T ct = std::cos(theta), st = std::sin(theta);
        const T ca = std::cos(p.alpha), sa = std::sin(p.alpha);

        Matrix<T> m(4, 4, T(0));
        m(0, 0) = ct;        m(0, 1) = -st * ca;  m(0, 2) = st * sa;   m(0, 3) = p.a * ct;
        m(1, 0) = st;        m(1, 1) = ct * ca;   m(1, 2) = -ct * sa;  m(1, 3) = p.a * st;
        m(2, 0) = T(0);      m(2, 1) = sa;        m(2, 2) = ca;        m(2, 3) = p.d;
        m(3, 0) = T(0);      m(3, 1) = T(0);      m(3, 2) = T(0);      m(3, 3) = T(1);
        return m;
    }

    // 順運動学: DH パラメータ列 + 関節角度 → エンドエフェクタの同次変換行列 (4×4)
    template<typename T>
    Matrix<T> forwardKinematics(const std::vector<DHParam<T>>& dh, const JointConfig<T>& q) {
        if (dh.size() != q.angles.size()) {
            throw std::invalid_argument("DH parameters and joint angles must have the same size");
        }

        // 単位行列 4×4
        Matrix<T> pose(4, 4, T(0));
        for (int i = 0; i < 4; ++i) pose(i, i) = T(1);

        for (size_t i = 0; i < dh.size(); ++i) {
            pose = pose * dhTransform(dh[i], q.angles[i]);
        }
        return pose;
    }

    // 幾何ヤコビアン (6×n, 上3行: 線速度, 下3行: 角速度)
    template<typename T>
    Matrix<T> jacobian(const std::vector<DHParam<T>>& dh, const JointConfig<T>& q) {
        const size_t n = dh.size();
        if (n != q.angles.size()) {
            throw std::invalid_argument("DH parameters and joint angles must have the same size");
        }

        // 各関節までの累積変換を計算
        std::vector<Matrix<T>> transforms(n + 1);
        transforms[0] = Matrix<T>(4, 4, T(0));
        for (int i = 0; i < 4; ++i) transforms[0](i, i) = T(1);
        for (size_t i = 0; i < n; ++i) {
            transforms[i + 1] = transforms[i] * dhTransform(dh[i], q.angles[i]);
        }

        // エンドエフェクタ位置
        const T px = transforms[n](0, 3);
        const T py = transforms[n](1, 3);
        const T pz = transforms[n](2, 3);

        Matrix<T> J(6, static_cast<int>(n), T(0));
        for (size_t i = 0; i < n; ++i) {
            // z 軸ベクトル (回転関節)
            const T zx = transforms[i](0, 2);
            const T zy = transforms[i](1, 2);
            const T zz = transforms[i](2, 2);

            // 原点位置
            const T ox = transforms[i](0, 3);
            const T oy = transforms[i](1, 3);
            const T oz = transforms[i](2, 3);

            // 線速度部: z_i × (p_e - o_i)
            const T dx = px - ox, dy = py - oy, dz = pz - oz;
            J(0, static_cast<int>(i)) = zy * dz - zz * dy;
            J(1, static_cast<int>(i)) = zz * dx - zx * dz;
            J(2, static_cast<int>(i)) = zx * dy - zy * dx;

            // 角速度部: z_i
            J(3, static_cast<int>(i)) = zx;
            J(4, static_cast<int>(i)) = zy;
            J(5, static_cast<int>(i)) = zz;
        }
        return J;
    }

    // 逆運動学 (ヤコビアン擬似逆行列による反復法)
    template<typename T>
    JointConfig<T> inverseKinematics(const std::vector<DHParam<T>>& dh,
                                      const Matrix<T>& targetPose,
                                      const JointConfig<T>& q0,
                                      size_t maxIter = 100, T tol = T(1e-6)) {
        const size_t n = dh.size();
        JointConfig<T> q = q0;

        for (size_t iter = 0; iter < maxIter; ++iter) {
            Matrix<T> currentPose = forwardKinematics(dh, q);

            // 位置誤差
            T ex = targetPose(0, 3) - currentPose(0, 3);
            T ey = targetPose(1, 3) - currentPose(1, 3);
            T ez = targetPose(2, 3) - currentPose(2, 3);

            // 姿勢誤差 (回転行列の差分からスキュー対称成分を抽出)
            T eRx = (targetPose(2, 1) * currentPose(2, 2) - targetPose(2, 2) * currentPose(2, 1)
                   + targetPose(1, 1) * currentPose(1, 2) - targetPose(1, 2) * currentPose(1, 1)
                   + targetPose(0, 1) * currentPose(0, 2) - targetPose(0, 2) * currentPose(0, 1)) / T(2);
            T eRy = (targetPose(0, 2) * currentPose(0, 0) - targetPose(0, 0) * currentPose(0, 2)
                   + targetPose(1, 2) * currentPose(1, 0) - targetPose(1, 0) * currentPose(1, 2)
                   + targetPose(2, 2) * currentPose(2, 0) - targetPose(2, 0) * currentPose(2, 2)) / T(2);
            T eRz = (targetPose(0, 0) * currentPose(0, 1) - targetPose(0, 1) * currentPose(0, 0)
                   + targetPose(1, 0) * currentPose(1, 1) - targetPose(1, 1) * currentPose(1, 0)
                   + targetPose(2, 0) * currentPose(2, 1) - targetPose(2, 1) * currentPose(2, 0)) / T(2);

            T errNorm = std::sqrt(ex * ex + ey * ey + ez * ez + eRx * eRx + eRy * eRy + eRz * eRz);
            if (errNorm < tol) break;

            // 誤差ベクトル (6×1)
            Matrix<T> e(6, 1, T(0));
            e(0, 0) = ex; e(1, 0) = ey; e(2, 0) = ez;
            e(3, 0) = eRx; e(4, 0) = eRy; e(5, 0) = eRz;

            // ヤコビアン J (6×n)
            Matrix<T> J = jacobian(dh, q);

            // 擬似逆行列: J^+ = Jᵀ(J Jᵀ)⁻¹ (ダンプド最小二乗)
            Matrix<T> Jt = J.transpose();
            T lambda = T(1e-4);
            Matrix<T> JJt = J * Jt;
            for (int i = 0; i < 6; ++i) JJt(i, i) += lambda;

            // JJt を解いて dq = Jᵀ (JJᵀ + λI)⁻¹ e
            // 簡易実装: 6×6 のガウス消去
            Matrix<T> augmented(6, 7, T(0));
            for (int r = 0; r < 6; ++r) {
                for (int c = 0; c < 6; ++c) augmented(r, c) = JJt(r, c);
                augmented(r, 6) = e(r, 0);
            }

            // 前進消去
            for (int col = 0; col < 6; ++col) {
                int pivotRow = col;
                T pivotVal = std::abs(augmented(col, col));
                for (int r = col + 1; r < 6; ++r) {
                    if (std::abs(augmented(r, col)) > pivotVal) {
                        pivotVal = std::abs(augmented(r, col));
                        pivotRow = r;
                    }
                }
                if (pivotRow != col) {
                    for (int c = 0; c < 7; ++c) {
                        T tmp = augmented(col, c);
                        augmented(col, c) = augmented(pivotRow, c);
                        augmented(pivotRow, c) = tmp;
                    }
                }
                for (int r = col + 1; r < 6; ++r) {
                    T factor = augmented(r, col) / augmented(col, col);
                    for (int c = col; c < 7; ++c) {
                        augmented(r, c) -= factor * augmented(col, c);
                    }
                }
            }

            // 後退代入 → y = (JJᵀ + λI)⁻¹ e
            Matrix<T> y(6, 1, T(0));
            for (int r = 5; r >= 0; --r) {
                T sum = augmented(r, 6);
                for (int c = r + 1; c < 6; ++c) sum -= augmented(r, c) * y(c, 0);
                y(r, 0) = sum / augmented(r, r);
            }

            // dq = Jᵀ y
            Matrix<T> dq = Jt * y;
            for (size_t i = 0; i < n; ++i) {
                q.angles[i] += dq(static_cast<int>(i), 0);
            }
        }
        return q;
    }

} // namespace calx::robotics

#endif // CALX_ROBOTICS_KINEMATICS_HPP
