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

// random.hpp
//
// Probability distributions and random number engine
//
// Class list:
//   RandomEngine — std::mt19937 seed-management wrapper
//   UniformDistribution<T> — uniform distribution [a, b]
//   NormalDistribution<T> — normal distribution (μ, σ)
//   MultivariateNormal<T> — multivariate normal distribution (μ, Σ)

#ifndef SANGI_RANDOM_HPP
#define SANGI_RANDOM_HPP

#include <math/core/common.hpp>
#include <math/core/vector.hpp>
#include <math/core/matrix.hpp>
#include <math/concepts/algebraic_concepts.hpp>
#include <math/linalg/decomposition.hpp>
#include <math/random/normal.hpp>  // Ziggurat normal random numbers
#include <math/random/random_sampling.hpp>  // Float random-number foundation (uniform01/sampleNormalStd etc.)
#include <algorithm>  // std::shuffle
#include <cmath>
#include <random>
#include <stdexcept>
#include <type_traits>
#include <vector>

namespace sangi {

// =====================================================================
// RandomEngine — seed management and reproducibility guarantee
// =====================================================================

/**
 * @brief Wrapper for a random number engine
 *
 * Holds a std::mt19937 internally and provides seed management and reproducibility guarantees.
 * Used for consistent seed control in statistical simulations and Monte Carlo methods.
 */
class RandomEngine {
public:
    explicit RandomEngine(unsigned int seed = 42) : gen_(seed) {}

    void seed(unsigned int s) { gen_.seed(s); }

    [[nodiscard]] std::mt19937& generator() { return gen_; }
    [[nodiscard]] const std::mt19937& generator() const { return gen_; }

    /// Uniform random number in [0, 1)
    [[nodiscard]] double uniform() {
        std::uniform_real_distribution<double> dist(0.0, 1.0);
        return dist(gen_);
    }

    /// Uniform random number in [a, b)
    [[nodiscard]] double uniform(double a, double b) {
        std::uniform_real_distribution<double> dist(a, b);
        return dist(gen_);
    }

    /// Normal random number (μ, σ). Internally the Ziggurat method (sangi::random::detail::ziggurat_normal)
    [[nodiscard]] double normal(double mu = 0.0, double sigma = 1.0) {
        return mu + sigma * sangi::random::detail::ziggurat_normal(gen_);
    }

    /// Uniform integer random number in [a, b] (both endpoints included)
    [[nodiscard]] int uniformInt(int a, int b) {
        std::uniform_int_distribution<int> dist(a, b);
        return dist(gen_);
    }

    /// Templated integer version (long, size_t, ...)
    template<typename Int>
        requires std::is_integral_v<Int>
    [[nodiscard]] Int uniformInt(Int a, Int b) {
        std::uniform_int_distribution<Int> dist(a, b);
        return dist(gen_);
    }

    /// In-place shuffle of the container's elements (Fisher-Yates)
    template<typename Container>
    void shuffle(Container& c) {
        std::shuffle(c.begin(), c.end(), gen_);
    }

    /// Uniform random choice from the container (returns a copy of the element)
    template<typename Container>
    [[nodiscard]] auto choice(const Container& c) -> typename Container::value_type {
        if (c.empty())
            throw std::invalid_argument("RandomEngine::choice: empty container");
        std::uniform_int_distribution<std::size_t> dist(0, c.size() - 1);
        return c[dist(gen_)];
    }

private:
    std::mt19937 gen_;
};

// =====================================================================
// UniformDistribution — uniform distribution [a, b]
// =====================================================================

/**
 * @brief Continuous uniform distribution U(a, b)
 *
 * pdf(x) = 1/(b-a)  (a ≤ x ≤ b)
 * cdf(x) = (x-a)/(b-a)
 */
template<concepts::OrderedField T>
class UniformDistribution {
public:
    UniformDistribution(T a = T(0), T b = T(1)) : a_(a), b_(b) {
        if (a >= b) {
            throw std::invalid_argument("UniformDistribution: a must be less than b");
        }
    }

    [[nodiscard]] T pdf(T x) const {
        if (x < a_ || x > b_) return T(0);
        return T(1) / (b_ - a_);
    }

    [[nodiscard]] T cdf(T x) const {
        if (x <= a_) return T(0);
        if (x >= b_) return T(1);
        return (x - a_) / (b_ - a_);
    }

    [[nodiscard]] T quantile(T p) const {
        if (p < T(0) || p > T(1)) {
            throw std::invalid_argument("UniformDistribution::quantile: p must be in [0, 1]");
        }
        return a_ + p * (b_ - a_);
    }

    [[nodiscard]] T mean() const { return (a_ + b_) / T(2); }
    [[nodiscard]] T variance() const { return (b_ - a_) * (b_ - a_) / T(12); }

    [[nodiscard]] T lower() const { return a_; }
    [[nodiscard]] T upper() const { return b_; }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        if constexpr (std::is_floating_point_v<T>) {
            std::uniform_real_distribution<T> dist(a_, b_);
            return dist(gen);
        } else {
            return a_ + (b_ - a_) * detail::uniform01<T>(gen);
        }
    }

    [[nodiscard]] std::vector<T> sample(std::mt19937& gen, std::size_t n) const {
        std::vector<T> result(n);
        for (auto& v : result) v = sample(gen);
        return result;
    }

private:
    T a_, b_;
};

// =====================================================================
// NormalDistribution — normal distribution N(μ, σ²)
// =====================================================================

namespace detail {

/// CDF of the standard normal distribution: Φ(x) = 0.5·erfc(-x/√2)
template<concepts::OrderedField T>
[[nodiscard]] T standard_normal_cdf(T x) {
    return T(0.5) * generic_erfc(-x / detail::generic_sqrt(T(2)));
}

/// Quantile of the standard normal distribution (inverse CDF): Peter Acklam's rational approximation
/// Accuracy: |ε| < 1.15e-9
template<concepts::OrderedField T>
[[nodiscard]] T standard_normal_quantile(T p) {
    if (p <= T(0) || p >= T(1)) {
        if (p == T(0)) return -std::numeric_limits<T>::infinity();
        if (p == T(1)) return std::numeric_limits<T>::infinity();
        throw std::invalid_argument("standard_normal_quantile: p must be in (0, 1)");
    }

    // Coefficients of the approximation (Peter Acklam)
    constexpr T a1 = T(-3.969683028665376e+01);
    constexpr T a2 = T( 2.209460984245205e+02);
    constexpr T a3 = T(-2.759285104469687e+02);
    constexpr T a4 = T( 1.383577518672690e+02);
    constexpr T a5 = T(-3.066479806614716e+01);
    constexpr T a6 = T( 2.506628277459239e+00);

    constexpr T b1 = T(-5.447609879822406e+01);
    constexpr T b2 = T( 1.615858368580409e+02);
    constexpr T b3 = T(-1.556989798598866e+02);
    constexpr T b4 = T( 6.680131188771972e+01);
    constexpr T b5 = T(-1.328068155288572e+01);

    constexpr T c1 = T(-7.784894002430293e-03);
    constexpr T c2 = T(-3.223964580411365e-01);
    constexpr T c3 = T(-2.400758277161838e+00);
    constexpr T c4 = T(-2.549732539343734e+00);
    constexpr T c5 = T( 4.374664141464968e+00);
    constexpr T c6 = T( 2.938163982698783e+00);

    constexpr T d1 = T( 7.784695709041462e-03);
    constexpr T d2 = T( 3.224671290700398e-01);
    constexpr T d3 = T( 2.445134137142996e+00);
    constexpr T d4 = T( 3.754408661907416e+00);

    constexpr T p_low  = T(0.02425);
    constexpr T p_high = T(1) - p_low;

    T q, r;

    if (p < p_low) {
        // Lower tail
        q = std::sqrt(T(-2) * std::log(p));
        return (((((c1*q + c2)*q + c3)*q + c4)*q + c5)*q + c6) /
               ((((d1*q + d2)*q + d3)*q + d4)*q + T(1));
    } else if (p <= p_high) {
        // Central region
        q = p - T(0.5);
        r = q * q;
        return (((((a1*r + a2)*r + a3)*r + a4)*r + a5)*r + a6) * q /
               (((((b1*r + b2)*r + b3)*r + b4)*r + b5)*r + T(1));
    } else {
        // Upper tail
        q = std::sqrt(T(-2) * std::log(T(1) - p));
        return -(((((c1*q + c2)*q + c3)*q + c4)*q + c5)*q + c6) /
                ((((d1*q + d2)*q + d3)*q + d4)*q + T(1));
    }
}

} // namespace detail

/**
 * @brief Normal distribution N(μ, σ²)
 *
 * pdf(x) = (1/(σ√(2π))) exp(-(x-μ)²/(2σ²))
 * cdf(x) = Φ((x-μ)/σ)
 */
template<concepts::OrderedField T>
class NormalDistribution {
public:
    NormalDistribution(T mu = T(0), T sigma = T(1)) : mu_(mu), sigma_(sigma) {
        if (sigma <= T(0)) {
            throw std::invalid_argument("NormalDistribution: sigma must be positive");
        }
    }

    [[nodiscard]] T pdf(T x) const {
        const T pi = std::acos(T(-1));
        T z = (x - mu_) / sigma_;
        return std::exp(T(-0.5) * z * z) / (sigma_ * std::sqrt(T(2) * pi));
    }

    [[nodiscard]] T cdf(T x) const {
        return detail::standard_normal_cdf((x - mu_) / sigma_);
    }

    [[nodiscard]] T quantile(T p) const {
        return mu_ + sigma_ * detail::standard_normal_quantile(p);
    }

    [[nodiscard]] T mean() const { return mu_; }
    [[nodiscard]] T variance() const { return sigma_ * sigma_; }
    [[nodiscard]] T stddev() const { return sigma_; }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        if constexpr (std::is_floating_point_v<T>) {
            return static_cast<T>(mu_ + sigma_ * static_cast<T>(
                sangi::random::detail::ziggurat_normal(gen)));
        } else {
            // Multi-precision Float: generate full-precision N(0,1) via Box-Muller
            return mu_ + sigma_ * detail::sampleNormalStd<T>(gen);
        }
    }

    [[nodiscard]] std::vector<T> sample(std::mt19937& gen, std::size_t n) const {
        std::vector<T> result(n);
        for (auto& v : result) v = sample(gen);
        return result;
    }

private:
    T mu_, sigma_;
};

// =====================================================================
// MultivariateNormal<T> — multivariate normal distribution
// =====================================================================

/**
 * @brief D-dimensional normal distribution N(μ, Σ)
 *
 * Provides the density function and sampling.
 * The constructor precomputes the Cholesky decomposition, inverse matrix, and normalization constant
 * to speed up calls to pdf/sample.
 *
 * Sampling: generate z ~ N(0, I) and x = μ + L·z (L: Cholesky factor)
 */
template<concepts::OrderedField T>
class MultivariateNormal {
public:
    MultivariateNormal(const Vector<T>& mean, const Matrix<T>& covariance)
        : dim_(mean.size()), mean_(mean), covariance_(covariance)
    {
        if (covariance.rows() != dim_ || covariance.cols() != dim_) {
            throw DimensionError(
                "MultivariateNormal: covariance must be " +
                std::to_string(dim_) + "x" + std::to_string(dim_));
        }

        // Cholesky decomposition (exception if not positive definite)
        L_ = algorithms::cholesky_decomposition(covariance_);

        // Inverse matrix (for PDF)
        covInv_ = algorithms::lu_inverse(covariance_);

        // Normalization constant: -0.5 * (D·log(2π) + log|Σ|)
        const T pi = std::acos(T(-1));
        T logDet = std::log(std::abs(algorithms::lu_determinant(covariance_)));
        logNormConst_ = T(-0.5) * (static_cast<T>(dim_) * std::log(T(2) * pi) + logDet);
    }

    /// Probability density function
    [[nodiscard]] T pdf(const Vector<T>& x) const {
        return std::exp(logPdf(x));
    }

    /// Log probability density function
    [[nodiscard]] T logPdf(const Vector<T>& x) const {
        if (x.size() != dim_) {
            throw DimensionError("MultivariateNormal::logPdf: dimension mismatch");
        }
        Vector<T> diff = x;
        diff -= mean_;
        Vector<T> tmp = covInv_ * diff;
        T mahal = diff.dot(tmp);
        return logNormConst_ - T(0.5) * mahal;
    }

    /// Generate 1 sample
    [[nodiscard]] Vector<T> sample(std::mt19937& gen) const {
        Vector<T> z(dim_);
        for (std::size_t i = 0; i < dim_; ++i) {
            if constexpr (std::is_floating_point_v<T>) {
                z[i] = static_cast<T>(sangi::random::detail::ziggurat_normal(gen));
            } else {
                z[i] = detail::sampleNormalStd<T>(gen);
            }
        }
        // x = μ + L·z
        Vector<T> result = L_ * z;
        result += mean_;
        return result;
    }

    /// Generate n samples
    [[nodiscard]] std::vector<Vector<T>> sample(std::mt19937& gen, std::size_t n) const {
        std::vector<Vector<T>> results;
        results.reserve(n);
        for (std::size_t i = 0; i < n; ++i) {
            results.push_back(sample(gen));
        }
        return results;
    }

    [[nodiscard]] std::size_t dim() const { return dim_; }
    [[nodiscard]] const Vector<T>& mean() const { return mean_; }
    [[nodiscard]] const Matrix<T>& covariance() const { return covariance_; }

private:
    std::size_t dim_;
    Vector<T> mean_;
    Matrix<T> covariance_;
    Matrix<T> L_;            // Cholesky factor
    Matrix<T> covInv_;       // Σ^{-1}
    T logNormConst_;         // -0.5 * (D·log(2π) + log|Σ|)
};

} // namespace sangi

#endif // SANGI_RANDOM_HPP