package field

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

// False-colour maps of the fields a run works from, as opposed to the field it produces.
//
// preview.png answers "does this look like a landscape". These answer the question that comes next, which is
// "why does it look like that" — the uplift field, the erodibility, the slope and the basins are the inputs
// and the structure, and when a run comes out wrong it is almost always one of them that says so first. The
// uplift map in particular is the one that would have shown, without any arithmetic, that the plains were
// being raised at mountain rates.

// DataMapOptions controls one false-colour render.
type DataMapOptions struct {
	// Sea marks cells to render as flat water rather than data. Optional.
	Sea []bool
	// Size is the output width in pixels; the field is point-sampled down to it and the height follows the
	// field's own aspect.
	Size int
	// Log renders log10 of the value, for anything with a heavy tail — drainage area spans seven decades and
	// is unreadable linearly.
	Log bool
	// Lo and Hi bound the colour ramp. Left at zero they are taken from the 1st and 99th percentile of the
	// data, which keeps one outlier cell from flattening the whole image.
	Lo, Hi float64
	// Palette maps 0..1 to a colour. Nil is Viridis.
	Palette func(float64) [3]float64
}

// WriteDataMap renders one scalar field as a false-colour PNG.
func WriteDataMap(path string, f *Field, opt DataMapOptions) error {
	size := opt.Size
	if size <= 0 || size > f.W {
		size = f.W
	}
	pal := opt.Palette
	if pal == nil {
		pal = Viridis
	}

	vals := make([]float64, len(f.Data))
	for i, v := range f.Data {
		x := float64(v)
		if opt.Log {
			if x < 1 {
				x = 1
			}
			x = math.Log10(x)
		}
		vals[i] = x
	}

	lo, hi := opt.Lo, opt.Hi
	if lo == 0 && hi == 0 {
		lo, hi = percentiles(vals, opt.Sea, 1, 99)
	}
	span := hi - lo
	if span < 1e-12 {
		span = 1
	}

	sizeH := aspectH(f, size)
	img := image.NewRGBA(image.Rect(0, 0, size, sizeH))
	for y := 0; y < sizeH; y++ {
		sy := y * f.H / sizeH
		for x := 0; x < size; x++ {
			sx := x * f.W / size
			i := sy*f.W + sx
			if opt.Sea != nil && opt.Sea[i] {
				img.Set(x, y, color.RGBA{24, 44, 74, 255})
				continue
			}
			t := (vals[i] - lo) / span
			if t < 0 {
				t = 0
			} else if t > 1 {
				t = 1
			}
			c := pal(t)
			img.Set(x, y, color.RGBA{clamp8(c[0]), clamp8(c[1]), clamp8(c[2]), 255})
		}
	}
	return encode(path, img, png.DefaultCompression)
}

// WriteBasinMap colours each drainage basin, which is the one picture that shows whether the solve produced a
// *network* rather than a set of scratches: real basins tile the land, meet along divides that sit where the
// two catchments either side put them, and come in a spread of sizes. A map of noisy speckle means the router
// is re-deciding where the water goes every few cells.
//
// receiver is the D8 receiver array; a cell whose receiver is itself is a basin root.
func WriteBasinMap(path string, w, h int, receiver []int32, sea []bool, size int) error {
	if size <= 0 || size > w {
		size = w
	}
	// Walk each cell down to its root with path compression, so the whole thing stays O(n).
	root := make([]int32, w*h)
	for i := range root {
		root[i] = -1
	}
	var stack []int32
	for i := range root {
		if root[i] >= 0 {
			continue
		}
		stack = stack[:0]
		c := int32(i)
		for root[c] < 0 && receiver[c] != c {
			stack = append(stack, c)
			c = receiver[c]
		}
		r := root[c]
		if r < 0 {
			r = c
			root[c] = r
		}
		for _, s := range stack {
			root[s] = r
		}
	}

	sizeH := int(float64(size)*float64(h)/float64(w) + 0.5)
	if sizeH < 1 {
		sizeH = 1
	}
	img := image.NewRGBA(image.Rect(0, 0, size, sizeH))
	for y := 0; y < sizeH; y++ {
		sy := y * h / sizeH
		for x := 0; x < size; x++ {
			sx := x * w / size
			i := sy*w + sx
			if sea != nil && sea[i] {
				img.Set(x, y, color.RGBA{24, 44, 74, 255})
				continue
			}
			// A hash of the root id, so neighbouring basins get unrelated colours and a divide is a hard
			// edge rather than a gradient.
			k := uint64(uint32(root[i]))*0x9e3779b97f4a7c15 + 0x2545f4914f6cdd1d
			k ^= k >> 29
			k *= 0xbf58476d1ce4e5b9
			k ^= k >> 32
			c := hsv(float64(k%3600)/3600, 0.45+float64((k>>12)%40)/100, 0.55+float64((k>>24)%40)/100)
			img.Set(x, y, color.RGBA{clamp8(c[0]), clamp8(c[1]), clamp8(c[2]), 255})
		}
	}
	return encode(path, img, png.DefaultCompression)
}

// Viridis, sampled at sixteen stops. Perceptually uniform and legible in greyscale, which matters because
// these get pasted into notes and printed.
var viridisStops = [][3]float64{
	{68, 1, 84}, {72, 26, 108}, {71, 47, 125}, {65, 68, 135},
	{57, 86, 140}, {49, 104, 142}, {42, 120, 142}, {35, 136, 142},
	{31, 152, 139}, {34, 168, 132}, {53, 183, 121}, {84, 197, 104},
	{122, 209, 81}, {165, 219, 54}, {210, 226, 27}, {253, 231, 37},
}

func Viridis(t float64) [3]float64 { return sampleStops(viridisStops, t) }

// Inferno, for anything where "how much" reads better as heat: slope and local relief.
var infernoStops = [][3]float64{
	{0, 0, 4}, {12, 8, 38}, {36, 12, 79}, {66, 10, 104},
	{93, 18, 110}, {120, 28, 109}, {147, 38, 103}, {174, 48, 92},
	{199, 62, 76}, {221, 81, 58}, {237, 105, 37}, {247, 133, 17},
	{251, 164, 10}, {249, 196, 41}, {243, 228, 96}, {252, 255, 164},
}

func Inferno(t float64) [3]float64 { return sampleStops(infernoStops, t) }

// Divergent is for a field with a meaningful zero and a sign: cool below, near-white at zero, warm above.
// The change map is the one that needs it — where the surf cut and where it laid are the same magnitude and
// opposite in meaning, and a sequential ramp renders them as the same colour.
func Divergent(t float64) [3]float64 { return sampleStops(divergentStops, t) }

var divergentStops = [][3]float64{
	{30, 64, 120}, {64, 126, 180}, {150, 196, 220}, {238, 238, 236},
	{236, 196, 140}, {206, 132, 62}, {140, 66, 22},
}

func sampleStops(s [][3]float64, t float64) [3]float64 {
	if t <= 0 {
		return s[0]
	}
	if t >= 1 {
		return s[len(s)-1]
	}
	x := t * float64(len(s)-1)
	i := int(x)
	u := x - float64(i)
	a, b := s[i], s[i+1]
	return [3]float64{a[0] + (b[0]-a[0])*u, a[1] + (b[1]-a[1])*u, a[2] + (b[2]-a[2])*u}
}

// HSV is exported because the region map colours its regions the same way the basin map colours its basins:
// a hash of the id straight to a hue, so neighbours get unrelated colours and a boundary is a hard edge.
func HSV(hue, sat, val float64) [3]float64 { return hsv(hue, sat, val) }

func hsv(hue, sat, val float64) [3]float64 {
	h6 := hue * 6
	i := int(h6)
	f := h6 - float64(i)
	p := val * (1 - sat)
	q := val * (1 - sat*f)
	t := val * (1 - sat*(1-f))
	var r, g, b float64
	switch i % 6 {
	case 0:
		r, g, b = val, t, p
	case 1:
		r, g, b = q, val, p
	case 2:
		r, g, b = p, val, t
	case 3:
		r, g, b = p, q, val
	case 4:
		r, g, b = t, p, val
	default:
		r, g, b = val, p, q
	}
	return [3]float64{r * 255, g * 255, b * 255}
}

func percentiles(vals []float64, sea []bool, loPct, hiPct float64) (float64, float64) {
	keep := make([]float64, 0, len(vals))
	for i, v := range vals {
		if sea != nil && sea[i] {
			continue
		}
		keep = append(keep, v)
	}
	if len(keep) == 0 {
		return 0, 1
	}
	sort.Float64s(keep)
	at := func(p float64) float64 {
		i := int(p / 100 * float64(len(keep)-1))
		return keep[i]
	}
	return at(loPct), at(hiPct)
}