backend/autotune/training_loop.py

"""
Training Loop for Autotune Module

Implements periodic training loop for policy optimization.
Manages training episodes and state-action-reward history.
"""

from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass
from datetime import datetime, timedelta
import logging
import threading
import time

from .schemas import (
    SystemState,
    ActionType,
    SchedulerAction,
    Reward,
    TrainingEpisode,
    AdaptiveConfig,
    PolicyUpdate,
)
from .state_encoder import StateEncoder
from .action_space import ActionSpace
from .reward_engine import RewardEngine
from .policy_model import PolicyModel
from .optimizer import Optimizer, OptimizerConfig

logger = logging.getLogger(__name__)


@dataclass
class TrainingConfig:
    """
    Configuration for training loop.
    """
    # Training frequency
    training_interval_seconds: int = 300  # 5 minutes
    episode_length_steps: int = 100
    
    # Batch settings
    batch_size: int = 32
    buffer_size: int = 10000
    
    # Exploration
    epsilon_start: float = 0.2
    epsilon_decay: float = 0.995
    epsilon_min: float = 0.01
    
    # Learning
    learning_rate: float = 0.01
    discount_factor: float = 0.95
    
    # Safety
    enable_rollback: bool = True
    rollback_threshold: float = 0.1
    
    # Logging
    log_interval_seconds: int = 60
    save_interval_episodes: int = 100


class TrainingLoop:
    """
    Manages the RL training loop for policy optimization.
    
    Runs periodic training updates:
    1. Collect state-action-reward tuples
    2. Apply scheduling policy
    3. Observe reward
    4. Update policy weights
    5. Log performance
    """
    
    def __init__(
        self,
        state_encoder: StateEncoder,
        action_space: ActionSpace,
        reward_engine: RewardEngine,
        policy_model: PolicyModel,
        optimizer: Optimizer,
        config: Optional[TrainingConfig] = None,
    ):
        """
        Initialize training loop.
        
        Args:
            state_encoder: State encoder instance
            action_space: Action space instance
            reward_engine: Reward engine instance
            policy_model: Policy model instance
            optimizer: Optimizer instance
            config: Training configuration
        """
        self.state_encoder = state_encoder
        self.action_space = action_space
        self.reward_engine = reward_engine
        self.policy_model = policy_model
        self.optimizer = optimizer
        self.config = config or TrainingConfig()
        
        # Training state
        self.current_episode: Optional[TrainingEpisode] = None
        self.episode_count = 0
        self.step_count = 0
        
        # Historical data
        self.episodes: List[TrainingEpisode] = []
        self.policy_updates: List[PolicyUpdate] = []
        
        # State tracking
        self.last_state: Optional[SystemState] = None
        self.last_action: Optional[SchedulerAction] = None
        self.last_reward: Optional[float] = None
        
        # Control flags
        self._running = False
        self._paused = False
        self._stop_event = threading.Event()
        
        # Callbacks
        self.on_episode_complete: Optional[Callable] = None
        self.on_training_update: Optional[Callable] = None
    
    def start(self) -> None:
        """Start the training loop."""
        if self._running:
            logger.warning("Training loop already running")
            return
        
        self._running = True
        self._stop_event.clear()
        
        # Start training thread
        self._training_thread = threading.Thread(
            target=self._training_loop,
            daemon=True
        )
        self._training_thread.start()
        
        logger.info("Training loop started")
    
    def stop(self) -> None:
        """Stop the training loop."""
        if not self._running:
            return
        
        self._stop_event.set()
        self._running = False
        
        if hasattr(self, '_training_thread'):
            self._training_thread.join(timeout=5.0)
        
        logger.info("Training loop stopped")
    
    def pause(self) -> None:
        """Pause the training loop."""
        self._paused = True
        logger.info("Training loop paused")
    
    def resume(self) -> None:
        """Resume the training loop."""
        self._paused = False
        logger.info("Training loop resumed")
    
    def record_step(
        self,
        state: SystemState,
        action: SchedulerAction,
        reward: float,
        next_state: Optional[SystemState] = None,
    ) -> None:
        """
        Record a training step.
        
        Args:
            state: Current state
            action: Action taken
            reward: Reward received
            next_state: Next state (if available)
        """
        # Start new episode if needed
        if self.current_episode is None:
            self._start_new_episode()
        
        # Add step to current episode
        self.current_episode.add_step(state, action, reward)
        self.step_count += 1
        
        # Update optimizer
        self.optimizer.update(
            state=state,
            action=action.action_type,
            reward=reward,
            next_state=next_state,
        )
        
        # Check if episode is complete
        if self.current_episode.episode_length >= self.config.episode_length_steps:
            self._complete_episode()
        
        # Store last transition
        self.last_state = state
        self.last_action = action
        self.last_reward = reward
    
    def get_current_state(self) -> SystemState:
        """
        Get current system state for decision making.
        
        Returns:
            Current system state
        """
        return self.last_state or self.state_encoder.get_default_state()
    
    def _training_loop(self) -> None:
        """Main training loop (runs in separate thread)."""
        last_log_time = time.time()
        
        while not self._stop_event.is_set():
            try:
                if self._paused:
                    time.sleep(1.0)
                    continue
                
                current_time = time.time()
                
                # Log progress periodically
                if current_time - last_log_time >= self.config.log_interval_seconds:
                    self._log_progress()
                    last_log_time = current_time
                
                # Sleep briefly to avoid busy waiting
                time.sleep(1.0)
                
            except Exception as e:
                logger.error(f"Error in training loop: {e}", exc_info=True)
        
        logger.info("Training loop thread exited")
    
    def _start_new_episode(self) -> None:
        """Start a new training episode."""
        self.episode_count += 1
        self.current_episode = TrainingEpisode(
            episode_id=self.episode_count,
            start_time=datetime.utcnow(),
        )
        
        logger.debug(f"Started episode {self.episode_count}")
    
    def _complete_episode(self) -> None:
        """Complete the current episode."""
        if self.current_episode is None:
            return
        
        # Finalize episode
        self.current_episode.finalize()
        
        # Store episode
        self.episodes.append(self.current_episode)
        
        # Maintain max episodes
        if len(self.episodes) > 1000:
            self.episodes.pop(0)
        
        # Log episode summary
        logger.info(
            f"Episode {self.episode_count} complete: "
            f"reward={self.current_episode.total_reward:.4f}, "
            f"length={self.current_episode.episode_length}, "
            f"unique_actions={self.current_episode.unique_actions}"
        )
        
        # Trigger callback
        if self.on_episode_complete:
            self.on_episode_complete(self.current_episode)
        
        # Start new episode
        self._start_new_episode()
    
    def _log_progress(self) -> None:
        """Log training progress."""
        stats = self.get_statistics()
        
        logger.info(
            f"Training progress: "
            f"episodes={stats['episode_count']}, "
            f"steps={stats['step_count']}, "
            f"epsilon={stats['epsilon']:.4f}, "
            f"avg_reward={stats['avg_reward']:.4f}"
        )
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        Get training statistics.
        
        Returns:
            Dictionary of statistics
        """
        # Calculate recent average reward
        recent_episodes = self.episodes[-10:]
        if recent_episodes:
            avg_reward = sum(e.avg_reward for e in recent_episodes) / len(recent_episodes)
        else:
            avg_reward = 0.0
        
        return {
            "episode_count": self.episode_count,
            "step_count": self.step_count,
            "epsilon": self.policy_model.strategy.epsilon,
            "avg_reward": avg_reward,
            "buffer_size": len(self.optimizer.experience_buffer),
            "total_updates": self.optimizer.total_updates,
            "is_running": self._running,
            "is_paused": self._paused,
        }
    
    def get_episode_history(
        self,
        count: int = 100,
    ) -> List[TrainingEpisode]:
        """
        Get recent episode history.
        
        Args:
            count: Number of recent episodes to return
            
        Returns:
            List of recent episodes
        """
        return self.episodes[-count:]
    
    def decay_epsilon(self) -> None:
        """Decay exploration rate."""
        self.policy_model.strategy.decay_epsilon()
        
        logger.debug(
            f"Epsilon decayed to {self.policy_model.strategy.epsilon:.4f}"
        )
    
    def reset(self) -> None:
        """Reset training state."""
        self.stop()
        
        self.current_episode = None
        self.episode_count = 0
        self.step_count = 0
        self.episodes.clear()
        self.policy_updates.clear()
        self.last_state = None
        self.last_action = None
        self.last_reward = None
        
        self.optimizer.reset()
        
        logger.info("Training loop reset")


# Global training loop instance
_training_loop: Optional[TrainingLoop] = None


def get_training_loop(
    state_encoder: Optional[StateEncoder] = None,
    action_space: Optional[ActionSpace] = None,
    reward_engine: Optional[RewardEngine] = None,
    policy_model: Optional[PolicyModel] = None,
    optimizer: Optional[Optimizer] = None,
    config: Optional[TrainingConfig] = None,
) -> TrainingLoop:
    """
    Get or create global TrainingLoop instance.
    
    Args:
        state_encoder: State encoder instance
        action_space: Action space instance
        reward_engine: Reward engine instance
        policy_model: Policy model instance
        optimizer: Optimizer instance
        config: Training configuration
        
    Returns:
        TrainingLoop instance
    """
    global _training_loop
    if _training_loop is None:
        # Create components if not provided
        if state_encoder is None:
            state_encoder = StateEncoder()
        if action_space is None:
            action_space = ActionSpace()
        if reward_engine is None:
            reward_engine = RewardEngine()
        if policy_model is None:
            policy_model = PolicyModel()
        if optimizer is None:
            opt_config = OptimizerConfig() if config is None else OptimizerConfig(
                learning_rate=config.learning_rate,
                discount_factor=config.discount_factor,
            )
            optimizer = Optimizer(policy_model=policy_model, config=opt_config)
        
        _training_loop = TrainingLoop(
            state_encoder=state_encoder,
            action_space=action_space,
            reward_engine=reward_engine,
            policy_model=policy_model,
            optimizer=optimizer,
            config=config,
        )
    
    return _training_loop


__all__ = [
    "TrainingLoop",
    "TrainingConfig",
    "get_training_loop",
]