Trixy3/trainer/config.py

"""
Configuration management and argument parsing for ML trainers.

This module provides comprehensive configuration management including
file-based configuration, command-line argument parsing, and validation
with support for inheritance and environment variable override.
"""

import os
import sys
import json
import yaml
import argparse
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, Type
from dataclasses import dataclass, field, fields, asdict
from abc import ABC, abstractmethod
import torch

from .base import TrainerConfig, ModelFormat, TrainingState


class ConfigError(Exception):
    """Exception raised for configuration-related errors."""
    pass


class ConfigValidator(ABC):
    """Abstract base class for configuration validators."""
    
    @abstractmethod
    def validate(self, config: Dict[str, Any]) -> List[str]:
        """Validate configuration and return list of error messages."""
        pass


class TrainerConfigValidator(ConfigValidator):
    """Validator for trainer configuration."""
    
    def validate(self, config: Dict[str, Any]) -> List[str]:
        """Validate trainer configuration."""
        errors = []
        
        # Required fields
        required_fields = ['trainer_name', 'model_name', 'data_dir']
        for field in required_fields:
            if field not in config or not config[field]:
                errors.append(f"Required field '{field}' is missing or empty")
        
        # Numeric validations
        numeric_validations = {
            'batch_size': (1, 1024),
            'learning_rate': (1e-8, 1.0),
            'num_epochs': (1, 10000),
            'validation_split': (0.0, 1.0),
            'test_split': (0.0, 1.0),
            'sample_rate': (8000, 48000),
            'n_mels': (10, 128),
            'audio_length': (0.1, 30.0)
        }
        
        for field, (min_val, max_val) in numeric_validations.items():
            if field in config:
                value = config[field]
                if not isinstance(value, (int, float)) or not (min_val <= value <= max_val):
                    errors.append(f"Field '{field}' must be between {min_val} and {max_val}")
        
        # Validation split + test split check
        if 'validation_split' in config and 'test_split' in config:
            if config['validation_split'] + config['test_split'] >= 1.0:
                errors.append("validation_split + test_split must be less than 1.0")
        
        # Directory existence checks
        directory_fields = ['data_dir', 'output_dir', 'config_dir']
        for field in directory_fields:
            if field in config and config[field]:
                path = Path(config[field])
                if field == 'data_dir' and not path.exists():
                    errors.append(f"Data directory does not exist: {config[field]}")
        
        # Model format validation
        if 'model_format' in config:
            format_value = config['model_format']
            if isinstance(format_value, str):
                valid_formats = [fmt.value for fmt in ModelFormat]
                if format_value not in valid_formats:
                    errors.append(f"Invalid model format: {format_value}. Valid options: {valid_formats}")
        
        # Device validation
        if 'device' in config:
            device = config['device']
            if device not in ['auto', 'cpu', 'cuda', 'mps']:
                errors.append(f"Invalid device: {device}. Valid options: auto, cpu, cuda, mps")
        
        return errors


class WakewordConfigValidator(ConfigValidator):
    """Validator for wakeword trainer configuration."""
    
    def validate(self, config: Dict[str, Any]) -> List[str]:
        """Validate wakeword-specific configuration."""
        errors = []
        
        # Wakeword-specific validations
        if 'wakeword_classes' in config:
            classes = config['wakeword_classes']
            if not isinstance(classes, dict):
                errors.append("wakeword_classes must be a dictionary")
            elif len(classes) < 2:
                errors.append("At least 2 wakeword classes required (including negative)")
        
        # Audio processing for wakewords
        if 'audio_length' in config:
            if config['audio_length'] < 0.5 or config['audio_length'] > 5.0:
                errors.append("Wakeword audio_length should be between 0.5 and 5.0 seconds")
        
        return errors


class VoiceRecognitionConfigValidator(ConfigValidator):
    """Validator for voice recognition trainer configuration."""
    
    def validate(self, config: Dict[str, Any]) -> List[str]:
        """Validate voice recognition specific configuration."""
        errors = []
        
        # Speaker-specific validations
        if 'num_speakers' in config:
            num_speakers = config['num_speakers']
            if not isinstance(num_speakers, int) or num_speakers < 2:
                errors.append("num_speakers must be at least 2")
            elif num_speakers > 1000:
                errors.append("num_speakers should not exceed 1000 for practical purposes")
        
        if 'embedding_dimension' in config:
            embedding_dim = config['embedding_dimension']
            if not isinstance(embedding_dim, int) or embedding_dim < 64:
                errors.append("embedding_dimension must be at least 64")
            elif embedding_dim > 2048:
                errors.append("embedding_dimension should not exceed 2048")
        
        # Audio processing for voice recognition
        if 'audio_length' in config:
            if config['audio_length'] < 1.0 or config['audio_length'] > 10.0:
                errors.append("Voice recognition audio_length should be between 1.0 and 10.0 seconds")
        
        return errors


class ConfigManager:
    """
    Comprehensive configuration manager for ML trainers.
    
    Supports loading from files, command-line arguments, environment variables,
    and provides validation and inheritance capabilities.
    """
    
    def __init__(self, config_dir: str = "./config", logger: Optional[logging.Logger] = None):
        """
        Initialize configuration manager.
        
        Args:
            config_dir: Directory containing configuration files
            logger: Optional logger instance
        """
        self.config_dir = Path(config_dir)
        self.logger = logger or logging.getLogger(__name__)
        self.validators: Dict[str, ConfigValidator] = {
            'base': TrainerConfigValidator(),
            'wakeword': WakewordConfigValidator(),
            'voice_recognition': VoiceRecognitionConfigValidator()
        }
        self.loaded_configs: Dict[str, Dict[str, Any]] = {}
    
    def load_config_file(self, filepath: Union[str, Path]) -> Dict[str, Any]:
        """
        Load configuration from file (JSON or YAML).
        
        Args:
            filepath: Path to configuration file
            
        Returns:
            Configuration dictionary
        """
        filepath = Path(filepath)
        
        if not filepath.exists():
            raise ConfigError(f"Configuration file not found: {filepath}")
        
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                if filepath.suffix.lower() in ['.yml', '.yaml']:
                    try:
                        config = yaml.safe_load(f)
                    except ImportError:
                        raise ConfigError("PyYAML not available for YAML configuration files")
                elif filepath.suffix.lower() == '.json':
                    config = json.load(f)
                else:
                    raise ConfigError(f"Unsupported configuration file format: {filepath.suffix}")
            
            self.logger.info(f"Loaded configuration from: {filepath}")
            return config or {}
            
        except (json.JSONDecodeError, yaml.YAMLError) as e:
            raise ConfigError(f"Error parsing configuration file {filepath}: {str(e)}")
        except Exception as e:
            raise ConfigError(f"Error loading configuration file {filepath}: {str(e)}")
    
    def save_config_file(self, config: Dict[str, Any], filepath: Union[str, Path]):
        """
        Save configuration to file.
        
        Args:
            config: Configuration dictionary
            filepath: Path to save configuration
        """
        filepath = Path(filepath)
        filepath.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                if filepath.suffix.lower() in ['.yml', '.yaml']:
                    try:
                        yaml.safe_dump(config, f, default_flow_style=False, allow_unicode=True)
                    except ImportError:
                        raise ConfigError("PyYAML not available for YAML configuration files")
                elif filepath.suffix.lower() == '.json':
                    json.dump(config, f, indent=2, ensure_ascii=False)
                else:
                    raise ConfigError(f"Unsupported configuration file format: {filepath.suffix}")
            
            self.logger.info(f"Saved configuration to: {filepath}")
            
        except Exception as e:
            raise ConfigError(f"Error saving configuration file {filepath}: {str(e)}")
    
    def load_trainer_config(self, trainer_name: str, config_name: str = None) -> TrainerConfig:
        """
        Load trainer configuration with inheritance and validation.
        
        Args:
            trainer_name: Name of the trainer (wakeword, voice_recognition, etc.)
            config_name: Specific configuration name (defaults to trainer_name)
            
        Returns:
            TrainerConfig instance
        """
        if config_name is None:
            config_name = trainer_name
        
        # Load base configuration
        base_config = self._load_base_config()
        
        # Load trainer-specific configuration
        trainer_config = self._load_trainer_specific_config(trainer_name, config_name)
        
        # Merge configurations (trainer-specific overrides base)
        merged_config = self._merge_configs(base_config, trainer_config)
        
        # Apply environment variable overrides
        merged_config = self._apply_env_overrides(merged_config)
        
        # Validate configuration
        self._validate_config(merged_config, trainer_name)
        
        # Convert to TrainerConfig object
        return self._dict_to_trainer_config(merged_config)
    
    def _load_base_config(self) -> Dict[str, Any]:
        """Load base trainer configuration."""
        base_config_path = self.config_dir / "trainer_base.json"
        
        if base_config_path.exists():
            return self.load_config_file(base_config_path)
        else:
            # Return default base configuration
            self.logger.info("Using default base configuration")
            return {
                "batch_size": 32,
                "learning_rate": 0.001,
                "num_epochs": 100,
                "validation_split": 0.2,
                "test_split": 0.1,
                "sample_rate": 16000,
                "audio_length": 1.5,
                "n_mels": 40,
                "use_augmentation": True,
                "log_level": "INFO",
                "device": "auto"
            }
    
    def _load_trainer_specific_config(self, trainer_name: str, config_name: str) -> Dict[str, Any]:
        """Load trainer-specific configuration."""
        config_paths = [
            self.config_dir / f"{config_name}.json",
            self.config_dir / f"{config_name}.yml",
            self.config_dir / f"{config_name}.yaml",
            self.config_dir / f"{trainer_name}.json",
            self.config_dir / f"{trainer_name}.yml",
            self.config_dir / f"{trainer_name}.yaml"
        ]
        
        for config_path in config_paths:
            if config_path.exists():
                return self.load_config_file(config_path)
        
        self.logger.warning(f"No specific configuration found for {trainer_name}/{config_name}")
        return {}
    
    def _merge_configs(self, base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
        """Merge configuration dictionaries (override takes precedence)."""
        merged = base.copy()
        
        for key, value in override.items():
            if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
                # Recursively merge nested dictionaries
                merged[key] = self._merge_configs(merged[key], value)
            else:
                merged[key] = value
        
        return merged
    
    def _apply_env_overrides(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Apply environment variable overrides."""
        env_prefix = "TRIXY_TRAINER_"
        
        for key, value in os.environ.items():
            if key.startswith(env_prefix):
                config_key = key[len(env_prefix):].lower()
                
                # Convert environment variable value to appropriate type
                try:
                    # Try to parse as JSON first (for complex types)
                    parsed_value = json.loads(value)
                except json.JSONDecodeError:
                    # Fall back to string/numeric parsing
                    if value.lower() in ['true', 'false']:
                        parsed_value = value.lower() == 'true'
                    elif value.isdigit():
                        parsed_value = int(value)
                    else:
                        try:
                            parsed_value = float(value)
                        except ValueError:
                            parsed_value = value
                
                config[config_key] = parsed_value
                self.logger.info(f"Applied environment override: {config_key} = {parsed_value}")
        
        return config
    
    def _validate_config(self, config: Dict[str, Any], trainer_name: str):
        """Validate configuration using appropriate validators."""
        all_errors = []
        
        # Always validate with base validator
        base_errors = self.validators['base'].validate(config)
        all_errors.extend(base_errors)
        
        # Validate with trainer-specific validator if available
        if trainer_name in self.validators:
            trainer_errors = self.validators[trainer_name].validate(config)
            all_errors.extend(trainer_errors)
        
        if all_errors:
            error_msg = f"Configuration validation failed:\n" + "\n".join(f"- {error}" for error in all_errors)
            raise ConfigError(error_msg)
    
    def _dict_to_trainer_config(self, config_dict: Dict[str, Any]) -> TrainerConfig:
        """Convert configuration dictionary to TrainerConfig object."""
        # Handle enum conversions
        if 'model_format' in config_dict:
            if isinstance(config_dict['model_format'], str):
                config_dict['model_format'] = ModelFormat(config_dict['model_format'])
        
        # Filter only valid TrainerConfig fields
        trainer_config_fields = {f.name for f in fields(TrainerConfig)}
        filtered_config = {k: v for k, v in config_dict.items() if k in trainer_config_fields}
        
        # Store additional config in custom_params
        custom_params = {k: v for k, v in config_dict.items() if k not in trainer_config_fields}
        if custom_params:
            filtered_config['custom_params'] = custom_params
        
        try:
            return TrainerConfig(**filtered_config)
        except Exception as e:
            raise ConfigError(f"Error creating TrainerConfig: {str(e)}")
    
    def create_argument_parser(self, trainer_name: str = "trainer") -> argparse.ArgumentParser:
        """
        Create command-line argument parser for trainer.
        
        Args:
            trainer_name: Name of the trainer for help text
            
        Returns:
            Configured ArgumentParser
        """
        parser = argparse.ArgumentParser(
            description=f"Trixy {trainer_name.title()} Trainer",
            formatter_class=argparse.ArgumentDefaultsHelpFormatter
        )
        
        # General arguments
        parser.add_argument(
            '--config', '-c', type=str,
            help='Configuration file path'
        )
        parser.add_argument(
            '--trainer-name', type=str, default=trainer_name,
            help='Name of the trainer'
        )
        parser.add_argument(
            '--model-name', type=str,
            help='Name of the model'
        )
        parser.add_argument(
            '--output-dir', type=str, default='./models',
            help='Output directory for models'
        )
        parser.add_argument(
            '--data-dir', type=str, default='./trainer/data',
            help='Data directory'
        )
        
        # Training parameters
        training_group = parser.add_argument_group('Training Parameters')
        training_group.add_argument(
            '--batch-size', type=int, default=32,
            help='Batch size for training'
        )
        training_group.add_argument(
            '--learning-rate', '--lr', type=float, default=0.001,
            help='Learning rate'
        )
        training_group.add_argument(
            '--num-epochs', '--epochs', type=int, default=100,
            help='Number of training epochs'
        )
        training_group.add_argument(
            '--min-epochs', type=int, default=10,
            help='Minimum number of epochs'
        )
        training_group.add_argument(
            '--max-epochs', type=int, default=1000,
            help='Maximum number of epochs'
        )
        training_group.add_argument(
            '--early-stopping-patience', type=int, default=10,
            help='Early stopping patience'
        )
        
        # Model parameters
        model_group = parser.add_argument_group('Model Parameters')
        model_group.add_argument(
            '--model-format', choices=['pth', 'pt', 'onnx'], default='pth',
            help='Model file format'
        )
        model_group.add_argument(
            '--use-mixed-precision', action='store_true', default=True,
            help='Use mixed precision training'
        )
        model_group.add_argument(
            '--no-mixed-precision', dest='use_mixed_precision', action='store_false',
            help='Disable mixed precision training'
        )
        
        # Data parameters
        data_group = parser.add_argument_group('Data Parameters')
        data_group.add_argument(
            '--validation-split', type=float, default=0.2,
            help='Validation split ratio'
        )
        data_group.add_argument(
            '--test-split', type=float, default=0.1,
            help='Test split ratio'
        )
        data_group.add_argument(
            '--sample-rate', type=int, default=16000,
            help='Audio sample rate'
        )
        data_group.add_argument(
            '--audio-length', type=float, default=1.5,
            help='Audio length in seconds'
        )
        data_group.add_argument(
            '--n-mels', type=int, default=40,
            help='Number of mel filterbanks'
        )
        
        # Augmentation
        aug_group = parser.add_argument_group('Data Augmentation')
        aug_group.add_argument(
            '--use-augmentation', action='store_true', default=True,
            help='Use data augmentation'
        )
        aug_group.add_argument(
            '--no-augmentation', dest='use_augmentation', action='store_false',
            help='Disable data augmentation'
        )
        aug_group.add_argument(
            '--noise-factor', type=float, default=0.1,
            help='Noise augmentation factor'
        )
        aug_group.add_argument(
            '--speed-factor', type=float, default=0.1,
            help='Speed augmentation factor'
        )
        
        # Security
        security_group = parser.add_argument_group('Security')
        security_group.add_argument(
            '--password', type=str,
            help='Password for model encryption'
        )
        security_group.add_argument(
            '--no-password-protection', dest='use_password_protection', 
            action='store_false', default=True,
            help='Disable password protection'
        )
        
        # Hardware
        hardware_group = parser.add_argument_group('Hardware')
        hardware_group.add_argument(
            '--device', choices=['auto', 'cpu', 'cuda', 'mps'], default='auto',
            help='Training device'
        )
        hardware_group.add_argument(
            '--num-workers', type=int, default=4,
            help='Number of data loader workers'
        )
        
        # Logging and checkpoints
        logging_group = parser.add_argument_group('Logging and Checkpoints')
        logging_group.add_argument(
            '--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default='INFO',
            help='Logging level'
        )
        logging_group.add_argument(
            '--save-checkpoints', action='store_true', default=True,
            help='Save training checkpoints'
        )
        logging_group.add_argument(
            '--no-checkpoints', dest='save_checkpoints', action='store_false',
            help='Disable checkpoint saving'
        )
        logging_group.add_argument(
            '--checkpoint-interval', type=int, default=10,
            help='Checkpoint save interval (epochs)'
        )
        
        # Resume training
        resume_group = parser.add_argument_group('Resume Training')
        resume_group.add_argument(
            '--resume', '--resume-from-checkpoint', type=str,
            help='Resume training from checkpoint'
        )
        
        # Actions
        action_group = parser.add_argument_group('Actions')
        action_group.add_argument(
            '--train', action='store_true', default=True,
            help='Train the model'
        )
        action_group.add_argument(
            '--test-only', action='store_true',
            help='Only test the model (requires trained model)'
        )
        action_group.add_argument(
            '--prepare-data-only', action='store_true',
            help='Only prepare data without training'
        )
        action_group.add_argument(
            '--validate-config-only', action='store_true',
            help='Only validate configuration'
        )
        
        return parser
    
    def parse_args_to_config(self, args: argparse.Namespace, trainer_name: str) -> TrainerConfig:
        """
        Convert command-line arguments to TrainerConfig.
        
        Args:
            args: Parsed command-line arguments
            trainer_name: Name of the trainer
            
        Returns:
            TrainerConfig instance
        """
        # Start with configuration file if provided
        if hasattr(args, 'config') and args.config:
            config = self.load_config_file(args.config)
        else:
            config = self.load_trainer_config(trainer_name).to_dict()
        
        # Override with command-line arguments
        arg_dict = vars(args)
        for key, value in arg_dict.items():
            if value is not None:
                # Convert hyphens to underscores for config keys
                config_key = key.replace('-', '_')
                config[config_key] = value
        
        # Handle special cases
        if hasattr(args, 'resume') and args.resume:
            config['resume_from_checkpoint'] = True
            config['checkpoint_path'] = args.resume
        
        if hasattr(args, 'model_format') and args.model_format:
            config['model_format'] = f".{args.model_format}"
        
        # Validate and convert to TrainerConfig
        self._validate_config(config, trainer_name)
        return self._dict_to_trainer_config(config)
    
    def generate_config_template(self, trainer_name: str, filepath: str = None) -> str:
        """
        Generate configuration file template.
        
        Args:
            trainer_name: Name of the trainer
            filepath: Optional file path to save template
            
        Returns:
            Configuration template as string
        """
        template = {
            "# Configuration for Trixy ML Trainer": None,
            "trainer_name": trainer_name,
            "model_name": f"{trainer_name}_model",
            "description": f"Configuration for {trainer_name} trainer",
            
            "# Data settings": None,
            "data_dir": f"./trainer/data/{trainer_name}",
            "output_dir": "./models",
            "validation_split": 0.2,
            "test_split": 0.1,
            
            "# Training parameters": None,
            "batch_size": 32,
            "learning_rate": 0.001,
            "num_epochs": 100,
            "early_stopping_patience": 10,
            
            "# Model settings": None,
            "model_format": ".pth",
            "use_mixed_precision": True,
            "use_password_protection": True,
            
            "# Audio processing": None,
            "sample_rate": 16000,
            "audio_length": 1.5,
            "n_mels": 40,
            "n_fft": 512,
            "hop_length": 160,
            
            "# Data augmentation": None,
            "use_augmentation": True,
            "noise_factor": 0.1,
            "speed_factor": 0.1,
            "pitch_factor": 0.1,
            "volume_factor": 0.2,
            
            "# Hardware": None,
            "device": "auto",
            "num_workers": 4,
            
            "# Logging": None,
            "log_level": "INFO",
            "save_checkpoints": True,
            "checkpoint_interval": 10
        }
        
        # Add trainer-specific settings
        if trainer_name == "wakeword":
            template.update({
                "# Wakeword specific": None,
                "wakeword_classes": {
                    0: "custom",
                    1: "system_command", 
                    2: "negative"
                },
                "detection_threshold": 0.5
            })
        elif trainer_name == "voice_recognition":
            template.update({
                "# Voice recognition specific": None,
                "num_speakers": 2,
                "embedding_dimension": 256,
                "architecture": "ecapa_tdnn"
            })
        
        # Convert to JSON string
        template_str = "{\n"
        for key, value in template.items():
            if value is None:  # Comment line
                template_str += f"  \"{key}\",\n"
            elif isinstance(value, str):
                template_str += f"  \"{key}\": \"{value}\",\n"
            elif isinstance(value, dict):
                template_str += f"  \"{key}\": {json.dumps(value, indent=4)},\n"
            else:
                template_str += f"  \"{key}\": {json.dumps(value)},\n"
        template_str = template_str.rstrip(',\n') + "\n}"
        
        # Save to file if requested
        if filepath:
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(template_str)
            self.logger.info(f"Configuration template saved to: {filepath}")
        
        return template_str
    
    def list_available_configs(self) -> List[str]:
        """List all available configuration files."""
        configs = []
        
        if self.config_dir.exists():
            for pattern in ['*.json', '*.yml', '*.yaml']:
                configs.extend([f.stem for f in self.config_dir.glob(pattern)])
        
        return sorted(list(set(configs)))
    
    def get_config_info(self, config_name: str) -> Dict[str, Any]:
        """Get information about a configuration file."""
        config_paths = [
            self.config_dir / f"{config_name}.json",
            self.config_dir / f"{config_name}.yml", 
            self.config_dir / f"{config_name}.yaml"
        ]
        
        for config_path in config_paths:
            if config_path.exists():
                config = self.load_config_file(config_path)
                return {
                    'name': config_name,
                    'path': str(config_path),
                    'format': config_path.suffix,
                    'size': config_path.stat().st_size,
                    'trainer_name': config.get('trainer_name', 'unknown'),
                    'model_name': config.get('model_name', 'unknown'),
                    'description': config.get('description', ''),
                    'keys': list(config.keys())
                }
        
        raise ConfigError(f"Configuration not found: {config_name}")


def create_trainer_cli(trainer_name: str) -> argparse.ArgumentParser:
    """
    Create command-line interface for a trainer.
    
    Args:
        trainer_name: Name of the trainer
        
    Returns:
        Configured ArgumentParser
    """
    config_manager = ConfigManager()
    return config_manager.create_argument_parser(trainer_name)


def load_config_from_args(args: argparse.Namespace, trainer_name: str) -> TrainerConfig:
    """
    Load configuration from command-line arguments.
    
    Args:
        args: Parsed command-line arguments
        trainer_name: Name of the trainer
        
    Returns:
        TrainerConfig instance
    """
    config_manager = ConfigManager()
    return config_manager.parse_args_to_config(args, trainer_name)