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

// Polynomial roots through the companion matrix. The roots of a
// polynomial are exactly the eigenvalues of its companion matrix, so
// the general nonsymmetric eigensolver answers the question directly:
// no Aberth iteration, no bracketing, one direct construction and a
// free ride on EigenGeneral's shifted QR. (The companion matrix is
// not balanced; coefficients spread over many magnitudes condition
// the roots through the eigenvalue problem as it stands.)

// PolynomialRoots returns the roots of the polynomial whose
// coefficients are given in ascending power order, lowest power first,
// the same convention EvaluatePolynomial uses. The answer is a complex
// vector sorted descending by magnitude, as EigenGeneral orders its
// values. Trailing zero coefficients raise nothing: they are stripped
// before the companion matrix is built, so the degree is the true one.
// A nonzero constant has no roots and answers an empty vector; the
// zero polynomial has every point as a root and is an error, as are
// coefficients that are not a vector and an empty coefficient list.
func PolynomialRoots(coeffs *core.Array) (*core.Array, error) {
	const name = "PolynomialRoots"
	if coeffs.NDim() != 1 {
		return nil, base.Errf("%s: coefficients must be a vector, got shape %s",
			name, base.ShapeText(coeffs.Shape()))
	}
	n := coeffs.Len()
	if n == 0 {
		return nil, base.Errf("%s: the coefficient vector must not be empty", name)
	}
	c := make([]complex128, n)
	for i := range n {
		if coeffs.Dtype() == core.Complex {
			c[i] = coeffs.ComplexAt(i)
		} else {
			c[i] = complex(coeffs.FloatAt(i), 0)
		}
	}
	// Strip trailing zeros to reach the true degree.
	for n > 0 && c[n-1] == 0 {
		n--
	}
	if n == 0 {
		return nil, base.Errf("%s: the zero polynomial has every point as a root", name)
	}
	if n == 1 {
		return core.FromComplexes(nil, 0)
	}
	// The Frobenius companion of the monic polynomial: ones on the
	// subdiagonal, the negated scaled coefficients down the last
	// column. Its characteristic polynomial is p(x)/c_{n-1}, so its
	// eigenvalues are the roots.
	degree := n - 1
	companion := make([]complex128, degree*degree)
	for row := 1; row < degree; row++ {
		companion[row*degree+row-1] = 1
	}
	for k := range degree {
		companion[k*degree+degree-1] = -c[k] / c[degree]
	}
	values, _, err := EigenGeneral(fromComplexesMust(companion, degree, degree))
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	return values, nil
}

// fromComplexesMust wraps a construction that cannot fail: the value
// count always matches the two-dimensional shape. It is unexported on
// purpose: a library that panics on a caller's input is a defect, and
// this caller cannot fail.
func fromComplexesMust(vals []complex128, rows, cols int) *core.Array {
	a, _ := core.FromComplexes(vals, rows, cols)
	return a
}