import torch
from torchrl.envs.transforms import R3MTransform
from torchrl.data import TensorDict

for shape in ([], [2], [2, 3]):
    td = TensorDict({"next_observation": torch.randint(255, (*shape, 224, 224, 3))}, shape)
    transform = R3MTransform(model_name="resnet50", keys_in=["next_observation"])
    transform(td)
    print(td)

from torchrl.trainers.helpers.envs import LIBS
from utils import MJEnv
from torchrl.envs import R3MTransform, TransformedEnv, ParallelEnv
import time
import torch

LIBS["mjenv"] = MJEnv

device = torch.device("cuda:0") if torch.has_cuda and torch.cuda.device_count() else torch.device("cpu")
if __name__ == "__main__":
    transform = R3MTransform(model_name="resnet18", keys_in=["next_pixels"])
    env1 = TransformedEnv(
        ParallelEnv(16, lambda: MJEnv("kitchen_micro_open-v3", from_pixels=True,
                                     device=device)
                    ),
        transform=transform.eval(),
    )
    assert env1.device == device
    # assert not env1.transform.training
    # assert not env1.transform[-1].convnet.training
    del transform

    env1.reset()
    t0 = time.time()
    print(env1.rollout(max_steps=20))
    print(time.time() - t0)
    t0 = time.time()
    print(env1.rollout(max_steps=20))
    print(time.time() - t0)
    env1.close()
    del env1

    env2 = ParallelEnv(
        16,
        lambda: MJEnv("kitchen_micro_open-v4", device=device)
    )
    assert env2.device == device
    env2.reset()
    t0 = time.time()
    print(env2.rollout(max_steps=20))
    print(time.time() - t0)
    t0 = time.time()
    print(env2.rollout(max_steps=20))
    print(time.time() - t0)
    env2.close()
    del env2