from pat.models.pat_tokenizer import PATConfig, PATTokenizer, TokenizerType

import torch


if __name__ == "__main__":

    cfg = PATConfig(
        n_tokens=12,
        tokenizer_type=TokenizerType.UAT,
        # n_quantizers=4,
        # vq_type="rvq",
        n_quantizers=1,
        vq_type="vq",
    )

    action_tokenizer = PATTokenizer(cfg)

    # 50, 20, chunk_size, dim
    chunk_size = 50
    action_dim = 20
    actions = torch.randn(2, chunk_size, action_dim)
    codes, _ = action_tokenizer.encode(actions=actions)

    print(actions)
    print(codes)
    print(len(codes))

    decoded_actions = action_tokenizer.decode([codes])
    print(decoded_actions)
