src/features/embeddings.py

"""
Embedding cache utilities for pre-computed features.

Handles saving/loading embeddings to disk for fast retrieval.
"""

import torch
from pathlib import Path
from typing import Tuple, List, Optional, Dict
import json


def save_embeddings(
    embeddings: torch.Tensor,
    metadata: Dict,
    output_dir: str,
    filename: Optional[str] = None
) -> Path:
    """
    Save embeddings and metadata to disk.
    
    Args:
        embeddings: Tensor of shape (N, embed_dim)
        metadata: Dict with keys like 'modality', 'sample_ids', 'class_labels'
        output_dir: Directory to save to
        filename: Optional custom filename (without extension)
        
    Returns:
        Path to saved file
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Generate filename if not provided
    if filename is None:
        modality = metadata.get("modality", "unknown")
        n_samples = embeddings.shape[0]
        filename = f"{modality}_embeddings_{n_samples}"
    
    # Save embeddings tensor
    embeddings_path = output_dir / f"{filename}.pt"
    torch.save(embeddings, embeddings_path)
    
    # Save metadata as JSON
    metadata_path = output_dir / f"{filename}_metadata.json"
    
    # Convert tensors in metadata to lists for JSON serialization
    serializable_metadata = {}
    for key, value in metadata.items():
        if isinstance(value, torch.Tensor):
            serializable_metadata[key] = value.tolist()
        else:
            serializable_metadata[key] = value
    
    with open(metadata_path, "w") as f:
        json.dump(serializable_metadata, f, indent=2)
    
    return embeddings_path


def load_embeddings(
    cache_path: str
) -> Tuple[torch.Tensor, Dict]:
    """
    Load embeddings and metadata from disk.
    
    Args:
        cache_path: Path to .pt embeddings file
        
    Returns:
        (embeddings tensor, metadata dict)
    """
    cache_path = Path(cache_path)
    
    # Load embeddings
    embeddings = torch.load(cache_path, weights_only=True)
    
    # Load metadata if exists
    metadata_path = cache_path.with_name(
        cache_path.stem + "_metadata.json"
    )
    
    metadata = {}
    if metadata_path.exists():
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
    
    return embeddings, metadata


def get_cache_path(
    output_dir: str,
    modality: str,
    split: str = "gallery",
    embed_dim: int = 768
) -> Path:
    """
    Generate standard cache file path.
    
    Args:
        output_dir: Base output directory
        modality: Modality type
        split: Dataset split (query/gallery)
        embed_dim: Embedding dimension
        
    Returns:
        Path object for cache file
    """
    output_dir = Path(output_dir)
    filename = f"{modality}_{split}_embeddings.pt"
    return output_dir / filename


def verify_embeddings(
    embeddings: torch.Tensor,
    expected_dim: Optional[int] = None,
    l2_normalized: bool = True
) -> bool:
    """
    Verify embeddings are valid.
    
    Args:
        embeddings: Embedding tensor
        expected_dim: Expected embedding dimension
        l2_normalized: Whether embeddings should be L2-normalized
        
    Returns:
        True if valid
    """
    if embeddings.dim() != 2:
        print(f"Expected 2D tensor, got {embeddings.dim()}D")
        return False
    
    if expected_dim is not None and embeddings.shape[1] != expected_dim:
        print(f"Expected dim {expected_dim}, got {embeddings.shape[1]}")
        return False
    
    if l2_normalized:
        norms = torch.norm(embeddings, dim=1)
        # Check if norms are close to 1 (allowing for floating point)
        if not torch.allclose(norms, torch.ones_like(norms), atol=1e-3):
            print(f"Embeddings not L2-normalized. Norms: {norms[:5]}")
            return False
    
    return True


# Self-check
if __name__ == "__main__":
    import tempfile
    
    print("Testing embedding cache utilities...")
    
    # Create dummy embeddings
    embeddings = torch.randn(100, 768)
    embeddings = torch.nn.functional.normalize(embeddings, dim=1)  # L2 normalize
    
    metadata = {
        "modality": "optical",
        "sample_ids": list(range(100)),
        "class_labels": [i % 10 for i in range(100)],
    }
    
    # Test save/load roundtrip
    with tempfile.TemporaryDirectory() as tmpdir:
        # Save
        save_path = save_embeddings(embeddings, metadata, tmpdir, "test")
        print(f"Saved to: {save_path}")
        
        # Load
        loaded_embeddings, loaded_metadata = load_embeddings(save_path)
        print(f"Loaded embeddings shape: {loaded_embeddings.shape}")
        print(f"Loaded metadata keys: {list(loaded_metadata.keys())}")
        
        # Verify
        assert torch.allclose(embeddings, loaded_embeddings), "Embeddings mismatch!"
        assert loaded_metadata["modality"] == "optical", "Metadata mismatch!"
        
        # Test verify
        assert verify_embeddings(loaded_embeddings, expected_dim=768), "Verification failed!"
        
    print("\nEmbedding cache test passed!")