"""
CounterfactualProbe - Find minimal changes that flip detections.

Uses PGD (Projected Gradient Descent) to find the smallest perturbation
that removes, changes class, or shifts a detection.
"""

from typing import Any, Optional, Tuple
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image

from ..core.base import BaseModelAdapter, Detection


@dataclass
class CounterfactualResult:
    """Result from counterfactual analysis."""
    
    # The perturbation that flips the detection
    perturbation: torch.Tensor
    
    # Mask showing critical regions (high perturbation areas)
    critical_mask: torch.Tensor
    
    # Perturbed image
    perturbed_image: torch.Tensor
    
    # Original vs new detection state
    original_detection: Detection
    new_detections: list  # What model detects after perturbation
    
    # Success metrics
    success: bool  # Did we achieve the target?
    perturbation_norm: float  # L2 norm of perturbation
    iterations: int  # How many PGD steps taken


class CounterfactualProbe:
    """
    Find minimal changes that flip model decisions.
    
    Uses PGD optimization to answer:
    - "What would remove this detection?"
    - "What would change the class?"  
    - "What would shift the bounding box?"
    
    Usage:
        probe = CounterfactualProbe(adapter)
        result = probe.find_critical_change(image, detection_idx=0, target="remove")
    """
    
    def __init__(
        self, 
        adapter: BaseModelAdapter,
        epsilon: float = 0.1,  # Max perturbation magnitude
        step_size: float = 0.01,
        max_iterations: int = 100,
    ):
        self.adapter = adapter
        self.epsilon = epsilon
        self.step_size = step_size
        self.max_iterations = max_iterations
    
    def find_critical_change(
        self,
        image: Any,
        detection_idx: int,
        target: str = "remove",  # "remove", "change_class", "shift_bbox"
        conf_threshold: float = 0.25,
    ) -> CounterfactualResult:
        """
        Find minimal perturbation that achieves target.
        
        Args:
            image: Input image
            detection_idx: Which detection to target
            target: What change to achieve
            conf_threshold: Confidence threshold for detection
            
        Returns:
            CounterfactualResult with perturbation and analysis
        """
        # Get original detections
        original_detections = self.adapter.predict(image)
        if detection_idx >= len(original_detections):
            raise IndexError(f"detection_idx {detection_idx} out of range")
        
        target_detection = original_detections[detection_idx]
        
        # Preprocess image
        input_tensor = self.adapter.preprocess(image)
        original_tensor = input_tensor.clone()
        
        # Initialize perturbation
        perturbation = torch.zeros_like(input_tensor, requires_grad=True)
        
        optimizer = torch.optim.Adam([perturbation], lr=self.step_size)
        
        success = False
        for iteration in range(self.max_iterations):
            optimizer.zero_grad()
            
            # Apply perturbation
            perturbed = original_tensor + perturbation
            perturbed = torch.clamp(perturbed, 0, 1)
            
            # Compute loss based on target
            loss = self._compute_target_loss(
                perturbed, 
                target_detection,
                detection_idx,
                target
            )
            
            # Add regularization for minimal perturbation
            reg_loss = perturbation.norm(p=2)
            total_loss = loss + 0.01 * reg_loss
            
            total_loss.backward()
            optimizer.step()
            
            # Project perturbation to epsilon ball
            with torch.no_grad():
                perturbation.data = self._project_to_ball(
                    perturbation.data, self.epsilon
                )
            
            # Check if target achieved
            if self._check_success(perturbed, target_detection, detection_idx, target):
                success = True
                break
        
        # Final perturbed image
        final_perturbed = torch.clamp(original_tensor + perturbation, 0, 1)
        
        # Get new detections
        new_detections = self._get_detections_from_tensor(final_perturbed)
        
        # Compute critical mask (where perturbation is high)
        critical_mask = self._compute_critical_mask(perturbation)
        
        return CounterfactualResult(
            perturbation=perturbation.detach(),
            critical_mask=critical_mask,
            perturbed_image=final_perturbed.detach(),
            original_detection=target_detection,
            new_detections=new_detections,
            success=success,
            perturbation_norm=float(perturbation.norm(p=2)),
            iterations=iteration + 1,
        )
    
    def _compute_target_loss(
        self,
        perturbed: torch.Tensor,
        target_detection: Detection,
        detection_idx: int,
        target: str
    ) -> torch.Tensor:
        """Compute loss for target objective."""
        
        # Forward pass through model
        pytorch_model = self.adapter.get_pytorch_model()
        pytorch_model.eval()
        
        output = pytorch_model(perturbed)
        
        if target == "remove":
            # Maximize confidence drop
            # We want the detection confidence to go below threshold
            return self._compute_removal_loss(output, target_detection, detection_idx)
        
        elif target == "change_class":
            # Push towards different class
            return self._compute_class_change_loss(output, target_detection, detection_idx)
        
        elif target == "shift_bbox":
            # Push bbox coordinates away
            return self._compute_bbox_shift_loss(output, target_detection, detection_idx)
        
        else:
            raise ValueError(f"Unknown target: {target}")
    
    def _compute_removal_loss(
        self,
        output: Any,
        target_detection: Detection,
        detection_idx: int
    ) -> torch.Tensor:
        """Loss that encourages detection removal (confidence drop)."""
        # This is model-specific
        # For demonstration, we use a proxy loss
        
        if isinstance(output, (list, tuple)):
            pred = output[0]
        else:
            pred = output
        
        # We want to minimize confidence at the detection location
        # Negative loss = we minimize (-confidence) = maximize confidence drop
        if hasattr(pred, 'shape') and len(pred.shape) >= 2:
            # Use prediction at detection index as proxy
            idx = detection_idx % pred.shape[1] if pred.shape[1] > 0 else 0
            confidence_proxy = pred[0, idx].mean()
            return -confidence_proxy  # Minimize = push confidence down
        
        return torch.tensor(0.0, requires_grad=True)
    
    def _compute_class_change_loss(
        self,
        output: Any,
        target_detection: Detection,
        detection_idx: int
    ) -> torch.Tensor:
        """Loss that encourages class change."""
        if isinstance(output, (list, tuple)):
            pred = output[0]
        else:
            pred = output
        
        # Push away from current class, towards others
        if hasattr(pred, 'shape') and len(pred.shape) >= 2:
            idx = detection_idx % pred.shape[1] if pred.shape[1] > 0 else 0
            # Increase entropy / push to different class
            return pred[0, idx].mean()
        
        return torch.tensor(0.0, requires_grad=True)
    
    def _compute_bbox_shift_loss(
        self,
        output: Any,
        target_detection: Detection,
        detection_idx: int
    ) -> torch.Tensor:
        """Loss that encourages bbox coordinate change."""
        # Similar to above - model specific
        if isinstance(output, (list, tuple)):
            pred = output[0]
        else:
            pred = output
        
        if hasattr(pred, 'shape') and len(pred.shape) >= 2:
            idx = detection_idx % pred.shape[1] if pred.shape[1] > 0 else 0
            return pred[0, idx, :4].mean() if pred.shape[-1] >= 4 else pred[0, idx].mean()
        
        return torch.tensor(0.0, requires_grad=True)
    
    def _project_to_ball(
        self, 
        perturbation: torch.Tensor, 
        epsilon: float
    ) -> torch.Tensor:
        """Project perturbation to L-infinity ball."""
        return torch.clamp(perturbation, -epsilon, epsilon)
    
    def _check_success(
        self,
        perturbed: torch.Tensor,
        target_detection: Detection,
        detection_idx: int,
        target: str
    ) -> bool:
        """Check if counterfactual objective achieved."""
        new_dets = self._get_detections_from_tensor(perturbed)
        
        if target == "remove":
            # Check if detection is gone or confidence dropped significantly
            return len(new_dets) < detection_idx + 1 or (
                len(new_dets) > detection_idx and 
                new_dets[detection_idx].confidence < 0.1
            )
        
        elif target == "change_class":
            # Check if class changed
            if len(new_dets) > detection_idx:
                return new_dets[detection_idx].class_id != target_detection.class_id
            return False
        
        elif target == "shift_bbox":
            # Check if IoU dropped below 0.5
            if len(new_dets) > detection_idx:
                iou = self._compute_iou(
                    target_detection.bbox, 
                    new_dets[detection_idx].bbox
                )
                return iou < 0.5
            return False
        
        return False
    
    def _get_detections_from_tensor(self, tensor: torch.Tensor) -> list:
        """Run inference on tensor and return detections."""
        # Convert tensor back to image format for prediction
        # This is a simplified version
        with torch.no_grad():
            img = tensor.squeeze().permute(1, 2, 0).cpu().numpy()
            img = (img * 255).astype(np.uint8)
            return self.adapter.predict(img)
    
    def _compute_critical_mask(self, perturbation: torch.Tensor) -> torch.Tensor:
        """Identify critical regions from perturbation magnitude."""
        # Compute per-pixel perturbation magnitude
        magnitude = perturbation.abs().mean(dim=1)  # Average across channels
        
        # Normalize to [0, 1]
        mag_min = magnitude.min()
        mag_max = magnitude.max()
        if mag_max - mag_min > 1e-8:
            mask = (magnitude - mag_min) / (mag_max - mag_min)
        else:
            mask = torch.zeros_like(magnitude)
        
        return mask.squeeze()
    
    def _compute_iou(
        self, 
        box1: Tuple[float, ...], 
        box2: Tuple[float, ...]
    ) -> float:
        """Compute IoU between two boxes."""
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])
        
        inter = max(0, x2 - x1) * max(0, y2 - y1)
        
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        
        union = area1 + area2 - inter
        
        return inter / union if union > 0 else 0.0
