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

package linalg

import (
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// The generalised symmetric eigenproblem A·v = λ·B·v, the standard form
// of vibrating-system and covariance questions: the eigenvalues of the
// pencil (A, B) with B symmetric positive definite. The Cholesky route
// reduces it to the ordinary symmetric problem without ever forming
// B⁻¹A, whose asymmetry would square the conditioning.

// EigenGeneralised solves A·v = λ·B·v for a symmetric a and a
// symmetric positive definite b, both real n×n. b = L·Lᵀ turns the
// pencil into the standard symmetric problem for C = L⁻¹·A·L⁻ᵀ, which
// shares the eigenvalues; its ordinary eigenvectors y transform back as
// v = L⁻ᵀ·y, which lands them B-orthonormal (vᵀ·B·v = 1) for free.
// Values come back ascending in a 1-D array with the eigenvectors as
// the matching columns, the convention Eigen uses. A complex input, a
// size mismatch, or a b that fails its Cholesky factorisation is an
// error; a itself must be symmetric, which is not verified.
func EigenGeneralised(a, b *core.Array) (values, vectors *core.Array, err error) {
	const name = "EigenGeneralised"
	if a.Dtype() == core.Complex || b.Dtype() == core.Complex {
		return nil, nil, base.Errf("%s: complex pencils are not supported", name)
	}
	if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
		return nil, nil, base.Errf("%s: a must be a square 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	if b.NDim() != 2 || b.Shape()[0] != b.Shape()[1] {
		return nil, nil, base.Errf("%s: b must be a square 2-D matrix, got shape %s", name, base.ShapeText(b.Shape()))
	}
	n := a.Shape()[0]
	if b.Shape()[0] != n {
		return nil, nil, base.Errf("%s: size mismatch, a is %d×%d and b is %d×%d",
			name, n, n, b.Shape()[0], b.Shape()[1])
	}
	if n == 0 {
		return nil, nil, base.Errf("%s: zero-sized pencil", name)
	}
	l, err := Cholesky(b)
	if err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	lFlat := denseFloats(l, n, n)
	// solveSystem consumes its matrix in place, so each solve gets a
	// fresh copy of L's rows as views over one flat backing slice.
	freshRows := func() [][]float64 {
		back := make([]float64, n*n)
		copy(back, lFlat)
		rows := make([][]float64, n)
		for i := range n {
			rows[i] = back[i*n : (i+1)*n]
		}
		return rows
	}
	// X = L⁻¹·A, one column of a per right-hand side.
	aCols := make([][]float64, n)
	for j := range n {
		aCols[j] = make([]float64, n)
		for i := range n {
			aCols[j][i] = a.FloatAt(i*n + j)
		}
	}
	if _, err := base.SolveSystem(name, freshRows(), aCols); err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	// C = X·L⁻ᵀ, gathered by solving L·Z = Xᵀ and transposing.
	cMat := make([]float64, n*n)
	{
		xT := make([][]float64, n)
		for j := range n {
			xT[j] = make([]float64, n)
			for i := range n {
				xT[j][i] = aCols[i][j] // column j of Xᵀ is row j of X
			}
		}
		if _, err := base.SolveSystem(name, freshRows(), xT); err != nil {
			return nil, nil, base.Errf("%s: %w", name, err)
		}
		for i := range n {
			for j := range n {
				cMat[i*n+j] = xT[j][i]
			}
		}
	}
	// Rounding leaves C a hair off symmetric; the eigensolver wants the
	// exact form, so take the symmetric part.
	for i := range n {
		for j := i + 1; j < n; j++ {
			m := (cMat[i*n+j] + cMat[j*n+i]) / 2
			cMat[i*n+j] = m
			cMat[j*n+i] = m
		}
	}
	cArr := floatsToArray(cMat, []int{n, n})
	values, yArr, err := Eigen(cArr)
	if err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	// V = L⁻ᵀ·Y: each eigenvector column solves Lᵀ·v = y.
	ltBack := make([]float64, n*n)
	lt := make([][]float64, n)
	for i := range n {
		lt[i] = ltBack[i*n : (i+1)*n]
		for j := range n {
			lt[i][j] = lFlat[j*n+i]
		}
	}
	yCols := make([][]float64, n)
	for j := range n {
		yCols[j] = make([]float64, n)
		for i := range n {
			yCols[j][i] = yArr.FloatAt(i*n + j)
		}
	}
	if _, err := base.SolveSystem(name, lt, yCols); err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	vMat := make([]float64, n*n)
	for j := range n {
		for i := range n {
			vMat[i*n+j] = yCols[j][i]
		}
	}
	return values, floatsToArray(vMat, []int{n, n}), nil
}