// Copyright ©2016 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package distmat

import (
	"math"
	"math/rand/v2"
	"sync"

	"gonum.org/v1/gonum/mat"
	"gonum.org/v1/gonum/mathext"
	"gonum.org/v1/gonum/stat/distuv"
)

// Wishart is a distribution over d×d positive symmetric definite matrices. It
// is parametrized by a scalar degrees of freedom parameter ν and a d×d positive
// definite matrix V.
//
// The Wishart PDF is given by
//
//	p(X) = [|X|^((ν-d-1)/2) * exp(-tr(V^-1 * X)/2)] / [2^(ν*d/2) * |V|^(ν/2) * Γ_d(ν/2)]
//
// where X is a d×d PSD matrix, ν > d-1, |·| denotes the determinant, tr is the
// trace and Γ_d is the multivariate gamma function.
//
// See https://en.wikipedia.org/wiki/Wishart_distribution for more information.
type Wishart struct {
	nu  float64
	src rand.Source

	dim     int
	cholv   mat.Cholesky
	logdetv float64
	upper   mat.TriDense

	once sync.Once
	v    *mat.SymDense // only stored if needed
}

// NewWishart returns a new Wishart distribution with the given shape matrix and
// degrees of freedom parameter. NewWishart returns whether the creation was
// successful.
//
// NewWishart panics if nu <= d - 1 where d is the order of v.
func NewWishart(v mat.Symmetric, nu float64, src rand.Source) (*Wishart, bool) {
	dim := v.SymmetricDim()
	if nu <= float64(dim-1) {
		panic("wishart: nu must be greater than dim-1")
	}
	var chol mat.Cholesky
	ok := chol.Factorize(v)
	if !ok {
		return nil, false
	}

	var u mat.TriDense
	chol.UTo(&u)

	w := &Wishart{
		nu:  nu,
		src: src,

		dim:     dim,
		cholv:   chol,
		logdetv: chol.LogDet(),
		upper:   u,
	}
	return w, true
}

// MeanSymTo calculates the mean matrix of the distribution in and stores it in dst.
// If dst is empty, it is resized to be an d×d symmetric matrix where d is the order
// of the receiver. When dst is non-empty, MeanSymTo panics if dst is not d×d.
func (w *Wishart) MeanSymTo(dst *mat.SymDense) {
	if dst.IsEmpty() {
		dst.ReuseAsSym(w.dim)
	} else if dst.SymmetricDim() != w.dim {
		panic(badDim)
	}
	w.setV()
	dst.CopySym(w.v)
	dst.ScaleSym(w.nu, dst)
}

// ProbSym returns the probability of the symmetric matrix x. If x is not positive
// definite (the Cholesky decomposition fails), it has 0 probability.
func (w *Wishart) ProbSym(x mat.Symmetric) float64 {
	return math.Exp(w.LogProbSym(x))
}

// LogProbSym returns the log of the probability of the input symmetric matrix.
//
// LogProbSym returns -∞ if the input matrix is not positive definite (the Cholesky
// decomposition fails).
func (w *Wishart) LogProbSym(x mat.Symmetric) float64 {
	dim := x.SymmetricDim()
	if dim != w.dim {
		panic(badDim)
	}
	var chol mat.Cholesky
	ok := chol.Factorize(x)
	if !ok {
		return math.Inf(-1)
	}
	return w.logProbSymChol(&chol)
}

// LogProbSymChol returns the log of the probability of the input symmetric matrix
// given its Cholesky decomposition.
func (w *Wishart) LogProbSymChol(cholX *mat.Cholesky) float64 {
	dim := cholX.SymmetricDim()
	if dim != w.dim {
		panic(badDim)
	}
	return w.logProbSymChol(cholX)
}

func (w *Wishart) logProbSymChol(cholX *mat.Cholesky) float64 {
	// The PDF is
	//  p(X) = [|X|^((ν-d-1)/2) * exp(-tr(V^-1 * X)/2)] / [2^(ν*d/2) * |V|^(ν/2) * Γ_d(ν/2)]
	// The LogPDF is thus
	//  (ν-d-1)/2 * log(|X|) - tr(V^-1 * X)/2  - (ν*d/2)*log(2) - ν/2 * log(|V|) - log(Γ_d(ν/2))
	logdetx := cholX.LogDet()

	// Compute tr(V^-1 * X), using the fact that X = Uᵀ * U.
	var u mat.TriDense
	cholX.UTo(&u)

	var vinvx mat.Dense
	err := w.cholv.SolveTo(&vinvx, u.T())
	if err != nil {
		return math.Inf(-1)
	}
	vinvx.Mul(&vinvx, &u)
	tr := mat.Trace(&vinvx)

	fnu := float64(w.nu)
	fdim := float64(w.dim)

	return 0.5*((fnu-fdim-1)*logdetx-tr-fnu*fdim*math.Ln2-fnu*w.logdetv) - mathext.MvLgamma(0.5*fnu, w.dim)
}

// RandSymTo generates a random symmetric matrix from the distribution.
// If dst is empty, it is resized to be an d×d symmetric matrix where d is the order
// of the receiver. When dst is non-empty, RandSymTo panics if dst is not d×d.
func (w *Wishart) RandSymTo(dst *mat.SymDense) {
	var c mat.Cholesky
	w.RandCholTo(&c)
	c.ToSym(dst)
}

// RandCholTo generates the Cholesky decomposition of a random matrix from the distribution.
// If dst is empty, it is resized to be an d×d symmetric matrix where d is the order
// of the receiver. When dst is non-empty, RandCholTo panics if dst is not d×d.
func (w *Wishart) RandCholTo(dst *mat.Cholesky) {
	// TODO(btracey): Modify the code if the underlying data from dst is exposed
	// to avoid the dim^2 allocation here.

	// Use the Bartlett Decomposition, which says that
	//  X ~ L A Aᵀ Lᵀ
	// Where A is a lower triangular matrix in which the diagonal of A is
	// generated from the square roots of χ^2 random variables, and the
	// off-diagonals are generated from standard normal variables.
	// The above gives the cholesky decomposition of X, where L_x = L A.
	//
	// mat works with the upper triagular decomposition, so we would like to do
	// the same. We can instead say that
	//  U_x = L_xᵀ = (L * A)ᵀ = Aᵀ * Lᵀ = Aᵀ * U
	// Instead, generate Aᵀ, by using the procedure above, except as an upper
	// triangular matrix.
	norm := distuv.Normal{
		Mu:    0,
		Sigma: 1,
		Src:   w.src,
	}

	t := mat.NewTriDense(w.dim, mat.Upper, nil)
	for i := 0; i < w.dim; i++ {
		v := distuv.ChiSquared{
			K:   w.nu - float64(i),
			Src: w.src,
		}.Rand()
		t.SetTri(i, i, math.Sqrt(v))
	}
	for i := 0; i < w.dim; i++ {
		for j := i + 1; j < w.dim; j++ {
			t.SetTri(i, j, norm.Rand())
		}
	}

	t.MulTri(t, &w.upper)
	dst.SetFromU(t)
}

// setV computes and stores the covariance matrix of the distribution.
func (w *Wishart) setV() {
	w.once.Do(func() {
		w.v = mat.NewSymDense(w.dim, nil)
		w.cholv.ToSym(w.v)
	})
}