// Package fluvial is the stream-power erosion solve: dh/dt = U - K * A^m * S^n, integrated implicitly up the
// drainage stack by the method of Braun & Willett (2013).
//
// This is the reason the generator exists (D-47). Particle erosion carves the path each droplet happens to
// take: it makes wear, but never a network. Stream power solves for drainage area first and erodes in
// proportion to it, which is what produces a branching hierarchy, valleys whose size matches the area they
// drain, and divides that sit where the basins either side of them put them.
//
// Four things happen per step, in this order:
//
//  1. Depressions are filled (priority-flood), because a D8 receiver graph containing a pit has no path to
//     base level and the implicit solve has nothing to descend to. This is the only part of a step that is
//     not O(n), so it runs every FillEvery steps, not every step.
//  2. Receivers and the stack are computed: steepest descent to one of eight neighbours, then a depth-first
//     ordering in which every node appears after its receiver.
//  3. Drainage area is accumulated down the stack in reverse.
//  4. The implicit update runs up the stack, so each node's receiver already holds its new height. This is
//     what makes the scheme unconditionally stable in dt, and it is why a naive explicit solver is not an
//     acceptable substitute at dt = 1500 yr.
package fluvial

import (
	"math"

	"salty/terrain/internal/field"
	"salty/terrain/internal/thermal"
)

// Params are the stream-power constants. K is per year with A in m².
type Params struct {
	K         float64
	M         float64
	N         float64
	DtYr      float64
	Steps     int
	Diffusion float64 // hillslope diffusivity, m²/yr
	FillEvery int

	// Landsliding. Stream power bounds nothing at small drainage area, so without this the hillslopes grow
	// as steep as the uplift rate asks them to and the map becomes needles. TalusSlope is rise over run; 0
	// disables. It runs inside the loop rather than after it, because a cap applied once at the end just
	// shaves the tops off, while a cap applied throughout changes where the sediment goes.
	TalusSlope    float64
	ThermalEvery  int
	ThermalPasses int

	// CriticalAreaM2 is where channels begin. Below it the cell is a hillslope: it still rises, and
	// diffusion and landsliding still shape it, but stream power does not incise it.
	//
	// This is not a tuning knob, it is a correctness fix. Stream power is a law about channels, and applying
	// it at A = one cell says a cell that drains only itself should stand at U/(K*cellM^(2m)) — 63 degrees at
	// the rates here. That is why the map came out as needles at a 39 degree median. It is also why relief
	// was resolution-dependent (1020 m at 512², 2605 m at 1786² on one seed): halving the cell size halves
	// the smallest A and steepens every divide, for ever. A critical area is a physical length, so the same
	// landscape comes out at any resolution, which is the property the full-resolution run depends on.
	CriticalAreaM2 float64

	// ChannelTaper is the exponent on (A/Ac) below the channel head. 0 is no taper, 2 is strong.
	ChannelTaper float64

	// CriticalSlope is Sc in the nonlinear hillslope law, rise over run. Above zero it selects
	// DiffuseNonlinear over plain linear diffusion and takes the repose clamp out of the step loop; see
	// hillslope.go and Run. Zero keeps the old pairing of linear diffusion and an in-loop clamp.
	CriticalSlope   float64
	SlopeCap        float64 // where the flux stops stiffening, as a fraction of Sc
	MaxHillslopeSub int     // the sub-step budget that bound buys

	// MFDExponent selects multiple-flow-direction drainage area over D8's single receiver, and is the
	// exponent on the partition. 0 keeps the old Accumulate, which is what every bake before this ran and
	// what the A/B comparison needs. See mfd.go for why one is the right default.
	MFDExponent float64
}

// Grid holds the flow topology and the scratch it is built from. Allocated once and reused across every
// step: at 3.2 M cells the allocations would otherwise dominate the solve.
type Grid struct {
	W, H  int
	CellM float64

	// Base marks cells fixed at base level: the ocean. The map border is an outlet too, so what actually
	// counts as "fixed" is Base plus the border, and that union is `fixed`. Keeping only Base in mind here is
	// how the outlets themselves ended up being uplifted: on a map with no ocean, every border cell has
	// Receiver == itself and Base == false, so base level rose two metres a step and the whole solve chased
	// it. The steady-state test is what caught it.
	Base  []bool
	fixed []bool

	Receiver []int32   // index of the cell this one drains to; itself for a base cell
	Length   []float32 // distance to that receiver, metres
	Stack    []int32   // every node after its receiver
	Area     []float32 // drainage area, m²

	cancel    <-chan struct{} // closed to abandon a run mid-solve; see SetCancel
	seed      uint64          // the jitter's seed; see jitter.go and SetSeed
	originX   int32           // where this grid sits on the planet; see SetFrame. Zero is "this grid is the world"
	originY   int32
	planetW   int32 // the cylinder's width, or 0 when there is no cylinder
	donorOff  []int32
	donorList []int32
	cursor    []int32
	closed    []bool
	pq        *bucketPQ
	fifo      []int32
	scratch   []float32

	// Multiple-flow accumulation. mfdPending is how many strictly higher neighbours a cell still owes before
	// it may be released; a byte, because a cell has eight neighbours and cannot owe more. See mfd.go.
	mfdPending []uint8
	mfdQueue   []int32
	mfdMode    mfdPow
	mfdExp     float64
}

// SetElevationRange sizes the flood's bucket queue. Called once, with the manifest's elevation range plus a
// margin, before the first step.
func (g *Grid) SetElevationRange(loM, hiM float64) {
	g.pq = newBucketPQ(loM, hiM)
}

var (
	// D8, in the order (-1,-1) .. (1,1) skipping the centre. The order is fixed so that a run is reproducible,
	// but it is no longer what decides a tie between two equally steep neighbours: a fixed order resolves
	// every tie the same way and prints its preferred axis across any near-flat ground. See jitter.go.
	dx8 = [8]int{-1, 0, 1, -1, 1, -1, 0, 1}
	dy8 = [8]int{-1, -1, -1, 0, 0, 1, 1, 1}
)

func NewGrid(w, h int, cellM float64, base []bool) *Grid {
	n := w * h
	g := &Grid{
		W: w, H: h, CellM: cellM, Base: base,
		Receiver: make([]int32, n), Length: make([]float32, n), Stack: make([]int32, 0, n),
		Area: make([]float32, n), donorOff: make([]int32, n+1), donorList: make([]int32, n),
		closed: make([]bool, n), fifo: make([]int32, 0, n), scratch: make([]float32, n),
		mfdPending: make([]uint8, n),
	}
	g.fixed = make([]bool, n)
	for i := range g.fixed {
		g.fixed[i] = g.isOutlet(i)
	}
	g.SetElevationRange(-2000, 4000)
	return g
}

// FillDepressions raises closed depressions to their spill point, in place, using Barnes' improved
// priority-flood with a plain FIFO beside the heap. The FIFO is the optimisation that matters: on real
// terrain most cells are reached while descending into an already-flooded pit, and those never touch the
// heap, which turns the cost from "half an hour over a run" into something affordable.
//
// The epsilon variant adds a millimetre of fall per cell across a flat, so filled lakes still route rather
// than becoming a plateau the flow accumulator cannot leave. That millimetre is scattered per cell by a hash
// of the index rather than applied uniformly: a uniform epsilon means the only gradient on a flat is the
// flood's own traversal order, and the router then draws that order as rivers. See jitter.go.
func (g *Grid) FillDepressions(h []float32, epsilon float32) {
	n := g.W * g.H
	for i := range g.closed {
		g.closed[i] = false
	}
	g.pq.reset()
	g.fifo = g.fifo[:0]

	for i := 0; i < n; i++ {
		if g.isOutlet(i) {
			g.closed[i] = true
			g.pq.push(h[i], int32(i))
		}
	}

	head := 0
	for g.pq.len() > 0 || head < len(g.fifo) {
		var c int32
		var celev float32
		// Drain the FIFO while it cannot violate the queue's ordering.
		if head < len(g.fifo) && (g.pq.len() == 0 || h[g.fifo[head]] <= g.pq.frontElev()) {
			c = g.fifo[head]
			head++
			celev = h[c]
		} else {
			c = g.pq.pop()
			if c < 0 {
				break
			}
			celev = h[c]
		}
		cx := int(c) % g.W
		cy := int(c) / g.W
		for k := 0; k < 8; k++ {
			nx, ny := cx+dx8[k], cy+dy8[k]
			if nx < 0 || ny < 0 || nx >= g.W || ny >= g.H {
				continue
			}
			ni := int32(ny*g.W + nx)
			if g.closed[ni] {
				continue
			}
			g.closed[ni] = true
			if h[ni] <= celev {
				h[ni] = celev + epsilon*(0.5+hashXY(g.seed, g.worldX(nx), g.worldY(ny), jitterFloodEpsilon))
				g.fifo = append(g.fifo, ni)
			} else {
				g.pq.push(h[ni], ni)
			}
		}
		// Compact the FIFO occasionally so it does not grow without bound over a whole flood.
		if head > n/2 {
			g.fifo = append(g.fifo[:0], g.fifo[head:]...)
			head = 0
		}
	}
}

func (g *Grid) isOutlet(i int) bool {
	if g.Base != nil && g.Base[i] {
		return true
	}
	x, y := i%g.W, i/g.W
	return x == 0 || y == 0 || x == g.W-1 || y == g.H-1
}

// ComputeReceivers picks the steepest downhill neighbour of each cell. A base cell, and any cell with no
// lower neighbour, receives itself, which makes it a root of the stack.
//
// Two neighbours of equal steepness are separated by a hash of the cell and the direction rather than by the
// fixed order of dx8/dy8. A fixed order always resolves a tie the same way, which on any near-flat surface
// puts a systematic preference on one grid axis and shows up as rivers that run straight along it. The
// perturbation is a tenth of a percent, so it decides near-ties and nothing else: a neighbour that is
// genuinely steeper than another by more than that is still chosen.
func (g *Grid) ComputeReceivers(h []float32) {
	diag := float32(g.CellM * math.Sqrt2)
	card := float32(g.CellM)
	field.Rows(g.H, func(y0, y1 int) {
		for y := y0; y < y1; y++ {
			for x := 0; x < g.W; x++ {
				i := int32(y*g.W + x)
				if g.isOutlet(int(i)) {
					g.Receiver[i] = i
					g.Length[i] = card
					continue
				}
				best := int32(-1)
				bestJitter := float32(0)
				bestLen := card
				for k := 0; k < 8; k++ {
					nx, ny := x+dx8[k], y+dy8[k]
					if nx < 0 || ny < 0 || nx >= g.W || ny >= g.H {
						continue
					}
					ni := int32(ny*g.W + nx)
					l := card
					if dx8[k] != 0 && dy8[k] != 0 {
						l = diag
					}
					s := (h[i] - h[ni]) / l
					if s <= 0 {
						continue
					}
					// The tie-break, not a change of gradient: the comparison is jittered, the slope that
					// is kept is not, so Length and the stream-power update see the true geometry.
					sj := s * (1 + 1e-3*(hashXY(g.seed, g.worldX(x), g.worldY(y), int32(k)+jitterReceiverTie)-0.5))
					if sj > bestJitter {
						bestJitter, best, bestLen = sj, ni, l
					}
				}
				if best < 0 {
					g.Receiver[i] = i
					g.Length[i] = card
				} else {
					g.Receiver[i] = best
					g.Length[i] = bestLen
				}
			}
		}
	})
}

// BuildStack orders every node after its receiver, by counting donors into a CSR list and then walking it
// depth-first from the roots. O(n), no recursion, and the order is fully determined by the receiver array,
// so it does not vary between runs.
func (g *Grid) BuildStack() {
	n := g.W * g.H
	for i := 0; i <= n; i++ {
		g.donorOff[i] = 0
	}
	for i := 0; i < n; i++ {
		r := g.Receiver[i]
		if int(r) != i {
			g.donorOff[r+1]++
		}
	}
	for i := 0; i < n; i++ {
		g.donorOff[i+1] += g.donorOff[i]
	}
	cursor := g.scratchInt32()
	copy(cursor, g.donorOff[:n])
	for i := 0; i < n; i++ {
		r := g.Receiver[i]
		if int(r) != i {
			g.donorList[cursor[r]] = int32(i)
			cursor[r]++
		}
	}

	g.Stack = g.Stack[:0]
	for i := 0; i < n; i++ {
		if int(g.Receiver[i]) == i {
			g.Stack = append(g.Stack, int32(i))
		}
	}
	// Depth-first: everything already in the stack expands its donors, which land after it.
	for read := 0; read < len(g.Stack); read++ {
		c := g.Stack[read]
		for d := g.donorOff[c]; d < g.donorOff[c+1]; d++ {
			g.Stack = append(g.Stack, g.donorList[d])
		}
	}
}

// Scratch hands out the grid's spare float32 buffer, which is the width of the grid and is dead between
// steps. It is here so a pass that runs once after the solve - the edge-preserving smooth - does not allocate
// a second copy of the height field at planet scale just to have somewhere to write.
func (g *Grid) Scratch() []float32 { return g.scratch }

// scratchInt32 reuses the float32 scratch as int32 storage; same width, and it saves a 12 MB allocation per
// step at the geology grid.
func (g *Grid) scratchInt32() []int32 {
	if cap(g.cursor) < g.W*g.H {
		g.cursor = make([]int32, g.W*g.H)
	}
	return g.cursor[:g.W*g.H]
}

// Accumulate sums drainage area down the stack in reverse, so every node has collected its whole upstream
// catchment before its receiver is reached.
func (g *Grid) Accumulate() {
	cell := float32(g.CellM * g.CellM)
	for i := range g.Area {
		g.Area[i] = cell
	}
	for k := len(g.Stack) - 1; k >= 0; k-- {
		i := g.Stack[k]
		r := g.Receiver[i]
		if r != i {
			g.Area[r] += g.Area[i]
		}
	}
}

// StreamPower is the implicit update, walked up the stack. uplift is in metres per year and k is the local
// erodibility; either may be nil for a uniform value.
func (g *Grid) StreamPower(h []float32, uplift, k []float32, p Params) {
	dt := p.DtYr
	linear := math.Abs(p.N-1) < 1e-9
	for _, i := range g.Stack {
		if g.fixed[i] {
			continue // base level is fixed: no uplift, no erosion
		}
		u := 0.0
		if uplift != nil {
			u = float64(uplift[i])
		}
		r := g.Receiver[i]
		if r == i {
			// A local minimum that the last flood has not reached yet. It still rises: skipping uplift here
			// freezes exactly the cells that differential uplift is busy pushing up, which quietly removes
			// the basins from the landscape between floods.
			h[i] += float32(dt * u)
			continue
		}
		kk := p.K
		if k != nil {
			kk *= float64(k[i])
		}
		hr := float64(h[r])
		hi := float64(h[i]) + dt*u
		area := float64(g.Area[i])
		if p.CriticalAreaM2 > 0 && area < p.CriticalAreaM2 {
			// Below the channel head, incision is tapered rather than switched off. Switching it off
			// entirely is what broke: the material had nowhere to go, because the only remaining transport
			// was landsliding, which caps slope but not height, and the map grew until 22 % of it clipped
			// the elevation range. A taper suppresses the fine dissection that makes lowlands look like
			// small mountains while still letting the hillslope shed its uplift into the network.
			kk *= math.Pow(area/p.CriticalAreaM2, p.ChannelTaper)
		}
		a := math.Pow(area, p.M)
		l := float64(g.Length[i])

		var next float64
		if linear {
			f := kk * dt * a / l
			next = (hi + f*hr) / (1 + f)
		} else {
			next = newtonStreamPower(hi, hr, kk*dt*a, l, p.N)
		}
		// A node may never fall below what it drains into; the implicit form only guarantees that while
		// uplift has not raised the receiver past it.
		if next < hr {
			next = hr
		}
		h[i] = float32(next)
	}
}

// newtonStreamPower solves h - hi + c*((h-hr)/l)^n = 0 for n != 1. Five iterations from the linear answer is
// comfortably enough at the exponents anyone actually uses; it is here so the manifest's n is not a lie.
func newtonStreamPower(hi, hr, c, l, n float64) float64 {
	f := c / l
	h := (hi + f*hr) / (1 + f) // the n = 1 answer, as a starting point
	for iter := 0; iter < 5; iter++ {
		d := h - hr
		if d < 0 {
			d = 0
		}
		s := d / l
		fx := h - hi + c*math.Pow(s, n)/1
		dfx := 1 + c*n*math.Pow(s, n-1)/l
		if dfx == 0 {
			break
		}
		step := fx / dfx
		h -= step
		if math.Abs(step) < 1e-6 {
			break
		}
	}
	if h < hr {
		h = hr
	}
	return h
}

// Diffuse is hillslope diffusion, which rounds the divides and stops every channel head from being a needle.
//
// Explicit five-point diffusion is stable only while D*dt/dx² <= 0.25, and the defaults sit above that
// (0.02 m²/yr at dt 1500 on 8 m cells is 0.47), so it sub-steps rather than quietly going unstable. This is
// the sort of thing that shows up as a checkerboard three thousand steps in.
func (g *Grid) Diffuse(h []float32, d, dt float64) {
	if d <= 0 || dt <= 0 {
		return
	}
	dx2 := g.CellM * g.CellM
	total := d * dt / dx2
	sub := int(math.Ceil(total / 0.2))
	if sub < 1 {
		sub = 1
	}
	alpha := float32(total / float64(sub))
	src := h
	tmp := g.scratch[:len(h)]
	for s := 0; s < sub; s++ {
		field.Rows(g.H, func(y0, y1 int) {
			for y := y0; y < y1; y++ {
				for x := 0; x < g.W; x++ {
					i := y*g.W + x
					if g.fixed[i] {
						tmp[i] = src[i]
						continue
					}
					c := src[i]
					lap := clampAt(src, g.W, g.H, x-1, y) + clampAt(src, g.W, g.H, x+1, y) +
						clampAt(src, g.W, g.H, x, y-1) + clampAt(src, g.W, g.H, x, y+1) - 4*c
					tmp[i] = c + alpha*lap
				}
			}
		})
		copy(src, tmp)
	}
}

func clampAt(a []float32, w, h, x, y int) float32 {
	if x < 0 {
		x = 0
	} else if x >= w {
		x = w - 1
	}
	if y < 0 {
		y = 0
	} else if y >= h {
		y = h - 1
	}
	return a[y*w+x]
}

// SetCancel gives the solve a way to be abandoned part way through.
//
// It is checked once a step rather than inside one, which is the right granularity: a step is milliseconds on
// a small region and a couple of seconds on a big one, so the longest a caller waits is one step, and nothing
// inside a step is safe to leave half done. The height field is left wherever the solve had got to, which is
// what a cancelled run means - it is not a checkpoint and nothing downstream should read it as one.
func (g *Grid) SetCancel(ch <-chan struct{}) { g.cancel = ch }

// Run is the whole solve. Progress is reported through log, which is what a five-minute budget needs to be
// steerable: a run that is going wrong should say so at step 500, not at the end.
func (g *Grid) Run(h []float32, uplift, k []float32, p Params, log func(step int, total int, elapsedPct float64)) {
	fill := p.FillEvery
	if fill < 1 {
		fill = 1
	}
	for step := 0; step < p.Steps; step++ {
		if g.cancel != nil {
			select {
			case <-g.cancel:
				return
			default:
			}
		}
		if step%fill == 0 {
			g.FillDepressions(h, 1e-3)
		}
		g.ComputeReceivers(h)
		g.BuildStack()
		if p.MFDExponent > 0 {
			g.AccumulateMFD(h, p.MFDExponent)
		} else {
			g.Accumulate()
		}
		g.StreamPower(h, uplift, k, p)
		if p.CriticalSlope > 0 {
			// The clamp still runs, and it still has to: a belt rising at millimetres a year asks for slopes
			// no bounded-flux transport law can hold, which is a fact about the forcing and not about the
			// scheme. What changes is the order and who gets the last word. The clamp cuts along the eight
			// D8 directions and leaves grid-aligned pyramid faces; nonlinear diffusion then runs over the
			// result with a symmetric five-point stencil and rounds them off before the next step sees them.
			//
			// Running the clamp only once at the end was tried and is worse: a thousand steps of unclamped
			// growth arrive at it together, so it cuts deeply, and nothing runs afterwards to soften what it
			// cut. The facets came back in the summits. Little and often, with diffusion last, is what keeps
			// the constraint without printing the stencil.
			if p.TalusSlope > 0 && p.ThermalEvery > 0 && step%p.ThermalEvery == 0 {
				g.ClampToRepose(h, p.TalusSlope)
			}
			g.DiffuseNonlinear(h, p.Diffusion, p.CriticalSlope, p.SlopeCap, p.DtYr, p.MaxHillslopeSub)
		} else {
			g.Diffuse(h, p.Diffusion, p.DtYr)
			if p.TalusSlope > 0 && p.ThermalEvery > 0 && step%p.ThermalEvery == 0 {
				// The constraint first, which actually binds, then the transport, which puts scree at the
				// foot of what the constraint cut.
				g.ClampToRepose(h, p.TalusSlope)
				if p.ThermalPasses > 0 {
					thermal.Apply(h, g.W, g.H, g.CellM, p.TalusSlope, p.ThermalPasses, g.fixed, g.scratch)
				}
			}
		}
		if log != nil && p.Steps >= 10 && step%(p.Steps/10) == 0 {
			log(step, p.Steps, float64(step)/float64(p.Steps)*100)
		}
	}
	// One last fill so the result has no closed pits to hand to the detail passes, and one last routing so
	// Area and Receiver describe the surface that is actually returned.
	g.FillDepressions(h, 1e-3)
	g.ComputeReceivers(h)
	g.BuildStack()
	g.Accumulate()
}