package lrls

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

func recursive(u, y []float64) ([][]float64, []float64, error) {
	dl := len(u)
	if len(y) < dl {
		dl = len(y)
	}

	// number of params.
	p := 4
	e := make([]float64, dl)

	covMs, err := covMatrices(dl, p)
	if err != nil {
		return nil, nil, err
	}

	theta, err := newTheta(dl, p)
	if err != nil {
		return nil, nil, err
	}

	phi, err := newPhi(dl, p)
	if err != nil {
		return nil, nil, err
	}

	// the guts.
	// in matlab:
	// phi(:, k) = [-y(k-1) -y(k-2) u(k-1) u(k-2)]';
	// e(k) = y(k) - phi(:, k)' * thm(:, k-1);
	// C(:,:,k) = C(:,:,k-1) - (C(:,:,k-1)*phi(:,k)*phi(:,k)'*C(:,:,k-1))/(1+phi(:,k)'*C(:,:,k-1)*phi(:,k));
	// thm(:,k) = thm(:,k-1) + C(:,:,k)*phi(:,k)*e(k);
	for k := 2; k <= dl-1; k++ {
		// phi[k] = []float64{-1.0 * y[k-1], -1 * y[k-2], -1 * u[k-1], -1 * u[k-2]}
		// phi[k] = []float64{y[k-1], y[k-2], u[k-1], u[k-2]}
		phi[k] = []float64{-1 * (y[k-1]), -1 * (y[k-2]), u[k-1], u[k-2]}
		phiK := phi[k]
		th := theta[k-1]

		// phiT[k].theta[k-1].
		phiThetaDotProd := phiThetaDotProd(p, phiK, th)

		// get current error.
		currentError := y[k] - phiThetaDotProd

		// save current error.
		e[k] = currentError

		// next CM:
		// cm(k) = cm(k-1) - [cm(k-1).phi(k).phiT(k).cm(k-1) / 1 + phiT(k).cm(k-1).phi(k)]
		// next theta:
		// theta(k) = theta(k-1) + cm(k).phi(k).e(k)

		cmPrev := &mat.Dense{}
		cmPrev.CloneFrom(covMs[k-1])

		cmK := nextCM(p, cmPrev, phiK)

		covMs[k] = cmK

		tmp := &mat.Dense{}
		phiVec := mat.NewVecDense(p, phiK)

		tmp.Mul(cmK, phiVec)

		// tmp.Scale(currentError, cmK) // -> this errs with dimens mismatch.
		for i := range tmp.RawMatrix().Data {
			tmp.RawMatrix().Data[i] *= currentError
		}

		tmpV := mat.NewVecDense(p, tmp.RawMatrix().Data)

		thVec := mat.NewVecDense(p, th)
		nuTh := mat.NewVecDense(p, nil)

		nuTh.AddVec(thVec, tmpV)

		theta[k] = nuTh.RawVector().Data
	}

	return theta, e, nil
}

func phiThetaDotProd(p int, phi, th []float64) float64 {
	var res float64

	for i := 0; i < p; i++ {
		res += phi[i] * th[i]
	}

	return res
}

// nextCM calculates the value of the next Covariance Matrix.
// cm(k) = cm(k-1) - [ cm(k-1).phi(k).phiT(k).cm(k-1) / 1 + phiT(k).cm(k-1).phi(k) ]
// matlab:
// C(k) = C(k-1) - (C(:,:,k-1)*phi(:,k)*phi(:,k)'*C(:,:,k-1))/(1+phi(:,k)'*C(:,:,k-1)*phi(:,k)).
func nextCM(p int, cmPrev *mat.Dense, phi []float64) *mat.Dense {
	phiVec := mat.NewVecDense(p, phi)
	phiVecTransposed := phiVec.TVec()
	cmTmp := &mat.Dense{}

	// r, c := cmPrev.Dims()
	// log.Printf("cmPrev dims: %d, %d", r, c)
	// r, c = phiVec.Dims()
	// log.Printf("phiVec dims: %d, %d", r, c)

	cmTmp.Mul(phiVec, phiVecTransposed)

	cmTmp.Mul(cmTmp, cmPrev)

	cmTmp.Mul(cmPrev, cmTmp)

	// -----

	cmTmp2 := &mat.Dense{}

	cmTmp2.Mul(phiVecTransposed, cmPrev)

	vecTmp := mat.NewVecDense(p, cmTmp2.RawMatrix().Data)

	denominator := 1 + doVector(
		vecTmp.RawVector().Data,
		phiVec.RawVector().Data,
	)

	for i := range cmTmp.RawMatrix().Data {
		cmTmp.RawMatrix().Data[i] /= denominator
	}

	cm := &mat.Dense{}

	cm.Sub(cmPrev, cmTmp)
	cm = subMat(cmPrev, cmTmp)

	return cm
}

// doVector returns a scalar product of column vector `cv` and regular vector
// `rv`.
func doVector(cv, rv []float64) float64 {
	var res float64

	for i := range cv {
		res += cv[i] * rv[i]
	}

	return res
}

func subMat(m1, m2 *mat.Dense) *mat.Dense {
	r, c := m1.Dims()
	m := mat.NewDense(r, c, nil)

	for i := range m1.RawMatrix().Data {
		m.RawMatrix().Data[i] = m1.RawMatrix().Data[i] - m2.RawMatrix().Data[i]
	}

	return m
}

func TransposeTheta(theta [][]float64) [][]float64 {
	thLen := len(theta)
	elemLen := len(theta[0])
	nuTheta := make([][]float64, elemLen)

	for i := range theta[0] {
		p := make([]float64, thLen)

		for j := range theta {
			p[j] = theta[j][i]
		}

		nuTheta[i] = p
	}

	return nuTheta
}