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

// post_quantum.hpp — ポスト量子暗号 基礎 (CRYPTO-4)
// Simplified LWE-based encryption (Regev scheme, educational)

#ifndef CALX_CRYPTO_POST_QUANTUM_HPP
#define CALX_CRYPTO_POST_QUANTUM_HPP

#include <vector>
#include <cmath>
#include <random>
#include <stdexcept>

namespace calx::crypto {

    // LWE パラメータ
    struct LWEParams {
        int n = 256;        // 格子次元
        int q = 7681;       // 係数体の法 (素数)
        double sigma = 3.2; // 誤差分布の標準偏差
    };

    // LWE 鍵ペア
    struct LWEKeypair {
        std::vector<std::vector<int>> A;  // 公開行列 n×n
        std::vector<int> b;               // 公開ベクトル (A*s + e mod q)
        std::vector<int> s;               // 秘密鍵ベクトル
    };

    namespace detail {

        // mod q を [0, q) の範囲に正規化
        inline int modQ(int x, int q) {
            int r = x % q;
            return (r < 0) ? r + q : r;
        }

        // 離散ガウス分布からサンプリング (rounded Gaussian)
        inline int sampleGaussian(double sigma, std::mt19937& rng) {
            std::normal_distribution<double> dist(0.0, sigma);
            return static_cast<int>(std::round(dist(rng)));
        }

    } // namespace detail

    // LWE 鍵生成
    // 公開鍵: (A, b = A*s + e mod q), 秘密鍵: s
    inline LWEKeypair lweKeyGen(const LWEParams& params, unsigned seed = 42) {
        const int n = params.n;
        const int q = params.q;

        std::mt19937 rng(seed);
        std::uniform_int_distribution<int> uniformQ(0, q - 1);

        LWEKeypair kp;

        // 公開行列 A: n×n, 一様ランダム mod q
        kp.A.resize(n, std::vector<int>(n));
        for (int i = 0; i < n; ++i)
            for (int j = 0; j < n; ++j)
                kp.A[i][j] = uniformQ(rng);

        // 秘密鍵 s: 一様ランダム mod q
        kp.s.resize(n);
        for (int i = 0; i < n; ++i)
            kp.s[i] = uniformQ(rng);

        // 誤差ベクトル e: 離散ガウス分布
        std::vector<int> e(n);
        for (int i = 0; i < n; ++i)
            e[i] = detail::sampleGaussian(params.sigma, rng);

        // b = A*s + e mod q
        kp.b.resize(n);
        for (int i = 0; i < n; ++i) {
            long long sum = 0;
            for (int j = 0; j < n; ++j)
                sum += static_cast<long long>(kp.A[i][j]) * kp.s[j];
            kp.b[i] = detail::modQ(static_cast<int>((sum % q) + e[i]), q);
        }

        return kp;
    }

    // LWE 暗号化
    // message: 0 または 1 のビット
    // 暗号文: (u, v) を連結したベクトル (長さ n+1)
    //   u = A^T * r mod q,  v = b^T * r + floor(q/2)*m + e' mod q
    inline std::vector<int> lweEncrypt(const LWEKeypair& kp, const LWEParams& params,
                                       int message, unsigned seed = 42) {
        if (message != 0 && message != 1)
            throw std::invalid_argument("lweEncrypt: message must be 0 or 1");

        const int n = params.n;
        const int q = params.q;

        std::mt19937 rng(seed);

        // ランダムベクトル r: {0, 1}^n (ベルヌーイ)
        std::vector<int> r(n);
        std::uniform_int_distribution<int> coin(0, 1);
        for (int i = 0; i < n; ++i)
            r[i] = coin(rng);

        // u = A^T * r mod q (長さ n)
        std::vector<int> u(n, 0);
        for (int j = 0; j < n; ++j) {
            long long sum = 0;
            for (int i = 0; i < n; ++i)
                sum += static_cast<long long>(kp.A[i][j]) * r[i];
            u[j] = detail::modQ(static_cast<int>(sum % q), q);
        }

        // v = b^T * r + floor(q/2) * message + e' mod q
        long long btr = 0;
        for (int i = 0; i < n; ++i)
            btr += static_cast<long long>(kp.b[i]) * r[i];

        int ePrime = detail::sampleGaussian(params.sigma, rng);
        int v = detail::modQ(
            static_cast<int>((btr % q) + (q / 2) * message + ePrime), q);

        // 暗号文: u (長さ n) + v (1要素) = 長さ n+1
        std::vector<int> ciphertext(n + 1);
        for (int i = 0; i < n; ++i)
            ciphertext[i] = u[i];
        ciphertext[n] = v;

        return ciphertext;
    }

    // LWE 復号
    // 暗号文 (長さ n+1) から元のビットを復元
    // d = v - s^T * u mod q → floor(q/2) に近ければ 1, 0 に近ければ 0
    inline int lweDecrypt(const LWEKeypair& kp, const LWEParams& params,
                           const std::vector<int>& ciphertext) {
        const int n = params.n;
        const int q = params.q;

        if (static_cast<int>(ciphertext.size()) != n + 1)
            throw std::invalid_argument("lweDecrypt: ciphertext size mismatch");

        // u = ciphertext[0..n-1], v = ciphertext[n]
        long long stu = 0;
        for (int i = 0; i < n; ++i)
            stu += static_cast<long long>(kp.s[i]) * ciphertext[i];

        int d = detail::modQ(static_cast<int>(ciphertext[n] - (stu % q)), q);

        // d が floor(q/2) に近ければ 1, 0 に近ければ 0
        // 判定閾値: q/4
        int half = q / 2;
        int dist0 = std::min(d, q - d);            // 0 からの距離
        int dist1 = std::abs(d - half);             // q/2 からの距離
        return (dist1 < dist0) ? 1 : 0;
    }

} // namespace calx::crypto

#endif // CALX_CRYPTO_POST_QUANTUM_HPP
