package imageproc

import (
	"fmt"

	"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 {
	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)
	if err != nil {
		panic(err)
	}

	return coeffs.RawMatrix().Data
}

// 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
}

func main() {
	// Input data
	x := []float64{1.0, 2.0, 3.0, 4.0, 5.0}
	y := []float64{1.0, 1.5, 2.5, 3.5, 5.0}

	// Fit a second degree polynomial (quadratic)
	degree := 2
	coeffs := PolynomialFit(x, y, degree)

	// Print the coefficients
	fmt.Printf("Fitted polynomial coefficients: %v\n", coeffs)
	if degree == 2 {
		fmt.Printf("Fitted quadratic: y = %.4fx^2 + %.4fx + %.4f\n", coeffs[2], coeffs[1], coeffs[0])
	}
}