import torch

from pat import (
    ActionSequenceDataset,
    OptimizerConfig,
    TokenizerModelConfig,
    TokenizerTrainer,
    TrainConfig,
    build_tokenizer,
    train_val_split,
)


def test_tokenizer_roundtrip_shapes():
    actions = torch.randn(32, 16, 6)
    ds = ActionSequenceDataset(actions)

    cfg = TokenizerModelConfig(
        action_dim=ds.action_dim,
        horizon=ds.horizon,
        emb_dim=64,
        head_dim=32,
        encoder_depth=1,
        decoder_depth=1,
        num_registers=4,
        latent_levels=[5, 5, 5],
        token_dropout_mode="uniform",
    )
    model = build_tokenizer(cfg)
    normalizer = ds.get_normalizer()

    batch = {"action": normalizer["action"].normalize(actions[:8])}
    loss = model(batch)
    assert loss.ndim == 0

    tokens = model.tokenize(normalizer["action"].normalize(actions[:8]))
    recon = model.detokenize(tokens)
    recon = normalizer["action"].unnormalize(recon)
    assert tokens.shape == (8, cfg.num_registers)
    assert recon.shape == (8, ds.horizon, ds.action_dim)


def test_trainer_runs_one_epoch(tmp_path):
    actions = torch.randn(64, 16, 6)
    ds = ActionSequenceDataset(actions)
    train_set, val_set = train_val_split(ds, val_ratio=0.2, seed=0)

    cfg = TokenizerModelConfig(
        action_dim=ds.action_dim,
        horizon=ds.horizon,
        emb_dim=64,
        head_dim=32,
        encoder_depth=1,
        decoder_depth=1,
        num_registers=4,
        latent_levels=[5, 5, 5],
        token_dropout_mode="uniform",
    )
    model = build_tokenizer(cfg)

    trainer = TokenizerTrainer(
        model,
        train_cfg=TrainConfig(
            num_epochs=1,
            batch_size=8,
            val_batch_size=8,
            num_workers=0,
            val_every=1,
            sample_every=1,
            checkpoint_every=1,
            save_dir=str(tmp_path / "run"),
            use_ema=True,
            mixed_precision="none",
        ),
        optim_cfg=OptimizerConfig(learning_rate=1e-4),
    )
    trainer.fit(train_set, val_set)

    assert (tmp_path / "run" / "checkpoints" / "last.ckpt").exists()
    assert (tmp_path / "run" / "checkpoints" / "final.ckpt").exists()
