import svetoviz_webgpu as sv

# 1. Define Modules
downscale_factor = 2
# Reverses PixelShuffle: (*, C, H*r, W*r) to (*, C*r^2, H, W)
unshuffle = nn.PixelUnshuffle(downscale_factor=downscale_factor)
shuffle = nn.PixelShuffle(upscale_factor=downscale_factor)

container = nn.ModuleList([unshuffle, shuffle])

def terminal_callback(buffer, message, images, files):
    # Step A: Unshuffle (Spatial to Channels)
    unshuffled = unshuffle(input_tensor)

    # Step B: Shuffle back (Recovery)
    reconstructed = shuffle(unshuffled)

    buffer.send_system_message(f"Input Resolution: {input_tensor.shape[2:]}")
    buffer.send_system_message(f"Unshuffled (Feature Depth): {list(unshuffled.shape)}")

    mse = torch.mean((input_tensor - reconstructed) ** 2)
    buffer.send_system_message(f"Reconstruction Error: {mse.item():.6f}")

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