backend/utils/enhanced_ml_service.py

# pylint:  disable=wrong-import-order, unused-import
"""
Enhanced API endpoints with explainability features.
Extends the existing FastAPI backend with SHAP-based model explanations
and improved prediction capabilities.
"""
from backend.config import TARGET_LEN  # Import TARGET_LEN for model loading
import numpy as np
import torch
from typing import Dict, Any, List, Optional
from fastapi import HTTPException  # Keep HTTPException for API errors
# PredictionResult is not directly returned by this service
from backend.pydantic_models import SpectrumData
from backend.models.registry import build as build_model, choices, registry_spec
from backend.utils.preprocessing_fixed import SpectrumPreprocessor

import os
# Import moved here to the toplevel
from backend.utils.model_manager import model_manager


class EnhancedMLService:
    """
    Enhanced ML service with explainability features.
    Provides predictions with feature importance and model confidence.
    """

    def __init__(self):
        self.model_manager = model_manager
        # Local cache for loaded models (model, preprocessor)
        self._model_cache = {}
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu")
        print(f"✅ Enhanced ML Service initialized on {self.device}")


    def cache_model(self, model_name: str, model_instance, preprocessor):
        """Public method to cache a model and its preprocessor."""
        self._model_cache[model_name] = {
            'model': model_instance,
            'preprocessor': preprocessor
        }

    def predict_with_explanation(
        self,
        spectrum_data: SpectrumData,
        model_name: str,
        modality: str = "raman",
        include_feature_importance: bool = True
    ) -> Dict[str, Any]:
        """
        Make prediction with explainability features.

        Args:
            spectrum_data (SpectrumData): Input spectrum data
            model_name (str): Name of model to use
            modality (str): The spectroscopy modality ('raman' or 'ftir')
            include_feature_importance (bool): Whether to compute feature importance

        Returns:
            dict: Prediction results with explanations
        """
        if model_name not in self._model_cache:
            # Attempt to load model via centralized manager if not in local cache
            model_instance, weights_loaded, _ = self.model_manager.load_model(
                model_name)
            if model_instance is None or not weights_loaded:
                raise HTTPException(
                    status_code=400,
                    detail=f"Model {model_name} not loaded or weights not found"
                )

            # Determine model input length robustly: prefer model attribute,
            # fallback to registry/spec, then TARGET_LEN
            input_len = getattr(model_instance, 'input_length', None)
            if input_len is None:
                try:
                    spec_info = registry_spec(model_name)
                    input_len = int(spec_info.get("input_length", TARGET_LEN))
                except Exception:
                    input_len = TARGET_LEN

            # Create preprocessor for this model (use resolved input_len)
            preprocessor = SpectrumPreprocessor(
                target_len=input_len,
                do_baseline=True,
                do_smooth=True,
                do_normalize=True,
                modality=modality  # Use the provided modality
            )
            self._model_cache[model_name] = {
                'model': model_instance, 'preprocessor': preprocessor}

        model_entry = self._model_cache.get(model_name)
        if not model_entry:  # Should not happen if previous block executed
            raise HTTPException(
                status_code=400,
                detail=f"Model {model_name} not loaded"
            )
        model = model_entry['model']

        # --- FIX: Ensure preprocessor has the correct modality ---
        # The preprocessor might have been cached with a default or different modality.
        # We must ensure it matches the one from the current request.
        if model_entry['preprocessor'].modality != modality:
            print(
                f"🔄 Updating preprocessor modality for '{model_name}' from '{model_entry['preprocessor'].modality}' to '{modality}'")
            model_entry['preprocessor'] = SpectrumPreprocessor(
                target_len=model.input_length,
                do_baseline=True, do_smooth=True, do_normalize=True,
                modality=modality
            )

        preprocessor = model_entry['preprocessor']

        try:
            # Preprocess input data
            x_data = np.array(spectrum_data.x_values)
            y_data = np.array(spectrum_data.y_values)

            # Preprocess spectrum
            processed_spectrum = preprocessor.preprocess_single_spectrum(
                x_data, y_data, use_fitted_stats=False
            )

            # Convert to tensor
            input_tensor = torch.tensor(
                processed_spectrum, dtype=torch.float32)
            # Add batch and channel dimensions
            input_tensor = input_tensor.unsqueeze(0)
            input_tensor = input_tensor.unsqueeze(0)
            input_tensor = input_tensor.to(self.device)

            # Make prediction
            with torch.no_grad():
                outputs = model(input_tensor)
                probabilities = torch.softmax(outputs, dim=1)
                predicted_class = torch.argmax(probabilities, dim=1).item()
                confidence = torch.max(probabilities).item()

            # Basic prediction result
            result = {
                'prediction': predicted_class,
                'confidence': confidence,
                'probabilities': probabilities.cpu().numpy().tolist()[0],
                'class_labels': ['stable', 'weathered'],
                'model_used': model_name,
                'spectrum_filename': spectrum_data.filename
            }

            # Add feature importance if requested
            if include_feature_importance:
                feature_importance = self._compute_feature_importance(
                    model, input_tensor, processed_spectrum
                )
                result['feature_importance'] = feature_importance

            return result

        except (RuntimeError, ValueError, TypeError) as e:
            raise HTTPException(
                status_code=500,
                detail=f"Prediction failed: {str(e)}"
            ) from e

    def _compute_feature_importance(
        self,
        model: torch.nn.Module,
        input_tensor: torch.Tensor,
        processed_spectrum: np.ndarray
    ) -> Dict[str, Any]:
        """
        Compute feature importance using gradient-based methods.

        Args:
            model: PyTorch model
            input_tensor: Preprocessed input tensor
            processed_spectrum: Original processed spectrum

        Returns:
            dict: Feature importance information
        """
        try:
            # Enable gradient computation
            input_tensor.requires_grad_(True)
            torch.set_grad_enabled(True)

            # Forward pass
            output = model(input_tensor)
            predicted_class = torch.argmax(output, dim=1).item()

            # Compute gradients with respect to input
            class_score = output[0, predicted_class]
            class_score.backward()

            if input_tensor.grad is not None:
                gradients = input_tensor.grad.data.cpu().numpy().squeeze()
            else:
                raise RuntimeError(
                    "Gradients were not computed. Ensure requires_grad is set "
                    "and gradient computation is enabled."
                )
            gradients = input_tensor.grad.data.cpu().numpy().squeeze()

            # Compute importance metrics
            importance_abs = np.abs(gradients)

            # Find most important regions
            top_indices = np.argsort(importance_abs)[-20:]  # Top 20 features

            # Create interpretable output
            feature_importance = {
                'method': 'gradient_saliency',
                'importance_scores': importance_abs.tolist(),
                'top_features': {
                    'indices': top_indices.tolist(),
                    'values': importance_abs[top_indices].tolist()
                },
                'summary': {
                    'max_importance': float(np.max(importance_abs)),
                    'mean_importance': float(np.mean(importance_abs)),
                    'important_region_start': int(top_indices[0]),
                    'important_region_end': int(top_indices[-1])
                }
            }

            return feature_importance

        except (RuntimeError, ValueError, TypeError) as e:
            print(f"⚠️ Feature importance computation failed: {e}")
            return {
                'method': 'gradient_saliency',
                'error': str(e),
                'importance_scores': [0.0] * len(processed_spectrum)
            }

    def get_model_info(self) -> List[Dict[str, Any]]:
        """
        Get information about loaded models.

        Returns:
            list: List of ModelInfo objects from the centralized manager.
        """
        return self.model_manager.get_available_models()

    def batch_predict_with_explanation(
        self,
        spectra: List[SpectrumData],
        model_name: str,
        modality: str,  # Add modality for preprocessor
        include_feature_importance: bool = True
    ) -> List[Dict[str, Any]]:
        """
        Batch prediction with explanations.

        Args:
            spectra (list): List of spectrum data
            model_name (str): Model to use
            modality (str): Spectroscopy modality
            include_feature_importance (bool): Whether to include explanations

        Returns:
            list: List of prediction results
        """
        results = []

        for spectrum in spectra:
            try:
                result = self.predict_with_explanation(
                    spectrum,
                    model_name,
                    modality=modality,  # Pass modality down
                    include_feature_importance=include_feature_importance
                )
                results.append(result)
            except (HTTPException, ValueError, RuntimeError) as e:
                results.append({
                    'error': str(e),
                    'spectrum_filename': spectrum.filename
                })

        return results


# Global enhanced service instance
enhanced_ml_service = EnhancedMLService()


def initialize_enhanced_service():
    """Initialize the enhanced ML service with available models."""
    print("Initializing Enhanced ML Service models...")
    # Iterate through all known models in the registry by calling choices() directly
    for model_name in choices():
        try:
            # Attempt to load each model via the centralized manager
            model_instance, weights_loaded, _ = enhanced_ml_service.model_manager.load_model(
                model_name, TARGET_LEN)
            if model_instance and weights_loaded:
                preprocessor = SpectrumPreprocessor(
                    target_len=TARGET_LEN,
                    do_baseline=True,
                    do_smooth=True,
                    do_normalize=True,
                    modality="raman"
                )
                enhanced_ml_service.cache_model(model_name, model_instance, preprocessor)
                print(f"✅ Enhanced ML Service: Prepared model '{model_name}' with preprocessor.")
            else:
                print(
                    f"⚠️ Enhanced ML Service: Model '{model_name}' not fully loaded or weights missing.")
        except (RuntimeError, ValueError, ImportError) as e:
            print(
                f"❌ Enhanced ML Service: Error initializing model '{model_name}': {e}")


# Initialize on import
initialize_enhanced_service()