class_name TerrainHeightField
extends RefCounted

const SonarMaterial = preload("res://game/domain/sonar/sonar_material.gd")
const SiteTerrainData = preload("res://game/domain/sites/site_terrain_data.gd")

var _data: SiteTerrainData
var _minimum_elevation_m := 0.0
var _maximum_elevation_m := 0.0


static func from_data(data: SiteTerrainData) -> TerrainHeightField:
	if data == null or not data.validation_error().is_empty():
		return null
	var field := TerrainHeightField.new()
	field._data = data
	field._minimum_elevation_m = data.elevations_m[0]
	field._maximum_elevation_m = data.elevations_m[0]
	for elevation: float in data.elevations_m:
		field._minimum_elevation_m = minf(field._minimum_elevation_m, elevation)
		field._maximum_elevation_m = maxf(field._maximum_elevation_m, elevation)
	return field


func columns() -> int:
	return _data.columns


func rows() -> int:
	return _data.rows


func cell_size_m() -> float:
	return _data.cell_size_m


func origin_m() -> Vector2:
	return Vector2(_data.origin_x_m, _data.origin_z_m)


func bounds_m() -> Rect2:
	return Rect2(
		_data.origin_x_m,
		_data.origin_z_m,
		float(_data.columns - 1) * _data.cell_size_m,
		float(_data.rows - 1) * _data.cell_size_m
	)


func contains(x_m: float, z_m: float) -> bool:
	return bounds_m().has_point(Vector2(x_m, z_m))


func minimum_elevation_m() -> float:
	return _minimum_elevation_m


func maximum_elevation_m() -> float:
	return _maximum_elevation_m


## Row-major elevations (index = row * columns + column), the HeightMapShape3D layout.
func elevations_m() -> PackedFloat32Array:
	return _data.elevations_m


## Row-major sonar materials, same layout as elevations_m.
func materials() -> PackedByteArray:
	return _data.materials


func elevation_at_cell(column: int, row: int) -> float:
	return _data.elevations_m[_data.index_of(_clamp_column(column), _clamp_row(row))]


func material_at_cell(column: int, row: int) -> int:
	return _data.materials[_data.index_of(_clamp_column(column), _clamp_row(row))]


func cell_world_position(column: int, row: int) -> Vector3:
	return Vector3(
		_data.origin_x_m + float(column) * _data.cell_size_m,
		elevation_at_cell(column, row),
		_data.origin_z_m + float(row) * _data.cell_size_m
	)


func normal_at_cell(column: int, row: int) -> Vector3:
	return _normal_at(_data.elevations_m, _clamp_column(column), _clamp_row(row))


## World positions of the cells [start_column..end_column] x [start_row..end_row]
## (bounds included, clamped to the grid), row-major. Same values as cell_world_position.
func cell_block_positions(start_column: int, end_column: int, start_row: int, end_row: int) -> PackedVector3Array:
	var first_column := _clamp_column(start_column)
	var last_column := _clamp_column(end_column)
	var first_row := _clamp_row(start_row)
	var last_row := _clamp_row(end_row)
	var elevations := _data.elevations_m
	var columns := _data.columns
	var cell_size := _data.cell_size_m
	var positions := PackedVector3Array()
	positions.resize((last_column - first_column + 1) * (last_row - first_row + 1))
	var index := 0
	for row in range(first_row, last_row + 1):
		var z := _data.origin_z_m + float(row) * cell_size
		var row_start := row * columns
		for column in range(first_column, last_column + 1):
			positions[index] = Vector3(
				_data.origin_x_m + float(column) * cell_size,
				elevations[row_start + column],
				z
			)
			index += 1
	return positions


## Normals of the same cell block, row-major. Same values as normal_at_cell.
func cell_block_normals(start_column: int, end_column: int, start_row: int, end_row: int) -> PackedVector3Array:
	var first_column := _clamp_column(start_column)
	var last_column := _clamp_column(end_column)
	var first_row := _clamp_row(start_row)
	var last_row := _clamp_row(end_row)
	var elevations := _data.elevations_m
	var normals := PackedVector3Array()
	normals.resize((last_column - first_column + 1) * (last_row - first_row + 1))
	var index := 0
	for row in range(first_row, last_row + 1):
		for column in range(first_column, last_column + 1):
			normals[index] = _normal_at(elevations, column, row)
			index += 1
	return normals


func elevation_at(x_m: float, z_m: float) -> float:
	var column_f := clampf(
		(x_m - _data.origin_x_m) / _data.cell_size_m,
		0.0,
		float(_data.columns - 1)
	)
	var row_f := clampf(
		(z_m - _data.origin_z_m) / _data.cell_size_m,
		0.0,
		float(_data.rows - 1)
	)
	var column := mini(int(floorf(column_f)), _data.columns - 2)
	var row := mini(int(floorf(row_f)), _data.rows - 2)
	var tx := column_f - float(column)
	var tz := row_f - float(row)
	var top := lerpf(elevation_at_cell(column, row), elevation_at_cell(column + 1, row), tx)
	var bottom := lerpf(
		elevation_at_cell(column, row + 1),
		elevation_at_cell(column + 1, row + 1),
		tx
	)
	return lerpf(top, bottom, tz)


## Smooth (Catmull-Rom bicubic) elevation: passes through every cell value, so it
## agrees with elevation_at on the grid nodes and bends between them instead of
## creasing. Reproduces planes exactly, including at the grid edges, where the field
## is extrapolated linearly. Each pass is clamped between its two bracketing nodes,
## so the surface never overshoots a step (no dip at the foot of a block, no bump
## above a ledge). Used to render steep tiles at a finer step.
func smooth_elevation_at(x_m: float, z_m: float) -> float:
	var column_f := clampf((x_m - _data.origin_x_m) / _data.cell_size_m, 0.0, float(_data.columns - 1))
	var row_f := clampf((z_m - _data.origin_z_m) / _data.cell_size_m, 0.0, float(_data.rows - 1))
	var column := mini(int(floorf(column_f)), _data.columns - 2)
	var row := mini(int(floorf(row_f)), _data.rows - 2)
	var weights_x := _catmull_rom_weights(column_f - float(column))
	var weights_z := _catmull_rom_weights(row_f - float(row))
	var along := PackedFloat32Array()
	along.resize(4)
	for offset_z in 4:
		var sample_row := row - 1 + offset_z
		var p1 := _extrapolated_elevation(column, sample_row)
		var p2 := _extrapolated_elevation(column + 1, sample_row)
		var along_x := (
			weights_x[0] * _extrapolated_elevation(column - 1, sample_row)
			+ weights_x[1] * p1
			+ weights_x[2] * p2
			+ weights_x[3] * _extrapolated_elevation(column + 2, sample_row)
		)
		along[offset_z] = clampf(along_x, minf(p1, p2), maxf(p1, p2))
	var result := weights_z[0] * along[0] + weights_z[1] * along[1] + weights_z[2] * along[2] + weights_z[3] * along[3]
	return clampf(result, minf(along[1], along[2]), maxf(along[1], along[2]))


## Positions and normals of one cell block on the smooth surface.
class SurfaceBlock extends RefCounted:
	var positions := PackedVector3Array()
	var normals := PackedVector3Array()


## The cell block [start_column..end_column] x [start_row..end_row] sampled on the smooth
## surface every 1 / subdivision cell, row-major: (cells * subdivision + 1) points per
## axis. The points on the grid nodes equal cell_world_position; the normals are
## central differences on the smooth surface, continuous across tile borders thanks to
## a one-cell margin. The smooth grid is built once for both arrays.
func cell_block_subdivided(
	start_column: int, end_column: int, start_row: int, end_row: int, subdivision: int
) -> SurfaceBlock:
	var block := SurfaceBlock.new()
	if subdivision <= 1:
		block.positions = cell_block_positions(start_column, end_column, start_row, end_row)
		block.normals = cell_block_normals(start_column, end_column, start_row, end_row)
		return block
	var first_column := _clamp_column(start_column)
	var last_column := _clamp_column(end_column)
	var first_row := _clamp_row(start_row)
	var last_row := _clamp_row(end_row)
	var fine := _fine_block(first_column, last_column, first_row, last_row, subdivision)
	var points_x := (last_column - first_column) * subdivision + 1
	var points_z := (last_row - first_row) * subdivision + 1
	var padded_x := points_x + 2 * subdivision
	var step_m := _data.cell_size_m / float(subdivision)
	var x0 := _data.origin_x_m + float(first_column) * _data.cell_size_m
	var z0 := _data.origin_z_m + float(first_row) * _data.cell_size_m
	var twice_step_m := 2.0 * step_m
	var positions := PackedVector3Array()
	var normals := PackedVector3Array()
	positions.resize(points_x * points_z)
	normals.resize(points_x * points_z)
	var index := 0
	for sub_row in points_z:
		var fine_row_start := (sub_row + subdivision) * padded_x + subdivision
		var z := z0 + float(sub_row) * step_m
		for sub_column in points_x:
			var at := fine_row_start + sub_column
			positions[index] = Vector3(x0 + float(sub_column) * step_m, fine[at], z)
			var slope_x := (fine[at + 1] - fine[at - 1]) / twice_step_m
			var slope_z := (fine[at + padded_x] - fine[at - padded_x]) / twice_step_m
			normals[index] = Vector3(-slope_x, 1.0, -slope_z).normalized()
			index += 1
	block.positions = positions
	block.normals = normals
	return block


func cell_block_positions_subdivided(
	start_column: int, end_column: int, start_row: int, end_row: int, subdivision: int
) -> PackedVector3Array:
	return cell_block_subdivided(start_column, end_column, start_row, end_row, subdivision).positions


func cell_block_normals_subdivided(
	start_column: int, end_column: int, start_row: int, end_row: int, subdivision: int
) -> PackedVector3Array:
	return cell_block_subdivided(start_column, end_column, start_row, end_row, subdivision).normals


## Smooth heights of the block padded by one cell on every side, row-major with
## (cells + 2) * subdivision + 1 points per axis. Separable Catmull-Rom with the
## weights of each sub-cell phase computed once: no per-point function calls.
func _fine_block(first_column: int, last_column: int, first_row: int, last_row: int, subdivision: int) -> PackedFloat32Array:
	var cells_x := last_column - first_column + 2
	var cells_z := last_row - first_row + 2
	var source_x := cells_x + 4
	var source_z := cells_z + 4
	var source_column0 := first_column - 2
	var source_row0 := first_row - 2
	var fine_x := cells_x * subdivision + 1
	var fine_z := cells_z * subdivision + 1
	var phase_weights: Array[PackedFloat32Array] = []
	for phase in subdivision:
		phase_weights.append(_catmull_rom_weights(float(phase) / float(subdivision)))
	# Source samples with linear extrapolation outside the grid.
	var source := PackedFloat32Array()
	source.resize(source_x * source_z)
	for sz in source_z:
		for sx in source_x:
			source[sz * source_x + sx] = _extrapolated_elevation(source_column0 + sx, source_row0 + sz)
	# Pass 1: along x, for every source row.
	var along := PackedFloat32Array()
	along.resize(fine_x * source_z)
	for sz in source_z:
		var row_start := sz * source_x
		var out_start := sz * fine_x
		for cell in cells_x:
			var p0 := source[row_start + cell]
			var p1 := source[row_start + cell + 1]
			var p2 := source[row_start + cell + 2]
			var p3 := source[row_start + cell + 3]
			var low := minf(p1, p2)
			var high := maxf(p1, p2)
			for phase in subdivision:
				var w := phase_weights[phase]
				along[out_start + cell * subdivision + phase] = clampf(w[0] * p0 + w[1] * p1 + w[2] * p2 + w[3] * p3, low, high)
		along[out_start + cells_x * subdivision] = source[row_start + cells_x + 1]
	# Pass 2: along z, for every fine column.
	var fine := PackedFloat32Array()
	fine.resize(fine_x * fine_z)
	for cell in cells_z:
		for phase in subdivision:
			var w := phase_weights[phase]
			var out_row := (cell * subdivision + phase) * fine_x
			var r0 := cell * fine_x
			var r1 := r0 + fine_x
			var r2 := r1 + fine_x
			var r3 := r2 + fine_x
			for fx in fine_x:
				var a1 := along[r1 + fx]
				var a2 := along[r2 + fx]
				fine[out_row + fx] = clampf(
					w[0] * along[r0 + fx] + w[1] * a1 + w[2] * a2 + w[3] * along[r3 + fx],
					minf(a1, a2),
					maxf(a1, a2)
				)
	var last_out := cells_z * subdivision * fine_x
	var last_source := (cells_z + 1) * fine_x
	for fx in fine_x:
		fine[last_out + fx] = along[last_source + fx]
	return fine


## Cell elevation, extended past the grid by linear extrapolation so the smooth surface
## keeps reproducing planes at the edges.
func _extrapolated_elevation(column: int, row: int) -> float:
	var columns := _data.columns
	var rows := _data.rows
	var c := clampi(column, 0, columns - 1)
	var r := clampi(row, 0, rows - 1)
	var value := _data.elevations_m[r * columns + c]
	if column < 0:
		value += float(column) * (_data.elevations_m[r * columns + 1] - _data.elevations_m[r * columns])
	elif column > columns - 1:
		value += float(column - columns + 1) * (_data.elevations_m[r * columns + columns - 1] - _data.elevations_m[r * columns + columns - 2])
	if row < 0:
		value += float(row) * (_data.elevations_m[columns + c] - _data.elevations_m[c])
	elif row > rows - 1:
		value += float(row - rows + 1) * (_data.elevations_m[(rows - 1) * columns + c] - _data.elevations_m[(rows - 2) * columns + c])
	return value


static func _catmull_rom_weights(t: float) -> PackedFloat32Array:
	var t2 := t * t
	var t3 := t2 * t
	return PackedFloat32Array([
		0.5 * (-t + 2.0 * t2 - t3),
		0.5 * (2.0 - 5.0 * t2 + 3.0 * t3),
		0.5 * (t + 4.0 * t2 - 3.0 * t3),
		0.5 * (-t2 + t3),
	])


func material_at(x_m: float, z_m: float) -> int:
	var column := int(roundf((x_m - _data.origin_x_m) / _data.cell_size_m))
	var row := int(roundf((z_m - _data.origin_z_m) / _data.cell_size_m))
	return material_at_cell(column, row)


func reflectivity_for(material: int) -> float:
	match material:
		SonarMaterial.Value.SAND:
			return _data.sand_reflectivity
		_:
			return _data.rock_reflectivity


## Central differences, one-sided on the grid edges. column and row must be in range.
func _normal_at(elevations: PackedFloat32Array, column: int, row: int) -> Vector3:
	var columns := _data.columns
	var left := maxi(column - 1, 0)
	var right := mini(column + 1, columns - 1)
	var north := maxi(row - 1, 0)
	var south := mini(row + 1, _data.rows - 1)
	var slope_x := (
		(elevations[row * columns + right] - elevations[row * columns + left])
		/ (float(right - left) * _data.cell_size_m)
	)
	var slope_z := (
		(elevations[south * columns + column] - elevations[north * columns + column])
		/ (float(south - north) * _data.cell_size_m)
	)
	return Vector3(-slope_x, 1.0, -slope_z).normalized()


func _clamp_column(column: int) -> int:
	return clampi(column, 0, _data.columns - 1)


func _clamp_row(row: int) -> int:
	return clampi(row, 0, _data.rows - 1)