Trixy3/trainer/metadata.py

"""
Metadata management system for ML models.

This module provides comprehensive metadata management for trained models,
including version tracking, training statistics, model information, and
speaker/class management for different model types.
"""

import os
import json
import time
import hashlib
import socket
import platform
from datetime import datetime, timezone
from typing import Dict, Any, List, Optional, Union
from dataclasses import dataclass, field, asdict
from enum import Enum
import logging

import torch
import torch.nn as nn
import numpy as np


class ModelType(Enum):
    """Types of supported models."""
    WAKEWORD = "wakeword"
    VOICE_RECOGNITION = "voice_recognition"
    SPEAKER_EMBEDDING = "speaker_embedding"
    AUDIO_CLASSIFICATION = "audio_classification"
    CUSTOM = "custom"


class TrainingPhase(Enum):
    """Training phase indicators."""
    INITIAL = "initial"
    FINE_TUNING = "fine_tuning"
    TRANSFER_LEARNING = "transfer_learning"
    CONTINUED = "continued"


@dataclass
class TrainingStatistics:
    """Statistics from the training process."""
    total_epochs: int = 0
    final_train_loss: float = 0.0
    final_train_accuracy: float = 0.0
    best_val_loss: float = float('inf')
    best_val_accuracy: float = 0.0
    best_epoch: int = 0
    training_duration_seconds: float = 0.0
    total_parameters: int = 0
    trainable_parameters: int = 0
    convergence_epoch: Optional[int] = None
    early_stopping_triggered: bool = False
    final_learning_rate: float = 0.0
    max_memory_usage_mb: Optional[float] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'TrainingStatistics':
        """Create from dictionary."""
        return cls(**data)


@dataclass
class HardwareInfo:
    """Information about training hardware."""
    hostname: str = field(default_factory=socket.gethostname)
    platform_system: str = field(default_factory=platform.system)
    platform_release: str = field(default_factory=platform.release)
    platform_version: str = field(default_factory=platform.version)
    cpu_count: int = field(default_factory=os.cpu_count)
    python_version: str = field(default_factory=platform.python_version)
    pytorch_version: str = field(default_factory=lambda: torch.__version__)
    cuda_available: bool = field(default_factory=torch.cuda.is_available)
    cuda_version: Optional[str] = None
    gpu_count: int = 0
    gpu_names: List[str] = field(default_factory=list)
    total_memory_gb: Optional[float] = None
    
    def __post_init__(self):
        """Initialize CUDA and GPU information."""
        if self.cuda_available:
            self.cuda_version = torch.version.cuda
            self.gpu_count = torch.cuda.device_count()
            self.gpu_names = [torch.cuda.get_device_name(i) for i in range(self.gpu_count)]
        
        # Try to get memory information
        try:
            import psutil
            self.total_memory_gb = psutil.virtual_memory().total / (1024**3)
        except ImportError:
            pass
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'HardwareInfo':
        """Create from dictionary."""
        return cls(**data)


@dataclass
class AudioConfig:
    """Audio processing configuration."""
    sample_rate: int = 16000
    bit_depth: int = 16
    channels: int = 1
    audio_length_seconds: float = 1.5
    chunk_length_samples: int = 24000  # 1.5s at 16kHz
    n_mels: int = 40
    n_fft: int = 512
    hop_length: int = 160
    win_length: int = 400
    fmin: float = 80.0
    fmax: float = 8000.0
    feature_type: str = "mel_spectrogram"  # mel_spectrogram, mfcc, raw
    normalization: str = "db"  # db, power, none
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'AudioConfig':
        """Create from dictionary."""
        return cls(**data)


@dataclass
class DatasetInfo:
    """Information about the training dataset."""
    total_samples: int = 0
    train_samples: int = 0
    val_samples: int = 0
    test_samples: int = 0
    num_classes: int = 0
    class_names: List[str] = field(default_factory=list)
    class_distribution: Dict[str, int] = field(default_factory=dict)
    data_source_paths: List[str] = field(default_factory=list)
    augmentation_applied: bool = False
    augmentation_config: Dict[str, Any] = field(default_factory=dict)
    preprocessing_steps: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'DatasetInfo':
        """Create from dictionary."""
        return cls(**data)


@dataclass
class ModelArchitectureInfo:
    """Information about model architecture."""
    architecture_name: str = "unknown"
    architecture_type: str = "unknown"  # cnn, rnn, transformer, hybrid
    input_shape: List[int] = field(default_factory=list)
    output_shape: List[int] = field(default_factory=list)
    total_parameters: int = 0
    trainable_parameters: int = 0
    model_size_mb: float = 0.0
    layer_count: int = 0
    activation_functions: List[str] = field(default_factory=list)
    optimizer_config: Dict[str, Any] = field(default_factory=dict)
    loss_function: str = "unknown"
    custom_components: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ModelArchitectureInfo':
        """Create from dictionary."""
        return cls(**data)


class ModelMetadata:
    """
    Comprehensive metadata management for ML models.
    
    This class handles all metadata aspects including training information,
    model architecture, dataset details, and format-specific requirements.
    """
    
    def __init__(self, model_type: ModelType = ModelType.CUSTOM):
        """
        Initialize model metadata.
        
        Args:
            model_type: Type of the model (wakeword, voice_recognition, etc.)
        """
        self.model_type = model_type
        self.creation_timestamp = datetime.now(timezone.utc).isoformat()
        self.last_modified = self.creation_timestamp
        self.version = "1.0.0"
        self.format_version = "1.0"
        
        # Core information
        self.model_name = "untitled_model"
        self.description = ""
        self.author = "Unknown"
        self.organization = ""
        self.license = "Private"
        self.tags = []
        
        # Training information
        self.training_phase = TrainingPhase.INITIAL
        self.training_statistics = TrainingStatistics()
        self.hardware_info = HardwareInfo()
        self.audio_config = AudioConfig()
        self.dataset_info = DatasetInfo()
        self.architecture_info = ModelArchitectureInfo()
        
        # File information
        self.file_format = "unknown"
        self.file_hash = ""
        self.file_size_bytes = 0
        self.compression_used = False
        self.encryption_used = False
        
        # Model-specific data
        self.model_specific = self._initialize_model_specific()
        
        # Performance metrics
        self.performance_metrics = {}
        self.benchmark_results = {}
        
        # Deployment information
        self.deployment_targets = []
        self.optimization_level = "none"  # none, quantized, pruned, distilled
        self.inference_requirements = {}
        
        # Validation and testing
        self.validation_results = {}
        self.test_results = {}
        self.cross_validation_scores = []
        
        # Compatibility
        self.framework_version = torch.__version__
        self.python_version = platform.python_version()
        self.dependencies = {}
        
        # Custom fields (extensible)
        self.custom_fields = {}
    
    def _initialize_model_specific(self) -> Dict[str, Any]:
        """Initialize model-type specific metadata."""
        if self.model_type == ModelType.WAKEWORD:
            return {
                "wakewords": ["custom", "system_command"],
                "wakeword_classes": {
                    0: "custom",
                    1: "system_command",
                    2: "negative"
                },
                "detection_threshold": 0.5,
                "false_positive_rate": 0.0,
                "false_negative_rate": 0.0,
                "wake_word_durations": {},
                "background_noise_types": []
            }
        elif self.model_type == ModelType.VOICE_RECOGNITION:
            return {
                "num_speakers": 0,
                "speakers": {},  # speaker_id: speaker_name
                "embedding_dimension": 256,
                "verification_threshold": 0.5,
                "identification_threshold": 0.3,
                "equal_error_rate": 0.0,
                "speaker_enrollment_samples": {},
                "cross_speaker_accuracy": 0.0
            }
        elif self.model_type == ModelType.SPEAKER_EMBEDDING:
            return {
                "embedding_dimension": 256,
                "pooling_method": "statistical",
                "normalization": "l2",
                "cosine_similarity_threshold": 0.5,
                "embedding_model_type": "ecapa_tdnn"
            }
        elif self.model_type == ModelType.AUDIO_CLASSIFICATION:
            return {
                "num_classes": 0,
                "class_names": [],
                "class_weights": {},
                "multi_label": False,
                "hierarchy": {}
            }
        else:
            return {}
    
    def update_from_training_config(self, config: 'TrainerConfig'):
        """Update metadata from training configuration."""
        self.audio_config.sample_rate = config.sample_rate
        self.audio_config.audio_length_seconds = config.audio_length
        self.audio_config.n_mels = config.n_mels
        self.audio_config.n_fft = config.n_fft
        self.audio_config.hop_length = config.hop_length
        self.audio_config.win_length = config.win_length
        
        self.file_format = config.model_format.value
        self.encryption_used = config.use_password_protection
        
        # Store training hyperparameters
        self.custom_fields['training_config'] = {
            'batch_size': config.batch_size,
            'learning_rate': config.learning_rate,
            'num_epochs': config.num_epochs,
            'early_stopping_patience': config.early_stopping_patience,
            'optimizer': 'AdamW',
            'scheduler_type': config.scheduler_type,
            'weight_decay': config.weight_decay,
            'gradient_clip_norm': config.gradient_clip_norm,
            'use_mixed_precision': config.use_mixed_precision
        }
    
    def update_from_model(self, model: nn.Module):
        """Update metadata from model architecture."""
        # Count parameters
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        self.architecture_info.total_parameters = total_params
        self.architecture_info.trainable_parameters = trainable_params
        self.training_statistics.total_parameters = total_params
        self.training_statistics.trainable_parameters = trainable_params
        
        # Get model size
        param_size = sum(p.nelement() * p.element_size() for p in model.parameters())
        buffer_size = sum(b.nelement() * b.element_size() for b in model.buffers())
        self.architecture_info.model_size_mb = (param_size + buffer_size) / (1024 * 1024)
        
        # Architecture details
        self.architecture_info.architecture_name = model.__class__.__name__
        self.architecture_info.layer_count = len(list(model.modules()))
        
        # Extract activation functions
        activations = []
        for module in model.modules():
            if hasattr(module, '__class__'):
                module_name = module.__class__.__name__
                if any(act in module_name.lower() for act in ['relu', 'gelu', 'tanh', 'sigmoid', 'leaky', 'silu', 'swish']):
                    activations.append(module_name)
        self.architecture_info.activation_functions = list(set(activations))
    
    def update_from_dataset(self, train_loader, val_loader=None, test_loader=None):
        """Update metadata from dataset information."""
        if train_loader:
            self.dataset_info.train_samples = len(train_loader.dataset)
            
            # Try to get class information
            if hasattr(train_loader.dataset, 'classes'):
                self.dataset_info.class_names = train_loader.dataset.classes
                self.dataset_info.num_classes = len(train_loader.dataset.classes)
            elif hasattr(train_loader.dataset, 'targets'):
                targets = train_loader.dataset.targets
                if isinstance(targets, torch.Tensor):
                    targets = targets.numpy()
                unique_classes = np.unique(targets)
                self.dataset_info.num_classes = len(unique_classes)
                
                # Count class distribution
                unique, counts = np.unique(targets, return_counts=True)
                self.dataset_info.class_distribution = {str(cls): int(count) for cls, count in zip(unique, counts)}
        
        if val_loader:
            self.dataset_info.val_samples = len(val_loader.dataset)
        
        if test_loader:
            self.dataset_info.test_samples = len(test_loader.dataset)
        
        self.dataset_info.total_samples = (
            self.dataset_info.train_samples + 
            self.dataset_info.val_samples + 
            self.dataset_info.test_samples
        )
    
    def update_from_training_results(self, metrics: 'TrainingMetrics'):
        """Update metadata from training results."""
        summary = metrics.get_summary()
        
        self.training_statistics.total_epochs = summary['current_epoch']
        self.training_statistics.best_epoch = summary['best_epoch']
        self.training_statistics.best_val_loss = summary['best_val_loss']
        self.training_statistics.best_val_accuracy = summary['best_val_accuracy']
        self.training_statistics.training_duration_seconds = summary['total_training_time']
        
        if summary['final_train_loss'] is not None:
            self.training_statistics.final_train_loss = summary['final_train_loss']
        if summary['final_val_loss'] is not None:
            self.training_statistics.final_val_loss = summary['final_val_loss']
        
        # Store detailed metrics
        self.performance_metrics['training_loss_history'] = metrics.metrics['train_loss']
        self.performance_metrics['validation_loss_history'] = metrics.metrics['val_loss']
        self.performance_metrics['training_accuracy_history'] = metrics.metrics['train_accuracy']
        self.performance_metrics['validation_accuracy_history'] = metrics.metrics['val_accuracy']
        self.performance_metrics['learning_rate_history'] = metrics.metrics['learning_rates']
        self.performance_metrics['epoch_times'] = metrics.metrics['epoch_times']
    
    def update_wakeword_info(self, wakeword_classes: Dict[int, str], thresholds: Dict[str, float] = None):
        """Update wakeword-specific metadata."""
        if self.model_type != ModelType.WAKEWORD:
            return
        
        self.model_specific["wakeword_classes"] = wakeword_classes
        self.model_specific["wakewords"] = list(wakeword_classes.values())
        
        if thresholds:
            self.model_specific.update(thresholds)
    
    def update_speaker_info(self, speakers: Dict[str, str], embedding_dim: int = None):
        """Update voice recognition specific metadata."""
        if self.model_type != ModelType.VOICE_RECOGNITION:
            return
        
        # Convert speaker dict to proper format (speaker_id: speaker_name)
        if isinstance(list(speakers.keys())[0], str) and not list(speakers.keys())[0].isdigit():
            # speakers = {name: id} -> convert to {id: name}
            speaker_dict = {f"speaker{str(idx).zfill(3)}": name for idx, name in enumerate(speakers.keys())}
        else:
            speaker_dict = speakers
        
        self.model_specific["speakers"] = speaker_dict
        self.model_specific["num_speakers"] = len(speaker_dict)
        
        if embedding_dim:
            self.model_specific["embedding_dimension"] = embedding_dim
    
    def add_test_results(self, test_results: Dict[str, Any]):
        """Add test results to metadata."""
        self.test_results.update(test_results)
        self.last_modified = datetime.now(timezone.utc).isoformat()
    
    def add_benchmark_results(self, benchmark_name: str, results: Dict[str, Any]):
        """Add benchmark results."""
        self.benchmark_results[benchmark_name] = {
            'results': results,
            'timestamp': datetime.now(timezone.utc).isoformat()
        }
        self.last_modified = datetime.now(timezone.utc).isoformat()
    
    def add_deployment_target(self, target: str, requirements: Dict[str, Any] = None):
        """Add deployment target information."""
        if target not in self.deployment_targets:
            self.deployment_targets.append(target)
        
        if requirements:
            self.inference_requirements[target] = requirements
        
        self.last_modified = datetime.now(timezone.utc).isoformat()
    
    def set_file_info(self, filepath: str):
        """Update file-related information."""
        if os.path.exists(filepath):
            self.file_size_bytes = os.path.getsize(filepath)
            
            # Generate file hash
            hash_sha256 = hashlib.sha256()
            with open(filepath, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_sha256.update(chunk)
            self.file_hash = hash_sha256.hexdigest()
        
        self.last_modified = datetime.now(timezone.utc).isoformat()
    
    def increment_version(self, version_type: str = "patch"):
        """Increment version number."""
        major, minor, patch = map(int, self.version.split('.'))
        
        if version_type == "major":
            major += 1
            minor = 0
            patch = 0
        elif version_type == "minor":
            minor += 1
            patch = 0
        else:  # patch
            patch += 1
        
        self.version = f"{major}.{minor}.{patch}"
        self.last_modified = datetime.now(timezone.utc).isoformat()
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert metadata to dictionary format."""
        return {
            # Core information
            'model_type': self.model_type.value,
            'creation_timestamp': self.creation_timestamp,
            'last_modified': self.last_modified,
            'version': self.version,
            'format_version': self.format_version,
            'model_name': self.model_name,
            'description': self.description,
            'author': self.author,
            'organization': self.organization,
            'license': self.license,
            'tags': self.tags,
            
            # Training information
            'training_phase': self.training_phase.value,
            'training_statistics': self.training_statistics.to_dict(),
            'hardware_info': self.hardware_info.to_dict(),
            'audio_config': self.audio_config.to_dict(),
            'dataset_info': self.dataset_info.to_dict(),
            'architecture_info': self.architecture_info.to_dict(),
            
            # File information
            'file_format': self.file_format,
            'file_hash': self.file_hash,
            'file_size_bytes': self.file_size_bytes,
            'compression_used': self.compression_used,
            'encryption_used': self.encryption_used,
            
            # Model-specific data
            'model_specific': self.model_specific,
            
            # Performance and testing
            'performance_metrics': self.performance_metrics,
            'benchmark_results': self.benchmark_results,
            'validation_results': self.validation_results,
            'test_results': self.test_results,
            'cross_validation_scores': self.cross_validation_scores,
            
            # Deployment
            'deployment_targets': self.deployment_targets,
            'optimization_level': self.optimization_level,
            'inference_requirements': self.inference_requirements,
            
            # Compatibility
            'framework_version': self.framework_version,
            'python_version': self.python_version,
            'dependencies': self.dependencies,
            
            # Custom
            'custom_fields': self.custom_fields
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ModelMetadata':
        """Create metadata from dictionary."""
        metadata = cls(ModelType(data.get('model_type', 'custom')))
        
        # Core information
        metadata.creation_timestamp = data.get('creation_timestamp', metadata.creation_timestamp)
        metadata.last_modified = data.get('last_modified', metadata.last_modified)
        metadata.version = data.get('version', metadata.version)
        metadata.format_version = data.get('format_version', metadata.format_version)
        metadata.model_name = data.get('model_name', metadata.model_name)
        metadata.description = data.get('description', metadata.description)
        metadata.author = data.get('author', metadata.author)
        metadata.organization = data.get('organization', metadata.organization)
        metadata.license = data.get('license', metadata.license)
        metadata.tags = data.get('tags', metadata.tags)
        
        # Training information
        if 'training_phase' in data:
            metadata.training_phase = TrainingPhase(data['training_phase'])
        if 'training_statistics' in data:
            metadata.training_statistics = TrainingStatistics.from_dict(data['training_statistics'])
        if 'hardware_info' in data:
            metadata.hardware_info = HardwareInfo.from_dict(data['hardware_info'])
        if 'audio_config' in data:
            metadata.audio_config = AudioConfig.from_dict(data['audio_config'])
        if 'dataset_info' in data:
            metadata.dataset_info = DatasetInfo.from_dict(data['dataset_info'])
        if 'architecture_info' in data:
            metadata.architecture_info = ModelArchitectureInfo.from_dict(data['architecture_info'])
        
        # File information
        metadata.file_format = data.get('file_format', metadata.file_format)
        metadata.file_hash = data.get('file_hash', metadata.file_hash)
        metadata.file_size_bytes = data.get('file_size_bytes', metadata.file_size_bytes)
        metadata.compression_used = data.get('compression_used', metadata.compression_used)
        metadata.encryption_used = data.get('encryption_used', metadata.encryption_used)
        
        # Model-specific data
        metadata.model_specific = data.get('model_specific', metadata.model_specific)
        
        # Performance and testing
        metadata.performance_metrics = data.get('performance_metrics', metadata.performance_metrics)
        metadata.benchmark_results = data.get('benchmark_results', metadata.benchmark_results)
        metadata.validation_results = data.get('validation_results', metadata.validation_results)
        metadata.test_results = data.get('test_results', metadata.test_results)
        metadata.cross_validation_scores = data.get('cross_validation_scores', metadata.cross_validation_scores)
        
        # Deployment
        metadata.deployment_targets = data.get('deployment_targets', metadata.deployment_targets)
        metadata.optimization_level = data.get('optimization_level', metadata.optimization_level)
        metadata.inference_requirements = data.get('inference_requirements', metadata.inference_requirements)
        
        # Compatibility
        metadata.framework_version = data.get('framework_version', metadata.framework_version)
        metadata.python_version = data.get('python_version', metadata.python_version)
        metadata.dependencies = data.get('dependencies', metadata.dependencies)
        
        # Custom
        metadata.custom_fields = data.get('custom_fields', metadata.custom_fields)
        
        return metadata
    
    def save_to_file(self, filepath: str):
        """Save metadata to JSON file."""
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
    
    @classmethod
    def load_from_file(cls, filepath: str) -> 'ModelMetadata':
        """Load metadata from JSON file."""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return cls.from_dict(data)
    
    def get_summary(self) -> str:
        """Get a human-readable summary of the metadata."""
        summary = f"Model: {self.model_name}\n"
        summary += f"Type: {self.model_type.value}\n"
        summary += f"Version: {self.version}\n"
        summary += f"Author: {self.author}\n"
        summary += f"Created: {self.creation_timestamp}\n"
        summary += f"Description: {self.description}\n\n"
        
        summary += "Architecture:\n"
        summary += f"  Name: {self.architecture_info.architecture_name}\n"
        summary += f"  Parameters: {self.architecture_info.total_parameters:,}\n"
        summary += f"  Trainable: {self.architecture_info.trainable_parameters:,}\n"
        summary += f"  Size: {self.architecture_info.model_size_mb:.2f} MB\n\n"
        
        summary += "Training:\n"
        summary += f"  Epochs: {self.training_statistics.total_epochs}\n"
        summary += f"  Best Epoch: {self.training_statistics.best_epoch}\n"
        summary += f"  Best Val Accuracy: {self.training_statistics.best_val_accuracy:.4f}\n"
        summary += f"  Duration: {self.training_statistics.training_duration_seconds:.2f}s\n\n"
        
        summary += "Dataset:\n"
        summary += f"  Total Samples: {self.dataset_info.total_samples:,}\n"
        summary += f"  Classes: {self.dataset_info.num_classes}\n"
        summary += f"  Train/Val/Test: {self.dataset_info.train_samples}/{self.dataset_info.val_samples}/{self.dataset_info.test_samples}\n\n"
        
        if self.model_type == ModelType.VOICE_RECOGNITION:
            summary += "Speakers:\n"
            for speaker_id, speaker_name in self.model_specific.get('speakers', {}).items():
                summary += f"  {speaker_id}: {speaker_name}\n"
        elif self.model_type == ModelType.WAKEWORD:
            summary += "Wakewords:\n"
            for wakeword in self.model_specific.get('wakewords', []):
                summary += f"  - {wakeword}\n"
        
        return summary


class MetadataManager:
    """Manager class for handling model metadata operations."""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        self.logger = logger or logging.getLogger(__name__)
    
    def create_metadata(self, model_type: ModelType, model_name: str, 
                       author: str = "Unknown") -> ModelMetadata:
        """Create new metadata instance."""
        metadata = ModelMetadata(model_type)
        metadata.model_name = model_name
        metadata.author = author
        
        self.logger.info(f"Created metadata for {model_type.value} model: {model_name}")
        return metadata
    
    def merge_metadata(self, base_metadata: ModelMetadata, 
                      update_metadata: ModelMetadata) -> ModelMetadata:
        """Merge two metadata instances."""
        # Create a copy of base metadata
        merged = ModelMetadata.from_dict(base_metadata.to_dict())
        
        # Update with new information
        update_dict = update_metadata.to_dict()
        
        # Merge training statistics (keep best results)
        if update_metadata.training_statistics.best_val_accuracy > merged.training_statistics.best_val_accuracy:
            merged.training_statistics = update_metadata.training_statistics
        
        # Always update architecture and dataset info
        merged.architecture_info = update_metadata.architecture_info
        merged.dataset_info = update_metadata.dataset_info
        
        # Merge performance metrics
        merged.performance_metrics.update(update_metadata.performance_metrics)
        merged.benchmark_results.update(update_metadata.benchmark_results)
        
        # Update modification time
        merged.last_modified = datetime.now(timezone.utc).isoformat()
        merged.increment_version("patch")
        
        return merged
    
    def validate_metadata(self, metadata: ModelMetadata) -> List[str]:
        """Validate metadata and return list of issues."""
        issues = []
        
        if not metadata.model_name or metadata.model_name == "untitled_model":
            issues.append("Model name not set")
        
        if not metadata.description:
            issues.append("Model description is empty")
        
        if metadata.training_statistics.total_epochs == 0:
            issues.append("No training statistics available")
        
        if metadata.dataset_info.total_samples == 0:
            issues.append("No dataset information available")
        
        if metadata.architecture_info.total_parameters == 0:
            issues.append("Model architecture information missing")
        
        # Model-specific validations
        if metadata.model_type == ModelType.VOICE_RECOGNITION:
            if metadata.model_specific.get('num_speakers', 0) == 0:
                issues.append("No speakers defined for voice recognition model")
        elif metadata.model_type == ModelType.WAKEWORD:
            if not metadata.model_specific.get('wakewords'):
                issues.append("No wakewords defined for wakeword model")
        
        return issues
    
    def export_metadata(self, metadata: ModelMetadata, format: str = "json") -> str:
        """Export metadata to various formats."""
        if format.lower() == "json":
            return json.dumps(metadata.to_dict(), indent=2, ensure_ascii=False)
        elif format.lower() == "yaml":
            try:
                import yaml
                return yaml.dump(metadata.to_dict(), default_flow_style=False, allow_unicode=True)
            except ImportError:
                raise ImportError("PyYAML not available for YAML export")
        elif format.lower() == "summary":
            return metadata.get_summary()
        else:
            raise ValueError(f"Unsupported export format: {format}")
    
    def create_training_report(self, metadata: ModelMetadata) -> str:
        """Create a comprehensive training report."""
        report = f"# Training Report: {metadata.model_name}\n\n"
        
        # Basic Information
        report += "## Basic Information\n"
        report += f"- **Model Type**: {metadata.model_type.value}\n"
        report += f"- **Version**: {metadata.version}\n"
        report += f"- **Author**: {metadata.author}\n"
        report += f"- **Created**: {metadata.creation_timestamp}\n"
        report += f"- **Last Modified**: {metadata.last_modified}\n"
        report += f"- **Description**: {metadata.description}\n\n"
        
        # Architecture
        report += "## Model Architecture\n"
        arch = metadata.architecture_info
        report += f"- **Architecture**: {arch.architecture_name}\n"
        report += f"- **Total Parameters**: {arch.total_parameters:,}\n"
        report += f"- **Trainable Parameters**: {arch.trainable_parameters:,}\n"
        report += f"- **Model Size**: {arch.model_size_mb:.2f} MB\n"
        report += f"- **Layers**: {arch.layer_count}\n"
        if arch.activation_functions:
            report += f"- **Activations**: {', '.join(arch.activation_functions)}\n"
        report += "\n"
        
        # Training Results
        report += "## Training Results\n"
        stats = metadata.training_statistics
        report += f"- **Total Epochs**: {stats.total_epochs}\n"
        report += f"- **Best Epoch**: {stats.best_epoch}\n"
        report += f"- **Final Train Loss**: {stats.final_train_loss:.6f}\n"
        report += f"- **Final Train Accuracy**: {stats.final_train_accuracy:.4f}\n"
        report += f"- **Best Validation Loss**: {stats.best_val_loss:.6f}\n"
        report += f"- **Best Validation Accuracy**: {stats.best_val_accuracy:.4f}\n"
        report += f"- **Training Duration**: {stats.training_duration_seconds:.2f} seconds\n"
        if stats.early_stopping_triggered:
            report += "- **Early Stopping**: Triggered\n"
        report += "\n"
        
        # Dataset Information
        report += "## Dataset Information\n"
        dataset = metadata.dataset_info
        report += f"- **Total Samples**: {dataset.total_samples:,}\n"
        report += f"- **Training Samples**: {dataset.train_samples:,}\n"
        report += f"- **Validation Samples**: {dataset.val_samples:,}\n"
        report += f"- **Test Samples**: {dataset.test_samples:,}\n"
        report += f"- **Number of Classes**: {dataset.num_classes}\n"
        if dataset.class_distribution:
            report += "- **Class Distribution**:\n"
            for class_name, count in dataset.class_distribution.items():
                report += f"  - {class_name}: {count:,} samples\n"
        report += "\n"
        
        # Hardware Information
        report += "## Hardware Information\n"
        hw = metadata.hardware_info
        report += f"- **Hostname**: {hw.hostname}\n"
        report += f"- **Platform**: {hw.platform_system} {hw.platform_release}\n"
        report += f"- **CPU Cores**: {hw.cpu_count}\n"
        report += f"- **CUDA Available**: {hw.cuda_available}\n"
        if hw.cuda_available:
            report += f"- **CUDA Version**: {hw.cuda_version}\n"
            report += f"- **GPU Count**: {hw.gpu_count}\n"
            if hw.gpu_names:
                report += f"- **GPUs**: {', '.join(hw.gpu_names)}\n"
        if hw.total_memory_gb:
            report += f"- **Total Memory**: {hw.total_memory_gb:.1f} GB\n"
        report += "\n"
        
        # Model-specific Information
        if metadata.model_type == ModelType.VOICE_RECOGNITION:
            report += "## Speaker Information\n"
            speakers = metadata.model_specific.get('speakers', {})
            report += f"- **Number of Speakers**: {len(speakers)}\n"
            if speakers:
                report += "- **Speakers**:\n"
                for speaker_id, speaker_name in speakers.items():
                    report += f"  - {speaker_id}: {speaker_name}\n"
            embedding_dim = metadata.model_specific.get('embedding_dimension', 'Unknown')
            report += f"- **Embedding Dimension**: {embedding_dim}\n"
            report += "\n"
        
        elif metadata.model_type == ModelType.WAKEWORD:
            report += "## Wakeword Information\n"
            wakewords = metadata.model_specific.get('wakewords', [])
            report += f"- **Wakewords**: {', '.join(wakewords)}\n"
            classes = metadata.model_specific.get('wakeword_classes', {})
            if classes:
                report += "- **Classes**:\n"
                for class_id, class_name in classes.items():
                    report += f"  - {class_id}: {class_name}\n"
            threshold = metadata.model_specific.get('detection_threshold', 'Not set')
            report += f"- **Detection Threshold**: {threshold}\n"
            report += "\n"
        
        # Test Results
        if metadata.test_results:
            report += "## Test Results\n"
            for metric, value in metadata.test_results.items():
                if isinstance(value, float):
                    report += f"- **{metric}**: {value:.6f}\n"
                else:
                    report += f"- **{metric}**: {value}\n"
            report += "\n"
        
        return report