from kennel.engine.components.component import Component, ComponentType


class MarkovTransitionState(Component):
    def __init__(
        self,
        state_names: dict[int, str],
        initial_state_vector: list[float],
        transition_matrix: list[list[float]],
    ):
        # TODO: Poll rate per state?
        # TODO: State being an enum instead of a vector, just choose max and map
        self.state_names = state_names
        self.state = initial_state_vector
        self.transition_matrix = transition_matrix

        super().__init__(ComponentType.MARKOV)

    def get_max_state_name(self, state_vector: list[float]):
        max_val = max(state_vector)
        return self.state_names[state_vector.index(max_val)]

    def to_dict(self):
        return {"state": self.get_max_state_name(self.state)}