import torch
from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler


labels = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1])
features = torch.arange(len(labels), dtype=torch.float32).unsqueeze(1)
dataset = TensorDataset(features, labels)

class_counts = torch.bincount(labels)
class_weights = 1.0 / class_counts.float()
sample_weights = class_weights[labels]

generator = torch.Generator().manual_seed(42)
sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=24,
    replacement=True,
    generator=generator,
)
loader = DataLoader(dataset, batch_size=6, sampler=sampler)

sampled_labels = []
batch_labels = []
for _, label_batch in loader:
    batch = label_batch.tolist()
    batch_labels.append(batch)
    sampled_labels.extend(batch)

sampled_counts = torch.bincount(torch.tensor(sampled_labels), minlength=2)

print(f"torch_version={torch.__version__}")
print(f"original_counts={class_counts.tolist()}")
print(f"original_minority_share={class_counts[1].item() / len(labels):.2f}")
print(f"class_weights={[round(value, 4) for value in class_weights.tolist()]}")
print(f"sample_weights_first_last={[round(sample_weights[0].item(), 4), round(sample_weights[-1].item(), 4)]}")
print(f"sampled_batches={batch_labels}")
print(f"sampled_counts={sampled_counts.tolist()}")
print(f"sampled_minority_share={sampled_counts[1].item() / len(sampled_labels):.2f}")
print(f"minority_original={class_counts[1].item()}")
print(f"minority_sampled={sampled_counts[1].item()}")
print(f"minority_oversampled={sampled_counts[1].item() > class_counts[1].item()}")