// Copyright ©2017 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 mat

import (
	"errors"

	"gonum.org/v1/gonum/blas/blas64"
)

// HOGSVD is a type for creating and using the Higher Order Generalized Singular Value
// Decomposition (HOGSVD) of a set of matrices.
//
// The factorization is a linear transformation of the data sets from the given
// variable×sample spaces to reduced and diagonalized "eigenvariable"×"eigensample"
// spaces.
type HOGSVD struct {
	n int
	v *Dense
	b []Dense

	err error
}

// succFact returns whether the receiver contains a successful factorization.
func (gsvd *HOGSVD) succFact() bool {
	return gsvd.n != 0
}

// Factorize computes the higher order generalized singular value decomposition (HOGSVD)
// of the n input r_i×c column tall matrices in m. HOGSV extends the GSVD case from 2 to n
// input matrices.
//
//	M_0 = U_0 * Σ_0 * Vᵀ
//	M_1 = U_1 * Σ_1 * Vᵀ
//	.
//	.
//	.
//	M_{n-1} = U_{n-1} * Σ_{n-1} * Vᵀ
//
// where U_i are r_i×c matrices of singular vectors, Σ are c×c matrices singular values, and V
// is a c×c matrix of singular vectors.
//
// Factorize returns whether the decomposition succeeded. If the decomposition
// failed, routines that require a successful factorization will panic.
func (gsvd *HOGSVD) Factorize(m ...Matrix) (ok bool) {
	// Factorize performs the HOGSVD factorisation
	// essentially as described by Ponnapalli et al.
	// https://doi.org/10.1371/journal.pone.0028072

	if len(m) < 2 {
		panic("hogsvd: too few matrices")
	}
	gsvd.n = 0

	r, c := m[0].Dims()
	a := make([]Cholesky, len(m))
	var ts SymDense
	for i, d := range m {
		rd, cd := d.Dims()
		if rd < cd {
			gsvd.err = ErrShape
			return false
		}
		if rd > r {
			r = rd
		}
		if cd != c {
			panic(ErrShape)
		}
		ts.Reset()
		ts.SymOuterK(1, d.T())
		ok = a[i].Factorize(&ts)
		if !ok {
			gsvd.err = errors.New("hogsvd: cholesky decomposition failed")
			return false
		}
	}

	s := getDenseWorkspace(c, c, true)
	defer putDenseWorkspace(s)
	sij := getDenseWorkspace(c, c, false)
	defer putDenseWorkspace(sij)
	for i, ai := range a {
		for _, aj := range a[i+1:] {
			gsvd.err = ai.SolveCholTo(sij, &aj)
			if gsvd.err != nil {
				return false
			}
			s.Add(s, sij)

			gsvd.err = aj.SolveCholTo(sij, &ai)
			if gsvd.err != nil {
				return false
			}
			s.Add(s, sij)
		}
	}
	s.Scale(1/float64(len(m)*(len(m)-1)), s)

	var eig Eigen
	ok = eig.Factorize(s.T(), EigenRight)
	if !ok {
		gsvd.err = errors.New("hogsvd: eigen decomposition failed")
		return false
	}
	var vc CDense
	eig.VectorsTo(&vc)
	// vc is guaranteed to have real eigenvalues.
	rc, cc := vc.Dims()
	v := NewDense(rc, cc, nil)
	for i := 0; i < rc; i++ {
		for j := 0; j < cc; j++ {
			a := vc.At(i, j)
			v.set(i, j, real(a))
		}
	}
	// Rescale the columns of v by their Frobenius norms.
	// Work done in cv is reflected in v.
	var cv VecDense
	for j := 0; j < c; j++ {
		cv.ColViewOf(v, j)
		cv.ScaleVec(1/blas64.Nrm2(cv.mat), &cv)
	}

	b := make([]Dense, len(m))
	biT := getDenseWorkspace(c, r, false)
	defer putDenseWorkspace(biT)
	for i, d := range m {
		// All calls to reset will leave an emptied
		// matrix with capacity to store the result
		// without additional allocation.
		biT.Reset()
		gsvd.err = biT.Solve(v, d.T())
		if gsvd.err != nil {
			return false
		}
		b[i].CloneFrom(biT.T())
	}

	gsvd.n = len(m)
	gsvd.v = v
	gsvd.b = b
	return true
}

// Err returns the reason for a factorization failure.
func (gsvd *HOGSVD) Err() error {
	return gsvd.err
}

// Len returns the number of matrices that have been factorized. If Len returns
// zero, the factorization was not successful.
func (gsvd *HOGSVD) Len() int {
	return gsvd.n
}

// UTo extracts the matrix U_n from the singular value decomposition, storing
// the result in-place into dst. U_n is size r×c.
//
// If dst is empty, UTo will resize dst to be r×c. When dst is
// non-empty, UTo will panic if dst is not r×c. UTo will also
// panic if the receiver does not contain a successful factorization.
func (gsvd *HOGSVD) UTo(dst *Dense, n int) {
	if !gsvd.succFact() {
		panic(badFact)
	}
	if n < 0 || gsvd.n <= n {
		panic("hogsvd: invalid index")
	}
	r, c := gsvd.b[n].Dims()
	if dst.IsEmpty() {
		dst.ReuseAs(r, c)
	} else {
		r2, c2 := dst.Dims()
		if r != r2 || c != c2 {
			panic(ErrShape)
		}
	}
	dst.Copy(&gsvd.b[n])
	var v VecDense
	for j, f := range gsvd.Values(nil, n) {
		v.ColViewOf(dst, j)
		v.ScaleVec(1/f, &v)
	}
}

// Values returns the nth set of singular values of the factorized system.
// If the input slice is non-nil, the values will be stored in-place into the slice.
// In this case, the slice must have length c, and Values will panic with
// ErrSliceLengthMismatch otherwise. If the input slice is nil,
// a new slice of the appropriate length will be allocated and returned.
//
// Values will panic if the receiver does not contain a successful factorization.
func (gsvd *HOGSVD) Values(s []float64, n int) []float64 {
	if !gsvd.succFact() {
		panic(badFact)
	}
	if n < 0 || gsvd.n <= n {
		panic("hogsvd: invalid index")
	}

	_, c := gsvd.b[n].Dims()
	if s == nil {
		s = make([]float64, c)
	} else if len(s) != c {
		panic(ErrSliceLengthMismatch)
	}
	var v VecDense
	for j := 0; j < c; j++ {
		v.ColViewOf(&gsvd.b[n], j)
		s[j] = blas64.Nrm2(v.mat)
	}
	return s
}

// VTo extracts the matrix V from the singular value decomposition, storing
// the result in-place into dst. V is size c×c.
//
// If dst is empty, VTo will resize dst to be c×c. When dst is
// non-empty, VTo will panic if dst is not c×c. VTo will also
// panic if the receiver does not contain a successful factorization.
func (gsvd *HOGSVD) VTo(dst *Dense) {
	if !gsvd.succFact() {
		panic(badFact)
	}
	r, c := gsvd.v.Dims()
	if dst.IsEmpty() {
		dst.ReuseAs(r, c)
	} else {
		r2, c2 := dst.Dims()
		if r != r2 || c != c2 {
			panic(ErrShape)
		}
	}
	dst.Copy(gsvd.v)
}