from joblib import Parallel, delayed
from matplotlib.colors import TwoSlopeNorm
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from granatpy.core import naturalize_rgb_image

__all__ = ["overlay_naturalness_heatmap"]

def _plot_heatmap(img_array, heatmap_data, figsize, grid_rows, grid_cols, cmap, alpha, vmin, vmax, title, label, ax=None, use_two_slope=False):

    skip_label = False

    H, W, _ = img_array.shape
    
    # Compute cell dimensions
    cell_height = H // grid_rows
    cell_width = W // grid_cols

    # If no axis provided, create new figure
    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)
    else:
        skip_label = True
        fig = ax.figure

    ax.imshow(img_array)
        
    # Normalize heatmap data for coloring
    if use_two_slope:
        norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax)
    else:
        norm = Normalize(vmin=heatmap_data.min(), vmax=heatmap_data.max())

    # Draw a rectangle for each cell with correct position and color
    for i in range(grid_rows):
        for j in range(grid_cols):
            y_start = i * cell_height
            y_end = min((i + 1) * cell_height, H)
            x_start = j * cell_width
            x_end = min((j + 1) * cell_width, W)
            
            # Get color from colormap based on normalized value
            color = plt.cm.get_cmap(cmap)(norm(heatmap_data[i, j]))
            
            # Create and add rectangle with transparency
            rect = plt.Rectangle(
                (x_start, y_start),  # bottom-left corner
                x_end - x_start,     # width
                y_end - y_start,     # height
                linewidth=0.5,
                edgecolor='black',   # thin border for clarity
                facecolor=color,     # colored based on naturalness
                alpha=alpha          # transparency applied here!
            )
            ax.add_patch(rect)
    
    # Set limits to avoid white borders
    ax.set_xlim(0, W)
    ax.set_ylim(H, 0)  # Invert y-axis to match image coordinates
    
    if not skip_label:
        # Add colorbar
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
        sm.set_array([])
        # sm.autoscale_None()
        # sm.set_clim(0.5, 4.5)
        cbar = plt.colorbar(sm, ax=ax, shrink=0.8)
        cbar.set_label(label, rotation=270, labelpad=20)
    
    ax.set_title(title)
    ax.axis('off')  # Hide axes for cleaner look
    
    plt.tight_layout()

    return fig, ax

def _process_cell(args):
    """Processes a single cell of a grid using the naturalisation algorithm. This function is extracted so that we can multithread it."""
    i, j, img_array, grid_rows, grid_cols, H, W, cell_height, cell_width, verbose = args
    
    # Define cell boundaries
    y_start = i * cell_height
    y_end = min((i + 1) * cell_height, H)
    x_start = j * cell_width
    x_end = min((j + 1) * cell_width, W)
    
    # Extract cell region
    cell = img_array[y_start:y_end, x_start:x_end]
    
    # Skip if cell is empty
    if cell.size == 0:
        if verbose:
            print(f"Cell size is 0 for cell {i}, {j}")
        return i, j, 0.0
    
    # Compute naturalness stats for this cell
    try:
        _, nfs, _, _, _ = naturalize_rgb_image(cell)
        # Average the three Nfs values (R, G, B)
        avg_nfs = np.mean(nfs)
        print(f"Thread for cell {i} {j} finished, returning avg Nfs now.")
        return i, j, avg_nfs
    except Exception as e:
        print(f"Warning: Failed to compute naturalness for cell ({i}, {j}): {e}")
        return i, j, 0.0

def overlay_naturalness_heatmap(img_array, grid_rows=10, grid_cols=15, cmap='YlOrRd', alpha=0.5, figsize=(12, 8), verbose = False):
    """
    Overlay a heatmap of naturalness factors on the original image.
    
    Divides the image into a grid of (grid_rows x grid_cols) cells.
    For each cell, computes the average naturalness factor (mean of Nfs across RGB channels)
    using the naturalize_rgb_image function, then plots a semi-transparent heatmap overlay.
    
    Args:
        img_array: numpy array of shape (H, W, 3) representing an RGB image
        grid_rows: int, number of rows in the grid (default: 10)
        grid_cols: int, number of columns in the grid (default: 15)
        cmap: str, matplotlib colormap for the heatmap (default: 'YlOrRd')
        alpha: float, transparency of the heatmap overlay (0=transparent, 1=opaque)
        figsize: tuple, figure size for plotting
    
    Returns:
        fig, ax: matplotlib figure and axes objects
        heatmap_data: 2D numpy array of shape (grid_rows, grid_cols) with average Nfs per cell
    """
    
    H, W, _ = img_array.shape
    
    print(f"image shape is height {H} x width {W}")

    # Compute cell dimensions
    cell_height = H // grid_rows
    cell_width = W // grid_cols
    
    # Initialize heatmap data
    heatmap_data = np.zeros((grid_rows, grid_cols))
    
    # Prepare arguments for each cell
    args_list = []
    for i in range(grid_rows):
        for j in range(grid_cols):
            args_list.append((i, j, img_array, grid_rows, grid_cols, H, W, cell_height, cell_width, verbose))

    print("Assembled args list for multithreaded processing.")

    # Process cells in parallel
    results = Parallel(n_jobs=-1)(delayed(_process_cell)(args) for args in args_list)

    # Fill heatmap_data with results
    for i, j, value in results:
        heatmap_data[i, j] = value
    
    fig, ax = _plot_heatmap(img_array, heatmap_data, figsize, grid_rows, grid_cols, cmap, alpha, heatmap_data.min(), heatmap_data.max(),
                           'Image with Naturalness Factor Heatmap Overlay',
                           'Average Naturalness Factor (Nfs)')

    return fig, ax, heatmap_data