AzurLaneAutoScript/module/campaign/campaign_status.py

import re

import cv2
import numpy as np

import module.config.server as server

from module.base.timer import Timer
from module.base.utils import color_similar, get_color
from module.campaign.assets import OCR_COIN, OCR_EVENT_PT, OCR_OIL, OCR_COIN_LIMIT, OCR_OIL_LIMIT, OCR_OIL_CHECK
from module.log_res.log_res import LogRes
from module.logger import logger
from module.ocr.ocr import Digit, Ocr
from module.ui.ui import UI

if server.server != 'jp':
    OCR_COIN = Digit(OCR_COIN, name='OCR_COIN', letter=(239, 239, 239), threshold=128)
else:
    OCR_COIN = Digit(OCR_COIN, name='OCR_COIN', letter=(201, 201, 201), threshold=128)

class PtOcr(Ocr):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, lang='azur_lane', alphabet='X0123456789', **kwargs)

    def pre_process(self, image):
        """
        Args:
            image (np.ndarray): Shape (height, width, channel)

        Returns:
            np.ndarray: Shape (width, height)
        """
        # Use MAX(r, g, b)
        r, g, b = cv2.split(cv2.subtract((255, 255, 255, 0), image))
        image = cv2.min(cv2.min(r, g), b)
        # Remove background, 0-192 => 0-255
        image = cv2.multiply(image, 255 / 192)

        return image.astype(np.uint8)


OCR_PT = PtOcr(OCR_EVENT_PT)


class CampaignStatus(UI):
    def get_event_pt(self, update=False):
        """
        Returns:
            int: PT amount, or 0 if unable to parse
        """
        pt = OCR_PT.ocr(self.device.image)

        res = re.search(r'X(\d+)', pt)
        if res:
            pt = int(res.group(1))
            logger.attr('Event_PT', pt)
            LogRes(self.config).Pt = pt
        else:
            logger.warning(f'Invalid pt result: {pt}')
            pt = 0
        if update:
            self.config.update()
        return pt

    def _get_num(self, _button, name, letter):
        # Update offset
        _ = self.appear(OCR_OIL_CHECK)

        color = get_color(self.device.image, OCR_OIL_CHECK.button)
        if color_similar(color, OCR_OIL_CHECK.color):
            # Original color
            if isinstance(_button, Ocr):
                ocr = _button
            else:
                if server.server != 'jp':
                    ocr = Digit(_button, name=name, letter=letter, threshold=128)
                else:
                    ocr = Digit(_button, name=name, letter=(201, 201, 201), threshold=128)
        elif color_similar(color, (59, 59, 64)):
            # With black overlay
            if isinstance(_button, Ocr):
                ocr = Digit(_button.buttons, name=name, letter=(165, 165, 165), threshold=128)
            else:
                ocr = Digit(_button, name=name, letter=(165, 165, 165), threshold=128)
        else:
            logger.warning(f'Unexpected OCR_OIL_CHECK color')
            if isinstance(_button, Ocr):
                ocr = Digit(_button.buttons, name=name, letter=(247, 247, 247), threshold=128)
            else:
                ocr = Digit(_button, name=name, letter=(247, 247, 247), threshold=128)

        return ocr.ocr(self.device.image)

    def _get_oil_limit(self):
        if not hasattr(self, '_oil_limit_cache'):
            limit = self._get_num(OCR_OIL_LIMIT, 'OCR_OIL_LIMIT', (247, 247, 247))
            if 1000 < limit <= 25000:
                self._oil_limit_cache = limit
            else:
                logger.warning(f'Invalid oil limit: {limit}')
                return 0
        return self._oil_limit_cache

    def _get_coin_limit(self):
        if not hasattr(self, '_coin_limit_cache'):
            limit = self._get_num(OCR_COIN_LIMIT, 'OCR_COIN_LIMIT', (239, 239, 239))
            if 10000 < limit <= 150000:
                self._coin_limit_cache = limit
            else:
                logger.warning(f'Invalid coin limit: {limit}')
                return 0
        return self._coin_limit_cache

    def get_oil(self, skip_first_screenshot=True, update=False):
        """
        Returns:
            int: Oil amount
        """
        oil = 0
        last_oil = 0
        timeout = Timer(1, count=2).start()
        while 1:
            if skip_first_screenshot:
                skip_first_screenshot = False
            else:
                self.device.screenshot()

            if not self.appear(OCR_OIL_CHECK, offset=(10, 2)):
                logger.info('No oil icon')
                continue

            if timeout.reached():
                logger.warning('Get oil timeout')
                oil = 0
                break

            oil = self._get_num(OCR_OIL, 'OCR_OIL', (247, 247, 247))
            if last_oil == oil and oil >= 100:
                break
            last_oil = oil

        if oil:
            LogRes(self.config).Oil = {'Value': oil, 'Limit': self._get_oil_limit()}
        if update:
            self.config.update()

        return oil

    def get_oil_and_coin(self, skip_first_screenshot=True, update=False):
        """
        Read oil and coin sharing a single screenshot.

        Returns:
            tuple[int, int]: (oil, coin)
        """
        oil = self.get_oil(skip_first_screenshot=skip_first_screenshot, update=False)
        coin = 0
        if oil:
            coin = self._get_num(OCR_COIN, 'OCR_COIN', (239, 239, 239))
            LogRes(self.config).Coin = {'Value': coin, 'Limit': self._get_coin_limit()}
        if update:
            self.config.update()
        return oil, coin

    def is_balancer_task(self):
        """
        Returns:
             bool: If is event task but not daily event task
        """
        tasks = [
            'Event',
            'Event2',
            'Raid',
            'Coalition',
            'GemsFarming',
        ]
        command = self.config.Scheduler_Command
        if command in tasks:
            if self.config.Campaign_Event == 'campaign_main':
                return False
            else:
                return True
        else:
            return False