package fluvial

import (
	"math"

	"salty/terrain/internal/field"
)

// Nonlinear hillslope transport: q = D*S / (1 - (S/Sc)^2), the Roering form.
//
// It replaces the pair of patches that stood in for a hillslope law — linear diffusion, which does not care
// how steep the ground is, and ClampToRepose, which cares about nothing else — with one process that does
// both jobs. As S goes to zero it *is* linear diffusion, q -> D*S, so the divides in the lowlands round over
// exactly as before. As S approaches Sc the flux diverges, so the slope approaches Sc and never reaches it.
//
// The difference that shows is not in the numbers, it is in the shape. ClampToRepose cuts each cell down to
// talus*distance along one of eight neighbour directions and pops cells in grid order within an elevation
// bucket, so what it leaves is pyramids with faces aligned to the grid — the blocky, ruler-cut facets that
// are visible in any preview of a mountain belt here. Nothing about that is geology; it is the D8 stencil
// printed onto the landscape. Nonlinear diffusion approaches the same limiting angle *asymptotically* and
// through a symmetric five-point stencil, so there is no cut, no facet and no preferred direction.
//
// It is also mass-conserving, which the clamp is not: the flux out of one cell is the flux into its
// neighbour by construction, so material shed from a divide arrives at the foot of the slope rather than
// being deleted.
//
// # The cost, and the honest limit
//
// The catch is stiffness. The tangent of the flux law, which is what sets the explicit time-step limit, is
//
//	D_eff(S) = D * (1 + u^2) / (1 - u^2)^2,  u = S/Sc
//
// and that goes to infinity at u = 1. At the defaults the linear Courant number D*dt/dx^2 is already 0.29,
// so u = 0.9 alone asks for about seventy sub-steps a step and u = 0.95 for nearly three hundred. That is
// not affordable over a thousand steps, so the stiffening is bounded: u is capped at SlopeCap, and if even
// that exceeds MaxSubSteps the cap is lowered further to whatever the budget affords. The sub-step count is
// then derived from the cap that was actually used, so the scheme stays inside its stability limit whatever
// happens — it degrades by transporting less on the steepest ground, never by going unstable.
//
// Ground steeper than the cap therefore relaxes at a finite rate instead of an unbounded one, and on a
// mountain belt rising at 2 mm/yr that is not fast enough on its own. ClampToRepose stays for exactly that,
// as a safety pass rather than as the process that shapes the land: see Run.
func (g *Grid) DiffuseNonlinear(h []float32, d, sc, slopeCap, dt float64, maxSub int) {
	if d <= 0 || dt <= 0 {
		return
	}
	if sc <= 0 {
		g.Diffuse(h, d, dt) // no critical slope configured: the linear law, unchanged
		return
	}
	if slopeCap <= 0 || slopeCap >= 1 {
		slopeCap = 0.9
	}
	if maxSub < 1 {
		maxSub = 1
	}
	dx := g.CellM
	dx2 := dx * dx

	// The steepest ground on the grid bounds D_eff for the whole call. Uplift is not applied in here and
	// diffusion only relaxes slopes, so nothing can get steeper part-way through and invalidate the bound.
	u := math.Min(g.maxSlopeRatio(h, sc), slopeCap)
	f := stiffness(u)
	// What the sub-step budget can pay for. Lowering the cap rather than truncating the sub-step count is
	// what keeps this stable: a truncated count leaves alpha above 0.25 and the surface checkerboards a few
	// hundred steps later, which is precisely the sort of failure that does not announce itself.
	if budget := float64(maxSub) * 0.2 * dx2 / (d * dt); f > budget {
		f = budget
		u = invStiffness(f)
	}
	if f < 1 {
		// The budget cannot buy even the linear law. It is not optional: D*dt/dx^2 alone may need several
		// sub-steps and going without them is an unstable scheme, so MaxSubSteps bounds the nonlinear
		// *enhancement* and never the stability floor underneath it.
		f = 1
		u = 0
	}
	sub := int(math.Ceil(d * f * dt / dx2 / 0.2))
	if sub < 1 {
		sub = 1
	}
	dtSub := dt / float64(sub)
	coeff := float32(d * dtSub / dx2)
	uCap := float32(u)
	// The height difference across one cell that *is* Sc. flux works in height differences rather than
	// slopes, so the cell spacing has to be folded into the critical value here; leaving it out makes u a
	// factor of dx too large, which pins every face against the cap and quietly turns the whole law into
	// linear diffusion with a constant multiplier.
	dhCrit := float32(sc * dx)

	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] // base level: held, and whatever arrives here has left the system
						continue
					}
					c := src[i]
					// The net inflow over the four faces. Each face is evaluated from both of its cells,
					// which costs twice and buys a gather: no two goroutines ever write the same cell.
					net := flux(clampAt(src, g.W, g.H, x-1, y)-c, dhCrit, uCap) +
						flux(clampAt(src, g.W, g.H, x+1, y)-c, dhCrit, uCap) +
						flux(clampAt(src, g.W, g.H, x, y-1)-c, dhCrit, uCap) +
						flux(clampAt(src, g.W, g.H, x, y+1)-c, dhCrit, uCap)
					tmp[i] = c + coeff*net
				}
			}
		})
		copy(src, tmp)
	}
}

// flux is q/D for one face, in height differences rather than slopes: one factor of the cell spacing cancels
// against the divergence and is carried in coeff instead. dhCrit is the height difference that corresponds to
// Sc across one cell, so dh/dhCrit is exactly S/Sc. u is capped so the denominator cannot reach zero.
func flux(dh, dhCrit, uCap float32) float32 {
	u := dh / dhCrit
	if u < 0 {
		u = -u
	}
	if u > uCap {
		u = uCap
	}
	return dh / (1 - u*u)
}

// stiffness is D_eff/D at a given u = S/Sc: the factor by which the nonlinear law shortens the stable step.
func stiffness(u float64) float64 {
	q := 1 - u*u
	return (1 + u*u) / (q * q)
}

// invStiffness inverts it. Bisection because stiffness is monotone on [0,1) and this runs once per call, so
// there is nothing to gain from being cleverer and something to lose from being wrong.
func invStiffness(f float64) float64 {
	if f <= 1 {
		return 0
	}
	lo, hi := 0.0, 0.999999
	for i := 0; i < 60; i++ {
		mid := (lo + hi) / 2
		if stiffness(mid) < f {
			lo = mid
		} else {
			hi = mid
		}
	}
	return lo
}

// maxSlopeRatio is the steepest face on the grid as a fraction of Sc. Cardinal neighbours only, because those
// are the faces the five-point stencil actually transports across.
func maxSlopeRatio(h []float32, w, hgt int, sc, cellM float64) float64 {
	var maxDiff float32
	for y := 0; y < hgt; y++ {
		for x := 0; x < w; x++ {
			i := y*w + x
			c := h[i]
			if x+1 < w {
				if dv := abs32(h[i+1] - c); dv > maxDiff {
					maxDiff = dv
				}
			}
			if y+1 < hgt {
				if dv := abs32(h[i+w] - c); dv > maxDiff {
					maxDiff = dv
				}
			}
		}
	}
	return float64(maxDiff) / cellM / sc
}

func (g *Grid) maxSlopeRatio(h []float32, sc float64) float64 {
	return maxSlopeRatio(h, g.W, g.H, sc, g.CellM)
}

func abs32(v float32) float32 {
	if v < 0 {
		return -v
	}
	return v
}