import cv2
import numpy as np

class DICOM_Utils(object):
    def apply_windowing(image_array, window_center, window_width):
        """
        Apply windowing to a DICOM image array.

        Parameters:
        - image_array: numpy array of the DICOM image
        - window_center: center of the window
        - window_width: width of the window

        Returns:
        - Windowed image array
        """
        lower_bound = window_center - (window_width / 2)
        upper_bound = window_center + (window_width / 2)

        # Apply windowing
        windowed_image = image_array.copy()
        windowed_image[windowed_image < lower_bound] = lower_bound
        windowed_image[windowed_image > upper_bound] = upper_bound

        # Normalize to [0, 255]
        windowed_image = ((windowed_image - lower_bound) / window_width) * 255

        return windowed_image.astype('uint8')

    def apply_CAM_overlay(heatmap, windowed_image, overlay_alpha=0.4):
        """
        Apply CAM (Class Activation Map) overlay to a given image.

        Parameters:
        - heatmap: torch.Tensor, the heatmap generated by CAM.
        - windowed_image: numpy.ndarray, the windowed image to overlay the heatmap on.
        - overlay_alpha: float, the transparency for overlaying heatmap. Default is 0.4.

        Returns:
        - overlayed: numpy.ndarray, the resulting image after overlaying the heatmap.
        """
        # Convert the heatmap tensor to a numpy array
        heatmap_np = heatmap.cpu().numpy().squeeze()

        # Normalize the heatmap to [0, 255]
        heatmap_normalized = ((heatmap_np - heatmap_np.min()) / 
                            (heatmap_np.max() - heatmap_np.min()) * 255).astype(np.uint8)

        # Convert the normalized heatmap to a colormap (for example, using the "jet" colormap)
        heatmap_colormap = cv2.applyColorMap(heatmap_normalized, cv2.COLORMAP_JET)

        # Resize the colormap to the original image size
        heatmap_resized = cv2.resize(heatmap_colormap, 
                                    (windowed_image.shape[1], windowed_image.shape[0]))

        # Convert the grayscale windowed_image to 3 channels
        windowed_image_colored = cv2.cvtColor(windowed_image, cv2.COLOR_GRAY2BGR)

        # Overlay the heatmap on the original image with a certain transparency
        overlayed = cv2.addWeighted(windowed_image_colored, 1 - overlay_alpha, 
                                    heatmap_resized, overlay_alpha, 0)

        return overlayed