import torch.nn

from torchrl.modules.tensordict_module.actors import BatchedActionWrapper, Actor

from torchrl.envs import CatFrames, GymEnv, TransformedEnv, SerialEnv, InitTracker, Compose
from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod

time_steps = 6
n_obs = 4
n_action = 2
batch = 5


def reshape_cat(data: torch.Tensor):
    return data.unflatten(-1, (time_steps, n_obs))


actor_base = Seq(
    Mod(reshape_cat, in_keys=["obs_cat"], out_keys=["obs_cat_reshape"]),
    Mod(torch.nn.Linear(n_obs, n_action), in_keys=["obs_cat_reshape"], out_keys=["action"])
)
actor = BatchedActionWrapper(actor_base, n_steps=time_steps)

env = TransformedEnv(SerialEnv(batch, lambda: GymEnv("CartPole-v1")),
                     Compose(InitTracker(),
                             CatFrames(N=time_steps, in_keys=["observation"], out_keys=["obs_cat"], dim=-1)
                             ))

print(env.rollout(100, policy=actor, break_when_any_done=False))