src/models/advanced_forensic_auditor.py

import logging
import torch
import numpy as np
from PIL import Image
from typing import List, Dict, Any, Optional
import os
import torchvision.transforms as transforms
from facenet_pytorch import MTCNN, InceptionResnetV1
from .ctm_forensic_agent import get_ctm_agent

logger = logging.getLogger(__name__)

class AdvancedForensicAuditor:
    """
    Advanced Forensic Auditor
    
    Performs deep verification by cross-referencing a probe image against
    candidate source identities, guided by CTM attention.
    """
    
    def __init__(self, device: str = None):
        self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Load CTM Agent (Lazy)
        self.ctm_agent = get_ctm_agent()
        
        # Load Comparison Model (FaceNet for region embeddings)
        # We use a separate instance or reuse? Reusing might be complex if IdentityMatcher is singleton.
        # Let's create a small instance just for feature extraction on crops.
        self.resnet = InceptionResnetV1(pretrained='vggface2').eval().to(self.device)
        self.mtcnn = MTCNN(keep_all=True, device=self.device)
        
        # Standard regions for fallback
        self.standard_regions = {
            'eyes': (0.2, 0.45, 0.2, 0.8), # y1, y2, x1, x2 (approx relative coords)
            'nose': (0.4, 0.65, 0.3, 0.7),
            'mouth': (0.65, 0.9, 0.25, 0.75)
        }
        
    def _get_crop(self, img: Image.Image, region_name: str) -> Optional[Image.Image]:
        """Crop a specific region from the image based on approximate relative coordinates"""
        coords = self.standard_regions.get(region_name)
        if not coords:
            return None
            
        y1, y2, x1, x2 = coords
        w, h = img.size
        
        left = int(x1 * w)
        top = int(y1 * h)
        right = int(x2 * w)
        bottom = int(y2 * h)
        
        return img.crop((left, top, right, bottom))
        
    def _compute_similarity(self, img1: Image.Image, img2: Image.Image) -> float:
        """Compute cosine similarity between two image crops using FaceNet"""
        try:
            # Resize to expected input size
            t1 = transforms.functional.resize(img1, (160, 160))
            t2 = transforms.functional.resize(img2, (160, 160))
            
            # To Tensor
            t1 = transforms.functional.to_tensor(t1).unsqueeze(0).to(self.device)
            t2 = transforms.functional.to_tensor(t2).unsqueeze(0).to(self.device)
            
            # Normalize? InceptionResnet expects specific normalization usually? 
            # Facenet-pytorch usually handles it if using their preprocess, but here we do manual.
            # Let's assume standard normalization.
            mean = torch.as_tensor([0.5, 0.5, 0.5], device=self.device).view(1, 3, 1, 1)
            std = torch.as_tensor([0.5, 0.5, 0.5], device=self.device).view(1, 3, 1, 1)
            t1 = (t1 - mean) / std
            t2 = (t2 - mean) / std
            
            with torch.no_grad():
                emb1 = self.resnet(t1)
                emb2 = self.resnet(t2)
                
            sim = torch.nn.functional.cosine_similarity(emb1, emb2).item()
            return (sim + 1) / 2 # Normalize -1..1 to 0..1
            
        except Exception as e:
            logger.error(f"Error computing similarity: {e}")
            return 0.0

    def audit(self, probe_img: Image.Image, candidate_paths: List[str]) -> Dict[str, Any]:
        """
        Perform the full audit.
        
        Args:
            probe_img: The morphed/probe image (PIL)
            candidate_paths: List of file paths to candidate images
            
        Returns:
            Dict containing report and details
        """
        report = {
            'ctm_analysis': None,
            'matches': [],
            'conclusion': "No conclusion."
        }
        
        # 1. Run CTM on Probe
        logger.info("Running CTM analysis on probe...")
        ctm_result = self.ctm_agent.analyze(probe_img, generate_evidence=False)
        
        # Map CTM regions to our crop regions (simplified mapping)
        relevant_regions = []
        for region in ctm_result.attention_regions:
            if 'eye' in region: relevant_regions.append('eyes')
            if 'mouth' in region: relevant_regions.append('mouth')
            if 'nose' in region: relevant_regions.append('nose')
        
        # Deduplicate
        relevant_regions = list(set(relevant_regions))
        if not relevant_regions:
            relevant_regions = ['eyes', 'mouth'] # Default fallback
            
        report['ctm_analysis'] = {
            'is_morphed': ctm_result.is_morphed,
            'suspicious_regions': relevant_regions, # mapped regions
            'raw_regions': ctm_result.attention_regions
        }
        
        # 2. Compare against Candidates
        logger.info(f"Auditing against {len(candidate_paths)} candidates...")
        for cand_path in candidate_paths:
            try:
                cand_img = Image.open(cand_path).convert('RGB')
                cand_name = os.path.basename(cand_path)
                
                match_details = {
                    'filename': cand_name,
                    'region_scores': {}
                }
                
                for region in relevant_regions:
                    # Crop
                    probe_crop = self._get_crop(probe_img, region)
                    cand_crop = self._get_crop(cand_img, region)
                    
                    if probe_crop and cand_crop:
                        score = self._compute_similarity(probe_crop, cand_crop)
                        match_details['region_scores'][region] = score
                
                report['matches'].append(match_details)
                
            except Exception as e:
                logger.error(f"Failed to audit candidate {cand_path}: {e}")
        
        # 3. Generate Conclusion
        # Simple logic: High match in one region for Cand A and high match in another for Cand B?
        conclusion = "Audit completed. "
        if ctm_result.is_morphed:
            conclusion += "CTM flagged this image as MORPHED. "
        else:
            conclusion += "CTM did not flag major morphing artifacts. "
            
        # Analyze regional matches
        # Find best match for each region
        best_matches = {}
        for region in relevant_regions:
            best_score = -1
            best_cand = None
            for m in report['matches']:
                s = m['region_scores'].get(region, 0)
                if s > best_score:
                    best_score = s
                    best_cand = m['filename']
            
            if best_score > 0.75: # Threshold
                best_matches[region] = (best_cand, best_score)
        
        if best_matches:
            conclusion += "Regional Analysis suggests: "
            parts = []
            for reg, (name, score) in best_matches.items():
                parts.append(f"the {reg} resemble {name} ({score:.0%})")
            conclusion += ", and ".join(parts) + "."
        else:
            conclusion += "No strong regional matches found among candidates."
            
        report['conclusion'] = conclusion
        return report