import svetoviz_webgpu as sv

# 1. Prepare 1D data and MaxPool with indices
t = torch.linspace(0, 10, 100)
input_tensor = (torch.sin(t) + torch.randn(100) * 0.2).unsqueeze(0).unsqueeze(0)

# Return_indices=True is mandatory for Unpool
pool1d = nn.MaxPool1d(kernel_size=4, stride=4, return_indices=True)
unpool1d = nn.MaxUnpool1d(kernel_size=4, stride=4)
container = nn.ModuleList([pool1d, unpool1d])

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

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

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