// Package field is the one array type the whole generator passes around: a square-ish grid of float32 in a
// known unit, with the cell size in metres attached so no pass has to be told the scale twice.
//
// Determinism (cross-cutting rule 12) is a property of this package as much as of the passes. Everything
// parallel here partitions rows into disjoint, contiguous ranges and writes only into its own range, so the
// result does not depend on how the goroutines were scheduled. Nothing reduces through a channel.
package field

import (
	"math"
	"runtime"
	"sort"
	"sync"
)

// Field is a W x H grid, row-major, with CellM metres between neighbouring samples.
type Field struct {
	W, H  int
	CellM float64
	Data  []float32
}

func New(w, h int, cellM float64) *Field {
	return &Field{W: w, H: h, CellM: cellM, Data: make([]float32, w*h)}
}

// NewLike is an empty field with another's shape and scale.
func NewLike(f *Field) *Field { return New(f.W, f.H, f.CellM) }

func (f *Field) Idx(x, y int) int        { return y*f.W + x }
func (f *Field) At(x, y int) float32     { return f.Data[y*f.W+x] }
func (f *Field) Set(x, y int, v float32) { f.Data[y*f.W+x] = v }
func (f *Field) Len() int                { return len(f.Data) }

// AtClamped samples with edge clamping, which is what every stencil in the generator wants at the border.
func (f *Field) AtClamped(x, y int) float32 {
	if x < 0 {
		x = 0
	} else if x >= f.W {
		x = f.W - 1
	}
	if y < 0 {
		y = 0
	} else if y >= f.H {
		y = f.H - 1
	}
	return f.Data[y*f.W+x]
}

func (f *Field) Clone() *Field {
	c := New(f.W, f.H, f.CellM)
	copy(c.Data, f.Data)
	return c
}

func (f *Field) Fill(v float32) {
	for i := range f.Data {
		f.Data[i] = v
	}
}

func (f *Field) MinMax() (float32, float32) {
	if len(f.Data) == 0 {
		return 0, 0
	}
	lo, hi := f.Data[0], f.Data[0]
	for _, v := range f.Data {
		if v < lo {
			lo = v
		}
		if v > hi {
			hi = v
		}
	}
	return lo, hi
}

func (f *Field) Mean() float64 {
	if len(f.Data) == 0 {
		return 0
	}
	// Summed as float64 in index order: the same total every run, whatever the machine.
	var sum float64
	for _, v := range f.Data {
		sum += float64(v)
	}
	return sum / float64(len(f.Data))
}

// Percentile sorts a copy, so it costs a copy and a sort; used for thresholds, not in inner loops.
func (f *Field) Percentile(p float64) float32 {
	if len(f.Data) == 0 {
		return 0
	}
	c := make([]float32, len(f.Data))
	copy(c, f.Data)
	sort.Slice(c, func(i, j int) bool { return c[i] < c[j] })
	i := int(p / 100 * float64(len(c)-1))
	if i < 0 {
		i = 0
	} else if i >= len(c) {
		i = len(c) - 1
	}
	return c[i]
}

// Normalise maps the field onto [0, 1]. A flat field becomes zero rather than a division by nothing.
func (f *Field) Normalise() {
	lo, hi := f.MinMax()
	span := float64(hi - lo)
	if span < 1e-9 {
		f.Fill(0)
		return
	}
	for i, v := range f.Data {
		f.Data[i] = float32((float64(v) - float64(lo)) / span)
	}
}

// Slope returns rise over run per cell, the central difference used by the layer rules and the statistics.
func (f *Field) Slope() *Field {
	out := NewLike(f)
	inv := float32(1.0 / (2.0 * f.CellM))
	Rows(f.H, func(y0, y1 int) {
		for y := y0; y < y1; y++ {
			for x := 0; x < f.W; x++ {
				gx := (f.AtClamped(x+1, y) - f.AtClamped(x-1, y)) * inv
				gy := (f.AtClamped(x, y+1) - f.AtClamped(x, y-1)) * inv
				out.Data[out.Idx(x, y)] = float32(math.Hypot(float64(gx), float64(gy)))
			}
		}
	})
	return out
}

// Curvature is the Laplacian in metres per cell squared: positive on ridges and convex shoulders, negative in
// gullies and sediment traps. Ported from heightmap_erosion.curvature, which blurs lightly first.
func (f *Field) Curvature() *Field {
	h := f.Blur(2)
	out := NewLike(f)
	inv := float32(1.0 / f.CellM)
	Rows(f.H, func(y0, y1 int) {
		for y := y0; y < y1; y++ {
			for x := 0; x < f.W; x++ {
				lap := h.AtClamped(x-1, y) + h.AtClamped(x+1, y) + h.AtClamped(x, y-1) + h.AtClamped(x, y+1) - 4*h.At(x, y)
				out.Data[out.Idx(x, y)] = lap * inv
			}
		}
	})
	return out
}

// Blur is the five-point box blur the numpy pipeline used, repeated. Edge-clamped, so it does not darken
// the border the way a zero-padded one would.
func (f *Field) Blur(passes int) *Field {
	cur := f.Clone()
	if passes <= 0 {
		return cur
	}
	next := NewLike(f)
	for p := 0; p < passes; p++ {
		Rows(f.H, func(y0, y1 int) {
			for y := y0; y < y1; y++ {
				for x := 0; x < cur.W; x++ {
					s := cur.At(x, y) + cur.AtClamped(x-1, y) + cur.AtClamped(x+1, y) + cur.AtClamped(x, y-1) + cur.AtClamped(x, y+1)
					next.Data[next.Idx(x, y)] = s / 5
				}
			}
		})
		cur, next = next, cur
	}
	return cur
}

// Rows runs fn over disjoint contiguous row ranges, one per core. The ranges are fixed before any goroutine
// starts and each writes only into its own, so the output is identical at any GOMAXPROCS. Every parallel
// loop in the generator goes through here; none spawns goroutines of its own.
func Rows(h int, fn func(y0, y1 int)) {
	RowsIndexed(h, func(_, y0, y1 int) { fn(y0, y1) })
}

// RowsIndexed is Rows with the band number, which is what a parallel loop needs when it has to reduce
// something rather than only write into its own rows: it gives each goroutine a pre-allocated indexed
// slot to accumulate into, so the reduction can be replayed in band order afterwards instead of
// depending on which goroutine finished first. Size the slots with BandCount.
func RowsIndexed(h int, fn func(band, y0, y1 int)) {
	step := rowStep(h)
	if step >= h {
		fn(0, 0, h)
		return
	}
	var wg sync.WaitGroup
	band := 0
	for y0 := 0; y0 < h; y0 += step {
		y1 := y0 + step
		if y1 > h {
			y1 = h
		}
		wg.Add(1)
		go func(k, a, b int) {
			defer wg.Done()
			fn(k, a, b)
		}(band, y0, y1)
		band++
	}
	wg.Wait()
}

// FixedBands is RowsIndexed with a partition that does not depend on the core count: bands of exactly rows
// rows, run by however many workers there are.
//
// It exists for one reason. A parallel loop that only writes into its own rows can be partitioned any way at
// all, which is what Rows does. A loop that *reduces* into overlapping buffers cannot: floating-point addition
// is not associative, so summing a cell's contributions in a different grouping gives a different last bit,
// and the result would depend on GOMAXPROCS. The particle pass is that loop. Fix the partition and the
// arithmetic is fixed with it.
func FixedBands(h, rows int, fn func(band, y0, y1 int)) {
	if rows < 1 {
		rows = 1
	}
	n := FixedBandCount(h, rows)
	workers := runtime.GOMAXPROCS(0)
	if workers > n {
		workers = n
	}
	if workers <= 1 {
		for b := 0; b < n; b++ {
			y0 := b * rows
			y1 := min(y0+rows, h)
			fn(b, y0, y1)
		}
		return
	}
	next := make(chan int)
	go func() {
		for b := 0; b < n; b++ {
			next <- b
		}
		close(next)
	}()
	var wg sync.WaitGroup
	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for b := range next {
				y0 := b * rows
				y1 := min(y0+rows, h)
				fn(b, y0, y1)
			}
		}()
	}
	wg.Wait()
}

// FixedBandCount is how many bands FixedBands will make.
func FixedBandCount(h, rows int) int {
	if rows < 1 {
		rows = 1
	}
	if h <= 0 {
		return 0
	}
	return (h + rows - 1) / rows
}

// BandCount is how many ranges Rows and RowsIndexed split h into. It is fixed by h and GOMAXPROCS, so it
// can be called to size a reduction before the loop starts.
func BandCount(h int) int {
	step := rowStep(h)
	if step >= h {
		return 1
	}
	return (h + step - 1) / step
}

func rowStep(h int) int {
	workers := runtime.GOMAXPROCS(0)
	if workers > h {
		workers = h
	}
	if workers <= 1 {
		return h
	}
	return (h + workers - 1) / workers
}