Source code for smif.metadata.timestep

"""smif runs models over a series of planning timesteps.

:class:`~smif.metadata.timestep.RelativeTimestep` is used to specify one timestep relative to
another in the context of the set of planning timesteps.

For example, if a model needs to know its own outputs from a previous year, we can specify the
self-dependency using a relative timestep.
"""
from enum import Enum
from typing import List, Union

from smif.exception import SmifTimestepResolutionError


class RelativeTimestep(Enum):
    """Specify current, previous or base year timestep"""

    CURRENT = "CURRENT"  # current planning timestep
    PREVIOUS = "PREVIOUS"  # previous planning timestep
    BASE = "BASE"  # base year planning timestep
    ALL = "ALL"  # all planning timesteps

    @classmethod
    def from_name(cls, name):
        if name == "CURRENT":
            return cls.CURRENT
        elif name == "PREVIOUS":
            return cls.PREVIOUS
        elif name == "BASE":
            return cls.BASE
        elif name == "ALL":
            return cls.ALL
        raise ValueError("Relative timestep '%s' is not recognised" % name)

    def resolve_relative_to(
        self, timestep: int, timesteps: List[int]
    ) -> Union[None, int]:
        """Resolve a relative timestep with respect to a given timestep and
        sequence of timesteps.

        Parameters
        ----------
        timestep : int
        timesteps : list of int
        """
        if timestep not in timesteps:
            raise SmifTimestepResolutionError(
                "Timestep {} is not present in {}".format(timestep, timesteps)
            )

        # default None, e.g. for ALL
        relative_timestep = None

        if self.name == "CURRENT":
            relative_timestep = timestep
        elif self.name == "PREVIOUS":
            try:
                relative_timestep = element_before(timestep, timesteps)
            except ValueError:
                raise SmifTimestepResolutionError(
                    "{} has no previous timestep in {}".format(timestep, timesteps)
                )
        elif self.name == "BASE":
            relative_timestep = timesteps[0]

        return relative_timestep


def element_before(element: int, list_: List[int]) -> int:
    """Return the element before a given element in a list, or None if the
    given element is first or not in the list.
    """
    if element not in list_ or element == list_[0]:
        raise ValueError("No element before {} in {}".format(element, list_))
    else:
        index = list_.index(element)
        return list_[index - 1]