from abc import ABC, abstractmethod
from dataclasses import dataclass

import torch

from tianshou.highlevel.env import Environments
from tianshou.highlevel.module.core import ModuleFactory, TDevice
from tianshou.utils.string import ToStringMixin


@dataclass
class IntermediateModule:
    """Container for a module which computes an intermediate representation (with a known dimension)."""

    module: torch.nn.Module
    output_dim: int


class IntermediateModuleFactory(ToStringMixin, ModuleFactory, ABC):
    """Factory for the generation of a module which computes an intermediate representation."""

    @abstractmethod
    def create_intermediate_module(self, envs: Environments, device: TDevice) -> IntermediateModule:
        pass

    def create_module(self, envs: Environments, device: TDevice) -> torch.nn.Module:
        return self.create_intermediate_module(envs, device).module