src/models/ctm_forensic_agent.py

"""
CTM Forensic Agent for MorphGuard
=================================

Deep Forensic Reasoning Engine using Sakana AI's Continuous Thought Machine.
This provides "System 2" (deliberative) analysis for detecting geometric inconsistencies
in high-quality morphed images that standard models miss.

Key Features:
- Temporal reasoning loop (10-20 "ticks" of thinking)
- Attention-based sequential image scanning
- Evidence video generation showing detection reasoning
- Adaptive early exit when confidence exceeds threshold
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os
import time
import logging
from typing import Dict, Any, Optional, List, Tuple
from dataclasses import dataclass
from PIL import Image
import torchvision.transforms as transforms

# Import CTM wrapper
try:
    from external import CTMModel, CTMConfig, CTMOutput, plot_attention_video, CTM_AVAILABLE
except ImportError:
    CTM_AVAILABLE = False
    CTMModel = None
    CTMConfig = None

logger = logging.getLogger(__name__)


@dataclass
class ForensicResult:
    """Result from CTM forensic analysis"""
    is_morphed: bool
    confidence: float
    forensic_steps: int
    evidence_video_path: Optional[str]
    attention_regions: List[str]
    processing_time_ms: float
    tier: str = "Tier2_Forensic"
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'is_morphed': self.is_morphed,
            'confidence': self.confidence,
            'forensic_steps': self.forensic_steps,
            'evidence_video_path': self.evidence_video_path,
            'attention_regions': self.attention_regions,
            'processing_time_ms': self.processing_time_ms,
            'tier': self.tier
        }


class CTMForensicAgent:
    """
    Deep Forensic Reasoning Agent using Continuous Thought Machine.
    
    This agent performs multi-step reasoning on face images to detect
    morphing artifacts that standard single-pass models miss, such as:
    - Geometric inconsistencies (ear symmetry, hairline logic)
    - Light/shadow coherence across facial features
    - Micro-texture continuity at blend boundaries
    
    Usage:
        agent = CTMForensicAgent('models/ctm_forensic.pth')
        result = agent.analyze(image_tensor)
        print(f"Is Morphed: {result.is_morphed}, Steps: {result.forensic_steps}")
    """
    
    def __init__(
        self,
        checkpoint_path: Optional[str] = None,
        device: str = 'cuda' if torch.cuda.is_available() else 'cpu',
        max_thoughts: int = 16,
        confidence_threshold: float = 0.95,
        early_exit_confidence: float = 0.99,
        evidence_output_dir: str = 'static/evidence'
    ):
        """
        Initialize CTM Forensic Agent.
        
        Args:
            checkpoint_path: Path to CTM model checkpoint
            device: Device to run inference on ('cuda' or 'cpu')
            max_thoughts: Maximum number of thinking steps
            confidence_threshold: Threshold above which fast detector result is trusted
            early_exit_confidence: Confidence level for early exit during thinking
            evidence_output_dir: Directory to save evidence videos
        """
        self.device = device
        self.max_thoughts = max_thoughts
        self.confidence_threshold = confidence_threshold
        self.early_exit_confidence = early_exit_confidence
        self.evidence_output_dir = evidence_output_dir
        self.checkpoint_path = checkpoint_path
        
        # Lazy loading - model loaded on first use
        self._model: Optional[nn.Module] = None
        self._is_loaded = False
        
        # Image preprocessing (ImageNet normalization)
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])
        ])
        
        # Ensure evidence output directory exists
        os.makedirs(evidence_output_dir, exist_ok=True)
        
        # Facial region labels for interpretability
        self.region_labels = [
            'left_eye', 'right_eye', 'nose', 'mouth', 'left_ear', 
            'right_ear', 'forehead', 'chin', 'left_cheek', 'right_cheek',
            'jawline', 'hairline', 'neck', 'background'
        ]
        
        logger.info(f"CTMForensicAgent initialized (lazy loading enabled)")
        logger.info(f"  Device: {device}, Max thoughts: {max_thoughts}")
        logger.info(f"  CTM Available: {CTM_AVAILABLE}")
    
    def load_model(self) -> bool:
        """Load CTM model (called lazily on first inference)"""
        if self._is_loaded:
            return True
        
        if not CTM_AVAILABLE or CTMModel is None:
            logger.warning("CTM not available, skipping CTM load")
            return False
        
        try:
            # Load or initialize model
            if self.checkpoint_path:
                if not os.path.exists(self.checkpoint_path):
                    raise FileNotFoundError(f"CTM checkpoint not found at {self.checkpoint_path}")
                
                logger.info(f"Loading CTM from {self.checkpoint_path}")
                self._model = CTMModel.load_from_checkpoint(
                    self.checkpoint_path, 
                    device=self.device
                )
            else:
                raise ValueError("CTM checkpoint path must be provided to avoid random weights")
            
            self._model.eval()
            self._is_loaded = True
            
            # Warmup inference
            self._warmup()
            
            logger.info("CTM model loaded successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to load CTM model: {e}")
            return False
    
    def _warmup(self):
        """Warmup model with dummy input"""
        try:
            dummy = torch.randn(1, 3, 224, 224, device=self.device)
            with torch.no_grad():
                self._model(dummy, max_thoughts=2)
            logger.info("CTM warmup complete")
        except Exception as e:
            logger.warning(f"CTM warmup failed: {e}")
    
    def unload_model(self):
        """Unload model to free GPU memory"""
        if self._model is not None:
            del self._model
            self._model = None
            self._is_loaded = False
            torch.cuda.empty_cache()
            logger.info("CTM model unloaded")
    
    def preprocess(self, image: Any) -> torch.Tensor:
        """
        Preprocess image for CTM inference.
        
        Args:
            image: PIL Image, numpy array, or file path
            
        Returns:
            Preprocessed tensor (1, 3, 224, 224)
        """
        if isinstance(image, str):
            image = Image.open(image).convert('RGB')
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert('RGB')
        elif isinstance(image, torch.Tensor):
            # Already a tensor, just ensure correct shape
            if image.dim() == 3:
                image = image.unsqueeze(0)
            return image.to(self.device)
        
        tensor = self.transform(image).unsqueeze(0)
        return tensor.to(self.device)
    
    def analyze(
        self,
        image: Any,
        generate_evidence: bool = True,
        request_id: Optional[str] = None
    ) -> ForensicResult:
        """
        Perform deep forensic reasoning on a face image.
        
        This is the main entry point for CTM analysis. The model will:
        1. Initialize neural state
        2. Iteratively attend to different facial regions
        3. Update internal representation with each "thought"
        4. Exit early if highly confident
        5. Generate evidence video if requested
        
        Args:
            image: Input image (PIL, numpy, path, or tensor)
            generate_evidence: Whether to generate attention replay GIF
            request_id: Optional ID for naming evidence files
            
        Returns:
            ForensicResult with prediction, confidence, and evidence
        """
        start_time = time.time()
        
        # Ensure model is loaded
        if not self._is_loaded:
            if not self.load_model():
                return ForensicResult(
                    is_morphed=False,
                    confidence=0.0,
                    forensic_steps=0,
                    evidence_video_path=None,
                    attention_regions=[],
                    processing_time_ms=0.0
                )
        
        # Preprocess image
        image_tensor = self.preprocess(image)
        
        # Keep original for visualization
        if isinstance(image, str):
            original_image = Image.open(image).convert('RGB')
            original_tensor = transforms.ToTensor()(original_image).unsqueeze(0)
        elif isinstance(image, np.ndarray):
            original_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float() / 255.0
        else:
            original_tensor = image_tensor.clone()
        
        # Run CTM inference
        with torch.no_grad():
            output, attention_history = self._model(
                image_tensor,
                max_thoughts=self.max_thoughts,
                early_exit_threshold=self.early_exit_confidence
            )
        
        # Analyze which regions were attended to
        attention_regions = self._analyze_attention_regions(attention_history)
        
        # Generate evidence video
        evidence_path = None
        if generate_evidence and attention_history:
            evidence_filename = f"forensic_{request_id or int(time.time())}.gif"
            evidence_path = os.path.join(self.evidence_output_dir, evidence_filename)
            try:
                plot_attention_video(attention_history, original_tensor, evidence_path)
            except Exception as e:
                logger.warning(f"Failed to generate evidence video: {e}")
                evidence_path = None
        
        processing_time = (time.time() - start_time) * 1000  # Convert to ms
        
        # Create result
        result = ForensicResult(
            is_morphed=output.prediction == 1,  # 1 = morphed, 0 = real
            confidence=output.confidence,
            forensic_steps=len(attention_history),
            evidence_video_path=evidence_path,
            attention_regions=attention_regions,
            processing_time_ms=processing_time
        )
        
        logger.info(f"CTM analysis complete: morphed={result.is_morphed}, "
                   f"confidence={result.confidence:.3f}, steps={result.forensic_steps}")
        
        return result
    
    def _analyze_attention_regions(
        self, 
        attention_history: List[torch.Tensor]
    ) -> List[str]:
        """
        Analyze attention maps to determine which facial regions were focused on.
        
        Maps attention peak locations to semantic facial regions.
        """
        if not attention_history:
            return []
        
        regions_attended = set()
        
        for attention in attention_history:
            if isinstance(attention, torch.Tensor):
                attn = attention.numpy() if attention.dim() == 2 else attention[0].numpy()
            else:
                attn = attention
            
            # Find peak attention location
            peak_idx = np.unravel_index(np.argmax(attn), attn.shape)
            
            # Map to facial region based on location in 14x14 grid
            # Approximate facial region mapping
            y, x = peak_idx
            region = self._map_location_to_region(x, y, attn.shape[1], attn.shape[0])
            regions_attended.add(region)
        
        return list(regions_attended)
    
    def _map_location_to_region(
        self, 
        x: int, 
        y: int, 
        width: int, 
        height: int
    ) -> str:
        """Map grid location to facial region name"""
        # Normalize to 0-1 range
        nx = x / width
        ny = y / height
        
        # Simple region mapping based on typical face proportions
        if ny < 0.25:
            return 'forehead' if 0.25 < nx < 0.75 else 'hairline'
        elif ny < 0.45:
            if nx < 0.35:
                return 'left_eye'
            elif nx > 0.65:
                return 'right_eye'
            else:
                return 'nose'
        elif ny < 0.65:
            if nx < 0.2:
                return 'left_ear'
            elif nx > 0.8:
                return 'right_ear'
            elif nx < 0.35:
                return 'left_cheek'
            elif nx > 0.65:
                return 'right_cheek'
            else:
                return 'nose'
        elif ny < 0.8:
            if 0.3 < nx < 0.7:
                return 'mouth'
            else:
                return 'jawline'
        else:
            return 'chin' if 0.3 < nx < 0.7 else 'neck'
    
    def should_use_ctm(self, fast_confidence: float, mode: str = 'auto') -> bool:
        """
        Determine whether to escalate to CTM analysis.
        
        Args:
            fast_confidence: Confidence from fast Tier 1 detector
            mode: Detection mode ('fast', 'deep_forensics', 'auto')
            
        Returns:
            True if CTM should be used
        """
        if mode == 'deep_forensics':
            return True
        elif mode == 'fast':
            return False
        else:  # auto mode
            return fast_confidence < self.confidence_threshold
    
    @property
    def is_loaded(self) -> bool:
        """Check if model is currently loaded"""
        return self._is_loaded
    
    def get_status(self) -> Dict[str, Any]:
        """Get agent status for monitoring"""
        return {
            'loaded': self._is_loaded,
            'device': self.device,
            'max_thoughts': self.max_thoughts,
            'confidence_threshold': self.confidence_threshold,
            'ctm_available': CTM_AVAILABLE,
            'checkpoint_path': self.checkpoint_path
        }


# Singleton instance for lazy initialization
_ctm_agent: Optional[CTMForensicAgent] = None


def get_ctm_agent(
    checkpoint_path: Optional[str] = None,
    **kwargs
) -> CTMForensicAgent:
    """
    Get or create singleton CTMForensicAgent instance.
    
    Args:
        checkpoint_path: Path to CTM checkpoint
        **kwargs: Additional arguments for CTMForensicAgent
        
    Returns:
        CTMForensicAgent instance
    """
    global _ctm_agent
    
    if _ctm_agent is None:
        # Get default checkpoint path from environment
        if checkpoint_path is None:
            checkpoint_path = os.environ.get(
                'CTM_CHECKPOINT_PATH',
                'models/ctm_forensic.pth'
            )
        
        # Get other config from environment
        max_thoughts = int(os.environ.get('CTM_MAX_THOUGHTS', 16))
        confidence_threshold = float(os.environ.get('CTM_CONFIDENCE_THRESHOLD', 0.95))
        
        _ctm_agent = CTMForensicAgent(
            checkpoint_path=checkpoint_path,
            max_thoughts=max_thoughts,
            confidence_threshold=confidence_threshold,
            **kwargs
        )
    
    return _ctm_agent


def reset_ctm_agent():
    """Reset singleton agent (useful for testing)"""
    global _ctm_agent
    if _ctm_agent is not None:
        _ctm_agent.unload_model()
        _ctm_agent = None