"""MNIST job module.

Exports make_model() and make_dataloader() for the scheduler + worker.
Falls back to synthetic data if torchvision is not installed.
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Subset, TensorDataset


class MNISTNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)


def make_model() -> nn.Module:
    return MNISTNet()


def make_dataloader(
    rank: int, world_size: int, batch_size: int = 64
) -> DataLoader:
    try:
        from torchvision import datasets, transforms

        dataset = datasets.MNIST(
            "/tmp/mnist_data",
            train=True,
            download=True,
            transform=transforms.ToTensor(),
        )
    except (ImportError, Exception):
        # synthetic fallback — same shape as MNIST
        n = 4096
        x = torch.randn(n, 1, 28, 28)
        y = torch.randint(0, 10, (n,))
        dataset = TensorDataset(x, y)

    # shard by rank: interleaved assignment
    indices = list(range(rank, len(dataset), world_size))
    subset = Subset(dataset, indices)

    return DataLoader(subset, batch_size=batch_size, shuffle=True, drop_last=True)