imitation_datasets

class Controller:
    """Controller for running experiments."""

    def __init__(
            self,
            enjoy: EnjoyFunction,
            collate: CollateFunction,
            amount: int,
            threads: int = 1,
            path: str = './dataset/'
    ) -> None:
        """Initialize the controller.

        Args:
            enjoy (EnjoyFunction): Function to run the expert.
            collate (CollateFunction): Function to collate the data.
            amount (int): Amount of episodes to run.
            threads (int, optional): Amount of threads to use. Defaults to 1.
            path (str, optional): Path to save the dataset. Defaults to './dataset/'.
        """
        self.enjoy = enjoy
        self.collate = collate
        self.threads = CPUS(threads)
        self.experiments = Experiment(amount)
        self.path = path

        self.pbar = None
        set_start_method('spawn', force=True)

    def create_folder(self, path: str) -> None:
        """Create a folder if it does not exist.

        Args:
            path (str): Path to the folder.
        """
        if not os.path.exists(path):
            os.makedirs(path)

    async def set_cpu(self, cpu: int) -> None:
        """Set the cpu affinity for the current process.

        Args:
            cpu (int): CPU index to use.
        """
        try:
            proc = psutil.Process()
            proc.cpu_affinity([int(cpu)])
            if 'linux' in platform:
                os.sched_setaffinity(proc.pid, [int(cpu)])
        except OSError:
            pass

    def enjoy_closure(self, opt: Namespace) -> EnjoyFunction:
        """Create a closure for the enjoy function.

        Args:
            opt (Namespace): Namespace with the arguments.

        Returns:
            EnjoyFunction: Enjoy function with part of the arguments.
        """
        os.system("set LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libGLEW.so")
        os.system("set LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia")
        return partial(self.enjoy, expert=Experts.get_expert(opt.game))

    def collate_closure(self, opt: Namespace) -> CollateFunction:
        """Create a closure for the collate function.

        Args:
            opt (Namespace): Namespace with the arguments.

        Returns:
            CollateFunction: Collate function with part of the arguments.
        """
        path = f'{self.path}{opt.game}/'
        files = list(listdir(path))
        return partial(self.collate, data=files, path=path)

    async def enjoy_sequence(self, future: EnjoyFunction, executor: ProcessPoolExecutor) -> bool:
        """_summary_

        Args:
            future (EnjoyFunction): Enjoy function already with async future.
            executor (ProcessPoolExecutor): Executor to run the future.

        Returns:
            bool: Result of the future.
                  True if the expert was able to solve the game. False otherwise.
        """
        # Pre
        cpu = await self.threads.cpu_allock()
        await self.experiments.start()
        await self.set_cpu(cpu)

        # Enjoy
        result = await asyncio.get_event_loop().run_in_executor(executor, future)

        # Post
        self.threads.cpu_release(cpu)
        await self.experiments.stop(result)
        self.pbar.update(1 if result else 0)

        return result if result else await asyncio.gather(self.enjoy_sequence(future, executor))

    async def run(self, opt) -> None:
        """Run the experiments.

        Args:
            opt (Namespace): Namespace with the arguments.
        """
        path = f'{self.path}{opt.game}/'
        self.create_folder(path)

        tasks = []
        with ProcessPoolExecutor() as executor:
            for idx in range(self.experiments.amount):
                enjoy = self.enjoy_closure(opt)
                enjoy = partial(enjoy, path=path, context=Context(self.experiments, idx))
                task = asyncio.ensure_future(
                    self.enjoy_sequence(
                        enjoy,
                        executor
                    )
                )
                tasks.append(task)
            await asyncio.gather(*tasks)

    def start(self, opt: Namespace):
        """Start the experiments.

        Args:
            opt (Namespace): Namespace with the arguments.

        Raises:
            exception: Exception (general) raised during the execution.
        """
        try:
            if opt.mode in ['all', 'play']:
                self.pbar = tqdm(range(self.experiments.amount), desc='Running episodes')
                asyncio.run(self.run(opt))

            if opt.mode in ['all', 'collate']:
                self.pbar = tqdm(range(self.experiments.amount), desc='Running collate')
                collate = self.collate_closure(opt)
                collate()
        except Exception as exception:
            self.experiments.add_log(-99, exception)
            raise exception
        finally:
            self.experiments.write_log()

@dataclass
class Policy:
    """Policy dataclass to load and use expert policies."""

    name: str
    repo_id: str
    filename: str
    threshold: float
    algo: BaseAlgorithm
    policy: BaseAlgorithm = field(init=False, default=None)
    internal_state: Any = field(init=False, default=None)
    environment: Any = field(init=False, default=None)

    def load(self) -> BaseAlgorithm:
        """
        Load policy from HuggingFace hub.
        It uses a custom_object to replicate stable_baselines behaviour.

        custom_objects = {
            "learning_rate": 0.0,
            "lr_schedule": lambda _: 0.0,
            "clip_range": lambda _: 0.0
        }

        Returns:
            BaseAlgorithm: Stable baseline policy loaded from HuggingFace hub.
        """
        checkpoint = load_from_hub(
            repo_id=self.repo_id,
            filename=self.filename,
        )

        custom_objects = {
            "learning_rate": 0.0,
            "lr_schedule": lambda _: 0.0,
            "clip_range": lambda _: 0.0
        }

        self.policy = self.algo.load(
            checkpoint,
            custom_objects=custom_objects
        )
        return self.policy

    def predict(
            self,
            obs: List[Union[int, float]],
            deterministic: bool = True
    ) -> Tuple[
        Union[int, float, List[Union[int, float]]],
        Union[int, float, List[Union[int, float]]]
    ]:
        """
        Predict action given observation.

        Args:
            obs (List[int | float]): observation from environment.
            deterministic (bool, optional): Use exploration to predict action. Defaults to True.

        Returns:
            action (Union[int, float, List[Union[int, float]]]):
                action predicted by the policy.
            internal_states (Union[int, float, List[Union[int, float]]]):
                internal states of the policy.

        Note: typing depends on the environment.
        """
        action, internal_states = self.policy.predict(
            obs,
            state=self.internal_state,
            deterministic=deterministic,
        )
        self.internal_state = internal_states
        return action, internal_states

    def get_environment(self) -> str:
        """Return environment name.

        Returns:
            str: environment name.
        """
        if self.environment is None:
            self.environment = gym.make(self.name, render_mode="rgb_array")
        return self.environment

class Experts:
    """Helper class to register and get expert policies."""
    experts: Dict[str, Policy] = {
        key: Policy(**value) for env in [atari, classic, mujoco] for key, value in env.items()
    }

    @classmethod
    def register(cls, identifier: str, policy: Policy) -> None:
        """Register a new policy."""
        if not isinstance(policy.threshold, float):
            policy.threshold = float(policy.threshold)

        cls.experts[identifier] = policy

    @classmethod
    def get_expert(cls, identifier: str) -> Policy:
        """Return expert policy.

        Args:
            identifier (str): identifier of the policy.

        Returns:
            Policy: dataclass with expert policy information.
        """
        try:
            return cls.experts[identifier]
        except KeyError:
            return None

    @classmethod
    def get_register(cls) -> None:
        """Print entire register of expert policies."""
        return cls.experts

@dataclass
class Experiment:
    """Experiment dataclass to keep track of the experiments."""

    amount: int
    path: str = './logs.txt'
    waiting: int = field(
        init=False,
        default_factory=int
    )
    logs: DefaultDict[int, list] = field(
        init=False,
        default_factory=lambda: defaultdict(list)
    )
    experiment_semaphore: asyncio.Lock = field(
        init=False,
        default=asyncio.BoundedSemaphore(value=1)
    )

    def __post_init__(self) -> None:
        """Write in log file that the dataset creation has started."""
        if os.path.exists(self.path):
            os.remove(self.path)

        if not os.path.exists(self.path):
            with open(self.path, 'w', encoding='utf8') as log_file:
                log_file.write('#### Starting dataset creation ####\n')

    def is_done(self) -> bool:
        """Check if the experiment is done.

        Returns:
            bool: True if the experiment is done, False otherwise.
        """
        return self.amount == 0

    async def start(self, amount: int = 1) -> Tuple[bool, int]:
        """Start an experiment.

        Args:
            amount (int, optional): How many experiments are left to run. Defaults to 1.

        Returns:
            status (bool): True if the experiment can be started, False otherwise.
            amount (int): How many experiments are left to run.
        """
        await self.experiment_semaphore.acquire()
        if self.amount > 0:
            self.waiting += amount
            self.amount -= amount
            self.experiment_semaphore.release()
            return True, self.amount

        self.experiment_semaphore.release()
        return False, -1

    async def stop(self, status: bool, amount: int = 1) -> None:
        """Stop an experiment.

        Args:
            status (bool): True if the experiment was successful, False otherwise.
            amount (int, optional): How many experiments are left to run. Defaults to 1.
        """
        await self.experiment_semaphore.acquire()
        self.amount += 0 if status else amount
        self.waiting -= amount
        self.experiment_semaphore.release()

    def add_log(self, experiment: int, log: str) -> None:
        """Add a log to the experiment.

        Args:
            experiment (int): Experiment index.
            log (str): Log to add.
        """
        self.logs[experiment].append(log)

    def write_log(self) -> None:
        """Write the logs in the log file."""
        with open('./logs.txt', 'a', encoding='utf8') as log_file:
            for idx, logs in self.logs.items():
                for log in logs:
                    log_file.write(f'\nExperiment {idx}: {log}')
                log_file.write('\n')

@dataclass
class Context:
    """Context dataclass to keep track of the context of the experiment."""
    experiments: Experiment
    index: int

    def add_log(self, log: str) -> None:
        """Add a log to the experiment."""
        self.experiments.add_log(self.index, log)

@dataclass
class CPUS(metaclass=Singleton):
    """CPUS dataclass to keep track of the available CPUs."""

    available_cpus: int = field(default_factory=multiprocessing.cpu_count())
    cpus: DefaultDict[int, bool] = field(init=False, default_factory=lambda: defaultdict(bool))
    cpu_semaphore: asyncio.Lock = field(init=False)

    def __post_init__(self) -> None:
        """Initialize the cpu_semaphore."""
        if self.available_cpus > multiprocessing.cpu_count() - 1:
            self.available_cpus = multiprocessing.cpu_count() - 1
        self.cpu_semaphore = asyncio.BoundedSemaphore(value=self.available_cpus)

    async def cpu_allock(self) -> int:
        """Acquire a CPU.

        Returns:
            int: CPU index.
        """
        await self.cpu_semaphore.acquire()
        for idx in range(self.available_cpus):
            if not self.cpus[idx]:
                self.cpus[idx] = True
                return idx

    def cpu_release(self, cpu_idx: int) -> None:
        """Release a CPU.

        Args:
            cpu_idx (int): CPU index.
        """
        try:
            self.cpus[cpu_idx] = False
            self.cpu_semaphore.release()
        except ValueError:
            pass

class GymWrapper:
    """
        Wrapper for gym environment. Since Gymnasium and Gym version 0.26
        there are some environments that were working under Gym-v.0.21 stopped 
        working. This wrapper just makes sure that the output for the environment 
        will always work with the version the user wants.
    """

    def __init__(self, environment: Any, version: str = "newest") -> None:
        """
        Args:
            name: gym environment name
            version: ["newest", "older"] refers to the compatibility version.

        In this case, "newest" is 0.26 and "older" is 0.21.
        """
        if version not in ["newest", "older"]:
            raise ValueError("Version has to be : ['newest', 'older']")

        self.env = environment
        state = environment.reset()
        if version == "older" and not isinstance(state[0], np.floating):
            raise WrapperException("Incopatible environment version and wrapper version.")
        if version == "newest" and not isinstance(state[0], np.ndarray):
            raise WrapperException("Incopatible environment version and wrapper version.")

        self.version = version

    @property
    def action_space(self):
        """Map gym action_space attribute to wrapper."""
        return self.env.action_space

    @property
    def observation_space(self):
        """Map gym env_space attribute to wrapper."""
        return self.env.observation_space

    def set_seed(self, seed: int) -> None:
        """Set seed for all packages (Pytorch, Numpy and Python).

        Args:
            seed (optional, int): seed number to use for the random generator.
        """
        torch.manual_seed(seed)
        np.random.seed(seed)
        random.seed(seed)

    def reset(self) -> Union[Tuple[List[float], Dict[str, Any]], List[float]]:
        """Resets the framework and return the appropriate return."""
        state = self.env.reset()
        if self.version == "newest":
            return state[0]
        return state

    def step(
            self,
            action: Union[float, int]
    ) -> Union[
        Tuple[List[float], float, bool, bool, Dict[str, Any]],
        Tuple[List[float], float, bool, Dict[str, Any]]
    ]:
        """
        Perform an action in the environment and return the appropriate return
        according to version.
        """
        gym_return = self.env.step(action)
        if self.version == "newest":
            state, reward, terminated, truncated, info = gym_return
            return state, reward, terminated or truncated, info

        return gym_return

    def render(self, mode="rgb_array"):
        """Return the render for the environment."""
        if self.version == "newest":
            state = self.env.render()
            if state is None and self.env.render_mode != "human":
                raise WrapperException("No render mode set.")
            return state

        return self.env.render(mode)

    def close(self) -> None:
        """Close the environment."""
        self.env.close()

class WrapperException(Exception):
    """Wrapper exception for all exceptions related to the wrapper."""

    def __init__(self, message: str) -> None:
        self.message = message
        super().__init__(self.message)

class BaselineDataset(Dataset):
    """Teacher dataset for IL methods."""

    def __init__(
        self,
        path: str,
        source: str = "local",
        hf_split: str = "train",
        split: str = "train",
        n_episodes: int = None,
        transform: Callable[[torch.Tensor], torch.Tensor] = None
    ) -> None:
        """Initialize dataset.

        Args:
            path (Str): path to the dataset.
            source (str): whether is a HuggingFace or a local dataset.
                Defaults to 'local'.
            hf_split (str): HuggingFace split to use. Defaults to 'train'.
            split (str): split to use. Defaults to 'train'.
            n_episodes (int): number of episodes to use. Defaults to None.
            transform (Callable[[torch.Tensor], torch.Tensor]): transform to apply to the data. Defaults to None.

        Raises:
            ValueError: if path does not exist.
        """
        self.transform = transform

        if source == "local" and not os.path.exists(path):
            raise ValueError(f"No dataset at: {path}")

        if source == "local":
            self.data = np.load(path, allow_pickle=True)
            self.average_reward = np.mean(self.data["episode_returns"])
        else:
            dataset = load_dataset(path, split=hf_split)
            self.data = huggingface_to_baseline(dataset)
            if len(self.data["obs"].shape) == 1:
                self.data["obs"] = self.data["obs"].reshape((-1, 1))
            self.average_reward = []

        shape = [1] if isinstance(self.data["obs"][0], str) else self.data["obs"].shape[1:]
        self.states = np.ndarray(shape=(0, *shape))
        self.next_states = np.ndarray(shape=(0, *shape))

        if len(self.data["actions"].shape) == 1:
            action_size = 1
        else:
            action_size = self.data["actions"].shape[-1]
        self.actions = np.ndarray(shape=(0, action_size))

        episode_starts = list(np.where(self.data["episode_starts"] == 1)[0])
        episode_starts.append(len(self.data["episode_starts"]))

        if n_episodes is not None:
            if split == "train":
                episode_starts = episode_starts[:n_episodes + 1]
            else:
                episode_starts = episode_starts[n_episodes:]

        for start, end in zip(episode_starts, tqdm(episode_starts[1:], desc="Creating dataset")):
            episode = self.data["obs"][start:end]
            actions = self.data["actions"][start:end].reshape((-1, 1))
            self.actions = np.append(self.actions, actions[:-1], axis=0)
            self.states = np.append(self.states, episode[:-1], axis=0)
            self.next_states = np.append(self.next_states, episode[1:], axis=0)

            if source != "local":
                self.average_reward.append(self.data["rewards"][start:end].sum())

        if isinstance(self.average_reward, list):
            self.average_reward = np.mean(self.average_reward)

        assert self.states.shape[0] == self.actions.shape[0] == self.next_states.shape[0]

        if not isinstance(self.states[0, 0], str):
            self.states = torch.from_numpy(self.states)
            self.next_states = torch.from_numpy(self.next_states)
        self.actions = torch.from_numpy(self.actions)

    def __len__(self) -> int:
        """Dataset length.

        Returns:
            length (int): length.
        """
        return self.states.shape[0]

    def __getitem__(self, index: int) -> Tuple[torch.Tensor]:
        """Get item from dataset.

        Args:
            index (int): index.

        Returns:
            state (torch.Tensor): state for timestep t.
            action (torch.Tensor): action for timestep t.
            next_state (torch.Tensor): state for timestep t + 1.
        """
        state = self.states[index]
        next_state = self.next_states[index]
        if isinstance(state[0], str):
            state = ToTensor()(Image.open(state[0]))
            next_state = ToTensor()(Image.open(next_state[0]))

        if self.transform is not None:
            state = self.transform(state)
            next_state = self.transform(next_state)

        action = self.actions[index]
        return state, action, next_state