gen2d/backend/pkg/splitsprite/splitsprite.go

// Package splitsprite provides PNG sprite sheet splitting utilities.
//
// Pipeline: white/green background removal → projection-based gap detection →
// split into tiles → filter out low-fill tiles → trim transparent edges.
package splitsprite

import (
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"math"
	"sort"
)

// Options configures the sprite sheet splitting pipeline.
type Options struct {
	// WhiteBg enables white background removal.
	WhiteBg bool
	// WhiteThreshold is the max distance from pure white (0–255, default 40).
	WhiteThreshold uint8

	// GreenScreen enables green background removal.
	GreenScreen bool
	// GreenTolerance controls how aggressively green pixels are removed (0–1, default 0.2).
	GreenTolerance float64

	// GridRows / GridCols enable fixed-grid splitting (overrides projection detection).
	// When >0, the image is divided equally into Rows×Cols cells.
	GridRows int
	GridCols int
	// GridPadding is the gap between cells in pixels (default 2).
	GridPadding int

	// GapThreshold is the max fraction of non-transparent pixels a row/column
	// can have to be considered a gap (0–1, default 0.03).
	GapThreshold float64
	// MinGapWidth is the minimum width in pixels a gap must have.
	MinGapWidth int

	// MinFillRatio is the minimum fraction of non-transparent pixels a tile
	// must have to be kept (0–1, default 0.3).
	MinFillRatio float64

	// Trim removes transparent borders from output tiles.
	Trim bool
	// CenterAlign centers content across all frames so characters stay in place.
	// All output frames get the same dimensions with content centered.
	CenterAlign bool
	// OutW / OutH specify the output tile size (0 = keep original).
	OutW, OutH int
}

// DefaultOptions returns sensible default splitting options (white background mode).
func DefaultOptions() *Options {
	return &Options{
		WhiteBg:        true,
		WhiteThreshold: 40,
		GapThreshold:   0.03,
		MinGapWidth:    2,
		MinFillRatio:   0.14,
		Trim:           true,
		CenterAlign:    true,
	}
}

// DefaultGreenOptions returns options tuned for green screen sprite sheets.
func DefaultGreenOptions() *Options {
	return &Options{
		GreenScreen:    true,
		GreenTolerance: 0.2,
		GapThreshold:   0.03,
		MinGapWidth:    2,
		MinFillRatio:   0.14,
		Trim:           true,
		CenterAlign:    true,
	}
}

// Process splits a sprite sheet image into individual cleaned tile images.
// It runs the full pipeline: background removal → split → trim → resize.
// When GridRows/GridCols > 0, fixed-grid splitting is used instead of projection detection.
func Process(img image.Image, opts *Options) ([]image.Image, error) {
	if opts == nil {
		opts = DefaultOptions()
	}

	src := toRGBA(img)

	if opts.WhiteBg {
		src = removeWhiteBg(src, opts.WhiteThreshold)
	} else if opts.GreenScreen {
		src = removeGreenScreen(src, opts.GreenTolerance)
	}

	var tiles []tile
	if opts.GridRows > 0 && opts.GridCols > 0 {
		tiles = fixedGridSplit(src, opts.GridRows, opts.GridCols, opts.GridPadding, opts.MinFillRatio)
	} else {
		tiles = projectionSplit(src, opts.GapThreshold, opts.MinGapWidth, opts.MinFillRatio)
	}
	if len(tiles) == 0 {
		return nil, fmt.Errorf("no tiles detected — try lowering GapThreshold or setting GridRows/GridCols")
	}

	results := make([]image.Image, len(tiles))
	for i, t := range tiles {
		sub := image.NewRGBA(image.Rect(0, 0, t.w, t.h))
		draw.Draw(sub, sub.Bounds(), src, image.Point{t.x, t.y}, draw.Src)

		if opts.Trim {
			sub = trimAlpha(sub)
		}
		if opts.OutW > 0 && opts.OutH > 0 {
			sub = resize(sub, opts.OutW, opts.OutH)
		}
		results[i] = sub
	}

	if opts.CenterAlign && len(results) > 1 {
		results = alignCenter(results)
	}

	return results, nil
}

// RemoveWhiteBg removes near-white background pixels, making them transparent.
func RemoveWhiteBg(img image.Image, threshold uint8) image.Image {
	return removeWhiteBg(toRGBA(img), threshold)
}

// CenterFrames centers the content of each frame within a uniform canvas so
// characters stay in place across frames.
func CenterFrames(frames []image.Image) []image.Image {
	return alignCenter(frames)
}

// alignCenter aligns all frames to a uniform canvas with a fixed reference point.
// Uses bottom-center alignment so characters share a common ground plane across frames,
// preventing drift/jitter in animation playback.
func alignCenter(frames []image.Image) []image.Image {
	type contentBox struct {
		minX, minY, maxX, maxY int
	}
	boxes := make([]contentBox, len(frames))
	maxCW, maxCH := 0, 0

	for i, f := range frames {
		b := f.Bounds()
		minX, minY := b.Max.X, b.Max.Y
		maxX, maxY := b.Min.X, b.Min.Y
		hasContent := false
		for y := b.Min.Y; y < b.Max.Y; y++ {
			for x := b.Min.X; x < b.Max.X; x++ {
				_, _, _, a := f.At(x, y).RGBA()
				if a > 0 {
					hasContent = true
					if x < minX {
						minX = x
					}
					if x > maxX {
						maxX = x
					}
					if y < minY {
						minY = y
					}
					if y > maxY {
						maxY = y
					}
				}
			}
		}
		if !hasContent {
			boxes[i] = contentBox{0, 0, b.Dx(), b.Dy()}
		} else {
			boxes[i] = contentBox{minX, minY, maxX, maxY}
		}
		cw := boxes[i].maxX - boxes[i].minX + 1
		ch := boxes[i].maxY - boxes[i].minY + 1
		if cw > maxCW {
			maxCW = cw
		}
		if ch > maxCH {
			maxCH = ch
		}
	}

	// Uniform canvas with 10% padding
	canvasW := maxCW * 11 / 10
	canvasH := maxCH * 11 / 10
	// Fixed X center reference: anchor all frames to the same horizontal center
	fixedCenterX := canvasW / 2

	out := make([]image.Image, len(frames))
	for i, f := range frames {
		cb := boxes[i]
		cw := cb.maxX - cb.minX + 1
		ch := cb.maxY - cb.minY + 1
		// All frames share the same center-X and bottom-Y anchor
		ox := fixedCenterX - cw/2 // consistent horizontal center
		oy := canvasH - ch        // bottom-align: feet planted at same Y

		canvas := image.NewRGBA(image.Rect(0, 0, canvasW, canvasH))
		draw.Draw(canvas,
			image.Rect(ox, oy, ox+cw, oy+ch),
			f,
			image.Point{cb.minX, cb.minY},
			draw.Src,
		)
		out[i] = canvas
	}
	return out
}

// RemoveGreenScreen removes green-dominant background pixels, making them transparent.
func RemoveGreenScreen(img image.Image, tol float64) image.Image {
	return removeGreenScreen(toRGBA(img), tol)
}

// TrimAlpha removes fully transparent borders from an image.
func TrimAlpha(img image.Image) image.Image {
	return trimAlpha(toRGBA(img))
}

// Resize resizes an image using nearest-neighbor interpolation.
func Resize(img image.Image, w, h int) image.Image {
	return resize(toRGBA(img), w, h)
}

// ============================
// internal
// ============================

type tile struct {
	x, y, w, h int
}

// removeWhiteBg removes pixels close to pure white (R,G,B all within threshold of 255).
func removeWhiteBg(rgba *image.RGBA, threshold uint8) *image.RGBA {
	if threshold == 0 {
		threshold = 40
	}
	b := rgba.Bounds()
	dst := image.NewRGBA(b)
	draw.Draw(dst, b, rgba, b.Min, draw.Src)

	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			r, g, bl, a := rgba.At(x, y).RGBA()
			if a == 0 {
				continue
			}
			r8, g8, b8 := uint8(r>>8), uint8(g>>8), uint8(bl>>8)
			// Distance from pure white
			dist := max(int(255-r8), max(int(255-g8), int(255-b8)))
			if dist < int(threshold)/2 {
				// Very close to white — fully transparent
				dst.SetRGBA(x, y, color.RGBA{R: r8, G: g8, B: b8, A: 0})
			} else if dist < int(threshold) {
				// Semi-white — fade alpha
				alpha := float64(dist-int(threshold)/2) / float64(int(threshold)/2)
				dst.SetRGBA(x, y, color.RGBA{R: r8, G: g8, B: b8, A: uint8(alpha * 255)})
			}
		}
	}
	return dst
}

func removeGreenScreen(rgba *image.RGBA, tol float64) *image.RGBA {
	b := rgba.Bounds()
	dst := image.NewRGBA(b)
	absTol := tol * 255

	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			r16, g16, bl16, a16 := rgba.At(x, y).RGBA()
			if a16 == 0 {
				continue
			}
			r, g, bl := float64(r16>>8), float64(g16>>8), float64(bl16>>8)
			gDominance := g - (r+bl)/2
			if gDominance > absTol {
				alpha := 1.0 - math.Min(gDominance/(absTol*2), 1.0)
				dst.SetRGBA(x, y, color.RGBA{
					R: uint8(r), G: uint8(g), B: uint8(bl),
					A: uint8(alpha * 255),
				})
			} else {
				dst.Set(x, y, rgba.At(x, y))
			}
		}
	}
	return dst
}

// fixedGridSplit divides the image into Rows×Cols equally-sized cells,
// accounting for a fixed padding between cells.
func fixedGridSplit(rgba *image.RGBA, rows, cols, padding int, minFill float64) []tile {
	b := rgba.Bounds()
	W, H := b.Dx(), b.Dy()

	if padding < 0 {
		padding = 0
	}
	totalPadW := padding * (cols + 1)
	totalPadH := padding * (rows + 1)
	cellW := (W - totalPadW) / cols
	cellH := (H - totalPadH) / rows
	if cellW <= 0 || cellH <= 0 {
		return nil
	}

	var tiles []tile
	for r := 0; r < rows; r++ {
		for c := 0; c < cols; c++ {
			x := padding + c*(cellW+padding)
			y := padding + r*(cellH+padding)
			if tileFillRatio(rgba, x, y, cellW, cellH) >= minFill {
				tiles = append(tiles, tile{x: x, y: y, w: cellW, h: cellH})
			}
		}
	}
	return tiles
}

func projectionSplit(rgba *image.RGBA, gapThreshold float64, minGap int, minFill float64) []tile {
	bounds := rgba.Bounds()
	W, H := bounds.Dx(), bounds.Dy()

	rowRatio := make([]float64, H)
	colRatio := make([]float64, W)
	for y := 0; y < H; y++ {
		n := 0
		for x := 0; x < W; x++ {
			if alphaAt(rgba, x, y) > 0 {
				n++
			}
		}
		rowRatio[y] = float64(n) / float64(W)
	}
	for x := 0; x < W; x++ {
		n := 0
		for y := 0; y < H; y++ {
			if alphaAt(rgba, x, y) > 0 {
				n++
			}
		}
		colRatio[x] = float64(n) / float64(H)
	}

	rowCuts := findCuts(rowRatio, gapThreshold, minGap)
	colCuts := findCuts(colRatio, gapThreshold, minGap)

	if len(rowCuts) < 2 || len(colCuts) < 2 {
		return nil
	}

	var tiles []tile
	for ri := 0; ri < len(rowCuts)-1; ri++ {
		for ci := 0; ci < len(colCuts)-1; ci++ {
			x := colCuts[ci]
			y := rowCuts[ri]
			w := colCuts[ci+1] - x
			h := rowCuts[ri+1] - y
			if tileFillRatio(rgba, x, y, w, h) >= minFill {
				tiles = append(tiles, tile{x: x, y: y, w: w, h: h})
			}
		}
	}
	return tiles
}

func tileFillRatio(rgba *image.RGBA, x0, y0, w, h int) float64 {
	total := w * h
	if total == 0 {
		return 0
	}
	n := 0
	for y := y0; y < y0+h; y++ {
		for x := x0; x < x0+w; x++ {
			if alphaAt(rgba, x, y) > 0 {
				n++
			}
		}
	}
	return float64(n) / float64(total)
}

func findCuts(ratios []float64, threshold float64, minGap int) []int {
	n := len(ratios)
	if n == 0 {
		return nil
	}

	// Smooth the ratio curve with a moving average (kernel size = minGap)
	smoothed := make([]float64, n)
	kernel := max(minGap, 3)
	for i := 0; i < n; i++ {
		sum := 0.0
		count := 0
		for j := max(0, i-kernel/2); j < min(n, i+kernel/2+1); j++ {
			sum += ratios[j]
			count++
		}
		if count > 0 {
			smoothed[i] = sum / float64(count)
		}
	}

	// Compute mean to use as reference
	mean := 0.0
	for _, r := range smoothed {
		mean += r
	}
	mean /= float64(n)

	// Find peaks: contiguous regions where smoothed ratio > mean*1.2
	type segment struct{ start, end int }
	var peaks []segment
	i := 0
	for i < n {
		if smoothed[i] > mean*1.2 {
			start := i
			for i < n && smoothed[i] > mean*0.8 {
				i++
			}
			peaks = append(peaks, segment{start, i})
		} else {
			i++
		}
	}

	if len(peaks) < 2 {
		// Fallback: use threshold-based gap detection
		return findCutsByGap(ratios, threshold, minGap)
	}

	// Find valleys between adjacent peaks (minimum smoothed ratio between them)
	cuts := []int{0}
	for p := 0; p < len(peaks)-1; p++ {
		valleyStart := peaks[p].end
		valleyEnd := peaks[p+1].start
		if valleyStart >= valleyEnd {
			// Peaks adjacent — cut at midpoint
			cuts = append(cuts, (peaks[p].end+peaks[p+1].start)/2)
			continue
		}
		// Find minimum in the valley region
		minIdx := valleyStart
		minVal := smoothed[valleyStart]
		for j := valleyStart + 1; j < valleyEnd; j++ {
			if smoothed[j] < minVal {
				minVal = smoothed[j]
				minIdx = j
			}
		}
		cuts = append(cuts, minIdx)
	}
	cuts = append(cuts, n)
	sort.Ints(cuts)

	// Deduplicate
	dedup := cuts[:1]
	for j := 1; j < len(cuts); j++ {
		if cuts[j] != dedup[len(dedup)-1] {
			dedup = append(dedup, cuts[j])
		}
	}
	return dedup
}

// findCutsByGap is the original threshold-based fallback.
func findCutsByGap(ratios []float64, threshold float64, minGap int) []int {
	n := len(ratios)
	isGap := make([]bool, n)
	for i, r := range ratios {
		isGap[i] = r < threshold
	}

	type segment struct{ start, end int }
	var gaps []segment
	i := 0
	for i < n {
		if isGap[i] {
			start := i
			for i < n && isGap[i] {
				i++
			}
			if i-start >= minGap {
				gaps = append(gaps, segment{start, i})
			}
		} else {
			i++
		}
	}

	if len(gaps) == 0 {
		return []int{0, n}
	}

	cuts := []int{0}
	for _, seg := range gaps {
		cuts = append(cuts, seg.start+(seg.end-seg.start)/2)
	}
	cuts = append(cuts, n)
	sort.Ints(cuts)

	dedup := cuts[:1]
	for j := 1; j < len(cuts); j++ {
		if cuts[j] != dedup[len(dedup)-1] {
			dedup = append(dedup, cuts[j])
		}
	}
	return dedup
}

func alphaAt(rgba *image.RGBA, x, y int) uint8 {
	return rgba.Pix[rgba.PixOffset(x, y)+3]
}

func trimAlpha(rgba *image.RGBA) *image.RGBA {
	b := rgba.Bounds()
	minX, minY := b.Max.X, b.Max.Y
	maxX, maxY := b.Min.X, b.Min.Y
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			if alphaAt(rgba, x-b.Min.X, y-b.Min.Y) > 0 {
				if x < minX {
					minX = x
				}
				if x > maxX {
					maxX = x
				}
				if y < minY {
					minY = y
				}
				if y > maxY {
					maxY = y
				}
			}
		}
	}
	w := maxX - minX + 1
	h := maxY - minY + 1
	if w <= 0 || h <= 0 {
		return rgba
	}
	dst := image.NewRGBA(image.Rect(0, 0, w, h))
	draw.Draw(dst, dst.Bounds(), rgba, image.Point{minX, minY}, draw.Src)
	return dst
}

func resize(rgba *image.RGBA, w, h int) *image.RGBA {
	dst := image.NewRGBA(image.Rect(0, 0, w, h))
	sw, sh := rgba.Bounds().Dx(), rgba.Bounds().Dy()
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			sx := x * sw / w
			sy := y * sh / h
			dst.Set(x, y, rgba.At(sx, sy))
		}
	}
	return dst
}

func toRGBA(src image.Image) *image.RGBA {
	if rgba, ok := src.(*image.RGBA); ok {
		return rgba
	}
	b := src.Bounds()
	rgba := image.NewRGBA(b)
	draw.Draw(rgba, b, src, b.Min, draw.Src)
	return rgba
}