package utils

import (
	"math"

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

type PolynomialInterpolator struct {
	Degree int
	Coeffs []float64
}

func (p *PolynomialInterpolator) Fit(x, y []float64) error {
	if p.Degree == 0 {
		p.Degree = len(x) - 1
	}

	degree := p.Degree
	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, math.Pow(x[i], float64(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)
	p.Coeffs = coeffs.RawMatrix().Data
	return err
}

func (p PolynomialInterpolator) Predict(x float64) float64 {
	var y float64
	for i, coeff := range p.Coeffs {
		y += coeff * math.Pow(x, float64(i))
	}
	return y
}

func InterpPolynomial(x []float64, y []float64, xq float64) float64 {
	if len(x) != len(y) {
		return 0.0
	}

	start, end := FindClosestSubset(x, xq, 4)
	interp := &PolynomialInterpolator{Degree: 3}
	interp.Fit(x[start:end+1], y[start:end+1])
	return interp.Predict(xq)
}