swarm_engine_eval/environments/

maze.rs

//! MazeEnvironment - 迷路探索環境
//!
//! グリッドベースの迷路を探索するシミュレーション環境。
//!
//! # アクション
//!
//! - `Move`: 移動（target: north/south/east/west）
//! - `Look`: 周囲を観察（StepResult.data に情報を返す）
//! - `Wait`: 待機
//!
//! # 設計
//!
//! - Worker は `Look` アクションで周囲を観察
//! - 結果は `StepResult.data` に JSON で返る
//! - Worker/Manager は結果を見て次の行動を決定
//! - **迷路全体は見えない**（実用的なSwarm評価のため）
//!
//! # 使用例
//!
//! ```ignore
//! let map = MazeMap::parse("
//!     #####
//!     #S..#
//!     #.#.#
//!     #..G#
//!     #####
//! ");
//!
//! let env = MazeEnvironment::new(map, 2); // 2 workers
//!
//! // Look アクション
//! let action = Action::new("Look");
//! let result = env.step(worker_id, &action);
//! // result.data = { "position": {...}, "surroundings": {...}, "at_goal": false }
//!
//! // Move アクション
//! let action = Action::new("Move").with_arg("target", "east");
//! let result = env.step(worker_id, &action);
//! // result.success = true/false
//! ```

use std::collections::HashMap;
use std::sync::RwLock;

use serde::{Deserialize, Serialize};

use swarm_engine_core::agent::WorkResult;
use swarm_engine_core::environment::Environment;
use swarm_engine_core::types::{Action, WorkerId};

// ============================================================================
// Types
// ============================================================================

/// 位置
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Position {
    pub x: i32,
    pub y: i32,
}

impl Position {
    pub fn new(x: i32, y: i32) -> Self {
        Self { x, y }
    }

    /// 方向に移動した新しい位置を返す
    pub fn moved(&self, direction: &str) -> Self {
        match direction.to_lowercase().as_str() {
            "north" | "n" | "up" => Self::new(self.x, self.y - 1),
            "south" | "s" | "down" => Self::new(self.x, self.y + 1),
            "east" | "e" | "right" => Self::new(self.x + 1, self.y),
            "west" | "w" | "left" => Self::new(self.x - 1, self.y),
            _ => *self,
        }
    }
}

/// セルの種類
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Cell {
    /// 壁（通行不可）
    Wall,
    /// 床（通行可能）
    Floor,
    /// スタート地点
    Start,
    /// ゴール地点
    Goal,
}

impl Cell {
    /// 通行可能か
    pub fn is_passable(&self) -> bool {
        matches!(self, Cell::Floor | Cell::Start | Cell::Goal)
    }
}

/// 迷路マップ
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MazeMap {
    /// グリッド（[y][x]）
    pub grid: Vec<Vec<Cell>>,
    /// 幅
    pub width: usize,
    /// 高さ
    pub height: usize,
    /// スタート位置
    pub start: Position,
    /// ゴール位置
    pub goal: Position,
}

impl MazeMap {
    /// 文字列からマップを生成
    ///
    /// - `#`: 壁
    /// - `.`: 床
    /// - `S`: スタート
    /// - `G`: ゴール
    pub fn parse(s: &str) -> Self {
        let mut grid = Vec::new();
        let mut start = Position::new(0, 0);
        let mut goal = Position::new(0, 0);

        for (y, line) in s.lines().enumerate() {
            let trimmed = line.trim();
            if trimmed.is_empty() {
                continue;
            }

            let mut row = Vec::new();
            for (x, ch) in trimmed.chars().enumerate() {
                let cell = match ch {
                    '#' => Cell::Wall,
                    '.' => Cell::Floor,
                    'S' => {
                        start = Position::new(x as i32, y as i32);
                        Cell::Start
                    }
                    'G' => {
                        goal = Position::new(x as i32, y as i32);
                        Cell::Goal
                    }
                    _ => Cell::Floor,
                };
                row.push(cell);
            }
            grid.push(row);
        }

        let height = grid.len();
        let width = grid.first().map(|r| r.len()).unwrap_or(0);

        // スタート/ゴールの y 座標を調整（空行スキップ対応）
        let mut actual_y = 0;
        for line in s.lines() {
            let trimmed = line.trim();
            if trimmed.is_empty() {
                continue;
            }
            for (x, ch) in trimmed.chars().enumerate() {
                if ch == 'S' {
                    start = Position::new(x as i32, actual_y);
                } else if ch == 'G' {
                    goal = Position::new(x as i32, actual_y);
                }
            }
            actual_y += 1;
        }

        Self {
            grid,
            width,
            height,
            start,
            goal,
        }
    }

    /// 位置のセルを取得
    pub fn get(&self, pos: Position) -> Option<Cell> {
        if pos.x < 0 || pos.y < 0 {
            return None;
        }
        self.grid
            .get(pos.y as usize)
            .and_then(|row| row.get(pos.x as usize))
            .copied()
    }

    /// 位置が通行可能か
    pub fn is_passable(&self, pos: Position) -> bool {
        self.get(pos).map(|c| c.is_passable()).unwrap_or(false)
    }

    /// 位置がゴールか
    pub fn is_goal(&self, pos: Position) -> bool {
        self.get(pos) == Some(Cell::Goal)
    }
}

// ============================================================================
// MazeState (内部状態)
// ============================================================================

/// 迷路の内部状態（RwLock で保護）
#[derive(Debug)]
struct MazeState {
    /// エージェント位置 (WorkerId -> Position)
    agents: HashMap<WorkerId, Position>,
    /// ゴール到達したエージェント
    reached_goal: Vec<WorkerId>,
}

// ============================================================================
// MazeEnvironment
// ============================================================================

/// 迷路探索環境
pub struct MazeEnvironment {
    /// マップ
    map: MazeMap,
    /// 内部状態（RwLock で保護）
    state: RwLock<MazeState>,
    /// Worker 数
    worker_count: usize,
}

impl MazeEnvironment {
    /// 新しい迷路環境を作成
    pub fn new(map: MazeMap, worker_count: usize) -> Self {
        let mut agents = HashMap::new();

        // 全 Worker をスタート位置に配置
        for i in 0..worker_count {
            agents.insert(WorkerId(i), map.start);
        }

        Self {
            map,
            state: RwLock::new(MazeState {
                agents,
                reached_goal: Vec::new(),
            }),
            worker_count,
        }
    }

    /// 文字列からマップを読み込んで環境を作成
    pub fn from_str(map_str: &str, worker_count: usize) -> Self {
        let map = MazeMap::parse(map_str);
        Self::new(map, worker_count)
    }

    // ------------------------------------------------------------------------
    // Action Handlers
    // ------------------------------------------------------------------------

    fn handle_move(&self, worker_id: WorkerId, action: &Action) -> WorkResult {
        let direction = action
            .params
            .args
            .get("target")
            .map(|s| s.as_str())
            .unwrap_or("north");

        let mut state = self.state.write().unwrap();

        let current_pos = match state.agents.get(&worker_id) {
            Some(pos) => *pos,
            None => return WorkResult::env_failure("Worker not found in maze"),
        };

        let new_pos = current_pos.moved(direction);

        // 通行可能チェック
        if !self.map.is_passable(new_pos) {
            return WorkResult::env_failure(format!(
                "Cannot move {}: wall or out of bounds",
                direction
            ));
        }

        // 位置を更新
        state.agents.insert(worker_id, new_pos);

        // ゴール到達チェック
        if self.map.is_goal(new_pos) {
            if !state.reached_goal.contains(&worker_id) {
                state.reached_goal.push(worker_id);
            }

            // 全員ゴール到達？
            let all_reached = state.reached_goal.len() >= self.worker_count;
            if all_reached {
                return WorkResult::done_success(format!(
                    "Moved {} to goal! All workers reached goal!",
                    direction
                ));
            } else {
                return WorkResult::env_success(format!("Moved {} to goal!", direction));
            }
        }

        WorkResult::env_success(format!(
            "Moved {} to ({}, {})",
            direction, new_pos.x, new_pos.y
        ))
    }

    /// Look アクション: 周囲の情報を ActionResult.output に返す
    fn handle_look(&self, worker_id: WorkerId) -> WorkResult {
        let state = self.state.read().unwrap();

        let current_pos = match state.agents.get(&worker_id) {
            Some(pos) => *pos,
            None => return WorkResult::env_failure("Worker not found in maze"),
        };

        // 周囲4方向の情報を取得
        let mut surroundings = HashMap::new();
        for (dir, offset) in &[
            ("north", (0, -1)),
            ("south", (0, 1)),
            ("east", (1, 0)),
            ("west", (-1, 0)),
        ] {
            let check_pos = Position::new(current_pos.x + offset.0, current_pos.y + offset.1);
            let cell_info = match self.map.get(check_pos) {
                Some(Cell::Wall) => "wall",
                Some(Cell::Floor) => "floor",
                Some(Cell::Start) => "start",
                Some(Cell::Goal) => "goal",
                None => "void",
            };
            surroundings.insert(*dir, cell_info);
        }

        let data = serde_json::json!({
            "position": { "x": current_pos.x, "y": current_pos.y },
            "surroundings": surroundings,
            "at_goal": self.map.is_goal(current_pos),
        });

        WorkResult::env_success_structured(data)
    }

    fn handle_wait(&self, _worker_id: WorkerId) -> WorkResult {
        WorkResult::env_success("Waiting...")
    }
}

impl Environment for MazeEnvironment {
    fn step(&self, worker_id: WorkerId, action: &Action) -> WorkResult {
        match action.name.as_str() {
            "Move" => self.handle_move(worker_id, action),
            "Look" => self.handle_look(worker_id),
            "Wait" => self.handle_wait(worker_id),
            _ => WorkResult::unsupported(&action.name),
        }
    }

    fn reset(&self) {
        let mut state = self.state.write().unwrap();

        // 全 Worker をスタート位置に戻す
        state.agents.clear();
        for i in 0..self.worker_count {
            state.agents.insert(WorkerId(i), self.map.start);
        }
        state.reached_goal.clear();
    }

    fn name(&self) -> &str {
        "MazeEnvironment"
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    const SIMPLE_MAZE: &str = "
        #####
        #S..#
        #.#.#
        #..G#
        #####
    ";

    #[test]
    fn test_maze_map_from_str() {
        let map = MazeMap::parse(SIMPLE_MAZE);
        assert_eq!(map.width, 5);
        assert_eq!(map.height, 5);
        assert_eq!(map.start, Position::new(1, 1));
        assert_eq!(map.goal, Position::new(3, 3));
    }

    #[test]
    fn test_maze_passable() {
        let map = MazeMap::parse(SIMPLE_MAZE);
        assert!(map.is_passable(Position::new(1, 1))); // Start
        assert!(map.is_passable(Position::new(2, 1))); // Floor
        assert!(!map.is_passable(Position::new(0, 0))); // Wall
        assert!(!map.is_passable(Position::new(2, 2))); // Wall (center)
    }

    /// WorkResult から成功/失敗を取得するヘルパー
    fn is_success(result: &WorkResult) -> bool {
        match result {
            WorkResult::Acted { action_result, .. } => action_result.success,
            WorkResult::Done { success, .. } => *success,
            _ => false,
        }
    }

    #[test]
    fn test_maze_environment_move() {
        let env = MazeEnvironment::from_str(SIMPLE_MAZE, 1);
        let worker = WorkerId(0);

        // Move east
        let action = Action {
            name: "Move".to_string(),
            params: swarm_engine_core::types::ActionParams {
                target: None,
                args: [("target".to_string(), "east".to_string())]
                    .into_iter()
                    .collect(),
                data: vec![],
            },
        };

        let result = env.step(worker, &action);
        assert!(is_success(&result));

        // Check position
        let state = env.state.read().unwrap();
        assert_eq!(state.agents.get(&worker), Some(&Position::new(2, 1)));
    }

    #[test]
    fn test_maze_environment_wall_collision() {
        let env = MazeEnvironment::from_str(SIMPLE_MAZE, 1);
        let worker = WorkerId(0);

        // Try to move north (into wall)
        let action = Action {
            name: "Move".to_string(),
            params: swarm_engine_core::types::ActionParams {
                target: None,
                args: [("target".to_string(), "north".to_string())]
                    .into_iter()
                    .collect(),
                data: vec![],
            },
        };

        let result = env.step(worker, &action);
        assert!(!is_success(&result)); // Should fail
    }

    #[test]
    fn test_maze_environment_look() {
        let env = MazeEnvironment::from_str(SIMPLE_MAZE, 1);
        let worker = WorkerId(0);

        let action = Action {
            name: "Look".to_string(),
            params: Default::default(),
        };

        let result = env.step(worker, &action);
        assert!(is_success(&result));
        // output should contain surroundings (as String)
        if let WorkResult::Acted { action_result, .. } = result {
            assert!(action_result.output.is_some());
        } else {
            panic!("Expected WorkResult::Acted");
        }
    }
}