AzurLaneAutoScript/module/map/map_fleet_preparation.py

import numpy as np
from PIL import ImageStat
from scipy import signal

from module.base.base import ModuleBase
from module.base.button import Button
from module.base.timer import Timer
from module.base.utils import area_offset, color_similarity_2d, rgb2gray
from module.logger import logger
from module.map.assets import *


class FleetOperator:
    FLEET_BAR_SHAPE_Y = 33
    FLEET_BAR_MARGIN_Y = 9
    FLEET_BAR_ACTIVE_STD = 45  # Active: 67, inactive: 12.
    FLEET_IN_USE_STD = 27  # In use 52, not in use (3, 6).

    def __init__(self, choose, bar, clear, in_use, main):
        """
        Args:
            choose (Button): Button to activate or deactivate dropdown menu.
            bar (Button): Dropdown menu for fleet selection。
            clear (Button): Button to clear current fleet.
            in_use (Button): Button to detect if it's using current fleet.
            main (ModuleBase): Alas module.
        """
        self._choose = choose
        self._bar = bar
        self._clear = clear
        self._in_use = in_use
        self.main = main

    def __str__(self):
        return str(self._choose)[:-7]

    def parse_fleet_bar(self, image):
        """
        Args:
            image (PIL.Image.Image): Image of dropdown menu.

        Returns:
            list: List of int. Currently selected fleet ranges from 1 to 6.
        """
        result = []
        for index, y in enumerate(range(0, image.size[1], self.FLEET_BAR_SHAPE_Y + self.FLEET_BAR_MARGIN_Y)):
            area = (0, y, image.size[0], y + self.FLEET_BAR_SHAPE_Y)
            stat = ImageStat.Stat(image.crop(area))
            if np.std(stat.mean, ddof=1) > self.FLEET_BAR_ACTIVE_STD:
                result.append(index + 1)
        logger.info('Current selected: %s' % str(result))
        return result

    def get_button(self, index):
        """
        Convert fleet index to the Button object on dropdown menu.

        Args:
            index (int): Fleet index, 1-6.

        Returns:
            Button: Button instance.
        """
        area = area_offset(area=(
            0,
            (self.FLEET_BAR_SHAPE_Y + self.FLEET_BAR_MARGIN_Y) * (index - 1),
            self._bar.area[2] - self._bar.area[0],
            (self.FLEET_BAR_SHAPE_Y + self.FLEET_BAR_MARGIN_Y) * (index - 1) + self.FLEET_BAR_SHAPE_Y
        ), offset=(self._bar.area[0:2]))
        return Button(area=(), color=(), button=area, name='%s_INDEX_%s' % (str(self._bar), str(index)))

    def allow(self):
        """
        Returns:
            bool: If current fleet is allow to be chosen.
        """
        return self.main.appear(self._choose)

    def clear(self, skip_first_screenshot=True):
        """
        Clear chosen fleet.
        """
        main = self.main
        click_timer = Timer(3, count=6)
        while 1:
            if skip_first_screenshot:
                skip_first_screenshot = False
            else:
                main.device.screenshot()

            # End
            if not self.in_use():
                break

            # Click
            if click_timer.reached():
                main.device.click(self._clear)
                click_timer.reset()

    def open(self, skip_first_screenshot=True):
        """
        Activate dropdown menu for fleet selection.
        """
        main = self.main
        click_timer = Timer(3, count=6)
        while 1:
            if skip_first_screenshot:
                skip_first_screenshot = False
            else:
                main.device.screenshot()

            # End
            if self.bar_opened():
                break

            # Click
            if click_timer.reached():
                main.device.click(self._choose)
                click_timer.reset()

    def close(self, skip_first_screenshot=True):
        """
        Deactivate dropdown menu for fleet selection.
        """
        main = self.main
        click_timer = Timer(3, count=6)
        while 1:
            if skip_first_screenshot:
                skip_first_screenshot = False
            else:
                main.device.screenshot()

            # End
            if not self.bar_opened():
                break

            # Click
            if click_timer.reached():
                main.device.click(self._choose)
                click_timer.reset()

    def click(self, index, skip_first_screenshot=True):
        """
        Choose a fleet on dropdown menu, and dropdown deactivated.

        Args:
            index (int): Fleet index, 1-6.
            skip_first_screenshot (bool):
        """
        main = self.main
        button = self.get_button(index)
        click_timer = Timer(3, count=6)
        while 1:
            if skip_first_screenshot:
                skip_first_screenshot = False
            else:
                main.device.screenshot()

            if not self.bar_opened():
                # End
                if self.in_use():
                    break
                else:
                    self.open()

            # Click
            if click_timer.reached():
                main.device.click(button)
                click_timer.reset()

    def selected(self):
        """
        Returns:
            list: List of int. Currently selected fleet ranges from 1 to 6.
        """
        data = self.parse_fleet_bar(self.main.device.image.crop(self._bar.area))
        return data

    def in_use(self):
        """
        Returns:
            bool: If has selected to any fleet.
        """
        # Handle the info bar of auto search info.
        # if area_cross_area(self._in_use.area, INFO_BAR_1.area):
        #     self.main.handle_info_bar()

        # Cropping FLEET_*_IN_USE to avoid detecting info_bar, also do the trick.
        # It also avoids wasting time on handling the info_bar.
        image = rgb2gray(np.array(self.main.image_area(self._in_use)))
        return np.std(image.flatten(), ddof=1) > self.FLEET_IN_USE_STD

    def bar_opened(self):
        """
        Returns:
            bool: If dropdown menu appears.
        """
        # Check the brightness of the rightest column of the bar area.
        luma = rgb2gray(np.array(self.main.image_area(self._bar)))[:, -1]
        return np.sum(luma > 127) / luma.size > 0.5

    def ensure_to_be(self, index):
        """
        Set to a specific fleet.

        Args:
            index (int): Fleet index, 1-6.
        """
        self.open()
        if index in self.selected():
            self.close()
        else:
            self.click(index)


class FleetPreparation(ModuleBase):
    map_fleet_checked = False
    map_is_hard_mode = False

    def get_map_is_hard_mode(self):
        """
        Detect how many light orange lines are there.
        Having lines means current map has stat limits and user has satisfied at least one of them,
        so this is a hard map.

        Returns:
            bool:
        """
        area = (208, 130, 226, 551)
        image = color_similarity_2d(self.image_area(area), color=(249, 199, 0))
        height = np.max(image, axis=1)
        parameters = {'height': 180, 'distance': 5}
        peaks, _ = signal.find_peaks(height, **parameters)
        lines = len(peaks)
        logger.attr('Light_orange_line', lines)
        return lines > 0

    def fleet_preparation(self):
        """Change fleets.

        Returns:
            bool: True if changed.
        """
        logger.info(f'Using fleet: {[self.config.Fleet_Fleet1, self.config.Fleet_Fleet2, self.config.Submarine_Fleet]}')
        if self.map_fleet_checked:
            return False

        self.map_is_hard_mode = self.get_map_is_hard_mode()
        if self.map_is_hard_mode:
            logger.info('Hard Campaign. No fleet preparation')
            return False

        fleet_1 = FleetOperator(
            choose=FLEET_1_CHOOSE, bar=FLEET_1_BAR, clear=FLEET_1_CLEAR, in_use=FLEET_1_IN_USE, main=self)
        fleet_2 = FleetOperator(
            choose=FLEET_2_CHOOSE, bar=FLEET_2_BAR, clear=FLEET_2_CLEAR, in_use=FLEET_2_IN_USE, main=self)
        submarine = FleetOperator(
            choose=SUBMARINE_CHOOSE, bar=SUBMARINE_BAR, clear=SUBMARINE_CLEAR, in_use=SUBMARINE_IN_USE, main=self)

        # Submarine.
        if submarine.allow():
            if self.config.Submarine_Fleet:
                submarine.ensure_to_be(self.config.Submarine_Fleet)
            else:
                submarine.clear()

        # No need, this may clear FLEET_2 by mistake, clear FLEET_2 in map config.
        # if not fleet_2.allow():
        #     self.config.FLEET_2 = 0

        if self.config.Fleet_Fleet2:
            # Using both fleets.
            # Force to set it again.
            # Fleets may reversed, because AL no longer treat the fleet with smaller index as first fleet
            fleet_2.clear()
            fleet_1.ensure_to_be(self.config.Fleet_Fleet1)
            fleet_2.ensure_to_be(self.config.Fleet_Fleet2)
        else:
            # Not using fleet 2.
            if fleet_2.allow():
                fleet_2.clear()
            fleet_1.ensure_to_be(self.config.Fleet_Fleet1)

        # Check if submarine is empty again.
        if submarine.allow():
            if self.config.Submarine_Fleet:
                pass
            else:
                submarine.clear()

        self.map_fleet_checked = True
        return True