#!/usr/bin/env python3
"""
SpectralDetector - End-to-End Image Classification
==================================================

Script completo per classificare nuove immagini usando il sistema V3.5.

Workflow:
1. Carica immagine RAW
2. Estrae features spettrali (FFT, DCT, compressione, ecc.)
3. HEAD A: REAL vs AI classification (p_real)
4. HEAD D: GPT-IMAGE-1 detection (head_d_score)
5. Meta-Router V3.5: Decisione finale con soft-veto
6. Output: REAL / AI (con generator type se disponibile)

Usage:
    python detect_image.py path/to/image.jpg
    python detect_image.py path/to/image.jpg --verbose
    python detect_image.py path/to/image.jpg --save-features features.json

Author: Denis Billi
Date: November 6, 2025
Version: 1.0 (V3.5 Production)
"""

import sys
import os
import json
import argparse
import time
from pathlib import Path
from typing import Dict, Tuple, Optional
import warnings
warnings.filterwarnings('ignore')

# Add project root and src to path
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / 'src'))

import cv2
import numpy as np
import pandas as pd

# Import joblib from archive/legacy_src if needed
try:
    import joblib
except ImportError:
    sys.path.insert(0, str(project_root / 'archive' / 'legacy_src'))
    import joblib

# Import feature extractor
try:
    # Try from archive/legacy_src (where it actually is)
    legacy_src_path = project_root / 'archive' / 'legacy_src'
    if legacy_src_path.exists():
        sys.path.insert(0, str(legacy_src_path))
    from fft_features import FFTFeatureExtractor
except ImportError as e:
    raise ImportError(f"Cannot find fft_features module: {e}. Make sure archive/legacy_src/ exists.")


class SpectralDetectorV35:
    """
    End-to-end detector for REAL vs AI image classification.
    
    Architecture:
    - HEAD A: REAL vs AI binary classification (p_real)
    - HEAD D: GPT-IMAGE-1 vs other AI generators (head_d_score)
    - Meta-Router V3.5: Logistic regression with soft-veto
    - Isotonic calibration for both HEAD A and HEAD D
    """
    
    VERSION = "3.5-PRODUCTION"
    BUILD_DATE = "2025-11-06"
    
    # Default model paths (production deployment)
    DEFAULT_MODELS = {
        'head_a': 'outputs/hierarchical_pipeline/head_a_model_20251023_091017.pkl',
        'head_d': 'outputs/hierarchical_pipeline/head_d_model_20251023_091017.pkl',  # Corrotto, ma proviamo
        'meta_router': 'outputs/hybrid_soft_veto_v35/SpectralDetector_V3_5_PROD.pkl',  # Corrotto, fallback
        'v35_alternative': 'outputs/soft_veto_robust/soft_veto_robust_C0.1_20251105_141740.pkl'
    }
    
    def __init__(
        self,
        head_a_path: Optional[str] = None,
        head_d_path: Optional[str] = None,
        meta_router_path: Optional[str] = None,
        verbose: bool = False
    ):
        """
        Initialize detector.
        
        Args:
            head_a_path: Path to HEAD A model (.pkl)
            head_d_path: Path to HEAD D model (.pkl)
            meta_router_path: Path to Meta-Router V3.5 model (.pkl)
            verbose: Enable verbose logging
        """
        self.verbose = verbose
        
        # Feature extractor configuration (aligned with training)
        self.feature_extractor = FFTFeatureExtractor(
            target_size=1024,
            grayscale=True,
            saturation_features=True,
            chroma_features=False,
            exclude_jpeg_freqs=True
        )
        
        # Load models
        self.head_a = self._load_head_a(head_a_path or self.DEFAULT_MODELS['head_a'])
        self.head_d = self._load_head_d(head_d_path or self.DEFAULT_MODELS['head_d'])
        self.meta_router = self._load_meta_router(
            meta_router_path or self.DEFAULT_MODELS['meta_router'],
            fallback=self.DEFAULT_MODELS['v35_alternative']
        )
        
        # Decision thresholds (from V3.5 production config)
        self.threshold_real_high = 0.42  # REAL decision threshold
        self.threshold_real_low = 0.28   # Gray zone lower bound
        self.threshold_ai = 0.5          # P(AI) decision threshold (soft-veto)
        
        # Backward compatibility aliases (for old meta_router pickle files)
        self.tau_real_high = self.threshold_real_high
        self.tau_real_low = self.threshold_real_low
        self.tau_ai = self.threshold_ai
        self.p_ai_threshold = self.threshold_ai  # Alternative name for threshold_ai
        
        if self.verbose:
            self._print_config()
    
    def _log(self, message: str):
        """Print if verbose enabled"""
        if self.verbose:
            # Fix Windows encoding issues with emoji
            try:
                print(message)
            except UnicodeEncodeError:
                # Remove emoji and retry
                message_clean = message.encode('ascii', 'ignore').decode('ascii')
                print(message_clean)
    
    def _load_head_a(self, model_path: str) -> Optional[Dict]:
        """Load HEAD A model"""
        if not os.path.exists(model_path):
            self._log(f"  HEAD A not found: {model_path}")
            return None
        
        try:
            # Load model with joblib
            model = joblib.load(model_path)
            
            # Verify model size (corrotti se < 10KB)
            file_size = os.path.getsize(model_path)
            if file_size < 10000:
                self._log(f"  HEAD A model corrupted ({file_size} bytes)")
                return None
            
            # Load feature names from separate file (if exists)
            model_dir = Path(model_path).parent
            feature_names_path = model_dir / 'feature_names.txt'
            feature_names = None
            
            if feature_names_path.exists():
                with open(feature_names_path, 'r') as f:
                    feature_names = [line.strip() for line in f if line.strip()]
                self._log(f"  Loaded {len(feature_names)} feature names")
            
            # Load scaler if exists
            scaler_path = model_dir / 'head_a_scaler.pkl'
            scaler = None
            if scaler_path.exists():
                scaler = joblib.load(scaler_path)
                self._log(f"  Loaded scaler")
            
            self._log(f" HEAD A loaded: {model_path} ({file_size:,} bytes)")
            
            return {
                'model': model,
                'scaler': scaler,
                'feature_names': feature_names
            }
        
        except Exception as e:
            self._log(f" HEAD A loading failed: {e}")
            return None
    
    def _load_head_d(self, model_path: str) -> Optional[Dict]:
        """Load HEAD D model"""
        if not os.path.exists(model_path):
            self._log(f"  HEAD D not found: {model_path}")
            return None
        
        try:
            # Load model with joblib
            model = joblib.load(model_path)
            
            file_size = os.path.getsize(model_path)
            if file_size < 10000:
                self._log(f"  HEAD D model corrupted ({file_size} bytes)")
                return None
            
            # Load feature names from separate file (if exists)
            model_dir = Path(model_path).parent
            feature_names_path = model_dir / 'feature_names.txt'
            feature_names = None
            
            if feature_names_path.exists():
                with open(feature_names_path, 'r') as f:
                    feature_names = [line.strip() for line in f if line.strip()]
            
            # Load scaler if exists
            scaler_path = model_dir / 'head_d_scaler.pkl'
            scaler = None
            if scaler_path.exists():
                scaler = joblib.load(scaler_path)
            
            self._log(f" HEAD D loaded: {model_path} ({file_size:,} bytes)")
            
            return {
                'model': model,
                'scaler': scaler,
                'feature_names': feature_names
            }
        
        except Exception as e:
            self._log(f" HEAD D loading failed: {e}")
            return None
    
    def _load_meta_router(self, model_path: str, fallback: str) -> Optional[Dict]:
        """Load Meta-Router V3.5"""
        # Try primary path
        if os.path.exists(model_path):
            try:
                with open(model_path, 'rb') as f:
                    meta = joblib.load(f)
                
                file_size = os.path.getsize(model_path)
                if file_size < 1000:  # Changed from 10000 to 1000 (meta-router is small)
                    self._log(f"⚠️  Meta-Router corrupted ({file_size} bytes), trying fallback...")
                else:
                    self._log(f"✅ Meta-Router loaded: {model_path}")
                    return meta
            except Exception as e:
                self._log(f"⚠️  Meta-Router loading failed: {e}, trying fallback...")
        
        # Try fallback
        if os.path.exists(fallback):
            try:
                with open(fallback, 'rb') as f:
                    meta = joblib.load(f)
                
                self._log(f"✅ Meta-Router loaded (fallback): {fallback}")
                return meta
            except Exception as e:
                self._log(f"❌ Fallback Meta-Router failed: {e}")
        
        self._log("❌ No Meta-Router available - using rule-based fallback")
        return None
    
    def _print_config(self):
        """Print detector configuration"""
        print(f"\n{'='*60}")
        print(f"SpectralDetector {self.VERSION}")
        print(f"Build: {self.BUILD_DATE}")
        print(f"{'='*60}")
        print(f"\nModels Loaded:")
        print(f"  HEAD A:       {'OK' if self.head_a else 'MISSING'}")
        print(f"  HEAD D:       {'OK' if self.head_d else 'MISSING'}")
        print(f"  Meta-Router:  {'OK' if self.meta_router else 'MISSING (using rule-based)'}")
        print(f"\nThresholds:")
        print(f"  tau_real_high: {self.threshold_real_high}")
        print(f"  tau_real_low:  {self.threshold_real_low}")
        print(f"  P(AI) thresh:  {self.threshold_ai}")
        print(f"{'='*60}\n")
    
    def extract_features(self, image_path: str) -> Optional[Dict]:
        """
        Extract spectral features from image.
        
        Args:
            image_path: Path to image file
            
        Returns:
            Dictionary with extracted features, or None if failed
        """
        if not os.path.exists(image_path):
            self._log(f"❌ Image not found: {image_path}")
            return None
        
        try:
            # Load image
            image = cv2.imread(image_path)
            if image is None:
                self._log(f"❌ Could not load image: {image_path}")
                return None
            
            self._log(f"📸 Image loaded: {image.shape} - {Path(image_path).name}")
            
            # Extract features
            start_time = time.time()
            features = self.feature_extractor.extract_features(image)
            extraction_time = time.time() - start_time
            
            if features is None:
                self._log("❌ Feature extraction failed")
                return None
            
            self._log(f"✅ Features extracted: {len(features)} features in {extraction_time:.2f}s")
            return features
        
        except Exception as e:
            self._log(f"❌ Feature extraction error: {e}")
            return None
    
    def predict_head_a(self, features: Dict) -> Tuple[float, str]:
        """
        Predict REAL vs AI using HEAD A.
        
        Args:
            features: Extracted features dict
            
        Returns:
            (p_real, status_message)
        """
        if self.head_a is None:
            self._log("  HEAD A not available - using default p_real=0.5")
            return 0.5, "HEAD A not available"
        
        try:
            # Get model, scaler, and feature names
            model = self.head_a['model']
            scaler = self.head_a.get('scaler', None)
            feature_names = self.head_a.get('feature_names', None)
            
            if feature_names is None:
                self._log("  No feature names - using all features")
                # Fallback: use all features
                feature_values = []
                for key, val in features.items():
                    if isinstance(val, (dict, list)):
                        val = next(iter(val.values())) if isinstance(val, dict) else (val[0] if val else 0.0)
                    feature_values.append(float(val))
            else:
                # Align features with expected feature names
                feature_values = []
                for fname in feature_names:
                    if fname in features:
                        val = features[fname]
                        if isinstance(val, (dict, list)):
                            val = next(iter(val.values())) if isinstance(val, dict) else (val[0] if val else 0.0)
                        feature_values.append(float(val))
                    else:
                        feature_values.append(0.0)  # Missing feature
            
            feature_array = np.array(feature_values).reshape(1, -1)
            
            # Apply scaler if available
            if scaler:
                feature_array = scaler.transform(feature_array)
            
            # Predict (class 1 = REAL, class 0 = AI)
            p_real = model.predict_proba(feature_array)[0][1]
            
            return float(p_real), "OK"
        
        except Exception as e:
            self._log(f"  HEAD A prediction error: {e}")
            import traceback
            self._log(traceback.format_exc())
            return 0.5, f"Error: {e}"
    
    def predict_head_d(self, features: Dict) -> Tuple[float, str]:
        """
        Predict GPT-IMAGE-1 score using HEAD D.
        
        Args:
            features: Extracted features dict
            
        Returns:
            (head_d_score, status_message)
        """
        if self.head_d is None:
            self._log("  HEAD D not available - using default score=0.0")
            return 0.0, "HEAD D not available"
        
        try:
            # Get model, scaler, and feature names
            model = self.head_d['model']
            scaler = self.head_d.get('scaler', None)
            feature_names = self.head_d.get('feature_names', None)
            
            if feature_names is None:
                return 0.0, "No feature names available"
            
            # Align features with expected feature names
            feature_values = []
            for fname in feature_names:
                if fname in features:
                    val = features[fname]
                    if isinstance(val, (dict, list)):
                        val = next(iter(val.values())) if isinstance(val, dict) else (val[0] if val else 0.0)
                    feature_values.append(float(val))
                else:
                    feature_values.append(0.0)
            
            feature_array = np.array(feature_values).reshape(1, -1)
            
            # Apply scaler if available
            if scaler:
                feature_array = scaler.transform(feature_array)
            
            # Predict (class 1 = GPT-IMAGE-1, class 0 = Other AI)
            head_d_score = model.predict_proba(feature_array)[0][1]
            
            return float(head_d_score), "OK"
        
        except Exception as e:
            self._log(f"  HEAD D prediction error: {e}")
            return 0.0, f"Error: {e}"
    
    def soft_veto_decision(self, p_real: float, head_d_score: float) -> Tuple[str, float, str]:
        """
        Make final decision using soft-veto meta-router.
        
        Args:
            p_real: HEAD A probability of REAL
            head_d_score: HEAD D GPT-IMAGE-1 score
            
        Returns:
            (decision, p_ai_final, zone)
            decision: 'REAL' or 'AI'
            p_ai_final: Final P(AI) after soft-veto
            zone: 'SAFE_REAL', 'SAFE_AI', or 'GRAY'
        """
        # Compute compression awareness (simplified - assume JPEG compressed)
        f_compress = 1.0  # Default: assume compressed
        
        # Compute margin (distance from threshold)
        w_margin = abs(p_real - self.threshold_real_high)
        
        # Build meta-features
        meta_features = np.array([
            [
                -np.log(p_real + 1e-10),      # neg_logit_real
                np.log(1 - p_real + 1e-10),   # logit_ai
                f_compress,                    # compression awareness
                w_margin                       # margin from threshold
            ]
        ])
        
        if self.meta_router:
            try:
                # Use meta-router
                if isinstance(self.meta_router, dict):
                    model = self.meta_router.get('model', None)
                    if model:
                        p_ai_final = model.predict_proba(meta_features)[0][1]
                    else:
                        p_ai_final = 1 - p_real  # Fallback
                else:
                    p_ai_final = self.meta_router.predict_proba(meta_features)[0][1]
            except Exception as e:
                self._log(f"⚠️  Meta-router error: {e}, using fallback")
                p_ai_final = 1 - p_real
        else:
            # Rule-based fallback (Option 4 logic)
            if head_d_score > 0.58 and abs(p_real - self.threshold_real_high) <= 0.05:
                # Margin-gate veto activated
                p_ai_final = 1.0  # Force AI
            else:
                p_ai_final = 1 - p_real
        
        # Decision
        if p_ai_final >= self.threshold_ai:
            decision = "AI"
        else:
            decision = "REAL"
        
        # Zone classification
        if p_real >= self.threshold_real_high:
            zone = "SAFE_REAL"
        elif p_real <= self.threshold_real_low:
            zone = "SAFE_AI"
        else:
            zone = "GRAY"
        
        return decision, float(p_ai_final), zone
    
    def classify(self, image_path: str, save_features: Optional[str] = None) -> Dict:
        """
        Complete classification pipeline.
        
        Args:
            image_path: Path to image file
            save_features: Optional path to save extracted features JSON
            
        Returns:
            Dictionary with classification results
        """
        start_time = time.time()
        
        # Step 1: Extract features
        self._log(f"\n{'='*60}")
        self._log(f"CLASSIFYING: {Path(image_path).name}")
        self._log(f"{'='*60}")
        
        features = self.extract_features(image_path)
        if features is None:
            return {
                'error': 'Feature extraction failed',
                'image_path': image_path
            }
        
        # Save features if requested
        if save_features:
            with open(save_features, 'w') as f:
                # Convert numpy/complex types to JSON-serializable
                features_clean = {}
                for k, v in features.items():
                    if isinstance(v, (np.ndarray, list)):
                        features_clean[k] = [float(x) if isinstance(x, (np.floating, np.integer)) else x for x in v]
                    elif isinstance(v, dict):
                        features_clean[k] = {kk: float(vv) if isinstance(vv, (np.floating, np.integer)) else vv for kk, vv in v.items()}
                    elif isinstance(v, (np.floating, np.integer)):
                        features_clean[k] = float(v)
                    else:
                        features_clean[k] = v
                
                json.dump(features_clean, f, indent=2)
            self._log(f"💾 Features saved: {save_features}")
        
        # Step 2: HEAD A prediction
        self._log("\n[HEAD A] REAL vs AI Classification...")
        p_real, head_a_status = self.predict_head_a(features)
        self._log(f"  P(REAL) = {p_real:.4f} ({head_a_status})")
        
        # Step 3: HEAD D prediction
        self._log("\n[HEAD D] GPT-IMAGE-1 Detection...")
        head_d_score, head_d_status = self.predict_head_d(features)
        self._log(f"  HEAD D Score = {head_d_score:.4f} ({head_d_status})")
        
        # Step 4: Meta-Router decision
        self._log("\n[META-ROUTER V3.5] Final Decision...")
        decision, p_ai_final, zone = self.soft_veto_decision(p_real, head_d_score)
        
        total_time = time.time() - start_time
        
        # Build result
        result = {
            'image_path': image_path,
            'decision': decision,
            'p_real': round(p_real, 4),
            'p_ai_final': round(p_ai_final, 4),
            'head_d_score': round(head_d_score, 4),
            'zone': zone,
            'processing_time': round(total_time, 3),
            'head_a_status': head_a_status,
            'head_d_status': head_d_status,
            'version': self.VERSION
        }
        
        # Generator identification (improved heuristic)
        if decision == "AI":
            # GPT-IMAGE-1 signature detection (primary)
            if head_d_score > 0.5:  # Lowered from 0.7 for better detection
                result['likely_generator'] = "GPT-IMAGE-1"
            # High confidence AI with typical ChatGPT patterns
            elif p_ai_final > 0.85:  # Lowered from 0.95
                result['likely_generator'] = "ChatGPT/DALL-E-3"
            # Medium-high AI probability - likely SDXL/Gemini
            elif p_ai_final > 0.65:
                result['likely_generator'] = "SDXL/Gemini/MidJourney"
            else:
                result['likely_generator'] = "AI (Unknown)"
        else:
            result['likely_generator'] = "REAL"
        
        self._log(f"\n{'='*60}")
        self._log(f"FINAL DECISION: {decision} ({zone})")
        if decision == "AI":
            self._log(f"Likely Generator: {result['likely_generator']}")
        self._log(f"P(AI) Final: {p_ai_final:.4f}")
        self._log(f"Processing Time: {total_time:.2f}s")
        self._log(f"{'='*60}\n")
        
        return result


def main():
    """CLI interface"""
    parser = argparse.ArgumentParser(
        description="SpectralDetector V3.5 - REAL vs AI Image Classification",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python detect_image.py photo.jpg
  python detect_image.py photo.jpg --verbose
  python detect_image.py photo.jpg --save-features features.json
  python detect_image.py photo.jpg --save-result result.json
        """
    )
    
    parser.add_argument('image', help="Path to image file")
    parser.add_argument('--verbose', '-v', action='store_true', help="Enable verbose output")
    parser.add_argument('--save-features', help="Save extracted features to JSON file")
    parser.add_argument('--save-result', help="Save classification result to JSON file")
    parser.add_argument('--head-a', help="Path to HEAD A model (optional)")
    parser.add_argument('--head-d', help="Path to HEAD D model (optional)")
    parser.add_argument('--meta-router', help="Path to Meta-Router model (optional)")
    
    args = parser.parse_args()
    
    # Initialize detector
    detector = SpectralDetectorV35(
        head_a_path=args.head_a,
        head_d_path=args.head_d,
        meta_router_path=args.meta_router,
        verbose=args.verbose
    )
    
    # Classify image
    result = detector.classify(
        args.image,
        save_features=args.save_features
    )
    
    # Check for errors
    if 'error' in result:
        print(f"\n❌ ERROR: {result['error']}")
        return 1
    
    # Print result
    if not args.verbose:
        print(f"\n{'='*60}")
        print(f"Image: {Path(args.image).name}")
        print(f"{'='*60}")
        print(f"Decision:         {result['decision']}")
        print(f"Likely Generator: {result['likely_generator']}")
        print(f"P(AI) Final:      {result['p_ai_final']:.4f}")
        print(f"P(REAL):          {result['p_real']:.4f}")
        print(f"HEAD D Score:     {result['head_d_score']:.4f}")
        print(f"Zone:             {result['zone']}")
        print(f"Processing Time:  {result['processing_time']:.2f}s")
        print(f"{'='*60}\n")
    
    # Save result if requested
    if args.save_result:
        with open(args.save_result, 'w') as f:
            json.dump(result, f, indent=2)
        print(f"💾 Result saved: {args.save_result}\n")
    
    return 0


class SpectralDetectorV36(SpectralDetectorV35):
    """
    Enhanced detector with multi-generator identification.
    
    Architecture V3.6:
    - HEAD A: REAL vs AI binary classification (inherited from V3.5)
    - HEAD B: SDXL vs NON-SDXL detection (NEW)
    - HEAD C: ChatGPT vs Gemini classification (NEW)
    - HEAD D: GPT-IMAGE-1 veto (inherited from V3.5)
    - Meta-Router V3.6: 4-HEAD hierarchical integration
    
    Improvements over V3.5:
    - Accurate per-generator identification (SDXL, ChatGPT, Gemini, GPT-IMAGE-1)
    - Reduced "AI (Unknown)" rate from 85% to <20%
    - Maintains 92%+ REAL vs AI accuracy
    """
    
    VERSION = "3.6-PRODUCTION"
    BUILD_DATE = "2025-11-06"
    
    # Default model paths for V3.6
    DEFAULT_MODELS_V36 = {
        'head_a': 'outputs/hierarchical_pipeline/head_a_model_20251023_091017.pkl',
        'head_b': 'outputs/heads_v36/head_b_model.pkl',
        'head_c': 'outputs/heads_v36/head_c_model.pkl',
        'head_d': 'outputs/hierarchical_pipeline/head_d_model_20251023_091017.pkl',
        'meta_router': 'outputs/meta_router_v36/meta_router_v36.pkl',
    }
    
    def __init__(
        self,
        head_a_path: Optional[str] = None,
        head_b_path: Optional[str] = None,
        head_c_path: Optional[str] = None,
        head_d_path: Optional[str] = None,
        meta_router_path: Optional[str] = None,
        verbose: bool = False
    ):
        """
        Initialize V3.6 detector.
        
        Args:
            head_a_path: Path to HEAD A model (.pkl)
            head_b_path: Path to HEAD B model (.pkl)
            head_c_path: Path to HEAD C model (.pkl)
            head_d_path: Path to HEAD D model (.pkl)
            meta_router_path: Path to Meta-Router V3.6 model (.pkl)
            verbose: Enable verbose logging
        """
        # Initialize parent V3.5 (loads HEAD A, HEAD D, feature extractor)
        super().__init__(
            head_a_path=head_a_path or self.DEFAULT_MODELS_V36['head_a'],
            head_d_path=head_d_path or self.DEFAULT_MODELS_V36['head_d'],
            meta_router_path=None,  # We'll load V3.6 meta-router separately
            verbose=verbose
        )
        
        # Load HEAD B (SDXL detection)
        self.head_b = self._load_head_b(head_b_path or self.DEFAULT_MODELS_V36['head_b'])
        
        # Load HEAD C (ChatGPT vs Gemini)
        self.head_c = self._load_head_c(head_c_path or self.DEFAULT_MODELS_V36['head_c'])
        
        # Load Meta-Router V3.6
        self.meta_router_v36 = self._load_meta_router_v36(
            meta_router_path or self.DEFAULT_MODELS_V36['meta_router']
        )
        
        if self.verbose:
            self._log(f"\n{'='*60}")
            self._log(f"SpectralDetector V3.6 initialized")
            self._log(f"{'='*60}")
    
    def _load_head_b(self, model_path: str) -> Optional[Dict]:
        """Load HEAD B (SDXL detection) model"""
        if not os.path.exists(model_path):
            self._log(f"⚠️  HEAD B not found: {model_path}")
            return None
        
        try:
            model = joblib.load(model_path)
            file_size = os.path.getsize(model_path)
            
            # Load scaler
            model_dir = Path(model_path).parent
            scaler_path = model_dir / 'head_b_scaler.pkl'
            scaler = None
            if scaler_path.exists():
                scaler = joblib.load(scaler_path)
            
            # Load feature names
            feature_names_path = model_dir / 'head_b_feature_names.txt'
            feature_names = None
            if feature_names_path.exists():
                with open(feature_names_path, 'r') as f:
                    feature_names = [line.strip() for line in f if line.strip()]
            
            self._log(f"✅ HEAD B loaded: {model_path} ({file_size:,} bytes)")
            
            return {
                'model': model,
                'scaler': scaler,
                'feature_names': feature_names
            }
        
        except Exception as e:
            self._log(f"❌ HEAD B loading failed: {e}")
            return None
    
    def _load_head_c(self, model_path: str) -> Optional[Dict]:
        """Load HEAD C (ChatGPT vs Gemini) model"""
        if not os.path.exists(model_path):
            self._log(f"⚠️  HEAD C not found: {model_path}")
            return None
        
        try:
            model = joblib.load(model_path)
            file_size = os.path.getsize(model_path)
            
            # Load scaler
            model_dir = Path(model_path).parent
            scaler_path = model_dir / 'head_c_scaler.pkl'
            scaler = None
            if scaler_path.exists():
                scaler = joblib.load(scaler_path)
            
            # Load feature names
            feature_names_path = model_dir / 'head_c_feature_names.txt'
            feature_names = None
            if feature_names_path.exists():
                with open(feature_names_path, 'r') as f:
                    feature_names = [line.strip() for line in f if line.strip()]
            
            # Load class mapping
            class_mapping_path = model_dir / 'head_c_class_mapping.json'
            class_mapping = None
            if class_mapping_path.exists():
                with open(class_mapping_path, 'r') as f:
                    class_mapping = json.load(f)
            
            self._log(f"✅ HEAD C loaded: {model_path} ({file_size:,} bytes)")
            if class_mapping:
                self._log(f"   Class mapping: {class_mapping}")
            
            return {
                'model': model,
                'scaler': scaler,
                'feature_names': feature_names,
                'class_mapping': class_mapping
            }
        
        except Exception as e:
            self._log(f"❌ HEAD C loading failed: {e}")
            return None
    
    def _load_meta_router_v36(self, model_path: str) -> Optional[object]:
        """Load Meta-Router V3.6"""
        if not os.path.exists(model_path):
            self._log(f"⚠️  Meta-Router V3.6 not found: {model_path}")
            return None
        
        try:
            router = joblib.load(model_path)
            file_size = os.path.getsize(model_path)
            
            self._log(f"✅ Meta-Router V3.6 loaded: {model_path} ({file_size:,} bytes)")
            
            return router
        
        except Exception as e:
            self._log(f"❌ Meta-Router V3.6 loading failed: {e}")
            return None
    
    def classify(self, image_path: str) -> Dict:
        """
        Classify image with V3.6 multi-HEAD architecture.
        
        Returns:
        --------
        dict with keys:
            - final_label: str (REAL, SDXL, ChatGPT, Gemini, GPT-IMAGE-1, AI (Unknown))
            - confidence: float
            - confidence_zone: str (HIGH, MEDIUM, LOW)
            - p_real: float
            - p_ai: float
            - p_sdxl: float (NEW in V3.6)
            - p_chatgpt: float (NEW in V3.6)
            - p_gemini: float (NEW in V3.6)
            - p_gpt_image_1: float
            - processing_time_ms: float
            - version: str
        """
        start_time = time.time()
        
        # Extract features
        features = self.extract_features(image_path)
        
        if features is None or 'error' in features:
            return {
                'error': 'Feature extraction failed',
                'details': features.get('error') if features else 'Unknown error',
                'version': self.VERSION
            }
        
        # Prepare feature vector (exclude metadata)
        exclude_cols = ['image_path', 'filepath', 'source', 'label', 'model_version']
        feature_names = [k for k in features.keys() if k not in exclude_cols]
        X = np.array([features[k] for k in feature_names]).reshape(1, -1)
        
        # HEAD A: REAL vs AI
        if self.head_a is None:
            return {'error': 'HEAD A not loaded', 'version': self.VERSION}
        
        if self.head_a['scaler']:
            X_a = self.head_a['scaler'].transform(X)
        else:
            X_a = X
        
        p_ai = self.head_a['model'].predict_proba(X_a)[0, 1]
        p_real = 1 - p_ai
        
        # HEAD D: GPT-IMAGE-1 detection
        p_gpt_image_1 = 0.0
        if self.head_d is not None:
            if self.head_d['scaler']:
                X_d = self.head_d['scaler'].transform(X)
            else:
                X_d = X
            
            p_gpt_image_1 = self.head_d['model'].predict_proba(X_d)[0, 1]
        
        # HEAD B: SDXL detection
        p_sdxl = 0.0
        if self.head_b is not None:
            if self.head_b['scaler']:
                X_b = self.head_b['scaler'].transform(X)
            else:
                X_b = X
            
            p_sdxl = self.head_b['model'].predict_proba(X_b)[0, 1]
        
        # HEAD C: ChatGPT vs Gemini
        p_chatgpt = 0.0
        p_gemini = 0.0
        if self.head_c is not None:
            if self.head_c['scaler']:
                X_c = self.head_c['scaler'].transform(X)
            else:
                X_c = X
            
            proba = self.head_c['model'].predict_proba(X_c)[0]
            
            # Map to class names (class_mapping: {'chatgpt': 0, 'gemini': 1})
            if self.head_c['class_mapping']:
                chatgpt_idx = self.head_c['class_mapping'].get('chatgpt', 0)
                gemini_idx = self.head_c['class_mapping'].get('gemini', 1)
                p_chatgpt = proba[chatgpt_idx]
                p_gemini = proba[gemini_idx]
            else:
                # Fallback: assume index 0=chatgpt, 1=gemini
                p_chatgpt = proba[0]
                p_gemini = proba[1]
        
        # Meta-Router V3.6: Hierarchical decision
        if self.meta_router_v36 is not None:
            result = self.meta_router_v36.predict(
                head_a_proba=p_ai,
                head_b_proba=p_sdxl,
                head_c_proba=[p_chatgpt, p_gemini],
                head_d_proba=p_gpt_image_1
            )
        else:
            # Fallback to simple heuristic if meta-router not loaded
            result = self._fallback_decision(p_real, p_ai, p_sdxl, p_chatgpt, p_gemini, p_gpt_image_1)
        
        # Add timing and version
        processing_time = (time.time() - start_time) * 1000
        result['processing_time_ms'] = processing_time
        result['version'] = self.VERSION
        
        return result
    
    def _fallback_decision(self, p_real, p_ai, p_sdxl, p_chatgpt, p_gemini, p_gpt_image_1):
        """Fallback decision logic if meta-router not loaded"""
        
        # STEP 1: REAL vs AI
        if p_real > 0.5:
            confidence_zone = 'HIGH' if p_real >= 0.7 else 'MEDIUM' if p_real >= 0.3 else 'LOW'
            return {
                'final_label': 'REAL',
                'confidence': p_real,
                'confidence_zone': confidence_zone,
                'p_real': p_real,
                'p_ai': p_ai,
                'p_sdxl': p_sdxl,
                'p_chatgpt': p_chatgpt,
                'p_gemini': p_gemini,
                'p_gpt_image_1': p_gpt_image_1
            }
        
        # STEP 2: AI detected - cascade
        
        # 2a. GPT-IMAGE-1 veto
        if p_gpt_image_1 >= 0.7:
            confidence_zone = 'HIGH' if p_gpt_image_1 >= 0.85 else 'MEDIUM'
            return {
                'final_label': 'GPT-IMAGE-1',
                'confidence': p_gpt_image_1,
                'confidence_zone': confidence_zone,
                'p_real': p_real,
                'p_ai': p_ai,
                'p_sdxl': p_sdxl,
                'p_chatgpt': p_chatgpt,
                'p_gemini': p_gemini,
                'p_gpt_image_1': p_gpt_image_1
            }
        
        # 2b. SDXL detection
        if p_sdxl >= 0.8:
            confidence_zone = 'HIGH' if p_sdxl >= 0.9 else 'MEDIUM'
            return {
                'final_label': 'SDXL',
                'confidence': p_sdxl,
                'confidence_zone': confidence_zone,
                'p_real': p_real,
                'p_ai': p_ai,
                'p_sdxl': p_sdxl,
                'p_chatgpt': p_chatgpt,
                'p_gemini': p_gemini,
                'p_gpt_image_1': p_gpt_image_1
            }
        
        # 2c. ChatGPT vs Gemini
        if p_chatgpt > p_gemini:
            if p_chatgpt >= 0.6:
                confidence_zone = 'HIGH' if p_chatgpt >= 0.8 else 'MEDIUM'
                return {
                    'final_label': 'ChatGPT',
                    'confidence': p_chatgpt,
                    'confidence_zone': confidence_zone,
                    'p_real': p_real,
                    'p_ai': p_ai,
                    'p_sdxl': p_sdxl,
                    'p_chatgpt': p_chatgpt,
                    'p_gemini': p_gemini,
                    'p_gpt_image_1': p_gpt_image_1
                }
        else:
            if p_gemini >= 0.6:
                confidence_zone = 'HIGH' if p_gemini >= 0.8 else 'MEDIUM'
                return {
                    'final_label': 'Gemini',
                    'confidence': p_gemini,
                    'confidence_zone': confidence_zone,
                    'p_real': p_real,
                    'p_ai': p_ai,
                    'p_sdxl': p_sdxl,
                    'p_chatgpt': p_chatgpt,
                    'p_gemini': p_gemini,
                    'p_gpt_image_1': p_gpt_image_1
                }
        
        # 2d. Fallback: AI (Unknown)
        max_ai_confidence = max(p_ai, p_sdxl, p_chatgpt, p_gemini, p_gpt_image_1)
        
        return {
            'final_label': 'AI (Unknown)',
            'confidence': max_ai_confidence,
            'confidence_zone': 'LOW',
            'p_real': p_real,
            'p_ai': p_ai,
            'p_sdxl': p_sdxl,
            'p_chatgpt': p_chatgpt,
            'p_gemini': p_gemini,
            'p_gpt_image_1': p_gpt_image_1
        }


class SpectralDetectorV36Cascaded:
    """
    V3.6 Cascaded Router with 4-HEAD system (A, B, C, D).
    
    Architecture:
    - HEAD A: REAL vs AI (gate)
    - HEAD B: SDXL detection
    - HEAD C: ChatGPT vs Gemini
    - HEAD D: GPT-IMAGE-1 detection
    - Meta-Router V3.6: Cascaded decision with fallback
    
    Feature extraction:
    - HEAD A/B/C: 67 features
    - HEAD D: 90 features (includes extra global/patch features)
    """
    
    VERSION = "3.6-CASCADED"
    BUILD_DATE = "2025-11-07"
    
    DEFAULT_MODELS = {
        'head_a': 'outputs/heads_v36_67features/head_a_model.pkl',
        'head_b': 'outputs/heads_v36_67features/head_b_model.pkl',
        'head_c': 'outputs/heads_v36_67features/head_c_model.pkl',
        'head_d': 'outputs/heads_v36_67features/head_d_model.pkl',
        'calibrators': 'outputs/heads_v36_67features/'
    }
    
    def __init__(
        self,
        models_dir: Optional[str] = None,
        verbose: bool = False
    ):
        """Initialize V3.6 detector with all 4 HEADs + calibrators."""
        self.verbose = verbose
        self.models_dir = models_dir or self.DEFAULT_MODELS['calibrators']
        
        # Feature extractor (same config as training)
        self.feature_extractor = FFTFeatureExtractor(
            target_size=1024,
            grayscale=True,
            saturation_features=True,
            chroma_features=False,
            exclude_jpeg_freqs=True
        )
        
        # Load models and feature names
        self._load_models()
        
        # Initialize meta-router
        from meta_router_v36_cascaded import create_router_option3
        self.router = create_router_option3()
        
        if self.verbose:
            self._log(f"✅ SpectralDetector V{self.VERSION} initialized")
    
    def _log(self, message: str):
        """Print if verbose"""
        if self.verbose:
            try:
                print(message)
            except UnicodeEncodeError:
                print(message.encode('ascii', 'ignore').decode('ascii'))
    
    def _load_models(self):
        """Load all 4 HEADs + calibrators + feature names"""
        import joblib
        
        models_path = Path(self.models_dir)
        
        # Load HEAD A/B/C (67 features)
        self.head_a = joblib.load(models_path / 'head_a_model.pkl')
        self.head_b = joblib.load(models_path / 'head_b_model.pkl')
        self.head_c = joblib.load(models_path / 'head_c_model.pkl')
        self.head_d = joblib.load(models_path / 'head_d_model.pkl')
        
        # Load scalers
        self.scaler_a = joblib.load(models_path / 'head_a_scaler.pkl')
        self.scaler_b = joblib.load(models_path / 'head_b_scaler.pkl')
        self.scaler_c = joblib.load(models_path / 'head_c_scaler.pkl')
        self.scaler_d = joblib.load(models_path / 'head_d_scaler.pkl')
        
        # Load feature names
        with open(models_path / 'head_a_features.txt', 'r') as f:
            self.features_67 = [line.strip() for line in f]
        
        with open(models_path / 'head_d_features.txt', 'r') as f:
            self.features_90 = [line.strip() for line in f]
        
        # Load calibrators (temperature scalers)
        try:
            self.calibrator_a = joblib.load(models_path / 'head_a_temperature_scaler.pkl')
            self.calibrator_b = joblib.load(models_path / 'head_b_temperature_scaler.pkl')
            self.calibrator_c = joblib.load(models_path / 'head_c_temperature_scaler.pkl')
            self._log("  Loaded temperature calibrators")
        except FileNotFoundError:
            self._log("  ⚠️ Temperature calibrators not found (using uncalibrated)")
            self.calibrator_a = self.calibrator_b = self.calibrator_c = None
        
        self._log(f"  Loaded 4 HEADs: {len(self.features_67)} features (A/B/C), {len(self.features_90)} features (D)")
    
    def classify_image(self, image_path: str) -> Dict:
        """
        Classify a single image.
        
        Args:
            image_path: Path to image file
            
        Returns:
            Dict with decision, confidence, route, probabilities
        """
        # Load image
        img = cv2.imread(str(image_path))
        if img is None:
            raise ValueError(f"Cannot load image: {image_path}")
        
        # Extract features (full set for 90-feature extraction)
        features_dict = self.feature_extractor.extract_features(img)
        
        # Convert to DataFrame
        features_df = pd.DataFrame([features_dict])
        
        # Extract 67 features for HEAD A/B/C
        X_67 = features_df[self.features_67].values
        X_67_scaled = self.scaler_a.transform(X_67)  # Scaler A (same for B/C)
        
        # Extract 90 features for HEAD D
        X_90 = features_df[self.features_90].values
        X_90_scaled = self.scaler_d.transform(X_90)
        
        # HEAD A: REAL vs AI
        proba_a = self.head_a.predict_proba(X_67_scaled)[0]
        p_real_raw = proba_a[0]  # Class 0 = REAL
        p_ai_raw = proba_a[1]    # Class 1 = AI
        
        # Calibrate HEAD A
        if self.calibrator_a:
            logits_a = np.log(proba_a / (1 - proba_a + 1e-10))
            proba_a_cal = self.calibrator_a.transform(logits_a.reshape(1, -1))[0]
            p_real = proba_a_cal[0]
            p_ai = proba_a_cal[1]
        else:
            p_real = p_real_raw
            p_ai = p_ai_raw
        
        # HEAD B: SDXL detection
        proba_b = self.head_b.predict_proba(X_67_scaled)[0]
        p_sdxl_raw = proba_b[1]  # Assuming class 1 = SDXL
        
        if self.calibrator_b:
            logits_b = np.log(proba_b / (1 - proba_b + 1e-10))
            proba_b_cal = self.calibrator_b.transform(logits_b.reshape(1, -1))[0]
            p_sdxl = proba_b_cal[1]
        else:
            p_sdxl = p_sdxl_raw
        
        # HEAD C: ChatGPT vs Gemini
        proba_c = self.head_c.predict_proba(X_67_scaled)[0]
        p_chatgpt_raw = proba_c[0]  # Assuming class 0 = ChatGPT
        p_gemini_raw = proba_c[1]
        
        if self.calibrator_c:
            logits_c = np.log(proba_c / (1 - proba_c + 1e-10))
            proba_c_cal = self.calibrator_c.transform(logits_c.reshape(1, -1))[0]
            p_chatgpt = proba_c_cal[0]
            p_gemini = proba_c_cal[1]
        else:
            p_chatgpt = p_chatgpt_raw
            p_gemini = p_gemini_raw
        
        # HEAD D: GPT-IMAGE-1 detection
        proba_d = self.head_d.predict_proba(X_90_scaled)[0]
        p_gpt_image_1 = proba_d[1]  # Assuming class 1 = GPT-IMAGE-1
        
        # Meta-Router V3.6 decision
        decision = self.router.predict(
            p_real=p_real,
            p_ai=p_ai,
            p_gpt_image_1=p_gpt_image_1,
            p_sdxl=p_sdxl,
            p_chatgpt=p_chatgpt,
            p_gemini=p_gemini
        )
        
        return {
            'final_label': decision.final_label,
            'confidence': decision.confidence,
            'route_taken': decision.route_taken,
            'exit_head': decision.exit_head,
            'threshold_hit': decision.threshold_hit,
            'probabilities': decision.probabilities,
            'scores': decision.scores
        }


if __name__ == '__main__':
    sys.exit(main())