package main

import (
	"image"
	"image/color"
	"math/rand"
)

// provided x, y, and color at location, return a color
type quantizerFunction func(int, int, color.Color) color.Color

// apply sequentially applies a quantizing function to an image and returns the result
func apply(i image.Image, f quantizerFunction) image.Image {
	out := image.NewRGBA(image.Rect(0, 0, i.Bounds().Max.X, i.Bounds().Max.Y))

	b := out.Bounds()
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			out.Set(x, y, f(x, y, i.At(x, y)))
		}
	}

	return out
}

// noOp just clones colors from one image to another, to validate file handling.
func noOp(_, _ int, c color.Color) color.Color {
	return c
}

// naiveBW smashes each pixel to black or white based on lumosity.
func naiveBW(_, _ int, c color.Color) color.Color {
	l := luminence(c)
	if l > 0.5 {
		return color.White
	}
	return color.Black
}

// randomNoise injects random noise into the quantization step
func randomNoise(_, _ int, c color.Color) color.Color {
	l := luminence(c)
	if (l + (rand.Float64() - 0.5)) > 0.5 {
		return color.White
	}
	return color.Black
}

// bayer dithering applies a "value map" to our brightness range, instead of messing with the luminence itself.
// the basic one is a matrix of
// 0, 2
// 1, 3
// normalized by the number of cells in the matrix (i.e. divided by 4, in this case) and then compared to the luminosity.
// it takes the current coordinates as input to find your location in the (tiled) matrix.
type bayer struct {
	side   int
	matrix map[coord]float64
}

type coord struct {
	x, y int
}

func (b *bayer) valueAt(x, y int) float64 {
	return b.matrix[coord{x: x % b.side, y: y % b.side}]
}

// I could do this recursively but don't feel like it
func newBayer(level int) *bayer {
	if level == 1 {
		return &bayer{
			side: 4,
			matrix: map[coord]float64{
				{0, 0}: 0 / 16.0,
				{1, 0}: 8 / 16.0,
				{0, 1}: 12 / 16.0,
				{1, 1}: 4 / 16.0,

				{2, 0}: 2 / 16.0,
				{3, 0}: 10 / 16.0,
				{2, 1}: 14 / 16.0,
				{3, 1}: 6 / 16.0,

				{0, 2}: 3 / 16.0,
				{1, 2}: 11 / 16.0,
				{0, 3}: 15 / 16.0,
				{1, 3}: 7 / 16.0,

				{2, 2}: 1 / 16.0,
				{3, 2}: 9 / 16.0,
				{2, 3}: 13 / 16.0,
				{3, 3}: 5 / 16.0,
			},
		}
	}
	return &bayer{
		side: 2,
		matrix: map[coord]float64{
			{0, 0}: 0 / 4.0,
			{1, 0}: 2 / 4.0,
			{0, 1}: 3 / 4.0,
			{1, 1}: 1 / 4.0,
		},
	}
}

func bayerDithering(level int, invert bool) quantizerFunction {
	b := newBayer(level)
	return func(x int, y int, c color.Color) color.Color {
		l := luminence(c)
		v := b.valueAt(x, y)
		if invert {
			if l > 1-v {
				return color.White
			}
			return color.Black
		}
		if l > v {
			return color.White
		}
		return color.Black
	}
}

type diffusion struct {
	matrix  map[coord]float64
	divisor float64
}

func applyError(diffusion diffusion, quantError float64, currentPixel coord, errMap map[coord]float64) {
	for c, i := range diffusion.matrix {
		target := coord{x: currentPixel.x + c.x, y: currentPixel.y + c.y}
		errMap[target] = errMap[target] + (quantError * (i / diffusion.divisor))
	}
}

func simpleErrorDiffusion() quantizerFunction {
	errMap := make(map[coord]float64)
	d := diffusion{
		divisor: 2.0,
		matrix: map[coord]float64{
			{x: 1, y: 0}: 1.0,
			{x: 0, y: 1}: 1.0,
		},
	}
	return func(x int, y int, c color.Color) color.Color {
		p := coord{x: x, y: y}
		l := luminence(c) + errMap[p]
		delete(errMap, p) // don't let the error map grow too big
		if l > 0.5 {
			applyError(d, l-1.0, p, errMap)
			return color.White
		}
		applyError(d, l, p, errMap)
		return color.Black
	}
}

func floydSteinberg() quantizerFunction {
	errMap := make(map[coord]float64)
	d := diffusion{
		divisor: 16.0,
		matrix: map[coord]float64{
			{x: 1, y: 0}:  7.0,
			{x: -1, y: 1}: 3.0,
			{x: 0, y: 1}:  5.0,
			{x: 1, y: 1}:  1.0,
		},
	}
	return func(x int, y int, c color.Color) color.Color {
		p := coord{x: x, y: y}
		l := luminence(c) + errMap[p]
		delete(errMap, p) // don't let the error map grow too big
		if l > 0.5 {
			applyError(d, l-1.0, p, errMap)
			return color.White
		}
		applyError(d, l, p, errMap)
		return color.Black
	}
}

// That is, "relative luminance": https://en.wikipedia.org/wiki/Relative_luminance.
// go's color library doesn't give any information on what "color space" the RGBA is derived from,
// so we convert to Y'CbCr, which returns luminence directly as the Y component.
func luminence(c color.Color) float64 {
	nr, ok := color.NRGBAModel.Convert(c).(color.NRGBA)
	if !ok {
		return 0
	}
	y, _, _ := color.RGBToYCbCr(nr.R, nr.G, nr.B)
	return float64(y) / 255.0
}