import svetoviz_webgpu as sv

# 1. Define Similarity Parameters
# dim=1: Compare across the feature dimension
cos = nn.CosineSimilarity(dim=1, eps=1e-6)

# 2. Prepare Input Tensors (Batch=1, Features=128)
input1 = torch.randn(1, 128)
# Vector B: Slightly modified version of A (high similarity)
input2 = input1 + (torch.randn(1, 128) * 0.2)
# Vector C: Independent random vector (low similarity)
input3 = torch.randn(1, 128)

def terminal_callback(buffer, message, images, files):
    # 3. Calculate similarities
    sim_high = cos(input1, input2)
    sim_low = cos(input1, input3)

    buffer.send_system_message(f"Similarity (Related): {sim_high.item():.4f}")
    buffer.send_system_message(f"Similarity (Random): {sim_low.item():.4f}")
    buffer.send_system_message("Logic: Closer to 1.0 means vectors are more aligned in orientation.")

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