Module AmpliVision.src.objs

class Grid:
    def __init__(self, img: np.ndarray):
        # scanned image
        self.img = img.copy()

        # setup ratios used in the grid
        # such as the plus minus, etc.
        self.setup_ratios()

        # saving blocks here and in grid creates 2 sources of truth.
        # This list should keep index or something.
        self.blocks = []

        # represents the grid in the image as a 2D array of squares. Initialized as 2D array of None to represent empty.
        self.grid = [
            [None for _ in range(self.MAX_INDEX + 1)]
            for _ in range(self.MAX_INDEX + 1)
        ]

        self.create_grid()

    ## Setup functions ##
    def setup_ratios(self):
        """
        ### Setup rations
        Function that sets up the ratios used in the grid.
        """
        # max x and y coordinates of the image
        self.MAX_XY = self.img.shape[0]  # assumes image is square

        # ratios measured experimentally as a percentage of the grid
        # size of pin diameter/grid size
        self.PIN_RATIO = int(self.MAX_XY * 0.012)
        # size of edge/grid size. Edge is the "lines" around squares
        self.EDGE_RATIO = int(self.MAX_XY * 0.01)
        # squares are the places where you can insert the ampli blocks
        self.SQUARE_RATIO = int(self.MAX_XY * 0.089)
        # an arbitrary general tolerance
        self.PLUS_MINUS = int(self.MAX_XY * 0.005)
        self.SQUARE_LENGTH = self.SQUARE_RATIO + self.EDGE_RATIO
        self.MAX_INDEX = 9  # assumes grid is square 10x10

    def create_grid(self):
        """ 
        ### Create grid
        Function that creates the grid of squares in the image.
        """
        # These are the stop values for the for loops.
        STOP_XY = self.MAX_XY - self.EDGE_RATIO
        STEP = self.SQUARE_LENGTH

        # iterate through the grid by moving in steps of SQUARE_LENGTH
        # until the max x and y values are reached.
        # x and y are the top left points of the squares
        # x_index and y_index are the index of the square in the grid
        for y, x in itertools.product(range(0, STOP_XY, STEP), range(0, STOP_XY, STEP)):

            # get the corresponding index of the square in the grid
            # flipped coordinates to make it (row, column)
            y_index, x_index = Utils.xy_to_index(self, x, y)

            # coordinates of the top left and bottom right points of the square
            top_left = (
                x + (self.EDGE_RATIO),
                y + (self.EDGE_RATIO)
            )
            bottom_right = (
                x + self.SQUARE_RATIO + (self.EDGE_RATIO),
                y + self.SQUARE_RATIO + (self.EDGE_RATIO)
            )

            # create a square object sq
            sq = Square(
                top_left,
                bottom_right,
                (x_index, y_index),
                self.PIN_RATIO,
                self.PLUS_MINUS,
                self.img
            )

            # add the square to the grid list
            self.grid[x_index][y_index] = sq

    ### Find functions ###

    def find_pins(self, contours: list[np.ndarray]):
        """
        ### Find potential pins
        ---------------
        Function that finds the pins and adds them to their bounding squares.
        The pins are found by finding the square structures in the image.

        #### Args:
        * contours: list of contours around non-grayscale (colorful) edges in image
        """

        # Square structures are 4 points (in this case pins),
        # arranged in the shape of a square
        square_structures, pin_list = self.get_square_structures(contours)

        # adds the 4 potential pins structured as a square shape to the
        # square in the grid where the middle of the structure is located
        for square_structure, pins in zip(square_structures, pin_list):

            # get the middle of the structure
            center = find_center_of_points(square_structure)

            # get the index of sq based on the center of the structure
            # flipped coordinates to make it (row, column)
            y_index, x_index = Utils.xy_to_index(self, center[0], center[1])

            # add pins to the appropriate square in the grid
            for pin in pins:
                self.grid[x_index][y_index].add_pin(pin)

    def find_blocks(self, contours: list[np.ndarray]):
        """
        ### Find blocks
        ---------------
        Function that determines which squares are blocks in the grid.
        It does this by finding the potential pins (p_pins) 
        then checking if the pin is in the corners of the square and adding it to the square if it is.

        #### Args:
        * contours: list of contours around non-grayscale (colorful) edges in image

        #### Returns:
        None
        """

        # finds the potential pins (p_pins)
        # and adds them to their bounding squares.
        self.find_pins(contours)

        # checks if the potential pins are in one of the corners of square.
        # adds potential pin as a pin to the square if it is.
        self.process_pins()

        # checks if the square has x or more pins
        # if it does, it is considered a block.
        self.add_blocks()

    ## Helper functions ##
    def process_pins(self):
        """
        checks if the potential pins are in one of the corners of square.
        Adds potential pin as a pin to the square if it is.
        """
        for sq in itertools.chain(*self.grid):
            if len(sq.get_p_pins()) < 4:
                continue

            for p_pin in sq.get_p_pins():
                x, y, w, h = cv.boundingRect(p_pin)

                # checks if top left or bottom right point of pin
                # is inside corner of square within error range
                if sq.is_in_corners_skewed(x, y, w, h):
                    sq.add_pin(p_pin)

    ### Get functions ###

    def get_blocks(self):
        return self.blocks

    def get_contour_centers(self, contours: list[np.ndarray]):
        """
        Function that finds the center point of each contour

        #### Args:
        * contours: list of contours around non-grayscale (colorful) edges in image

        #### Returns:
        * center_to_contour_index: dictionary with center points as keys and their corresponding contour indices as values
        """
        center_to_contour_index = {}
        for i, contour in enumerate(contours):
            center = find_center_of_contour(contour)
            if center is not None:
                center_to_contour_index[center] = i

        # save the indexes bounding the centers of the
        # contours in a list and remove None values
        centers = list(center_to_contour_index.keys())
        centers = [x for x in centers if x != None]

        return center_to_contour_index, centers

    def get_square_structures(self, contours: list[np.ndarray]):
        """
        ### Square structures
        ---------------
        Function that finds the square structures in the image.  
        A square structure is defined as 4 points (in this case potential pins) arranged in the shape of a square.

        #### Args:
        * contours: list of contours around non-grayscale (colorful) edges in image

        #### Returns:
        * square_structures: list of square structures
        * p_pins: list of p_pins
        """
        square_structures = []
        pins = []

        # find the center of each contour
        center_to_contour_index, centers = self.get_contour_centers(contours)

        # Find all combinations of four points
        combinations = list(itertools.combinations(centers, 4))

        # useful values for debuggin
        """
        point0_step = math.comb(len(centers)-1, 3)
        point1_step = math.comb(len(centers)-2, 2)
        point2_step = math.comb(len(centers)-3, 1)"""
        # print("centers:", len(centers), "combinations:", len(combinations))
        # print("point0_step:", point0_step, "point1_step:", point1_step, "point2_step:", point2_step)
        index = 0
        debug_flag = 0
        #step_filter = point0_step

        # iterate through the combinations of points
        for comb in combinations:
            # (previously) missing block @ 6,4 in image 6066
            '''if index == 1171074: #(1179520 - point1_step - (point2_step*45) - 1):
                print("special: ", index)
                debug_flag = True

            else:
                debug_flag = False
            '''
            if is_arranged_as_square(comb, self.img, self.SQUARE_LENGTH, recursion_flag=0, debug=debug_flag):

                # Add the square to the list of
                # combinations if it is arranged as a square
                square_structures.append(list(comb))

                # Find the indices of the contours that form the square
                contour_indices = [center_to_contour_index[point]
                                   for point in comb]
                pins.append([contours[i] for i in contour_indices])

            index += 1

        return square_structures, pins

    ## Add functions ##

    def add_blocks(self):
        for sq in itertools.chain(*self.grid):
            if len(sq.get_pins()) >= 4:
                sq.is_block = True
                self.blocks.append(sq)

    def add_artificial_block(self, index: tuple[int, int], img, sq_img: np.ndarray):
        """
        ### Add artificial block
        ---------------
        Function that adds an artificial block to the grid.

        #### Args:
        * index: index of the square in the grid
        * sq_img: image of the square
        """

        # get the square in the grid
        i, j = index
        sq = self.grid[i][j]

        # No need to worry about all the square parameters,
        # After the colage, we can pass it through phase 1 again
        # sq.is_block = True, sq.block_type = ...
        return self.paste_block(img, sq_img, sq.tl, sq.br)

    ### Draw functions ###

    def paste_test_area(self, block):
        """ replaces the pixels at the block's location with the block's test area image """
        
        
        # get the top left and bottom right points of the block
        tl, br = block.tl, block.br

        # get the images
        img = self.img
        test_area_img = block.test_area_img

        # pixel coordinates
        corners = block.calculate_corners_pinbased()
        y_min = corners[0][1][1]
        y_max = corners[2][0][1]
        x_min = corners[0][1][0]
        x_max = corners[2][0][0]

        # paste the image of the block's test area on the grid        
        img[y_min:y_max, x_min:x_max] = test_area_img

        self.img = img

    def paste_block(self, img, sq_img, tl, br):
        " pastes the image of the block with transparent bkg on the grid "

        # paste the image of the square on the grid
        # at the top left and bottom right points of the square
        # tl and br are the top left and bottom right points of the square
        center_pt = self.calculate_center(tl, br)
        # sq_img = cv.resize(sq_img, sq_size)

        img = self.add_transparent_image(img, sq_img, center_pt)

        return img

    def add_transparent_image(self, background, foreground, center_pt):
        "Source: https://stackoverflow.com/questions/40895785/using-opencv-to-overlay-transparent-image-onto-another-image"

        bg_h, bg_w, bg_channels = background.shape
        fg_h, fg_w, fg_channels = foreground.shape

        assert bg_channels == 3, f"background image should have exactly 3 channels (RGB). found:{bg_channels}"
        assert fg_channels == 4, f"foreground image should have exactly 4 channels (RGBA). found:{fg_channels}"

        # center the foreground image on the background image according to the center point
        x_offset = center_pt[0] - fg_w // 2
        y_offset = center_pt[1] - fg_h // 2

        w = min(fg_w, bg_w, fg_w + x_offset, bg_w - x_offset)
        h = min(fg_h, bg_h, fg_h + y_offset, bg_h - y_offset)

        if w < 1 or h < 1:
            return

        # clip foreground and background images to the overlapping regions
        bg_x = max(0, x_offset)
        bg_y = max(0, y_offset)
        fg_x = max(0, x_offset * -1)
        fg_y = max(0, y_offset * -1)
        foreground = foreground[fg_y:fg_y + h, fg_x:fg_x + w]
        background_subsection = background[bg_y:bg_y + h, bg_x:bg_x + w]

        # separate alpha and color channels from the foreground image
        foreground_colors = foreground[:, :, :3]
        alpha_channel = foreground[:, :, 3] / 255  # 0-255 => 0.0-1.0

        # construct an alpha_mask that matches the image shape
        alpha_mask = np.dstack((alpha_channel, alpha_channel, alpha_channel))

        # combine the background with the overlay image weighted by alpha
        composite = background_subsection * \
            (1 - alpha_mask) + foreground_colors * alpha_mask

        # overwrite the section of the background image that has been updated
        background[bg_y:bg_y + h, bg_x:bg_x + w] = composite

        return background

    def draw_gridLines(self, img: np.ndarray):
        """
        ### draws grid lines
        ---------------
        Function that draws the grid lines on the image.
        """

        # draw grid lines
        start = 0 + self.EDGE_RATIO
        stop = self.MAX_XY - self.EDGE_RATIO + self.PLUS_MINUS
        step = self.SQUARE_RATIO + self.EDGE_RATIO

        for i, j in itertools.product(range(start, stop, step), repeat=2):
            # vertical lines
            cv.line(img, (i, 0), (i, self.MAX_XY), (0, 255, 0), 2)
            cv.line(img, (i + self.SQUARE_RATIO, 0),
                    (i + self.SQUARE_RATIO, self.MAX_XY), (0, 255, 0), 2)

            # horizontal lines
            cv.line(img, (0, i), (self.MAX_XY, i), (0, 255, 0), 2)
            cv.line(img, (0, i + self.SQUARE_RATIO),
                    (self.MAX_XY, i + self.SQUARE_RATIO), (0, 255, 0), 2)

    def draw_blocks(self, image_copy: np.ndarray, show_pins=False, show_corners=False):
        """
        Function that shows image with pins and corners drawn

        #### Args:
        * image_copy: copy of the original image
        * show_pins: boolean to show pins
        * show_corners: boolean to show corners
        """
        for blk in self.blocks:
            blk.draw_pins(image_copy) if show_pins else None
            blk.draw_corners(image_copy) if show_corners else None
            cv.rectangle(image_copy, blk.tl, blk.br, (0, 0, 255), 3)

    def calculate_center(self, lft_top, rgt_bot):
        """
        ### Calculate center
        ---------------
        Function that calculates the center of a square.

        #### Args:
        * lft_top: top left point of the square
        * rgt_bot: bottom right point of the square

        #### Returns:
        * center: center point of the square
        """
        x = (lft_top[0] + rgt_bot[0]) // 2
        y = (lft_top[1] + rgt_bot[1]) // 2
        return (x, y)

    def get_square(self, index):
        """
        ### Get square
        ---------------
        Function that gets the square at the given index.

        #### Args:
        * index: index of the square in the grid

        #### Returns:
        * square: square at the given index
        """
        return self.grid[index[0]][index[1]]

    def set_square(self, index, square):
        """
        ### Set square
        ---------------
        Function that sets the square at the given index.

        #### Args:
        * index: index of the square in the grid
        * square: square to set
        """
        self.grid[index[0]][index[1]] = square

class Square:
    """
    ### Square
    ---------------
    Class that represents a square in the grid_ds.

    #### Args:
    * tl: top left point of the square
    * br: bottom right point of the square
    * index: index of the square in the grid_ds

    #### Attributes:
    * tl: top left point of the square
    * br: bottom right point of the square
    * index: index of the square in the grid_ds
    * block: boolean that indicates if the square is a block
    * pin_count: number of pins in the square

    #### Methods:
    * add_pin: adds a pin to the square
    * draw_pins: draws the pins in the square
    * draw_corners: draws the corners of the square
    * createImg: creates an image of the square, a cutout of the image around the square
    * add_corners: adds the corners of the square to the square object
    * is_in_corners: checks if a point is in the corners of the square
    * which_corner_is_contour_in: finds which corner of square a contour is in
    * get_rgb_avg_of_contour: gets the average RGB of a contour in the image
    * get_pins_rgb: gets the average RGB of the pins in the square

    """

    def __init__(self, tl: int, br: int, index: Tuple, PIN_RATIO: int, PLUS_MINUS: int, img: np.ndarray) -> None:
        # potential pins
        self.p_pins = []

        # pins
        self.pins = []

        # block or not and type of block
        self.is_block = False
        self.block_type = ''

        # RBG values of the pins in the square (tl, tr, bl, br)
        self.rgb_sequence = []

        # coordinates and index in Grid
        self.tl = tl
        self.br = br
        self.index = index

        # image and image of the square for visualization if necessary
        self.img = img.copy()
        if img is not None:
            self.sq_img = self.createImg(img.copy())

        # corners of the square
        self.corners = []
        self.add_corners(PIN_RATIO, PLUS_MINUS)

        self.test_area_img = None

        # rotation of the block. 0 is vertical strip with bkg on bottom. 1 is horizontal strip with bkg on left side. etc.
        self.rotation = 0

        # ratios
        self.PIN_RATIO = PIN_RATIO
        self.PLUS_MINUS = PLUS_MINUS

    ## Get functions ##
    def get_index(self) -> Tuple:
        """ Returns the index of the square """
        return self.index

    def get_p_pins(self) -> list[int]:
        """ Returns the potential pins in the square """
        return self.p_pins

    def get_pins(self) -> list[int]:
        """ Returns the pins in the square """
        return self.pins

    def get_corners(self) -> list[int]:
        """ Returns the corners of the square """
        return self.corners

    def get_img(self) -> np.ndarray:
        """ Returns the image of the square """
        return self.img

    def get_sq_img(self) -> np.ndarray:
        """ Returns the image of the square """
        if self.sq_img is None:
            self.sq_img = self.createImg(self.img)
        return self.sq_img

    def get_test_area_img(self) -> np.ndarray:
        " Returns the image of squares test area (inner square where test strip can be)"
        if self.test_area_img is None:
            self.test_area_img = self.create_test_area_img(self.get_sq_img())
        return self.test_area_img

    def get_block_type(self) -> str:
        """ Returns the block type of the square """
        if self.is_block:
            return self.block_type
        else:
            return "Not a block"

    def get_rgb_sequence(self) -> list[int]:
        """ Returns the RGB sequence of the square """
        return self.rgb_sequence

    def createImg(self, img: np.ndarray) -> np.ndarray:
        """ Creates an image of the square, a cutout of the image around the square"""
        return img[(self.tl[1]-10):(self.br[1]+10), (self.tl[0]-10):(self.br[0]+10)]

    def create_test_area_img(self, sq_img: np.ndarray) -> np.ndarray:
        " Creates an image of the inner test spot"

        sq_img = self.img
        corners = self.calculate_corners_pinbased()

        """
        r = sq_img[a:b, c:d] means that the image r is a cutout of the image sq_img. 
        from the top left corner (a, c) to the bottom right corner (b, d).
        where a, b, c, d are the coordinates of the corners of the square.
        for example,
        a = corners[0][1][1] means that a is the y coordinate of the bottom right corner of the top right corner of the square.
        b = corners[2][0][1] means that b is the y coordinate of the top left corner of the bottom right corner of the square.
        c = corners[0][1][0] means that c is the x coordinate of the bottom right corner of the top right corner of the square.
        d = corners[2][0][0] means that d is the x coordinate of the top left corner of the bottom right corner of the square.
        """

        return sq_img[corners[0][1][1]:corners[2][0][1], corners[0][1][0]:corners[2][0][0]]

    ## Add functions ##
    def add_pin(self, pin: np.ndarray) -> None:
        """ Adds a pin to the square """
        self.pins.append(pin)

    def add_p_pin(self, pin: np.ndarray) -> None:
        """ Adds a potential pin to the square """
        self.p_pins.append(pin)

    def add_corners(self, PIN_RATIO: int, PLUS_MINUS: int, p: int = 3, a: float = 1.8) -> None:
        """ 
        Adds the corners of the square to the square object

        #### Args:
        * PIN_RATIO: ratio of the pin size to the square size
        * PLUS_MINUS: arbitrary tolerance value
        * p: "padding" value. Determines size of the corners.
        * a: skew value. Is the exponential determining how skewed the corners are.
        """

        # top left and bottom right coordinates of the square
        tl_x, tl_y = self.tl
        br_x, br_y = self.br

        # Skewing the corners in relation to the center of the grid to account for perspective.
        # the further away from the center, the more skewed the corners are (exponential).

        # Avoiding division by zero
        SKEW_x, SKEW_y = self.calculate_skew(a)

        # The following four values: top_right, top_left, bottom_right, bottom_left are the corners of the square.
        # Each corner contains its top left and bottom right coordinates.
        # Coordinates are calculated using:
        # top left and bottom right coordinates of the square, arbitrary plus minus value, the padding value and the skew value.

        self.corners = self.calculate_corners(
            tl_x, tl_y, br_x, br_y, PIN_RATIO, PLUS_MINUS, p, SKEW_x, SKEW_y)

    def calculate_skew(self, a: float) -> Tuple:
        """ 
        Calculates the skew originated from cellphone cameras

        |x-4|^a * (x-4)/|x-4|
        """
        if self.index[0] != 4:
            SKEW_x = int(
                (abs(self.index[0] - 4) ** a) * ((self.index[0] - 4) / abs(self.index[0] - 4)))
        else:
            SKEW_x = 0

        # Avoiding division by zero
        if self.index[1] != 4:
            SKEW_y = int(
                (abs(self.index[1] - 4) ** a) * ((self.index[1] - 4) / abs(self.index[1] - 4)))
        else:
            SKEW_y = 0

        return SKEW_x, SKEW_y

    def calculate_corners(self, tl_x: int, tl_y: int, br_x: int, br_y: int, PIN_RATIO: int, PLUS_MINUS: int, p: int, SKEW_x: int, SKEW_y: int) -> list[int]:
        """
        Calculates the corners of the square using magic. The "corners" here refer to the space in the ampli block where the pins are located.
        """
        top_right = (
            (tl_x - (p*PLUS_MINUS) + SKEW_x, tl_y - (p*PLUS_MINUS) + SKEW_y),
            (tl_x + PIN_RATIO + (p*PLUS_MINUS) + SKEW_x,
             tl_y + PIN_RATIO + (p*PLUS_MINUS) + SKEW_y)
        )

        top_left = (
            (br_x - PIN_RATIO - (p*PLUS_MINUS) +
             SKEW_x, tl_y - (p*PLUS_MINUS) + SKEW_y),
            (br_x + (p*PLUS_MINUS) + SKEW_x, tl_y +
             PIN_RATIO + (p*PLUS_MINUS) + SKEW_y)
        )

        bottom_right = (
            (tl_x - (p*PLUS_MINUS) + SKEW_x, br_y -
             PIN_RATIO - (p*PLUS_MINUS) + SKEW_y),
            (tl_x + PIN_RATIO+(p*PLUS_MINUS) +
             SKEW_x, br_y + (p*PLUS_MINUS) + SKEW_y)
        )

        bottom_left = (
            (br_x - PIN_RATIO - (p*PLUS_MINUS) + SKEW_x,
             br_y - PIN_RATIO - (p*PLUS_MINUS) + SKEW_y),
            (br_x + (p*PLUS_MINUS) + SKEW_x, br_y + (p*PLUS_MINUS) + SKEW_y)
        )
        return [top_right, top_left, bottom_right, bottom_left]

    def calculate_corners_pinbased(self) -> list[list[int]]:
        """
        Calculates the corners of the square based on the pins in the square.
        To be used after the pins have been added to the square.
        list[[corner_tl, corner_br], ...] in clockwise order starting from top left.
        """
        corners = []
        # pin is a list of contours
        for pin in self.pins:
            x, y, w, h = cv.boundingRect(pin)

            # add extra padding to the corners
            px, py = self.calculate_skew(0.2)
            px = int(px)
            py = int(py)

            # append top left and bottom right points of the test area
            corners.append([(x-px, y-py), (x+w+px, y+h+py)])

        return self.order_corner_points(corners)

    ## Drawing functions ##

    def draw_p_pins(self, image: np.ndarray) -> None:
        """ Draws the potential pins in the square """
        for pin in self.p_pins:
            cv.drawContours(image, pin, -1, (0, 255, 0), 3)

    def draw_pins(self, image: np.ndarray) -> None:
        """ Draws the pins in the square """
        for pin in self.pins:
            cv.drawContours(image, pin, -1, (0, 255, 0), 3)

    def draw_corners(self, img: np.ndarray) -> None:
        """ Draws the corners of the square """
        for corner in self.corners:
            cv.rectangle(img, corner[0], corner[1], (0, 0, 255), 1)

    def draw_corners_pinbased(self, img: np.ndarray) -> None:
        """ Draws the corners of the square based on the pins in the square
        To be used after the pins have been added to the square."""
        # pin is a list of contours
        for x, y in self.calculate_corners_pinbased():
            cv.rectangle(img, x, y, (0, 0, 255), 1)

    def draw_test_area(self, img: np.ndarray) -> None:
        "Draws the test area of the square"

        corners = self.calculate_corners_pinbased()
        cv.rectangle(img, corners[0][1], corners[2][0], (0, 0, 255), 1)

    ### Boolean functions ###
    def is_in_test_bounds(self, x: int, y: int) -> bool:
        "checks if coordinate is within test bounds (inner square where strip is)"

        pass

    def is_in_corners(self, x: int, y: int) -> bool:
        """ 
        Checks if a point is in the corners of the square. 
        """
        # corn = ["top_left", "top_right", "bottom_left", "bottom_right"]

        i = 0
        for corner in self.corners:
            if x >= corner[0][0] and x <= corner[1][0]:
                if y >= corner[0][1] and y <= corner[1][1]:
                    # print(corn[i], ": ", round(self.get_rgb_avg_of_contour(contour)))
                    return True
            i += 1

        return False

    def is_in_corners_skewed(self, x: int, y: int, w: float, h: float) -> bool:
        """ Checks if a point is in the corners of the square, 
        taking into consideration the skewing that happens."""
        return (self.is_in_corners(x, y) or
                self.is_in_corners(x+int(w), y+int(h)) or
                self.is_in_corners(x-int(w), y-int(h)) or
                self.is_in_corners(x+int(w), y-int(h)) or
                self.is_in_corners(x-int(w), y+int(h)))

    def which_corner_is_contour_in(self, contour: np.ndarray = None, xy=None) -> str:
        """
        Function that finds which corner of square a contour is in.
        """
        corn = ["top_left", "top_right", "bottom_left", "bottom_right"]

        if xy is None:
            x, y = cv.boundingRect(contour)[:2]
        else:
            x, y = xy
            x = int(x)
            y = int(y)

        i = 0
        for corner in self.corners:
            if x >= corner[0][0] and x <= corner[1][0]:
                if y >= corner[0][1] and y <= corner[1][1]:
                    return corn[i]
            i += 1

        # might be unecessary after corner skewing
        i = 0
        for corner in self.corners:
            if x + (2*self.PLUS_MINUS) >= corner[0][0] and x - (2*self.PLUS_MINUS) <= corner[1][0]:
                if y + (2*self.PLUS_MINUS) >= corner[0][1] and y - (2*self.PLUS_MINUS) <= corner[1][1]:
                    return corn[i]
            i += 1

    def order_corner_points(self, corners: list[int]) -> list[int]:
        """
        Orders the corners of the square in a clockwise manner starting from the top-left corner.
        """
        # top right, top left, bottom right, bottom left
        ordered_corners = [None, None, None, None]

        for xy in corners:
            mid_x = (xy[0][0] + xy[1][0]) / 2
            mid_y = (xy[0][1] + xy[1][1]) / 2
            s = self.which_corner_is_contour_in(xy=(mid_x,  mid_y))

            if s == "top_left":
                ordered_corners[0] = xy
            elif s == "top_right":
                ordered_corners[1] = xy
            elif s == "bottom_right":
                ordered_corners[2] = xy
            elif s == "bottom_left":
                ordered_corners[3] = xy

        if None in ordered_corners:
            print("\nError in ordering corners\n")
            return None

        return ordered_corners

    # set functions

    def set_rgb_sequence(self) -> None:
        """
        ### Set rgb sequence
        ---------------
        Function that sets the rgb sequence of the square.

        #### Returns:
        * None
        """

        # get the RGB values of the pins in the square
        pins_rgb, corner_key = get_pins_rgb(self)

        # fixing the order from tr,tl,br,bl to clockwise starting from top-right. This might be the ugliest code I've ever written. But it works!
        set_rgb_sequence_clockwise(self, pins_rgb, corner_key)

    def set_test_area_img(self, img):
        " Sets the image of the test area"
        self.test_area_img = img

class TestAnalyzer:
    "This class is responsible for getting and analyzing test results a.k.a phase B"

    def __init__(self, block):

        self.block = block

        # look only at the inner test square:
        self.test_square_img = block.get_test_area_img()

        # square used in csv export
        self.grid_index = block.index
        self.block_type = block.get_block_type()

        self.strip_sections = {
            "bkg": StripSection(self.test_square_img, 'bkg', block.rotation),
            "spot1": StripSection(self.test_square_img, 'spot1', block.rotation),
            "spot2": StripSection(self.test_square_img, 'spot2', block.rotation)
        }

    def analyze_test_result(self, double_thresh = False, display: bool = False):  # should I name it main?
        "gets test results from a block, analyses them, and export them to csv"

        # find the positive spots with hsv mask
        # need to think about cases where mask for example return one pixel.
        #   do you check for minimum contour size? do you only look for it manually? food for thought

        if display:
            print("rotation: ", self.block.rotation)

        # thresholds optimized for marker data
        rgb_spots = ColorContourExtractor.process_image(
            self.test_square_img, 
            hsv_lower=[0, 40, 20],
            double_thresh=double_thresh, 
            display=display
        )
        

        if display:
            cv.waitKey(100)
            plt.close()
            cv.destroyAllWindows()

        self.add_positives_to_sections(rgb_spots, display=display)

        # find the negative spots "manually" through ratios
        self.add_negatives_to_sections(display=display) 

        # get background color noise so we can remove it from other sections
        self.strip_sections['bkg'].set_total_avg_rgb()
        bkg_rgb_avg = self.strip_sections['bkg'].total_avg_rgb

        # remove background noise from other sections
        corrected_rgbs = []
        for section in self.strip_sections.values():
            if section.strip_type != 'bkg':

                if display:
                    print(f"{section.strip_type}AVG RGB: {section.total_avg_rgb}")
                    print("correcting: ", section.strip_type)

                corrected_rgbs.append(section.subtract_bkg(bkg_rgb_avg))

        if display:
            print("\n")

        # validate results to catch any potential errors in the test
        "TODO: adapt validate_results to work with the new strip configuration"
        # self.validate_results()

        # export results to csv
        row = self.create_csv_row(corrected_rgbs)
        return row
    

    def add_positives_to_sections(self, rgb_spots, display: int = 0) -> None:
        "used to add positive result spots to appropriate strip section"

        # adds each spot to its strip section
        for spot in rgb_spots:
            # display the spot
            cpy = cv.drawContours(self.test_square_img.copy(), [spot], -1, (0, 255, 0), 1)
            #plt.imshow(f'TA/add_positives_to_sections', cv.resize(cpy, (400, 400)))
            

            for section in self.strip_sections.values():
                if section.bounds_contour(spot):
                    section.add_spot(self.block, spot, True, debug=display)
                    # break # only adds to one section

    def add_negatives_to_sections(self, display: int = 0) -> None:
        "used to find negative result spots to appropriate strip section"
        for type, section in zip(self.strip_sections.keys(), self.strip_sections.values()):

            if len(section.spots) == 0:
                section.set_spots_manually(self.block, debug=display)

    def validate_results(self) -> None:
        "deals with test result potential positive, negative, false positive, error scenarios"

        results = self.get_section_results()

        # 1 test is properly positive (bkg, test, and spot2 line rgbs are > threshold)
        if results[1] & results[2]:
            print("Test worked properly and result is positive")

        # 2 test is properly negative (spot2 line rgb is > threshold)
        elif (not results[1]) & results[2]:
            print("Test worked properly and result is negative")

        # 3 spot2 error (bkg, and maybe test line rgbs are > threshold)
        else:
            print("Test may have not worked properly")

    def get_section_results(self) -> list[bool]:
        "returs a list of booleans representing the result (positive or negative) of each section bkg, test, spot2"
        results = []  # bkg, test, spot2
        for strip in self.strip_sections.values():
            strip_result = False

            # display
            # strip.print_spots()

            for spot in strip.spots:
                if spot["positive"] == True:
                    strip_result = True
                    break
            results.append(strip_result)

        print("\n")

        return results

    def create_csv_row(self, corrected_rgbs: list[list]) -> str:
        """
        writes the test results to csv file row in format:\n
        date, time, grid_index, block_type, bkg_r, bkg_g, bkg_b, test_r, test_g, test_b, cntrl_r, cntrl_g, cntrl_b"""

        # get current date and time
        now = datetime.now()

        # format date and time
        date = now.strftime("%m/%d/%Y")
        time = now.strftime("%H:%M:%S")

        # setting all rgb values to None
        bkg_r = " None"
        bkg_g, bkg_b = bkg_r, bkg_r
        spot1_r, spot1_g, spot1_b = bkg_r, bkg_r, bkg_r
        spot2_r, spot2_g, spot2_b = bkg_r, bkg_r, bkg_r

        # get rgb values of each section
        if self.strip_sections['bkg'].total_avg_rgb != None:
            bkg_b, bkg_g, bkg_r = self.strip_sections['bkg'].total_avg_rgb
            # print("bkg rgb: ", bkg_r, bkg_g, bkg_b)

        if self.strip_sections['spot1'].total_avg_rgb != None:
            spot1_b, spot1_g, spot1_r = self.strip_sections['spot1'].total_avg_rgb
            #print("spot1 rgb: ", test_r, test_g, test_b)

        if self.strip_sections["spot2"].total_avg_rgb != None:
            spot2_b, spot2_g, spot2_r = self.strip_sections['spot2'].total_avg_rgb

        spot1_corr_b, spot1_corr_g, spot1_corr_r = corrected_rgbs[0]
        spot2_corr_b, spot2_corr_g, spot2_corr_r = corrected_rgbs[1]

        # create data to be written to csv
        data = [
            date, time, self.grid_index, self.block_type,
            spot1_r, spot1_g, spot1_b,
            spot2_r, spot2_g, spot2_b,
            bkg_r, bkg_g, bkg_b,
            spot1_corr_r, spot1_corr_g, spot1_corr_b,
            spot2_corr_r, spot2_corr_g, spot2_corr_b
        ]

        return data

    def paint_spots(self, rgb_spot_results: dict[list]):
        """ colors the spots in the image accourding to the results 

        Args:
            rgb_spot_results (dict[list]): dictionary containing the rgb values of the spots in the image
            format: {r: [mean1, std1, mean2, std2], g: [mean1, std1, mean2, std2], b: [mean1, std1, mean2, std2]}
        """
        from numpy import random
        import time 

        # can show improved performance if get_rgb_avg_of_contour TODO is done
        t = time.time()
        self.analyze_test_result()
        #print(f"        analyze_test_result in {round(time.time() - t,2)}")

        image = self.test_square_img
        image_ = image.copy()
        for type, section in self.strip_sections.items():
            
            # dont paint bkg
            if section.strip_type == 'bkg':
                continue
            
            # get the correct index for result
            i = 0 if type == 'spot1' else 2

            rgb = []
            means = []
            for c in ('b', 'g', 'r'):
                mean, std = rgb_spot_results[c][i:i+2]
                means.append(mean)
                
                rgb.append(int(random.normal(mean, std)))
            rgb = tuple(rgb)
            
            if means == [0,0,0]:
                continue
            
            #print(f"Paiting {type} with {rgb}")
            image_ = section.paint_spot(image, rgb, display=False)
        
        self.block.set_test_area_img(image_)

        return self.block

class TestGraph:
    def __init__(self, blocks):
        self.graph = self.build_graph(blocks)

    def build_graph(self, blocks):
        G = nx.DiGraph()

        # add nodes (block types) to graph.
        nodes = [block.block_type for block in blocks]

        # make sure duplicate names are changed. Ex [wick, wick] -> [wick_1, wick_2]
        nodes = self.make_unique_names(nodes)
  
        G.add_nodes_from(nodes)
        # add edges (sequence of blocks) to graph
        positions = [block.index for block in blocks]
        
        try:
            head = nodes.index('sample_block')
        except ValueError:
            # block identification error
            return G
        
        edges = self.get_edges(nodes, positions, head)
        G.add_edges_from(edges)

        return G

    def get_edges(self, nodes: list, positions: list, head_node_index: int) -> list:
        """Finds the sequence of blocks. eg. sample -> test 1 -> test2 -> wick block"""

        # Initialize the search
        edges = []
        head_node = nodes.pop(head_node_index)
        head_pos = positions.pop(head_node_index)

        # Start the recursive search for edges
        edges.extend(self.search_edges(head_node, head_pos, nodes, positions))

        return edges

    def search_edges(self, current_node, current_pos, remaining_nodes, remaining_positions):
        edges = []
        y, x = current_pos

        # Define potential neighbors
        neighbors = [
            (y, x + 1),
            (y, x - 1),
            (y + 1, x),
            (y - 1, x)
        ]

        for neighbor in neighbors:

            # list the grid positions. Ex: (5, 2)

            if neighbor in remaining_positions:

                # append found edge
                neighbor_index = remaining_positions.index(neighbor)
                remaining_positions.remove(neighbor)

                neighbor_node = remaining_nodes.pop(neighbor_index)
                edges.append((current_node, neighbor_node))

                # recursively search for neighboring edges and add to list
                edges.extend(
                    self.search_edges(
                        neighbor_node, neighbor, remaining_nodes, remaining_positions
                    )
                )

        return edges

    def display(self):
        G = self.graph
        pos = nx.spring_layout(G)  # Positions for all nodes
        nx.draw(G, pos, with_labels=True, node_size=2000, node_color="skyblue",
                font_size=15, font_weight="bold", arrows=True)
        plt.show()

    def make_unique_names(self, nodes):
        count = {}
        unique_names = []
        for name in nodes:
            if name in count:
                count[name] += 1
                unique_name = f"{name}_{count[name]}"
            else:
                count[name] = 0
                unique_name = name
            unique_names.append(unique_name)
        return unique_names

    def example_usage(self):

        # add nodes
        sequence = [1, 2, 3, 4]
        G = nx.DiGraph()
        G.add_nodes_from(sequence)

        # add edges
        edges = [(1, 2), (2, 3), (1, 4)]
        G.add_edges_from(edges)

        # Draw the graph
        pos = nx.spring_layout(G)  # Positions for all nodes
        nx.draw(G, pos, with_labels=True, node_size=2000, node_color="skyblue",
                font_size=15, font_weight="bold", arrows=True)

        # how to get graph information
        print(G.nodes)
        print(G.out_edges)

        plt.show()