import torch
import torch.nn as nn
from torch.utils.data import Dataset

from torchlitex.launcher import DistributedConfig, launch
from torchlitex.trainer import Trainer


class TinyDataset(Dataset):
    def __len__(self):
        return 8

    def __getitem__(self, idx):
        x = torch.ones(4) * idx
        y = torch.tensor([float(idx)], dtype=torch.float32)
        return x, y


def _train_fn(rank, world_size, ckpt_path):
    torch.manual_seed(0)
    dataset = TinyDataset()
    model = nn.Linear(4, 1)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    trainer = Trainer(
        model,
        dataset,
        loss_fn,
        optimizer,
        scaler=None,
        amp_enabled=False,
        log_every=None,
    )
    trainer.ddp_train_loop(rank, world_size, batch_size=2, epochs=1, ckpt_path=ckpt_path, num_workers=0)


def test_launcher_gloo_smoke(tmp_path):
    ckpt_path = tmp_path / "ckpt.pt"
    cfg = DistributedConfig(gpus=2, backend="gloo")

    launch(_train_fn, cfg, ckpt_path=ckpt_path)

    assert ckpt_path.exists()
