"""
MedEmbed Embedding Generator
Medical-specific embedding model for Alzheimer's Disease research

Model: abhinand/MedEmbed-large-v0.1
- Fine-tuned on PubMed Central clinical notes
- Outperforms general embeddings on medical benchmarks
- Free and local (no API costs)
- Apache 2.0 license
"""

import logging
from typing import List, Union

from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)


class MedEmbedEmbedder:
    """
    Wrapper for MedEmbed-large-v0.1 embedding model.

    Features:
    - Medical domain-specific embeddings
    - Batch processing support
    - Automatic normalization
    - GPU acceleration if available
    """

    def __init__(self, model_name: str = "abhinand/MedEmbed-large-v0.1", device: str = None):
        """
        Initialize MedEmbed embedder.

        Args:
            model_name: HuggingFace model identifier
            device: Device to run on ('cuda', 'cpu', or None for auto)
        """
        logger.info(f"Loading MedEmbed model: {model_name}")
        self.model = SentenceTransformer(model_name, device=device)
        self.dimension = self.model.get_sentence_embedding_dimension()
        self.max_seq_length = self.model.max_seq_length

        logger.info("✅ MedEmbed loaded successfully")
        logger.info(f"   Dimension: {self.dimension}")
        logger.info(f"   Max sequence length: {self.max_seq_length}")
        logger.info(f"   Device: {self.model.device}")

    def embed_documents(self, texts: List[str], batch_size: int = 32,
                       show_progress: bool = False) -> List[List[float]]:
        """
        Generate embeddings for a list of documents.

        Args:
            texts: List of text strings to embed
            batch_size: Batch size for processing
            show_progress: Show progress bar

        Returns:
            List of embedding vectors
        """
        if not texts:
            return []

        logger.debug(f"Embedding {len(texts)} documents (batch_size={batch_size})")

        # Generate embeddings
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            convert_to_numpy=True,
            normalize_embeddings=True  # L2 normalization for cosine similarity
        )

        # Convert to list of lists
        return embeddings.tolist()

    def embed_query(self, text: str) -> List[float]:
        """
        Generate embedding for a single query.

        Args:
            text: Query text

        Returns:
            Embedding vector
        """
        logger.debug(f"Embedding query: {text[:100]}...")

        embedding = self.model.encode(
            text,
            convert_to_numpy=True,
            normalize_embeddings=True
        )

        return embedding.tolist()

    def get_dimension(self) -> int:
        """Get embedding dimension."""
        return self.dimension

    def get_max_length(self) -> int:
        """Get maximum sequence length."""
        return self.max_seq_length

    def truncate_text(self, text: str, max_tokens: int = None) -> str:
        """
        Truncate text to fit within model's max sequence length.

        Args:
            text: Input text
            max_tokens: Maximum tokens (defaults to model's max_seq_length)

        Returns:
            Truncated text
        """
        if max_tokens is None:
            max_tokens = self.max_seq_length

        # Rough approximation: 1 token ≈ 4 characters
        max_chars = max_tokens * 4

        if len(text) <= max_chars:
            return text

        logger.warning(f"Truncating text from {len(text)} to {max_chars} chars")
        return text[:max_chars]


def create_medembed_embedder(device: str = None) -> MedEmbedEmbedder:
    """
    Factory function to create MedEmbed embedder.

    Args:
        device: Device to run on ('cuda', 'cpu', or None for auto)

    Returns:
        MedEmbedEmbedder instance
    """
    return MedEmbedEmbedder(device=device)


# Langchain-compatible wrapper
class LangchainMedEmbedEmbeddings:
    """
    Langchain-compatible wrapper for MedEmbed embeddings.

    Usage:
        embeddings = LangchainMedEmbedEmbeddings()
        vector_store = Chroma(
            embedding_function=embeddings,
            persist_directory="./chroma_db"
        )
    """

    def __init__(self, device: str = None):
        """Initialize with optional device specification."""
        self.embedder = MedEmbedEmbedder(device=device)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed a list of documents."""
        return self.embedder.embed_documents(texts)

    def embed_query(self, text: str) -> List[float]:
        """Embed a single query."""
        return self.embedder.embed_query(text)

    def __call__(self, input: Union[str, List[str]]) -> Union[List[float], List[List[float]]]:
        """
        Callable interface for ChromaDB compatibility.
        ChromaDB 0.4+ requires parameter named 'input' not 'text'.

        Args:
            input: Single string or list of strings

        Returns:
            Single embedding or list of embeddings
        """
        if isinstance(input, str):
            return self.embed_query(input)
        else:
            return self.embed_documents(input)


if __name__ == "__main__":
    # Test the embedder
    logging.basicConfig(level=logging.INFO)

    print("=" * 60)
    print("MedEmbed Embedder Test")
    print("=" * 60)

    # Initialize
    embedder = MedEmbedEmbedder()

    # Test single query
    print("\n1. Testing single query embedding...")
    query = "Alzheimer's disease amyloid beta tau protein neurodegeneration"
    query_embedding = embedder.embed_query(query)
    print(f"   Query: {query}")
    print(f"   Embedding dimension: {len(query_embedding)}")
    print(f"   Embedding sample: {query_embedding[:5]}")

    # Test batch documents
    print("\n2. Testing batch document embedding...")
    documents = [
        "Amyloid plaques and neurofibrillary tangles are hallmarks of Alzheimer's disease.",
        "Gut microbiome composition affects cognitive function in elderly patients.",
        "APOE4 is a major genetic risk factor for late-onset Alzheimer's disease.",
    ]
    doc_embeddings = embedder.embed_documents(documents, show_progress=True)
    print(f"   Embedded {len(doc_embeddings)} documents")
    print(f"   Each embedding dimension: {len(doc_embeddings[0])}")

    # Test Langchain wrapper
    print("\n3. Testing Langchain-compatible wrapper...")
    lc_embeddings = LangchainMedEmbedEmbeddings()
    test_embedding = lc_embeddings("Test medical text")
    print(f"   Single text embedding dimension: {len(test_embedding)}")

    batch_embeddings = lc_embeddings(["Text 1", "Text 2"])
    print(f"   Batch embeddings: {len(batch_embeddings)} documents")

    print("\n" + "=" * 60)
    print("✅ All tests passed!")
    print("=" * 60)
