import svetoviz_webgpu as sv

# 1. Load sample_image.jpg
img = Image.open("sample_image.jpg").convert("L")
width, height = img.size
img_np = np.array(img).astype(np.float32) / 255.0

# Minimal batch input: (B=1, C=2, H, W)
pred = torch.from_numpy(img_np).view(1, 1, height, width).repeat(1, 2, 1, 1)
target = pred * 0.9

# Specific target for CrossEntropy: (B, H, W)
target_cls = torch.randint(0, 2, (1, height, width)).long()

# 2. Containers
standard_losses = nn.ModuleDict({
    "L1Loss": nn.L1Loss(),
    "MSELoss": nn.MSELoss(),
    "SmoothL1Loss": nn.SmoothL1Loss(),
    "HuberLoss": nn.HuberLoss(delta=2.0),
    "CrossEntropyLoss": nn.CrossEntropyLoss(),
    "BCEWithLogitsLoss": nn.BCEWithLogitsLoss(pos_weight=torch.tensor([2.0])),
    "KLDivLoss": nn.KLDivLoss(reduction="batchmean"),
    "PoissonNLLLoss": nn.PoissonNLLLoss(),
    "MultiLabelSoftMarginLoss": nn.MultiLabelSoftMarginLoss()
})

ranking_losses = nn.ModuleDict({
    "MarginRankingLoss": nn.MarginRankingLoss(margin=0.5),
    "CosineEmbeddingLoss": nn.CosineEmbeddingLoss()
})

triplet_losses = nn.ModuleDict({
    "TripletMarginLoss": nn.TripletMarginLoss(margin=1.0),
    "TripletMarginWithDistanceLoss": nn.TripletMarginWithDistanceLoss()
})

losses = nn.ModuleDict({
    "standard": standard_losses,
    "ranking": ranking_losses,
    "triplet": triplet_losses,
})

def terminal_callback(buffer, message, images, files):
    # --- 3. Standard Losses ---
    for name, module in standard_losses.items():
        if name == "CrossEntropyLoss":
            loss = module(pred, target_cls)
        elif name == "KLDivLoss":
            loss = module(pred, target)
        else:
            loss = module(pred, target)
        buffer.send_system_message(f"{name}: {loss.item():.4f}")

    # --- 4. Ranking Losses ---
    feat_dim = 128
    r_input1 = torch.randn(1, feat_dim)
    r_input2 = torch.randn(1, feat_dim)

    for name, module in ranking_losses.items():
        rt = torch.ones(1, 1) if name == "MarginRankingLoss" else torch.ones(1)
        loss = module(r_input1, r_input2, rt)
        buffer.send_system_message(f"{name}: {loss.item():.4f}")

    # --- 5. Triplet Losses ---
    anchor, pos, neg = pred, target, 1.0 - pred
    for name, module in triplet_losses.items():
        loss = module(anchor, pos, neg)
        buffer.send_system_message(f"{name}: {loss.item():.4f}")

sv.pytorch_web(module=losses, terminal_callback=terminal_callback)