detectors/frequency.py

"""
Frequency-domain detector for GAN/diffusion artifacts.
Computes 2D FFT, extracts spectral features, uses lightweight classifier.
Always available, no model download needed.
"""

import torch
import numpy as np
import cv2
from PIL import Image
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
import joblib
import os
from .base import BaseDetector, DetectorResult


class FrequencyDetector(BaseDetector):
    """
    Frequency-domain analysis for periodic upsampling artifacts.
    Detects checkerboard patterns, grid artifacts from GAN/diffusion upsampling.
    """

    def __init__(self, device: torch.device):
        super().__init__(device)
        self._classifier = None
        self._scaler = None
        self._fitted = False

    @property
    def model_name(self) -> str:
        return "Frequency Analysis (FFT + LR)"

    def load(self) -> None:
        """Initialize or load pre-trained lightweight classifier."""
        # Try to load pre-fitted model
        model_path = os.path.join(os.path.dirname(__file__), "..", "frequency_model.pkl")
        if os.path.exists(model_path):
            try:
                data = joblib.load(model_path)
                self._classifier = data["classifier"]
                self._scaler = data["scaler"]
                self._fitted = True
                print(f"[Detector] Loaded pre-trained {self.model_name}")
                return
            except Exception as e:
                print(f"[Detector] Could not load frequency model: {e}")
        
        # Initialize fresh - will be fitted on first use with synthetic data
        self._classifier = LogisticRegression(
            C=1.0, max_iter=1000, random_state=42, class_weight="balanced"
        )
        self._scaler = StandardScaler()
        self._fitted = False
        print(f"[Detector] Initialized {self.model_name} (will fit on first use)")

    def _extract_spectral_features(self, image: Image.Image) -> np.ndarray:
        """Extract spectral features from 2D FFT magnitude spectrum."""
        # Convert to grayscale
        gray = np.array(image.convert("L"), dtype=np.float32)
        
        # Resize to fixed size for consistent FFT
        gray = cv2.resize(gray, (256, 256))
        
        # Apply window to reduce edge artifacts
        window = np.outer(np.hanning(256), np.hanning(256))
        gray = gray * window
        
        # 2D FFT
        fft = np.fft.fft2(gray)
        fft_shift = np.fft.fftshift(fft)
        magnitude = np.log1p(np.abs(fft_shift))
        
        # Extract radial profile (average magnitude at each radius)
        center = (128, 128)
        y, x = np.ogrid[:256, :256]
        r = np.sqrt((x - center[0])**2 + (y - center[1])**2).astype(int)
        
        max_radius = 128
        radial_profile = np.zeros(max_radius)
        for i in range(max_radius):
            mask = (r == i)
            if mask.any():
                radial_profile[i] = magnitude[mask].mean()
        
        # Features from radial profile
        features = []
        
        # 1. Low-frequency energy (center)
        features.append(radial_profile[:10].mean())
        
        # 2. Mid-frequency energy
        features.append(radial_profile[10:50].mean())
        
        # 3. High-frequency energy
        features.append(radial_profile[50:100].mean())
        
        # 4. Spectral slope (log-log fit)
        valid = radial_profile[1:100] > 0
        if valid.any():
            radii = np.arange(1, 100)[valid]
            mags = radial_profile[1:100][valid]
            if len(radii) > 2:
                slope = np.polyfit(np.log(radii), np.log(mags), 1)[0]
                features.append(slope)
            else:
                features.append(0.0)
        else:
            features.append(0.0)
        
        # 5. Peak-to-average ratio (detects periodic peaks)
        if radial_profile[10:100].max() > 0:
            features.append(radial_profile[10:100].max() / (radial_profile[10:100].mean() + 1e-8))
        else:
            features.append(1.0)
        
        # 6. Check for concentric rings (common in diffusion)
        ring_scores = []
        for ring_radius in [16, 32, 48, 64]:
            ring_mask = (r >= ring_radius - 2) & (r <= ring_radius + 2)
            if ring_mask.any():
                ring_scores.append(magnitude[ring_mask].mean())
        features.extend(ring_scores[:4])  # Pad to 4 values
        
        while len(features) < 10:
            features.append(0.0)
            
        return np.array(features[:10], dtype=np.float32)

    def _detect_artifacts(self, image: Image.Image, magnitude: np.ndarray) -> list:
        """Detect specific artifact types from FFT magnitude."""
        artifacts = []
        
        # Checkerboard pattern - peaks at regular intervals
        center = (128, 128)
        y, x = np.ogrid[:256, :256]
        r = np.sqrt((x - center[0])**2 + (y - center[1])**2)
        
        # Look for periodic peaks in radial profile
        radial = np.zeros(128)
        for i in range(128):
            mask = (r >= i - 0.5) & (r < i + 0.5)
            if mask.any():
                radial[i] = magnitude[mask].mean()
        
        # Find peaks
        from scipy.signal import find_peaks
        peaks, _ = find_peaks(radial[10:100], distance=8, prominence=0.3)
        if len(peaks) >= 2:
            artifacts.append({
                "type": "gan_upsampling",
                "description": "Periodic spectral peaks indicating checkerboard/upsampling artifacts",
                "severity": min(len(peaks) / 5.0, 1.0),
                "locations": [int(p + 10) for p in peaks[:5]]
            })
        
        # High-frequency grid artifacts (diffusion)
        hf_energy = radial[80:120].mean()
        if hf_energy > radial[10:80].mean() * 1.5:
            artifacts.append({
                "type": "diffusion_hf_residue",
                "description": "Elevated high-frequency energy consistent with diffusion model residue",
                "severity": min(hf_energy / (radial[10:80].mean() + 1e-8), 1.0),
                "locations": []
            })
        
        return artifacts

    def predict(self, image: Image.Image) -> DetectorResult:
        self._ensure_loaded()

        # Extract features
        gray = np.array(image.convert("L"), dtype=np.float32)
        gray = cv2.resize(gray, (256, 256))
        window = np.outer(np.hanning(256), np.hanning(256))
        gray = gray * window
        
        fft = np.fft.fft2(gray)
        fft_shift = np.fft.fftshift(fft)
        magnitude = np.log1p(np.abs(fft_shift))
        
        features = self._extract_spectral_features(image).reshape(1, -1)
        artifacts = self._detect_artifacts(image, magnitude)

        # If not fitted, use heuristic based on spectral slope and peaks
        if not self._fitted:
            # Heuristic: AI images tend to have steeper spectral decay and periodic peaks
            slope = features[0, 3]
            peak_ratio = features[0, 4]
            
            # These thresholds are approximate - would be calibrated with real data
            fake_score = 0.5
            if slope < -2.5:  # Steeper decay
                fake_score += 0.2
            if peak_ratio > 3.0:  # Strong periodic peaks
                fake_score += 0.3
            if len(artifacts) > 0:
                fake_score += 0.2 * len(artifacts)
            
            fake_prob = min(max(fake_score, 0.0), 1.0)
        else:
            # Use trained classifier
            features_scaled = self._scaler.transform(features)
            fake_prob = self._classifier.predict_proba(features_scaled)[0, 1]

        label = "FAKE" if fake_prob > 0.5 else "REAL"
        confidence = fake_prob if label == "FAKE" else 1.0 - fake_prob

        return DetectorResult(
            label=label,
            confidence=confidence,
            raw_score=fake_prob,
            model_name=self.model_name,
            metadata={
                "artifacts": artifacts,
                "spectral_slope": float(features[0, 3]),
                "peak_ratio": float(features[0, 4]),
            }
        )

    def get_gradcam_target_layer(self) -> None:
        """Frequency detector doesn't support Grad-CAM."""
        return None

    def get_artifacts(self, image: Image.Image) -> list:
        """Return detected artifacts for Stage 2."""
        gray = np.array(image.convert("L"), dtype=np.float32)
        gray = cv2.resize(gray, (256, 256))
        window = np.outer(np.hanning(256), np.hanning(256))
        gray = gray * window
        
        fft = np.fft.fft2(gray)
        fft_shift = np.fft.fftshift(fft)
        magnitude = np.log1p(np.abs(fft_shift))
        
        return self._detect_artifacts(image, magnitude)