controller.py

from abc import ABC, abstractmethod
from enum import Enum
import numpy as np

class Action(Enum):
    LEFT = -1
    NONE = 0
    RIGHT = 1
    SHOOT = 2

class Controller(ABC):
    @abstractmethod
    def get_action(self, game_state: dict) -> Action:
        pass

class HumanController(Controller):
    def __init__(self):
        self.current_action = Action.NONE
        self.shoot_requested = False
    
    def set_action(self, action: Action):
        self.current_action = action
    
    def set_shoot(self, shoot: bool):
        self.shoot_requested = shoot
    
    def get_action(self, game_state: dict) -> Action:
        if self.shoot_requested:
            self.shoot_requested = False  # Reset after one use
            return Action.SHOOT
        return self.current_action

class MLController(Controller):
    def __init__(self, model=None):
        self.model = model
    
    def get_action(self, game_state: dict) -> Action:
        if self.model is None:
            return Action.NONE
        
        # Extract features from game state
        features = self._extract_features(game_state)
        
        # Get prediction from model (placeholder)
        action_idx = 0  # self.model.predict(features)
        
        return list(Action)[action_idx + 1]  # Convert to Action enum
    
    def _extract_features(self, game_state: dict) -> np.ndarray:
        player_x = game_state.get('player_x', 0)
        player_y = game_state.get('player_y', 0)
        velocity_x = game_state.get('velocity_x', 0)
        velocity_y = game_state.get('velocity_y', 0)
        platforms = game_state.get('nearby_platforms', [])
        
        features = [player_x, player_y, velocity_x, velocity_y]
        
        # Add nearby platform positions (up to 5 platforms)
        for i in range(5):
            if i < len(platforms):
                features.extend([platforms[i]['x'], platforms[i]['y']])
            else:
                features.extend([0, 0])
        
        return np.array(features, dtype=np.float32)