controller.py

"""
MC-Planner Controller Module
Migrated to gymnasium environment without Java dependencies
"""

import torch
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path
import cv2

from src.gymnasium_env import MineDojoEnv
from src.minedojo_core import mc, SimulationState


class CraftAgent:
    """
    Craft agent for gymnasium-based Minecraft environment
    Simplified version without Java dependencies
    """
    
    def __init__(self, env: MineDojoEnv):
        self.env = env
        self.simulation_state = SimulationState()
        
    def no_op(self, times: int = 20) -> List[Dict]:
        """Perform no-operation actions"""
        actions = []
        for _ in range(times):
            action = {
                'movement': 0,
                'camera': np.array([0.0, 0.0]),
                'use': 0,
                'attack': 0,
                'jump': 0
            }
            actions.append(action)
        return actions
    
    def take_forward(self, times: int = 3) -> List[Dict]:
        """Move forward"""
        actions = []
        for _ in range(times):
            action = {
                'movement': 0,  # forward
                'camera': np.array([0.0, 0.0]),
                'use': 0,
                'attack': 0,
                'jump': 0
            }
            actions.append(action)
        return actions
    
    def index_slot(self, goal: str) -> int:
        """Get inventory slot index for item"""
        # Simplified slot indexing
        item_slots = {
            'wooden_pickaxe': 0,
            'stone_pickaxe': 1,
            'iron_pickaxe': 2,
            'wooden_planks': 3,
            'stick': 4,
            'cobblestone': 5,
        }
        return item_slots.get(goal, 0)
    
    def equip(self, goal: str) -> List[Dict]:
        """Equip item"""
        slot = self.index_slot(goal)
        # In real implementation, this would change selected hotbar slot
        return self.no_op(1)
    
    def pillar_jump(self, stepping_stone: str = "cobblestone") -> List[Dict]:
        """Perform pillar jumping"""
        actions = []
        # Place block, jump, repeat
        for _ in range(3):
            # Place block
            actions.extend(self.place(stepping_stone))
            # Jump
            actions.extend(self.jump())
        return actions
    
    def go_surface(self) -> List[Dict]:
        """Go to surface (simplified)"""
        # Look up and move forward
        actions = []
        for _ in range(10):
            action = {
                'movement': 0,  # forward
                'camera': np.array([0.0, 0.3]),  # look up
                'use': 0,
                'attack': 0,
                'jump': 1  # jump to go up
            }
            actions.append(action)
        return actions
    
    def acquire_info(self) -> Dict:
        """Get current environment information"""
        return {
            'inventory': self.simulation_state.get_inventory(),
            'position': self.simulation_state.position.copy(),
            'health': self.simulation_state.health,
            'hunger': self.simulation_state.hunger
        }
    
    def use(self) -> List[Dict]:
        """Use/interact action"""
        action = {
            'movement': 0,
            'camera': np.array([0.0, 0.0]),
            'use': 1,  # use item
            'attack': 0,
            'jump': 0
        }
        return [action]
    
    def look_to(self, deg: float = 0) -> List[Dict]:
        """Look in specific direction"""
        # Convert degrees to camera movement
        camera_x = np.sin(np.radians(deg)) * 0.5
        camera_y = np.cos(np.radians(deg)) * 0.5
        
        action = {
            'movement': 0,
            'camera': np.array([camera_x, camera_y]),
            'use': 0,
            'attack': 0,
            'jump': 0
        }
        return [action]
    
    def jump(self) -> List[Dict]:
        """Jump action"""
        action = {
            'movement': 0,
            'camera': np.array([0.0, 0.0]),
            'use': 0,
            'attack': 0,
            'jump': 1
        }
        return [action]
    
    def place(self, goal: str) -> List[Dict]:
        """Place block"""
        # Equip item first, then use
        actions = self.equip(goal)
        actions.extend(self.use())
        return actions
    
    def place_down(self, goal: str) -> List[Dict]:
        """Place block downward"""
        # Look down and place
        actions = self.look_to(-90)  # Look down
        actions.extend(self.place(goal))
        return actions
    
    def attack(self, times: int = 20) -> List[Dict]:
        """Attack/mine action"""
        actions = []
        for _ in range(times):
            action = {
                'movement': 0,
                'camera': np.array([0.0, 0.0]),
                'use': 0,
                'attack': 1,  # attack/mine
                'jump': 0
            }
            actions.append(action)
        return actions
    
    def recycle(self, goal: str, times: int = 20) -> List[Dict]:
        """Recycle/break item"""
        return self.attack(times)
    
    def craft_wo_table(self, goal: str) -> List[Dict]:
        """Craft without crafting table (2x2 grid)"""
        # Open inventory and craft
        actions = self.use()  # Open inventory
        
        # Simulate crafting
        if goal in mc.RECIPES:
            recipe = mc.RECIPES[goal]
            if self.simulation_state.can_craft(goal):
                self.simulation_state.craft_item(goal)
        
        return actions
    
    def forward(self, times: int = 5) -> List[Dict]:
        """Move forward"""
        return self.take_forward(times)
    
    def craft_w_table(self, goal: str) -> List[Dict]:
        """Craft with crafting table (3x3 grid)"""
        # Find/place crafting table, then craft
        actions = self.place("crafting_table")
        actions.extend(self.use())  # Use crafting table
        
        # Simulate crafting
        if goal in mc.RECIPES:
            recipe = mc.RECIPES[goal]
            if self.simulation_state.can_craft(goal):
                self.simulation_state.craft_item(goal)
        
        return actions
    
    def smelt_w_furnace(self, goal: str) -> List[Dict]:
        """Smelt with furnace"""
        # Place/use furnace
        actions = self.place("furnace")
        actions.extend(self.use())  # Use furnace
        
        # Simulate smelting
        if goal in mc.SMELTING:
            recipe = mc.SMELTING[goal]
            for ingredient, amount in recipe.items():
                if self.simulation_state.inventory.get(ingredient, 0) >= amount:
                    self.simulation_state.update_inventory(ingredient, -amount)
                    self.simulation_state.update_inventory(goal, 1)
        
        return actions
    
    def smelt_wo_furnace(self, goal: str) -> List[Dict]:
        """Smelt without furnace (not possible, fallback to with furnace)"""
        return self.smelt_w_furnace(goal)
    
    def get_action(self, preconditions: Dict, goal_type: str, goal: str) -> List[Dict]:
        """Get action sequence for goal"""
        if goal_type == "mine":
            # Mine goal
            tool_needed = preconditions.get('tool', 'wooden_pickaxe')
            actions = self.equip(tool_needed)
            actions.extend(self.attack(20))
            
            # Update simulation state
            if "cobblestone" in goal:
                self.simulation_state.update_inventory("cobblestone", 3)
            elif "iron_ore" in goal:
                self.simulation_state.update_inventory("iron_ore", 2)
            elif "diamond" in goal:
                self.simulation_state.update_inventory("diamond", 1)
            
        elif goal_type == "craft":
            # Craft goal
            if goal in ["wooden_slab", "wooden_planks", "stick"]:
                actions = self.craft_wo_table(goal)
            else:
                actions = self.craft_w_table(goal)
                
        elif goal_type == "smelt":
            actions = self.smelt_w_furnace(goal)
            
        else:
            # Default action
            actions = self.no_op(5)
        
        return actions


class MineAgent:
    """
    Mining agent for specific mining tasks
    """
    
    def __init__(self, env: MineDojoEnv, device: str = "cuda"):
        self.env = env
        self.device = device
        self.craft_agent = CraftAgent(env)
        
    def mine_item(self, item: str, quantity: int = 1) -> bool:
        """Mine specific item"""
        # Determine required tool
        tool_map = {
            'cobblestone': 'wooden_pickaxe',
            'iron_ore': 'stone_pickaxe', 
            'diamond': 'iron_pickaxe',
            'coal': 'wooden_pickaxe'
        }
        
        required_tool = tool_map.get(item, 'wooden_pickaxe')
        
        # Check if we have the tool
        inventory = self.craft_agent.simulation_state.get_inventory()
        if inventory.get(required_tool, 0) == 0:
            print(f"Missing required tool: {required_tool}")
            return False
        
        # Execute mining actions
        actions = self.craft_agent.get_action(
            preconditions={'tool': required_tool},
            goal_type='mine',
            goal=item
        )
        
        # Execute actions in environment
        for action in actions:
            obs, reward, terminated, truncated, info = self.env.step(action)
            if terminated or truncated:
                break
        
        return True


class MineAgentWrapper:
    """
    Wrapper for mining agent with goal-specific behavior
    """
    
    script_goals = ['cobblestone', 'stone', 'coal', 'iron_ore', 'diamond']
    
    def __init__(self, env: MineDojoEnv, mine_agent: MineAgent, max_ranking: int = 15):
        self.env = env
        self.mine_agent = mine_agent
        self.max_ranking = max_ranking
        
    def can_handle_goal(self, goal: str) -> bool:
        """Check if this wrapper can handle the goal"""
        goal_item = goal.replace("mine_", "").replace("obtain_", "")
        return goal_item in self.script_goals
    
    def execute_goal(self, goal: str) -> bool:
        """Execute the goal"""
        if not self.can_handle_goal(goal):
            return False
        
        goal_item = goal.replace("mine_", "").replace("obtain_", "")
        return self.mine_agent.mine_item(goal_item)


def resize_image(img: np.ndarray, target_resolution: Tuple[int, int] = (128, 128)) -> np.ndarray:
    """Resize image to target resolution"""
    return cv2.resize(img, dsize=target_resolution, interpolation=cv2.INTER_LINEAR)


def accquire_goal_embeddings(clip_path: str, goal_list: List[str], device: str = "cuda") -> Dict:
    """
    Acquire goal embeddings (simplified version without actual CLIP model)
    """
    # Placeholder implementation - in real version, this would use actual CLIP model
    embeddings = {}
    
    for goal in goal_list:
        # Generate dummy embedding
        embedding = torch.randn(512).to(device)
        embeddings[goal] = embedding
    
    return embeddings