1 files changed, 72 insertions, 33 deletions

diff --git a/go/game/ai.go b/go/game/ai.go
index 7b20d187..07440899 100644
--- a/go/game/ai.go
+++ b/go/game/ai.go
@@ -122,11 +122,16 @@ func selectRandomTargets(rand *rand.Rand, targets *Targets) error {
 
 var ErrNoPath = errors.New("no path found")
 
-func (m *Map) generateMapGraphFor(u *Unit) *dijkstra.Graph {
+func generateGraph(tiles []*Tile) *dijkstra.Graph {
 	graph := dijkstra.NewGraph()
-	for _, t := range m.AllTiles() {
+	for _, t := range tiles {
 		graph.AddMappedVertex(t.Position.String())
 	}
+	return graph
+}
+
+func (m *Map) generateMapGraphFor(u *Unit) *dijkstra.Graph {
+	graph := generateGraph(m.AllTiles())
 
 	for _, t := range m.AllTiles() {
 		tId := t.Position.String()
@@ -146,35 +151,65 @@ func (m *Map) generateMapGraphFor(u *Unit) *dijkstra.Graph {
 	return graph
 }
 
-// func (m *Map) generateMovementRangeGraphFor(u *Unit) *dijkstra.Graph {
-// origin := TileOrContainingPermTile(u).Position
-// graph := dijkstra.NewGraph()
-// for _, t := range tilesInRangeFromOrigin(m, origin, u.Movement.Range, true) {
-// graph.AddMappedVertex(t.Position.String())
-// }
-
-// for _, t := range TilesInRange(m, u, u.Movement.Range) {
-// tId := t.Position.String()
-// for _, neighbour := range tilesInRangeFromOrigin(m, t.Position, 1, true) {
-// if !IsPositionInRange(neighbour.Position, origin, u.Movement.Range) {
-// continue
-// }
-
-// if u.IsAvailableTile(neighbour) {
-// cost := 2
-// if t.OnDiagonal(neighbour) {
-// cost = 3
-// }
-// log.Printf("add %v -> %v: %d", tId, neighbour.Position.String(), int64(cost))
-// err := graph.AddMappedArc(tId, neighbour.Position.String(), int64(cost))
-// if err != nil {
-// log.Panicf("Failed to add mapped arc: %s", err)
-// }
-// }
-// }
-// }
-// return graph
-// }
+func (m *Map) generateMovementRangeGraphFor(u *Unit) *dijkstra.Graph {
+	origin := TileOrContainingPermTile(u).Position
+	graph := generateGraph(tilesInRangeFromOrigin(m, origin, u.Movement.Range, true))
+
+	// Add connections for all tiles and their neighbours
+	for _, t := range tilesInRangeFromOrigin(m, origin, u.Movement.Range, true) {
+		tId := t.Position.String()
+		for _, neighbour := range TilesInRangeFromOrigin(m, t.Position, 1) {
+			// Skip neighbours not in movement range
+			if !IsPositionInRange(neighbour.Position, origin, u.Movement.Range) {
+				continue
+			}
+
+			if u.IsAvailableTile(neighbour) {
+				cost := 2
+				if t.OnDiagonal(neighbour) {
+					cost = 3
+				}
+				err := graph.AddMappedArc(tId, neighbour.Position.String(), int64(cost))
+				if err != nil {
+					log.Panicf("Failed to add mapped arc: %s", err)
+				}
+			}
+		}
+	}
+	return graph
+}
+
+func (m *Map) generateConnectedMovementRangeGraphFor(u *Unit, conn TileType) *dijkstra.Graph {
+	origin := TileOrContainingPermTile(u).Position
+	graph := generateGraph(tilesInRangeFromOrigin(m, origin, u.Movement.Range, true))
+
+	// Add connections for all tiles and their neighbours
+	for _, t := range tilesInRangeFromOrigin(m, origin, u.Movement.Range, true) {
+		tId := t.Position.String()
+		// Skip tiles with different types
+		if t.Type != conn {
+			continue
+		}
+		for _, neighbour := range TilesInRangeFromOrigin(m, t.Position, 1) {
+			// Skip neighbours not in movement range or with different type
+			if !IsPositionInRange(neighbour.Position, origin, u.Movement.Range) || neighbour.Type != conn {
+				continue
+			}
+
+			if u.IsAvailableTile(neighbour) {
+				cost := 2
+				if t.OnDiagonal(neighbour) {
+					cost = 3
+				}
+				err := graph.AddMappedArc(tId, neighbour.Position.String(), int64(cost))
+				if err != nil {
+					log.Panicf("Failed to add mapped arc: %s", err)
+				}
+			}
+		}
+	}
+	return graph
+}
 
 func findPathTo(graph *dijkstra.Graph, u *Unit, pos Position) (dijkstra.BestPath, error) {
 	srcId, err := graph.GetMapping(TileOrContainingPermTile(u).Position.String())
@@ -276,14 +311,17 @@ func moveAwayFromEnemies(ai *UnitAI) Action {
 		return nil
 	}
 
+	// A Movement graph is not sufficient because we have to check the distance to units potentially not in range.
 	graph := ai.s.Map().generateMapGraphFor(ai.u)
 	var farest *Tile
 	var farestEnemyDistance int64
 
-	for _, availableTile := range ai.u.MoveRangeTiles(ai.s.Map()) {
+	// Check all reachable tiles
+	for _, availableTile := range ai.u.MoveRangeTiles() {
 		var nearestEnemy int64 = -1
 
 		for _, enemyUnit := range ai.s.EnemyUnits(ai.u.Controller()) {
+			// Get distance to the enemy unit
 			p, err := findPathTo(graph, enemyUnit, availableTile.Position)
 			if err != nil {
 				continue
@@ -299,6 +337,7 @@ func moveAwayFromEnemies(ai *UnitAI) Action {
 		}
 	}
 
+	// No move possible
 	if farest == nil {
 		return nil
 	}
@@ -331,7 +370,7 @@ func actionFromPaths(ai *UnitAI, graph *dijkstra.Graph, paths []*dijkstra.BestPa
 		var x, y int
 		_, _ = fmt.Sscanf(mapping, "(%d, %d)", &x, &y)
 		nextStep := m.TileAt(Position{X: x, Y: y})
-		if slices.Contains(ai.u.MoveRangeTiles(m), nextStep) {
+		if slices.Contains(ai.u.MoveRangeTiles(), nextStep) {
 			target = nextStep
 		} else {
 			break