package epd7in3e

import (
	"image"
	"image/draw"
)

// Quantise maps img to a *image.Paletted on the panel's Palette. With
// dither it applies Floyd-Steinberg error diffusion; otherwise each
// pixel becomes its nearest palette colour.
func Quantise(img image.Image, dither bool) *image.Paletted {
	if dither {
		return floydSteinberg(img)
	}
	r := img.Bounds()
	p := image.NewPaletted(r, Palette)
	// Drawing into a Paletted quantises each pixel to the nearest
	// palette colour.
	draw.Draw(p, r, img, r.Min, draw.Src)
	return p
}

// floydSteinberg dithers img to the panel's Palette using
// Floyd-Steinberg error diffusion.
func floydSteinberg(img image.Image) *image.Paletted {
	r := img.Bounds()
	w, h := r.Dx(), r.Dy()

	// Work on a plain NRGBA copy so the loop reads raw bytes.
	src := image.NewNRGBA(r)
	draw.Draw(src, r, img, r.Min, draw.Src)

	// Float copy so diffused error may overshoot 0..255.
	f := make([]float32, w*h*3)
	for y := 0; y < h; y++ {
		i := (y-r.Min.Y)*src.Stride + 4*r.Min.X
		for x := 0; x < w; x++ {
			k := (y*w + x) * 3
			f[k] = float32(src.Pix[i])
			f[k+1] = float32(src.Pix[i+1])
			f[k+2] = float32(src.Pix[i+2])
			i += 4
		}
	}

	p := image.NewPaletted(r, Palette)
	for y := 0; y < h; y++ {
		i := (y-r.Min.Y)*p.Stride + r.Min.X
		for x := 0; x < w; x++ {
			k := (y*w + x) * 3
			cr, cg, cb := f[k], f[k+1], f[k+2]
			idx := nearestIndex(cr, cg, cb)
			p.Pix[i+x] = idx
			pp := paletteRGB[idx]
			er := cr - float32(pp[0])
			eg := cg - float32(pp[1])
			eb := cb - float32(pp[2])
			if x+1 < w { // right: 7/16
				j := k + 3
				f[j] += er * 7 / 16
				f[j+1] += eg * 7 / 16
				f[j+2] += eb * 7 / 16
			}
			if y+1 < h {
				jb := k + 3*w // below: 5/16
				f[jb] += er * 5 / 16
				f[jb+1] += eg * 5 / 16
				f[jb+2] += eb * 5 / 16
				if x > 0 { // bottom-left: 3/16
					j := jb - 3
					f[j] += er * 3 / 16
					f[j+1] += eg * 3 / 16
					f[j+2] += eb * 3 / 16
				}
				if x+1 < w { // bottom-right: 1/16
					j := jb + 3
					f[j] += er / 16
					f[j+1] += eg / 16
					f[j+2] += eb / 16
				}
			}
		}
	}
	return p
}

// nearestIndex returns the Palette index whose RGB value is closest
// to (r, g, b).
func nearestIndex(r, g, b float32) uint8 {
	var best uint8
	bestD := float32(1 << 30)
	for i := range paletteRGB {
		p := paletteRGB[i]
		dr := r - float32(p[0])
		dg := g - float32(p[1])
		db := b - float32(p[2])
		if d := dr*dr + dg*dg + db*db; d < bestD {
			best, bestD = uint8(i), d
		}
	}
	return best
}