// 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"
)

// ExponentialDraws returns n draws from the exponential distribution
// with the given rate (mean 1/rate), by inverse CDF.
func ExponentialDraws(g *core.Generator, n int, rate float64) (*core.Array, error) {
	if n < 1 {
		return nil, base.Errf("ExponentialDraws: n must be ≥ 1")
	}
	if !(rate > 0) {
		return nil, base.Errf("ExponentialDraws: rate must be positive, got %v", rate)
	}
	out := core.New(core.Float, []int{n}...)
	for i := range n {
		u := 1 - g.Unit()
		out.RawFloats()[i] = -math.Log(u) / rate
	}
	return out, nil
}

// GammaDraws returns n draws from the gamma distribution with shape
// α > 0 and rate β > 0, by Marsaglia-Tsang for α ≥ 1 and by boosting
// with the exponential for α < 1.
func GammaDraws(g *core.Generator, n int, alpha, beta float64) (*core.Array, error) {
	if n < 1 {
		return nil, base.Errf("GammaDraws: n must be ≥ 1")
	}
	if !(alpha > 0 && beta > 0) {
		return nil, base.Errf("GammaDraws: shape and rate must be positive")
	}
	out := core.New(core.Float, []int{n}...)
	d := alpha - 1.0/3
	c := 1 / math.Sqrt(9*d)
	for i := range n {
		if alpha >= 1 {
			out.RawFloats()[i] = gammaMarsagliaTsang(g, alpha, d, c) / beta
			continue
		}
		// α < 1: boost to α + 1 and scale by a uniform^(1/α).
		boost := alpha + 1
		dd := boost - 1.0/3
		cc := 1 / math.Sqrt(9*dd)
		v := gammaMarsagliaTsang(g, boost, dd, cc)
		out.RawFloats()[i] = v * math.Pow(g.Unit(), 1/alpha) / beta
	}
	return out, nil
}

// gammaMarsagliaTsang draws one gamma(α, 1) for α ≥ 1 by the
// Marsaglia-Tsang squeeze: a normal draw shapes the cube root of the
// scale, an exponential-tilted accept/reject polishes the tail.
func gammaMarsagliaTsang(g *core.Generator, alpha, d, c float64) float64 {
	for {
		x := g.NormalUnit()
		v := 1 + c*x
		if v <= 0 {
			continue
		}
		vvv := v * v * v
		u := g.Unit()
		if u < 1-0.0331*x*x*x*x {
			return d * vvv
		}
		if math.Log(u) < 0.5*x*x+d*(1-vvv+math.Log(vvv)) {
			return d * vvv
		}
	}
}

// ChiSquareDraws returns n draws from the chi-squared distribution
// with df degrees of freedom, the gamma(df/2, 2) distribution.
func ChiSquareDraws(g *core.Generator, n int, df int) (*core.Array, error) {
	if df < 1 {
		return nil, base.Errf("ChiSquareDraws: df must be ≥ 1")
	}
	return GammaDraws(g, n, float64(df)/2, 0.5)
}

// StudentTDraws returns n draws from Student's t with df degrees of
// freedom, as N(0,1)/√(χ²_df/df).
func StudentTDraws(g *core.Generator, n int, df int) (*core.Array, error) {
	if df < 1 {
		return nil, base.Errf("StudentTDraws: df must be ≥ 1")
	}
	chi, err := ChiSquareDraws(g, n, df)
	if err != nil {
		return nil, err
	}
	out := core.New(core.Float, []int{n}...)
	for i := range n {
		z := g.NormalUnit()
		den := math.Sqrt(chi.FloatAt(i) / float64(df))
		if den == 0 {
			// The χ² draw underflowed to exactly zero: the ratio is an
			// infinity carrying the numerator's sign, not an unsigned
			// one (and not the NaN a 0/0 numerator of zero would make).
			out.RawFloats()[i] = math.Copysign(math.Inf(1), z)
		} else {
			out.RawFloats()[i] = z / den
		}
	}
	return out, nil
}

// PoissonDraws returns n draws from the Poisson distribution with the
// given λ, by the Knuth multiplication method for λ < 30 and the
// normal approximation above.
func PoissonDraws(g *core.Generator, n int, lambda float64) (*core.Array, error) {
	if n < 1 {
		return nil, base.Errf("PoissonDraws: n must be ≥ 1")
	}
	if !(lambda >= 0) {
		return nil, base.Errf("PoissonDraws: λ must be ≥ 0")
	}
	out := core.New(core.Float, []int{n}...)
	L := math.Exp(-lambda)
	for i := range n {
		if lambda < 30 {
			// λ = 0 (or tiny enough that exp(−λ) rounds to 1) is the
			// degenerate distribution at 0: the multiplication loop
			// would never run and hand back k−1 = −1.
			if L >= 1 {
				out.RawFloats()[i] = 0
				continue
			}
			k := 0.0
			p := 1.0
			for p > L {
				k++
				p *= g.Unit()
			}
			out.RawFloats()[i] = k - 1
		} else {
			// Normal approximation for large λ.
			z := g.NormalUnit()
			out.RawFloats()[i] = max(0, math.Floor(lambda+math.Sqrt(lambda)*z+0.5))
		}
	}
	return out, nil
}

// BinomialDraws returns n draws from the binomial distribution with
// the given number of trials and success probability: each draw runs
// trials uniform comparisons against p and counts the successes, the
// exact per-trial loop. It is exact but costs O(trials) uniforms per
// draw, so keep trials modest.
func BinomialDraws(g *core.Generator, n int, trials int, p float64) (*core.Array, error) {
	if n < 1 {
		// The sibling draws all refuse n < 1; without the guard this one
		// returned a nil array with a nil error, which a caller that only
		// checks the error then dereferenced.
		return nil, base.Errf("BinomialDraws: n must be ≥ 1")
	}
	if trials < 1 {
		return nil, base.Errf("BinomialDraws: trials must be ≥ 1")
	}
	if !(p >= 0 && p <= 1) {
		return nil, base.Errf("BinomialDraws: p must be in [0, 1]")
	}
	out := core.New(core.Float, []int{n}...)
	for i := range n {
		count := 0.0
		for range trials {
			if g.Unit() < p {
				count++
			}
		}
		out.RawFloats()[i] = count
	}
	return out, nil
}