import os

import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler


def main():
    dist.init_process_group("gloo")

    rank = dist.get_rank()
    world_size = dist.get_world_size()
    local_rank = int(os.environ["LOCAL_RANK"])

    torch.manual_seed(17)
    model = nn.Linear(2, 1)
    ddp_model = DistributedDataParallel(model)

    features = torch.tensor(
        [
            [0.0, 0.0],
            [1.0, 1.0],
            [2.0, 2.0],
            [3.0, 3.0],
        ],
        dtype=torch.float32,
    )
    targets = torch.tensor([[0.0], [2.0], [4.0], [6.0]], dtype=torch.float32)
    dataset = TensorDataset(features, targets)
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=False,
    )
    loader = DataLoader(dataset, batch_size=2, sampler=sampler)

    batch_features, batch_targets = next(iter(loader))
    sample_ids = [int(value) for value in batch_features[:, 0].tolist()]

    optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.1)
    loss_fn = nn.MSELoss()

    optimizer.zero_grad()
    loss = loss_fn(ddp_model(batch_features), batch_targets)
    loss.backward()
    optimizer.step()

    weight = ddp_model.module.weight.detach()[0, 0].item()
    message = (
        f"rank={rank} local_rank={local_rank} world_size={world_size} "
        f"samples={sample_ids} loss={loss.item():.4f} weight={weight:.4f}"
    )

    gathered = [None] * world_size if rank == 0 else None
    dist.gather_object(message, gathered, dst=0)

    if rank == 0:
        for line in gathered:
            print(line)
        print("ddp_step_complete=True")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()