Palette from Image — Go source

// Package paletteextract is the Go twin of CosmoDev's src/lib/palette-extract.ts
// (dual source: the web lib is TypeScript, the CLI lib is Go — kept in
// lock-step). Pure + deterministic: median-cut quantization over RGBA pixels
// using stable sorts only. The table-driven tests in palette-from-image_test.go
// share vectors with src/lib/palette-extract.test.ts so the two implementations
// are held to the same contract.
package paletteextract

import (
	"fmt"
	"math"
	"sort"
)

// Swatch is one quantized palette color plus the number of sampled pixels it
// covers. Mirrors the TS Swatch interface.
type Swatch struct {
	R, G, B    uint8
	Population int
}

// pixel is an opaque sample pixel (alpha-0 pixels are dropped at sampling).
type pixel struct {
	r, g, b uint8
}

// maxSamples down-samples so large images quantize in bounded time.
const maxSamples = 16384

// ExtractPalette runs median-cut quantization over an RGBA byte buffer.
// maxColors bounds the palette size (the TS default is 8). It returns nil for
// empty or fully transparent input, and never panics.
func ExtractPalette(rgba []byte, maxColors int) []Swatch {
	total := len(rgba) / 4
	if total == 0 {
		return nil
	}

	pixels := make([]pixel, 0, total)
	stride := max(1, total/maxSamples)
	for i := 0; i < total; i += stride {
		o := i * 4
		if rgba[o+3] == 0 {
			continue // fully transparent pixels never contribute
		}
		pixels = append(pixels, pixel{r: rgba[o], g: rgba[o+1], b: rgba[o+2]})
	}
	if len(pixels) == 0 {
		return nil
	}

	buckets := [][]pixel{pixels}
	for len(buckets) < maxColors {
		// Widest-range bucket with more than one distinct value wins the split.
		// Strict > keeps the lowest-index bucket on ties, as in the TS loop.
		bestIdx := -1
		bestRange := 1 // range 1 (exact duplicates only) never splits further
		for i := range buckets {
			if r := channelRange(buckets[i]); r > bestRange {
				bestRange = r
				bestIdx = i
			}
		}
		if bestIdx == -1 {
			break
		}
		bucket := buckets[bestIdx]
		// Mirror the TS splice(bestIdx, 1) + push(halves): remove, then append.
		buckets = append(buckets[:bestIdx], buckets[bestIdx+1:]...)
		lo, hi := splitBucket(bucket)
		buckets = append(buckets, lo, hi)
	}

	swatches := make([]Swatch, 0, len(buckets))
	for _, bucket := range buckets {
		if len(bucket) == 0 {
			continue
		}
		var sr, sg, sb int
		for _, p := range bucket {
			sr += int(p.r)
			sg += int(p.g)
			sb += int(p.b)
		}
		swatches = append(swatches, Swatch{
			R:          roundAvg(sr, len(bucket)),
			G:          roundAvg(sg, len(bucket)),
			B:          roundAvg(sb, len(bucket)),
			Population: len(bucket),
		})
	}
	// Population-descending, stable — mirrors the TS .sort(b - a) comparator.
	sort.SliceStable(swatches, func(i, j int) bool {
		return swatches[i].Population > swatches[j].Population
	})
	return swatches
}

// ToHex formats a swatch as #rrggbb (lowercase, always two digits per channel).
func ToHex(sw Swatch) string {
	return fmt.Sprintf("#%02x%02x%02x", sw.R, sw.G, sw.B)
}

// roundAvg averages a channel sum over n with TS Math.round semantics
// (half rounds up; sums are non-negative so math.Round matches exactly).
func roundAvg(sum, n int) uint8 {
	return uint8(math.Round(float64(sum) / float64(n)))
}

// channelRange returns the widest per-channel min..max span in the bucket.
func channelRange(bucket []pixel) int {
	minR, maxR := uint8(255), uint8(0)
	minG, maxG := uint8(255), uint8(0)
	minB, maxB := uint8(255), uint8(0)
	for _, p := range bucket {
		minR, maxR = min(minR, p.r), max(maxR, p.r)
		minG, maxG = min(minG, p.g), max(maxG, p.g)
		minB, maxB = min(minB, p.b), max(maxB, p.b)
	}
	return max(int(maxR-minR), int(maxG-minG), int(maxB-minB))
}

// splitBucket sorts the bucket by its widest channel (ties prefer r, then g —
// mirroring the TS reduce) and halves it at the median.
func splitBucket(bucket []pixel) ([]pixel, []pixel) {
	minR, maxR := uint8(255), uint8(0)
	minG, maxG := uint8(255), uint8(0)
	minB, maxB := uint8(255), uint8(0)
	for _, p := range bucket {
		minR, maxR = min(minR, p.r), max(maxR, p.r)
		minG, maxG = min(minG, p.g), max(maxG, p.g)
		minB, maxB = min(minB, p.b), max(maxB, p.b)
	}
	ranges := [3]int{int(maxR - minR), int(maxG - minG), int(maxB - minB)}
	ch := 0 // r
	if ranges[1] > ranges[ch] {
		ch = 1 // g
	}
	if ranges[2] > ranges[ch] {
		ch = 2 // b
	}

	sort.SliceStable(bucket, func(i, j int) bool {
		switch ch {
		case 0:
			return bucket[i].r < bucket[j].r
		case 1:
			return bucket[i].g < bucket[j].g
		default:
			return bucket[i].b < bucket[j].b
		}
	})
	mid := len(bucket) / 2
	return bucket[:mid], bucket[mid:]
}