# pylint: disable=multiple-statements
import pathlib
from typing import Any, Dict, Iterator, List, Optional, Sequence, Union

import numpy as np

class Transition:
    def __init__(
        self,
        observation_shape: Sequence[int],
        action_size: int,
        observation: np.ndarray,
        action: Union[int, np.ndarray],
        reward: float,
        next_observation: np.ndarray,
        termination: float,
        truncation: float,
        prev_transition: Optional["Transition"] = ...,
        next_transition: Optional["Transition"] = ...,
    ): ...
    def get_observation_shape(self) -> Sequence[int]: ...
    def get_action_size(self) -> int: ...
    @property
    def is_discrete(self) -> bool: ...
    @property
    def observation(self) -> np.ndarray: ...
    @property
    def action(self) -> Union[int, np.ndarray]: ...
    @property
    def reward(self) -> float: ...
    @property
    def next_observation(self) -> np.ndarray: ...
    @property
    def termination(self) -> float: ...
    @property
    def truncation(self) -> float: ...
    @property
    def prev_transition(self) -> Optional["Transition"]: ...
    @prev_transition.setter
    def prev_transition(self, transition: "Transition") -> None: ...
    @property
    def next_transition(self) -> Optional["Transition"]: ...
    @next_transition.setter
    def next_transition(self, transition: "Transition") -> None: ...
    def clear_links(self) -> None: ...

class TransitionMiniBatch:
    def __init__(
        self,
        transitions: List[Transition],
        n_frames: int = ...,
        n_steps: int = ...,
        gamma: float = ...,
    ): ...
    @property
    def observations(self) -> np.ndarray: ...
    @property
    def actions(self) -> np.ndarray: ...
    @property
    def rewards(self) -> np.ndarray: ...
    @property
    def next_observations(self) -> np.ndarray: ...
    @property
    def transitions(self) -> List[Transition]: ...
    @property
    def terminations(self) -> np.ndarray: ...
    @property
    def truncations(self) -> np.ndarray: ...
    @property
    def n_steps(self) -> np.ndarray: ...
    def __len__(self) -> int: ...
    def __iter__(self) -> Iterator[Transition]: ...

def trace_back_and_clear(transition: Transition) -> None: ...

class Episode:
    def __init__(
        self,
        observation_shape: Sequence[int],
        action_size: int,
        observations: np.ndarray,
        actions: np.ndarray,
        rewards: np.ndarray,
        terminated: bool = ...,
        truncated: bool = ...,
    ): ...
    @property
    def observations(self) -> np.ndarray: ...
    @property
    def actions(self) -> np.ndarray: ...
    @property
    def rewards(self) -> np.ndarray: ...
    @property
    def terminated(self) -> bool: ...
    @property
    def truncated(self) -> bool: ...
    @property
    def transitions(self) -> List[Transition]: ...
    def build_transitions(self) -> None: ...
    def size(self) -> int: ...
    def get_observation_shape(self) -> Sequence[int]: ...
    def get_action_size(self) -> int: ...
    def compute_return(self) -> float: ...
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> Transition: ...
    def __iter__(self) -> Iterator[Transition]: ...

class MinariDataset:
    def __init__(
        self,
        dataset_name: str,
        algorithm_name: str,
        environment_name: str,
        environment_stack: Union[str, None],
        seed_used: int,
        code_permalink: Union[str, None],
        author: Union[str, None],
        author_email: Union[str, None],
        observations: np.ndarray,
        actions: np.ndarray,
        rewards: np.ndarray,
        terminations: np.ndarray,
        truncations: np.ndarray,
        episode_terminals: Optional[np.ndarray] = ...,
        discrete_action: Optional[bool] = ...,
    ): ...
    @property
    def dataset_name(self) -> str: ...
    @property
    def algorithm_name(self) -> str: ...
    @property
    def environment_name(self) -> str: ...
    @property
    def environment_stack(self) -> str: ...
    @property
    def seed_used(self) -> int: ...
    @property
    def code_permalink(self) -> Union[str, None]: ...
    @property
    def author(self) -> Union[str, None]: ...
    @property
    def author_email(self) -> Union[str, None]: ...
    @property
    def observations(self) -> np.ndarray: ...
    @property
    def actions(self) -> np.ndarray: ...
    @property
    def rewards(self) -> np.ndarray: ...
    @property
    def terminations(self) -> np.ndarray: ...
    @property
    def truncations(self) -> np.ndarray: ...
    @property
    def episode_terminals(self) -> np.ndarray: ...
    @property
    def episodes(self) -> List[Episode]: ...
    def size(self) -> int: ...
    def get_action_size(self) -> int: ...
    def get_observation_shape(self) -> Sequence[int]: ...
    def is_action_discrete(self) -> bool: ...
    def compute_stats(self) -> Dict[str, Any]: ...
    def clip_reward(self, low: Optional[float], high: Optional[float]) -> None: ...
    def append(
        self,
        observations: np.ndarray,
        actions: np.ndarray,
        rewards: np.ndarray,
        terminations: np.ndarray,
        truncations: np.ndarray,
        episode_terminals: Optional[np.ndarray] = ...,
    ) -> None: ...
    def extend(self, dataset: MinariDataset) -> None: ...
    def save(self) -> None: ...
    @classmethod
    def load(cls, fname: Union[str, pathlib.Path]) -> MinariDataset: ...
    def build_episodes(self) -> None: ...
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> Episode: ...
    def __iter__(self) -> Iterator[Episode]: ...
