"""
src/gradcam.py
Grad-CAM implementation for visualizing model decisions
"""
import cv2
import numpy as np
import tensorflow as tf
from typing import Optional, Tuple, Union
class GradCAM:
    """
    Gradient-weighted Class Activation Mapping (Grad-CAM)
    
    Visualizes which regions of an image the model focuses on for its predictions.
    Based on the original Grad-CAM paper: "Grad-CAM: Visual Explanations from Deep Neural Networks via Gradient-based Localization"
    """
    
    def __init__(self, model, last_conv_layer_name: str = 'out_relu'):
        """
        Initialize Grad-CAM with a Keras model.
        
        Args:
            model: Keras model
            last_conv_layer_name: Name of the last convolutional layer to use
        """
        self.model = model
        self.last_conv_layer_name = last_conv_layer_name
        
        # Create the gradient model once
        self.grad_model = self._create_gradient_model()
    
    def _create_gradient_model(self):
        """
        Create a model that connects the last conv layer and model output.
        
        Returns:
            Keras Model that outputs both the last conv layer and predictions
        """
        return tf.keras.Model(
            [self.model.inputs],
            [self.model.get_layer(self.last_conv_layer_name).output, self.model.output]
        )
    
    def generate_heatmap(self, image_array: np.ndarray, class_index: int) -> np.ndarray:
        """
        Generate a Grad-CAM heatmap for a given image and class.
        
        Args:
            image_array: Preprocessed image array (with batch dimension)
            class_index: Index of the class to generate heatmap for
            
        Returns:
            Generated heatmap (numpy array, shape: H x W)
        """
        with tf.GradientTape() as tape:
            last_conv_layer_output, preds = self.grad_model(image_array)
            class_channel = preds[:, class_index]
        
        grads = tape.gradient(class_channel, last_conv_layer_output)
        
        # Pool gradients across spatial dimensions
        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
        
        # Get the output of the last conv layer
        last_conv_layer_output = last_conv_layer_output[0]
        
        # Compute the weighted combination of feature maps
        heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
        heatmap = tf.squeeze(heatmap)
        
        # Apply ReLU to get positive weights
        heatmap = tf.maximum(heatmap, 0)
        
        # Normalize heatmap to [0, 1]
        max_val = tf.math.reduce_max(heatmap)
        if max_val > 0:
            heatmap = heatmap / max_val
        
        return heatmap.numpy()
    
    def overlay_heatmap(self, original_image: np.ndarray, heatmap: np.ndarray, 
                       alpha: float = 0.5) -> np.ndarray:
        """
        Overlay heatmap on original image using JET colormap.
        
        Args:
            original_image: Original image (numpy array, RGB)
            heatmap: Grad-CAM heatmap (normalized to [0, 1])
            alpha: Weight for heatmap in overlay (0-1)
            
        Returns:
            Image with heatmap overlay (numpy array, RGB)
        """
        # Resize heatmap to match original image dimensions
        heatmap = cv2.resize(heatmap, (original_image.shape[1], original_image.shape[0]))
        
        # Convert heatmap to 8-bit format and apply JET colormap
        heatmap = np.uint8(255 * heatmap)
        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        
        # Convert original image to float for blending
        original_float = original_image.astype(np.float32)
        heatmap_float = heatmap.astype(np.float32)
        
        # Blend original image with heatmap
        superimposed_img = heatmap_float * alpha + original_float
        superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8)
        
        return superimposed_img
    
    @staticmethod
    def preprocess_frame(frame: np.ndarray, target_size: Tuple[int, int] = (224, 224)):
        """
        Preprocess a video frame for model inference.
        
        Args:
            frame: Input frame (numpy array, BGR format from OpenCV)
            target_size: Target size for model input
            
        Returns:
            Preprocessed frame (numpy array)
        """
        # Resize to model input size
        resized = cv2.resize(frame, target_size)
        # Normalize to [0, 1]
        normalized = resized / 255.0
        # Add batch dimension
        batched = np.expand_dims(normalized, axis=0)
        
        return batched
    
    @staticmethod
    def get_confidence_and_prediction(predictions: np.ndarray, 
                                     class_labels: list) -> Tuple[str, float, np.ndarray]:
        """
        Get prediction label, confidence, and full probability array.
        
        Args:
            predictions: Model output probabilities (array of shape [1, 3])
            class_labels: List of class labels
            
        Returns:
            Tuple of (predicted_label, confidence, probability_array)
        """
        prob_array = predictions[0]
        max_idx = np.argmax(prob_array)
        confidence = prob_array[max_idx]
        label = class_labels[max_idx]
        
        return label, confidence, prob_array
    
    @staticmethod
    def get_text_overlay_params(label: str, color_map: dict):
        """
        Get parameters for text overlay on frames.
        
        Args:
            label: Predicted class label
            color_map: Dictionary mapping labels to colors
            
        Returns:
            Tuple of (text_color, background_color)
        """
        text_color = color_map.get(label, (255, 255, 255))  # White default
        background_color = (0, 0, 0)  # Black background
        return text_color, background_color
    
    @staticmethod
    def create_text_overlay(frame: np.ndarray, text: str, 
                           text_color: Tuple[int, int, int], 
                           bg_color: Tuple[int, int, int] = (0, 0, 0)):
        """
        Create a text overlay with background on the frame.
        
        Args:
            frame: Input frame
            text: Text to display
            text_color: Color of the text
            bg_color: Background color
            
        Returns:
            Frame with text overlay
        """
        # Get text size
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.7
        thickness = 2
        (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, thickness)
        
        # Define text position with padding
        text_x, text_y = 10, 40
        rect_x1, rect_y1 = text_x - 5, text_y - text_height - baseline - 5
        rect_x2, rect_y2 = text_x + text_width + 5, text_y + baseline + 5
        
        # Draw background rectangle
        cv2.rectangle(frame, (rect_x1, rect_y1), (rect_x2, rect_y2), bg_color, -1)
        
        # Draw text
        cv2.putText(frame, text, (text_x, text_y), font, font_scale, text_color, thickness)
        
        return frame