Atomizer/optimization_engine/utils/logger.py

"""
Atomizer Structured Logging System - Phase 1.3

Provides consistent, production-ready logging across all optimization studies.

Usage:
    from optimization_engine.utils.logger import get_logger

    logger = get_logger(__name__)
    logger.info("Starting optimization...")
    logger.error("Simulation failed", exc_info=True)

    # Study-specific logger with automatic file logging
    logger = get_logger("drone_gimbal_arm", study_dir="studies/drone_gimbal_arm/2_results")
    logger.trial_start(trial_number=5, design_vars={"thickness": 2.5})
    logger.trial_complete(trial_number=5, objectives={"mass": 120, "freq": 155})

Features:
    - Automatic file logging to study_dir/optimization.log
    - Console output with color-coded levels (if supported)
    - Structured trial logging for dashboard integration
    - Log rotation (50MB max, 3 backups)
    - No external dependencies (stdlib only)
"""

import logging
import sys
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, Any
from logging.handlers import RotatingFileHandler


# ANSI color codes for console output (Windows 10+ and Unix)
class LogColors:
    """ANSI color codes for console output."""
    RESET = '\033[0m'
    BOLD = '\033[1m'

    # Levels
    DEBUG = '\033[36m'      # Cyan
    INFO = '\033[32m'       # Green
    WARNING = '\033[33m'    # Yellow
    ERROR = '\033[31m'      # Red
    CRITICAL = '\033[35m'   # Magenta

    # Custom
    TRIAL = '\033[94m'      # Bright Blue
    SUCCESS = '\033[92m'    # Bright Green


class ColoredFormatter(logging.Formatter):
    """Formatter that adds color to console output."""

    COLORS = {
        logging.DEBUG: LogColors.DEBUG,
        logging.INFO: LogColors.INFO,
        logging.WARNING: LogColors.WARNING,
        logging.ERROR: LogColors.ERROR,
        logging.CRITICAL: LogColors.CRITICAL,
    }

    def __init__(self, fmt: str, use_colors: bool = True):
        super().__init__(fmt)
        self.use_colors = use_colors and self._supports_color()

    def _supports_color(self) -> bool:
        """Check if terminal supports ANSI colors."""
        # Windows 10+ supports ANSI
        if sys.platform == 'win32':
            try:
                import ctypes
                kernel32 = ctypes.windll.kernel32
                kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
                return True
            except:
                return False
        # Unix-like systems
        return hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()

    def format(self, record: logging.LogRecord) -> str:
        if self.use_colors:
            levelname = record.levelname
            color = self.COLORS.get(record.levelno, '')
            record.levelname = f"{color}{levelname}{LogColors.RESET}"

        return super().format(record)


class AtomizerLogger(logging.Logger):
    """Extended logger with trial-specific methods."""

    def trial_start(self, trial_number: int, design_vars: Dict[str, float]):
        """Log trial start with design variables."""
        self.info(f"{'='*60}")
        self.info(f"Trial #{trial_number} START")
        self.info(f"{'='*60}")
        self.info("Design Variables:")
        for name, value in design_vars.items():
            if isinstance(value, float):
                self.info(f"  {name}: {value:.4f}")
            else:
                self.info(f"  {name}: {value}")

    def trial_complete(self, trial_number: int, objectives: Dict[str, float],
                      constraints: Optional[Dict[str, float]] = None,
                      feasible: bool = True):
        """Log trial completion with results."""
        self.info(f"\nTrial #{trial_number} COMPLETE")
        self.info("Objectives:")
        for name, value in objectives.items():
            if isinstance(value, float):
                self.info(f"  {name}: {value:.4f}")
            else:
                self.info(f"  {name}: {value}")

        if constraints:
            self.info("Constraints:")
            for name, value in constraints.items():
                if isinstance(value, float):
                    self.info(f"  {name}: {value:.4f}")
                else:
                    self.info(f"  {name}: {value}")

        status = "[OK] Feasible" if feasible else "[WARNING] Infeasible"
        self.info(f"{status}")
        self.info(f"{'='*60}\n")

    def trial_failed(self, trial_number: int, error: str):
        """Log trial failure."""
        self.error(f"\nTrial #{trial_number} FAILED")
        self.error(f"Error: {error}")
        self.error(f"{'='*60}\n")

    def study_start(self, study_name: str, n_trials: int, sampler: str):
        """Log study initialization."""
        self.info("=" * 80)
        self.info(f"OPTIMIZATION STUDY: {study_name}")
        self.info("=" * 80)
        self.info(f"Trials: {n_trials}")
        self.info(f"Sampler: {sampler}")
        self.info(f"Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        self.info("=" * 80)
        self.info("")

    def study_complete(self, study_name: str, n_trials: int, n_successful: int):
        """Log study completion."""
        self.info("")
        self.info("=" * 80)
        self.info(f"STUDY COMPLETE: {study_name}")
        self.info("=" * 80)
        self.info(f"Total trials: {n_trials}")
        self.info(f"Successful: {n_successful}")
        self.info(f"Failed/Pruned: {n_trials - n_successful}")
        self.info(f"Completed: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        self.info("=" * 80)


# Register custom logger class
logging.setLoggerClass(AtomizerLogger)


def get_logger(
    name: str,
    level: int = logging.INFO,
    study_dir: Optional[Path] = None,
    console: bool = True,
    file_logging: bool = True
) -> AtomizerLogger:
    """
    Get or create a logger instance.

    Args:
        name: Logger name (typically __name__ or study name)
        level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
        study_dir: If provided, creates log file at study_dir/optimization.log
        console: Enable console output (default: True)
        file_logging: Enable file logging (default: True, requires study_dir)

    Returns:
        AtomizerLogger instance

    Example:
        # Simple logger
        logger = get_logger(__name__)
        logger.info("Starting optimization...")

        # Study logger with file output
        logger = get_logger(
            "drone_gimbal_arm",
            study_dir=Path("studies/drone_gimbal_arm/2_results")
        )
        logger.study_start("drone_gimbal_arm", n_trials=30, sampler="NSGAIISampler")
    """
    logger = logging.getLogger(name)

    # Only configure if not already configured
    if not logger.handlers:
        logger.setLevel(level)

        # Console handler with colors
        if console:
            console_handler = logging.StreamHandler(sys.stdout)
            console_handler.setLevel(level)
            console_formatter = ColoredFormatter(
                fmt='[%(levelname)s] %(message)s',
                use_colors=True
            )
            console_handler.setFormatter(console_formatter)
            logger.addHandler(console_handler)

        # File handler with rotation
        if file_logging and study_dir:
            study_dir = Path(study_dir)
            study_dir.mkdir(parents=True, exist_ok=True)
            log_file = study_dir / "optimization.log"

            file_handler = RotatingFileHandler(
                log_file,
                maxBytes=50 * 1024 * 1024,  # 50MB
                backupCount=3,
                encoding='utf-8'
            )
            file_handler.setLevel(level)
            file_formatter = logging.Formatter(
                fmt='%(asctime)s | %(levelname)-8s | %(name)s | %(message)s',
                datefmt='%Y-%m-%d %H:%M:%S'
            )
            file_handler.setFormatter(file_formatter)
            logger.addHandler(file_handler)

            # Log that file logging is enabled
            logger.debug(f"File logging enabled: {log_file}")

        # Prevent propagation to root logger
        logger.propagate = False

    return logger


def configure_root_logger(level: int = logging.WARNING):
    """
    Configure root logger to catch unconfigured loggers.

    Call this once at application startup to set up default logging behavior.
    """
    root_logger = logging.getLogger()
    root_logger.setLevel(level)

    if not root_logger.handlers:
        handler = logging.StreamHandler(sys.stdout)
        handler.setLevel(level)
        formatter = logging.Formatter('[%(levelname)s] %(name)s: %(message)s')
        handler.setFormatter(formatter)
        root_logger.addHandler(handler)


# Example usage and testing
if __name__ == "__main__":
    # Test basic logging
    print("Testing Atomizer Logging System")
    print("=" * 80)

    # Simple logger
    logger = get_logger("test_module")
    logger.debug("This is a debug message")
    logger.info("This is an info message")
    logger.warning("This is a warning message")
    logger.error("This is an error message")

    print()

    # Study logger with file output
    test_dir = Path("test_logs")
    test_dir.mkdir(exist_ok=True)

    study_logger = get_logger("test_study", study_dir=test_dir)
    study_logger.study_start("test_study", n_trials=5, sampler="TPESampler")

    # Simulate trial
    study_logger.trial_start(1, {"thickness": 2.5, "width": 10.0})
    study_logger.info("Running simulation...")
    study_logger.trial_complete(
        1,
        objectives={"mass": 120.5, "stiffness": 1500.2},
        constraints={"max_stress": 85.3},
        feasible=True
    )

    # Failed trial
    study_logger.trial_start(2, {"thickness": 1.0, "width": 5.0})
    study_logger.trial_failed(2, "Simulation convergence failure")

    study_logger.study_complete("test_study", n_trials=5, n_successful=4)

    print()
    print(f"Log file created at: {test_dir / 'optimization.log'}")
    print("Check the file to see structured logging output!")