utils.py

"""
Utility functions for the Pong AI project.
"""

import os
from pathlib import Path
from typing import Optional


def ensure_models_directory(models_dir: str = "models") -> str:
    """
    Ensure the models directory exists and return its path.
    
    Args:
        models_dir: Path to models directory
    
    Returns:
        Absolute path to models directory
    """
    models_path = Path(models_dir)
    models_path.mkdir(parents=True, exist_ok=True)
    return str(models_path.resolve())


def ensure_logs_directory(logs_dir: str = "pong_tensorboard") -> str:
    """
    Ensure the tensorboard logs directory exists and return its path.
    
    Args:
        logs_dir: Path to logs directory
    
    Returns:
        Absolute path to logs directory
    """
    logs_path = Path(logs_dir)
    logs_path.mkdir(parents=True, exist_ok=True)
    return str(logs_path.resolve())


def get_latest_model(models_dir: str = "models") -> Optional[str]:
    """
    Get the path to the latest trained model.
    
    Args:
        models_dir: Path to models directory
    
    Returns:
        Path to latest model or None if no models exist
    """
    models_path = Path(models_dir)
    if not models_path.exists():
        return None
    
    # Look for .zip files (SB3 model format)
    zip_files = list(models_path.glob("*.zip"))
    if not zip_files:
        return None
    
    # Return the most recently modified file
    return str(max(zip_files, key=os.path.getmtime))


def get_available_models(models_dir: str = "models") -> dict:
    """
    Get all available trained models with their info.
    
    Args:
        models_dir: Path to models directory
    
    Returns:
        Dictionary mapping model path to model info
    """
    models_path = Path(models_dir)
    if not models_path.exists():
        return {}
    
    models = {}
    for model_file in sorted(models_path.glob("*.zip")):
        try:
            # Extract step count from filename if present
            name = model_file.stem
            models[str(model_file)] = {
                'name': name,
                'path': str(model_file),
                'size_mb': model_file.stat().st_size / (1024 * 1024),
            }
        except Exception as e:
            print(f"[Utils] Error reading model {model_file}: {e}")
    
    return models


if __name__ == "__main__":
    print("[Utils] Pong AI Utilities")
    print(f"[Utils] Models directory: {ensure_models_directory()}")
    print(f"[Utils] Logs directory: {ensure_logs_directory()}")
    print(f"[Utils] Available models:")
    for path, info in get_available_models().items():
        print(f"  - {info['name']} ({info['size_mb']:.1f} MB)")