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

// distributions_extra.hpp
// Noncentral distributions (5-3g) + additional continuous distributions (5-3h)
//
// Noncentral distributions:
//   NoncentralChiSquaredDistribution, NoncentralTDistribution, NoncentralFDistribution
//
// Additional distributions:
//   ArcsineDistribution, InverseGammaDistribution, InverseGaussianDistribution,
//   SkewNormalDistribution, TriangularDistribution, DirichletDistribution

#ifndef SANGI_DISTRIBUTIONS_EXTRA_HPP
#define SANGI_DISTRIBUTIONS_EXTRA_HPP

#include <math/random/distributions.hpp>
#include <cmath>
#include <random>
#include <stdexcept>
#include <vector>
#include <numbers>
#include <limits>
#include <numeric>

namespace sangi {

// ================================================================
// ArcsineDistribution — arcsine distribution Arcsine(a, b)
// ================================================================

template<FloatingPointType T>
class ArcsineDistribution {
public:
    ArcsineDistribution(T a = T(0), T b = T(1)) : a_(a), b_(b) {
        if (a >= b)
            throw std::invalid_argument("ArcsineDistribution: a must be < b");
    }

    [[nodiscard]] T pdf(T x) const {
        if (x <= a_ || x >= b_) return T(0);
        return T(1) / (detail::generic_pi<T>() * std::sqrt((x - a_) * (b_ - x)));
    }

    [[nodiscard]] T cdf(T x) const {
        if (x <= a_) return T(0);
        if (x >= b_) return T(1);
        return T(2) / detail::generic_pi<T>() * std::asin(std::sqrt((x - a_) / (b_ - a_)));
    }

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

    [[nodiscard]] T sample(std::mt19937& gen) const {
        // Use Beta(0.5, 0.5)
        T p;
        if constexpr (std::is_floating_point_v<T>) {
            std::uniform_real_distribution<T> u(T(0), T(1));
            p = u(gen);
        } else {
            p = detail::uniform01<T>(gen);
        }
        T s = detail::generic_sin(detail::generic_pi<T>() * p / T(2));
        return a_ + (b_ - a_) * s * s;
    }

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

private:
    T a_, b_;
};

// ================================================================
// InverseGammaDistribution — inverse-gamma distribution InvGamma(α, β)
// ================================================================

template<FloatingPointType T>
class InverseGammaDistribution {
public:
    InverseGammaDistribution(T alpha = T(1), T beta = T(1)) : alpha_(alpha), beta_(beta) {
        if (alpha <= T(0) || beta <= T(0))
            throw std::invalid_argument("InverseGammaDistribution: alpha, beta must be > 0");
    }

    [[nodiscard]] T pdf(T x) const {
        if (x <= T(0)) return T(0);
        return std::pow(beta_, alpha_) / detail::generic_tgamma(alpha_)
             * std::pow(x, -alpha_ - T(1)) * std::exp(-beta_ / x);
    }

    [[nodiscard]] T cdf(T x) const {
        if (x <= T(0)) return T(0);
        // P(X ≤ x) = 1 - P(α, β/x) (upper regularized incomplete gamma)
        return T(1) - special::gammaP<T>(alpha_, beta_ / x);
    }

    [[nodiscard]] T mean() const {
        if (alpha_ <= T(1)) return std::numeric_limits<T>::infinity();
        return beta_ / (alpha_ - T(1));
    }

    [[nodiscard]] T variance() const {
        if (alpha_ <= T(2)) return std::numeric_limits<T>::infinity();
        return beta_ * beta_ / ((alpha_ - T(1)) * (alpha_ - T(1)) * (alpha_ - T(2)));
    }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        if constexpr (std::is_floating_point_v<T>) {
            std::gamma_distribution<T> g(alpha_, T(1) / beta_);
            return T(1) / g(gen);
        } else {
            // X = 1 / Gamma(α, scale 1/β)
            return T(1) / (detail::sampleGammaStd<T>(gen, alpha_) / beta_);
        }
    }

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

private:
    T alpha_, beta_;
};

// ================================================================
// InverseGaussianDistribution — inverse Gaussian (Wald) distribution IG(μ, λ)
// ================================================================

template<FloatingPointType T>
class InverseGaussianDistribution {
public:
    InverseGaussianDistribution(T mu = T(1), T lambda = T(1)) : mu_(mu), lambda_(lambda) {
        if (mu <= T(0) || lambda <= T(0))
            throw std::invalid_argument("InverseGaussianDistribution: mu, lambda must be > 0");
    }

    [[nodiscard]] T pdf(T x) const {
        if (x <= T(0)) return T(0);
        T d = x - mu_;
        return std::sqrt(lambda_ / (T(2) * detail::generic_pi<T>() * x * x * x))
             * std::exp(-lambda_ * d * d / (T(2) * mu_ * mu_ * x));
    }

    [[nodiscard]] T cdf(T x) const {
        if (x <= T(0)) return T(0);
        T sq = std::sqrt(lambda_ / x);
        T t1 = sq * (x / mu_ - T(1));
        T t2 = sq * (x / mu_ + T(1));
        return T(0.5) * detail::generic_erfc(-t1 / std::sqrt(T(2)))
             + std::exp(T(2) * lambda_ / mu_)
             * T(0.5) * detail::generic_erfc(t2 / std::sqrt(T(2)));
    }

    [[nodiscard]] T mean() const { return mu_; }
    [[nodiscard]] T variance() const { return mu_ * mu_ * mu_ / lambda_; }

    // Michael-Schucany-Haas algorithm
    [[nodiscard]] T sample(std::mt19937& gen) const {
        T y, uu;
        if constexpr (std::is_floating_point_v<T>) {
            std::normal_distribution<T> norm(T(0), T(1));
            std::uniform_real_distribution<T> u(T(0), T(1));
            y = norm(gen);
            uu = u(gen);
        } else {
            y = detail::sampleNormalStd<T>(gen);
            uu = detail::uniform01<T>(gen);
        }
        y = y * y;
        T x = mu_ + (mu_ * mu_ * y) / (T(2) * lambda_)
            - (mu_ / (T(2) * lambda_))
            * detail::generic_sqrt(T(4) * mu_ * lambda_ * y + mu_ * mu_ * y * y);
        if (uu <= mu_ / (mu_ + x))
            return x;
        else
            return mu_ * mu_ / x;
    }

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

private:
    T mu_, lambda_;
};

// ================================================================
// SkewNormalDistribution — skew-normal distribution SN(ξ, ω, α)
// ================================================================

template<FloatingPointType T>
class SkewNormalDistribution {
public:
    SkewNormalDistribution(T xi = T(0), T omega = T(1), T alpha = T(0))
        : xi_(xi), omega_(omega), alpha_(alpha) {
        if (omega <= T(0))
            throw std::invalid_argument("SkewNormalDistribution: omega must be > 0");
    }

    [[nodiscard]] T pdf(T x) const {
        T z = (x - xi_) / omega_;
        T phi = std::exp(-T(0.5) * z * z) / std::sqrt(T(2) * detail::generic_pi<T>());
        T Phi = T(0.5) * detail::generic_erfc(-alpha_ * z / std::sqrt(T(2)));
        return T(2) / omega_ * phi * Phi;
    }

    [[nodiscard]] T cdf(T x) const {
        // The exact CDF via Owen's T function is complex; approximate by numerical integration.
        // Simple: trapezoidal rule (100 subdivisions)
        T z = (x - xi_) / omega_;
        int N = 200;
        T lo = -T(10);
        if (z < lo) return T(0);
        T hi = z;
        T h = (hi - lo) / N;
        T sum = T(0);
        for (int i = 0; i <= N; ++i) {
            T t = lo + i * h;
            T phi = std::exp(-T(0.5) * t * t) / std::sqrt(T(2) * detail::generic_pi<T>());
            T Phi = T(0.5) * detail::generic_erfc(-alpha_ * t / std::sqrt(T(2)));
            T val = T(2) * phi * Phi;
            if (i == 0 || i == N)
                sum += val;
            else
                sum += T(2) * val;
        }
        return sum * h / T(2);
    }

    [[nodiscard]] T mean() const {
        T delta = alpha_ / std::sqrt(T(1) + alpha_ * alpha_);
        return xi_ + omega_ * delta * std::sqrt(T(2) / detail::generic_pi<T>());
    }

    [[nodiscard]] T variance() const {
        T delta = alpha_ / std::sqrt(T(1) + alpha_ * alpha_);
        return omega_ * omega_ * (T(1) - T(2) * delta * delta / detail::generic_pi<T>());
    }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        T u0, u1;
        if constexpr (std::is_floating_point_v<T>) {
            std::normal_distribution<T> norm(T(0), T(1));
            u0 = norm(gen); u1 = norm(gen);
        } else {
            u0 = detail::sampleNormalStd<T>(gen); u1 = detail::sampleNormalStd<T>(gen);
        }
        T delta = alpha_ / detail::generic_sqrt(T(1) + alpha_ * alpha_);
        T z = delta * detail::generic_abs(u0) + detail::generic_sqrt(T(1) - delta * delta) * u1;
        return xi_ + omega_ * z;
    }

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

private:
    T xi_, omega_, alpha_;
};

// ================================================================
// TriangularDistribution — triangular distribution Triangular(a, b, c)
// ================================================================

template<FloatingPointType T>
class TriangularDistribution {
public:
    // a = lower bound, b = upper bound, c = mode
    TriangularDistribution(T a = T(0), T b = T(1), T c = T(0.5)) : a_(a), b_(b), c_(c) {
        if (a >= b || c < a || c > b)
            throw std::invalid_argument(
                "TriangularDistribution: must have a < b, a <= c <= b");
    }

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

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

    [[nodiscard]] T mean() const { return (a_ + b_ + c_) / T(3); }
    [[nodiscard]] T variance() const {
        return (a_ * a_ + b_ * b_ + c_ * c_ - a_ * b_ - a_ * c_ - b_ * c_) / T(18);
    }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        T p;
        if constexpr (std::is_floating_point_v<T>) {
            std::uniform_real_distribution<T> u(T(0), T(1));
            p = u(gen);
        } else {
            p = detail::uniform01<T>(gen);
        }
        T fc = (c_ - a_) / (b_ - a_);
        if (p < fc)
            return a_ + detail::generic_sqrt(p * (b_ - a_) * (c_ - a_));
        else
            return b_ - detail::generic_sqrt((T(1) - p) * (b_ - a_) * (b_ - c_));
    }

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

private:
    T a_, b_, c_;
};

// ================================================================
// DirichletDistribution — Dirichlet distribution Dir(α)
// ================================================================

template<FloatingPointType T>
class DirichletDistribution {
public:
    explicit DirichletDistribution(const std::vector<T>& alpha) : alpha_(alpha) {
        if (alpha.size() < 2)
            throw std::invalid_argument("DirichletDistribution: need at least 2 components");
        for (auto a : alpha)
            if (a <= T(0))
                throw std::invalid_argument("DirichletDistribution: all alpha must be > 0");
    }

    [[nodiscard]] size_t dim() const { return alpha_.size(); }

    [[nodiscard]] T logPdf(const std::vector<T>& x) const {
        if (x.size() != alpha_.size())
            throw std::invalid_argument("DirichletDistribution::logPdf: dimension mismatch");
        T logB = T(0);
        T sumAlpha = T(0);
        for (auto a : alpha_) {
            logB += detail::generic_lgamma(a);
            sumAlpha += a;
        }
        logB -= detail::generic_lgamma(sumAlpha);

        T logP = -logB;
        for (size_t i = 0; i < alpha_.size(); ++i)
            logP += (alpha_[i] - T(1)) * std::log(x[i]);
        return logP;
    }

    [[nodiscard]] T pdf(const std::vector<T>& x) const {
        return std::exp(logPdf(x));
    }

    [[nodiscard]] std::vector<T> mean() const {
        T s = T(0);
        for (auto a : alpha_) s += a;
        std::vector<T> m(alpha_.size());
        for (size_t i = 0; i < alpha_.size(); ++i) m[i] = alpha_[i] / s;
        return m;
    }

    [[nodiscard]] std::vector<T> sample(std::mt19937& gen) const {
        std::vector<T> r(alpha_.size());
        T sum = T(0);
        for (size_t i = 0; i < alpha_.size(); ++i) {
            if constexpr (std::is_floating_point_v<T>) {
                std::gamma_distribution<T> g(alpha_[i], T(1));
                r[i] = g(gen);
            } else {
                r[i] = detail::sampleGammaStd<T>(gen, alpha_[i]);
            }
            sum += r[i];
        }
        for (auto& v : r) v /= sum;
        return r;
    }

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

private:
    std::vector<T> alpha_;
};

// ================================================================
// NoncentralChiSquaredDistribution — noncentral chi-squared distribution χ²(k, λ)
// ================================================================

template<FloatingPointType T>
class NoncentralChiSquaredDistribution {
public:
    NoncentralChiSquaredDistribution(T k = T(1), T lambda = T(0))
        : k_(k), lambda_(lambda) {
        if (k <= T(0))
            throw std::invalid_argument("NoncentralChiSquaredDistribution: k must be > 0");
        if (lambda < T(0))
            throw std::invalid_argument("NoncentralChiSquaredDistribution: lambda must be >= 0");
    }

    // PDF: series expansion of the noncentral chi-squared
    [[nodiscard]] T pdf(T x) const {
        if (x <= T(0)) return T(0);
        T half_k = k_ / T(2);
        T half_lam = lambda_ / T(2);
        T sum = T(0);
        T term = std::exp(-half_lam);
        for (int j = 0; j < 200; ++j) {
            T a = half_k + T(j);
            T p = std::pow(x / T(2), a - T(1)) * std::exp(-x / T(2))
                / (T(2) * detail::generic_tgamma(a));
            sum += term * p;
            term *= half_lam / T(j + 1);
            if (term * p < sum * std::numeric_limits<T>::epsilon()) break;
        }
        return sum;
    }

    // CDF: series expansion
    [[nodiscard]] T cdf(T x) const {
        if (x <= T(0)) return T(0);
        T half_lam = lambda_ / T(2);
        T sum = T(0);
        T weight = std::exp(-half_lam);
        for (int j = 0; j < 200; ++j) {
            T a = k_ / T(2) + T(j);
            T p = special::gammaP<T>(a, x / T(2));
            sum += weight * p;
            weight *= half_lam / T(j + 1);
            if (weight < sum * std::numeric_limits<T>::epsilon()) break;
        }
        return sum;
    }

    [[nodiscard]] T mean() const { return k_ + lambda_; }
    [[nodiscard]] T variance() const { return T(2) * (k_ + T(2) * lambda_); }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        if constexpr (std::is_floating_point_v<T>) {
            // X ~ χ²(k-1) + (Z + √λ)² where Z ~ N(0,1)
            if (k_ >= T(1)) {
                std::chi_squared_distribution<T> chi2(k_ - T(1));
                std::normal_distribution<T> norm(std::sqrt(lambda_), T(1));
                T z = norm(gen);
                return chi2(gen) + z * z;
            } else {
                std::poisson_distribution<int> pois(static_cast<double>(lambda_ / T(2)));
                int j = pois(gen);
                std::chi_squared_distribution<T> chi2(k_ + T(2) * T(j));
                return chi2(gen);
            }
        } else {
            if (k_ >= T(1)) {
                // χ²(k-1)=2·Gamma((k-1)/2),  Z~N(√λ,1)
                T c = detail::sampleGammaStd<T>(gen, (k_ - T(1)) / T(2)) * T(2);
                T z = detail::sampleNormalStd<T>(gen) + detail::generic_sqrt(lambda_);
                return c + z * z;
            } else {
                int j = detail::samplePoissonKnuth<T>(gen, lambda_ / T(2));
                return detail::sampleGammaStd<T>(gen, (k_ + T(2) * T(j)) / T(2)) * T(2);
            }
        }
    }

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

private:
    T k_, lambda_;
};

// ================================================================
// NoncentralTDistribution — noncentral t distribution t(ν, δ)
// ================================================================

template<FloatingPointType T>
class NoncentralTDistribution {
public:
    NoncentralTDistribution(T nu = T(1), T delta = T(0))
        : nu_(nu), delta_(delta) {
        if (nu <= T(0))
            throw std::invalid_argument("NoncentralTDistribution: nu must be > 0");
    }

    [[nodiscard]] T mean() const {
        if (nu_ <= T(1)) return std::numeric_limits<T>::quiet_NaN();
        return delta_ * std::sqrt(nu_ / T(2)) * detail::generic_tgamma((nu_ - T(1)) / T(2))
             / detail::generic_tgamma(nu_ / T(2));
    }

    [[nodiscard]] T variance() const {
        if (nu_ <= T(2)) return std::numeric_limits<T>::quiet_NaN();
        T m = mean();
        return nu_ * (T(1) + delta_ * delta_) / (nu_ - T(2)) - m * m;
    }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        if constexpr (std::is_floating_point_v<T>) {
            std::normal_distribution<T> norm(delta_, T(1));
            std::chi_squared_distribution<T> chi2(nu_);
            return norm(gen) / std::sqrt(chi2(gen) / nu_);
        } else {
            T z = detail::sampleNormalStd<T>(gen) + delta_;          // N(δ,1)
            T v = detail::sampleGammaStd<T>(gen, nu_ / T(2)) * T(2);  // χ²(ν)
            return z / detail::generic_sqrt(v / nu_);
        }
    }

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

private:
    T nu_, delta_;
};

// ================================================================
// NoncentralFDistribution — noncentral F distribution F(d1, d2, λ)
// ================================================================

template<FloatingPointType T>
class NoncentralFDistribution {
public:
    NoncentralFDistribution(T d1 = T(1), T d2 = T(1), T lambda = T(0))
        : d1_(d1), d2_(d2), lambda_(lambda) {
        if (d1 <= T(0) || d2 <= T(0))
            throw std::invalid_argument("NoncentralFDistribution: d1, d2 must be > 0");
        if (lambda < T(0))
            throw std::invalid_argument("NoncentralFDistribution: lambda must be >= 0");
    }

    [[nodiscard]] T mean() const {
        if (d2_ <= T(2)) return std::numeric_limits<T>::infinity();
        return d2_ * (d1_ + lambda_) / (d1_ * (d2_ - T(2)));
    }

    [[nodiscard]] T sample(std::mt19937& gen) const {
        // (χ²_nc(d1, λ) / d1) / (χ²(d2) / d2)
        NoncentralChiSquaredDistribution<T> nc(d1_, lambda_);
        if constexpr (std::is_floating_point_v<T>) {
            std::chi_squared_distribution<T> chi2(d2_);
            return (nc.sample(gen) / d1_) / (chi2(gen) / d2_);
        } else {
            T c2 = detail::sampleGammaStd<T>(gen, d2_ / T(2)) * T(2);  // χ²(d2)
            return (nc.sample(gen) / d1_) / (c2 / d2_);
        }
    }

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

private:
    T d1_, d2_, lambda_;
};

} // namespace sangi

#endif // SANGI_DISTRIBUTIONS_EXTRA_HPP