package edt

import "math"

type Bounds struct {
	X int
	Y int
}

type Coordinate struct {
	X int
	Y int
}

/*
Represents a Euclidean Distance Transform of dimensions Dimensions with the transform
stored in D and a lookup in L. The lookup represents the nearest occupied coordinate (x, y)
relative to the current coordinate.
*/
type EDT struct {
	Dimensions Bounds       //spatial dimensions
	D          []float64    //distance transform
	L          []Coordinate //transform lookup
	gx         []float64    //buffer for our vector function, per row
	gy         []float64    //buffer for our vector function, per column
}

func NewEDT(dimensions Bounds) *EDT {
	edt := &EDT{
		Dimensions: dimensions,
	}

	edt.InitializeD()
	return edt
}

func (edt *EDT) InitializeD() {

	n := edt.Dimensions.X * edt.Dimensions.Y
	if n <= 0 {
		return
	}

	edt.D = make([]float64, n)
	edt.L = make([]Coordinate, n)
	edt.gx = make([]float64, edt.Dimensions.X)
	edt.gy = make([]float64, edt.Dimensions.Y)

	for i := range n {
		edt.D[i] = math.Inf(+1)
	}
}

/*
Reinitializes EDT using occupancy values. Occupancy dimensions must match EDT.
*/
func (edt *EDT) AssignOccupancy(occupancy []bool) {

	n := edt.Dimensions.X * edt.Dimensions.Y
	if len(occupancy) != n {
		return
	}

	//assign occupancy values
	for i := range n {
		if occupancy[i] {
			edt.D[i] = 0
		} else {
			edt.D[i] = math.Inf(+1)
		}
	}
}

/*
EDT must already have been initialized with occupancy assigned
*/
func (edt *EDT) ComputeDistanceTransform() {
	if len(edt.D) == 0 {
		return
	}
	//compute edt for each row
	for j := range edt.Dimensions.Y {
		//prepare our row vector
		edt.ConstructGx(j)

		//compute distance transform for this row
		//rowD, rowL := edt.DistanceTransform1D(edt.gx)
		rowD, rowL := edt1D_with_labels(edt.gx)

		//update corresponding row
		edt.UpdateRow(j, rowD, rowL)
	}

	//compute edt for each column
	for i := range edt.Dimensions.X {
		//prepare our column vector
		edt.ConstructGy(i)

		//compute distance transform for this column
		//colD, colL := edt.DistanceTransform1D(edt.gy)
		colD, colL := edt1D_with_labels(edt.gy)

		//update corresponding column
		edt.UpdateCol(i, colD, colL)
	}
}

/*Computes EDT and assigns lookup coordinates for given vector g*/
func (edt *EDT) DistanceTransform1D(g []float64) (d []float64, l []int) {
	//perform 1d distance transform
	n := len(g)
	d = make([]float64, n)
	l = make([]int, n)

	v := make([]int, n)
	// envelope sites (seed indices)
	z := make([]float64, n+1) // changeover abscissae
	k := 0

	// top
	v[0] = 0
	z[0] = math.Inf(-1)
	z[1] = math.Inf(+1)

	// initialize first valid seed
	k = -1
	for q := 0; q < n; q++ {
		if g[q] == math.Inf(+1) {
			continue
		}

		// push q: pop while it starts before current segment
		var s float64
		for {
			if k < 0 {
				k = 0
				v[0] = q
				z[0] = math.Inf(-1)
				z[1] = math.Inf(+1)
				break
			}
			i := v[k]
			// intersection s(i,q)
			s = ((float64(q*q) + g[q]) - (float64(i*i) + g[i])) / float64(2*(q-i))
			if s <= z[k] {
				k--
				if k < 0 {
					continue
				} // force a push with -inf
			} else {
				k++
				v[k] = q
				z[k] = s
				z[k+1] = math.Inf(+1)
				break
			}
		}
	}

	// evaluate
	k = 0
	for x := 0; x < n; x++ {
		for z[k+1] <= float64(x) {
			k++
		}
		i := v[k]
		l[x] = i
		dx := float64(x - i)
		d[x] = dx*dx + g[i]
	}

	return d, l
}

func edt1D_with_labels(g []float64) ([]float64, []int) {
	n := len(g)
	d := make([]float64, n)
	l := make([]int, n)

	// --- 1. Gather valid seeds ---
	type site struct {
		i  int
		gi float64
	}
	seeds := make([]site, 0, n)
	for i := 0; i < n; i++ {
		if !math.IsInf(g[i], +1) { // it's a seed
			seeds = append(seeds, site{i, g[i]})
		}
	}

	// --- 2. Handle empty case directly ---
	if len(seeds) == 0 {
		for i := 0; i < n; i++ {
			d[i] = math.Inf(+1)
			l[i] = -1
		}
		return d, l
	}

	// --- 3. Lower-envelope construction (Felzenszwalb–Huttenlocher) ---
	v := make([]int, len(seeds))
	z := make([]float64, len(seeds)+1)
	k := 0
	v[0] = seeds[0].i
	z[0] = math.Inf(-1)
	z[1] = math.Inf(+1)

	for q := 1; q < len(seeds); q++ {
		i := seeds[q].i
		gi := seeds[q].gi
		for {
			j := v[k]
			gj := g[j]
			s := ((float64(i*i) + gi) - (float64(j*j) + gj)) / float64(2*(i-j))
			if s <= z[k] {
				k--
				if k < 0 {
					k = 0
					break
				}
				continue
			}
			k++
			v[k] = i
			z[k] = s
			z[k+1] = math.Inf(+1)
			break
		}
	}

	// --- 4. Evaluate ---
	k = 0
	for x := 0; x < n; x++ {
		for z[k+1] <= float64(x) {
			k++
		}
		i := v[k]
		l[x] = i
		dx := float64(x - i)
		d[x] = dx*dx + g[i]
	}

	return d, l
}

/*
Prepares row vector (gx) from specified column in D
*/
func (edt *EDT) ConstructGx(col int) {
	if col < edt.Dimensions.Y {
		//we can be very efficient due to slice referencing
		startIdx := col * edt.Dimensions.Y
		edt.gx = edt.D[startIdx : startIdx+edt.Dimensions.X]
	}
}

/*
Prepares column vector (gy) from specified row in D
*/
func (edt *EDT) ConstructGy(row int) {
	if row < edt.Dimensions.X {
		for j := range edt.Dimensions.Y {
			edt.gy[j] = edt.D[j*edt.Dimensions.X+row]
		}
	}
}

/*
Writes data and lookup values into EDT row specified by rowidx
*/
func (edt *EDT) UpdateRow(rowidx int, data []float64, lookup []int) {
	if rowidx < edt.Dimensions.Y {
		startidx := rowidx * edt.Dimensions.X
		for i := range edt.Dimensions.X {
			edt.D[startidx+i] = data[i]
			edt.L[startidx+i].X = lookup[i]
		}
	}
}

/*
Writes data and lookup values into EDT column specified by colidx
*/
func (edt *EDT) UpdateCol(colidx int, data []float64, lookup []int) {
	if colidx < edt.Dimensions.X {
		for j := range edt.Dimensions.Y {
			edt.D[j*edt.Dimensions.X+colidx] = data[j]
			edt.L[j*edt.Dimensions.X+colidx].X = edt.L[lookup[j]*edt.Dimensions.X+colidx].X
			edt.L[j*edt.Dimensions.X+colidx].Y = lookup[j]
		}
	}
}