package planet

import (
	"image/png"
	"math"
	"path/filepath"

	"salty/terrain/internal/field"
	"salty/terrain/internal/uplift"
)

// The maps the plan command writes, and why each one earns its place.
//
// They are rendered by point-sampling the planet's own arrays straight into an image rather than by building
// a full-resolution Field and handing it to field.WriteDataMap. At 78 million cells a Field is 312 MB, and
// the whole point of the plan command is that it costs a minute and nothing else.

// water is the flat blue every map uses for cells that are not land, so the land reads against it.
var water = [3]uint8{24, 44, 74}

// renderRGB point-samples the painted rows of the planet down to width pixels, keeping the aspect, and asks
// at() for a colour per sampled cell. The polar pad is not drawn: it is scaffolding, not world.
// The callback is given the image pixel as well as the planet cell, because a map may have a field of its
// own built at the image's resolution rather than the planet's - the uplift map does, since the massif fabric
// is a field and not a per-class constant.
func renderRGB(in *Inputs, width int, at func(planetIdx, imgIdx int) [3]uint8) (px []uint8, w, h int) {
	p := in.P
	if width <= 0 || width > p.W {
		width = p.W
	}
	paintH := p.PaintH()
	height := int(float64(width)*float64(paintH)/float64(p.W) + 0.5)
	if height < 1 {
		height = 1
	}
	px = make([]uint8, width*height*3)
	field.Rows(height, func(y0, y1 int) {
		for y := y0; y < y1; y++ {
			sy := p.PadY + y*paintH/height
			for x := 0; x < width; x++ {
				sx := x * p.W / width
				c := at(sy*p.W+sx, y*width+x)
				o := (y*width + x) * 3
				px[o], px[o+1], px[o+2] = c[0], c[1], c[2]
			}
		}
	})
	return px, width, height
}

func write(path string, px []uint8, w, h int) error {
	return field.WriteRGB(path, w, h, px, png.DefaultCompression)
}

// WriteClassMap is the first thing to look at when a template comes out wrong: the legend's own colours,
// after the strokes have been dissolved and the poles rescued. If this is not the painting, nothing
// downstream can be.
func WriteClassMap(dir string, in *Inputs, width int) error {
	cols := make([][3]uint8, len(in.Legend.Classes))
	for i, c := range in.Legend.Classes {
		cols[i] = [3]uint8{uint8(c.RGB[0]), uint8(c.RGB[1]), uint8(c.RGB[2])}
	}
	px, w, h := renderRGB(in, width, func(i, _ int) [3]uint8 { return cols[in.Map.Class[i]] })
	return write(filepath.Join(dir, "map_class.png"), px, w, h)
}

// WriteRegionMap shows how the planet was cut up. Each region gets its own hue: its land saturated, the
// water it carries as a margin the same hue dimmed. Water owned by nobody is flat blue.
//
// What to read from it: that the seam-straddling landmass is one colour and not two, that no two landmasses
// which should be independent have been merged into one enormous box, and that the margins are not so wide
// that the regions have swallowed the ocean.
func WriteRegionMap(dir string, in *Inputs, width int) error {
	hues := in.RegionHues()
	px, w, h := renderRGB(in, width, func(i, _ int) [3]uint8 {
		r := in.Part.Owner[i]
		if r < 0 {
			return water
		}
		c := hues[r]
		if in.Map.Sea[i] {
			// The margin: the same region, dimmed, so the box it will be solved in is visible.
			return [3]uint8{c[0]/3 + water[0]/2, c[1]/3 + water[1]/2, c[2]/3 + water[2]/2}
		}
		return c
	})
	return write(filepath.Join(dir, "map_regions.png"), px, w, h)
}

// goldenAngle is the fraction of a turn between consecutive region hues: 1/phi, or 137.507 degrees.
const goldenAngle = 0.6180339887498949

// RegionHues is the colour each region is drawn in, indexed the same way in.Part.Regions is.
//
// Exported so that a caller drawing a key beside the map gets the colours from here rather than
// reimplementing it. A legend that is a second copy of the thing it describes is a legend that will
// eventually be wrong about it.
//
// The hue walks by the golden angle rather than coming out of a hash of the index, which is what it used to
// do. A hash gives *independent* hues, and independent hues collide: the closest pair of the hash's first
// twenty was 8.5 apart in RGB, which is two colours nobody can tell apart, on a map whose entire job is
// answering "is that one landmass or two". Stepping 137.5 degrees is the arrangement that keeps every prefix
// of the sequence as far apart as a sequence can be, and saturation and value then cycle on 3 and 2 so that
// two regions coming round to the same hue still differ in something else. Measured over the same walk: the
// closest pair is 44.0 at twenty regions, 41.9 at twenty-six and 37.7 at forty.
//
// Neither cycle is pushed far. The ocean margin is drawn as this colour thirded and mixed with water, so a
// region that starts dim dims to the same grey-blue as every other dim one.
func (in *Inputs) RegionHues() [][3]uint8 {
	hues := make([][3]uint8, len(in.Part.Regions))
	for i := range hues {
		c := field.HSV(math.Mod(float64(i)*goldenAngle, 1), 0.48+0.17*float64(i%3), 0.96-0.16*float64(i%2))
		hues[i] = [3]uint8{clamp8(c[0]), clamp8(c[1]), clamp8(c[2])}
	}
	return hues
}

// RegionLabels is where to write each region's id over the region map: the centroid of its land, as fractions
// of the drawn map - 0..1 across, 0..1 down the painted rows, the same frame renderRGB draws into.
//
// Exported for the reason RegionHues is. The alternative is a caller reproducing the polar row offset and the
// seam wrap in a second language, and a label half a region away from the region it names is worse than no
// label at all. Colour alone cannot carry this: even at 137.5 degrees a step, forty regions are forty hues and
// a person matching a hue to a swatch is doing work a two-digit number does for them.
//
// Two details it would be wrong to leave out. The mean across is *circular*, because a landmass over the seam
// has land at x=0 and at x=W-1 and a plain average puts its number on the opposite side of the planet. And it
// is sampled on a stride rather than walked whole: this is a place to put a number, the planet is seventy-six
// million cells, and a quarter of a cell of accuracy is not worth a sixteenth of a plan. A region too small to
// catch a sample falls back to the middle of its frame, which is the only thing left to say about it.
func (in *Inputs) RegionLabels() [][2]float64 {
	p := in.P
	out := make([][2]float64, len(in.Part.Regions))
	paintH := p.PaintH()
	if len(out) == 0 || paintH <= 0 {
		return out
	}

	const stride = 4
	cosX := make([]float64, p.W)
	sinX := make([]float64, p.W)
	for x := 0; x < p.W; x++ {
		a := 2 * math.Pi * float64(x) / float64(p.W)
		cosX[x], sinX[x] = math.Cos(a), math.Sin(a)
	}

	sumC := make([]float64, len(out))
	sumS := make([]float64, len(out))
	sumY := make([]float64, len(out))
	n := make([]float64, len(out))
	for y := p.PadY; y < p.PadY+paintH; y += stride {
		row := y * p.W
		for x := 0; x < p.W; x += stride {
			i := row + x
			r := in.Part.Owner[i]
			if r < 0 || in.Map.Sea[i] {
				continue
			}
			sumC[r] += cosX[x]
			sumS[r] += sinX[x]
			sumY[r] += float64(y - p.PadY)
			n[r]++
		}
	}

	for r := range out {
		if n[r] == 0 {
			f := in.Part.Regions[r].Frame
			// X0 can run past W on a seam region and Y0 can reach into the polar pad, so both are brought
			// back into the drawn frame rather than trusted.
			u := math.Mod(float64(f.X0)+float64(f.W)/2, float64(p.W)) / float64(p.W)
			v := (float64(f.Y0-p.PadY) + float64(f.H)/2) / float64(paintH)
			out[r] = [2]float64{clamp01(u), clamp01(v)}
			continue
		}
		a := math.Atan2(sumS[r]/n[r], sumC[r]/n[r])
		if a < 0 {
			a += 2 * math.Pi
		}
		out[r] = [2]float64{a / (2 * math.Pi), clamp01(sumY[r] / n[r] / float64(paintH))}
	}
	return out
}

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

// UpliftScale is the top of the uplift map's ramp, in mm/yr, and the colours along it. Same argument as
// RegionHues: the key comes from the code that drew the picture.
func (in *Inputs) UpliftScale(stops int) (hi float64, ramp [][3]uint8) {
	for _, r := range in.Legend.Rates() {
		if v := float64(r) * 1000; v > hi {
			hi = v
		}
	}
	if hi <= 0 {
		hi = 1
	}
	return hi, sampleRamp(stops, field.Inferno)
}

// ErodibilityScale is the low and high end of the erodibility map's ramp and the colours along it.
func (in *Inputs) ErodibilityScale(stops int) (lo, hi float64, ramp [][3]uint8) {
	lo, hi = in.erodibilityRange()
	return lo, hi, sampleRamp(stops, field.Viridis)
}

// erodibilityRange is what the erodibility map spans: every land class's own multiplier, widened by the rock
// field's extremes wherever a class lets them through. It has to account for the lithology or the ramp would
// clip exactly the variation the field was added to show.
func (in *Inputs) erodibilityRange() (lo, hi float64) {
	lo, hi = 1, 1
	mult := in.M.Pipeline.Lithology.KMultipliers
	rockLo, rockHi := 1.0, 1.0
	if in.M.Planet.LithologyCells() > 0 && len(mult) > 1 {
		rockLo, rockHi = mult[0], mult[0]
		for _, v := range mult {
			rockLo = math.Min(rockLo, v)
			rockHi = math.Max(rockHi, v)
		}
	}
	for i := range in.Legend.Classes {
		c := in.Legend.Classes[i]
		if !c.Land() {
			continue
		}
		k, mix := c.K(), c.LithMix()
		lo = math.Min(lo, k*(1+mix*(rockLo-1)))
		hi = math.Max(hi, k*(1+mix*(rockHi-1)))
	}
	return lo, hi
}

func sampleRamp(stops int, f func(float64) [3]float64) [][3]uint8 {
	if stops < 2 {
		stops = 2
	}
	out := make([][3]uint8, stops)
	for i := range out {
		c := f(float64(i) / float64(stops-1))
		out[i] = [3]uint8{clamp8(c[0]), clamp8(c[1]), clamp8(c[2])}
	}
	return out
}

// WriteUpliftMap is the field everything else is a consequence of, in mm/yr. On the square canvas this map
// would have shown, at a glance and with no arithmetic, that the plains were being raised at mountain rates.
// On a painted planet it is the direct check that the legend's numbers landed where the author painted them.
func WriteUpliftMap(dir string, in *Inputs, width int) error {
	rates := in.Legend.Rates()
	massifFloor, massifFraction := in.Legend.Massifs()
	hi := 0.0
	for _, r := range rates {
		if v := float64(r) * 1000; v > hi {
			hi = v
		}
	}
	if hi <= 0 {
		hi = 1
	}

	// The fabric, at the image's resolution rather than the planet's. It has to be drawn, not left out: with
	// massifs the rate is a field and not a per-class constant, and a map that showed the class rate flat
	// across a landmass would be showing the one thing that is no longer true about it. Building it here
	// costs a couple of million noise samples rather than the planet's seventy-eight.
	var rank *field.Field
	if in.Legend.HasMassifs() {
		u, v := in.renderUV(width)
		rank = uplift.MassifRank(in.P, in.M.Source.Seed, in.M.Planet.MassifCells(), u, v)
	}

	px, w, h := renderRGB(in, width, func(i, img int) [3]uint8 {
		if in.Map.Sea[i] {
			return water
		}
		cl := in.Map.Class[i]
		r := float64(rates[cl])
		if rank != nil && massifFraction[cl] > 0 {
			r = uplift.MassifRate(float64(massifFloor[cl]), r, float64(rank.Data[img]), massifFraction[cl])
		}
		c := field.Inferno(r * 1000 / hi)
		return [3]uint8{clamp8(c[0]), clamp8(c[1]), clamp8(c[2])}
	})
	drawFaults(in, px, w, h)
	return write(filepath.Join(dir, "map_uplift.png"), px, w, h)
}

// faultInk is the colour traces are drawn in: cyan, which appears nowhere in the Inferno ramp underneath, so
// a trace cannot be mistaken for a value.
var faultInk = [3]uint8{80, 240, 255}

// drawFaults strokes every fault trace over a map, as a line.
//
// The *line* rather than the rate it contributes, deliberately. A fault's escarpment is a couple of hundred
// metres wide and this image is a hundred kilometres across, so the thing it changes is a twentieth of a
// pixel and rendering the field would show nothing at all. What an author wants from this map is where the
// faults are and which way they run - the same question `map_regions` answers about the region cuts - and a
// stroked polyline answers it exactly.
func drawFaults(in *Inputs, px []uint8, w, h int) {
	if len(in.Faults) == 0 {
		return
	}
	p := in.P
	sx := float64(w) / p.CircumferenceM()
	sy := float64(h) / p.HeightM()
	set := func(x, y int) {
		if y < 0 || y >= h {
			return
		}
		x = ((x % w) + w) % w // X wraps, because the traces do
		o := (y*w + x) * 3
		px[o], px[o+1], px[o+2] = faultInk[0], faultInk[1], faultInk[2]
	}
	for _, f := range in.Faults {
		for j := 0; j+1 < len(f.PointsM); j++ {
			ax, ay := f.PointsM[j][0]*sx, f.PointsM[j][1]*sy
			bx, by := f.PointsM[j+1][0]*sx, f.PointsM[j+1][1]*sy
			steps := int(math.Hypot(bx-ax, by-ay)) + 1
			for k := 0; k <= steps; k++ {
				t := float64(k) / float64(steps)
				set(int(ax+(bx-ax)*t), int(ay+(by-ay)*t))
			}
		}
	}
}

// renderUV is the world coordinates of the cells renderRGB will point-sample, in the same order it writes
// them. Exactly the sampled cells rather than an even walk across the map: a diagnostic that showed the
// fabric half a pixel from where the class was read would be a diagnostic nobody could trust to the pixel.
func (in *Inputs) renderUV(width int) (u, v *field.Field) {
	p := in.P
	if width <= 0 || width > p.W {
		width = p.W
	}
	paintH := p.PaintH()
	height := int(float64(width)*float64(paintH)/float64(p.W) + 0.5)
	if height < 1 {
		height = 1
	}
	cellM := p.CircumferenceM() / float64(width)
	u = field.New(width, height, cellM)
	v = field.New(width, height, cellM)
	for y := 0; y < height; y++ {
		sy := p.PadY + y*paintH/height
		vy := float32(p.YM(sy) / p.NoisePeriodM)
		for x := 0; x < width; x++ {
			i := y*width + x
			u.Data[i] = float32(p.XM(x*p.W/width) / p.NoisePeriodM)
			v.Data[i] = vy
		}
	}
	return u, v
}

// WriteErodibilityMap is where texture inside a range comes from: the multiplier on stream-power K.
func WriteErodibilityMap(dir string, in *Inputs, width int) error {
	ks := in.Legend.Erodibilities()
	mix := in.Legend.LithologyMixes()
	lo, hi := in.erodibilityRange()
	span := hi - lo
	if span < 1e-9 {
		span = 1
	}

	// The rock field at the image's resolution rather than the planet's, the same way and for the same reason
	// the uplift map builds the massif fabric: with lithology the erodibility is a *field*, and a map drawing
	// the class multiplier flat across a landmass would be showing the one thing that is no longer true of it.
	var rock *field.Field
	if cells := in.M.Planet.LithologyCells(); cells > 0 && in.Legend.HasLithology() {
		u, v := in.renderUV(width)
		rock = uplift.RockK(in.P, in.M.Source.Seed, cells, in.M.Pipeline.Lithology.KMultipliers, u, v)
	}

	px, w, h := renderRGB(in, width, func(i, img int) [3]uint8 {
		if in.Map.Sea[i] {
			return water
		}
		cl := in.Map.Class[i]
		k := float64(ks[cl])
		if rock != nil && mix[cl] > 0 {
			k *= 1 + mix[cl]*(float64(rock.Data[img])-1)
		}
		c := field.Viridis((k - lo) / span)
		return [3]uint8{clamp8(c[0]), clamp8(c[1]), clamp8(c[2])}
	})
	return write(filepath.Join(dir, "map_erodibility.png"), px, w, h)
}

func clamp8(v float64) uint8 {
	if v <= 0 {
		return 0
	}
	if v >= 255 {
		return 255
	}
	return uint8(v + 0.5)
}

// WriteOverlayMap draws the annotation layer over a dimmed class map, which is the only way to judge it: a
// mark means nothing on its own and everything relative to the coastline or the range it was drawn against.
//
// It samples the overlay at its own resolution rather than the planet's. Everything else here reads a planet
// array; the overlay is registered to the *template*, so going through the planet grid would resample it
// twice and lose thin strokes on the way.
func WriteOverlayMap(dir string, in *Inputs, width int) error {
	if in.OverlayRaster == nil {
		return nil
	}
	cols := make([][3]uint8, len(in.Overlay.Marks)+1)
	for i, m := range in.Overlay.Marks {
		cols[i+1] = [3]uint8{uint8(m.RGB[0]), uint8(m.RGB[1]), uint8(m.RGB[2])}
	}
	class := make([][3]uint8, len(in.Legend.Classes))
	for i, c := range in.Legend.Classes {
		// Halved towards black, so a full-strength mark on top of it cannot be mistaken for the ground.
		class[i] = [3]uint8{uint8(c.RGB[0] / 2), uint8(c.RGB[1] / 2), uint8(c.RGB[2] / 2)}
	}

	ov := in.OverlayRaster
	p := in.P
	paintH := p.PaintH()
	px, w, h := renderRGB(in, width, func(i, img int) [3]uint8 {
		// The planet cell this pixel came from, turned back into an overlay pixel. Both rasters cover the
		// same painted rows, so the conversion is two ratios and no interpolation.
		x := i % p.W
		y := i/p.W - p.PadY
		ox := x * ov.W / p.W
		oy := y * ov.H / paintH
		if m := ov.At(ox, oy); m != 0 && int(m) < len(cols) {
			return cols[m]
		}
		return class[in.Map.Class[i]]
	})
	return write(filepath.Join(dir, "map_overlay.png"), px, w, h)
}