import svetoviz_webgpu as sv

# 1. Prepare 3D MRI volume
ds = load_dataset("g4m3r/T1w_MRI_Brain_Slices", split="train", streaming=True)

slices = []
for example in ds.skip(2).take(32):
    img = example['image'].convert("L")
    img_np = np.array(img.resize((128, 128))).astype(np.float32) / 255.0
    slices.append(img_np)

# Final Shape: (1, 1, 32, 128, 128)
input_volume = torch.from_numpy(np.stack(slices)).unsqueeze(0).unsqueeze(0)

# 2. Define MaxPool3d
# This pools across Depth, Height, and Width simultaneously.
pool3d = nn.MaxPool3d(kernel_size=6, stride=2)

def terminal_callback(buffer, message, images, files):
    # 3. Process the volume through the module
    output = pool3d(input_volume)

    buffer.send_system_message("Voxel-wise max pooling completed.")
    buffer.send_system_message(f"Input Volume: {list(input_volume.shape)}")
    buffer.send_system_message(f"Pooled Volume: {list(output.shape)}")

# 4. Start the interactive session
sv.pytorch_web(module=pool3d, terminal_callback=terminal_callback)