Add: class NestedNavbar for island shop

module/island/ui.py

@@ -1,16 +1,219 @@
 import cv2
 import numpy as np
+from typing import List
 
+from module.base.base import ModuleBase
 from module.base.button import ButtonGrid
 from module.base.decorator import cached_property
-from module.base.utils import color_similarity_2d, rgb2luma
+from module.base.timer import Timer
+from module.base.utils import area_offset, color_similarity_2d, rgb2luma
 from module.combat.assets import GET_ITEMS_1
 from module.island.assets import *
+from module.logger import logger
 from module.ui.navbar import Navbar
 from module.ui.ui import UI
 from module.ui_white.assets import BACK_ARROW_WHITE, POPUP_CONFIRM_WHITE_ISLAND
 
 
+class NestedNavbar:
+    # Navbar with level 2 option buttons which will span and make level 1 buttons move apart.
+    # For example if level 1 has 5 buttons with submenu of options 0, 2, 1, 2, 1, then when the 2nd submenu of the second button is chosen, it will be
+    # 1
+    # [[2]]
+    # 2.1
+    # [2.2]
+    # 3
+    # 4
+    # 5
+    # When the 1st submenu of the 3rd button is chosen, it will be
+    # 1
+    # 2
+    # [[3]]
+    # [3.1]
+    # 4
+    # 5
+    # Since the level 1 buttons will move, before clicking the level 1 button we will need to check the current acvive level 1 button and the number of active submenu options to calculate the position of the level 1 button to click.
+    def __init__(self, grids: ButtonGrid,
+                 subgrid_delta: tuple = None, subgrid_button_shape: tuple = None,
+                 subgrid_shapes: List[tuple] = None, direction: str = 'vertical',
+                 main_active_color=(57, 189, 255), main_inactive_color=(38, 39, 40),
+                 main_active_threshold=221, main_inactive_threshold=221,
+                 main_active_count=2000, main_inactive_count=2000,
+                 sub_active_color=(125, 126, 126), sub_inactive_color=(38, 39, 40),
+                 sub_active_threshold=221, sub_inactive_threshold=221,
+                 sub_active_count=500, sub_inactive_count=500, name: str = None):
+        """
+        Parameters:
+            grids (ButtonGrid): The ButtonGrid instance for the main level buttons.
+            subbutton (Button): The Button instance for the submenu options.
+            subgrid_shapes (list of tuples): A list of grid shapes for each submenu. Each tuple represents the grid shape (rows, cols) for the corresponding submenu.
+            main_active_color (tuple): The RGB color for active main buttons.
+            main_inactive_color (tuple): The RGB color for inactive main buttons.
+            main_active_threshold (int): The threshold for detecting active main buttons.
+            main_inactive_threshold (int): The threshold for detecting inactive main buttons.
+            main_active_count (int): The minimum count of pixels to confirm an active main button.
+            main_inactive_count (int): The maximum count of pixels to confirm an inactive main button.
+            sub_active_color (tuple): The RGB color for active submenu options.
+            sub_inactive_color (tuple): The RGB color for inactive submenu options.
+            sub_active_threshold (int): The threshold for detecting active submenu options.
+            sub_inactive_threshold (int): The threshold for detecting inactive submenu options.
+            sub_active_count (int): The minimum count of pixels to confirm an active submenu option.
+            sub_inactive_count (int): The maximum count of pixels to confirm an inactive submenu option.
+        """
+        self.grids = grids
+        self.subgrid_delta = subgrid_delta if subgrid_delta is not None else grids.delta
+        self.subgrid_button_shape = subgrid_button_shape if subgrid_button_shape is not None else grids.button_shape
+        self.subgrid_shapes = subgrid_shapes if subgrid_shapes is not None else [(0, 0)] * len(grids.buttons)
+        self.direction = direction
+        self.main_active_color = main_active_color
+        self.main_inactive_color = main_inactive_color
+        self.main_active_threshold = main_active_threshold
+        self.main_inactive_threshold = main_inactive_threshold
+        self.main_active_count = main_active_count
+        self.main_inactive_count = main_inactive_count
+        self.sub_active_color = sub_active_color
+        self.sub_inactive_color = sub_inactive_color
+        self.sub_active_threshold = sub_active_threshold
+        self.sub_inactive_threshold = sub_inactive_threshold
+        self.sub_active_count = sub_active_count
+        self.sub_inactive_count = sub_inactive_count
+        self.name = name if name is not None else self.grids._name
+        self._construct_subgrids()
+
+    def _construct_subgrids(self):
+        self.subgrids = []
+        for idx, (main_button, subgrid_shape) in enumerate(zip(self.grids.buttons, self.subgrid_shapes)):
+            if self.direction == 'vertical':
+                subgrid_origin = (main_button.area[0], main_button.area[1] + self.grids.delta[1])
+            else:
+                subgrid_origin = (main_button.area[0] + self.grids.delta[0], main_button.area[1])
+            subgrid = ButtonGrid(origin=subgrid_origin, delta=self.subgrid_delta,
+                                 button_shape=self.subgrid_button_shape, grid_shape=subgrid_shape,
+                                 name=f'{self.name}_SUB_GRID_{idx}')
+            self.subgrids.append(subgrid)
+
+    def is_main_button_active(self, button, main):
+        return main.image_color_count(
+                    button, color=self.main_active_color, threshold=self.main_active_threshold, count=self.main_active_count)
+
+    def is_main_button_inactive(self, button, main):
+        return main.image_color_count(
+                    button, color=self.main_inactive_color, threshold=self.main_inactive_threshold, count=self.main_inactive_count)
+
+    def is_subbutton_active(self, button, main):
+        return main.image_color_count(
+                    button, color=self.sub_active_color, threshold=self.sub_active_threshold, count=self.sub_active_count)
+
+    def is_subbutton_inactive(self, button, main):
+        return main.image_color_count(
+                    button, color=self.sub_inactive_color, threshold=self.sub_inactive_threshold, count=self.sub_inactive_count)
+
+    def get_info(self, main):
+        """
+        Get the index of the active main button and the active submenu option.
+        Parameters:
+            main (ModuleBase): The main module instance to capture screenshots.
+        Returns:
+            tuple: A tuple containing the index of the active main button and the index of the active submenu option. If no submenu is active, the second element will be None.
+        """
+        total_main = []
+        active_main = []
+        total_sub = []
+        active_sub = []
+        current_offset = (0, 0)
+        for index, button in enumerate(self.grids.buttons):
+            button._button_offset = area_offset(button._button, offset=current_offset)
+            if self.is_main_button_active(button, main=main):
+                total_main.append(index)
+                active_main.append(index)
+                subgrid = self.subgrids[index]
+                for sub_index, sub_button in enumerate(subgrid.buttons):
+                    if self.is_subbutton_active(sub_button, main=main):
+                        total_sub.append((index, sub_index))
+                        active_sub.append((index, sub_index))
+                    elif self.is_subbutton_inactive(sub_button, main=main):
+                        total_sub.append((index, sub_index))
+                if self.direction == 'vertical':
+                    current_offset = (
+                        current_offset[0],
+                        current_offset[1] + self.subgrid_button_shape[1] * subgrid.grid_shape[1],
+                    )
+                else:
+                    current_offset = (
+                        current_offset[0] + self.subgrid_button_shape[0] * subgrid.grid_shape[0],
+                        current_offset[1],
+                    )
+            elif self.is_main_button_inactive(button, main=main):
+                total_main.append(index)
+
+        if len(active_main) == 0:
+            active_main_index = None
+        elif len(active_main) == 1:
+            active_main_index = active_main[0]
+        else:
+            logger.warning(f'Too many active nav items found in {self.name}, items: {active_main}')
+            active_main_index = active_main[0]
+
+        active_sub_index = None
+        if active_main_index is not None:
+            active_subs_for_active_main = [sub for sub in active_sub if sub[0] == active_main_index]
+            if len(active_subs_for_active_main) == 0:
+                active_sub_index = None
+            elif len(active_subs_for_active_main) == 1:
+                active_sub_index = active_subs_for_active_main[0][1]
+            else:
+                logger.warning(f'Too many active sub nav items found in {self.name} for main index {active_main_index}, items: {active_subs_for_active_main}')
+                active_sub_index = active_subs_for_active_main[0][1]
+
+        if len(total_main) < 2:
+            logger.warning(f'Too few nav items found in {self.name}, items: {total_main}')
+        if len(total_main) == 0:
+            main_begin, main_end = None, None
+        else:
+            main_begin, main_end = total_main[0], total_main[-1]
+
+        return active_main_index, active_sub_index, main_begin, main_end
+
+    def set(self, main: ModuleBase, main_index: int, sub_index: int = None, skip_first_screenshot: bool = False):
+        """
+        Click the main button and submenu option based on the provided indices.
+        Should be used after calling get_info to get the current active main and submenu indices to calculate the position to click.
+        Parameters:
+            main (ModuleBase): The main module instance to perform clicks and capture screenshots.
+            main_index (int): The index of the main button to click. Starts from 0.
+            sub_index (int, optional): The index of the submenu option to click. If None, no submenu option will be clicked. Starts from 0.
+            skip_first_screenshot (bool):
+        """
+        logger.info(f'Setting {self.name} to main index {main_index} and sub index {sub_index}')
+        click_timer = Timer(2, count=4).reset()
+        confirm_timer = Timer(2, count=4).reset()
+        for _ in main.loop(skip_first=skip_first_screenshot, timeout=10):
+            active_main_index, active_sub_index, main_begin, main_end = self.get_info(main)
+            if active_main_index == main_index and (sub_index is None or active_sub_index == sub_index):
+                if confirm_timer.reached():
+                    logger.info(f'Successfully set {self.name} to main index {main_index} and sub index {sub_index}')
+                    return True
+                continue
+            if active_main_index is None or main_begin is None or main_end is None:
+                confirm_timer.reset()
+                continue
+            if not (main_begin <= main_index <= main_end):
+                logger.warning(f'Main index {main_index} out of range for {self.name}, appears ({main_begin}, {main_end})')
+                continue
+            if click_timer.reached_and_reset():
+                if active_main_index != main_index:
+                    main.device.click(self.grids.buttons[main_index])
+                    confirm_timer.reset()
+                else:
+                    subgrid = self.subgrids[main_index]
+                    if sub_index is not None:
+                        main.device.click(subgrid.buttons[sub_index])
+                        confirm_timer.reset()
+        else:
+            logger.warning(f'Failed to set {self.name} to main index {main_index} and sub index {sub_index}')
+            return False
+
+
 class IslandUI(UI):
     def ui_back(self, check_button, appear_button=None, offset=(30, 30), retry_wait=10, skip_first_screenshot=False):
         return self.ui_click(
@@ -120,4 +323,4 @@ class IslandUI(UI):
                 2 for restaurant,
                 3 for production
         """
-        return self._island_manage_side_navbar.set(self, upper=upper, skip_first_screenshot=skip_first_screenshot)
+        return self._island_manage_side_navbar.set(self, upper=upper, skip_first_screenshot=skip_first_screenshot)