"""
lib/query_expansion.py — V6 Query Expansion Engine

Expands JD skill terms with synonyms, related technologies, and
domain-specific variations. This improves retrieval coverage without
hardcoding every possible skill name.

Example:
  "semantic search" -> ["semantic search", "dense retrieval", "vector search",
                        "ANN", "HNSW", "FAISS", "Pinecone", "Milvus", "Qdrant"]

The expansion happens ONCE at JD parse time, not per-candidate.
Expanded terms are stored in the JD understanding for use by all
downstream modules.
"""
from __future__ import annotations

# Expansion map: canonical term -> expanded synonyms
# These cover common synonyms and related technologies in ML/AI/search
EXPANSION_MAP = {
    # Embeddings
    "embedding": ["embedding", "dense vector", "vector representation",
                  "dense representation", "embedding model"],
    "sentence-transformers": ["sentence-transformers", "sentence transformers",
                              "sbert", "sentence bert"],
    "openai embeddings": ["openai embeddings", "text-embedding-ada",
                          "text-embedding-3", "openai embedding"],
    "bge": ["bge", "bge embedding", "bge-m3", "bge-large", "bge-base"],
    "e5": ["e5", "e5 embedding", "e5-large", "e5-base", "e5-mistral"],

    # Vector DBs
    "pinecone": ["pinecone", "pinecone vector", "pinecone db"],
    "weaviate": ["weaviate", "weaviate vector db"],
    "qdrant": ["qdrant", "qdrant vector"],
    "milvus": ["milvus", "milvus vector db", "zilliz"],
    "faiss": ["faiss", "facebook ai similarity search"],
    "opensearch": ["opensearch", "aws opensearch", "elastic opensearch"],
    "elasticsearch": ["elasticsearch", "elastic search", "es cluster"],
    "vector database": ["vector database", "vector db", "vector store",
                        "vector index", "embedding store", "ann index"],

    # Search/Retrieval
    "semantic search": ["semantic search", "dense retrieval", "vector search",
                        "neural search", "meaning-based search"],
    "hybrid search": ["hybrid search", "hybrid retrieval", "sparse+dense",
                      "bm25+vector", "lexical+semantic"],
    "bm25": ["bm25", "okapi bm25", "tf-idf search", "lexical search",
             "keyword search", "sparse retrieval"],
    "reranking": ["reranking", "re-ranking", "cross-encoder rerank",
                  "reranker", "second-stage ranking"],
    "rag": ["rag", "retrieval-augmented generation", "rag pipeline",
            "rag system", "retrieval augmented"],

    # Ranking
    "ranking": ["ranking", "learning to rank", "ltr", "ranker",
                "ranking model", "rank ordering"],
    "learning to rank": ["learning to rank", "ltr", "lambdarank",
                         "lambdamart", "listwise ranking", "pairwise ranking"],
    "ndcg": ["ndcg", "normalized discounted cumulative gain",
             "ndcg@k", "ndcg evaluation"],

    # LLM/Fine-tuning
    "fine-tuning": ["fine-tuning", "finetuning", "fine tuning", "fine-tune",
                    "supervised fine-tuning", "sft"],
    "lora": ["lora", "low-rank adaptation", "low rank adaptation"],
    "qlora": ["qlora", "quantized lora", "quantized low-rank"],
    "peft": ["peft", "parameter efficient fine-tuning",
             "parameter-efficient fine-tuning"],
    "rlhf": ["rlhf", "reinforcement learning from human feedback"],
    "dpo": ["dpo", "direct preference optimization"],

    # Production
    "production": ["production", "prod", "deployed", "live", "shipped"],
    "at scale": ["at scale", "large scale", "production scale",
                 "internet scale", "high scale"],

    # Evaluation
    "evaluation framework": ["evaluation framework", "eval framework",
                             "evaluation pipeline", "eval pipeline",
                             "ranking evaluation", "quality evaluation"],
    "a/b test": ["a/b test", "ab test", "ab testing", "a/b testing",
                 "online experiment", "online evaluation"],
    "mrr": ["mrr", "mean reciprocal rank"],

    # NLP
    "nlp": ["nlp", "natural language processing", "natural language",
            "text understanding", "language model"],
    "llm": ["llm", "large language model", "large language models",
            "foundation model", "generative ai"],

    # Infrastructure
    "distributed system": ["distributed system", "distributed systems",
                           "distributed computing", "distributed architecture"],
    "low latency": ["low latency", "sub-millisecond", "real-time",
                    "low latency inference", "fast inference"],
}

# Domain-level expansions (applied to domain names from JD parser)
DOMAIN_EXPANSIONS = {
    "search": ["information retrieval", "search engine", "search system",
               "search relevance", "query understanding", "search ranking"],
    "ranking": ["rank ordering", "relevance ranking", "learning to rank",
                "ranking algorithm", "position ranking", "rank model"],
    "embeddings": ["vector representations", "dense vectors", "embedding models",
                   "sentence embeddings", "passage embeddings"],
    "vector_db": ["vector index", "approximate nearest neighbor", "ann",
                  "similarity search", "vector search engine"],
    "llm": ["large language model", "foundation model", "generative model",
            "language model", "transformer model", "gpt", "llama", "mistral"],
    "engineering": ["software engineering", "systems engineering",
                    "backend engineering", "machine learning engineering"],
    "evaluation": ["ranking evaluation", "quality assessment",
                   "relevance evaluation", "search quality"],
    "infrastructure": ["cloud infrastructure", "ml infrastructure",
                       "data infrastructure", "platform engineering"],
    "nlp": ["natural language processing", "text processing",
            "language understanding", "computational linguistics"],
}


def expand_skill(skill: str) -> list[str]:
    """Expand a single skill into a set of synonyms and related terms."""
    s = skill.lower().strip()

    # Check direct match
    if s in EXPANSION_MAP:
        return EXPANSION_MAP[s]

    # Check partial match (e.g., "fine-tun" matches "fine-tuning")
    for key, expansions in EXPANSION_MAP.items():
        if key in s or s in key:
            return expansions

    return [s]  # no expansion, return original


def expand_domain(domain: str) -> list[str]:
    """Expand a domain name into related terms."""
    return DOMAIN_EXPANSIONS.get(domain.lower(), [domain])


def expand_jd_skills(required: dict[str, list[str]],
                     preferred: dict[str, list[str]]) -> dict[str, list[str]]:
    """
    Expand all JD skills with synonyms. Returns a merged expansion dict.

    Each original skill maps to its expanded set. This is used to improve
    retrieval and skill matching without changing the core JD understanding.
    """
    expanded = {}

    for domain, skills in {**required, **preferred}.items():
        for skill in skills:
            expansions = expand_skill(skill)
            # Also add domain-level expansions
            domain_exps = expand_domain(domain)
            all_terms = list(set(expansions + domain_exps))
            expanded[skill] = all_terms

    return expanded


def get_expanded_text(jd_text: str) -> str:
    """
    Create an expanded version of JD text with synonyms appended.
    Used as the query for TF-IDF/SVD similarity computation.
    """
    extra_terms = []
    for skill, expansions in EXPANSION_MAP.items():
        if skill in jd_text.lower():
            extra_terms.extend(expansions)

    # Deduplicate and add to original
    seen = set(jd_text.lower().split())
    new_terms = [t for t in extra_terms if t not in seen]
    return jd_text + " " + " ".join(new_terms)