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

// The multivariate normal: densities through a Cholesky solve and
// draws through the same triangular factor, the two operations any
// Bayesian or Monte Carlo workflow reaches for.

// mvnSymmetryEps is the relative tolerance the covariance's mirror
// check allows. A covariance assembled as A·Aᵀ can differ from its
// mirror by an ulp of rounding; a matrix that is genuinely asymmetric
// differs by far more than this, and reading only its lower triangle
// would describe a different distribution from the one handed in.
const mvnSymmetryEps = 1e-12

// mvnCholesky factors a symmetric positive-definite matrix into the
// lower triangular L with A = L·Lᵀ; a non-positive pivot names the row
// in the error, and a mirror pair that disagrees beyond a relative
// mvnSymmetryEps names the entry, because only the lower triangle is
// read.
func mvnCholesky(name string, cov *core.Array, d int) ([][]float64, error) {
	// The covariance is read once into a flat local copy: the symmetry
	// sweep and the O(d³) factorisation then index a plain slice. The
	// copy carries the values FloatAt returns, so every difference,
	// product and pivot below keeps its exact bits.
	covVals := make([]float64, d*d)
	if fs := rawFloats(cov); fs != nil {
		copy(covVals, fs)
	} else {
		for i := range covVals {
			covVals[i] = cov.FloatAt(i)
		}
	}
	return mvnCholeskyFlat(name, covVals, d)
}

// mvnCholeskyFlat factors a matrix already held flat row-major, the
// shared body of the array entry point and of a caller that assembles
// the matrix in place. The slice is read, never written.
func mvnCholeskyFlat(name string, covVals []float64, d int) ([][]float64, error) {
	l := make([][]float64, d)
	for i := range d {
		l[i] = make([]float64, d)
	}
	if err := mvnCholeskyFlatInto(name, covVals, l, d); err != nil {
		return nil, err
	}
	return l, nil
}

// mvnCholeskyFlatInto is mvnCholeskyFlat on a caller-owned d×d
// destination, for a sweep that factors one covariance after another:
// the rows are cleared and refilled, and every entry the factorisation
// reads is one it has already written in the same pass. The arithmetic
// is mvnCholeskyFlat's own, unchanged.
func mvnCholeskyFlatInto(name string, covVals []float64, l [][]float64, d int) error {
	for i := range d {
		for j := range i {
			lo, hi := covVals[i*d+j], covVals[j*d+i]
			if math.Abs(lo-hi) > mvnSymmetryEps*math.Max(math.Abs(lo), math.Abs(hi)) {
				return base.Errf("%s: the covariance is not symmetric at (%d, %d): %g against %g",
					name, i+1, j+1, lo, hi)
			}
		}
	}
	for i := range d {
		row := l[i]
		clear(row)
		for j := range i + 1 {
			total := covVals[i*d+j]
			for k := range j {
				total -= l[i][k] * l[j][k]
			}
			if i == j {
				if total <= 0 {
					return base.Errf("%s: the covariance is not positive definite at row %d", name, i+1)
				}
				row[j] = math.Sqrt(total)
			} else {
				row[j] = total / l[j][j]
			}
		}
	}
	return nil
}

// MultivariateNormalLogDensity evaluates the density of the
// d-dimensional normal N(mean, cov) at one point x, all three rank-1
// and cov symmetric positive definite, through the Cholesky factor:
// the log determinant is twice the sum of the factor's diagonal and
// the quadratic form the squared norm of the forward solve. Every
// entry of all three inputs must be finite, and the covariance must
// mirror itself: a non-finite value or an asymmetric pair is refused
// by name rather than answered with a NaN or with the density of a
// different distribution.
func MultivariateNormalLogDensity(mean, cov, x *core.Array) (float64, error) {
	const name = "MultivariateNormalLogDensity"
	if mean.NDim() != 1 || x.NDim() != 1 || cov.NDim() != 2 {
		return 0, base.Errf("%s: mean and x must be rank 1 and cov rank 2", name)
	}
	if mean.Dtype() == core.Complex || cov.Dtype() == core.Complex || x.Dtype() == core.Complex {
		return 0, base.Errf("%s: complex inputs are not supported", name)
	}
	d := mean.Len()
	if x.Len() != d || cov.Shape()[0] != d || cov.Shape()[1] != d {
		return 0, base.Errf("%s: mean holds %d entries, x %d, cov %s", name, d, x.Len(), base.ShapeText(cov.Shape()))
	}
	if err := checkFinite(name, "the mean", mean); err != nil {
		return 0, err
	}
	if err := checkFinite(name, "the covariance", cov); err != nil {
		return 0, err
	}
	if err := checkFinite(name, "the point", x); err != nil {
		return 0, err
	}
	l, err := mvnCholesky(name, cov, d)
	if err != nil {
		return 0, err
	}
	diff := make([]float64, d)
	solve := make([]float64, d)
	logDet := 0.0
	for i := range d {
		diff[i] = x.FloatAt(i) - mean.FloatAt(i)
		logDet += math.Log(l[i][i])
		total := diff[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 -0.5*float64(d)*math.Log(2*math.Pi) - logDet - 0.5*quad, nil
}

// MultivariateNormalDraws returns n draws from N(mean, cov) as an
// (n × d) array, one draw per row: standard normals through the
// generator, coloured by the Cholesky factor of the covariance. The
// draws are deterministic for a given generator state. The mean and
// the covariance must be finite and the covariance symmetric, the
// same contract MultivariateNormalLogDensity enforces.
func MultivariateNormalDraws(g *core.Generator, n int, mean, cov *core.Array) (*core.Array, error) {
	const name = "MultivariateNormalDraws"
	if g == nil {
		return nil, base.Errf("%s: the generator is nil", name)
	}
	if n < 1 {
		return nil, base.Errf("%s: the count must be at least 1, got %d", name, n)
	}
	if mean.NDim() != 1 || cov.NDim() != 2 {
		return nil, base.Errf("%s: mean must be rank 1 and cov rank 2", name)
	}
	if mean.Dtype() == core.Complex || cov.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex inputs are not supported", name)
	}
	d := mean.Len()
	if cov.Shape()[0] != d || cov.Shape()[1] != d {
		return nil, base.Errf("%s: mean holds %d entries but cov is %s", name, d, base.ShapeText(cov.Shape()))
	}
	if err := checkFinite(name, "the mean", mean); err != nil {
		return nil, err
	}
	if err := checkFinite(name, "the covariance", cov); err != nil {
		return nil, err
	}
	l, err := mvnCholesky(name, cov, d)
	if err != nil {
		return nil, err
	}
	out := core.New(core.Float, n, d)
	vals := out.RawFloats()
	z := make([]float64, d)
	// The mean is read once: every draw reuses the same values the
	// accessor walk returned.
	meanVals := make([]float64, d)
	if fs := rawFloats(mean); fs != nil {
		copy(meanVals, fs)
	} else {
		for i := range meanVals {
			meanVals[i] = mean.FloatAt(i)
		}
	}
	for r := range n {
		// One standard normal vector per draw, coloured by L: the
		// shared z is what makes the off-diagonal covariance appear.
		for j := range d {
			z[j] = g.NormalUnit()
		}
		for i := range d {
			total := meanVals[i]
			for j := range i + 1 {
				total += l[i][j] * z[j]
			}
			vals[r*d+i] = total
		}
	}
	return out, nil
}