"""NLP embedding and feature extraction utilities."""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer


def build_bow(
    train_texts: list[str],
    test_texts: list[str] | None = None,
    max_features: int = 50000,
    ngram_range: tuple[int, int] = (1, 1),
    **kwargs: Any,
) -> tuple[Any, np.ndarray, np.ndarray | None]:
    """Build Bag-of-Words with CountVectorizer. Returns (vectorizer, X_train, X_test or None)."""
    vec = CountVectorizer(max_features=max_features, ngram_range=ngram_range, **kwargs)
    X_train = vec.fit_transform(train_texts)
    X_test = vec.transform(test_texts) if test_texts is not None else None
    return vec, X_train, X_test


def build_tfidf(
    train_texts: list[str],
    test_texts: list[str] | None = None,
    max_features: int = 50000,
    ngram_range: tuple[int, int] = (1, 2),
    **kwargs: Any,
) -> tuple[Any, np.ndarray, np.ndarray | None]:
    """Build TF-IDF with TfidfVectorizer. Returns (vectorizer, X_train, X_test or None)."""
    vec = TfidfVectorizer(max_features=max_features, ngram_range=ngram_range, **kwargs)
    X_train = vec.fit_transform(train_texts)
    X_test = vec.transform(test_texts) if test_texts is not None else None
    return vec, X_train, X_test


def extract_text_stats(df: pd.DataFrame, text_col: str) -> pd.DataFrame:
    """Extract character count, word count, avg word length, punctuation count, capital ratio."""
    texts = df[text_col].fillna("").astype(str)
    out = pd.DataFrame(index=df.index)
    out["char_count"] = texts.str.len()
    out["word_count"] = texts.str.split().str.len().fillna(0)
    out["avg_word_length"] = out["char_count"] / out["word_count"].replace(0, np.nan)
    out["avg_word_length"] = out["avg_word_length"].fillna(0)
    out["punct_count"] = texts.str.count(r"[^\w\s]")
    out["capital_ratio"] = texts.apply(lambda s: sum(1 for c in s if c.isupper()) / max(len(s), 1))
    return out


def get_sbert_embeddings(
    texts: list[str],
    model_name: str = "all-MiniLM-L6-v2",
    batch_size: int = 32,
    **kwargs: Any,
) -> np.ndarray:
    """Compute sentence embeddings with sentence-transformers. Requires optional dep: sentence-transformers."""
    try:
        from sentence_transformers import SentenceTransformer
    except ImportError as e:
        raise ImportError(
            "get_sbert_embeddings requires 'sentence-transformers'. Install with: pip install sentence-transformers"
        ) from e
    model = SentenceTransformer(model_name, **kwargs)
    return model.encode(texts, batch_size=batch_size, show_progress_bar=False)
