"""
Standalone usage example for compressed EmbeddingGemma models.

This script demonstrates how to load and use the compressed models from HuggingFace Hub
without requiring the full codebase.

For the depth+embedding width compressed model:
  Model ID: Pieces/embeddinggemma-300m-distilled-depth10pct-10-768dim-best

For the width-only compressed model:
  Model ID: Pieces/embeddinggemma-300m-distilled-width10pct-768dim-best
"""

import torch
from typing import List
from sentence_transformers import SentenceTransformer


def encode_texts_simple(model: SentenceTransformer, texts: List[str]) -> torch.Tensor:
    """
    Encode texts using SentenceTransformer model.
    
    Args:
        model: SentenceTransformer model
        texts: List of text strings to encode
        
    Returns:
        Tensor of shape (len(texts), embedding_dim)
    """
    with torch.no_grad():
        embeddings = model.encode(texts, convert_to_tensor=True)
    return embeddings


def compute_similarities(
    query_embeddings: torch.Tensor, tag_embeddings: torch.Tensor
) -> torch.Tensor:
    """
    Compute cosine similarities between queries and tags.
    
    Args:
        query_embeddings: Tensor of shape (num_queries, embedding_dim)
        tag_embeddings: Tensor of shape (num_tags, embedding_dim)
        
    Returns:
        Tensor of shape (num_queries, num_tags) with cosine similarities
    """
    # Normalize embeddings
    query_norm = query_embeddings / (
        query_embeddings.norm(dim=1, keepdim=True) + 1e-8
    )
    tag_norm = tag_embeddings / (tag_embeddings.norm(dim=1, keepdim=True) + 1e-8)
    
    # Compute cosine similarity
    similarities = torch.mm(query_norm, tag_norm.t())
    return similarities


def main():
    """Example usage of compressed models."""
    
    # Choose which model to use
    # Option 1: Depth + Embedding Width Compressed (2 layers, ~33M params)
    model_id_depth_emb = "Pieces/embeddinggemma-300m-distilled-depth10pct-10-768dim-best"
    
    # Option 2: Width Only Compressed (24 layers, ~20M params)
    model_id_width = "Pieces/embeddinggemma-300m-distilled-width10pct-768dim-best"
    
    # Load model (using width-only as example)
    print(f"Loading model: {model_id_width}")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = SentenceTransformer(model_id_width, device=device)
    
    print(f"✅ Model loaded on {device}")
    print(f"   Max sequence length: {model.max_seq_length}")
    print(f"   Embedding dimension: {model.get_sentence_embedding_dimension()}")
    
    # Example queries and tags
    queries = [
        "How to implement authentication in a web application?",
        "What are the best practices for database optimization?",
        "How to deploy a machine learning model to production?",
    ]
    
    tags = [
        "authentication",
        "security",
        "web-development",
        "database",
        "sql",
        "performance",
        "optimization",
        "machine-learning",
        "deployment",
        "production",
        "api",
        "backend",
        "frontend",
        "javascript",
        "python",
    ]
    
    # Encode queries and tags
    print(f"\nEncoding {len(queries)} queries and {len(tags)} tags...")
    query_embeddings = encode_texts_simple(model, queries)
    tag_embeddings = encode_texts_simple(model, tags)
    
    print(f"  Query embeddings shape: {query_embeddings.shape}")
    print(f"  Tag embeddings shape: {tag_embeddings.shape}")
    
    # Compute similarities
    print(f"\nComputing similarities...")
    similarities = compute_similarities(query_embeddings, tag_embeddings)
    print(f"  Similarities shape: {similarities.shape}")
    
    # Display top tags for each query
    print(f"\n{'='*80}")
    print("Retrieval Results")
    print(f"{'='*80}")
    
    for query_idx, query in enumerate(queries):
        top_k = 3
        top_similarities, top_indices = torch.topk(
            similarities[query_idx], k=top_k, dim=0
        )
        
        print(f"\nQuery: {query}")
        for rank, (tag_idx, sim) in enumerate(
            zip(top_indices.cpu().tolist(), top_similarities.cpu().tolist()), start=1
        ):
            print(f"  {rank}. {tags[tag_idx]} (similarity: {sim:.4f})")


if __name__ == "__main__":
    main()