Trixy3/trixy_core/ml/wakeword/model_loader.py

"""
Model loading and management system with password protection.

This module provides secure loading of PyTorch models from password-protected
ZIP files (.pth) with metadata validation and model verification.
"""

import json
import zipfile
import tempfile
import hashlib
import io
import warnings
from pathlib import Path
from typing import Dict, Any, Optional, Union, List, Tuple
from dataclasses import dataclass, asdict
from datetime import datetime

import torch
import torch.nn as nn


@dataclass
class ModelMetadata:
    """Metadata structure for wakeword models."""
    
    # Model identification
    model_name: str
    model_version: str
    model_type: str  # "wakeword", "voice_recognition"
    architecture: str  # "RepCNN", "ImprovedRepCNN", etc.
    
    # Training information
    training_date: str
    author: str
    description: str
    computer_name: str
    
    # Model configuration
    num_classes: int
    input_shape: List[int]  # [channels, height, width]
    class_labels: List[str]  # ["custom", "system_command", "negative"]
    
    # Training results
    final_accuracy: float
    final_loss: float
    num_epochs: int
    training_duration_hours: float
    
    # Audio configuration
    sample_rate: int
    n_mels: int
    time_frames: int
    chunk_duration_ms: float
    
    # Model performance
    inference_time_ms: float
    model_size_mb: float
    flops_estimate: int
    
    # Security and validation
    model_hash: str
    created_timestamp: str
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ModelMetadata':
        """Create ModelMetadata from dictionary."""
        return cls(**data)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert ModelMetadata to dictionary."""
        return asdict(self)
    
    def validate(self) -> bool:
        """Validate metadata completeness and consistency."""
        try:
            # Check required fields
            required_fields = [
                'model_name', 'model_version', 'model_type', 'architecture',
                'num_classes', 'input_shape', 'class_labels'
            ]
            
            for field in required_fields:
                value = getattr(self, field)
                if value is None or (isinstance(value, (str, list)) and len(value) == 0):
                    return False
            
            # Validate consistency
            if len(self.class_labels) != self.num_classes:
                return False
                
            if len(self.input_shape) != 3:  # Should be [C, H, W]
                return False
                
            # Check for wakeword-specific requirements
            if self.model_type == "wakeword":
                expected_labels = ["custom", "system_command", "negative"]
                if not all(label in self.class_labels for label in expected_labels):
                    return False
            
            return True
            
        except Exception:
            return False


class ModelLoadError(Exception):
    """Exception raised when model loading fails."""
    pass


class SecureModelLoader:
    """
    Secure model loader for password-protected PyTorch models.
    
    Handles loading of .pth files that are actually ZIP archives containing:
    - model.pth: The actual PyTorch model state dict
    - metadata.json: Model metadata and configuration
    - Optional additional files (training logs, etc.)
    """
    
    def __init__(self, device: Optional[torch.device] = None):
        """
        Initialize the secure model loader.
        
        Args:
            device: PyTorch device for loading models
        """
        self.device = device or torch.device("cpu")
        self._loaded_models: Dict[str, Dict[str, Any]] = {}
        
    def load_model(self, model_path: Union[str, Path], password: str,
                  model_class: Optional[type] = None,
                  verify_hash: bool = True) -> Tuple[nn.Module, ModelMetadata]:
        """
        Load a password-protected PyTorch model.
        
        Args:
            model_path: Path to the .pth file (ZIP archive)
            password: Password for decrypting the model
            model_class: Model class to instantiate (if None, attempts auto-detection)
            verify_hash: Whether to verify model integrity using hash
            
        Returns:
            Tuple of (model, metadata)
            
        Raises:
            ModelLoadError: If loading fails
        """
        model_path = Path(model_path)
        
        if not model_path.exists():
            raise ModelLoadError(f"Model file not found: {model_path}")
        
        try:
            # Load from ZIP archive
            with zipfile.ZipFile(model_path, 'r') as zip_file:
                # Try to extract with password
                try:
                    zip_file.setpassword(password.encode('utf-8'))
                    
                    # Load metadata
                    metadata = self._load_metadata(zip_file)
                    
                    # Load model state dict
                    model_state = self._load_model_state(zip_file)
                    
                    # Verify hash if requested
                    if verify_hash and metadata.model_hash:
                        if not self._verify_model_hash(model_state, metadata.model_hash):
                            raise ModelLoadError("Model hash verification failed")
                    
                    # Create model instance
                    model = self._create_model_instance(metadata, model_class)

                    # Reparameterize model if it has the method (RepCNN models)
                    # This converts the model structure to match the saved state dict
                    if hasattr(model, 'reparameterize'):
                        model.reparameterize()

                    # Load state dict
                    model.load_state_dict(model_state, strict=True)
                    model.to(self.device)
                    model.eval()
                    
                    # Cache the loaded model
                    cache_key = str(model_path)
                    self._loaded_models[cache_key] = {
                        'model': model,
                        'metadata': metadata,
                        'load_time': datetime.now()
                    }
                    
                    return model, metadata
                    
                except RuntimeError as e:
                    if "Bad password" in str(e) or "Bad CRC" in str(e):
                        raise ModelLoadError("Invalid password for model file")
                    else:
                        raise ModelLoadError(f"Failed to extract model: {e}")
                        
        except zipfile.BadZipFile:
            raise ModelLoadError("Invalid ZIP file format")
        except Exception as e:
            raise ModelLoadError(f"Failed to load model: {e}")
    
    def _load_metadata(self, zip_file: zipfile.ZipFile) -> ModelMetadata:
        """Load and validate metadata from ZIP file."""
        try:
            # Try metafile.json first (used by trainer), then metadata.json (legacy)
            try:
                with zip_file.open('metafile.json') as meta_file:
                    raw_metadata = json.load(meta_file)
                    # Convert trainer format to ModelMetadata format
                    metadata_dict = self._convert_trainer_metadata(raw_metadata)
            except KeyError:
                with zip_file.open('metadata.json') as meta_file:
                    metadata_dict = json.load(meta_file)

            metadata = ModelMetadata.from_dict(metadata_dict)

            if not metadata.validate():
                raise ModelLoadError("Invalid metadata structure")

            return metadata

        except KeyError:
            raise ModelLoadError("metadata.json/metafile.json not found in model file")
        except json.JSONDecodeError as e:
            raise ModelLoadError(f"Invalid JSON in metadata: {e}")

    def _convert_trainer_metadata(self, trainer_meta: Dict[str, Any]) -> Dict[str, Any]:
        """Convert trainer metafile.json format to ModelMetadata format."""
        # Extract nested data
        training_stats = trainer_meta.get('training_statistics', {})
        audio_cfg = trainer_meta.get('audio_config', {})
        arch_info = trainer_meta.get('architecture_info', {})
        hardware = trainer_meta.get('hardware_info', {})
        model_specific = trainer_meta.get('model_specific', {})

        # Calculate time_frames from audio config
        audio_length_sec = audio_cfg.get('audio_length_seconds', 1.5)
        hop_length = audio_cfg.get('hop_length', 160)
        sample_rate = audio_cfg.get('sample_rate', 16000)
        time_frames = int((audio_length_sec * sample_rate) / hop_length) + 1

        # Get class labels from model_specific.wakewords or fallback
        class_labels = model_specific.get('wakewords', ['custom', 'system_command', 'negative'])
        num_classes = len(class_labels)

        # Build input_shape
        n_mels = audio_cfg.get('n_mels', 40)
        input_shape = [1, n_mels, time_frames]

        # Map to ModelMetadata fields
        return {
            'model_name': trainer_meta.get('model_name', 'unknown'),
            'model_version': trainer_meta.get('version', '1.0.0'),
            'model_type': trainer_meta.get('model_type', 'wakeword'),
            'architecture': arch_info.get('architecture_name', 'RepCNN'),

            'training_date': trainer_meta.get('creation_timestamp', ''),
            'author': trainer_meta.get('author', 'Unknown'),
            'description': trainer_meta.get('description', ''),
            'computer_name': hardware.get('hostname', 'Unknown'),

            'num_classes': num_classes,
            'input_shape': input_shape,
            'class_labels': class_labels,

            'final_accuracy': training_stats.get('best_val_accuracy', 0.0),
            'final_loss': training_stats.get('best_val_loss', 0.0),
            'num_epochs': training_stats.get('total_epochs', 0),
            'training_duration_hours': training_stats.get('training_duration_seconds', 0.0) / 3600.0,

            'sample_rate': sample_rate,
            'n_mels': n_mels,
            'time_frames': time_frames,
            'chunk_duration_ms': audio_length_sec * 1000,

            'inference_time_ms': 0.0,
            'model_size_mb': training_stats.get('total_parameters', 0) * 4 / (1024 * 1024),
            'flops_estimate': 0,

            'model_hash': '',  # Not available in trainer format
            'created_timestamp': trainer_meta.get('creation_timestamp', ''),
        }
    
    def _load_model_state(self, zip_file: zipfile.ZipFile) -> Dict[str, torch.Tensor]:
        """Load model state dict from ZIP file."""
        try:
            with zip_file.open('model.pth') as model_file:
                # Load state dict directly from bytes
                model_bytes = model_file.read()
                state_dict = torch.load(
                    io.BytesIO(model_bytes), 
                    map_location=self.device,
                    weights_only=True
                )
                return state_dict
                
        except KeyError:
            raise ModelLoadError("model.pth not found in model file")
        except Exception as e:
            raise ModelLoadError(f"Failed to load model state dict: {e}")
    
    def _verify_model_hash(self, model_state: Dict[str, torch.Tensor], 
                          expected_hash: str) -> bool:
        """Verify model integrity using hash."""
        try:
            # Create hash of model parameters
            hasher = hashlib.sha256()
            
            # Sort keys for consistent hashing
            for key in sorted(model_state.keys()):
                tensor = model_state[key]
                hasher.update(key.encode('utf-8'))
                hasher.update(tensor.detach().cpu().numpy().tobytes())
            
            calculated_hash = hasher.hexdigest()
            return calculated_hash == expected_hash
            
        except Exception:
            return False
    
    def _create_model_instance(self, metadata: ModelMetadata, 
                             model_class: Optional[type] = None) -> nn.Module:
        """Create model instance based on metadata."""
        if model_class is not None:
            # Use provided model class
            return self._instantiate_model_class(model_class, metadata)
        
        # Attempt auto-detection based on metadata
        architecture = metadata.architecture.lower()
        
        if architecture in ["repcnn", "standard"]:
            from trainer.wakeword.models import RepCNN
            return self._instantiate_repcnn(RepCNN, metadata)
        elif architecture in ["improvedrepcnn", "improved"]:
            from trainer.wakeword.models import ImprovedRepCNN
            return self._instantiate_repcnn(ImprovedRepCNN, metadata)
        elif architecture in ["lightweightrepcnn", "lightweight"]:
            from trainer.wakeword.models import LightweightRepCNN
            return self._instantiate_lightweight_repcnn(metadata)
        else:
            raise ModelLoadError(f"Unknown architecture: {metadata.architecture}")
    
    def _instantiate_repcnn(self, model_class: type, metadata: ModelMetadata) -> nn.Module:
        """Instantiate RepCNN-based models."""
        return model_class(
            num_classes=metadata.num_classes,
            input_channels=metadata.input_shape[0],
            num_mels=metadata.n_mels,
            time_frames=metadata.time_frames
        )
    
    def _instantiate_lightweight_repcnn(self, metadata: ModelMetadata) -> nn.Module:
        """Instantiate LightweightRepCNN."""
        from trainer.wakeword.models import LightweightRepCNN
        return LightweightRepCNN(
            num_classes=metadata.num_classes,
            input_channels=metadata.input_shape[0],
            num_mels=metadata.n_mels,
            time_frames=metadata.time_frames
        )
    
    def _instantiate_model_class(self, model_class: type, metadata: ModelMetadata) -> nn.Module:
        """Instantiate arbitrary model class with metadata parameters."""
        try:
            # Try different initialization patterns
            init_params = {
                'num_classes': metadata.num_classes,
                'input_channels': metadata.input_shape[0],
                'num_mels': metadata.n_mels,
                'time_frames': metadata.time_frames
            }
            
            return model_class(**init_params)
            
        except Exception as e:
            raise ModelLoadError(f"Failed to instantiate model class {model_class.__name__}: {e}")
    
    def get_loaded_models(self) -> List[str]:
        """Get list of currently loaded model paths."""
        return list(self._loaded_models.keys())
    
    def unload_model(self, model_path: str) -> bool:
        """
        Unload a model from cache.
        
        Args:
            model_path: Path of model to unload
            
        Returns:
            True if model was unloaded, False if not found
        """
        if model_path in self._loaded_models:
            del self._loaded_models[model_path]
            return True
        return False
    
    def clear_cache(self):
        """Clear all cached models."""
        self._loaded_models.clear()
    
    def get_model_info(self, model_path: str) -> Optional[Dict[str, Any]]:
        """Get information about a loaded model."""
        return self._loaded_models.get(model_path)


class ModelLoader:
    """
    Standard model loader for unprotected PyTorch models.
    
    This is a simpler loader for development and testing purposes.
    """
    
    def __init__(self, device: Optional[torch.device] = None):
        """Initialize the model loader."""
        self.device = device or torch.device("cpu")
    
    def load_model(self, model_path: Union[str, Path], 
                  model_class: type, **model_kwargs) -> nn.Module:
        """
        Load a standard PyTorch model.
        
        Args:
            model_path: Path to the .pth file
            model_class: Model class to instantiate
            **model_kwargs: Arguments for model initialization
            
        Returns:
            Loaded PyTorch model
        """
        model_path = Path(model_path)
        
        if not model_path.exists():
            raise ModelLoadError(f"Model file not found: {model_path}")
        
        try:
            # Create model instance
            model = model_class(**model_kwargs)
            
            # Load state dict
            state_dict = torch.load(model_path, map_location=self.device, weights_only=True)
            model.load_state_dict(state_dict, strict=True)
            
            model.to(self.device)
            model.eval()
            
            return model
            
        except Exception as e:
            raise ModelLoadError(f"Failed to load model: {e}")


def validate_model_file(model_path: Union[str, Path], password: Optional[str] = None) -> bool:
    """
    Validate a model file without fully loading it.

    Args:
        model_path: Path to model file
        password: Password for encrypted models (optional)

    Returns:
        True if model file is valid, False otherwise
    """
    model_path = Path(model_path)

    if not model_path.exists():
        return False

    try:
        with zipfile.ZipFile(model_path, 'r') as zip_file:
            if password:
                zip_file.setpassword(password.encode('utf-8'))

            # Check for required files
            file_list = zip_file.namelist()

            # Check for model.pth
            if 'model.pth' not in file_list:
                return False

            # Check for metadata file (metafile.json or metadata.json)
            has_metadata = 'metafile.json' in file_list or 'metadata.json' in file_list
            if not has_metadata:
                return False

            # Try to load metadata
            try:
                meta_file_name = 'metafile.json' if 'metafile.json' in file_list else 'metadata.json'
                with zip_file.open(meta_file_name) as meta_file:
                    metadata_dict = json.load(meta_file)
                    metadata = ModelMetadata.from_dict(metadata_dict)
                    return metadata.validate()
            except KeyError:
                return False

    except (zipfile.BadZipFile, json.JSONDecodeError, RuntimeError):
        return False
    except Exception:
        return False


def create_model_metadata(model: nn.Module, training_config: Dict[str, Any],
                         training_results: Dict[str, Any],
                         audio_config: Dict[str, Any]) -> ModelMetadata:
    """
    Create model metadata from training information.
    
    Args:
        model: Trained PyTorch model
        training_config: Training configuration
        training_results: Training results and metrics
        audio_config: Audio processing configuration
        
    Returns:
        ModelMetadata instance
    """
    # Calculate model hash
    hasher = hashlib.sha256()
    for name, param in model.named_parameters():
        hasher.update(name.encode('utf-8'))
        hasher.update(param.detach().cpu().numpy().tobytes())
    model_hash = hasher.hexdigest()
    
    # Estimate model size
    param_count = sum(p.numel() for p in model.parameters())
    model_size_mb = param_count * 4 / (1024 * 1024)  # Assuming float32
    
    return ModelMetadata(
        model_name=training_config.get('model_name', 'wakeword_model'),
        model_version=training_config.get('model_version', '1.0.0'),
        model_type="wakeword",
        architecture=model.__class__.__name__,
        
        training_date=datetime.now().isoformat(),
        author=training_config.get('author', 'Unknown'),
        description=training_config.get('description', ''),
        computer_name=training_config.get('computer_name', 'Unknown'),
        
        num_classes=training_config.get('num_classes', 3),
        input_shape=[1, audio_config.get('n_mels', 40), audio_config.get('time_frames', 151)],
        class_labels=training_config.get('class_labels', ['custom', 'system_command', 'negative']),
        
        final_accuracy=training_results.get('accuracy', 0.0),
        final_loss=training_results.get('loss', 0.0),
        num_epochs=training_results.get('epochs', 0),
        training_duration_hours=training_results.get('duration_hours', 0.0),
        
        sample_rate=audio_config.get('sample_rate', 16000),
        n_mels=audio_config.get('n_mels', 40),
        time_frames=audio_config.get('time_frames', 151),
        chunk_duration_ms=audio_config.get('chunk_duration_ms', 1000),
        
        inference_time_ms=0.0,  # To be measured during deployment
        model_size_mb=model_size_mb,
        flops_estimate=0,  # To be calculated if needed
        
        model_hash=model_hash,
        created_timestamp=datetime.now().isoformat()
    )