package gnarlring

import "math"

// lllReduce reduces a lattice basis in-place using the LLL algorithm
// with parameter delta (typically 0.99). The input basis is a slice
// of d columns, each a vector of length d.
func lllReduce(basis [][]int64, delta float64) {
	d := len(basis)
	if d == 0 {
		return
	}

	mu := make([][]float64, d)
	for i := range mu {
		mu[i] = make([]float64, d)
	}
	bStar := make([][]float64, d)
	bStarNorm := make([]float64, d)

	computeGS(basis, mu, bStar, bStarNorm)

	k := 1
	for k < d {
		for j := k - 1; j >= 0; j-- {
			r := int(math.Round(mu[k][j]))
			if r != 0 {
				addScaled(basis[k], basis[j], -r)
				updateMuRow(mu, basis, bStar, bStarNorm, k, j)
				updateBStar(bStar, mu, basis, bStarNorm, k)
			}
		}

		if bStarNorm[k] >= (delta-mu[k][k-1]*mu[k][k-1])*bStarNorm[k-1] {
			k++
		} else {
			basis[k], basis[k-1] = basis[k-1], basis[k]
			for i := k - 1; i < d; i++ {
				updateBStar(bStar, mu, basis, bStarNorm, i)
				updateMuCol(mu, basis, bStar, bStarNorm, i, d)
			}
			if k > 1 {
				k--
			}
		}
	}
}

func computeGS(basis [][]int64, mu [][]float64, bStar [][]float64, bStarNorm []float64) {
	d := len(basis)
	for i := 0; i < d; i++ {
		bStar[i] = make([]float64, d)
		for j := 0; j < d; j++ {
			bStar[i][j] = float64(basis[i][j])
		}
	}
	for i := 0; i < d; i++ {
		for j := 0; j < i; j++ {
			mu[i][j] = dotFloat(bStar[i], bStar[j]) / bStarNorm[j]
			for k := 0; k < d; k++ {
				bStar[i][k] -= mu[i][j] * bStar[j][k]
			}
		}
		bStarNorm[i] = dotFloat(bStar[i], bStar[i])
	}
}

func updateBStar(bStar [][]float64, mu [][]float64, basis [][]int64, bStarNorm []float64, i int) {
	d := len(basis)
	for k := 0; k < d; k++ {
		bStar[i][k] = float64(basis[i][k])
	}
	for j := 0; j < i; j++ {
		mu[i][j] = dotFloat(bStar[i], bStar[j]) / bStarNorm[j]
		for k := 0; k < d; k++ {
			bStar[i][k] -= mu[i][j] * bStar[j][k]
		}
	}
	bStarNorm[i] = dotFloat(bStar[i], bStar[i])
}

func updateMuCol(mu [][]float64, basis [][]int64, bStar [][]float64, bStarNorm []float64, i, d int) {
	for t := i + 1; t < d; t++ {
		if bStarNorm[i] > 1e-20 {
			mu[t][i] = int64DotFloat(basis[t], bStar[i]) / bStarNorm[i]
		} else {
			mu[t][i] = 0
		}
	}
}

func updateMuRow(mu [][]float64, basis [][]int64, bStar [][]float64, bStarNorm []float64, k, j int) {
	for jj := 0; jj <= j; jj++ {
		if bStarNorm[jj] > 1e-20 {
			mu[k][jj] = int64DotFloat(basis[k], bStar[jj]) / bStarNorm[jj]
		} else {
			mu[k][jj] = 0
		}
	}
}

func dotFloat(a, b []float64) float64 {
	var sum float64
	for i := range a {
		sum += a[i] * b[i]
	}
	return sum
}

func int64DotFloat(a []int64, b []float64) float64 {
	var sum float64
	for i := range a {
		sum += float64(a[i]) * b[i]
	}
	return sum
}

func addScaled(dst, src []int64, r int) {
	for i := range dst {
		dst[i] += int64(r) * src[i]
	}
}