from typing import Optional

import torch
from tensordict import TensorDict, TensorDictBase
from torchrl.envs import EnvBase
from torchrl.data import (
    CompositeSpec,
    UnboundedContinuousTensorSpec,
    OneHotDiscreteTensorSpec,
)
from torchrl.envs.utils import check_env_specs


class DemoEnv(EnvBase):
    batch_locked = True

    def __init__(self):
        super().__init__(run_type_checks=True)

        # For demo purpose, params are hard coded.
        self.node_num = 3
        self.batch_size = (2,)

        self.set_seed(torch.empty((), dtype=torch.int64).random_().item())
        self._make_spec()

    def _make_spec(self):
        """Set the observation, action, and reward specs."""
        self.observation_spec = CompositeSpec(
            current_node=OneHotDiscreteTensorSpec(
                n=self.node_num,
                shape=(*self.batch_size, self.node_num),
                dtype=torch.bool,
            ),
            shape=self.batch_size,
        )
        self.state_spec = self.observation_spec.clone()
        self.action_spec = UnboundedContinuousTensorSpec(shape=self.batch_size)
        self.reward_spec = UnboundedContinuousTensorSpec(shape=self.batch_size)

    def _set_seed(self, seed: Optional[int]):
        rng = torch.manual_seed(seed)
        self.rng = rng

    def _step(self, td: TensorDictBase) -> TensorDict:
        # Trivial step function for demo purpose.
        out = TensorDict(
            {
                "current_node": td["current_node"],
                "reward": td["action"],
                "done": td["done"],
            },
            td.shape,
        )
        return out

    def _reset(self, td) -> TensorDict:
        current_node = torch.zeros(*self.batch_size, self.node_num).bool()
        current_node[:, 0] = True

        out = TensorDict(
            {
                "current_node": current_node,
            },
            self.batch_size,
        )

        return out

env = DemoEnv()
check_env_specs(env)