"""
Training callbacks for KWS models
"""

from tensorflow import keras
from typing import Optional, Dict, Any


class MLFlowCallback(keras.callbacks.Callback):
    """Keras callback for logging metrics to MLFlow"""
    
    def __init__(self, mlflow_api, run_id: str):
        """
        Initialize MLFlow callback.
        
        Args:
            mlflow_api: MLFlow API client instance
            run_id: MLFlow run ID
        """
        super().__init__()
        self.mlflow_api = mlflow_api
        self.run_id = run_id
        self.epoch = 0
    
    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None):
        """
        Log metrics at the end of each epoch.
        
        Args:
            epoch: Current epoch number
            logs: Dictionary of metrics
        """
        if logs is None:
            logs = {}
        self.epoch = epoch
        
        # Log metrics to MLFlow
        try:
            metrics = {
                "train_loss": float(logs.get('loss', 0)),
                "train_accuracy": float(logs.get('accuracy', 0)),
            }
            if 'val_loss' in logs:
                metrics["val_loss"] = float(logs['val_loss'])
            if 'val_accuracy' in logs:
                metrics["val_accuracy"] = float(logs['val_accuracy'])
            
            self.mlflow_api.log_metrics(
                run_id=self.run_id,
                metrics=metrics,
                step=epoch + 1
            )
        except Exception as e:
            print(f"Warning: Failed to log metrics to MLFlow: {e}")
