import svetoviz_webgpu as sv

# 1. Prepare 2D image data
img = Image.open("sample_image.jpg").convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
input_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0)

# 2. Define Pooling and Unpooling
pool2d = nn.MaxPool2d(kernel_size=6, stride=4, return_indices=True)
unpool2d = nn.MaxUnpool2d(kernel_size=6, stride=4)
container = nn.ModuleList([pool2d, unpool2d])

def terminal_callback(buffer, message, images, files):
    # Forward pass to get indices
    pooled, indices = pool2d(input_tensor)

    # FIX: Pass the original input size to unpool2d
    output = unpool2d(pooled, indices, output_size=input_tensor.shape)

    buffer.send_system_message(f"Input Shape: {list(input_tensor.shape)}")
    buffer.send_system_message(f"Pooled Shape: {list(pooled.shape)}")
    buffer.send_system_message(f"Unpooled Shape: {list(output.shape)}")

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