import numpy as np
import soundfile as sf
import tempfile
from datasets import Dataset, Audio

# Create a stereo audio file
sample_rate = 16000
duration = 1.0
num_samples = int(sample_rate * duration)

left_channel = np.sin(2 * np.pi * 440 * np.linspace(0, duration, num_samples))
right_channel = np.sin(2 * np.pi * 880 * np.linspace(0, duration, num_samples))
stereo_audio = np.stack([left_channel, right_channel], axis=1).astype(np.float32)

temp_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
sf.write(temp_file.name, stereo_audio, sample_rate)

# Create HuggingFace dataset
dataset_dict = {"audio": [temp_file.name]}
ds = Dataset.from_dict(dataset_dict)

# Test with num_channels=2
ds_stereo = ds.cast_column("audio", Audio(num_channels=2))
audio_data = ds_stereo[0]["audio"]

print(f"Original file shape (via soundfile): {sf.read(temp_file.name)[0].shape}")
# Output: (16000, 2) ✓ Stereo

print(f"HF datasets shape with num_channels=2: {audio_data['array'].shape}")
# Output: (16000,) ✗ Mono (should be (2, 16000))