package producer

import (
	"gonum.org/v1/gonum/mat"
)

// PolynomialFit fits a polynomial of given degree to the data points (x, y).
func PolynomialFit(x, y []float64, degree int) ([]float64, error) {
	n := len(x)
	// Create the Vandermonde matrix
	vander := mat.NewDense(n, degree+1, nil)
	for i := 0; i < n; i++ {
		for j := 0; j <= degree; j++ {
			vander.Set(i, j, pow(x[i], j))
		}
	}

	// Create the right-hand side vector
	yVec := mat.NewVecDense(n, y)

	// Solve the least squares problem
	var qr mat.QR
	qr.Factorize(vander)
	coeffs := mat.NewDense(degree+1, 1, nil)
	err := qr.SolveTo(coeffs, false, yVec)
	return coeffs.RawMatrix().Data, err
}

// pow is a helper function to calculate power of a number.
func pow(x float64, n int) float64 {
	if n == 0 {
		return 1
	}
	res := x
	for i := 1; i < n; i++ {
		res *= x
	}
	return res
}