cocostrat/app/ui/map_overlay.py

from itertools import product
import numpy as np
from app.verkolst import verkolst
from app.colormap.colormap import apply_colormap
from app.trade.trade_utils import trade_distance


class Overlay:

    def __init__(self, ui, hexmap_layer, opacity=120):
        self._ui = ui
        self._overlay = hexmap_layer
        self._enabled = False
        self._opacity = opacity

    def enable(self):
        self._enabled = True
        self.update()

    def disable(self):
        self._enabled = False

    def toggle(self):
        if self._enabled:
            self.disable()
        else:
            self.enable()

    @property
    def enabled(self):
        return self._enabled

    def _recalculate(self):
        pass

    def update(self):
        if not self._enabled:
            return
        self._recalculate()
        for i, j in product(range(verkolst.width), range(verkolst.height)):
            self._overlay.set_cell_color(i, j, self.get_cell_color(i, j))
            self._overlay.set_cell_opacity(i, j, self.get_cell_opacity(i, j))

    def get_cell_color(self, i, j):
        raise NotImplementedError()

    def get_cell_opacity(self, i, j):
        return self._opacity


class FoodOverlay(Overlay):

    def __init__(self, ui, hexmap_layer):
        super().__init__(ui, hexmap_layer, opacity=200)
        self._colors = None

    def _recalculate(self):
        self._colors = apply_colormap(verkolst.resources.food_stored/verkolst.resources.food_stored_capacity,
                                      'magma', min_=0)

    def get_cell_color(self, i, j):
        # res = verkolst.resources
        if verkolst.resources.any_starvation[i, j]:
            return 255, 0, 0  # Red
        else:
            return self._colors[i, j]


class PopsOverlay(Overlay):

    def __init__(self, ui, hexmap_layer, class_):
        super().__init__(ui, hexmap_layer, opacity=240)
        self._class = class_
        self._colors = None

    def _recalculate(self):
        pops = {
            'lc': verkolst.resources.lc_pops,
            'mc': verkolst.resources.mc_pops,
            'uc': verkolst.resources.uc_pops,
        }[self._class]
        self._colors = apply_colormap(pops, 'magma', min_=0)

    def get_cell_color(self, i, j):
        return self._colors[i, j]


class TradeDistanceOverlay(Overlay):

    def __init__(self, ui, hexmap_layer):
        super().__init__(ui, hexmap_layer, opacity=180)
        self._colors = None
        self.trade_distance = np.full_like(verkolst.resources.lc_pops, 0)
        self.exp_partner = None

    def _recalculate(self):
        if self._ui.hl_province is not None:
            i, j = self._ui.hl_province.i, self._ui.hl_province.j
            self.exp_partner, self.trade_distance = trade_distance(i, j,
                                                                   verkolst.resources.trade_range[i, j],
                                                                   verkolst.resources.trade_impedance,
                                                                   verkolst.resources.trade_value)
        else:
            self.trade_distance[:, :] = 0.0
        self._colors = apply_colormap(self.trade_distance, 'viridis', min_=0)

    def get_cell_color(self, i, j):
        if self._ui.hl_province is None:
            return 0, 0, 0
        elif (i, j) == self.exp_partner:
            return 255, 204, 0  # Gold
        else:
            return self._colors[i, j]

    def get_cell_opacity(self, i, j):
        if self._ui.hl_province is None:
            return 0
        elif self.trade_distance[i, j] < 0:  # out of range
            return 0
        elif (i, j) == self.exp_partner:
            return 255
        else:
            return self._opacity