"""Koine-Greek-BERT integration — semantic search for Greek NT texts.

OPTIONAL experimental module. Requires manual installation of dependencies:

    pip install transformers torch

Reference: Koine-Greek-BERT v1.0 (ABeZet, 2026, MIT)
  - First polytonic Koine-specific BERT model on HuggingFace
  - 50K vocabulary, zero-shot on NT verses
  - Enables semantic similarity, verse clustering, intelligent Greek search

Without the dependencies, ``greek_bert_available()`` returns ``False`` and
``GreekBERT`` raises ``ImportError`` on construction.
"""

from __future__ import annotations

import logging

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Lazy import guard — module is importable even without transformers/torch
# ---------------------------------------------------------------------------
_GREEK_BERT_AVAILABLE = False
try:
    import torch
    import torch.nn.functional as F
    from transformers import AutoTokenizer, AutoModel

    _GREEK_BERT_AVAILABLE = True
except ImportError:
    pass

MODEL_NAME = "ABeZet/koine-greek-bert"


class GreekBERT:
    """Koine-Greek-BERT encoder for Greek NT semantic search.

    Usage:

        model = GreekBERT()
        emb = model.encode("Ἐν ἀρχῇ ἦν ὁ λόγος")
        sim = model.similarity("ἀγάπη", "φιλία")
        results = model.find_similar_verses(
            "μακάριοι οἱ πτωχοὶ τῷ πνεύματι",
            ["verse1 text", "verse2 text"],
            top_k=3,
        )
    """

    def __init__(self) -> None:
        if not _GREEK_BERT_AVAILABLE:
            raise ImportError(
                "Koine-Greek-BERT requires `transformers` and `torch`. "
                "Install with: pip install transformers torch"
            )
        logger.info("Loading Koine-Greek-BERT from HuggingFace (%s)...", MODEL_NAME)
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        self.model = AutoModel.from_pretrained(MODEL_NAME)
        self.model.eval()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def encode(self, text: str) -> list[float]:
        """Encode a single Greek text into a fixed-length embedding vector.

        Uses mean-pooling over the last hidden state of the BERT model
        (CLS token excluded from pooling for better sentence semantics).

        Args:
            text: Greek text to encode (single verse, phrase, or word).

        Returns:
            list[float]: Embedding vector (dimension depends on model,
            typically 768).
        """
        encoded = self.tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512,
        )
        with torch.no_grad():
            outputs = self.model(**encoded)

        # Mean-pool over token dimension (exclude [CLS] at position 0 if
        # multiple tokens, but for short inputs mean pooling over all
        # non-padding tokens is the standard sentence-BERT approach).
        attention_mask = encoded["attention_mask"]  # (1, seq_len)
        token_embeddings = outputs.last_hidden_state  # (1, seq_len, hidden)

        # Expand mask to hidden dimension for correct averaging
        mask_expanded = (
            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        )
        sum_embeddings = (token_embeddings * mask_expanded).sum(dim=1)
        sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)
        mean_pooled = sum_embeddings / sum_mask  # (1, hidden)

        return mean_pooled.squeeze(0).tolist()

    def similarity(self, text1: str, text2: str) -> float:
        """Compute cosine similarity between two Greek texts.

        Args:
            text1: First Greek text.
            text2: Second Greek text.

        Returns:
            float: Cosine similarity in [-1, 1] (1 = identical meaning,
            -1 = opposite, 0 = orthogonal).
        """
        emb1 = self._encode_tensor(text1)
        emb2 = self._encode_tensor(text2)
        sim = F.cosine_similarity(emb1, emb2).item()
        return float(sim)

    def find_similar_verses(
        self,
        query_text: str,
        candidates: list[str],
        top_k: int = 5,
    ) -> list[tuple[str, float]]:
        """Find the most semantically similar verses from a candidate list.

        Encodes the query once, then scores each candidate against it
        using cosine similarity.

        Args:
            query_text: Greek query verse/phrase.
            candidates: List of Greek verse texts to search.
            top_k: Number of top results to return (default 5).

        Returns:
            list[tuple[str, float]]: Sorted list of (text, score) pairs,
            highest similarity first.
        """
        query_emb = self._encode_tensor(query_text)
        scored: list[tuple[str, float]] = []

        for candidate in candidates:
            cand_emb = self._encode_tensor(candidate)
            score = F.cosine_similarity(query_emb, cand_emb).item()
            scored.append((candidate, float(score)))

        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _encode_tensor(self, text: str) -> torch.Tensor:
        """Encode text and return a (1, hidden) tensor for internal use."""
        encoded = self.tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512,
        )
        with torch.no_grad():
            outputs = self.model(**encoded)

        attention_mask = encoded["attention_mask"]
        token_embeddings = outputs.last_hidden_state

        mask_expanded = (
            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        )
        sum_embeddings = (token_embeddings * mask_expanded).sum(dim=1)
        sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)
        return sum_embeddings / sum_mask  # (1, hidden)


# ---------------------------------------------------------------------------
# Module-level availability check
# ---------------------------------------------------------------------------


def greek_bert_available() -> bool:
    """Check whether Koine-Greek-BERT dependencies are installed.

    Returns:
        bool: ``True`` if ``transformers`` and ``torch`` are importable,
        ``False`` otherwise. Safe to call any time — never raises.
    """
    return _GREEK_BERT_AVAILABLE
