special_operations.py

"""@package docstring
B-ASIC Special Operations Module.
TODO: More info.
"""

from numbers import Number
from typing import Optional, Sequence

from b_asic.operation import AbstractOperation, ResultKey, RegisterMap, MutableResultMap, MutableRegisterMap
from b_asic.graph_component import Name, TypeName
from b_asic.port import SignalSourceProvider


class Input(AbstractOperation):
    """Input operation.
    TODO: More info.
    """

    def __init__(self, name: Name = ""):
        super().__init__(input_count = 0, output_count = 1, name = name)
        self.set_param("value", 0)

    @property
    def type_name(self) -> TypeName:
        return "in"
    
    def evaluate(self):
        return self.param("value")

    @property
    def value(self) -> Number:
        """Get the current value of this input."""
        return self.param("value")

    @value.setter
    def value(self, value: Number) -> None:
        """Set the current value of this input."""
        self.set_param("value", value)


class Output(AbstractOperation):
    """Output operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 0, name = name, input_sources = [src0])

    @property
    def type_name(self) -> TypeName:
        return "out"

    def evaluate(self, _):
        return None


class Register(AbstractOperation):
    """Unit delay operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, initial_value: Number = 0, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])
        self.set_param("initial_value", initial_value)

    @property
    def type_name(self) -> TypeName:
        return "reg"

    def evaluate(self, a):
        return self.param("initial_value")

    def current_output(self, index: int, registers: Optional[RegisterMap] = None, prefix: str = "") -> Optional[Number]:
        if registers is not None:
            return registers.get(self.key(index, prefix), self.param("initial_value"))
        return self.param("initial_value")
    
    def evaluate_output(self, index: int, input_values: Sequence[Number], results: Optional[MutableResultMap] = None, registers: Optional[MutableRegisterMap] = None, prefix: str = "") -> Number:
        if index != 0:
            raise IndexError(f"Output index out of range (expected 0-0, got {index})")
        if len(input_values) != 1:
            raise ValueError(f"Wrong number of inputs supplied to SFG for evaluation (expected 1, got {len(input_values)})")
        
        key = self.key(index, prefix)
        value = self.param("initial_value")
        if registers is not None:
            value = registers.get(key, value)
            registers[key] = self.truncate_inputs(input_values)[0]
        if results is not None:
            results[key] = value
        return value