from argparse import ArgumentParser

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


class ChurnDataset(Dataset):
    def __init__(self, feature_rows, label_rows):
        if len(feature_rows) != len(label_rows):
            raise ValueError("feature_rows and label_rows must have the same length")
        self.features = torch.as_tensor(feature_rows, dtype=torch.float32)
        self.labels = torch.as_tensor(label_rows, dtype=torch.long)

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

    def __getitem__(self, index):
        return {
            "features": self.features[index],
            "label": self.labels[index],
        }


def build_dataset():
    feature_rows = [
        [22, 0, 0.10],
        [36, 1, 0.75],
        [41, 0, 0.35],
        [58, 1, 0.92],
    ]
    label_rows = [0, 1, 0, 1]
    return ChurnDataset(feature_rows, label_rows)


def print_sample():
    dataset = build_dataset()
    sample = dataset[1]
    print(f"dataset_length={len(dataset)}")
    print(f"sample_keys={list(sample.keys())}")
    print(f"sample_features={sample['features'].tolist()}")
    print(f"sample_label={sample['label'].item()}")


def print_batch():
    dataset = build_dataset()
    loader = DataLoader(dataset, batch_size=2, shuffle=False)
    batch = next(iter(loader))
    print(f"batch_features_shape={tuple(batch['features'].shape)}")
    print(f"batch_labels={batch['label'].tolist()}")


def main():
    parser = ArgumentParser()
    parser.add_argument("--check", choices=["sample", "batch"], default="sample")
    args = parser.parse_args()

    if args.check == "sample":
        print_sample()
    else:
        print_batch()


if __name__ == "__main__":
    main()