package odometry

import (
	"math"
	"math/cmplx"
)

// EstimateRotation estimates the rotation angle (in degrees) between two
// grayscale matrices using the log-polar phase correlation method:
//
//  1. Compute the magnitude spectra |F1|, |F2| (rotation in spatial domain =
//     rotation in frequency domain).
//  2. Apply a high-pass filter to suppress the DC-dominated center.
//  3. Resample to log-polar coordinates: rotation becomes a shift along the
//     angular axis, scale becomes a shift along the log-radius axis.
//  4. Phase-correlate in log-polar space → the peak along the angular axis
//     gives the rotation angle.
//
// Returns the rotation angle in degrees (positive = counter-clockwise rotation
// of the scene from g1 to g2), and a confidence in [0,1].
func EstimateRotation(g1, g2 [][]float64, rows, cols int) (angleDeg float64, conf float64) {
	// 1. FFT both images.
	a := windowedComplex(g1, rows, cols)
	b := windowedComplex(g2, rows, cols)
	fft2(a, false)
	fft2(b, false)

	// 2. Magnitude spectra with high-pass filtering.
	magA := make([][]float64, rows)
	magB := make([][]float64, rows)
	cy, cx := float64(rows)/2, float64(cols)/2
	maxR := math.Min(cy, cx)
	for i := 0; i < rows; i++ {
		magA[i] = make([]float64, cols)
		magB[i] = make([]float64, cols)
		for j := 0; j < cols; j++ {
			// Shift zero-frequency to center (fftshift).
			si := i
			sj := j
			if i < rows/2 {
				si = i + rows/2
			} else {
				si = i - rows/2
			}
			if j < cols/2 {
				sj = j + cols/2
			} else {
				sj = j - cols/2
			}
			ma := cmplx.Abs(a[si][sj])
			mb := cmplx.Abs(b[si][sj])
			// High-pass: suppress frequencies near DC.
			dist := math.Sqrt(float64((float64(i)-cy)*(float64(i)-cy)) + float64((float64(j)-cx)*(float64(j)-cx)))
			hp := 1.0 - math.Exp(-(dist*dist)/(2*(maxR*0.1)*(maxR*0.1)))
			magA[i][j] = ma * hp
			magB[i][j] = mb * hp
		}
	}

	// 3. Resample magnitude spectra to log-polar coordinates.
	nAngles := 256   // angular resolution
	nRadii := 128    // radial resolution
	lpA := logPolarResample(magA, rows, cols, nAngles, nRadii, maxR)
	lpB := logPolarResample(magB, rows, cols, nAngles, nRadii, maxR)

	// 4. Phase-correlate in log-polar space.
	dx, dy, confRaw := phaseCorrelationFloat(lpA, lpB, nAngles, nRadii)

	// dy is the shift along the angular axis → rotation angle.
	// Each row in lpA/lpB corresponds to 360/nAngles degrees.
	angleDeg = -float64(dy) * 360.0 / float64(nAngles)
	// Wrap to [-180, 180)
	for angleDeg < -180 {
		angleDeg += 360
	}
	for angleDeg >= 180 {
		angleDeg -= 360
	}

	_ = dx // log-radius shift (scale change) — not used for rotation-only estimation
	conf = confRaw
	return angleDeg, conf
}

// logPolarResample converts a Cartesian matrix to log-polar coordinates.
// The output has nAngles rows (angular samples) and nRadii columns (log-radius
// samples). This transforms a rotation in Cartesian space into a cyclic shift
// along the angular (row) axis.
func logPolarResample(mag [][]float64, rows, cols, nAngles, nRadii int, maxR float64) [][]float64 {
	out := make([][]float64, nAngles)
	cy, cx := float64(rows)/2, float64(cols)/2

	logMinR := math.Log(2.0)
	logMaxR := math.Log(maxR)
	rStep := (logMaxR - logMinR) / float64(nRadii-1)

	for a := 0; a < nAngles; a++ {
		theta := 2 * math.Pi * float64(a) / float64(nAngles)
		out[a] = make([]float64, nRadii)
		for r := 0; r < nRadii; r++ {
			radius := math.Exp(logMinR + float64(r)*rStep)
			y := cy + radius*math.Sin(theta)
			x := cx + radius*math.Cos(theta)
			// Bilinear interpolation.
			y0 := int(math.Floor(y))
			x0 := int(math.Floor(x))
			fy := y - float64(y0)
			fx := x - float64(x0)
			y0 = clampInt(y0, 0, rows-1)
			x0 = clampInt(x0, 0, cols-1)
			y1 := clampInt(y0+1, 0, rows-1)
			x1 := clampInt(x0+1, 0, cols-1)
			v := (1-fy)*(1-fx)*mag[y0][x0] +
				(1-fy)*fx*mag[y0][x1] +
				fy*(1-fx)*mag[y1][x0] +
				fy*fx*mag[y1][x1]
			out[a][r] = v
		}
	}
	return out
}

// phaseCorrelationFloat is a float-based phase correlation (same algorithm as
// phaseCorrelation but operates on float64 matrices instead of complex, using
// its own internal FFT buffers).
func phaseCorrelationFloat(g1, g2 [][]float64, rows, cols int) (dx, dy, conf float64) {
	a := windowedComplex(g1, rows, cols)
	b := windowedComplex(g2, rows, cols)

	fft2(a, false)
	fft2(b, false)

	r := make([][]complex128, rows)
	for i := 0; i < rows; i++ {
		r[i] = make([]complex128, cols)
		for j := 0; j < cols; j++ {
			cross := cmplx.Conj(a[i][j]) * b[i][j]
			mag := cmplx.Abs(cross)
			if mag > 1e-12 {
				r[i][j] = cross / complex(mag, 0)
			}
		}
	}

	fft2(r, true)

	px, py, peak := 0, 0, math.Inf(-1)
	mean := 0.0
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			v := real(r[i][j])
			mean += v
			if v > peak {
				peak = v
				px, py = j, i
			}
		}
	}
	mean /= float64(rows * cols)

	dx = float64(signedShift(px, cols))
	dy = float64(signedShift(py, rows))

	// Sub-pixel refinement.
	dx += parabola(at2(r, py, px-1, cols), peak, at2(r, py, px+1, cols))
	dy += parabola(at2(r, py-1, px, rows), peak, at2(r, py+1, px, rows))

	denom := peak - mean
	if denom > 1 {
		denom = 1
	}
	if denom < 0 {
		denom = 0
	}
	conf = denom
	return dx, dy, conf
}

func clampInt(v, lo, hi int) int {
	if v < lo {
		return lo
	}
	if v > hi {
		return hi
	}
	return v
}