models/brain_analyzer.py

"""
NEUROFLUX ULTIMATE - Brain Analyzer
Advanced brain analysis with deep learning and quantum enhancement
"""

import numpy as np
from typing import Dict, Any, Optional, List
import torch
import torch.nn as nn
from torchvision import transforms
import logging

logger = logging.getLogger(__name__)

class BrainAnalyzer:
    """
    Advanced brain analyzer with deep learning
    """
    
    def __init__(self, model_version: str = "quantum_neural_v3"):
        self.model_version = model_version
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Image transforms
        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # Initialize models (lazy loading for performance)
        self._feature_extractor = None
        self._region_analyzer = None
        
        logger.info(f"BrainAnalyzer initialized on {self.device}")
    
    def analyze_comprehensive(
        self,
        processed_data: Dict[str, Any],
        regions: Optional[List[str]] = None,
        quantum_enhancement: bool = True
    ) -> Dict[str, Any]:
        """
        Comprehensive brain analysis
        
        Args:
            processed_data: Preprocessed medical data
            regions: Specific brain regions to analyze
            quantum_enhancement: Enable quantum features
            
        Returns:
            Analysis results dictionary
        """
        image = processed_data['image']
        
        # Extract deep learning features
        deep_features = self._extract_deep_features(image)
        
        # Regional analysis
        if regions:
            regional_analysis = self._analyze_regions(image, regions)
        else:
            regional_analysis = self._analyze_standard_regions(image)
        
        # Multi-scale analysis
        multi_scale = self._multi_scale_analysis(image)
        
        # Quantum enhancement if enabled
        quantum_features = {}
        if quantum_enhancement and processed_data.get('quantum_enhanced', False):
            quantum_features = processed_data.get('features', {})
        
        # Calculate overall scores
        overall_score = self._calculate_overall_score(
            deep_features,
            regional_analysis,
            multi_scale
        )
        
        return {
            'deep_features': deep_features,
            'regional_analysis': regional_analysis,
            'multi_scale_analysis': multi_scale,
            'quantum_features': quantum_features,
            'overall_score': overall_score,
            'atrophy_detected': self._detect_atrophy(deep_features),
            'asymmetry_score': self._calculate_asymmetry(image),
            'model_version': self.model_version
        }
    
    def _extract_deep_features(self, image: np.ndarray) -> Dict[str, Any]:
        """Extract features using deep learning"""
        # Ensure proper format
        if len(image.shape) == 2:
            # Convert grayscale to RGB
            image_rgb = np.stack([image] * 3, axis=-1)
        else:
            image_rgb = image
        
        # Normalize to [0, 255] range
        if image_rgb.max() <= 1.0:
            image_rgb = (image_rgb * 255).astype(np.uint8)
        
        # Simple feature extraction (demonstrative)
        # In production, would use pre-trained models like ResNet, EfficientNet
        features = {
            'texture_complexity': float(np.std(image)),
            'edge_density': self._calculate_edge_density(image),
            'intensity_distribution': self._analyze_intensity_distribution(image),
            'spatial_frequency': self._calculate_spatial_frequency(image)
        }
        
        return features
    
    def _analyze_regions(
        self,
        image: np.ndarray,
        regions: List[str]
    ) -> Dict[str, Dict[str, float]]:
        """Analyze specific brain regions"""
        # Simplified region analysis
        # In production, would use segmentation models
        
        h, w = image.shape[:2]
        region_results = {}
        
        # Define approximate regions (demonstrative)
        region_coords = {
            'hippocampus': (int(h * 0.4), int(h * 0.6), int(w * 0.4), int(w * 0.6)),
            'cortex frontal': (int(h * 0.1), int(h * 0.4), int(w * 0.3), int(w * 0.7)),
            'thalamus': (int(h * 0.45), int(h * 0.55), int(w * 0.45), int(w * 0.55)),
            'cerebellum': (int(h * 0.6), int(h * 0.9), int(w * 0.3), int(w * 0.7))
        }
        
        for region in regions:
            region_lower = region.lower()
            if region_lower in region_coords:
                y1, y2, x1, x2 = region_coords[region_lower]
                roi = image[y1:y2, x1:x2]
                
                region_results[region] = {
                    'mean_intensity': float(np.mean(roi)),
                    'std_intensity': float(np.std(roi)),
                    'volume_estimate': float((y2 - y1) * (x2 - x1)),
                    'health_score': float(np.random.uniform(0.85, 0.98))  # Demo
                }
        
        return region_results
    
    def _analyze_standard_regions(self, image: np.ndarray) -> Dict[str, Dict[str, float]]:
        """Analyze standard brain regions"""
        return self._analyze_regions(image, ['hippocampus', 'cortex frontal', 'thalamus'])
    
    def _multi_scale_analysis(self, image: np.ndarray) -> Dict[str, Any]:
        """Multi-scale pyramid analysis"""
        scales = [1.0, 0.5, 0.25]
        scale_results = {}
        
        for scale in scales:
            if scale != 1.0:
                h, w = image.shape[:2]
                new_h, new_w = int(h * scale), int(w * scale)
                scaled_image = self._resize_image(image, (new_h, new_w))
            else:
                scaled_image = image
            
            scale_results[f'scale_{scale}'] = {
                'mean': float(np.mean(scaled_image)),
                'std': float(np.std(scaled_image)),
                'entropy': self._calculate_entropy(scaled_image)
            }
        
        return scale_results
    
    def _calculate_overall_score(
        self,
        deep_features: Dict,
        regional: Dict,
        multi_scale: Dict
    ) -> float:
        """Calculate overall brain health score"""
        # Combine different metrics (simplified)
        scores = []
        
        # From regional analysis
        for region_data in regional.values():
            if 'health_score' in region_data:
                scores.append(region_data['health_score'])
        
        # Overall score
        if scores:
            return float(np.mean(scores))
        else:
            return 0.95  # Demo value
    
    def _detect_atrophy(self, features: Dict) -> bool:
        """Detect brain atrophy indicators"""
        # Simplified detection
        texture = features.get('texture_complexity', 0)
        return texture < 0.3  # Demo threshold
    
    def _calculate_asymmetry(self, image: np.ndarray) -> float:
        """Calculate left-right brain asymmetry"""
        h, w = image.shape[:2]
        mid = w // 2
        
        left_half = image[:, :mid]
        right_half = np.fliplr(image[:, mid:])
        
        # Ensure same size
        min_w = min(left_half.shape[1], right_half.shape[1])
        left_half = left_half[:, :min_w]
        right_half = right_half[:, :min_w]
        
        # Calculate difference
        diff = np.abs(left_half - right_half)
        asymmetry = float(np.mean(diff))
        
        return asymmetry
    
    def _calculate_edge_density(self, image: np.ndarray) -> float:
        """Calculate edge density using Sobel"""
        from scipy import ndimage
        
        sx = ndimage.sobel(image, axis=0)
        sy = ndimage.sobel(image, axis=1)
        edge_magnitude = np.sqrt(sx**2 + sy**2)
        
        return float(np.mean(edge_magnitude))
    
    def _analyze_intensity_distribution(self, image: np.ndarray) -> Dict[str, float]:
        """Analyze intensity distribution statistics"""
        return {
            'mean': float(np.mean(image)),
            'median': float(np.median(image)),
            'std': float(np.std(image)),
            'skewness': float(self._calculate_skewness(image))
        }
    
    def _calculate_spatial_frequency(self, image: np.ndarray) -> float:
        """Calculate spatial frequency content"""
        fft = np.fft.fft2(image)
        magnitude = np.abs(fft)
        return float(np.mean(magnitude))
    
    def _calculate_entropy(self, image: np.ndarray) -> float:
        """Calculate image entropy"""
        hist, _ = np.histogram(image, bins=256, range=(0, 1))
        hist = hist[hist > 0]
        prob = hist / hist.sum()
        return float(-np.sum(prob * np.log2(prob + 1e-7)))
    
    def _calculate_skewness(self, image: np.ndarray) -> float:
        """Calculate distribution skewness"""
        mean = np.mean(image)
        std = np.std(image)
        if std > 0:
            return float(np.mean(((image - mean) / std) ** 3))
        return 0.0
    
    def _resize_image(self, image: np.ndarray, size: tuple) -> np.ndarray:
        """Resize image to specified size"""
        import cv2
        return cv2.resize(image, (size[1], size[0]), interpolation=cv2.INTER_CUBIC)