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

// LLL.hpp
//
// LLL 格子基底簡約 (Lenstra-Lenstra-Lovász Lattice Basis Reduction)
//
// 整数格子の基底を「短く、直交に近い」形に変換する。
// 暗号解析、整数関係式発見 (PSLQ)、Diophantine 近似に応用。
//
// 主な API:
//   lllReduce(basis, delta)     — LLL 簡約 (Gram-Schmidt 直交化ベース)
//   isLLLReduced(basis, delta)  — LLL 簡約条件の検証
//   shortestVector(basis)       — 最短ベクトル (簡約後の b_1)
//   findIntegerRelation(x, C)   — PSLQ 風の整数関係式発見

#ifndef CALX_LINALG_LLL_HPP
#define CALX_LINALG_LLL_HPP

#include <vector>
#include <cmath>
#include <cstdint>
#include <stdexcept>
#include <algorithm>
#include <numeric>
#include <cassert>

namespace calx {

/**
 * @brief LLL 簡約の結果
 */
struct LLLResult {
    std::vector<std::vector<double>> basis;  ///< 簡約された基底
    std::size_t swaps;                        ///< スワップ回数
};

namespace detail {

// 内積
inline double dot(const std::vector<double>& a, const std::vector<double>& b) {
    double s = 0;
    for (std::size_t i = 0; i < a.size(); ++i) s += a[i] * b[i];
    return s;
}

// ノルム²
inline double norm2(const std::vector<double>& v) {
    return dot(v, v);
}

// v -= c * u
inline void sub_scaled(std::vector<double>& v, const std::vector<double>& u, double c) {
    for (std::size_t i = 0; i < v.size(); ++i) v[i] -= c * u[i];
}

} // namespace detail

/**
 * @brief LLL 格子基底簡約
 *
 * 入力: 基底ベクトル b_0, ..., b_{n-1} ∈ R^d (行ベクトル)
 * 出力: LLL-簡約された基底
 *
 * LLL 条件:
 *   (1) |μ_{i,j}| ≤ 1/2  (サイズ簡約)
 *   (2) δ·||b*_k||² ≤ ||b*_{k+1} + μ_{k+1,k}·b*_k||²  (Lovász 条件)
 *
 * @param basis  基底ベクトルの配列 (n×d)
 * @param delta  Lovász パラメータ (1/4 < δ ≤ 1, デフォルト 3/4)
 * @return LLLResult
 */
[[nodiscard]]
inline LLLResult lllReduce(std::vector<std::vector<double>> basis, double delta = 0.75)
{
    const std::size_t n = basis.size();
    if (n == 0) return {basis, 0};
    const std::size_t d = basis[0].size();

    // Gram-Schmidt 直交化係数と直交ベクトルのノルム²
    std::vector<std::vector<double>> mu(n, std::vector<double>(n, 0.0));
    std::vector<double> B(n, 0.0);  // B[i] = ||b*_i||²
    std::vector<std::vector<double>> bstar(n, std::vector<double>(d, 0.0));

    // Gram-Schmidt を計算
    auto computeGS = [&]() {
        for (std::size_t i = 0; i < n; ++i) {
            bstar[i] = basis[i];
            for (std::size_t j = 0; j < i; ++j) {
                mu[i][j] = detail::dot(basis[i], bstar[j]) / B[j];
                detail::sub_scaled(bstar[i], bstar[j], mu[i][j]);
            }
            B[i] = detail::norm2(bstar[i]);
        }
    };

    computeGS();

    std::size_t k = 1;
    std::size_t swaps = 0;

    while (k < n) {
        // サイズ簡約: |μ_{k,j}| ≤ 1/2
        for (std::size_t j = k; j > 0; --j) {
            std::size_t jj = j - 1;
            if (std::abs(mu[k][jj]) > 0.5) {
                double r = std::round(mu[k][jj]);
                detail::sub_scaled(basis[k], basis[jj], r);
                // μ を更新
                for (std::size_t i = 0; i <= jj; ++i) {
                    mu[k][i] -= r * mu[jj][i];
                }
                mu[k][jj] -= r;
            }
        }

        // Lovász 条件チェック
        double lhs = delta * B[k - 1];
        double rhs = B[k] + mu[k][k - 1] * mu[k][k - 1] * B[k - 1];

        if (lhs > rhs) {
            // スワップ b_k と b_{k-1}
            std::swap(basis[k], basis[k - 1]);
            ++swaps;

            // Gram-Schmidt を再計算 (効率化のため部分更新)
            // 簡易版: 全体再計算
            computeGS();

            if (k > 1) --k;
        } else {
            ++k;
        }
    }

    return {std::move(basis), swaps};
}

/**
 * @brief LLL 簡約条件を検証する
 *
 * @param basis 基底ベクトル
 * @param delta Lovász パラメータ
 * @return true なら LLL 簡約済み
 */
[[nodiscard]]
inline bool isLLLReduced(const std::vector<std::vector<double>>& basis, double delta = 0.75)
{
    const std::size_t n = basis.size();
    if (n <= 1) return true;
    const std::size_t d = basis[0].size();

    // Gram-Schmidt
    std::vector<std::vector<double>> bstar(n, std::vector<double>(d));
    std::vector<std::vector<double>> mu(n, std::vector<double>(n, 0.0));
    std::vector<double> B(n);

    for (std::size_t i = 0; i < n; ++i) {
        bstar[i] = basis[i];
        for (std::size_t j = 0; j < i; ++j) {
            mu[i][j] = detail::dot(basis[i], bstar[j]) / B[j];
            detail::sub_scaled(bstar[i], bstar[j], mu[i][j]);
        }
        B[i] = detail::norm2(bstar[i]);
    }

    // (1) サイズ簡約: |μ_{i,j}| ≤ 1/2
    for (std::size_t i = 1; i < n; ++i)
        for (std::size_t j = 0; j < i; ++j)
            if (std::abs(mu[i][j]) > 0.5 + 1e-10)
                return false;

    // (2) Lovász 条件
    for (std::size_t k = 1; k < n; ++k) {
        double lhs = delta * B[k - 1];
        double rhs = B[k] + mu[k][k - 1] * mu[k][k - 1] * B[k - 1];
        if (lhs > rhs + 1e-10)
            return false;
    }

    return true;
}

/**
 * @brief 格子の最短ベクトル (簡約後の第1基底ベクトル)
 *
 * LLL 簡約後の b_0 は最短ベクトルの 2^{(n-1)/2} 近似。
 */
[[nodiscard]]
inline std::vector<double> shortestVector(const std::vector<std::vector<double>>& basis,
                                           double delta = 0.75)
{
    auto result = lllReduce(basis, delta);
    return result.basis[0];
}

/**
 * @brief 整数関係式発見 (PSLQ 風)
 *
 * 実数ベクトル x = (x_0, ..., x_{n-1}) に対して、
 * 整数ベクトル m = (m_0, ..., m_{n-1}) で Σ m_i x_i ≈ 0 となるものを探す。
 *
 * 方法: 拡大格子を LLL 簡約する。
 *   格子基底 = [[C·x_0, 1, 0, ..., 0],
 *               [C·x_1, 0, 1, ..., 0],
 *               ...
 *               [C·x_{n-1}, 0, 0, ..., 1]]
 * C は精度パラメータ (大きいほど高精度だが数値的に不安定)。
 *
 * @param x     実数ベクトル
 * @param C     精度パラメータ (デフォルト 1e10)
 * @return 整数係数ベクトル m (最短の関係式)
 */
[[nodiscard]]
inline std::vector<int64_t> findIntegerRelation(
    const std::vector<double>& x, double C = 1e10)
{
    const std::size_t n = x.size();
    if (n < 2)
        throw std::invalid_argument("findIntegerRelation: requires at least 2 elements");

    // 拡大格子の構築: n × (n+1)
    std::vector<std::vector<double>> basis(n, std::vector<double>(n + 1, 0.0));
    for (std::size_t i = 0; i < n; ++i) {
        basis[i][0] = C * x[i];
        basis[i][i + 1] = 1.0;
    }

    auto result = lllReduce(basis);

    // 最短ベクトルの整数部分 (index 1..n) を返す
    // 最初の列 (C·Σm_i·x_i) がゼロに近い行を選ぶ
    std::size_t best = 0;
    double bestNorm = std::abs(result.basis[0][0]);
    for (std::size_t i = 1; i < n; ++i) {
        double v = std::abs(result.basis[i][0]);
        if (v < bestNorm) {
            bestNorm = v;
            best = i;
        }
    }

    std::vector<int64_t> m(n);
    for (std::size_t i = 0; i < n; ++i) {
        m[i] = static_cast<int64_t>(std::round(result.basis[best][i + 1]));
    }

    return m;
}

} // namespace calx

#endif // CALX_LINALG_LLL_HPP