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 stencil, so there is no cut, no facet and no preferred direction.
//
// The stencil is nine-point, and it has to be. Run's design is that the clamp cuts and this rounds off what
// it cut before the next step sees it - but the clamp cuts along all eight neighbour directions and a
// five-point stencil transports across four, so it cannot touch a diagonally-cut facet at all. That was not a
// refinement, it was a hole in the stated design. The weights are 4/6 cardinal and 1/6 diagonal, which is the
// isotropic nine-point Laplacian: on h = (a/2)(x^2+y^2) the cardinal faces sum to 2ad^2 and the diagonals to
// 4ad^2, so (1/6)(8ad^2 + 4ad^2) = 2ad^2 = dx^2 * grad2(h), exactly what the five-point gave. coeff is
// therefore unchanged. A diagonal face is sqrt(2) further away, so it carries its own critical height
// difference; leaving that out would make every diagonal read as 1.41 times its true S/Sc.
//
// 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 Courant number the sub-stepping aims for. The worst mode is the checkerboard: on the five-point
	// stencil its cardinal faces sum to -8*amp and the amplification is 1 - 8*coeff, stable to coeff 0.25;
	// on the nine-point the diagonals cancel and (4/6)*(-8*amp) leaves 1 - 5.333*coeff, stable to 0.375. Both
	// targets keep the same 1.25x margin under their own limit, and the extra room is most of what pays for
	// the four extra faces. Raising the target without the 4/6 and 1/6 weights, or adding the faces without
	// raising the target, is a scheme that checkerboards a few hundred steps in - which is the failure the
	// budget note below is about, and it does not announce itself.

	// 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) * subTargetNine * 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 / subTargetNine))
	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)
	dhCritDiag := float32(sc * dx * math.Sqrt2)

	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 all eight 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.
					card := 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)
					diag := flux(clampAt(src, g.W, g.H, x-1, y-1)-c, dhCritDiag, uCap) +
						flux(clampAt(src, g.W, g.H, x+1, y-1)-c, dhCritDiag, uCap) +
						flux(clampAt(src, g.W, g.H, x-1, y+1)-c, dhCritDiag, uCap) +
						flux(clampAt(src, g.W, g.H, x+1, y+1)-c, dhCritDiag, uCap)
					tmp[i] = c + coeff*(nineCard*card+nineDiag*diag)
				}
			}
		})
		copy(src, tmp)
	}
}

// The isotropic nine-point Laplacian's weights, and the Courant target its stability allows. See
// DiffuseNonlinear.
const (
	nineCard      = float32(4.0 / 6.0)
	nineDiag      = float32(1.0 / 6.0)
	subTargetNine = 0.3
)

// 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, over every face the stencil transports
// across - which since the stencil went to nine points means the diagonals too. A diagonal face is compared
// against its own critical height difference, sqrt(2) larger, so what comes back is a slope ratio either way.
//
// What this number is for is worth being exact about, because it looks like physics and is not. It bounds the
// stiffening for the whole call, and the flux law only caps a face when that face exceeds the bound - so on a
// grid whose steepest face is the bound, no face is capped and the value has no effect on any cell. Its one
// real job is to decide how many sub-steps the call pays for, which is a cost question. Where it does reach
// the physics is when the sub-step budget cannot buy the grid's own maximum; the cap is then lowered to what
// the budget affords, and that value is the manifest's - D, dt, the cell and MaxHillslopeSub - and not the
// grid's, so a planet decomposed two ways still answers the same.
func maxSlopeRatio(h []float32, w, hgt int, sc, cellM float64) float64 {
	invDiag := float32(1 / math.Sqrt2)
	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+1]-c) * invDiag; dv > maxDiff {
						maxDiff = dv
					}
				}
				if y > 0 {
					if dv := abs32(h[i-w+1]-c) * invDiag; 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
}