// Package odometry estimates the pointing rotation of the telescope between two
// wide-angle frames using image-based phase correlation.
//
// It is deliberately NOT star/plate-solving based: phase correlation works on
// any textured scene (landscape, horizon, daytime sky, clouds, ...), which is
// exactly what the wide camera sees. For small telescope rotations the image
// content undergoes an almost pure translation (pan -> horizontal shift,
// tilt -> vertical shift), so the recovered image-plane shift maps directly to
// angular pan/tilt via the camera field of view.
package odometry

import (
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"math"
	"math/cmplx"
	"os"
)

// Result holds the motion estimated between two frames.
type Result struct {
	// PanPx / TiltPx are the image-plane translation (in resized pixels).
	PanPx  float64
	TiltPx float64
	// PanDeg / TiltDeg convert the pixel shift to degrees using the supplied FoV.
	PanDeg  float64
	TiltDeg float64
	// RotDeg is the estimated field rotation (degrees). Positive = scene rotated
	// counter-clockwise. Dominant when the telescope points near zenith on an
	// alt-az mount (azimuth rotation → field rotation, not translation).
	RotDeg float64
	// Confidence in [0,1].
	Confidence float64
	// Size is the FFT grid dimension actually used.
	Size int
}

func (r Result) String() string {
	return fmt.Sprintf(
		"pan=%+.3fpx (%+.4f°)  tilt=%+.3fpx (%+.4f°)  conf=%.3f",
		r.PanPx, r.PanDeg, r.TiltPx, r.TiltDeg, r.Confidence)
}

// DefaultSize is the default FFT grid (power of two). 256 is fast and accurate
// to ~0.1px; 512 improves sub-pixel resolution at ~4x the CPU cost.
const DefaultSize = 256

// DefaultFoVH/V are the documented DWARF Mini "wide" display FoV (degrees).
// NOTE: firmware-reported live FoV can differ widely (see analysis/DEVICE_MODELS.md),
// so these are only a starting point — calibrate with a known motor move.
const (
	DefaultFoVH = 8.0
	DefaultFoVV = 6.5
)

// Estimate computes the rotation between two decoded images.
//
// fovXDeg/fovYDeg are the camera horizontal/vertical field of view in degrees
// (they only scale the pixel shift into degrees; pass DefaultFoVH/V if unsure).
// size is the square FFT grid (must be a power of two, e.g. 256 or 512).
func Estimate(img1, img2 image.Image, fovXDeg, fovYDeg float64, size int) (Result, error) {
	if img1 == nil || img2 == nil {
		return Result{}, fmt.Errorf("odometry: nil image")
	}
	if !isPow2(size) {
		return Result{}, fmt.Errorf("odometry: size %d is not a power of two", size)
	}
	if size < 8 {
		return Result{}, fmt.Errorf("odometry: size %d too small", size)
	}
	if fovXDeg <= 0 || fovYDeg <= 0 {
		return Result{}, fmt.Errorf("odometry: fov must be positive (got %.3f x %.3f)", fovXDeg, fovYDeg)
	}

	g1 := toGrayDownscaled(img1, size, size)
	g2 := toGrayDownscaled(img2, size, size)

	// Pass 1: estimate rotation (log-polar phase correlation).
	rotDeg, confRot := EstimateRotation(g1, g2, size, size)

	// Pass 2: de-rotate g2, then estimate translation on de-rotated images.
	// This removes the rotation-induced aliasing that corrupts translation.
	var dx, dy, confTrans float64
	if absf(rotDeg) > 0.3 {
		g2Derot := rotateGray(g2, size, rotDeg)
		dx, dy, confTrans = phaseCorrelation(g1, g2Derot, size, size)
	} else {
		dx, dy, confTrans = phaseCorrelation(g1, g2, size, size)
	}

	// Use the higher-confidence metric as the overall confidence.
	conf := confTrans
	if confRot > conf {
		conf = confRot
	}

	r := Result{
		PanPx:      dx,
		TiltPx:     dy,
		PanDeg:     dx * fovXDeg / float64(size),
		TiltDeg:    dy * fovYDeg / float64(size),
		RotDeg:     rotDeg,
		Confidence: conf,
		Size:       size,
	}
	return r, nil
}

// EstimateFiles decodes two image files and runs Estimate.
func EstimateFiles(path1, path2 string, fovXDeg, fovYDeg float64, size int) (Result, error) {
	img1, err := decodeImage(path1)
	if err != nil {
		return Result{}, fmt.Errorf("decode %s: %w", path1, err)
	}
	img2, err := decodeImage(path2)
	if err != nil {
		return Result{}, fmt.Errorf("decode %s: %w", path2, err)
	}
	return Estimate(img1, img2, fovXDeg, fovYDeg, size)
}

func decodeImage(path string) (image.Image, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	img, _, err := image.Decode(f)
	return img, err
}

// phaseCorrelation returns the integer+subpixel shift (dx,dy) of the scene from
// g1 to g2, plus a confidence metric. The shift is expressed so that
// g2(x,y) ≈ g1(x-dx, y-dy), i.e. positive dx => content moved to the right.
func phaseCorrelation(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)

	// Normalized cross-power spectrum. conj(A)*B peaks at the shift that maps
	// image A (g1) onto image B (g2): positive dx = content moved right.
	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) // inverse transform -> correlation surface

	// Peak of the real part.
	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)

	// Wrap-around to signed shift.
	dx = float64(signedShift(px, cols))
	dy = float64(signedShift(py, rows))

	// Sub-pixel refinement via 3-point parabolic interpolation along each axis.
	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))

	// Confidence: peak height relative to the surface mean (clamped to [0,1]).
	denom := peak - mean
	if denom > 1 {
		denom = 1
	}
	if denom < 0 {
		denom = 0
	}
	conf = denom
	return dx, dy, conf
}

// windowedComplex converts a grayscale matrix to complex, subtracts the DC
// component and applies a 2-D Hann window to reduce spectral leakage at the
// image borders.
func windowedComplex(g [][]float64, rows, cols int) [][]complex128 {
	// Mean (DC) removal.
	dc := 0.0
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			dc += g[i][j]
		}
	}
	dc /= float64(rows * cols)

	m := make([][]complex128, rows)
	for i := 0; i < rows; i++ {
		m[i] = make([]complex128, cols)
		// Hann along rows.
		wy := 0.5 * (1 - math.Cos(2*math.Pi*float64(i)/float64(rows-1)))
		for j := 0; j < cols; j++ {
			wx := 0.5 * (1 - math.Cos(2*math.Pi*float64(j)/float64(cols-1)))
			m[i][j] = complex((g[i][j]-dc)*wx*wy, 0)
		}
	}
	return m
}

// at2 reads r[y][x] with horizontal (column) wrap-around.
func at2(r [][]complex128, y, x, cols int) float64 {
	x = ((x % cols) + cols) % cols
	if y < 0 {
		y += len(r)
	}
	if y >= len(r) {
		y -= len(r)
	}
	return real(r[y][x])
}

// parabola returns the sub-pixel offset of the true peak given the values at
// (left, center, right). Standard 3-point parabolic interpolation.
func parabola(left, center, right float64) float64 {
	denom := left - 2*center + right
	if math.Abs(denom) < 1e-12 {
		return 0
	}
	return 0.5 * (left - right) / denom
}

// signedShift converts a correlation peak index in [0,n) to a signed shift in
// [-n/2, n/2).
func signedShift(idx, n int) int {
	if idx >= n/2 {
		return idx - n
	}
	return idx
}

func isPow2(n int) bool {
	return n > 0 && (n&(n-1)) == 0
}

// toGrayDownscaled converts an image to a rows×cols grayscale matrix using
// area-averaging (box filter) downscaling. This preserves the full field of
// view so pixel shifts scale directly to the original FoV.
func toGrayDownscaled(img image.Image, rows, cols int) [][]float64 {
	b := img.Bounds()
	srcW := b.Dx()
	srcH := b.Dy()

	out := make([][]float64, rows)
	for i := range out {
		out[i] = make([]float64, cols)
	}

	xRatio := float64(srcW) / float64(cols)
	yRatio := float64(srcH) / float64(rows)

	for oy := 0; oy < rows; oy++ {
		y0 := b.Min.Y + int(math.Floor(float64(oy)*yRatio))
		y1 := b.Min.Y + int(math.Floor(float64(oy+1)*yRatio))
		if y1 <= y0 {
			y1 = y0 + 1
		}
		for ox := 0; ox < cols; ox++ {
			x0 := b.Min.X + int(math.Floor(float64(ox)*xRatio))
			x1 := b.Min.X + int(math.Floor(float64(ox+1)*xRatio))
			if x1 <= x0 {
				x1 = x0 + 1
			}
			sum := 0.0
			cnt := 0
			for y := y0; y < y1 && y < b.Max.Y; y++ {
				for x := x0; x < x1 && x < b.Max.X; x++ {
					r, g, bl, _ := img.At(x, y).RGBA()
					// 16-bit RGBA -> 8-bit luminance (Rec. 601).
					lum := 0.299*float64(r) + 0.587*float64(g) + 0.114*float64(bl)
					sum += lum / 257.0
					cnt++
				}
			}
			if cnt > 0 {
				out[oy][ox] = sum / float64(cnt)
			}
		}
	}
	return out
}