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

// DifferentialPrivacy.hpp — 微分プライバシー (DP-1)
// Laplace / Gaussian / Exponential メカニズムとプライバシー予算合成

#ifndef CALX_ML_PRIVACY_DIFFERENTIAL_PRIVACY_HPP
#define CALX_ML_PRIVACY_DIFFERENTIAL_PRIVACY_HPP

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

namespace calx::ml {

    // 微分プライバシー適用結果
    template<typename T>
    struct DPResult {
        std::vector<T> noisyData;
        T epsilon;
        T actualPrivacyLoss;
    };

    // ---------------------------------------------------------------
    // Laplace メカニズム
    // 各データ要素に Laplace(0, sensitivity/epsilon) ノイズを加える
    // ---------------------------------------------------------------
    template<typename T>
    std::vector<T> laplaceMechanism(const std::vector<T>& data, T sensitivity,
                                     T epsilon, unsigned seed = 42) {
        if (epsilon <= T(0))
            throw std::invalid_argument("laplaceMechanism: epsilon must be positive");
        if (sensitivity < T(0))
            throw std::invalid_argument("laplaceMechanism: sensitivity must be non-negative");

        std::mt19937 rng(seed);
        // Laplace 分布: U(-0.5, 0.5) から変換
        // X = -b * sign(U) * ln(1 - 2|U|),  b = sensitivity / epsilon
        std::uniform_real_distribution<double> uniform(-0.5, 0.5);
        const double b = static_cast<double>(sensitivity) / static_cast<double>(epsilon);

        std::vector<T> result(data.size());
        for (size_t i = 0; i < data.size(); ++i) {
            double u = uniform(rng);
            // 0 除算回避
            while (u == 0.0) u = uniform(rng);
            double noise = -b * ((u > 0.0) ? 1.0 : -1.0) * std::log(1.0 - 2.0 * std::abs(u));
            result[i] = data[i] + static_cast<T>(noise);
        }
        return result;
    }

    // ---------------------------------------------------------------
    // Gaussian メカニズム
    // (epsilon, delta)-DP: sigma = sensitivity * sqrt(2 * ln(1.25/delta)) / epsilon
    // ---------------------------------------------------------------
    template<typename T>
    std::vector<T> gaussianMechanism(const std::vector<T>& data, T sensitivity,
                                      T epsilon, T delta, unsigned seed = 42) {
        if (epsilon <= T(0))
            throw std::invalid_argument("gaussianMechanism: epsilon must be positive");
        if (delta <= T(0) || delta >= T(1))
            throw std::invalid_argument("gaussianMechanism: delta must be in (0, 1)");
        if (sensitivity < T(0))
            throw std::invalid_argument("gaussianMechanism: sensitivity must be non-negative");

        double sig = static_cast<double>(sensitivity)
                   * std::sqrt(2.0 * std::log(1.25 / static_cast<double>(delta)))
                   / static_cast<double>(epsilon);

        std::mt19937 rng(seed);
        std::normal_distribution<double> gauss(0.0, sig);

        std::vector<T> result(data.size());
        for (size_t i = 0; i < data.size(); ++i)
            result[i] = data[i] + static_cast<T>(gauss(rng));
        return result;
    }

    // ---------------------------------------------------------------
    // Exponential メカニズム (離散選択)
    // スコアベクトルから確率 ∝ exp(epsilon * score / (2 * sensitivity)) で選択
    // ---------------------------------------------------------------
    template<typename T>
    size_t exponentialMechanism(const std::vector<T>& scores, T sensitivity,
                                T epsilon, unsigned seed = 42) {
        if (scores.empty())
            throw std::invalid_argument("exponentialMechanism: scores must not be empty");
        if (epsilon <= T(0))
            throw std::invalid_argument("exponentialMechanism: epsilon must be positive");

        // オーバーフロー防止: 最大スコアを引く
        T maxScore = *std::max_element(scores.begin(), scores.end());
        double scale = static_cast<double>(epsilon) / (2.0 * static_cast<double>(sensitivity));

        std::vector<double> weights(scores.size());
        double sumWeights = 0.0;
        for (size_t i = 0; i < scores.size(); ++i) {
            weights[i] = std::exp(scale * static_cast<double>(scores[i] - maxScore));
            sumWeights += weights[i];
        }

        // 累積分布からサンプリング
        std::mt19937 rng(seed);
        std::uniform_real_distribution<double> uniform(0.0, sumWeights);
        double dart = uniform(rng);

        double cumulative = 0.0;
        for (size_t i = 0; i < weights.size(); ++i) {
            cumulative += weights[i];
            if (dart <= cumulative)
                return i;
        }
        return weights.size() - 1; // 浮動小数点丸め対策
    }

    // ---------------------------------------------------------------
    // プライバシー予算合成
    // ---------------------------------------------------------------

    // 基本合成: 各クエリの epsilon の単純和
    template<typename T>
    T basicComposition(const std::vector<T>& epsilons) {
        return std::accumulate(epsilons.begin(), epsilons.end(), T(0));
    }

    // 高度な合成 (Dwork et al. 2010)
    // k 回の (eps, delta)-DP クエリの合成:
    //   total_epsilon = sqrt(2k * ln(1/delta')) * eps + k * eps * (exp(eps) - 1)
    // ここで eps = max(epsilons), k = epsilons.size()
    template<typename T>
    T advancedComposition(const std::vector<T>& epsilons, T delta) {
        if (epsilons.empty()) return T(0);
        if (delta <= T(0) || delta >= T(1))
            throw std::invalid_argument("advancedComposition: delta must be in (0, 1)");

        size_t k = epsilons.size();
        // 各クエリが同一 epsilon を使う前提で最大値を使用
        double eps = static_cast<double>(*std::max_element(epsilons.begin(), epsilons.end()));
        double d = static_cast<double>(delta);

        double term1 = std::sqrt(2.0 * static_cast<double>(k) * std::log(1.0 / d)) * eps;
        double term2 = static_cast<double>(k) * eps * (std::exp(eps) - 1.0);

        return static_cast<T>(term1 + term2);
    }

} // namespace calx::ml

#endif // CALX_ML_PRIVACY_DIFFERENTIAL_PRIVACY_HPP
