// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package stats

import (
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

import (
	"math"
)

// Distribution additions to the families of distrib.go and cdf.go: the
// Weibull, lognormal, Pareto and negative binomial laws, and the
// Dirichlet. The continuous laws carry closed density, CDF and
// quantile forms; the negative binomial follows the discrete house
// shape of Poisson and the binomial, counting on the same integer axis
// the Poisson counts on. The support convention of the existing
// distributions holds throughout: a CDF is 0 below the support and a
// density 0 outside it, while a faulty parameter is an error and never
// a silent value.

// WeibullDensity returns the Weibull density with shape k > 0 and
// scale λ > 0 at x ≥ 0: (k/λ)(x/λ)^{k−1}e^{−(x/λ)^k}. Below the
// support, and at +∞, the density is 0; at x = 0 the formula speaks
// for itself, 0 for k > 1, the finite 1/λ for k = 1 and the +Inf the
// integrable singularity carries for k < 1.
func WeibullDensity(x, k, lambda float64) (float64, error) {
	if !(k > 0) || math.IsInf(k, 0) {
		return 0, base.Errf("WeibullDensity: shape k must be finite and positive, got %g", k)
	}
	if !(lambda > 0) || math.IsInf(lambda, 0) {
		return 0, base.Errf("WeibullDensity: scale λ must be finite and positive, got %g", lambda)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("WeibullDensity: x must be a number, got %g", x)
	}
	if x < 0 || math.IsInf(x, 1) {
		return 0, nil
	}
	z := x / lambda
	return k / lambda * math.Pow(z, k-1) * math.Exp(-math.Pow(z, k)), nil
}

// WeibullCDF returns P(X ≤ x) for X ~ Weibull(k, λ), the closed form
// 1 − e^{−(x/λ)^k}.
func WeibullCDF(x, k, lambda float64) (float64, error) {
	if !(k > 0) || math.IsInf(k, 0) {
		return 0, base.Errf("WeibullCDF: shape k must be finite and positive, got %g", k)
	}
	if !(lambda > 0) || math.IsInf(lambda, 0) {
		return 0, base.Errf("WeibullCDF: scale λ must be finite and positive, got %g", lambda)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("WeibullCDF: x must be a number, got %g", x)
	}
	if x <= 0 {
		return 0, nil
	}
	// Expm1 keeps the left tail, as in ExponentialCDF.
	return -math.Expm1(-math.Pow(x/lambda, k)), nil
}

// WeibullQuantile returns the q-quantile of Weibull(k, λ), the closed
// form λ(−ln(1−q))^{1/k}.
func WeibullQuantile(q, k, lambda float64) (float64, error) {
	if !(k > 0) || math.IsInf(k, 0) {
		return 0, base.Errf("WeibullQuantile: shape k must be finite and positive, got %g", k)
	}
	if !(lambda > 0) || math.IsInf(lambda, 0) {
		return 0, base.Errf("WeibullQuantile: scale λ must be finite and positive, got %g", lambda)
	}
	// NaN-rejecting on purpose, as in NormalQuantile.
	if !(q >= 0 && q <= 1) {
		return 0, base.Errf("WeibullQuantile: q must lie in [0, 1], got %g", q)
	}
	if q == 0 || q == 1 {
		return 0, base.Errf("WeibullQuantile: q = %g has no finite quantile", q)
	}
	return lambda * math.Pow(-math.Log1p(-q), 1/k), nil
}

// LognormalDensity returns the lognormal density with location μ and
// log-scale σ > 0 at x > 0: 1/(xσ√(2π))e^{−(ln x−μ)²/(2σ²)}. The
// support convention gives 0 at x ≤ 0 and at +∞.
func LognormalDensity(x, mu, sigma float64) (float64, error) {
	if math.IsNaN(mu) || math.IsInf(mu, 0) {
		return 0, base.Errf("LognormalDensity: location μ must be finite, got %g", mu)
	}
	if !(sigma > 0) || math.IsInf(sigma, 0) {
		return 0, base.Errf("LognormalDensity: log-scale σ must be finite and positive, got %g", sigma)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("LognormalDensity: x must be a number, got %g", x)
	}
	if x <= 0 || math.IsInf(x, 1) {
		return 0, nil
	}
	z := (math.Log(x) - mu) / sigma
	// The tail is assembled in log space: at a subnormal x both the
	// numerator and the denominator of the closed form underflow to
	// exact zero, and their division reports NaN for a point inside
	// the support whose density is an honest 0.
	return math.Exp(-z*z/2 - math.Log(x) - math.Log(sigma) - 0.5*math.Log(2*math.Pi)), nil
}

// LognormalCDF returns P(X ≤ x) for X ~ lognormal(μ, σ), the normal
// CDF at (ln x − μ)/σ, the reduction the law is named for.
func LognormalCDF(x, mu, sigma float64) (float64, error) {
	if math.IsNaN(mu) || math.IsInf(mu, 0) {
		return 0, base.Errf("LognormalCDF: location μ must be finite, got %g", mu)
	}
	if !(sigma > 0) || math.IsInf(sigma, 0) {
		return 0, base.Errf("LognormalCDF: log-scale σ must be finite and positive, got %g", sigma)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("LognormalCDF: x must be a number, got %g", x)
	}
	if x <= 0 {
		return 0, nil
	}
	return NormalCDF((math.Log(x) - mu) / sigma), nil
}

// LognormalQuantile returns the q-quantile of lognormal(μ, σ) through
// the existing normal quantile: e^{μ + σ·Φ^{−1}(q)}.
func LognormalQuantile(q, mu, sigma float64) (float64, error) {
	if math.IsNaN(mu) || math.IsInf(mu, 0) {
		return 0, base.Errf("LognormalQuantile: location μ must be finite, got %g", mu)
	}
	if !(sigma > 0) || math.IsInf(sigma, 0) {
		return 0, base.Errf("LognormalQuantile: log-scale σ must be finite and positive, got %g", sigma)
	}
	z, err := NormalQuantile(q)
	if err != nil {
		return 0, base.Errf("LognormalQuantile: %w", err)
	}
	return math.Exp(mu + sigma*z), nil
}

// ParetoDensity returns the Pareto density with scale x_m > 0 and tail
// index α > 0 at x ≥ x_m: α·x_m^α/x^{α+1}. Below the support, and at
// +∞, the density is 0.
func ParetoDensity(x, xm, alpha float64) (float64, error) {
	if !(xm > 0) || math.IsInf(xm, 0) {
		return 0, base.Errf("ParetoDensity: scale x_m must be finite and positive, got %g", xm)
	}
	if !(alpha > 0) || math.IsInf(alpha, 0) {
		return 0, base.Errf("ParetoDensity: tail index α must be finite and positive, got %g", alpha)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("ParetoDensity: x must be a number, got %g", x)
	}
	if x < xm || math.IsInf(x, 1) {
		return 0, nil
	}
	return alpha / x * math.Pow(xm/x, alpha), nil
}

// ParetoCDF returns P(X ≤ x) for X ~ Pareto(x_m, α), the closed form
// 1 − (x_m/x)^α, evaluated through Expm1 so the answers just above the
// support keep their digits.
func ParetoCDF(x, xm, alpha float64) (float64, error) {
	if !(xm > 0) || math.IsInf(xm, 0) {
		return 0, base.Errf("ParetoCDF: scale x_m must be finite and positive, got %g", xm)
	}
	if !(alpha > 0) || math.IsInf(alpha, 0) {
		return 0, base.Errf("ParetoCDF: tail index α must be finite and positive, got %g", alpha)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("ParetoCDF: x must be a number, got %g", x)
	}
	if x < xm {
		return 0, nil
	}
	return -math.Expm1(alpha * math.Log(xm/x)), nil
}

// ParetoQuantile returns the q-quantile of Pareto(x_m, α), the closed
// form x_m(1−q)^{−1/α}.
func ParetoQuantile(q, xm, alpha float64) (float64, error) {
	if !(xm > 0) || math.IsInf(xm, 0) {
		return 0, base.Errf("ParetoQuantile: scale x_m must be finite and positive, got %g", xm)
	}
	if !(alpha > 0) || math.IsInf(alpha, 0) {
		return 0, base.Errf("ParetoQuantile: tail index α must be finite and positive, got %g", alpha)
	}
	if !(q >= 0 && q <= 1) {
		return 0, base.Errf("ParetoQuantile: q must lie in [0, 1], got %g", q)
	}
	if q == 0 || q == 1 {
		return 0, base.Errf("ParetoQuantile: q = %g has no finite quantile", q)
	}
	return xm * math.Pow(1-q, -1/alpha), nil
}

// NegativeBinomialPMF returns P(X = k), the probability of k failures
// before the r-th success in independent trials of probability p: the
// law the Poisson draws and BinomialDraws count on the same integer
// axis. The mass is assembled in log space with lgamma, the form that
// keeps every term representable for large r and k.
func NegativeBinomialPMF(k, r int, p float64) (float64, error) {
	if r < 1 {
		return 0, base.Errf("NegativeBinomialPMF: r must be ≥ 1, got %d", r)
	}
	if !(p > 0 && p < 1) {
		return 0, base.Errf("NegativeBinomialPMF: p must lie in (0, 1), got %g", p)
	}
	if k < 0 {
		return 0, nil
	}
	lf := logGamma(float64(r+k)) - logGamma(float64(r)) - logGamma(float64(k+1)) +
		float64(r)*math.Log(p) + float64(k)*math.Log1p(-p)
	return math.Exp(lf), nil
}

// NegativeBinomialCDF returns P(X ≤ k) for the number of failures X
// before the r-th success, by direct summation of the PMF terms under
// the multiplicative recurrence term_{j+1} = term_j·(r+j)/(j+1)·(1−p).
// Every term is positive, so the sum carries no cancellation; the same
// probability equals the regularised beta I_p(r, k+1), which the tests
// hold the summation against. The summation seeds from p^r, and a seed
// the format cannot hold rounds to an exact zero the recurrence never
// recovers from: every later term would stay zero while the true mass
// sits further out. That far regime answers through the beta identity
// instead, which keeps the whole support live.
func NegativeBinomialCDF(k, r int, p float64) (float64, error) {
	if r < 1 {
		return 0, base.Errf("NegativeBinomialCDF: r must be ≥ 1, got %d", r)
	}
	if !(p > 0 && p < 1) {
		return 0, base.Errf("NegativeBinomialCDF: p must lie in (0, 1), got %g", p)
	}
	if k < 0 {
		return 0, nil
	}
	term := math.Exp(float64(r) * math.Log(p))
	if term == 0 {
		// p^r underflowed: the recurrence multiplies zeros, so the sum
		// would answer 0 at every k. The identity I_p(r, k+1) = P(X ≤ k)
		// evaluates the same probability through the incomplete beta,
		// accurate across this regime.
		return BetaIncomplete(p, float64(r), float64(k+1))
	}
	sum := 0.0
	for j := 0; j <= k; j++ {
		sum += term
		term *= (float64(r) + float64(j)) / float64(j+1) * (1 - p)
	}
	return sum, nil
}

// NegativeBinomialQuantile returns the smallest k with P(X ≤ k) ≥ q,
// through the same discrete bracketed search the Poisson and binomial
// quantiles use.
func NegativeBinomialQuantile(q, p float64, r int) (float64, error) {
	if r < 1 {
		return 0, base.Errf("NegativeBinomialQuantile: r must be ≥ 1, got %d", r)
	}
	if !(p > 0 && p < 1) {
		return 0, base.Errf("NegativeBinomialQuantile: p must lie in (0, 1), got %g", p)
	}
	return discreteQuantile("NegativeBinomialQuantile", q, func(k int) (float64, error) {
		return NegativeBinomialCDF(k, r, p)
	})
}

// logGamma wraps math.Lgamma, keeping the log-space call sites free of
// the ignored-error idiom.
func logGamma(x float64) float64 {
	l, _ := math.Lgamma(x)
	return l
}

// DirichletDensity returns the Dirichlet density with concentration
// vector α at the simplex point x: Πx_i^{α_i−1}/B(α), the normalising
// constant assembled with lgamma. Both vectors must be finite, every
// α_i positive and every x_i non-negative, and the x must sum to 1 (to
// within 1e-9); a boundary x_i = 0 gives +Inf below α_i = 1, the value
// 1 continues to contribute nothing at α_i = 1 exactly, and 0 above.
func DirichletDensity(alpha, x []float64) (float64, error) {
	const name = "DirichletDensity"
	if len(alpha) < 2 {
		return 0, base.Errf("%s: needs at least two components, got %d", name, len(alpha))
	}
	if len(x) != len(alpha) {
		return 0, base.Errf("%s: the concentration has %d components, the point %d", name, len(alpha), len(x))
	}
	total := 0.0
	for i, a := range alpha {
		if !(a > 0) || math.IsInf(a, 0) {
			return 0, base.Errf("%s: alpha[%d] must be finite and positive, got %g", name, i, a)
		}
		total += a
	}
	sum := 0.0
	for i, v := range x {
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return 0, base.Errf("%s: x[%d] must be finite, got %g", name, i, v)
		}
		if v < 0 {
			return 0, base.Errf("%s: x[%d] = %g lies outside the simplex", name, i, v)
		}
		sum += v
	}
	if math.Abs(sum-1) > 1e-9 {
		return 0, base.Errf("%s: the point must sum to 1, got %g", name, sum)
	}
	// ln B(α) = Σ lgamma(α_i) − lgamma(α₀).
	lb := -logGamma(total)
	for _, a := range alpha {
		lb += logGamma(a)
	}
	s := -lb
	for i, a := range alpha {
		switch {
		case x[i] == 0:
			if a < 1 {
				return math.Inf(1), nil
			}
			if a > 1 {
				return 0, nil
			}
		default:
			s += (a - 1) * math.Log(x[i])
		}
	}
	return math.Exp(s), nil
}

// DirichletMean returns the mean of the Dirichlet with concentration
// α: the normalised concentration α_i/α₀.
func DirichletMean(alpha []float64) ([]float64, error) {
	if len(alpha) < 2 {
		return nil, base.Errf("DirichletMean: needs at least two components, got %d", len(alpha))
	}
	total := 0.0
	for i, a := range alpha {
		if !(a > 0) || math.IsInf(a, 0) {
			return nil, base.Errf("DirichletMean: alpha[%d] must be finite and positive, got %g", i, a)
		}
		total += a
	}
	mean := make([]float64, len(alpha))
	for i, a := range alpha {
		mean[i] = a / total
	}
	return mean, nil
}

// DirichletMode returns the interior mode (α_i−1)/(α₀−k), which exists
// only when every concentration exceeds 1; any α_i ≤ 1 pushes the mode
// onto the boundary and is refused rather than answered with a vector
// that is not a mode.
func DirichletMode(alpha []float64) ([]float64, error) {
	if len(alpha) < 2 {
		return nil, base.Errf("DirichletMode: needs at least two components, got %d", len(alpha))
	}
	total := 0.0
	for i, a := range alpha {
		if !(a > 0) || math.IsInf(a, 0) {
			return nil, base.Errf("DirichletMode: alpha[%d] must be finite and positive, got %g", i, a)
		}
		if a <= 1 {
			return nil, base.Errf("DirichletMode: alpha[%d] = %g leaves no interior mode; every concentration must exceed 1", i, a)
		}
		total += a
	}
	den := total - float64(len(alpha))
	mode := make([]float64, len(alpha))
	for i, a := range alpha {
		mode[i] = (a - 1) / den
	}
	return mode, nil
}

// DirichletDraws returns n draws from the Dirichlet with concentration
// α, as an (n, k) array whose rows are the draws. Each row scales the
// k independent gamma(α_i, 1) draws of GammaDraws' generator to sum to
// one; the all-underflow row (every α_i far below the float64 floor)
// would otherwise divide by zero and falls back to the uniform row.
func DirichletDraws(g *core.Generator, n int, alpha []float64) (*core.Array, error) {
	const name = "DirichletDraws"
	if n < 1 {
		return nil, base.Errf("%s: n must be ≥ 1", name)
	}
	if len(alpha) < 2 {
		return nil, base.Errf("%s: needs at least two components, got %d", name, len(alpha))
	}
	for i, a := range alpha {
		if !(a > 0) || math.IsInf(a, 0) {
			return nil, base.Errf("%s: alpha[%d] must be finite and positive, got %g", name, i, a)
		}
	}
	k := len(alpha)
	flat := make([]float64, n*k)
	for r := range n {
		sum := 0.0
		row := flat[r*k : r*k+k]
		for c, a := range alpha {
			row[c] = gammaOne(g, a)
			sum += row[c]
		}
		if sum == 0 {
			// Every gamma draw underflowed: the uniform row is the
			// honest stand-in, an Inf row would poison the draw.
			for c := range row {
				row[c] = 1 / float64(k)
			}
			continue
		}
		for c := range row {
			row[c] /= sum
		}
	}
	return floatsToArray(flat, []int{n, k}), nil
}

// gammaOne draws one gamma(α, 1) variate, the single-draw form of the
// GammaDraws loop: Marsaglia-Tsang for α ≥ 1 and the boost with the
// exponential for α < 1, reusing the package squeeze.
func gammaOne(g *core.Generator, alpha float64) float64 {
	if alpha >= 1 {
		d := alpha - 1.0/3
		return gammaMarsagliaTsang(g, alpha, d, 1/math.Sqrt(9*d))
	}
	boost := alpha + 1
	d := boost - 1.0/3
	v := gammaMarsagliaTsang(g, boost, d, 1/math.Sqrt(9*d))
	return v * math.Pow(g.Unit(), 1/alpha)
}