import cv2

from mower.solvers.navigation import NavigationSolver
from mower.utils import config
from mower.utils.graph import SceneGraphSolver
from mower.utils.image import cropimg
from mower.utils.log import logger
from mower.utils.recognize import Scene


class ChooseSquadSolver(SceneGraphSolver):
    "选择编队"

    def run(self, level_name, squad: int):
        logger.info("Start: 选择编队")
        self.level_name = level_name
        self.squad = squad
        super().run()

    def current_squad(self):
        count = []
        for i in range(4):
            hsv = cropimg(
                config.recog.hsv, ((153 + i * 411, 990), (550 + i * 411, (1080)))
            )
            mask = cv2.inRange(hsv, (97, 0, 0), (101, 255, 255))
            count.append(cv2.countNonZero(mask))
        return count.index(max(count)) + 1

    def transition(self):
        if (scene := self.scene()) == Scene.OPERATOR_SELECT:
            if self.current_squad() != self.squad:
                self.tap((self.squad * 411 - 99, 1040))
            else:
                return True
        elif scene in self.waiting_scene:
            self.waiting_solver()
        else:
            NavigationSolver().run(self.level_name, mode="copy")
