Module likelihood.models.environments

def flatten_chain(matrix: List[List[Any]]) -> List[Any]:
    return list(chain.from_iterable(matrix))

class ActionSpace:
    def __init__(self, num_actions: int) -> None:
        self._num_actions = num_actions

    @property
    def n(self) -> int:
        return self._num_actions

class OptionCriticEnv:
    """
    An environment for Option Critic reinforcement learning that processes a dataset of episodes.

    Attributes
    ----------
    episodes : `Dict[str, Dict]`
        Dataset of episodes with state, action, selected_option, reward, next_state, and done information.
    observation_space : `np.ndarray`
        Initial observation space shape (from first episode's state)
    done : `bool`
        Whether the current episode has terminated
    num_options : `int`
        Number of distinct options available in the dataset
    actions_by_option : `defaultdict(set)`
        Maps selected options to sets of actions that were taken with them
    unique_actions_count : `List[int]`
        Count of unique actions per option index (used for action space definition)
    action_space : `ActionSpace`
        Custom action space defined by unique actions per option
    idx_episode : `int`
        Current episode index being processed
    current_state : `np.ndarray`
        Current state observation in the environment
    epsilon : `float`
        Probability of choosing the reward from the nearest state. By default it is set to `0.1`.
    """

    def __init__(
        self,
        episodes: Dict[int, Dict[str, List]],
    ) -> None:
        """
        Initializes the OptionCriticEnv with a dataset of episodes.

        Parameters
        ----------
        episodes : `Dict[int, Dict]`
            Dataset of episodes where keys are episode identifiers and values are episode data.
            Each episode must contain at least:
                - *state*: `List` of state observations
                - *selected_option*: `List[int]` or `List[List[int]]` of selected options
                - *action*: `List[int]` or `List[List[int]]` of actions taken
                - *reward*: `List` of rewards
                - *next_state*: `List` of next states
                - *done*: `List` of termination flags

        Raises
        ------
        ValueError
            If required fields ("state" or "selected_option") are missing from episode data
        """
        self.episodes = episodes
        self.multiple_option = False

        required_keys = ["state", "action", "selected_option", "reward", "next_state", "done"]
        for episode_id, data in episodes.items():
            if not all(k in data for k in required_keys):
                raise ValueError(
                    f"Episode {episode_id} missing keys: {set(required_keys) - set(data.keys())}"
                )

        self.observation_space = np.array(episodes[0]["state"][0])
        self.done = False
        self.idx_episode = 0
        self.current_state = None
        self.num_options = len(
            set(flatten_chain(episodes[0]["selected_option"]))
            if isinstance(episodes[0]["selected_option"][0], list)
            else episodes[0]["selected_option"]
        )
        self.actions_by_option = defaultdict(set)

        # Build fast lookup for transitions
        self.state_action_option_to_transition: Dict[Tuple, Dict[str, Any]] = {}

        for episode_id, data in episodes.items():
            states = data["state"]
            actions = data["action"]
            options = data["selected_option"]
            next_states = data["next_state"]
            rewards = data["reward"]
            dones = data["done"]

            for i in range(len(states)):
                key = self._make_transition_key(
                    states[i],
                    options[i],
                    actions[i],
                )

                self.state_action_option_to_transition[key] = {
                    "next_state": next_states[i],
                    "reward": rewards[i],
                    "done": dones[i],
                }

            for i, selected in enumerate(options):
                self.actions_by_option[
                    tuple(selected) if isinstance(options[0], list) else selected
                ].add(tuple(actions[i]) if isinstance(actions[i], list) else actions[i])

        check_type = list(set([key for key in self.actions_by_option.keys()]))
        keys_actions_by_option = list(sorted(self.actions_by_option.keys(), reverse=True))
        actions = self.actions_by_option[keys_actions_by_option[0]]

        self.unique_actions_count = [
            len(
                list(
                    OrderedDict.fromkeys(
                        flatten_chain(
                            [
                                list(action)
                                for action in self.actions_by_option.get(
                                    keys_actions_by_option[i], set()
                                )
                            ]
                        )
                    )
                )
                if isinstance(keys_actions_by_option[i], tuple)
                else self.actions_by_option.get(keys_actions_by_option[i], set())
            )
            for i in range(
                max(self.actions_by_option.keys()) + 1
                if not isinstance(check_type[0], tuple)
                else len(set(keys_actions_by_option))
            )
        ]

        self.action_space = ActionSpace(self.unique_actions_count)
        next_states = np.array(
            [trans["next_state"] for trans in self.state_action_option_to_transition.values()]
        )
        self.kdtree = KDTree(next_states)
        self.trans_list = list(self.state_action_option_to_transition.values())
        self.stats = 0.0
        self.count_step_found = 0
        self.count_step_nearest = 0
        self.threshold = 0.5
        self.rate = 0.1
        self.epsilon = 0.1
        self.states = np.array(
            [
                self.state_action_option_to_transition[k]["next_state"]
                for k in self.state_action_option_to_transition
            ]
        )
        self.std_per_dim = np.std(states, axis=0)

    def _make_transition_key(
        self,
        state: List[float],
        option: Union[int, List[int]],
        action: Union[int, List[int]],
        decimals: int = 6,
    ) -> Tuple:
        """
        Builds a canonical transition key:
        ((state_tuple), option(s)..., action(s)...)
        """
        state_key = tuple(round(float(x), decimals) for x in state)

        if isinstance(action, (list, tuple)):
            self.multiple_option = True
            return (state_key,) + tuple(int(o) for o in option) + tuple(int(a) for a in action)
        else:
            return (state_key, int(option), int(action))

    def _check_condition(self, value: float, other: float) -> float:
        if other < self.threshold:
            increment = self.rate * (self.threshold - other)
            value += increment
            value = min(value, 1.0)
        return value

    def reset(self) -> tuple[np.ndarray, dict]:
        """
        Resets the environment to a random episode and returns the initial state.

        Returns
        -------
        observation : `np.ndarray`
            Initial state observation
        info : `Dict`
            Empty dictionary (no additional information)
        """
        self.count_step_found = 0
        self.count_step_nearest = 0
        episode_id = np.random.choice(list(self.episodes.keys()))
        self.idx_episode = episode_id
        self.current_state = self.episodes[episode_id]["state"][0]
        return self.current_state, {}

    def step(
        self, action: int | List[int], option: int | List[int]
    ) -> tuple[np.ndarray, float, bool, bool, dict]:
        """
        Takes an action with a specific option and returns the next state, reward, and termination status.

        Parameters
        ----------
        action : `int` | `List[int]`
            Action index to execute
        option : `int` | `List[int]`
            Selected option index

        Returns
        -------
        next_state : `np.ndarray`
            State after taking the action
        reward : `float`
            Immediate reward for the transition
        done : `bool`
            Whether the episode has terminated (from episode data)
        terminated : `bool`
            Whether the action-option pair was found in the dataset
        info : `Dict`
            Empty dictionary (no additional information)

        Notes
        -----
        - Uses a KD-tree for efficient nearest-neighbor lookup of next states.
        """
        key = self._make_transition_key(
            self.current_state,
            option,
            action,
        )
        self.stats = round(
            self.count_step_found
            / (
                (self.count_step_found + self.count_step_nearest)
                if (self.count_step_found + self.count_step_nearest) > 0.0
                else 1.0
            ),
            6,
        )

        if key in self.state_action_option_to_transition:
            self.count_step_found += 1
            trans = self.state_action_option_to_transition[key]
            self.current_state = trans["next_state"]
            return trans["next_state"].copy(), trans["reward"], trans["done"], True, {}
        else:
            # Query KD-tree for nearest neighbor
            self.count_step_nearest += 1
            distance, idx = self.kdtree.query(self.current_state)
            closest_trans = self.trans_list[idx]
            closest_state = closest_trans["next_state"]
            closest_reward = closest_trans["reward"]
            closest_reward_prob = random.random()
            self.epsilon = self._check_condition(self.epsilon, self.stats)
            if np.array_equal(self.current_state, closest_state):
                if self.count_step_nearest > 1:
                    done = True
                    reward = -1.0
                else:
                    done = False
                    reward = 0.0
                    closest_state = closest_state + np.random.normal(
                        0, self.std_per_dim, size=closest_state.shape
                    )
            else:
                done = False
                reward = 0.0

            self.current_state = closest_state.copy()
            return (
                closest_state.copy(),
                reward if closest_reward_prob < self.epsilon else closest_reward,
                done if done else closest_trans.get("done", False),
                False,
                {},
            )