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

/**
 * @file sparse_qr.hpp
 * @brief QR factorization for sparse matrices (SparseQR)
 *
 * Applies Householder QR to sparse matrices. Supports rectangular matrices (m >= n).
 * Used to solve the least-squares problem min ||Ax - b||.
 *
 * The API follows Eigen's SparseQR:
 *   SparseQR<T> qr;
 *   qr.compute(A);
 *   Vector<T> x = qr.solve(b);
 */

#ifndef SANGI_SPARSE_QR_HPP
#define SANGI_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 sangi {

/**
 * @brief Sparse QR factorization
 *
 * A = Q * R (Q: m×m orthogonal, R: m×n upper triangular)
 * Internally keeps Q implicitly (Householder vectors + tau).
 */
template<typename T, typename IndexType = std::size_t>
class SparseQR {
public:
    SparseQR() = default;

    /// QR-factorizes a sparse matrix
    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;

        // Convert to dense form and apply 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) {
            // Norm of the k:m part of column k
            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;

            // Save the Householder vector
            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;

            // Apply Householder to the remaining columns of R
            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;
            }

            // Clear the lower triangle of R
            for (std::size_t i = k + 1; i < m_; ++i)
                R_[i * n_ + k] = T(0);
        }

        computed_ = true;
    }

    /// Solves the least-squares problem 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");

        // Compute Q^T * b (apply the Householder vectors in order)
        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;
        }

        // Back substitution for 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;
    }

    /// Whether the factorization has been computed
    [[nodiscard]] bool computed() const { return computed_; }

    /// Matrix dimensions
    [[nodiscard]] std::size_t rows() const { return m_; }
    [[nodiscard]] std::size_t cols() const { return n_; }

    /// Numerical rank
    [[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 (upper triangular)
    std::vector<T> tau_;              // n (Householder coefficients)
    std::vector<T> householder_vecs_; // m×n (Householder vectors)
};

} // namespace sangi

#endif // SANGI_SPARSE_QR_HPP
