package palgen

import (
	"errors"
	"image"
	"image/color"
	"math"
	"sort"
)

// SortablePalette is a slice of color.Color that can be sorted with sort.Sort, by euclidian distance for R, G and B
type SortablePalette []color.Color

// Length from RGB (0, 0, 0)
func colorLength(c color.Color) float64 {
	r := c.(color.RGBA)
	ri, gi, bi := uint32(r.R), uint32(r.G), uint32(r.B)
	return math.Sqrt(float64(ri*ri + gi*gi + bi*bi))
}

func (a SortablePalette) Len() int           { return len(a) }
func (a SortablePalette) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a SortablePalette) Less(i, j int) bool { return colorLength(a[i]) < colorLength(a[j]) }

// Median finds not the average but the median color
func Median(colors []color.Color) (color.Color, error) {
	length := len(colors)
	if length == 0 {
		return nil, errors.New("can't find the median of an empty slice of colors")
	}
	if len(colors) == 1 {
		return colors[0], nil
	}

	// 1. Sort the colors
	sp := SortablePalette(colors)
	sort.Sort(sp)

	// 2. Select the center one, if odd
	if length%2 != 0 {
		centerPos := length / 2
		return sp[centerPos], nil
	}
	// 3. If the numbers are even, select the two center one and take the average of those
	centerPos1 := (length / 2) - 1
	centerPos2 := length / 2
	c1 := sp[centerPos1].(color.RGBA)
	c2 := sp[centerPos2].(color.RGBA)
	r := uint8((uint16(c1.R) + uint16(c2.R)) / 2)
	g := uint8((uint16(c1.G) + uint16(c2.G)) / 2)
	b := uint8((uint16(c1.B) + uint16(c2.B)) / 2)
	a := uint8((uint16(c1.A) + uint16(c2.A)) / 2)
	// return the new color
	return color.RGBA{r, g, b, a}, nil
}

// Median3 finds not the average but the median color. Returns three colors if the number of colors is even (average, first and second color in the center).
func Median3(colors []color.Color) (color.Color, color.Color, color.Color, error) {
	length := len(colors)
	if length == 0 {
		return nil, nil, nil, errors.New("can't find the median of an empty slice of colors")
	}
	if len(colors) == 1 {
		return colors[0], colors[0], colors[0], nil
	}

	// 1. Sort the colors
	sp := SortablePalette(colors)
	sort.Sort(sp)

	// 2. Select the center one, if odd
	if length%2 != 0 {
		centerPos := length / 2
		// Return the center color, thrice
		return sp[centerPos], sp[centerPos], sp[centerPos], nil
	}
	// 3. If the numbers are even, select the two center one and take the average of those
	centerPos1 := (length / 2) - 1
	centerPos2 := length / 2
	c1 := sp[centerPos1].(color.RGBA)
	c2 := sp[centerPos2].(color.RGBA)
	r := uint8((uint16(c1.R) + uint16(c2.R)) / 2)
	g := uint8((uint16(c1.G) + uint16(c2.G)) / 2)
	b := uint8((uint16(c1.B) + uint16(c2.B)) / 2)
	a := uint8((uint16(c1.A) + uint16(c2.A)) / 2)
	averageColor := color.RGBA{r, g, b, a}

	// Also return the two center colors
	return averageColor, sp[centerPos2], sp[centerPos1], nil
}

// Generate can generate a palette with N colors, given an image
func Generate(img image.Image, N int) (color.Palette, error) {
	groups := make(map[int][]color.Color)
	already := make(map[color.Color]bool)

	numberOfPixels := (img.Bounds().Max.Y - img.Bounds().Min.Y) * (img.Bounds().Max.X - img.Bounds().Min.X)
	levelCounter := make(map[int]float64, numberOfPixels)
	colorCounter := make(map[color.RGBA]float64, numberOfPixels)

	// Pick out the colors from the image, per intensity level, and store them in the groups map
	for y := img.Bounds().Min.Y; y < img.Bounds().Max.Y; y++ {
		for x := img.Bounds().Min.X; x < img.Bounds().Max.X; x++ {
			c := img.At(x, y)
			gc := color.GrayModel.Convert(c).(color.Gray)
			rgba := color.RGBAModel.Convert(c).(color.RGBA)
			level := int(float64(gc.Y) / (255.0 / float64(N-1)))
			alreadyColor, ok := already[rgba]
			if !alreadyColor || !ok {
				groups[level] = append(groups[level], rgba)
				already[rgba] = true
			}
			// Count the different levels used
			if counter, ok := levelCounter[level]; ok {
				levelCounter[level] = counter + 1.0
			} else {
				levelCounter[level] = 1.0
			}
			// Count the different colors used
			if counter, ok := colorCounter[rgba]; ok {
				colorCounter[rgba] = counter + 1.0
			} else {
				colorCounter[rgba] = 1.0
			}

		}
	}

	// If the image has fewer unique colors than requested, just return them all
	if len(already) <= N {
		var pal color.Palette
		for c := range already {
			pal = append(pal, c)
		}
		return pal, nil
	}

	// Reset the map for if colors are already appended to a slice
	already = make(map[color.Color]bool)
	already2 := make(map[color.Color]bool)
	var extrapal color.Palette

	// Remove the group of colors with the least used intensity level
	minCoverage := 1.0
	minCoverageKey := 0
	found := false
	for intensityLevelKey := range groups {
		coverage := levelCounter[intensityLevelKey] / float64(numberOfPixels)
		if coverage < minCoverage {
			minCoverage = coverage
			minCoverageKey = intensityLevelKey
			found = true
		}
	}
	if found {
		delete(groups, minCoverageKey)
	}

	// Find the median color for each intensity level
	var pal color.Palette
	for _, colors := range groups {
		// Find the median color of a group of colors of a certain intensity
		//medianColor, err := Median(colors)
		medianColor1, medianColor2, medianColor3, err := Median3(colors)
		if err != nil {
			return nil, err
		}
		// Add the medianColor1 to the palette, if it's not already there
		alreadyColor, ok := already[medianColor1]
		if !alreadyColor || !ok {
			pal = append(pal, medianColor1)
			already[medianColor1] = true
		}
		// Add medianColor2 and medianColor3 to the extra palette, if they are not already in it
		alreadyColor2, ok := already2[medianColor2]
		if !alreadyColor2 || !ok {
			extrapal = append(extrapal, medianColor2)
			already2[medianColor2] = true
		}
		alreadyColor2, ok = already2[medianColor3]
		if !alreadyColor2 || !ok {
			extrapal = append(extrapal, medianColor3)
			already2[medianColor3] = true
		}
	}

	// If there are not enough colors in the generated palette (N+1, because we will remove one), add colors from extrapal
	for (len(pal) < N) && (len(extrapal) > 0) {
		var (
			maxUsage          = -1.0
			mostFrequentIndex = -1
		)

		// Loop through extrapal to find the most frequently used color
		for i, extraColor := range extrapal {
			rgbaExtra := color.RGBAModel.Convert(extraColor).(color.RGBA)
			if usage, ok := colorCounter[rgbaExtra]; ok {
				if usage > maxUsage {
					maxUsage = usage
					mostFrequentIndex = i
				}
			}
		}
		if mostFrequentIndex == -1 { // not found
			break
		}

		// Use the most frequently used color from the palette of extra colors
		c := extrapal[mostFrequentIndex]

		// Add the color to the palette, if it's not already there
		alreadyColor, ok := already[c]
		if !alreadyColor || !ok {
			pal = append(pal, c)
			already[c] = true
		}

		// Remove this color from extrapal
		extrapal = append(extrapal[:mostFrequentIndex], extrapal[mostFrequentIndex+1:]...)
	}

	// Now remove the one extra color that covers the least amount of pixels of the image that has been converted to the current palette
	if len(pal) > N {
		minCoverage := 1.0
		minCoverageIndex := -1
		for i := 0; i < len(pal); i++ {
			c := color.RGBAModel.Convert(pal[i]).(color.RGBA)
			coverage := colorCounter[c] / float64(numberOfPixels)
			if coverage < minCoverage {
				minCoverage = coverage
				minCoverageIndex = i
			}
		}
		if minCoverageIndex != -1 {
			// Replace i with the last element
			pal[minCoverageIndex] = pal[len(pal)-1]
			// Remove the last element
			pal = pal[:len(pal)-1]
		}
	}

	// Remove near-duplicate dark colors. Look at the darkest 10% of the
	// palette (by lightness) and remove entries that are within 5% lightness
	// of another dark color, keeping the most frequently used one.
	darkCount := len(pal) / 10
	if darkCount >= 2 {
		const lightnessThreshold = 255.0 * 0.05 // 5% of full range

		// Find lightness for each palette entry
		type palEntry struct {
			index     int
			lightness float64
		}
		entries := make([]palEntry, len(pal))
		for i, c := range pal {
			gc := color.GrayModel.Convert(c).(color.Gray)
			entries[i] = palEntry{i, float64(gc.Y)}
		}

		// Sort by lightness to find the darkest entries
		sort.Slice(entries, func(i, j int) bool {
			return entries[i].lightness < entries[j].lightness
		})

		// Among the darkest entries, mark near-duplicates for removal
		darkEntries := entries[:darkCount]
		removeSet := make(map[int]bool)
		for i := range darkEntries {
			if removeSet[darkEntries[i].index] {
				continue
			}
			for j := i + 1; j < len(darkEntries); j++ {
				if removeSet[darkEntries[j].index] {
					continue
				}
				if darkEntries[j].lightness-darkEntries[i].lightness > lightnessThreshold {
					break // sorted, so no more matches
				}
				// These two are near-duplicates; remove the less used one
				ci := color.RGBAModel.Convert(pal[darkEntries[i].index]).(color.RGBA)
				cj := color.RGBAModel.Convert(pal[darkEntries[j].index]).(color.RGBA)
				if colorCounter[ci] >= colorCounter[cj] {
					removeSet[darkEntries[j].index] = true
				} else {
					removeSet[darkEntries[i].index] = true
					break // i is removed, stop comparing from it
				}
			}
		}

		// Rebuild palette without the removed entries
		if len(removeSet) > 0 {
			var trimmed color.Palette
			for i, c := range pal {
				if !removeSet[i] {
					trimmed = append(trimmed, c)
				}
			}
			pal = trimmed

			// Backfill freed slots from extrapal
			for (len(pal) < N) && (len(extrapal) > 0) {
				var (
					maxUsage          = -1.0
					mostFrequentIndex = -1
				)
				for i, extraColor := range extrapal {
					rgbaExtra := color.RGBAModel.Convert(extraColor).(color.RGBA)
					if usage, ok := colorCounter[rgbaExtra]; ok {
						if usage > maxUsage {
							maxUsage = usage
							mostFrequentIndex = i
						}
					}
				}
				if mostFrequentIndex == -1 {
					break
				}
				c := extrapal[mostFrequentIndex]
				alreadyColor, ok := already[c]
				if !alreadyColor || !ok {
					pal = append(pal, c)
					already[c] = true
				}
				extrapal = append(extrapal[:mostFrequentIndex], extrapal[mostFrequentIndex+1:]...)
			}
		}
	}

	// Return the generated palette
	return pal, nil
}

// GenerateUpTo can generate a palette with up to N colors, given an image
func GenerateUpTo(img image.Image, N int) (color.Palette, error) {
	groups := make(map[int][]color.Color)
	already := make(map[color.Color]bool)

	numberOfPixels := (img.Bounds().Max.Y - img.Bounds().Min.Y) * (img.Bounds().Max.X - img.Bounds().Min.X)
	levelCounter := make(map[int]float64, numberOfPixels)
	colorCounter := make(map[color.RGBA]float64, numberOfPixels)

	// Pick out the colors from the image, per intensity level, and store them in the groups map
	for y := img.Bounds().Min.Y; y < img.Bounds().Max.Y; y++ {
		for x := img.Bounds().Min.X; x < img.Bounds().Max.X; x++ {
			c := img.At(x, y)
			gc := color.GrayModel.Convert(c).(color.Gray)
			rgba := color.RGBAModel.Convert(c).(color.RGBA)
			level := int(float64(gc.Y) / (255.0 / float64(N-1)))
			alreadyColor, ok := already[rgba]
			if !alreadyColor || !ok {
				groups[level] = append(groups[level], rgba)
				already[rgba] = true
			}
			// Count the different levels used
			if counter, ok := levelCounter[level]; ok {
				levelCounter[level] = counter + 1.0
			} else {
				levelCounter[level] = 1.0
			}
			// Count the different colors used
			if counter, ok := colorCounter[rgba]; ok {
				colorCounter[rgba] = counter + 1.0
			} else {
				colorCounter[rgba] = 1.0
			}
		}
	}

	// Remove the group of colors with the least used intensity level,
	// but only if there are more groups than requested colors.
	if len(groups) > N {
		minCoverage := 1.0
		minCoverageKey := 0
		found := false
		for intensityLevelKey := range groups {
			coverage := levelCounter[intensityLevelKey] / float64(numberOfPixels)
			if coverage < minCoverage {
				minCoverage = coverage
				minCoverageKey = intensityLevelKey
				found = true
			}
		}
		if found {
			delete(groups, minCoverageKey)
		}
	}

	// Reset the map for if colors are already appended to a slice
	already = make(map[color.Color]bool)

	// Find the median color for each intensity level
	var pal color.Palette
	for _, colors := range groups {
		// Find the median color of a group of colors of a certain intensity
		//medianColor, err := Median(colors)
		medianColor1, _, _, err := Median3(colors)
		if err != nil {
			return nil, err
		}
		// Add the medianColor1 to the palette, if it's not already there
		alreadyColor, ok := already[medianColor1]
		if !alreadyColor || !ok {
			pal = append(pal, medianColor1)
			already[medianColor1] = true
		}
	}

	// Return the generated palette
	return pal, nil
}
