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

// RandomizedSVD.hpp — Randomized SVD
//
// Halko-Martinsson-Tropp (2011) algorithm.
// Approximates the top k singular values/vectors of a large matrix in O(mnk).
//
// Reference: "Finding Structure with Randomness: Probabilistic Algorithms
//        for Constructing Approximate Matrix Decompositions"

#ifndef SANGI_MATH_LINALG_RANDOMIZED_SVD_HPP
#define SANGI_MATH_LINALG_RANDOMIZED_SVD_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <random>
#include <vector>

namespace sangi {

    template<typename T>
    struct RandomizedSVDResult {
        Matrix<T> U;          // (m × k) left singular vectors
        std::vector<T> S;     // (k) singular values
        Matrix<T> Vt;         // (k × n) right singular vectors (transposed)
    };

    /// Randomized SVD
    /// @param A           input matrix (m × n)
    /// @param k           number of singular values to compute
    /// @param oversampling oversampling count (improves accuracy, usually 5–10)
    /// @param powerIter   number of power iterations (useful when spectral decay is slow)
    /// @param seed        random seed
    template<typename T>
    RandomizedSVDResult<T> randomizedSVD(const BaseMatrix<T>& A, size_t k,
                                          size_t oversampling = 5,
                                          size_t powerIter = 1,
                                          unsigned seed = 42)
    {
        size_t m = A.rows(), n = A.cols();
        size_t l = std::min(k + oversampling, std::min(m, n));

        std::mt19937 rng(seed);
        std::normal_distribution<T> normal(T{0}, T{1});

        // Step 1: random projection Ω (n × l)
        Matrix<T> Omega(n, l);
        for (size_t i = 0; i < n; ++i)
            for (size_t j = 0; j < l; ++j)
                Omega(i, j) = normal(rng);

        // Step 2: Y = A * Ω  (m × l)
        Matrix<T> Y(m, l, T{0});
        for (size_t i = 0; i < m; ++i)
            for (size_t j = 0; j < l; ++j)
                for (size_t p = 0; p < n; ++p)
                    Y(i, j) += A(i, p) * Omega(p, j);

        // Step 3: power iteration (improves spectral decay)
        for (size_t q = 0; q < powerIter; ++q) {
            // Z = A^T * Y  (n × l)
            Matrix<T> Z(n, l, T{0});
            for (size_t i = 0; i < n; ++i)
                for (size_t j = 0; j < l; ++j)
                    for (size_t p = 0; p < m; ++p)
                        Z(i, j) += A(p, i) * Y(p, j);

            // Y = A * Z  (m × l)
            Y = Matrix<T>(m, l, T{0});
            for (size_t i = 0; i < m; ++i)
                for (size_t j = 0; j < l; ++j)
                    for (size_t p = 0; p < n; ++p)
                        Y(i, j) += A(i, p) * Z(p, j);
        }

        // Step 4: QR decomposition (Y = Q * R) — modified Gram-Schmidt
        Matrix<T> Q(m, l, T{0});
        for (size_t j = 0; j < l; ++j) {
            for (size_t i = 0; i < m; ++i) Q(i, j) = Y(i, j);

            for (size_t p = 0; p < j; ++p) {
                T dot = T{0};
                for (size_t i = 0; i < m; ++i) dot += Q(i, p) * Q(i, j);
                for (size_t i = 0; i < m; ++i) Q(i, j) -= dot * Q(i, p);
            }

            T norm = T{0};
            for (size_t i = 0; i < m; ++i) norm += Q(i, j) * Q(i, j);
            norm = std::sqrt(norm);
            if (norm > T{1e-15})
                for (size_t i = 0; i < m; ++i) Q(i, j) /= norm;
        }

        // Step 5: B = Q^T * A  (l × n)
        Matrix<T> B(l, n, T{0});
        for (size_t i = 0; i < l; ++i)
            for (size_t j = 0; j < n; ++j)
                for (size_t p = 0; p < m; ++p)
                    B(i, j) += Q(p, i) * A(p, j);

        // Step 6: full SVD of B (small: l × n, l << m)
        // Extract the top k via power method
        Matrix<T> Ub(l, k, T{0});
        Matrix<T> Vtb(k, n, T{0});
        std::vector<T> Sb(k, T{0});

        Matrix<T> Bwork = B;
        for (size_t r = 0; r < k; ++r) {
            std::vector<T> v(n);
            for (size_t j = 0; j < n; ++j) v[j] = normal(rng);

            for (int iter = 0; iter < 100; ++iter) {
                // u = B * v
                std::vector<T> u(l, T{0});
                for (size_t i = 0; i < l; ++i)
                    for (size_t j = 0; j < n; ++j)
                        u[i] += Bwork(i, j) * v[j];
                T nu = T{0};
                for (auto x : u) nu += x * x;
                nu = std::sqrt(nu);
                if (nu < T{1e-15}) break;
                for (auto& x : u) x /= nu;

                // v = B^T * u
                std::fill(v.begin(), v.end(), T{0});
                for (size_t j = 0; j < n; ++j)
                    for (size_t i = 0; i < l; ++i)
                        v[j] += Bwork(i, j) * u[i];
                T nv = T{0};
                for (auto x : v) nv += x * x;
                Sb[r] = std::sqrt(nv);
                if (Sb[r] < T{1e-15}) break;
                for (auto& x : v) x /= Sb[r];
            }

            // u = B * v / sigma
            std::vector<T> u(l, T{0});
            for (size_t i = 0; i < l; ++i)
                for (size_t j = 0; j < n; ++j)
                    u[i] += Bwork(i, j) * v[j];
            if (Sb[r] > T{1e-15})
                for (auto& x : u) x /= Sb[r];

            for (size_t i = 0; i < l; ++i) Ub(i, r) = u[i];
            for (size_t j = 0; j < n; ++j) Vtb(r, j) = v[j];

            // Deflation
            for (size_t i = 0; i < l; ++i)
                for (size_t j = 0; j < n; ++j)
                    Bwork(i, j) -= Sb[r] * u[i] * v[j];
        }

        // Step 7: U = Q * Ub  (m × k)
        Matrix<T> U(m, k, T{0});
        for (size_t i = 0; i < m; ++i)
            for (size_t r = 0; r < k; ++r)
                for (size_t p = 0; p < l; ++p)
                    U(i, r) += Q(i, p) * Ub(p, r);

        return { std::move(U), std::move(Sb), std::move(Vtb) };
    }

} // namespace sangi

#endif // SANGI_MATH_LINALG_RANDOMIZED_SVD_HPP