class ColorContourExtractor:
    """"   
    ## ColorContourExtractor
    
    This class is responsible for processing an image to isolate the color of the pins.
    
    ### Methods
    - `process_image(scanned_image: np.ndarray) -> np.ndarray`
        - This method pre-processes the image to isolate the color of the pins.
        
    - `show_result(edges: np.ndarray) -> None`
        - This method shows the result of the pre-processing.
    
    ### Example
    ```python
    import cv2 as cv
    import numpy as np
    from src.objs.image.processors.image_processor import ImageProcessor

    scanned_image = cv.imread('path/to/image.jpg')
    edges = ImageProcessor.process_image(scanned_image)
    ImageProcessor.show_result(edges)
    ```
    """

    # A function that pre-processes the image to isolate the color of the pins.
    @staticmethod
    def process_image(
        scanned_image: np.ndarray, 
        hsv_lower = [0, 55, 0], 
        hsv_upper = [360, 255,255], 
        double_thresh:bool = False, 
        display:bool=False) -> np.ndarray:
        """ 
        This method pre-processes the image to isolate the color of the pins in Grid to find blocks. 
        It is also used in TestAnalyzer to find the positive spots with hsv mask.
        Thread carefully
        """

        # Copy the image to avoid modifying the original image
        scanned_image_copy = scanned_image.copy()
        
        # Convert the image to HSV color space. Hue Saturation Value. 
        # Similar to RGB but more useful for color isolation.
        img_hsv = cv.cvtColor(scanned_image_copy, cv.COLOR_BGR2HSV)

        # Define the lower and upper bounds for the color you want to isolate
        # These values are the product of trial and error and are not necessarily perfect.
        hsv_lower_color = np.array(hsv_lower)
        hsv_upper_color = np.array(hsv_upper)

        # Create a mask to filter out the grayscale colors isolating the color of the pins.
        color_mask = cv.inRange(img_hsv, hsv_lower_color, hsv_upper_color)

        # Visualize the mask on top of the original image before thresholding
        #mask_before_thresholding = color_mask.copy()

        edges = cv.Canny(color_mask, 0, 255)

        if double_thresh:
            
            second_mask = cv.bitwise_and(scanned_image_copy, scanned_image_copy, mask=color_mask)
            #cv.imshow('bitwise and image + mask 1', cv.resize(color_mask,(200,200)))
            
            # make all black pixels white
            second_mask[second_mask == 0] = 255
            #cv.imshow('make black pixels white',  cv.resize(color_mask, (200, 200)))
            
            #  thresholding
            second_mask = cv.cvtColor(second_mask, cv.COLOR_BGR2GRAY)
            #cv.imshow('grey',  cv.resize(color_mask, (200, 200)))
            
            second_mask = cv.bitwise_not(second_mask) 
            #cv.imshow('bitwise not',  cv.resize(second_mask, (200, 200)))
            #cv.waitKey(0)

            cv.threshold(second_mask, 125, 255, cv.THRESH_BINARY, second_mask)
            
            edges = cv.Canny(second_mask, 0, 255)
            
            #cv.imshow('second thresholding ',  cv.resize(
             #   cv.bitwise_and(scanned_image_copy,scanned_image_copy, mask=second_mask), (400, 400)))   
            #cv.imshow('first thresholding ',  cv.resize(
              #  cv.bitwise_and(scanned_image_copy,scanned_image_copy, mask=color_mask), (400, 400)))
            #cv.waitKey(0)

            
        #cv.destroyAllWindows()
            
        contours, _ = cv.findContours(edges, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)
    
        if display:
            pass
            contours = ColorContourExtractor.show_result(contours, scanned_image_copy)

        return contours
    
    # Show the result of the pre-processing.
    @staticmethod
    def show_result(contours: np.ndarray, image) -> None:
        """ this method shows the result of the pre-processing."""

        copy = image.copy()

        # show only the pixels where sqrt(a^2 + b^2) > 10
        # this will remove the background noise
        lab = cv.cvtColor(copy, cv.COLOR_BGR2LAB)
        lab = cv.blur(lab, (3, 3))

        # split the image into L, A, and B channels
        l, a, b = cv.split(lab)
        r, g, b = cv.split(copy)

        plt.imshow(copy)
        plt.title('Color Contour Extractor - COPY')
        plt.show()
        

        for row in range(len(l)):
            for col in range(len(l[0])):
                if  l[row][col] >= 225:
                    l[row][col] = 0
                    a[row][col] = 0
                    b[row][col] = 0
                else:
                    l[row][col] = 255
                    a[row][col] = 255
                    b[row][col] = 255
        
        lab = cv.merge((l, a, b))
    
        mask = cv.bitwise_and(lab, copy)
        mask = cv.cvtColor(mask, cv.COLOR_BGR2GRAY)
        mask = cv.bitwise_not(mask)

        # draw the contours
        edges = cv.Canny(mask, 0, 255)
        contours, _ = cv.findContours(edges, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)
        cv.drawContours(copy, contours, -1, (0, 255, 0), 1)
        copy = cv.resize(copy, (400, 400))


        """
        copy = image.copy()
        cv.drawContours(copy, contours, -1, (0, 255, 0), 1)
        copy = cv.resize(copy, (400, 400))
        """
        plt.imshow(copy)
        plt.title('Color Contour Extractor')
        plt.show()
        
        return max(contours, key = cv.contourArea)

class GridImageNormalizer:
    """
    ### Image Normalizer
    Class to normalize the image of the grid by scanning the grid and making it square ratio.

    #### Methods:
    - `scan(id: int, image: ndarray, resize_factor: float = 1) -> (Image, int)`
        - This method scans the image and returns the scanned image.
    - `resize_2_std(img: ndarray, factor: float, w:int=None, h:int = None) -> ndarray`
        - This method resizes the image to a given percentage of the current size.
    """
    @classmethod
    def scan(cls, image_name: str, image: ndarray, do_white_balance:bool):
        """
        ### Scan image
        Scan the image and return the scanned image.

        #### Args:
        * id : id of the image
        * image : image to be scanned

        #### Returns:
        * scanned image
        """
        print(f"{image_name} loaded")

        # Scan the image isolating the grid
        Image_i = ImageScanner.scan(
            image,
            do_white_balance=do_white_balance 
        )
        print(f"{image_name} scanned!")

        # Resize image so that its height and width are the same
        w, h = Image_i.shape[:2]
        Image_i = cls.resize(Image_i, 1, w, w)

        return Image_i

    @staticmethod
    def resize(img: ndarray, factor: float, w: int = None, h: int = None):
        """
        ### Resize
        Resize image to a given percentage of current size.

        #### Args:
        * img : image to be resized
        * factor : percentage of current size to resize to
        * w : width of image
        * h : height of image

        #### Returns:
        * resized image
        """

        # If width and height are not given, get them from the image
        if w == None and h == None:
            w, h = img.shape[:2]

        resized_image = cv.resize(
            img, (int(w*factor), int(h*factor)), interpolation=cv.INTER_CUBIC)

        return resized_image

class ImageLoader:
    """
    ## ImageLoader

    This class is responsible for loading images from a given folder and converting HEIC images to JPG.

    ### Methods
    - `load_images(path_to_imgs: str) -> list`
        - This method loads all the images in a folder and returns a list of images.
    - `heic2jpg(path_to_heic: str) -> None`
        - This method creates .jpg images from the .HEIC images of given folder.

    ### Example
    ```python
    from src.objs.image.utils.image_loader import ImageLoader

    images = ImageLoader.load_images('path/to/images')
    or 
    images = ImageLoader.heic2jpg('path/to/heic')
    images = ImageLoader.load_images('path/to/images')
    ```
    """
    @staticmethod
    def load_images(path_to_imgs: str, return_paths_only:bool = False, display: int = 0):
        """
        ### Image loader
        Loads all the images in a folder and returns a list of images

        #### Args:
        path_to_images: path to image folder

        #### Returns:
        List of images
        """

        # acceptable image types
        types = ('.png', '.jpg', 'JPEG')

        # reading single image if path is only one image
        end = path_to_imgs[-4:]

        if end in types:
            return [cv.imread(path_to_imgs)]

        # reading all images of acceptable types from given directory
        imgs = []
        for f_type in types:
            files = [file for file in glob(f"{path_to_imgs}*{f_type}")]
            
            if return_paths_only:
                return files 

            if display:
                for i, f in enumerate(files):
                    name = f[f.rfind('\\') + 1:]
                    print(f"{i} -> {name}")

            imgs.extend([cv.imread(file) for file in files])

        return imgs

    @staticmethod
    def heic2png(path_to_heic: str):
        """
        ### HEIC to PNG converte
        Creates .png images from the .HEIC images of given folder.    

        #### Args:
        path_to_heic: path to image folder

        #### Returns:
        None
        """

        # finding all .HEIC images in the given folder
        # and converting them to .png
        paths = glob(f"{path_to_heic}*.HEIC")
        print(paths)
        for path in paths:
            pillow_heif.register_heif_opener()

            img = im.open(path)
            img.save(path[:-4] + 'png', format="png")
            print(f"{path} converted to PNG")

    @staticmethod
    def heic2jpg(path_to_heic: str):
        """
        ### HEIC to JPG converte
        Creates .jpg images from the .HEIC images of given folder.    

        #### Args:
        path_to_heic: path to image folder

        #### Returns:
        None
        """

        # finding all .HEIC images in the given folder
        # and converting them to .jpg
        paths = glob(f"{path_to_heic}*.HEIC")
        print(paths)
        for path in paths:
            pillow_heif.register_heif_opener()

            img = im.open(path)
            img.save(path[:-4] + 'jpg', format="jpeg")
            print(f"{path} converted to JPG")

class ImageScanner:
    """
    Class to scan the image and return the scanned image.

    ## Methods:
    - `scan(image_og: np.ndarray) -> np.ndarray`
        - This method scans the image and returns the scanned image.    

    - `morphological_transform(gpu_img: cv.cuda_GpuMat) -> cv.cuda_GpuMat`
            - This method applies morphological transformations to highlight the grid.

    - `remove_background(img: np.ndarray) -> np.ndarray`
            - This method gets rid of the background through masking + grabcut algorithm.

    - `find_contours(gpu_img: cv.cuda_GpuMat) -> list`
            - This method finds the contours of the image.

    - `detect_corners(contours: list, img: np.ndarray) -> list`
            - This method detects the corners of the grid.

    - `perspective_transform(img: np.ndarray, corners: list) -> np.ndarray`
            - This method applies perspective transform to the image.

    - `find_dest(pts: list) -> list`
            - This method finds the destination coordinates.

    - `order_points(pts: list) -> list`
            - This method orders the points.

    ## reference
        https://learnopencv.com/automatic-document-scanner-using-opencv/
    """

    @classmethod
    def scan(cls, img_og: np.ndarray, do_white_balance: bool = False) -> np.ndarray:
        # Applying morphological transformations to highlight the grid
        # Utilizing the GPU for faster processing
        img = cls.hsv_threshold(img_og.copy(), 100)

        morph_img = MorphologicalTransformer.apply_morph(img)

        # Isolate the grid by removing background (Only works with CPU)
        no_bkg_img = BackgroundRemover.remove_background(morph_img)

        # Adjusting the image to highlight the grid
        contours = ContourFinder.find_contours(no_bkg_img)

        """
        a = no_bkg_img.copy()
        cv.drawContours(a, contours, -1, (0, 255, 0), 3)
        display(a, 0)  # """
        corners = CornerDetector.detect_corners(contours, no_bkg_img)

        final_image = cls.perspective_transform(img_og, corners)

        if do_white_balance:
            final_image = WhiteBalanceAdjuster.adjust(final_image)

        return final_image
    # ----------------- Helper Functions ----------------- #

    @classmethod
    # function to turn everything that isnt kinda white to black
    def hsv_threshold(cls, img: np.ndarray, threshold: int) -> np.ndarray:
        # convert image to hsv
        hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
        # define range of white color in HSV
        lower_white = np.array([0, 0, 255-threshold])
        upper_white = np.array([255, threshold, 255])
        # create a mask
        mask = cv.inRange(hsv, lower_white, upper_white)
        # apply the mask to the image
        res = cv.bitwise_and(img, img, mask=mask)
        return res

    @classmethod
    def perspective_transform(cls, img: np.ndarray, corners: list) -> np.ndarray:
        # REARRANGING THE CORNERS
        destination_corners = cls.find_dest(corners)

        # Getting the homography. (aka scanning the image)
        M = cv.getPerspectiveTransform(np.float32(
            corners), np.float32(destination_corners))

        # Perspective transform using homography.
        final = cv.warpPerspective(
            img, M, (destination_corners[2][0], destination_corners[2][1]), flags=cv.INTER_LINEAR)

        return final

    @classmethod
    def find_dest(cls, pts: list) -> list:
        # DESTINATION COORDINATES
        (tl, tr, br, bl) = pts

        # Finding the maximum width.
        widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
        widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
        maxWidth = max(int(widthA), int(widthB))

        # Finding the maximum height.
        heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
        heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
        maxHeight = max(int(heightA), int(heightB))

        # Final destination co-ordinates.
        destination_corners = [[0, 0], [maxWidth, 0],
                               [maxWidth, maxHeight], [0, maxHeight]]
        return cls.order_points(destination_corners)

    @staticmethod
    def order_points(pts: list) -> list:
        # Initialising a list of coordinates that will be ordered.
        rect = np.zeros((4, 2), dtype='float32')
        pts = np.array(pts)
        s = pts.sum(axis=1)

        # Top-left point will have the smallest sum.
        rect[0] = pts[np.argmin(s)]

        # Bottom-right point will have the largest sum.
        rect[2] = pts[np.argmax(s)]

        # Computing the difference between the points.
        diff = np.diff(pts, axis=1)

        # Top-right point will have the smallest difference.
        rect[1] = pts[np.argmin(diff)]

        # Bottom-left will have the largest difference.
        rect[3] = pts[np.argmax(diff)]

        # Return the ordered coordinates.
        return rect.astype('int').tolist()

class WhiteBalanceAdjuster:
    """
    # WhiteBalanceAdjuster
    This class is responsible for adjusting the white balance of an image.

    ## Methods
    - `adjust(image: np.ndarray, reference_region: tuple[int, int, int, int] = (62, 80, 20, 20)) -> np.ndarray`
        - This method adjusts the white balance of the image.
    
    ### Example
    ```python
    import cv2 as cv
    import numpy as np
    from src.objs.image.utils.image_white_balancer import WhiteBalanceAdjuster

    scanned_image = cv.imread('path/to/image.jpg')
    adjusted_image = WhiteBalanceAdjuster.adjust(scanned_image)
    ```
    """

    @staticmethod
    def adjust( image: np.ndarray, 
                reference_region: tuple[int, int, int, int] = (62, 80, 20, 20)
            ) -> np.ndarray:
        """
        Adjust the white balance of the image.
        Args:
            image: The image to adjust.
            reference_region: The top-left coordinates and size of the reference region.
        Returns:
            The white-balanced image.
        """

        # Get the top-left coordinates and size of the reference region
        reference_top_left, reference_size = reference_region[:2], reference_region[2:]

        # Create the reference 10x10 square for the reference region for white balancing
        reference_region = image[reference_top_left[1]:reference_top_left[1] + reference_size[1],
                                reference_top_left[0]:reference_top_left[0] + reference_size[0]]

        # Calculate the mean RGB values of the reference region - image white baseline value
        mean_reference = np.mean(reference_region, axis=(0, 1))

        # Scaling factors for each channel
        scale_factors = 255.0 / mean_reference

        # Apply white balancing to the entire image by multiplying the image to the scale factor
        balanced_image = cv.merge([cv.multiply(image[:, :, i], scale_factors[i]) for i in range(3)])

        # Clip the values to the valid range [0, 255]
        balanced_image = np.clip(balanced_image, 0, 255).astype(np.uint8)
        
        return balanced_image