package main

// `mapart biomes` turns the two categorical maps of the planet into the smooth 0..1 masks the landscape's
// paint layers are built from: one greyscale PNG per biome, at the source's own resolution, already blurred.
//
// Why this is here and not in generate_region_tiles.py, which is what consumes it. The class source is the
// painting - 7738 x 3761 of RGB - and the engine's Python has no image library that can decode it: heightmap_io
// is greyscale-only and unfilters a byte at a time. The same constraint that made this tool Go in the first
// place. What it hands back is 8-bit greyscale, which is heightmap_io's fast path.
//
// Why blurred here rather than per tile. A biome boundary has to be a gradient or the ground has a drawn line
// on it, and a blur computed per tile is a blur that disagrees with itself across a tile seam unless every tile
// carries a margin the width of the blur - 200 vertices at 400 m and 2 m quads, a third more area on every one
// of ninety-eight tiles. Blurring once, globally, in the source's own pixels, makes the field smooth *before*
// anything samples it, so a tile can read it with plain bilinear interpolation at its global coordinates and
// two tiles agree at a shared vertex by construction. It is the same reasoning as sampling the height by global
// position, applied a step earlier.
//
// Both sources are read by identity in normalised u,v, which is measured rather than assumed: the painting is
// 7738 x 3761 and the heightmap is 8192 x 4096, and the two candidate registrations were tested against each
// other on land/sea agreement - identity scored 98.09% against 96.14% for the alternative, and won in every
// latitude band including the polar ones, which is where a vertical scale error shows first.

import (
	"fmt"
	"image"
	"math"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
)

// ---------------------------------------------------------------------------------------------------------
// The two legends

type legendClass struct {
	Name string `json:"name"`
	RGB  []int  `json:"rgb"`
	Sea  bool   `json:"sea"`
}

type paintLegend struct {
	Classes []legendClass `json:"classes"`
}

// koppenClass is one row of Tools/Orogen/js/koppen.js.
//
// Parsed out of the JavaScript rather than copied into a JSON beside it, because the browser twin is where that
// table is *used* and two copies of a palette is how one of them goes stale. The parse is strict and the caller
// checks the count: a table that has moved or been reformatted fails loudly here rather than silently matching
// every pixel to the wrong biome.
type koppenClass struct {
	Code string
	Name string
	SRGB [3]uint8
}

// The colours in koppen.js are linear 0..1 and the browser writes them through an sRGB encode, which is why a
// naive read of the exported PNG matches nothing: the observed "ocean" is 147,177,211 where the table says
// 0.29,0.44,0.65. Encoding the table the same way reproduces every observed colour to within 1.4/255.
func linearToSRGB8(c float64) uint8 {
	var s float64
	if c <= 0.0031308 {
		s = c * 12.92
	} else {
		s = 1.055*math.Pow(c, 1/2.4) - 0.055
	}
	return uint8(math.Round(math.Max(0, math.Min(1, s)) * 255))
}

var koppenRow = regexp.MustCompile(`\{\s*code:\s*'([^']+)'\s*,\s*name:\s*'([^']+)'\s*,\s*color:\s*\[([^\]]+)\]`)

func parseKoppen(path string) ([]koppenClass, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	matches := koppenRow.FindAllStringSubmatch(string(data), -1)
	out := make([]koppenClass, 0, len(matches))
	for _, m := range matches {
		parts := strings.Split(m[3], ",")
		if len(parts) != 3 {
			continue
		}
		var rgb [3]uint8
		ok := true
		for i, p := range parts {
			f, err := strconv.ParseFloat(strings.TrimSpace(p), 64)
			if err != nil {
				ok = false
				break
			}
			rgb[i] = linearToSRGB8(f)
		}
		if ok {
			out = append(out, koppenClass{Code: m[1], Name: m[2], SRGB: rgb})
		}
	}
	if len(out) < 20 {
		return nil, fmt.Errorf("%s: parsed only %d Koppen classes, expected about 31 - has the table been "+
			"reformatted? Matching against a partial palette would put every unmatched pixel in the wrong biome", path, len(out))
	}
	return out, nil
}

// ---------------------------------------------------------------------------------------------------------
// Classification

// classify turns an image into a per-pixel index into `palette`, by exact match where possible and nearest
// colour otherwise. It returns how far the worst pixel had to travel: on the painting that is 0, because a
// painted map is made of the legend's own colours and nothing else, and a number above a few units means the
// image is a *render* of a classification rather than the classification itself - which is the difference
// between data and a picture of data, and the reason the Orogen class export is not used here.
func classify(im image.Image, palette [][3]uint8) ([]uint8, float64, float64, error) {
	if len(palette) == 0 || len(palette) > 255 {
		return nil, 0, 0, fmt.Errorf("classify: %d palette entries, need 1..255", len(palette))
	}
	read, err := rgbAccess(im)
	if err != nil {
		return nil, 0, 0, err
	}
	b := im.Bounds()
	w, h := b.Dx(), b.Dy()
	out := make([]uint8, w*h)
	exact := map[[3]uint8]uint8{}
	for i, p := range palette {
		exact[p] = uint8(i)
	}
	// Cached per distinct colour, but `far` counts *pixels*: an image whose boundaries are anti-aliased has
	// few distinct intermediate colours and a great many pixels wearing them, and it is the pixel count that
	// says whether the classification can be trusted.
	type match struct {
		index uint8
		far   bool
	}
	worst := 0.0
	far := 0
	cache := map[[3]uint8]match{}
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			r, g, bl := read(b.Min.X+x, b.Min.Y+y)
			key := [3]uint8{r, g, bl}
			if idx, ok := exact[key]; ok {
				out[y*w+x] = idx
				continue
			}
			m, ok := cache[key]
			if !ok {
				best, bestD := uint8(0), math.MaxFloat64
				for i, p := range palette {
					d := sq(float64(r)-float64(p[0])) + sq(float64(g)-float64(p[1])) + sq(float64(bl)-float64(p[2]))
					if d < bestD {
						bestD, best = d, uint8(i)
					}
				}
				d := math.Sqrt(bestD)
				if d > worst {
					worst = d
				}
				m = match{index: best, far: d > farThreshold}
				cache[key] = m
			}
			if m.far {
				far++
			}
			out[y*w+x] = m.index
		}
	}
	return out, worst, 100 * float64(far) / float64(w*h), nil
}

// How far a colour may be from the nearest palette entry before it counts as "not really that class". Eight is
// well past the rounding a PNG encode can introduce and well short of the distance between two palette colours.
const farThreshold = 8.0

func sq(v float64) float64 { return v * v }

// ---------------------------------------------------------------------------------------------------------
// Masks

// mask builds the 0/1 indicator of a set of palette indices, then blurs it. Two box passes rather than one:
// a single box leaves visible straight edges where the kernel enters a blob, and two is a good enough triangle
// filter for ground that is about to be broken up by noise anyway.
//
// X wraps and Y clamps, because the source is a cylinder with no route over its poles - the same rule the map
// view's projection follows.
func mask(index []uint8, w, h int, want map[uint8]bool, radius int) []float32 {
	field := make([]float32, w*h)
	for i, v := range index {
		if want[v] {
			field[i] = 1
		}
	}
	if radius < 1 {
		return field
	}
	field = boxBlur(field, w, h, radius)
	return boxBlur(field, w, h, radius)
}

func boxBlur(src []float32, w, h, radius int) []float32 {
	tmp := make([]float32, w*h)
	out := make([]float32, w*h)
	window := float32(2*radius + 1)

	// Horizontal, wrapping. A running sum, so the cost is per pixel and not per pixel per tap.
	for y := 0; y < h; y++ {
		row := src[y*w : (y+1)*w]
		var sum float32
		for k := -radius; k <= radius; k++ {
			sum += row[((k%w)+w)%w]
		}
		dst := tmp[y*w : (y+1)*w]
		for x := 0; x < w; x++ {
			dst[x] = sum / window
			sum -= row[(((x-radius)%w)+w)%w]
			sum += row[(((x+radius+1)%w)+w)%w]
		}
	}

	// Vertical, clamping.
	at := func(x, y int) float32 {
		if y < 0 {
			y = 0
		} else if y >= h {
			y = h - 1
		}
		return tmp[y*w+x]
	}
	for x := 0; x < w; x++ {
		var sum float32
		for k := -radius; k <= radius; k++ {
			sum += at(x, k)
		}
		for y := 0; y < h; y++ {
			out[y*w+x] = sum / window
			sum -= at(x, y-radius)
			sum += at(x, y+radius+1)
		}
	}
	return out
}

// ---------------------------------------------------------------------------------------------------------
// The command

type biomeLayer struct {
	Name    string   `json:"name"`
	Rule    string   `json:"rule"`
	Classes []string `json:"classes"`
	Koppen  []string `json:"koppen"`
	Enabled *bool    `json:"enabled"`
}

type biomeConfig struct {
	ClassImage   string  `json:"class_image"`
	ClassLegend  string  `json:"class_legend"`
	ClimateImage string  `json:"climate_image"`
	KoppenJS     string  `json:"koppen_js"`
	MasksDir     string  `json:"masks_dir"`
	BlendM       float64 `json:"blend_m"`
}

type regionLayers struct {
	Biomes biomeConfig  `json:"biomes"`
	Paint  []biomeLayer `json:"paint"`
}

type regionWithLayers struct {
	Region
	Layers regionLayers `json:"layers"`
}

type maskReport struct {
	Layer    string  `json:"layer"`
	Rule     string  `json:"rule"`
	Of       string  `json:"of"`
	Output   string  `json:"output"`
	Width    int     `json:"width"`
	Height   int     `json:"height"`
	RadiusPx int     `json:"radius_px"`
	CoverPct float64 `json:"cover_pct"`
	Enabled  bool    `json:"enabled"`
}

type biomeReport struct {
	When          string       `json:"when"`
	BlendM        float64      `json:"blend_m"`
	ClassWorst    float64      `json:"class_worst_colour_distance"`
	ClassFarPct   float64      `json:"class_unmatched_pct"`
	ClimateWorst  float64      `json:"climate_worst_colour_distance"`
	ClimateFarPct float64      `json:"climate_unmatched_pct"`
	Masks         []maskReport `json:"masks"`
}

func biomes(root string) error {
	regionPath := filepath.Join(root, "RawContent", "World", "Region.json")
	var region regionWithLayers
	if err := readJSON(regionPath, &region); err != nil {
		return err
	}
	cfg := region.Layers.Biomes
	if cfg.ClassImage == "" {
		return fmt.Errorf("%s: layers.biomes has no class_image; nothing to classify", regionPath)
	}
	if cfg.MasksDir == "" {
		cfg.MasksDir = "RawContent/World/Biomes"
	}
	if cfg.BlendM <= 0 {
		cfg.BlendM = 400
	}
	outDir := filepath.Join(root, filepath.FromSlash(cfg.MasksDir))
	if err := os.MkdirAll(outDir, 0o755); err != nil {
		return err
	}

	rep := biomeReport{When: time.Now().UTC().Format(time.RFC3339), BlendM: cfg.BlendM}
	worldW := region.widthM()

	// --- the painted classes
	var legend paintLegend
	if err := readJSON(filepath.Join(root, filepath.FromSlash(cfg.ClassLegend)), &legend); err != nil {
		return err
	}
	classNames := make([]string, len(legend.Classes))
	classPalette := make([][3]uint8, len(legend.Classes))
	for i, c := range legend.Classes {
		if len(c.RGB) != 3 {
			return fmt.Errorf("class %q has no rgb", c.Name)
		}
		classNames[i] = c.Name
		classPalette[i] = [3]uint8{uint8(c.RGB[0]), uint8(c.RGB[1]), uint8(c.RGB[2])}
	}
	classIm, err := load(filepath.Join(root, filepath.FromSlash(cfg.ClassImage)))
	if err != nil {
		return err
	}
	classIndex, classWorst, classFar, err := classify(classIm, classPalette)
	if err != nil {
		return err
	}
	cw, ch := classIm.Bounds().Dx(), classIm.Bounds().Dy()
	rep.ClassWorst, rep.ClassFarPct = classWorst, classFar
	fmt.Printf("class   %-38s %5dx%-5d %d classes, worst distance %.1f, %.3f%% of pixels unmatched\n",
		filepath.Base(cfg.ClassImage), cw, ch, len(classNames), classWorst, classFar)
	if classWorst > 8 {
		fmt.Printf("  WARNING: a painted map is made of its legend's own colours, so this should be 0.\n")
		fmt.Printf("  A number this size means the image is a *render* of a classification rather than the\n")
		fmt.Printf("  classification itself, and every pixel is being snapped to whatever is nearest.\n")
	}

	// --- the Koppen climate, only if a layer asks for it
	var climateIndex []uint8
	var climateCodes []string
	var clw, clh int
	needsClimate := false
	for _, l := range region.Layers.Paint {
		if l.Rule == "climate" {
			needsClimate = true
		}
	}
	if needsClimate {
		if cfg.ClimateImage == "" || cfg.KoppenJS == "" {
			return fmt.Errorf("a layer has rule \"climate\" but layers.biomes has no climate_image/koppen_js")
		}
		kop, err := parseKoppen(filepath.Join(root, filepath.FromSlash(cfg.KoppenJS)))
		if err != nil {
			return err
		}
		climatePalette := make([][3]uint8, len(kop))
		climateCodes = make([]string, len(kop))
		for i, k := range kop {
			climatePalette[i] = k.SRGB
			climateCodes[i] = k.Code
		}
		climateIm, err := load(filepath.Join(root, filepath.FromSlash(cfg.ClimateImage)))
		if err != nil {
			return err
		}
		var worst, farPct float64
		climateIndex, worst, farPct, err = classify(climateIm, climatePalette)
		if err != nil {
			return err
		}
		clw, clh = climateIm.Bounds().Dx(), climateIm.Bounds().Dy()
		rep.ClimateWorst, rep.ClimateFarPct = worst, farPct
		fmt.Printf("climate %-38s %5dx%-5d %d Koppen classes, worst distance %.1f, %.3f%% of pixels unmatched\n",
			filepath.Base(cfg.ClimateImage), clw, clh, len(kop), worst, farPct)
		// A render of a classification has anti-aliased boundaries, and a pixel halfway between two palette
		// colours is snapped to whichever is nearer - arbitrary, but only ever a pixel or two wide, and the
		// mask is blurred by tens of pixels afterwards. A large fraction would mean something else is wrong.
		if farPct > 5 {
			fmt.Printf("  WARNING: %.2f%% of the climate map is not close to any Koppen colour. Boundary\n", farPct)
			fmt.Printf("  anti-aliasing accounts for a fraction of a per cent; this is too much for that.\n")
		}
	}

	indexOf := func(names []string, want string) int {
		for i, n := range names {
			if n == want {
				return i
			}
		}
		return -1
	}

	for _, layer := range region.Layers.Paint {
		if layer.Rule != "class" && layer.Rule != "climate" {
			continue // slope, altitude, beach and remainder are derived per tile from the height
		}
		enabled := layer.Enabled == nil || *layer.Enabled

		var index []uint8
		var w, h int
		var names, wanted []string
		if layer.Rule == "class" {
			index, w, h, names, wanted = classIndex, cw, ch, classNames, layer.Classes
		} else {
			index, w, h, names, wanted = climateIndex, clw, clh, climateCodes, layer.Koppen
		}
		want := map[uint8]bool{}
		for _, n := range wanted {
			i := indexOf(names, n)
			if i < 0 {
				return fmt.Errorf("layer %q asks for %q, which is not in the %s legend (%s)",
					layer.Name, n, layer.Rule, strings.Join(names, ", "))
			}
			want[uint8(i)] = true
		}
		if len(want) == 0 {
			return fmt.Errorf("layer %q has rule %q but names no classes", layer.Name, layer.Rule)
		}

		// The blur radius is in this image's own pixels, because the two sources are not the same resolution.
		metresPerPx := worldW / float64(w)
		radius := int(math.Round(cfg.BlendM / metresPerPx / 2))
		field := mask(index, w, h, want, radius)

		var cover float64
		for _, v := range field {
			cover += float64(v)
		}
		cover = 100 * cover / float64(len(field))

		out := image.NewGray(image.Rect(0, 0, w, h))
		for i, v := range field {
			out.Pix[i] = uint8(math.Round(float64(clamp01(v)) * 255))
		}
		name := "mask_" + strings.ToLower(layer.Name) + ".png"
		if err := writePNG(filepath.Join(outDir, name), out); err != nil {
			return err
		}
		state := ""
		if !enabled {
			state = "  (not enabled yet)"
		}
		fmt.Printf("  %-10s %-8s %-28s -> %-22s r=%3d px  %5.2f%% cover%s\n",
			layer.Name, layer.Rule, strings.Join(wanted, "+"), name, radius, cover, state)
		rep.Masks = append(rep.Masks, maskReport{
			Layer: layer.Name, Rule: layer.Rule, Of: strings.Join(wanted, "+"), Output: name,
			Width: w, Height: h, RadiusPx: radius, CoverPct: cover, Enabled: enabled,
		})
	}

	sort.Slice(rep.Masks, func(i, j int) bool { return rep.Masks[i].Layer < rep.Masks[j].Layer })
	if err := writeJSON(filepath.Join(outDir, "biomes.json"), rep); err != nil {
		return err
	}
	fmt.Printf("%d mask(s) into %s\n", len(rep.Masks), outDir)
	return nil
}

func clamp01(v float32) float32 {
	if v < 0 {
		return 0
	}
	if v > 1 {
		return 1
	}
	return v
}