UnrealPrototyping/Tools/Terrain/internal/planet/overlaygen.go

package planet

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"salty/terrain/internal/field"
	"salty/terrain/internal/fluvial"
	"salty/terrain/internal/manifest"
	"salty/terrain/internal/overlay"
)

// Proposing an overlay from a finished bake.
//
// The annotation layer is hand-painted and starts blank, and the three things most worth putting on it -
// woodland, settlements and the roads between them - are all consequences of ground the author cannot see
// while painting. The classes are painted before the solve exists, and once it does exist it is a
// seventy-six-million-cell heightmap. So this reads the bake back and hands internal/overlay the four
// fields it needs to have an opinion: height, slope, the sea, and where the water collects.
//
// It reads the bake from disk rather than hooking into one, for the same reason `terrain tiles` does: the
// geology is two hours and this is seconds, so an author can regenerate the sheet as often as they like
// against a bake they already have.
//
// **It works at the template's resolution, not the geology grid's.** The overlay is registered to the
// template and must be exactly its size, so generating anywhere else would mean resampling the output - and
// a resampled mark is a blend of two colours, which the classifier reads as a third mark or as nothing. The
// geology is pooled down to the template once, here, and everything downstream is at that scale.

// OverlayGenOptions is what the generator is pointed at.
type OverlayGenOptions struct {
	In      *Inputs
	BakeDir string

	// Replace ignores the overlay already on disk instead of filling in around it. The default is to keep
	// every painted pixel, because regenerating must never cost an author their work; this is the flag for
	// "throw away the last generation and start again", and it says so at the call site.
	Replace bool

	// Seed overrides the manifest's, which is what a re-roll is: the painting fixes where the land is and
	// the seed decides everything it does not - which patch of eligible ground becomes woodland, and which
	// of the equally good sites gets the town.
	Seed int64

	// Existing overrides the overlay loaded from disk. The studio sets it, because the sheet an author is
	// looking at includes strokes they have not saved, and generating around the file instead of around the
	// screen would put marks on top of work that is visibly there.
	Existing *overlay.Raster

	Log func(string, ...any)
}

// GenerateOverlay reads a bake and proposes the marks whose legend asks for them.
func GenerateOverlay(opt OverlayGenOptions) (*overlay.Raster, overlay.GenReport, error) {
	var none overlay.GenReport
	in := opt.In
	log := opt.Log
	if log == nil {
		log = func(string, ...any) {}
	}
	if in.Overlay == nil {
		return nil, none, fmt.Errorf("%s has no overlay legend; set planet.overlay_legend and say which "+
			"marks to generate", in.M.Path)
	}
	wants := false
	for i := range in.Overlay.Marks {
		if in.Overlay.Marks[i].Generate != nil {
			wants = true
			break
		}
	}
	if !wants {
		return nil, none, fmt.Errorf("no mark in %s has a `generate` block, so there is nothing to "+
			"generate. Generation is opt-in per mark; see the overlay section of the templates README",
			in.M.Planet.OverlayLegend)
	}

	ow, oh := in.PaintW, in.PaintH
	cellM := in.P.CircumferenceM() / float64(ow)

	if opt.BakeDir != "" {
		if ok, why := BakeIsOfThisPainting(opt.BakeDir, in.M); !ok {
			log("ignoring %s: %s, so its terrain is not this world's", filepath.Base(opt.BakeDir), why)
			opt.BakeDir = ""
		}
	}

	if opt.BakeDir == "" {
		// No bake: the painting is all there is.
		//
		// Worth supporting rather than refusing, because the first thing an author wants after drawing a
		// world is to see something placed on it, and the bake is two hours away. What survives without a
		// solve is everything the *painting* knows - where the land is, where the sea is, and which class
		// each cell was painted - so the coast still shapes where towns go and the class filters still keep
		// woodland off the ice. What is lost is everything the terrain knows: there are no rivers to sit on,
		// no slope to avoid, and therefore no reason for a road to bend. Say which of the two ran; a draft
		// made this way is a sketch, and reading it as the other is how somebody concludes the generator
		// ignores the terrain.
		log("no bake: generating from the painting alone, so there are no rivers and no slope to read")
		gin := flatInputs(in, ow, oh, cellM, opt.Replace)
		opt.apply(&gin)
		return in.Overlay.Generate(gin)
	}

	if err := CheckBake(opt.BakeDir, in.M, log); err != nil {
		return nil, none, err
	}
	hpath := filepath.Join(opt.BakeDir, "planet_height.png")
	values, gw, gh, err := field.ReadHeightmap(hpath, 0)
	if err != nil {
		return nil, none, fmt.Errorf("%s: %w (run `terrain bake` first)", hpath, err)
	}
	if gw != in.P.W || gh != in.P.PaintH() {
		return nil, none, fmt.Errorf("%s is %dx%d but the manifest describes a %dx%d planet; the bake and "+
			"the manifest have drifted apart", hpath, gw, gh, in.P.W, in.P.PaintH())
	}
	geology := in.M.Decode(values)
	log("read    %s: %d x %d at %.1f m", filepath.Base(hpath), gw, gh, in.P.CellM)

	heightM := poolTo(geology, gw, gh, ow, oh)
	sea := make([]bool, ow*oh)
	seaCells := 0
	for i, v := range heightM {
		if float64(v) < in.M.SeaLevelM {
			sea[i] = true
			seaCells++
		}
	}
	log("overlay grid %d x %d at %.1f m a pixel, %.0f%% sea",
		ow, oh, cellM, 100*float64(seaCells)/float64(ow*oh))

	flow := drainage(heightM, sea, ow, oh, cellM, in.M.Pipeline.Fluvial.MFDExponent)

	// The class each overlay cell was painted, so a mark can be kept off ground its author called ice or
	// desert. Nearest-sampled rather than averaged: a class is a name, and the mean of two names is not one.
	classAt := make([]uint8, ow*oh)
	for y := 0; y < oh; y++ {
		sy := y*gh/oh + in.P.PadY
		for x := 0; x < ow; x++ {
			classAt[y*ow+x] = in.Map.Class[sy*in.P.W+x*gw/ow]
		}
	}
	names := make([]string, len(in.Legend.Classes))
	for i, c := range in.Legend.Classes {
		names[i] = c.Name
	}

	gin := overlay.GenInputs{
		W: ow, H: oh, CellM: cellM,
		HeightM: heightM, Sea: sea, FlowM2: flow,
		ClassAt: classAt, ClassNames: names,
		Seed: in.M.Source.Seed,
	}
	if !opt.Replace {
		gin.Existing = in.OverlayRaster
	}
	opt.apply(&gin)
	return in.Overlay.Generate(gin)
}

// apply puts the caller's overrides onto the inputs, after the world has been read.
func (opt OverlayGenOptions) apply(gin *overlay.GenInputs) {
	if opt.Seed != 0 {
		gin.Seed = opt.Seed
	}
	if opt.Existing != nil {
		gin.Existing = opt.Existing
	}
}

// poolTo box-averages a field onto a smaller grid. The two grids cover exactly the same painted rows, so
// this is a straight ratio in each axis with no registration to work out.
//
// Averaging rather than sampling, because a single sample of an 8 m grid at 12.9 m spacing would alias every
// ridge it stepped over and put woodland in stripes.
func poolTo(src []float32, sw, sh, dw, dh int) []float32 {
	out := make([]float32, dw*dh)
	field.Rows(dh, func(y0, y1 int) {
		for dy := y0; dy < y1; dy++ {
			sy0 := dy * sh / dh
			sy1 := (dy + 1) * sh / dh
			if sy1 <= sy0 {
				sy1 = sy0 + 1
			}
			for dx := 0; dx < dw; dx++ {
				sx0 := dx * sw / dw
				sx1 := (dx + 1) * sw / dw
				if sx1 <= sx0 {
					sx1 = sx0 + 1
				}
				sum, n := 0.0, 0
				for sy := sy0; sy < sy1 && sy < sh; sy++ {
					row := sy * sw
					for sx := sx0; sx < sx1 && sx < sw; sx++ {
						sum += float64(src[row+sx])
						n++
					}
				}
				if n > 0 {
					out[dy*dw+dx] = float32(sum / float64(n))
				}
			}
		}
	})
	return out
}

// drainage is the catchment area per cell at the overlay's resolution, which is what tells a settlement
// where the water is.
//
// Recomputed rather than read from the bake: `map_flow.png` is a picture a few hundred pixels wide, scaled
// for looking at, and what is wanted here is a number per overlay cell. Re-deriving it from the pooled
// height is cheap - one fill and one multiple-flow accumulation - and it is self-consistent with the slope
// and the sea mask beside it, which a resampled flow map would not be.
func drainage(heightM []float32, sea []bool, w, h int, cellM, mfdExp float64) []float32 {
	g := fluvial.NewGrid(w, h, cellM, sea)

	// The fill runs on a copy: it raises every pit to its spill level, which is right for routing water and
	// wrong for everything else here. Slope and the treeline must see the surface the bake actually made.
	filled := make([]float32, len(heightM))
	copy(filled, heightM)
	g.FillDepressions(filled, 1e-3)
	g.AccumulateMFD(filled, mfdExp)

	out := make([]float32, len(heightM))
	copy(out, g.Area)
	return out
}

// OverlaySummary is the run's report, as lines to print.
func OverlaySummary(rep overlay.GenReport, ow, oh int) []string {
	var lines []string
	total := float64(ow * oh)
	if rep.Kept > 0 {
		lines = append(lines, fmt.Sprintf("kept    %d px already painted (%.2f%% of the sheet); "+
			"generation only fills blank ground", rep.Kept, 100*float64(rep.Kept)/total))
	}
	if rep.TreelineM > 0 {
		lines = append(lines, fmt.Sprintf("treeline %.0f m, from the land's own heights", rep.TreelineM))
	}
	for _, m := range rep.Marks {
		switch m.Kind {
		case overlay.GenSettlement:
			line := fmt.Sprintf("  %-14s %d placed, %d px", m.Name, m.Pieces, m.Cells)
			if m.Wanted > m.Pieces {
				// The two things that ration settlements are the spacing and how much flat ground there is,
				// and neither is visible in the output, so the shortfall is said here rather than left to be
				// counted off the sheet.
				line += fmt.Sprintf("  (asked for %d; the spacing or the flat ground ran out)", m.Wanted)
			}
			lines = append(lines, line)
		case overlay.GenRoad:
			lines = append(lines, fmt.Sprintf("  %-14s %d %s, %d px",
				m.Name, m.Pieces, plural(m.Pieces, "link", "links"), m.Cells))
		default:
			lines = append(lines, fmt.Sprintf("  %-14s %d px (%.2f%% of the sheet)",
				m.Name, m.Cells, 100*float64(m.Cells)/total))
		}
	}
	lines = append(lines, fmt.Sprintf("painted %d px (%.2f%% of the sheet)",
		rep.Painted, 100*float64(rep.Painted)/total))
	return lines
}

func plural(n int, one, many string) string {
	if n == 1 {
		return one
	}
	return many
}

// flatInputs builds the generator's world from the painting alone, for a planet that has not been baked.
//
// Height is a flat plateau on land and a flat floor at sea, which makes the slope field zero everywhere and
// the treeline meaningless - both correct rather than approximate, because a world with no solve genuinely
// has no relief to read. Drainage is nil rather than zero, which the generator treats as "rivers contribute
// nothing" instead of "every cell is equally dry"; the difference matters, because a score of zero
// everywhere would still be a score and would silently reweight the coast against it.
func flatInputs(in *Inputs, ow, oh int, cellM float64, replace bool) overlay.GenInputs {
	n := ow * oh
	heightM := make([]float32, n)
	sea := make([]bool, n)
	classAt := make([]uint8, n)

	// The class raster is already at the template's resolution, which is the overlay's, so this is a direct
	// read with no resampling at all.
	for i := 0; i < n && i < len(in.Raster.Class); i++ {
		c := in.Raster.Class[i]
		classAt[i] = c
		if int(c) < len(in.Legend.Classes) && in.Legend.Classes[c].Sea {
			sea[i] = true
			heightM[i] = float32(in.M.SeaLevelM - 100)
		} else {
			heightM[i] = float32(in.M.SeaLevelM + 10)
		}
	}

	names := make([]string, len(in.Legend.Classes))
	for i, c := range in.Legend.Classes {
		names[i] = c.Name
	}

	gin := overlay.GenInputs{
		W: ow, H: oh, CellM: cellM,
		HeightM: heightM, Sea: sea,
		ClassAt: classAt, ClassNames: names,
		Seed: in.M.Source.Seed,
	}
	if !replace {
		gin.Existing = in.OverlayRaster
	}
	return gin
}

// bakeTemplate is the template a bake was made from, or "" when the bake does not record one.
//
// It exists because CheckBake cannot catch this. That check compares the numbers a heightmap is *encoded*
// with - the elevation range, the circumference, the cell size, the seed - and two different paintings of
// the same planet agree on every one of them. So a bake of one world passes every test and is then read as
// another, and the marks come out placed against terrain that is not there: rivers in the wrong valleys,
// towns on coasts that do not exist. Nothing in the output says so, which is what makes it worth a check of
// its own rather than a note in a README.
func bakeTemplate(dir string) string {
	raw, err := os.ReadFile(filepath.Join(dir, "meta.json"))
	if err != nil {
		return ""
	}
	var meta struct {
		Plan struct {
			Template string `json:"template"`
		} `json:"plan"`
	}
	if json.Unmarshal(raw, &meta) != nil {
		return ""
	}
	return meta.Plan.Template
}

// BakeIsOfThisPainting reports whether a bake was made from the template the manifest now names, and why not
// when it was not. A bake that does not record its template is taken on trust, because it predates the
// field; that is said rather than assumed.
func BakeIsOfThisPainting(dir string, m *manifest.Manifest) (bool, string) {
	was := bakeTemplate(dir)
	if was == "" {
		return true, ""
	}
	if filepath.Base(was) == filepath.Base(m.Planet.Template) {
		return true, ""
	}
	return false, fmt.Sprintf("%s was baked from %s and this planet is painted on %s",
		filepath.Base(dir), filepath.Base(was), filepath.Base(m.Planet.Template))
}