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

package stats

import (
	"math"

	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"sourcedock.dev/petrbalvin/tensor/internal/engine"
)

// Clustering: k-means with k-means++ seeding and the Gaussian mixture
// fitted by expectation maximisation over the multivariate normal
// machinery of mvn.go. Both fits are deterministic for a given
// generator state: every random choice runs through the house
// generator, never through a global source.

// kmeansMaxIterations caps the Lloyd sweeps; a fit that has not
// settled by then is reported with Converged false rather than
// pretended to have converged.
const kmeansMaxIterations = 300

// kmeansTolerance is the convergence tolerance of the Lloyd loop: the
// sweep stops once the largest coordinate movement of any centre is at
// most kmeansTolerance scaled by the largest absolute coordinate in
// the data (floored at 1, so a degenerate all-zero cloud still
// converges on an absolute scale). The rule mirrors the fixed points
// of Lloyd's iteration: past the tolerance the objective moves below
// rounding, and the label assignment no longer changes.
const kmeansTolerance = 1e-10

// KMeansResult carries the fit of k-means over a sample.
type KMeansResult struct {
	// Centres are the k fitted centroids, one row of d coordinates
	// each.
	Centres [][]float64
	// Labels assigns every sample row its cluster, 0 to k-1.
	Labels []int
	// Inertia is the within-cluster sum of squared distances to the
	// centres, the objective the Lloyd loop minimises.
	Inertia float64
	// Iterations counts the Lloyd sweeps taken; Converged reports
	// whether the centre movement fell under kmeansTolerance before
	// the iteration budget ran out.
	Iterations int
	Converged  bool
}

// KMeans partitions the n rows of x (n rows, d columns) into k
// clusters by Lloyd's iteration, seeded with k-means++ over the house
// generator: the first centre is drawn uniformly among the samples and
// every later centre with probability proportional to the squared
// distance to the nearest centre already chosen, so a well-separated
// cluster cannot be left unseeded except against odds its separation
// sets. The fit is deterministic for a given generator state.
//
// Convergence: the sweeps stop once no centre moves by more than
// kmeansTolerance (scaled, as the constant documents) or after
// kmeansMaxIterations sweeps; Converged names which happened.
//
// Empty clusters: a sweep can strand a centre with no members. The
// stranded centre then moves onto the sample farthest from its
// current centre, the sample that the objective can improve on most;
// if several centres strand in one sweep, each takes the next
// farthest unused sample. The rule is deterministic and keeps every
// cluster live.
//
// A nil generator, a non-finite sample, k below 1 or k above n is an
// error.
func KMeans(g *core.Generator, x *core.Array, k int) (*KMeansResult, error) {
	const name = "KMeans"
	if g == nil {
		return nil, base.Errf("%s: the generator is nil", name)
	}
	data, n, d, err := clusterReadSample(name, x)
	if err != nil {
		return nil, err
	}
	if k < 1 {
		return nil, base.Errf("%s: k must be at least 1, got %d", name, k)
	}
	if k > n {
		return nil, base.Errf("%s: k = %d exceeds the %d samples", name, k, n)
	}
	seeds := kmeansSeeds(g, data, n, d, k)
	return kmeansLloyd(data, n, d, k, seeds)
}

// kmeansSeeds draws k initial centres by the k-means++ rule: the
// first uniformly, each later one with probability proportional to the
// squared distance to the nearest centre already drawn. The weighted
// draw walks the cumulative weights against one uniform, the same
// inverse-transform the package's discrete draws use.
func kmeansSeeds(g *core.Generator, data []float64, n, d, k int) [][]float64 {
	seeds := make([][]float64, k)
	first := min(int(g.Unit()*float64(n)), n-1)
	seeds[0] = append([]float64(nil), data[first*d:first*d+d]...)
	nearest := make([]float64, n)
	for i := range n {
		nearest[i] = sqDistance(data[i*d:i*d+d], seeds[0])
	}
	for c := 1; c < k; c++ {
		total := 0.0
		for i := range n {
			total += nearest[i]
		}
		// A total of zero means every sample sits exactly on a chosen
		// centre: nothing left to separate, and any further centre is
		// arbitrary. The position repeats, and the Lloyd loop's
		// empty-cluster rule resolves the tie.
		if total == 0 {
			seeds[c] = append([]float64(nil), data[0:d]...)
		} else {
			target := g.Unit() * total
			walk := 0.0
			pick := n - 1
			for i := range n {
				walk += nearest[i]
				if walk >= target {
					pick = i
					break
				}
			}
			seeds[c] = append([]float64(nil), data[pick*d:pick*d+d]...)
		}
		for i := range n {
			if dist := sqDistance(data[i*d:i*d+d], seeds[c]); dist < nearest[i] {
				nearest[i] = dist
			}
		}
	}
	return seeds
}

// kmeansChosen reports whether the point equals one of the centres
// already drawn, the guard of the degenerate all-coincident fallback.
func kmeansChosen(centres [][]float64, point []float64) bool {
	for _, c := range centres {
		same := true
		for j, v := range c {
			if v != point[j] {
				same = false
				break
			}
		}
		if same {
			return true
		}
	}
	return false
}

// sqDistance returns the squared Euclidean distance between two
// points of equal length.
func sqDistance(a, b []float64) float64 {
	total := 0.0
	for j, v := range a {
		diff := v - b[j]
		total += diff * diff
	}
	return total
}

// kmeansLloyd runs the assignment and update sweeps from the given
// centres to the documented tolerance, applying the empty-cluster
// rule on the way. It is the shared engine of the public entry point
// and of the seeding comparisons the tests make.
func kmeansLloyd(data []float64, n, d, k int, centres [][]float64) (*KMeansResult, error) {
	// The scale that makes the tolerance relative: the largest
	// absolute coordinate in the sample, floored at 1.
	spread := 1.0
	for _, v := range data {
		if s := math.Abs(v); s > spread {
			spread = s
		}
	}
	labels := make([]int, n)
	converged := false
	iterations := kmeansMaxIterations
	// The update sweep's accumulators live outside the iteration: a sweep
	// clears them instead of allocating, and the fold below then adds the
	// same rows into the same slots in the same order.
	counts := make([]int, k)
	sums := make([][]float64, k)
	for c := range k {
		sums[c] = make([]float64, d)
	}
	for iter := 1; iter <= kmeansMaxIterations; iter++ {
		for i := range n {
			best, bestDist := 0, math.Inf(1)
			for c := range k {
				if dist := sqDistance(data[i*d:i*d+d], centres[c]); dist < bestDist {
					best, bestDist = c, dist
				}
			}
			labels[i] = best
		}
		// The update sweep, with the empty-cluster rule: a centre
		// whose cluster stranded takes the sample farthest from its
		// own centre, skipping the samples earlier stranded centres
		// in the same sweep already claimed.
		clear(counts)
		for c := range k {
			clear(sums[c])
		}
		for i := range n {
			counts[labels[i]]++
			row := sums[labels[i]]
			for j := range d {
				row[j] += data[i*d+j]
			}
		}
		worst := 0.0
		for c := range k {
			switch {
			case counts[c] > 0:
				for j := range d {
					moved := sums[c][j] / float64(counts[c])
					if m := math.Abs(moved - centres[c][j]); m > worst {
						worst = m
					}
					centres[c][j] = moved
				}
			default:
				// The farthest sample from this centre: the search is
				// for a maximum, so the running best starts at -Inf.
				farthest, farDist := -1, math.Inf(-1)
				for i := range n {
					if labels[i] == c || kmeansChosen(centres, data[i*d:i*d+d]) {
						continue
					}
					if dist := sqDistance(data[i*d:i*d+d], centres[c]); dist > farDist {
						farthest, farDist = i, dist
					}
				}
				if farthest < 0 {
					return nil, base.Errf("KMeans: cluster %d cannot be repopulated: the sample holds fewer distinct points than k", c)
				}
				copy(centres[c], data[farthest*d:farthest*d+d])
				worst = math.Inf(1)
			}
		}
		if worst <= kmeansTolerance*spread {
			converged = true
			iterations = iter
			break
		}
	}
	return &KMeansResult{
		Centres:    centres,
		Labels:     labels,
		Inertia:    kmeansInertia(data, n, d, centres, labels),
		Iterations: iterations,
		Converged:  converged,
	}, nil
}

// kmeansInertia sums the squared distances of every sample to its
// assigned centre.
func kmeansInertia(data []float64, n, d int, centres [][]float64, labels []int) float64 {
	total := 0.0
	for i := range n {
		total += sqDistance(data[i*d:i*d+d], centres[labels[i]])
	}
	return total
}

// gmmMaxIterations caps the EM sweeps, with the same reporting
// convention kmeansMaxIterations sets.
const gmmMaxIterations = 200

// gmmTolerance is the EM convergence tolerance on the relative change
// of the log likelihood: the sweeps stop once it moves by less than
// the tolerance scaled by 1 + |log likelihood|, past which no
// parameter the M step can produce moves the objective materially.
const gmmTolerance = 1e-9

// gmmCovarianceFloor is the regularisation floor of the M step: after
// every update, gmmCovarianceFloor of the data's mean per-dimension
// variance is added to the covariance diagonal. A component that
// collapses onto a single sample computes a singular covariance and a
// Cholesky factor that does not exist; the floor keeps the factor
// alive at a width far below any sample's spread, so it regularises
// the algebra without moving any fit the data actually supports.
const gmmCovarianceFloor = 1e-6

// GaussianMixtureResult carries the fit of a Gaussian mixture.
type GaussianMixtureResult struct {
	// Components is the fitted component count.
	Components int
	// Weights are the mixture weights, in component order, summing to
	// 1.
	Weights []float64
	// Means holds one mean vector per component.
	Means [][]float64
	// Covariances holds one d-by-d covariance per component,
	// row-major, in component order.
	Covariances [][]float64
	// Responsibilities are the final E-step posteriors, n rows of
	// Components entries each, row-major: P(component j | sample i).
	Responsibilities []float64
	// LogLikelihood is the maximised observed-data log likelihood.
	LogLikelihood float64
	// BIC is -2·LogLikelihood + p·ln n at the fitted parameters,
	// p the count of free parameters.
	BIC float64
	// Iterations counts the EM sweeps; Converged reports whether the
	// log likelihood settled under gmmTolerance before the budget.
	Iterations int
	Converged  bool
	// BICGrid is filled by GaussianMixtureBIC only: the BIC of every
	// fit on the component grid 1..len(BICGrid), in grid order. The
	// plain fit leaves it nil.
	BICGrid []float64
}

// GaussianMixture fits a mixture of `components` multivariate normals
// over the rows of x by expectation maximisation, built on the MVN
// machinery of mvn.go: every component density goes through the
// Cholesky factor of its covariance.
//
// The E step runs in log space: each sample's responsibilities come
// from a log-sum-exp normalisation guarded by the row maximum, so a
// sample sitting far outside every component underflows to a clean 0/1
// split rather than to a NaN.
//
// The M step regularises every covariance by the documented
// gmmCovarianceFloor before the next factorisation.
//
// Initialisation is k-means over the same generator: the partitions
// seed the weights, means and covariances, so the fit is deterministic
// for a given seed and every component starts populated. A component
// whose total responsibility collapses below 1e-10 (the floor at work
// on a stray sample) keeps its previous parameters and a floored
// weight rather than dividing by zero; the renormalised weights keep
// the mixture a mixture.
//
// The sample must be rank 2 and finite, the component count at least 1
// and at most the sample size, and the sample at least two rows.
func GaussianMixture(g *core.Generator, x *core.Array, components int) (*GaussianMixtureResult, error) {
	const name = "GaussianMixture"
	if g == nil {
		return nil, base.Errf("%s: the generator is nil", name)
	}
	data, n, d, err := clusterReadSample(name, x)
	if err != nil {
		return nil, err
	}
	if n < 2 {
		return nil, base.Errf("%s: needs at least two samples, got %d", name, n)
	}
	if components < 1 {
		return nil, base.Errf("%s: the component count must be at least 1, got %d", name, components)
	}
	if components > n {
		return nil, base.Errf("%s: %d components exceed the %d samples", name, components, n)
	}
	weights, means, covs, serr := gmmSeed(g, data, n, d, components)
	if serr != nil {
		return nil, base.Errf("%s: %w", name, serr)
	}
	result, err := gmmEM(data, n, d, components, weights, means, covs)
	if err != nil {
		return nil, err
	}
	result.BIC = gaussianMixtureBIC(result.LogLikelihood, components, n, d)
	return result, nil
}

// GaussianMixtureBIC fits mixtures over the component grid 1 to
// maxComponents and keeps the one with the lowest BIC. The grid and
// the winner travel in the result: BICGrid holds every fit's BIC in
// grid order and Components names the selected count. Ties resolve to
// the smaller model, the parsimony the penalty exists to buy.
func GaussianMixtureBIC(g *core.Generator, x *core.Array, maxComponents int) (*GaussianMixtureResult, error) {
	const name = "GaussianMixtureBIC"
	if g == nil {
		return nil, base.Errf("%s: the generator is nil", name)
	}
	data, n, d, err := clusterReadSample(name, x)
	if err != nil {
		return nil, err
	}
	if n < 2 {
		return nil, base.Errf("%s: needs at least two samples, got %d", name, n)
	}
	if maxComponents < 1 {
		return nil, base.Errf("%s: the largest component count must be at least 1, got %d", name, maxComponents)
	}
	if maxComponents > n {
		return nil, base.Errf("%s: %d components exceed the %d samples", name, maxComponents, n)
	}
	grid := make([]float64, maxComponents)
	best := -1
	bestBIC := math.Inf(1)
	var bestResult *GaussianMixtureResult
	for k := 1; k <= maxComponents; k++ {
		weights, means, covs, serr := gmmSeed(g, data, n, d, k)
		if serr != nil {
			return nil, serr
		}
		result, err := gmmEM(data, n, d, k, weights, means, covs)
		if err != nil {
			return nil, base.Errf("%s: the fit with %d components failed (%w)", name, k, err)
		}
		grid[k-1] = gaussianMixtureBIC(result.LogLikelihood, k, n, d)
		// Strictly better only: a tie keeps the smaller model.
		if grid[k-1] < bestBIC {
			bestBIC = grid[k-1]
			best = k
			bestResult = result
		}
	}
	bestResult.BICGrid = grid
	bestResult.BIC = grid[best-1]
	return bestResult, nil
}

// gaussianMixtureBIC assembles -2·logL + p·ln n, with p the free
// parameters of a k-component mixture: k-1 weights, k·d mean entries
// and k·d(d+1)/2 covariance entries.
func gaussianMixtureBIC(logL float64, k, n, d int) float64 {
	p := (k - 1) + k*d + k*d*(d+1)/2
	return -2*logL + float64(p)*math.Log(float64(n))
}

// clusterReadSample validates x as a finite real rank-2 sample and
// returns it flattened row-major with its shape.
func clusterReadSample(name string, x *core.Array) ([]float64, int, int, error) {
	if x.NDim() != 2 {
		return nil, 0, 0, base.Errf("%s: the sample must be rank 2, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if x.Dtype() == core.Complex {
		return nil, 0, 0, base.Errf("%s: complex inputs are not supported", name)
	}
	if err := checkFinite(name, "the sample", x); err != nil {
		return nil, 0, 0, err
	}
	n, d := x.Shape()[0], x.Shape()[1]
	if d < 1 {
		return nil, 0, 0, base.Errf("%s: the sample needs at least one column", name)
	}
	data := make([]float64, n*d)
	if fs := rawFloats(x); fs != nil {
		copy(data, fs)
	} else {
		for i := range data {
			data[i] = x.FloatAt(i)
		}
	}
	return data, n, d, nil
}

// gmmSeed initialises the mixture from the k-means partitions of the
// same generator: each component takes its cluster's weight, mean and
// covariance, floored exactly as the M step floors, so the first E
// step factors live covariances even for a cluster that collapsed onto
// one sample.
func gmmSeed(g *core.Generator, data []float64, n, d, k int) ([]float64, [][]float64, [][]float64, error) {
	seeds := kmeansSeeds(g, data, n, d, k)
	fit, err := kmeansLloyd(data, n, d, k, seeds)
	if err != nil {
		// A sample with fewer distinct points than k drives the
		// empty-cluster rule out of answers: the caller turns this
		// into the entry point's own named refusal.
		return nil, nil, nil, err
	}
	meanVariance := clusterMeanVariance(data, n, d)
	weights := make([]float64, k)
	means := make([][]float64, k)
	covs := make([][]float64, k)
	counts := make([]int, k)
	for i := range n {
		counts[fit.Labels[i]]++
	}
	for c := range k {
		weights[c] = math.Max(float64(counts[c]), 1e-10)
		mean := make([]float64, d)
		if counts[c] > 0 {
			for i := range n {
				if fit.Labels[i] != c {
					continue
				}
				for j := range d {
					mean[j] += data[i*d+j]
				}
			}
			for j := range d {
				mean[j] /= float64(counts[c])
			}
		}
		cov := make([]float64, d*d)
		for i := range n {
			if fit.Labels[i] != c {
				continue
			}
			for a := range d {
				for b := range d {
					cov[a*d+b] += (data[i*d+a] - mean[a]) * (data[i*d+b] - mean[b])
				}
			}
		}
		den := math.Max(float64(counts[c])-1, 1)
		for a := range d {
			for b := range d {
				cov[a*d+b] /= den
			}
			cov[a*d+a] += gmmCovarianceFloor * meanVariance
		}
		weights[c] /= float64(n)
		means[c] = mean
		covs[c] = cov
	}
	// The floored weights renormalised: the mixture must sum to 1 even
	// when a component started empty and took the floor.
	total := 0.0
	for _, w := range weights {
		total += w
	}
	for c := range k {
		weights[c] /= total
	}
	return weights, means, covs, nil
}

// clusterMeanVariance returns the mean per-dimension variance of the
// sample, the scale the covariance floor is measured in.
func clusterMeanVariance(data []float64, n, d int) float64 {
	if n < 2 {
		return 1
	}
	total := 0.0
	for j := range d {
		mean := 0.0
		for i := range n {
			mean += data[i*d+j]
		}
		mean /= float64(n)
		v := 0.0
		for i := range n {
			diff := data[i*d+j] - mean
			v += diff * diff
		}
		total += v / float64(n-1)
	}
	return total / float64(d)
}

// gmmParallelMinPoints is the sample count one worker must carry before
// the expectation sweep splits across goroutines: a point costs the
// component count in density evaluations and another in exponentials,
// so a shorter chunk is cheaper on the calling goroutine than in a
// pool.
const gmmParallelMinPoints = 32

// gmmEM runs the expectation maximisation sweeps from the seeded
// parameters to the documented tolerance.
func gmmEM(data []float64, n, d, k int, weights []float64, means [][]float64, covs [][]float64) (*GaussianMixtureResult, error) {
	const name = "GaussianMixture"
	resp := make([]float64, n*k)
	logL := math.Inf(-1)
	converged := false
	iterations := gmmMaxIterations
	// The sweep's buffers live outside the sweep: one flat covariance
	// copy the factor step refills, one factor per component the
	// factorisation clears and refills, and one row of per-point
	// normalisers the log-likelihood fold reads, so a sweep allocates
	// none of them. The per-component mean and covariance slices the M
	// step accumulates into are the parameter storage itself: each is
	// cleared before its pass, so the accumulation sees the zero state
	// a fresh allocation carried and the result carries the last
	// sweep's values in the same slices.
	covVals := make([]float64, d*d)
	factors := make([][][]float64, k)
	for c := range k {
		rows := make([][]float64, d)
		for i := range d {
			rows[i] = make([]float64, d)
		}
		factors[c] = rows
		if len(means[c]) != d {
			means[c] = make([]float64, d)
		}
		if len(covs[c]) != d*d {
			covs[c] = make([]float64, d*d)
		}
	}
	consts := make([]float64, k)
	logWeights := make([]float64, k)
	counts := make([]float64, k)
	centred := make([]float64, d)
	logNorms := make([]float64, n)
	// The covariance floor's scale depends on the sample alone, and the
	// sample never moves, so it is measured once for the whole fit.
	meanVariance := clusterMeanVariance(data, n, d)
	for iter := 1; iter <= gmmMaxIterations; iter++ {
		// E step: the components are factored once per sweep, and each
		// factor contributes its log determinant and the density's
		// normalising constant as one value, read once per component
		// instead of once per point. The component's log weight is
		// likewise constant across the sweep. The sample loop that
		// follows then carries only the quadratic form.
		for c := range k {
			copy(covVals, covs[c])
			if err := mvnCholeskyFlatInto(name, covVals, factors[c], d); err != nil {
				return nil, base.Errf("%s: component %d failed to factor (%w)", name, c, err)
			}
			l := factors[c]
			logDet := 0.0
			for i := range d {
				logDet += math.Log(l[i][i])
			}
			consts[c] = -0.5*float64(d)*math.Log(2*math.Pi) - logDet
			logWeights[c] = math.Log(weights[c])
		}
		// A point owns its responsibility row and its own normaliser
		// entry and nothing else, so the sample loop splits across
		// goroutines with no lock and no shared scratch: the worker's two
		// buffers are overwritten per point.
		engine.ParallelMin(n, gmmParallelMinPoints, func(start, end int) {
			logps := make([]float64, k)
			solve := make([]float64, d)
			for i := start; i < end; i++ {
				point := data[i*d : i*d+d]
				for c := range k {
					logps[c] = logWeights[c] + mvnLogDensitySolve(point, means[c], factors[c], d, consts[c], solve)
				}
				rowMax := math.Inf(-1)
				for _, lp := range logps {
					if lp > rowMax {
						rowMax = lp
					}
				}
				// A floored weight can park a component at -Inf; the
				// max-guarded sum survives it and the row still
				// normalises. logps is reused for the exponentials once
				// the row maximum has served: the posterior is the
				// quotient of the stored value and the row total, so
				// one exponential per component serves the row and the
				// divisor is the total in [1, k] instead of an exponent
				// compounded with the logarithm's rounding.
				total := 0.0
				for c := range k {
					e := math.Exp(logps[c] - rowMax)
					logps[c] = e
					total += e
				}
				logNorm := rowMax + math.Log(total)
				for c := range k {
					resp[i*k+c] = logps[c] / total
				}
				logNorms[i] = logNorm
			}
		})
		// The observed-data log likelihood folds the normalisers in
		// ascending point order on the calling goroutine: the chain is
		// the one the serial sweep accumulated, so the split above never
		// moves a bit of it.
		nextLogL := 0.0
		for _, logNorm := range logNorms {
			nextLogL += logNorm
		}
		move := nextLogL - logL
		logL = nextLogL
		if math.Abs(move) <= gmmTolerance*(1+math.Abs(logL)) {
			converged = true
			iterations = iter
			break
		}
		// M step: responsibilities to weights, means and floored
		// covariances. A component whose total responsibility
		// collapses under 1e-10 keeps its previous parameters under a
		// floored weight: dividing by zero would poison the sweep,
		// and the floor lets the next E step give the component
		// another chance.
		clear(counts)
		for i := range n {
			for c := range k {
				counts[c] += resp[i*k+c]
			}
		}
		totalWeight := 0.0
		for c := range k {
			if counts[c] < 1e-10 {
				weights[c] = 1e-10
			} else {
				weights[c] = counts[c] / float64(n)
			}
			totalWeight += weights[c]
		}
		for c := range k {
			weights[c] /= totalWeight
		}
		for c := range k {
			if counts[c] < 1e-10 {
				continue
			}
			mean := means[c]
			clear(mean)
			for i := range n {
				r := resp[i*k+c]
				for j := range d {
					mean[j] += r * data[i*d+j]
				}
			}
			for j := range d {
				mean[j] /= counts[c]
			}
			cov := covs[c]
			clear(cov)
			// The centred row is built once per sample and read by every
			// entry the covariance accumulates: the subtraction is the one
			// the product below performed per entry, and the remaining
			// factorisation of the term, (r·ca)·cb added to the entry,
			// keeps its own order and shape.
			for i := range n {
				r := resp[i*k+c]
				point := data[i*d : i*d+d]
				for a := range d {
					centred[a] = point[a] - mean[a]
				}
				for a := range d {
					scaled := r * centred[a]
					row := cov[a*d : a*d+d]
					for b, cb := range centred {
						row[b] += scaled * cb
					}
				}
			}
			for a := range d {
				for b := range d {
					cov[a*d+b] /= counts[c]
				}
				cov[a*d+a] += gmmCovarianceFloor * meanVariance
			}
		}
	}
	return &GaussianMixtureResult{
		Components:       k,
		Weights:          weights,
		Means:            means,
		Covariances:      covs,
		Responsibilities: resp,
		LogLikelihood:    logL,
		Iterations:       iterations,
		Converged:        converged,
	}, nil
}

// mvnLogDensitySolve evaluates the log density of one point under a
// component already factored, with the caller's normalising constant
// (the log determinant already folded in) and its own forward-solve
// scratch. Every entry of solve is written before it is read, so the
// scratch carries no state across calls; the constant is the identical
// expression the caller would otherwise recompute from the same factor.
func mvnLogDensitySolve(point, mean []float64, l [][]float64, d int, normConst float64, solve []float64) float64 {
	for i := range d {
		total := point[i] - mean[i]
		for j := range i {
			total -= l[i][j] * solve[j]
		}
		solve[i] = total / l[i][i]
	}
	quad := 0.0
	for i := range d {
		quad += solve[i] * solve[i]
	}
	return normConst - 0.5*quad
}