package mathkit

import (
	"fmt"
	"math"
)

// Mean returns the arithmetic mean of data.
func Mean(data []float64) (float64, error) {
	if len(data) == 0 {
		return 0, fmt.Errorf("mean: empty data set")
	}
	var sum float64
	for _, v := range data {
		sum += v
	}
	return sum / float64(len(data)), nil
}

// Variance returns the variance of data. population=true divides by
// n (the population variance); population=false (the default for ML
// work, where `data` is almost always a sample) divides by n-1
// (Bessel's correction, the unbiased sample-variance estimator).
func Variance(data []float64, population bool) (float64, error) {
	n := len(data)
	if n == 0 {
		return 0, fmt.Errorf("variance: empty data set")
	}
	if !population && n < 2 {
		return 0, fmt.Errorf("variance: sample variance needs at least 2 points (use population variance for a single point)")
	}
	mean, _ := Mean(data)
	var ss float64
	for _, v := range data {
		d := v - mean
		ss += d * d
	}
	denom := float64(n)
	if !population {
		denom = float64(n - 1)
	}
	return ss / denom, nil
}

// StdDev returns the standard deviation (sqrt of Variance).
func StdDev(data []float64, population bool) (float64, error) {
	v, err := Variance(data, population)
	if err != nil {
		return 0, err
	}
	return math.Sqrt(v), nil
}

// Median returns the median of data (average of the two middle
// values for an even-length slice). Does not mutate the input.
func Median(data []float64) (float64, error) {
	if len(data) == 0 {
		return 0, fmt.Errorf("median: empty data set")
	}
	sorted := append([]float64(nil), data...)
	insertionSort(sorted)
	n := len(sorted)
	if n%2 == 1 {
		return sorted[n/2], nil
	}
	return (sorted[n/2-1] + sorted[n/2]) / 2, nil
}

// insertionSort is O(n^2) but this package only ever sorts small
// stats samples, where it beats sort.Float64s' overhead and keeps
// the package allocation-free apart from the clone itself.
func insertionSort(a []float64) {
	for i := 1; i < len(a); i++ {
		key := a[i]
		j := i - 1
		for j >= 0 && a[j] > key {
			a[j+1] = a[j]
			j--
		}
		a[j+1] = key
	}
}

// Covariance returns the sample (or population) covariance of two
// equal-length data sets — the building block of both Correlation and
// a PCA covariance matrix.
func Covariance(x, y []float64, population bool) (float64, error) {
	if len(x) != len(y) {
		return 0, fmt.Errorf("covariance: data sets have different lengths (%d vs %d)", len(x), len(y))
	}
	n := len(x)
	if n == 0 {
		return 0, fmt.Errorf("covariance: empty data set")
	}
	if !population && n < 2 {
		return 0, fmt.Errorf("covariance: sample covariance needs at least 2 points")
	}
	mx, _ := Mean(x)
	my, _ := Mean(y)
	var sum float64
	for i := 0; i < n; i++ {
		sum += (x[i] - mx) * (y[i] - my)
	}
	denom := float64(n)
	if !population {
		denom = float64(n - 1)
	}
	return sum / denom, nil
}

// Correlation returns the Pearson correlation coefficient of x and y,
// in [-1, 1].
func Correlation(x, y []float64) (float64, error) {
	cov, err := Covariance(x, y, true)
	if err != nil {
		return 0, fmt.Errorf("correlation: %w", err)
	}
	sx, err := StdDev(x, true)
	if err != nil {
		return 0, fmt.Errorf("correlation: %w", err)
	}
	sy, err := StdDev(y, true)
	if err != nil {
		return 0, fmt.Errorf("correlation: %w", err)
	}
	if sx == 0 || sy == 0 {
		return 0, fmt.Errorf("correlation: a constant data set has undefined correlation")
	}
	return cov / (sx * sy), nil
}

// ZScoreNormalize returns (x - mean) / stddev for each point — the
// standardization every ML feature-scaling pipeline applies before
// training a distance- or gradient-sensitive model.
func ZScoreNormalize(data []float64) ([]float64, error) {
	mean, err := Mean(data)
	if err != nil {
		return nil, fmt.Errorf("z_score_normalize: %w", err)
	}
	std, err := StdDev(data, true)
	if err != nil {
		return nil, fmt.Errorf("z_score_normalize: %w", err)
	}
	if std == 0 {
		return nil, fmt.Errorf("z_score_normalize: a constant data set has zero standard deviation")
	}
	out := make([]float64, len(data))
	for i, v := range data {
		out[i] = (v - mean) / std
	}
	return out, nil
}

// MinMaxNormalize rescales data to [0, 1].
func MinMaxNormalize(data []float64) ([]float64, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("min_max_normalize: empty data set")
	}
	lo, hi := data[0], data[0]
	for _, v := range data {
		if v < lo {
			lo = v
		}
		if v > hi {
			hi = v
		}
	}
	if hi == lo {
		return nil, fmt.Errorf("min_max_normalize: a constant data set has zero range")
	}
	out := make([]float64, len(data))
	for i, v := range data {
		out[i] = (v - lo) / (hi - lo)
	}
	return out, nil
}

// NormalPDF is the probability density function of N(mean, stdDev^2)
// at x.
func NormalPDF(x, mean, stdDev float64) (float64, error) {
	if stdDev <= 0 {
		return 0, fmt.Errorf("normal_pdf: std_dev must be positive")
	}
	z := (x - mean) / stdDev
	return math.Exp(-0.5*z*z) / (stdDev * math.Sqrt(2*math.Pi)), nil
}

// NormalCDF is the cumulative distribution function of N(mean,
// stdDev^2) at x, computed via math.Erf (the standard closed form:
// Φ(x) = (1 + erf((x-μ)/(σ√2))) / 2).
func NormalCDF(x, mean, stdDev float64) (float64, error) {
	if stdDev <= 0 {
		return 0, fmt.Errorf("normal_cdf: std_dev must be positive")
	}
	return 0.5 * (1 + math.Erf((x-mean)/(stdDev*math.Sqrt2))), nil
}

// BinomialPMF returns P(X = k) for X ~ Binomial(n, p).
func BinomialPMF(k, n int, p float64) (float64, error) {
	if n < 0 || k < 0 || k > n {
		return 0, fmt.Errorf("binomial_pmf: require 0 <= k <= n, got k=%d n=%d", k, n)
	}
	if p < 0 || p > 1 {
		return 0, fmt.Errorf("binomial_pmf: p must be in [0, 1], got %g", p)
	}
	return binomialCoeff(n, k) * math.Pow(p, float64(k)) * math.Pow(1-p, float64(n-k)), nil
}

// binomialCoeff computes C(n, k) via the multiplicative formula,
// which stays in float64 range far longer than computing n!/(k!(n-k)!)
// from raw factorials would.
func binomialCoeff(n, k int) float64 {
	if k < 0 || k > n {
		return 0
	}
	if k > n-k {
		k = n - k
	}
	result := 1.0
	for i := 0; i < k; i++ {
		result *= float64(n-i) / float64(i+1)
	}
	return result
}