"""Supervised contrastive loss on existing fixed-crop foreground embeddings."""
import torch
import torch.nn.functional as F

def pixel_contrast(features, target):
    embeddings = features.float().permute(0,2,3,1).reshape(-1,features.shape[1])
    labels = target.reshape(-1)
    selected=[]
    for cls in range(1,19):
        indices = (labels==cls).nonzero().flatten()
        if len(indices):
            selected.append(indices[torch.randperm(len(indices),device=indices.device)[:16]])
    if not selected: return features.sum()*0
    indices=torch.cat(selected)
    embeddings=F.normalize(embeddings[indices],dim=1,eps=1e-4)
    labels=labels[indices]
    logits=embeddings@embeddings.T/.1
    eye=torch.eye(len(indices),device=features.device,dtype=torch.bool)
    positives=(labels[:,None]==labels[None,:])&~eye
    valid=positives.any(1)
    if not valid.any(): return features.sum()*0
    log_probs=logits-logits.masked_fill(eye,float('-inf')).logsumexp(1,keepdim=True)
    return -((log_probs*positives).sum(1)/positives.sum(1).clamp_min(1))[valid].mean()