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

// BlockLanczos.hpp
// GF(2) 上の線形代数: Gauss 消去 + Block Lanczos
// SIQS / GNFS の共通インフラ
//
// 参考文献:
//   P.L. Montgomery, "A Block Lanczos Algorithm for Finding Dependencies
//   over GF(2)", Eurocrypt 1995, LNCS 921, pp.106-120

#pragma once

#include <cstdint>
#include <cstddef>
#include <cstring>
#include <vector>
#include <cassert>
#include <random>
#include <math/core/mp/Int/compiler_intrinsics.hpp>

namespace calx {
namespace gf2_linalg {

// ================================================================
// GF(2) 疎行列: CSR (Compressed Sparse Row) 形式
// ================================================================
// Block Lanczos は CSR 行列を使用して行列×ベクトル積を高速化する。
// 現在は将来の Block Lanczos 実装のための基盤として提供。

struct SparseMatrixGF2 {
    int nrows;
    int ncols;
    std::vector<int> row_ptr;  // row_ptr[i] = 行 i の開始インデックス (size = nrows+1)
    std::vector<int> col_idx;  // col_idx[k] = k 番目の非零要素の列番号

    SparseMatrixGF2() : nrows(0), ncols(0) {}

    // dense 行列 (vector<vector<uint8_t>>) から構築
    void buildFromDense(const std::vector<std::vector<uint8_t>>& mat, int nr, int nc) {
        nrows = nr;
        ncols = nc;
        row_ptr.resize(nr + 1);
        col_idx.clear();
        row_ptr[0] = 0;
        for (int i = 0; i < nr; ++i) {
            for (int j = 0; j < nc; ++j) {
                if (mat[i][j]) col_idx.push_back(j);
            }
            row_ptr[i + 1] = static_cast<int>(col_idx.size());
        }
    }

    // y = A * x (64 ベクトル同時)
    void mulVec(const uint64_t* x, uint64_t* y) const {
        std::memset(y, 0, nrows * sizeof(uint64_t));
        for (int i = 0; i < nrows; ++i) {
            uint64_t acc = 0;
            for (int k = row_ptr[i]; k < row_ptr[i + 1]; ++k)
                acc ^= x[col_idx[k]];
            y[i] = acc;
        }
    }

    // y = A^T * x (64 ベクトル同時)
    void mulVecTranspose(const uint64_t* x, uint64_t* y) const {
        std::memset(y, 0, ncols * sizeof(uint64_t));
        for (int i = 0; i < nrows; ++i) {
            uint64_t xi = x[i];
            if (xi == 0) continue;
            for (int k = row_ptr[i]; k < row_ptr[i + 1]; ++k)
                y[col_idx[k]] ^= xi;
        }
    }

    // 非零要素数
    int nnz() const { return static_cast<int>(col_idx.size()); }
};

// ================================================================
// 64×64 GF(2) 行列演算 (Block Lanczos のブロック演算用)
// ================================================================

struct Mat64 {
    uint64_t m[64];

    Mat64() { std::memset(m, 0, sizeof(m)); }

    static Mat64 identity() {
        Mat64 r;
        for (int i = 0; i < 64; ++i) r.m[i] = 1ULL << i;
        return r;
    }

    // C = A * B
    static Mat64 mul(const Mat64& a, const Mat64& b) {
        Mat64 c;
        for (int i = 0; i < 64; ++i) {
            uint64_t row = a.m[i], result = 0;
            while (row) {
                result ^= b.m[_tzcnt_u64(row)];
                row &= row - 1;
            }
            c.m[i] = result;
        }
        return c;
    }

    // B = A^T
    static Mat64 transpose(const Mat64& a) {
        Mat64 r;
        for (int i = 0; i < 64; ++i) {
            uint64_t col = 0;
            for (int j = 0; j < 64; ++j)
                if (a.m[j] & (1ULL << i)) col |= 1ULL << j;
            r.m[i] = col;
        }
        return r;
    }

    // 逆行列 (存在する場合)。成功なら true。
    static bool invert(const Mat64& a, Mat64& result) {
        Mat64 work;
        std::memcpy(work.m, a.m, sizeof(work.m));
        result = identity();
        for (int col = 0; col < 64; ++col) {
            uint64_t bit = 1ULL << col;
            int pivot = -1;
            for (int row = col; row < 64; ++row)
                if (work.m[row] & bit) { pivot = row; break; }
            if (pivot < 0) return false;
            if (pivot != col) {
                std::swap(work.m[col], work.m[pivot]);
                std::swap(result.m[col], result.m[pivot]);
            }
            for (int row = 0; row < 64; ++row) {
                if (row != col && (work.m[row] & bit)) {
                    work.m[row] ^= work.m[col];
                    result.m[row] ^= result.m[col];
                }
            }
        }
        return true;
    }
};

// v^T * w → 64×64 行列
inline Mat64 dotProduct(const uint64_t* v, const uint64_t* w, int n) {
    Mat64 result;
    for (int i = 0; i < n; ++i) {
        uint64_t vi = v[i], wi = w[i];
        while (vi) {
            result.m[_tzcnt_u64(vi)] ^= wi;
            vi &= vi - 1;
        }
    }
    return result;
}

// ================================================================
// Block Lanczos ヘルパー関数
// ================================================================

// 64-bit ワードを 64×64 行列の行として右から乗算
// v のビット k が立っている → M.m[k] を XOR
inline uint64_t applyRow(uint64_t v, const Mat64& M) {
    uint64_t result = 0;
    while (v) {
        result ^= M.m[_tzcnt_u64(v)];
        v &= v - 1;
    }
    return result;
}

// D の部分空間 S (列マスク) と S 上の擬似逆行列を計算
// D = v_i^T * B * v_i (64×64)
// S = D の線形独立な列のビットマスク
// D_inv = S に制限した D の逆行列 (S 外は 0)
inline void computeSubspace(const Mat64& D, uint64_t& S, Mat64& D_inv) {
    Mat64 work;
    std::memcpy(work.m, D.m, sizeof(work.m));
    Mat64 inv = Mat64::identity();
    S = 0;

    for (int col = 0; col < 64; ++col) {
        uint64_t bit = 1ULL << col;
        int pivot = -1;
        for (int row = col; row < 64; ++row)
            if (work.m[row] & bit) { pivot = row; break; }
        if (pivot < 0) continue;

        S |= bit;
        if (pivot != col) {
            std::swap(work.m[col], work.m[pivot]);
            std::swap(inv.m[col], inv.m[pivot]);
        }
        for (int row = 0; row < 64; ++row) {
            if (row != col && (work.m[row] & bit)) {
                work.m[row] ^= work.m[col];
                inv.m[row] ^= inv.m[col];
            }
        }
    }

    // S 外の行・列をゼロクリア
    D_inv = inv;
    for (int k = 0; k < 64; ++k) {
        if (!(S & (1ULL << k))) {
            D_inv.m[k] = 0;
        }
        // S 外の列をクリア
        D_inv.m[k] &= S;
    }
}

// ================================================================
// Block Lanczos 反復 (Montgomery 1995, Algorithm 2)
// ================================================================
// GF(2) 上の nrows×ncols 行列 A に対して、
// B = A·A^T (nrows×nrows 対称行列) の核を求める。
// nrows > ncols (SIQS/GNFS の典型) の場合、rank(M) ≤ ncols なので
// dim ker(M M^T) ≥ nrows - ncols。
// M が列フルランクなら ker(M M^T) = ker(M^T)。
// 結果は行インデックスのリスト (= gaussianEliminationGF2 と同じ形式)。
//
// ブロック幅 64 (マシンワード): 64 本のベクトルを同時に処理。
// 計算量: O(w · N / 64 · iter) where w=非零要素数, iter≈N/64

inline std::vector<std::vector<int>> blockLanczosGF2(
    const SparseMatrixGF2& A, uint64_t seed = 42)
{
    int nrows = A.nrows;
    int ncols = A.ncols;
    int N = nrows;  // B = A·A^T は nrows×nrows 対称行列

    // バッファ確保
    std::vector<uint64_t> v(N, 0), v_prev(N, 0);
    std::vector<uint64_t> Bv(N, 0), Bv_prev(N, 0), tmp(ncols, 0);
    std::vector<uint64_t> vnext(N, 0);
    std::vector<uint64_t> y(N, 0);
    std::vector<uint64_t> x_sol(N, 0);

    // ランダム初期ベクトル x_0 (nrows 次元)
    std::vector<uint64_t> x0(N, 0);
    std::mt19937_64 rng(seed);
    for (int j = 0; j < N; ++j)
        x0[j] = rng();

    // y = B * x_0 = A * (A^T * x_0)
    A.mulVecTranspose(x0.data(), tmp.data());
    A.mulVec(tmp.data(), y.data());

    // v_0 = y
    std::memcpy(v.data(), y.data(), N * sizeof(uint64_t));

    Mat64 D_inv_prev;
    uint64_t S_prev = 0;
    int max_iter = N / 64 + 128;

    for (int iter = 0; iter < max_iter; ++iter) {
        // 1. Bv = B * v = A * (A^T * v)
        A.mulVecTranspose(v.data(), tmp.data());
        A.mulVec(tmp.data(), Bv.data());

        // 2. D_i = v^T * Bv
        Mat64 D = dotProduct(v.data(), Bv.data(), N);

        // 3. 部分空間 S_i と D_i^{-1}
        uint64_t S;
        Mat64 D_inv;
        computeSubspace(D, S, D_inv);

        // 4. 収束判定: S が空 → 終了
        if (S == 0) break;

        // 5. 解の蓄積: x_sol ^= v * (D_inv * (v^T * y))
        Mat64 vTy = dotProduct(v.data(), y.data(), N);
        Mat64 coeff = Mat64::mul(D_inv, vTy);
        for (int j = 0; j < N; ++j)
            x_sol[j] ^= applyRow(v[j], coeff);

        // 6. 三項再帰の係数 (Montgomery 1995)
        // E_i = D_i^{-1} · V_i^T · B² · V_i  (B-orthogonality with v_i)
        Mat64 E = Mat64::mul(D_inv, dotProduct(Bv.data(), Bv.data(), N));

        // F_i = D_{i-1}^{-1} · V_{i-1}^T · B² · V_i  (B-orthogonality with v_{i-1})
        // V_{i-1}^T · B² · V_i = (B·V_{i-1})^T · (B·V_i)
        Mat64 F;
        if (iter > 0) {
            F = Mat64::mul(D_inv_prev, dotProduct(Bv_prev.data(), Bv.data(), N));
        }

        // 7. 三項再帰: v_{i+1} = S_i · (Bv ⊕ v·E ⊕ v_prev·F)
        // S_i は全体に適用 (Montgomery Algorithm 2 step vi)
        for (int j = 0; j < N; ++j) {
            uint64_t val = Bv[j] ^ applyRow(v[j], E);
            if (iter > 0) val ^= applyRow(v_prev[j], F);
            vnext[j] = val & S;
        }

        // 8. シフト (Bv_prev も保存)
        std::swap(v_prev, v);
        std::swap(v, vnext);
        std::swap(Bv_prev, Bv);

        D_inv_prev = D_inv;
        S_prev = S;
    }

    // x = x_0 ⊕ Σ v_i D_i^{-1} v_i^T y  (Montgomery's solution)
    // x ∈ ker(B) = ker(A·A^T) ⊃ ker(A^T)
    for (int j = 0; j < N; ++j)
        x_sol[j] ^= x0[j];

    // 検証: A^T · x_sol = 0 なら x_sol は行の線形従属を表す
    std::vector<uint64_t> ATx(ncols, 0);
    A.mulVecTranspose(x_sol.data(), ATx.data());

    uint64_t nonzero_bits = 0;
    for (int i = 0; i < ncols; ++i)
        nonzero_bits |= ATx[i];

    // 有効なビット b について、x_sol[j] のビット b が立っている行 j の集合が
    // 行従属関係を構成する (= Gauss の出力と同形式)
    std::vector<std::vector<int>> results;

    for (int b = 0; b < 64; ++b) {
        if (nonzero_bits & (1ULL << b)) continue;

        std::vector<int> row_indices;
        for (int j = 0; j < N; ++j)
            if (x_sol[j] & (1ULL << b)) row_indices.push_back(j);

        if (row_indices.size() >= 2)
            results.push_back(std::move(row_indices));
    }

    return results;
}

// ================================================================
// Gauss 消去 (64-bit ワードパッキング)
// ================================================================
// SIQS / GNFS 共通。既存の gaussianEliminationGF2 / gnfsGaussGF2 を統合。

inline std::vector<std::vector<int>> gaussianEliminationGF2(
    std::vector<std::vector<uint8_t>>& matrix,
    int nrows, int ncols)
{
    int col_words = (ncols + 63) / 64;
    int hist_words = (nrows + 63) / 64;
    std::vector<std::vector<uint64_t>> packed(nrows, std::vector<uint64_t>(col_words, 0));
    std::vector<std::vector<uint64_t>> history(nrows, std::vector<uint64_t>(hist_words, 0));
    for (int i = 0; i < nrows; ++i) {
        for (int c = 0; c < ncols; ++c)
            if (matrix[i][c])
                packed[i][c / 64] |= 1ULL << (c % 64);
        history[i][i / 64] = 1ULL << (i % 64);
    }

    int pivot_row = 0;
    for (int col = 0; col < ncols && pivot_row < nrows; ++col) {
        int wi = col / 64;
        uint64_t bm = 1ULL << (col % 64);
        int found = -1;
        for (int row = pivot_row; row < nrows; ++row)
            if (packed[row][wi] & bm) { found = row; break; }
        if (found < 0) continue;
        if (found != pivot_row) {
            std::swap(packed[found], packed[pivot_row]);
            std::swap(history[found], history[pivot_row]);
        }
        for (int row = 0; row < nrows; ++row) {
            if (row == pivot_row || !(packed[row][wi] & bm)) continue;
            for (int w = 0; w < col_words; ++w) packed[row][w] ^= packed[pivot_row][w];
            for (int w = 0; w < hist_words; ++w) history[row][w] ^= history[pivot_row][w];
        }
        ++pivot_row;
    }

    std::vector<std::vector<int>> ns;
    for (int row = 0; row < nrows; ++row) {
        bool zero = true;
        for (int w = 0; w < col_words; ++w) if (packed[row][w]) { zero = false; break; }
        if (zero) {
            std::vector<int> indices;
            for (int w = 0; w < hist_words; ++w) {
                uint64_t bits = history[row][w];
                while (bits) {
                    unsigned long idx;
                    _BitScanForward64(&idx, bits);
                    indices.push_back(w * 64 + static_cast<int>(idx));
                    bits &= bits - 1;
                }
            }
            if (!indices.empty()) ns.push_back(indices);
        }
    }
    return ns;
}

// ================================================================
// 統合インタフェース
// ================================================================
// 関係数に応じて Gauss 消去 / Block Lanczos をディスパッチ。

inline std::vector<std::vector<int>> findNullSpaceGF2(
    std::vector<std::vector<uint8_t>>& matrix,
    int nrows, int ncols)
{
    // Block Lanczos は疎行列に強く、大きな行列で Gauss O(N³/64) より高速
    constexpr int BLOCK_LANCZOS_THRESHOLD = 400;
    if (nrows >= BLOCK_LANCZOS_THRESHOLD && ncols >= BLOCK_LANCZOS_THRESHOLD) {
        SparseMatrixGF2 A;
        A.buildFromDense(matrix, nrows, ncols);
        // 複数シードでリトライ
        // ★ 最低 10 ベクトル必要。少なすぎたら Gauss にフォールバック
        constexpr int MIN_LANCZOS_VECTORS = 10;
        for (uint64_t seed = 42; seed < 42 + 5; ++seed) {
            auto result = blockLanczosGF2(A, seed);
            if ((int)result.size() >= MIN_LANCZOS_VECTORS) return result;
        }
        // フォールバック: Block Lanczos が十分な解を見つけられなかった場合
    }
    return gaussianEliminationGF2(matrix, nrows, ncols);
}

} // namespace gf2_linalg
} // namespace calx