import matplotlib.pyplot as plt
import numpy as np
import cv2

def plot_greenery_overlay(image_np, mask_classifications, output_path=None):
    """
    Plots the original image and an overlay where greenery is highlighted.
    """
    plt.figure(figsize=(10, 10))
    plt.imshow(image_np)

    ax = plt.gca()
    ax.set_autoscale_on(False)

    # Create an aggregate mask for greenery
    h, w = image_np.shape[:2]
    greenery_overlay = np.zeros((h, w, 4))

    green_color = np.array([0, 1, 0, 0.5]) # Semi-transparent green

    for item in mask_classifications:
        if item['is_green']:
            seg = item['segmentation']
            greenery_overlay[seg] = green_color
            
    ax.imshow(greenery_overlay)
    plt.axis('off')
    
    if output_path:
        plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
        plt.close()
    else:
        plt.show()


def create_greenery_overlay(image_np, mask_classifications):
    """
    Creates a semi-transparent greenery overlay composite image.

    Returns:
        tuple: (composite_rgba, green_masks) where composite_rgba is the
               overlay image and green_masks is the list of greenery items.
    """
    h, w = image_np.shape[:2]
    greenery_overlay = np.zeros((h, w, 4))
    green_color = np.array([0, 1, 0, 0.5])  # Semi-transparent green

    green_masks = []
    for item in mask_classifications:
        if item['is_green']:
            green_masks.append(item)
            seg = item['segmentation']
            greenery_overlay[seg] = green_color

    base_img_rgba = cv2.cvtColor(image_np, cv2.COLOR_RGB2RGBA)
    alpha_mask = greenery_overlay[:, :, 3] > 0
    composite = base_img_rgba.copy()
    composite[alpha_mask] = (
        base_img_rgba[alpha_mask] * 0.5 +
        greenery_overlay[alpha_mask] * 255 * 0.5
    ).astype(np.uint8)

    return composite, green_masks


def apply_grad_cam(model, input_tensor, target_class=None):
    """
    Generates a Grad-CAM heatmap for the given input tensor using the model's
    final convolutional layer (layer4 for ResNet-50).
    
    Args:
        model: A GreeneryClassifier instance (wraps ResNet-50 via model.model).
        input_tensor: A preprocessed image tensor of shape (1, 3, 64, 64).
        target_class: The class index to generate the heatmap for.
                      If None, uses the model's predicted class.
    
    Returns:
        heatmap: A numpy array (H, W) with values in [0, 1] representing
                 the class activation map.
        predicted_class: The class index the model predicted.
    """
    import torch
    import torch.nn.functional as F

    model.eval()
    device = next(model.parameters()).device
    input_tensor = input_tensor.to(device)

    # Storage for hooked values
    activations = []
    gradients = []

    # Hook into the last convolutional block of the inner ResNet-50
    target_layer = model.model.layer4

    def forward_hook(module, input, output):
        activations.append(output.detach())

    def backward_hook(module, grad_input, grad_output):
        gradients.append(grad_output[0].detach())

    fwd_handle = target_layer.register_forward_hook(forward_hook)
    bwd_handle = target_layer.register_full_backward_hook(backward_hook)

    try:
        # Forward pass
        output = model(input_tensor)
        predicted_class = output.argmax(dim=1).item()

        if target_class is None:
            target_class = predicted_class

        # Backward pass for the target class
        model.zero_grad()
        class_score = output[0, target_class]
        class_score.backward()

        # Compute Grad-CAM
        act = activations[0]   # (1, C, H', W')
        grad = gradients[0]    # (1, C, H', W')

        # Global average pool the gradients to get per-channel weights
        weights = grad.mean(dim=[2, 3], keepdim=True)  # (1, C, 1, 1)

        # Weighted combination of activation maps
        cam = (weights * act).sum(dim=1, keepdim=True)  # (1, 1, H', W')
        cam = F.relu(cam)  # Only positive contributions

        # Upsample to input size
        cam = F.interpolate(cam, size=input_tensor.shape[2:], mode='bilinear', align_corners=False)
        cam = cam.squeeze().cpu().numpy()

        # Normalize to [0, 1]
        if cam.max() > 0:
            cam = (cam - cam.min()) / (cam.max() - cam.min())

        return cam, predicted_class
    finally:
        # Guarantee hook cleanup even if an exception fires mid-computation
        fwd_handle.remove()
        bwd_handle.remove()


def save_grad_cam_overlay(image_np, heatmap, output_path, alpha=0.4):
    """
    Overlays a Grad-CAM heatmap on the original image and saves to disk.
    
    Args:
        image_np: Original RGB image as numpy array (H, W, 3), values 0-255.
        heatmap: Grad-CAM heatmap as numpy array (H, W), values in [0, 1].
        output_path: File path to save the result.
        alpha: Transparency of the heatmap overlay.
    """
    # Resize heatmap to match the image dimensions
    heatmap_resized = cv2.resize(heatmap, (image_np.shape[1], image_np.shape[0]))
    heatmap_colored = cv2.applyColorMap(np.uint8(255 * heatmap_resized), cv2.COLORMAP_JET)
    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)

    overlay = np.uint8(image_np * (1 - alpha) + heatmap_colored * alpha)

    plt.figure(figsize=(10, 10))
    plt.imshow(overlay)
    plt.axis('off')
    plt.title('Grad-CAM Heatmap Overlay')
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    plt.close()
    print(f"Grad-CAM overlay saved to {output_path}")