PrimAITE/src/primaite/environment/observations.py

"""Module for handling configurable observation spaces in PrimAITE."""
import logging
from abc import ABC, abstractmethod
from enum import Enum
from typing import List, Tuple

import numpy as np
from gym import spaces

from primaite.common.enums import FileSystemState, HardwareState, SoftwareState
from primaite.environment.primaite_env import Primaite
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode

_LOGGER = logging.getLogger(__name__)


class AbstractObservationComponent(ABC):
    """Represents a part of the PrimAITE observation space."""

    @abstractmethod
    def __init__(self, env: Primaite):
        _LOGGER.info(f"Initialising {self} observation component")
        self.env: Primaite = env
        self.space: spaces.Space
        self.current_observation: np.ndarray  # type might be too restrictive?
        return NotImplemented

    @abstractmethod
    def update(self):
        """Look at the environment and update the current observation value."""
        self.current_observation = NotImplemented


class NodeLinkTable(AbstractObservationComponent):
    """Table with nodes/links as rows and hardware/software status as cols.

    Initialise the observation space with the BOX option chosen.

    This will create the observation space formatted as a table of integers.
    There is one row per node, followed by one row per link.
    Columns are as follows:
        * node/link ID
        * node hardware status / 0 for links
        * node operating system status (if active/service) / 0 for links
        * node file system status (active/service only) / 0 for links
        * node service1 status / traffic load from that service for links
        * node service2 status / traffic load from that service for links
        * ...
        * node serviceN status / traffic load from that service for links

    For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be
    ``(12, 7)``
    #todo: clean up description

    """

    _FIXED_PARAMETERS = 4
    _MAX_VAL = 1_000_000
    _DATA_TYPE = np.int64

    def __init__(self, env: Primaite):
        super().__init__(env)

        # 1. Define the shape of your observation space component
        num_items = self.env.num_links + self.env.num_nodes
        num_columns = self.env.num_services + self._FIXED_PARAMETERS
        observation_shape = (num_items, num_columns)

        # 2. Create Observation space
        self.space = spaces.Box(
            low=0,
            high=self._MAX_VAL,
            shape=observation_shape,
            dtype=self._DATA_TYPE,
        )

        # 3. Initialise Observation with zeroes
        self.current_observation = np.zeroes(observation_shape, dtype=self._DATA_TYPE)

    def update_obs(self):
        """Update the observation.

        Update the environment's observation state based on the current status of nodes and links.

        The structure of the observation space is described in :func:`~_init_box_observations`
        This function can only be called if the observation space setting is set to BOX.
        todo: complete description..
        """
        item_index = 0
        nodes = self.env.nodes
        links = self.env.links
        # Do nodes first
        for _, node in nodes.items():
            self.current_observation[item_index][0] = int(node.node_id)
            self.current_observation[item_index][1] = node.hardware_state.value
            if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
                self.current_observation[item_index][2] = node.software_state.value
                self.current_observation[item_index][
                    3
                ] = node.file_system_state_observed.value
            else:
                self.current_observation[item_index][2] = 0
                self.current_observation[item_index][3] = 0
            service_index = 4
            if isinstance(node, ServiceNode):
                for service in self.env.services_list:
                    if node.has_service(service):
                        self.current_observation[item_index][
                            service_index
                        ] = node.get_service_state(service).value
                    else:
                        self.current_observation[item_index][service_index] = 0
                    service_index += 1
            else:
                # Not a service node
                for service in self.env.services_list:
                    self.current_observation[item_index][service_index] = 0
                    service_index += 1
            item_index += 1

        # Now do links
        for _, link in links.items():
            self.current_observation[item_index][0] = int(link.get_id())
            self.current_observation[item_index][1] = 0
            self.current_observation[item_index][2] = 0
            self.current_observation[item_index][3] = 0
            protocol_list = link.get_protocol_list()
            protocol_index = 0
            for protocol in protocol_list:
                self.current_observation[item_index][
                    protocol_index + 4
                ] = protocol.get_load()
                protocol_index += 1
            item_index += 1


class NodeStatuses(AbstractObservationComponent):
    """todo: complete description.

    This will create the observation space with node observations followed by link observations.
    Each node has 3 elements in the observation space plus 1 per service, more specifically:
        * hardware state
        * operating system state
        * file system state
        * service states (one per service)
    """

    _DATA_TYPE = np.int64

    def __init__(self, env: Primaite):
        super().__init__(env)

        # 1. Define the shape of your observation space component
        shape = [
            len(HardwareState) + 1,
            len(SoftwareState) + 1,
            len(FileSystemState) + 1,
        ]
        services_shape = [len(SoftwareState) + 1] * self.env.num_services
        shape = shape + services_shape

        # 2. Create Observation space
        self.space = spaces.MultiDiscrete(shape)

        # 3. Initialise observation with zeroes
        self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)

    def update_obs(self):
        """todo: complete description.

        Update the environment's observation state based on the current status of nodes and links.

        The structure of the observation space is described in :func:`~_init_multidiscrete_observations`
        This function can only be called if the observation space setting is set to MULTIDISCRETE.


        """
        obs = []
        for _, node in self.env.nodes.items():
            hardware_state = node.hardware_state.value
            software_state = 0
            file_system_state = 0
            service_states = [0] * self.env.num_services

            if isinstance(node, ActiveNode):
                software_state = node.software_state.value
                file_system_state = node.file_system_state_observed.value

            if isinstance(node, ServiceNode):
                for i, service in enumerate(self.env.services_list):
                    if node.has_service(service):
                        service_states[i] = node.get_service_state(service).value
        obs.extend([hardware_state, software_state, file_system_state, *service_states])
        self.current_observation[:] = obs


class LinkTrafficLevels(AbstractObservationComponent):
    """todo: complete description.

    Each link has one element in the observation space, corresponding to the traffic load,
    it can take the following values:
        0 = No traffic (0% of bandwidth)
        1 = No traffic (0%-33% of bandwidth)
        2 = No traffic (33%-66% of bandwidth)
        3 = No traffic (66%-100% of bandwidth)
        4 = No traffic (100% of bandwidth)
    """

    _DATA_TYPE = np.int64

    def __init__(
        self,
        env: Primaite,
        combine_service_traffic: bool = False,
        quantisation_levels: int = 5,
    ):
        super().__init__(env)
        self._combine_service_traffic: bool = combine_service_traffic
        self._quantisation_levels: int = quantisation_levels
        self._entries_per_link: int = 1

        if not self._combine_service_traffic:
            self._entries_per_link = self.env.num_services

        # 1. Define the shape of your observation space component
        shape = (
            [self._quantisation_levels] * self.env.num_links * self._entries_per_link
        )

        # 2. Create Observation space
        self.space = spaces.MultiDiscrete(shape)

        # 3. Initialise observation with zeroes
        self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)

    def update_obs(self):
        """todo: complete description."""
        obs = []
        for _, link in self.env.links.items():
            bandwidth = link.bandwidth
            if self._combine_service_traffic:
                loads = [link.get_current_load()]
            else:
                loads = [protocol.get_load() for protocol in link.protocol_list]

            for load in loads:
                if load <= 0:
                    traffic_level = 0
                elif load >= bandwidth:
                    traffic_level = self._quantisation_levels - 1
                else:
                    traffic_level = (load / bandwidth) // (
                        1 / (self._quantisation_levels - 1)
                    ) + 1

                obs.append(int(traffic_level))

        self.current_observation[:] = obs


class ObservationsHandler:
    """Component-based observation space handler."""

    class registry(Enum):
        """todo: complete description."""

        NODE_LINK_TABLE: NodeLinkTable
        NODE_STATUSES: NodeStatuses
        LINK_TRAFFIC_LEVELS: LinkTrafficLevels

    def __init__(self):
        """todo: complete description."""
        """Initialise the handler without any components yet. They"""
        self.registered_obs_components: List[AbstractObservationComponent] = []
        self.space: spaces.Space
        self.current_observation: Tuple[np.ndarray]
        # i can access the registry items like this:
        # self.registry.LINK_TRAFFIC_LEVELS

    def update_obs(self):
        """todo: complete description."""
        current_obs = []
        for obs in self.registered_obs_components:
            obs.update_obs()
            current_obs.append(obs.current_observation)
        self.current_observation = tuple(current_obs)

    def register(self, obs_component: AbstractObservationComponent):
        """todo: complete description."""
        self.registered_obs_components.append(obs_component)
        self.update_space()

    def deregister(self, obs_component: AbstractObservationComponent):
        """todo: complete description."""
        self.registered_obs_components.remove(obs_component)
        self.update_space()

    def update_space(self):
        """todo: complete description."""
        component_spaces = []
        for obs_comp in self.registered_obs_components:
            component_spaces.append(obs_comp.space)
        self.space = spaces.Tuple(component_spaces)

    @classmethod
    def from_config(cls, obs_space_config):
        """todo: complete description.

        This method parses config items related to the observation space, then
        creates the necessary components and adds them to the observation handler.
        """
        # Instantiate the handler
        handler = cls()

        for component_cfg in obs_space_config["components"]:
            # Figure out which class can instantiate the desired component
            comp_type = component_cfg["name"]
            comp_class = cls.registry[comp_type].value

            # Create the component with options from the YAML
            component = comp_class(**component_cfg["options"])

            handler.register(component)

        return handler