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

/**
 * @file sparse_qr.hpp
 * @brief 疎行列用 QR 分解 (SparseQR)
 *
 * Householder QR を疎行列に適用。長方行列 (m >= n) に対応。
 * 最小二乗問題 min ||Ax - b|| の解法に使用。
 *
 * API は Eigen の SparseQR に準拠:
 *   SparseQR<T> qr;
 *   qr.compute(A);
 *   Vector<T> x = qr.solve(b);
 */

#ifndef CALX_SPARSE_QR_HPP
#define CALX_SPARSE_QR_HPP

#include <vector>
#include <algorithm>
#include <cmath>
#include <numeric>
#include <stdexcept>
#include <limits>

#include "../core/sparse_matrix.hpp"
#include "../core/vector.hpp"
#include "../core/matrix.hpp"

namespace calx {

/**
 * @brief 疎行列 QR 分解
 *
 * A = Q * R (Q: m×m 直交, R: m×n 上三角)
 * 内部的には Q を暗黙的に保持 (Householder ベクトル + tau)。
 */
template<typename T, typename IndexType = std::size_t>
class SparseQR {
public:
    SparseQR() = default;

    /// 疎行列を QR 分解する
    void compute(const SparseMatrix<T, IndexType>& A) {
        m_ = static_cast<std::size_t>(A.rows());
        n_ = static_cast<std::size_t>(A.cols());
        if (m_ < n_)
            throw std::invalid_argument("SparseQR: requires m >= n (overdetermined)");

        computed_ = false;

        // 密形式に変換して Householder QR を適用
        R_.resize(m_ * n_);
        for (std::size_t i = 0; i < m_; ++i)
            for (std::size_t j = 0; j < n_; ++j)
                R_[i * n_ + j] = A.coeff(static_cast<IndexType>(i),
                                          static_cast<IndexType>(j));

        tau_.resize(n_);
        householder_vecs_.resize(m_ * n_, T(0));

        for (std::size_t k = 0; k < n_; ++k) {
            // 列 k の k:m 部分のノルム
            T sigma = T(0);
            for (std::size_t i = k; i < m_; ++i)
                sigma += R_[i * n_ + k] * R_[i * n_ + k];
            sigma = std::sqrt(sigma);

            if (sigma < std::numeric_limits<T>::epsilon()) {
                tau_[k] = T(0);
                continue;
            }

            T alpha = (R_[k * n_ + k] >= T(0)) ? -sigma : sigma;
            T beta = R_[k * n_ + k] - alpha;

            // Householder ベクトルを保存
            householder_vecs_[k * n_ + k] = T(1);
            T inv_beta = T(1) / beta;
            for (std::size_t i = k + 1; i < m_; ++i)
                householder_vecs_[i * n_ + k] = R_[i * n_ + k] * inv_beta;

            tau_[k] = -beta / alpha;
            R_[k * n_ + k] = alpha;

            // R の残りの列に Householder を適用
            for (std::size_t j = k + 1; j < n_; ++j) {
                T d = R_[k * n_ + j];
                for (std::size_t i = k + 1; i < m_; ++i)
                    d += householder_vecs_[i * n_ + k] * R_[i * n_ + j];
                d *= tau_[k];
                R_[k * n_ + j] -= d;
                for (std::size_t i = k + 1; i < m_; ++i)
                    R_[i * n_ + j] -= householder_vecs_[i * n_ + k] * d;
            }

            // R の下三角をクリア
            for (std::size_t i = k + 1; i < m_; ++i)
                R_[i * n_ + k] = T(0);
        }

        computed_ = true;
    }

    /// 最小二乗問題 min ||Ax - b|| を解く
    [[nodiscard]] Vector<T> solve(const Vector<T>& b) const {
        if (!computed_)
            throw std::runtime_error("SparseQR: compute() not called");
        if (b.size() != m_)
            throw std::invalid_argument("SparseQR::solve: dimension mismatch");

        // Q^T * b を計算 (Householder ベクトルを順に適用)
        Vector<T> qtb = b;
        for (std::size_t k = 0; k < n_; ++k) {
            if (std::abs(tau_[k]) < std::numeric_limits<T>::epsilon()) continue;
            T d = qtb[k];
            for (std::size_t i = k + 1; i < m_; ++i)
                d += householder_vecs_[i * n_ + k] * qtb[i];
            d *= tau_[k];
            qtb[k] -= d;
            for (std::size_t i = k + 1; i < m_; ++i)
                qtb[i] -= householder_vecs_[i * n_ + k] * d;
        }

        // R * x = (Q^T b)[:n] の後退代入
        Vector<T> x(n_);
        for (std::size_t ii = 0; ii < n_; ++ii) {
            std::size_t i = n_ - 1 - ii;
            x[i] = qtb[i];
            for (std::size_t j = i + 1; j < n_; ++j)
                x[i] -= R_[i * n_ + j] * x[j];
            if (std::abs(R_[i * n_ + i]) < std::numeric_limits<T>::epsilon())
                throw std::runtime_error("SparseQR: R is singular");
            x[i] /= R_[i * n_ + i];
        }

        return x;
    }

    /// 分解済みかどうか
    [[nodiscard]] bool computed() const { return computed_; }

    /// 行列サイズ
    [[nodiscard]] std::size_t rows() const { return m_; }
    [[nodiscard]] std::size_t cols() const { return n_; }

    /// 数値ランク
    [[nodiscard]] std::size_t rank(T threshold = T(-1)) const {
        if (!computed_) return 0;
        if (threshold < T(0)) {
            threshold = std::abs(R_[0]) *
                static_cast<T>(std::max(m_, n_)) *
                std::numeric_limits<T>::epsilon();
        }
        std::size_t r = 0;
        for (std::size_t i = 0; i < n_; ++i) {
            if (std::abs(R_[i * n_ + i]) > threshold) ++r;
            else break;
        }
        return r;
    }

private:
    std::size_t m_ = 0, n_ = 0;
    bool computed_ = false;
    std::vector<T> R_;                // m×n (上三角)
    std::vector<T> tau_;              // n (Householder 係数)
    std::vector<T> householder_vecs_; // m×n (Householder ベクトル)
};

} // namespace calx

#endif // CALX_SPARSE_QR_HPP