package field

import (
	"bufio"
	"image"
	"image/color"
	"image/png"
	"math"
	"os"
	"path/filepath"
	"sort"
)

// A colour preview of a height field: hypsometric tint, hillshade, and the drainage network drawn on top.
//
// The grey thumbnail is nearly useless for judging this generator, which is a problem, because the thing it
// exists to produce is a drainage network and a flat grey ramp is exactly what hides one. Rivers are drawn
// from the flow accumulation with a width that grows with drainage area, so a glance says whether the network
// branches like a river system or like noise.

type PreviewOptions struct {
	// Flow is drainage area per cell, m². Optional; without it no rivers are drawn.
	Flow *Field
	// Sea marks cells below sea level. Optional.
	Sea       []bool
	SeaLevelM float64
	// RiverKm2 is the drainage area at which a channel starts being drawn.
	RiverKm2 float64
	Size     int
	// Crop is a sub-rectangle in map coordinates (x0, y0, x1, y1 in 0..1), rendered at full resolution.
	// A whole continent at 1500 px puts ten kilometres into a hundred pixels, which is enough to see that
	// there is drainage and not nearly enough to see whether it is the right *kind* of drainage. Judging
	// hill country against real hill country needs a crop.
	Crop [4]float64
	// Hillshade exaggerates the vertical before shading. Lowland relief is a few tens of metres over
	// kilometres and disappears at true scale, which is the same reason every printed relief map lies.
	Exaggeration float64
}

// rgb is a colour in 0..255 kept as float64 so the hillshade can multiply it before it is clamped.
type rgb = [3]float64

type stop struct {
	t float64
	c rgb
}

var (
	// A hypsometric ramp: salt-marsh green at sea level through farmland and rock to snow. Stops are chosen
	// so the lowland does not read as one flat colour, which is where most of the map is.
	landStops = []stop{
		{0.00, rgb{72, 106, 68}},
		{0.08, rgb{104, 132, 74}},
		{0.20, rgb{142, 152, 88}},
		{0.38, rgb{164, 148, 104}},
		{0.58, rgb{150, 128, 106}},
		{0.75, rgb{138, 130, 128}},
		{0.88, rgb{176, 174, 174}},
		{1.00, rgb{246, 246, 250}},
	}
	seaShallow = rgb{56, 104, 136}
	seaDeep    = rgb{18, 40, 72}
	riverTint  = rgb{70, 132, 180}
)

func ramp(t float64) rgb {
	if t <= 0 {
		return landStops[0].c
	}
	for i := 1; i < len(landStops); i++ {
		if t <= landStops[i].t {
			a, b := landStops[i-1], landStops[i]
			u := (t - a.t) / (b.t - a.t)
			return rgb{
				a.c[0] + (b.c[0]-a.c[0])*u,
				a.c[1] + (b.c[1]-a.c[1])*u,
				a.c[2] + (b.c[2]-a.c[2])*u,
			}
		}
	}
	return landStops[len(landStops)-1].c
}

// WritePreview renders the field at opt.Size and writes an RGB PNG.
func WritePreview(path string, h *Field, opt PreviewOptions) error {
	size := opt.Size
	if size <= 0 {
		size = 1024
	}
	if size > h.W {
		size = h.W
	}
	if opt.Crop[2] > opt.Crop[0] && opt.Crop[3] > opt.Crop[1] {
		fullW, fullH := h.W, h.H
		h = h.Sub(opt.Crop)
		if opt.Flow != nil {
			opt.Flow = opt.Flow.Sub(opt.Crop)
		}
		if opt.Sea != nil {
			opt.Sea = subMask(opt.Sea, fullW, fullH, opt.Crop)
		}
		if size > h.W {
			size = h.W
		}
	}
	small := h.Resample(size, size)
	exag := opt.Exaggeration
	if exag <= 0 {
		exag = 1
	}

	// Land elevations only: letting the sea floor into the range squashes the whole land ramp.
	//
	// And the top of the ramp is a high percentile, not the maximum. One 2800 m summit over a continent whose
	// land is mostly under 300 m puts every other cell into the bottom tenth of the ramp, and the map reads as
	// uniform green with a white dot on it — which says far more about one pixel than about the terrain. The
	// percentile lets the tint span the distribution that is actually there; the few cells above it clamp to
	// snow, which is what they should look like anyway.
	sea := resampleMask(opt.Sea, h.W, h.H, size)
	landVals := make([]float64, 0, len(small.Data))
	for i, v := range small.Data {
		if sea != nil && sea[i] {
			continue
		}
		landVals = append(landVals, float64(v))
	}
	landMax := 1.0
	if len(landVals) > 0 {
		sort.Float64s(landVals)
		landMax = landVals[int(0.995*float64(len(landVals)-1))]
	}
	if landMax <= 0 {
		landMax = 1
	}
	var seaMin float64
	for i, v := range small.Data {
		if sea != nil && sea[i] && float64(v) < seaMin {
			seaMin = float64(v)
		}
	}

	var flow *Field
	riverA := opt.RiverKm2 * 1e6
	if opt.Flow != nil && riverA > 0 {
		flow = opt.Flow.Resample(size, size)
	}

	img := image.NewRGBA(image.Rect(0, 0, size, size))
	for y := 0; y < size; y++ {
		for x := 0; x < size; x++ {
			i := y*size + x
			elev := float64(small.Data[i])
			var c rgb

			if sea != nil && sea[i] {
				d := 0.0
				if seaMin < 0 {
					d = math.Min(1, (opt.SeaLevelM-elev)/(opt.SeaLevelM-seaMin))
				}
				c = rgb{
					seaShallow[0] + (seaDeep[0]-seaShallow[0])*d,
					seaShallow[1] + (seaDeep[1]-seaShallow[1])*d,
					seaShallow[2] + (seaDeep[2]-seaShallow[2])*d,
				}
			} else {
				c = ramp(math.Min(1, math.Max(0, elev)/landMax))
				// Hillshade from the north-west at 45 degrees, the DEM convention. Applied to land only;
				// shading the sea floor would draw attention to bathymetry nobody will ever see.
				gx := float64(small.AtClamped(x+1, y)-small.AtClamped(x-1, y)) * exag
				gy := float64(small.AtClamped(x, y+1)-small.AtClamped(x, y-1)) * exag
				slope := math.Atan(math.Hypot(gx, gy) / (2 * small.CellM))
				aspect := math.Atan2(gy, -gx)
				lum := math.Cos(slope)*math.Cos(math.Pi/4) +
					math.Sin(slope)*math.Sin(math.Pi/4)*math.Cos(3*math.Pi/4-aspect)
				lum = 0.45 + 0.75*math.Max(0, lum)
				for k := range c {
					c[k] *= lum
				}
			}

			// Rivers on top, their strength growing with the log of drainage area so a trunk reads darker
			// than a headwater without needing a width in pixels.
			if flow != nil {
				if a := float64(flow.Data[i]); a >= riverA {
					w := math.Min(1, math.Log10(a/riverA)/2.2)
					blend := 0.45 + 0.55*w
					for k := range c {
						c[k] = c[k]*(1-blend) + riverTint[k]*blend
					}
				}
			}

			img.Set(x, y, color.RGBA{clamp8(c[0]), clamp8(c[1]), clamp8(c[2]), 255})
		}
	}

	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	bw := bufio.NewWriterSize(f, 1<<20)
	enc := png.Encoder{CompressionLevel: png.DefaultCompression}
	if err := enc.Encode(bw, img); err != nil {
		return err
	}
	return bw.Flush()
}

// resampleMask takes a boolean mask down to the preview size by nearest neighbour; a mask has no meaningful
// average.
func resampleMask(mask []bool, w, h, size int) []bool {
	if mask == nil {
		return nil
	}
	out := make([]bool, size*size)
	for y := 0; y < size; y++ {
		sy := y * (h - 1) / (size - 1)
		for x := 0; x < size; x++ {
			sx := x * (w - 1) / (size - 1)
			out[y*size+x] = mask[sy*w+sx]
		}
	}
	return out
}

func clamp8(v float64) uint8 {
	if v <= 0 {
		return 0
	}
	if v >= 255 {
		return 255
	}
	return uint8(v + 0.5)
}

// subMask is Sub for a boolean mask.
func subMask(mask []bool, w, h int, crop [4]float64) []bool {
	clamp := func(v float64) float64 { return math.Min(1, math.Max(0, v)) }
	x0 := int(clamp(crop[0]) * float64(w-1))
	y0 := int(clamp(crop[1]) * float64(h-1))
	x1 := int(clamp(crop[2]) * float64(w-1))
	y1 := int(clamp(crop[3]) * float64(h-1))
	if x1 <= x0 {
		x1 = x0 + 1
	}
	if y1 <= y0 {
		y1 = y0 + 1
	}
	cw, ch := x1-x0+1, y1-y0+1
	out := make([]bool, cw*ch)
	for y := 0; y < ch; y++ {
		copy(out[y*cw:(y+1)*cw], mask[(y0+y)*w+x0:(y0+y)*w+x0+cw])
	}
	return out
}