import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

class SigLIPLoss(nn.Module):
    """
    Implementation of the Sigmoid Loss for Language-Image Pre-training (SigLIP).
    Replaces global softmax with pairwise binary classification.
    """
    def __init__(self, t_prime=0.0, b=0.0):
        super().__init__()
        # t_prime and b are learnable parameters
        # t = exp(t_prime) ensures temperature remains positive
        self.t_prime = nn.Parameter(torch.tensor(t_prime))
        self.b = nn.Parameter(torch.tensor(b))

    def forward(self, img_emb, txt_emb):
        # 1. L2 Normalize embeddings to project them onto a unit hypersphere
        z_img = F.normalize(img_emb, p=2, dim=-1) 
        z_txt = F.normalize(txt_emb, p=2, dim=-1) 

        # 2. Compute pairwise similarity matrix
        # logits[i, j] = dot(z_img[i], z_txt[j]) * exp(t_prime) + b
        t = torch.exp(self.t_prime)
        logits = torch.matmul(z_img, z_txt.T) * t + self.b 

        # 3. Create labels: 1 for diagonal (positive), -1 for off-diagonal (negative)
        n = img_emb.shape[0]
        labels = 2 * torch.eye(n, device=img_emb.device) - 1 

        # 4. Compute Sigmoid Loss
        # We use logsigmoid for numerical stability. 
        # labels * logits flips the sign for negatives, treating it as binary classification.
        loss = -F.logsigmoid(labels * logits).sum() / n
        
        return loss

class SimpleDualEncoder(nn.Module):
    def __init__(self, input_dim, embed_dim):
        super().__init__()
        self.image_encoder = nn.Linear(input_dim, embed_dim)
        self.text_encoder = nn.Linear(input_dim, embed_dim)

    def forward(self, img_data, txt_data):
        return self.image_encoder(img_data), self.text_encoder(txt_data)

class MockImageTextDataset(Dataset):
    def __init__(self, num_samples=1000, dim=128):
        self.num_samples = num_samples
        self.dim = dim
        self.base_vectors = torch.randn(num_samples, dim)
        
    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # Create pairs by adding slight noise to the same base vector
        img = self.base_vectors[idx] + torch.randn(self.dim) * 0.1
        txt = self.base_vectors[idx] + torch.randn(self.dim) * 0.1
        return img, txt

# --- Execution ---
if __name__ == '__main__':
    # Hyperparameters
    BATCH_SIZE, EMBED_DIM, INPUT_DIM, EPOCHS = 64, 128, 128, 20
    
    dataset = MockImageTextDataset(num_samples=2000, dim=INPUT_DIM)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

    model = SimpleDualEncoder(INPUT_DIM, EMBED_DIM)
    criterion = SigLIPLoss()
    optimizer = torch.optim.Adam(list(model.parameters()) + list(criterion.parameters()), lr=1e-3)

    model.train()
    for epoch in range(EPOCHS):
        total_loss = 0
        for img_batch, txt_batch in dataloader:
            optimizer.zero_grad()
            img_emb, txt_emb = model(img_batch, txt_batch)
            loss = criterion(img_emb, txt_emb)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        if (epoch + 1) % 5 == 0:
            print(f"Epoch [{epoch+1}/{EPOCHS}] | Avg Loss: {total_loss/len(dataloader):.4f}")

    # Evaluation
    model.eval()
    with torch.no_grad():
        test_img, test_txt = dataset[0]
        img_emb, txt_emb = model(test_img.unsqueeze(0), test_txt.unsqueeze(0))
        similarity = torch.matmul(F.normalize(img_emb, dim=-1), F.normalize(txt_emb, dim=-1).T).item()
        print(f"\nFinal Test Pair Cosine Similarity: {similarity:.4f}")