from argparse import ArgumentParser

import torch
from torch.utils.data import DataLoader, Dataset


class CustomerDataset(Dataset):
    def __init__(self, state):
        self.records = [
            ("data/customer_001.pt", [0.10, 0.20, 0.30], 0),
            ("data/customer_002.pt", [0.60, 0.40, 0.90], 1),
            ("data/customer_003.pt", None, 1),
            ("data/customer_004.pt", [0.20, 0.70, 0.50], 0),
        ]
        if state == "fixed":
            self.records[2] = ("data/customer_003.pt", [0.55, 0.80, 0.35], 1)

    def __len__(self):
        return len(self.records)

    def __getitem__(self, index):
        path, values, label = self.records[index]
        if values is None:
            raise FileNotFoundError(f"missing training sample: {path}")
        return torch.tensor(values, dtype=torch.float32), torch.tensor(label)


def main():
    parser = ArgumentParser()
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--state", choices=["broken", "fixed"], default="broken")
    args = parser.parse_args()

    loader = DataLoader(
        CustomerDataset(args.state),
        batch_size=2,
        shuffle=False,
        num_workers=args.workers,
    )

    batch_count = 0
    for batch_features, batch_labels in loader:
        batch_count += 1
        print(
            f"batch {batch_count}: "
            f"features={tuple(batch_features.shape)} "
            f"labels={batch_labels.tolist()}"
        )
    print(f"completed batches: {batch_count}")


if __name__ == "__main__":
    main()