"""
explainability/gradcam.py
--------------------------
Grad-CAM implementation for the CNN Branch (EfficientNet-B0).

If the CNN model is untrained or fails, generates a meaningful
image-based saliency heatmap so the visualization is always populated.
"""

import numpy as np
import base64
import io
import cv2


# ─────────────────────────────────────────────────────────────────
# Helper: encode numpy image to base64 JPEG string
# ─────────────────────────────────────────────────────────────────

def _encode_b64(img_rgb: np.ndarray) -> str:
    """Encode (H, W, 3) uint8 RGB image to base64 JPEG string."""
    bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
    _, buf = cv2.imencode(".jpg", bgr, [cv2.IMWRITE_JPEG_QUALITY, 90])
    return base64.b64encode(buf.tobytes()).decode("utf-8")


# ─────────────────────────────────────────────────────────────────
# Saliency-based fallback (works without any model)
# ─────────────────────────────────────────────────────────────────

def _saliency_heatmap(img_tensor_or_array) -> dict:
    """
    Generate a gradient-free saliency/edge prominence heatmap from the
    image itself using Laplacian + Sobel, colorized like a Grad-CAM.
    This ensures the Grad-CAM panel always shows something informative.
    """
    try:
        # Accept tf.Tensor or numpy array
        try:
            arr = img_tensor_or_array.numpy()
        except AttributeError:
            arr = np.array(img_tensor_or_array)

        # Remove batch dim if present
        if arr.ndim == 4:
            arr = arr[0]

        # Convert [0,1] float to uint8
        img_u8 = np.clip(arr * 255, 0, 255).astype(np.uint8)
        gray   = cv2.cvtColor(img_u8, cv2.COLOR_RGB2GRAY)

        # Laplacian edge prominence
        lap = cv2.Laplacian(gray, cv2.CV_64F)
        lap = np.abs(lap)

        # Sobel magnitude
        sx  = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sy  = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        sobel = np.sqrt(sx**2 + sy**2)

        # Combine & blur for smooth heatmap
        combined = 0.5 * lap + 0.5 * sobel
        combined = cv2.GaussianBlur(combined, (21, 21), 0)

        # Normalize to [0, 255]
        mn, mx = combined.min(), combined.max()
        if mx > mn:
            combined = (combined - mn) / (mx - mn)
        heatmap_u8 = np.uint8(255 * combined)

        # Apply JET colormap
        colored = cv2.applyColorMap(heatmap_u8, cv2.COLORMAP_JET)
        colored_rgb = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)

        # Overlay on original image
        overlay = cv2.addWeighted(img_u8, 0.5, colored_rgb, 0.5, 0)

        return {
            "heatmap_b64": _encode_b64(overlay),
            "available":   True,     # show the image in the panel
            "note":        "Saliency heatmap (CNN weights loading failed — showing edge-based proxy)",
        }
    except Exception as e:
        print(f"[GradCAM] Saliency fallback failed: {e}")
        return _fallback_heatmap()


def _fallback_heatmap() -> dict:
    """Minimal fallback — tells frontend CNN is unavailable."""
    return {"heatmap_b64": "", "available": False}


# ─────────────────────────────────────────────────────────────────
# True Grad-CAM (requires trained CNN model)
# ─────────────────────────────────────────────────────────────────

def compute_gradcam(feature_model, img_tensor, target_class: int = 1) -> dict:
    """
    Compute Grad-CAM heatmap using the CNN branch feature model.

    Args:
        feature_model : Keras Model returning (conv_output, predictions)
        img_tensor    : tf.Tensor (1, 224, 224, 3)
        target_class  : 1 = fake class

    Returns:
        dict with heatmap_b64, available, overlay_b64
    """
    try:
        import tensorflow as tf

        with tf.GradientTape() as tape:
            inputs = tf.cast(img_tensor, tf.float32)
            tape.watch(inputs)
            conv_outputs, predictions = feature_model(inputs, training=False)
            if predictions.shape[-1] == 1:
                loss = predictions[:, 0]
            else:
                loss = predictions[:, target_class]

        grads = tape.gradient(loss, conv_outputs)
        if grads is None:
            raise ValueError("Gradients are None — check feature_model output.")

        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
        conv_out     = conv_outputs[0]
        heatmap      = tf.reduce_sum(
            conv_out * pooled_grads[tf.newaxis, tf.newaxis, :], axis=-1
        )
        heatmap = np.maximum(heatmap.numpy(), 0)
        if heatmap.max() > 0:
            heatmap /= heatmap.max()

        return _build_gradcam_output(heatmap, img_tensor)

    except Exception as e:
        print(f"[GradCAM] compute_gradcam failed: {e}. Falling back to saliency.")
        return _saliency_heatmap(img_tensor)


def _build_gradcam_output(heatmap: np.ndarray, img_tensor) -> dict:
    """Resize heatmap, colorize, overlay on original, return base64."""
    try:
        # Get original image
        try:
            arr = img_tensor.numpy()
        except AttributeError:
            arr = np.array(img_tensor)
        if arr.ndim == 4:
            arr = arr[0]

        img_u8 = np.clip(arr * 255, 0, 255).astype(np.uint8)
        h, w   = img_u8.shape[:2]

        # Resize heatmap to image size
        hm_resized = cv2.resize(heatmap, (w, h))
        hm_u8      = np.uint8(255 * hm_resized)
        colored    = cv2.applyColorMap(hm_u8, cv2.COLORMAP_JET)
        colored_rgb = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)

        overlay = cv2.addWeighted(img_u8, 0.55, colored_rgb, 0.45, 0)

        return {
            "heatmap_b64": _encode_b64(overlay),
            "available":   True,
        }
    except Exception as e:
        print(f"[GradCAM] _build_gradcam_output failed: {e}")
        return _fallback_heatmap()