"""
Модуль для векторизации текстов для кластеризации.
Использует модели из ЛР2: Word2Vec, FastText, GloVe, а также TF-IDF и BM25.
"""

from __future__ import annotations

import os
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import normalize

try:
    from rank_bm25 import BM25Okapi
    BM25_AVAILABLE = True
except ImportError:
    BM25_AVAILABLE = False
    print("⚠️ rank-bm25 не установлен. BM25 недоступен. Установите: pip install rank-bm25")

from gensim.models import Word2Vec, FastText, Doc2Vec
from gensim.utils import simple_preprocess

from src.classical_vectorizers import ClassicalVectorizers, VectorizationConfig


def load_embedding_model(model_path: str):
    """
    Загружает обученную модель эмбеддингов из ЛР2.
    
    Args:
        model_path: Путь к модели
    
    Returns:
        Загруженная модель (Word2Vec, FastText или Doc2Vec)
    """
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Модель не найдена: {model_path}")
    
    # Пробуем загрузить как Word2Vec
    try:
        return Word2Vec.load(model_path)
    except:
        pass
    
    # Пробуем загрузить как FastText
    try:
        return FastText.load(model_path)
    except:
        pass
    
    # Пробуем загрузить как Doc2Vec
    try:
        return Doc2Vec.load(model_path)
    except:
        pass
    
    raise ValueError(f"Не удалось загрузить модель из {model_path}")


def vectorize_tfidf(texts: List[str],
                    max_features: Optional[int] = None,
                    ngram_range: Tuple[int, int] = (1, 2),
                    normalize_vectors: bool = True) -> Tuple[np.ndarray, Any]:
    """
    Векторизация текстов с помощью TF-IDF.
    
    Args:
        texts: Список текстов
        max_features: Максимальное количество признаков
        ngram_range: Диапазон n-грамм
        normalize_vectors: Нормализовать ли векторы (L2)
    
    Returns:
        Матрица векторов и векторизатор
    """
    vectorizer = TfidfVectorizer(
        max_features=max_features,
        ngram_range=ngram_range,
        lowercase=True,
        min_df=1
    )
    
    X = vectorizer.fit_transform(texts).toarray()
    
    if normalize_vectors:
        X = normalize(X, norm='l2')
    
    return X, vectorizer


def vectorize_bm25(texts: List[str],
                  tokenize: bool = True) -> Tuple[np.ndarray, Any]:
    """
    Векторизация текстов с помощью BM25.
    
    Args:
        texts: Список текстов
        tokenize: Токенизировать ли тексты
    
    Returns:
        Матрица векторов и BM25 объект
    """
    if not BM25_AVAILABLE:
        raise ImportError("rank-bm25 не установлен. Установите: pip install rank-bm25")
    
    if tokenize:
        tokenized_texts = [simple_preprocess(text, deacc=False, min_len=1) for text in texts]
    else:
        tokenized_texts = [text.split() for text in texts]
    
    bm25 = BM25Okapi(tokenized_texts)
    
    # Создаем матрицу BM25 для всех документов
    X = np.array([bm25.get_scores(doc) for doc in tokenized_texts])
    
    # Нормализуем
    X = normalize(X, norm='l2')
    
    return X, bm25


def vectorize_with_word2vec(texts: List[str],
                            model: Word2Vec,
                            aggregation: str = "mean",
                            normalize_vectors: bool = True) -> np.ndarray:
    """
    Векторизация текстов с помощью Word2Vec модели из ЛР2.
    
    Args:
        texts: Список текстов
        model: Обученная Word2Vec модель
        aggregation: Метод агрегации (mean, max, sum)
        normalize_vectors: Нормализовать ли векторы (L2)
    
    Returns:
        Матрица векторов документов
    """
    kv = model.wv
    vector_size = kv.vector_size
    vectors = []
    
    for text in texts:
        tokens = simple_preprocess(text, deacc=False, min_len=1)
        word_vectors = []
        
        for token in tokens:
            if token in kv:
                word_vectors.append(kv[token])
        
        if not word_vectors:
            vectors.append(np.zeros(vector_size))
            continue
        
        word_vectors = np.array(word_vectors)
        
        if aggregation == "mean":
            doc_vector = np.mean(word_vectors, axis=0)
        elif aggregation == "max":
            doc_vector = np.max(word_vectors, axis=0)
        elif aggregation == "sum":
            doc_vector = np.sum(word_vectors, axis=0)
        else:
            doc_vector = np.mean(word_vectors, axis=0)
        
        vectors.append(doc_vector)
    
    X = np.array(vectors)
    
    if normalize_vectors:
        X = normalize(X, norm='l2')
    
    return X


def vectorize_with_fasttext(texts: List[str],
                           model: FastText,
                           aggregation: str = "mean",
                           normalize_vectors: bool = True) -> np.ndarray:
    """
    Векторизация текстов с помощью FastText модели из ЛР2.
    
    Args:
        texts: Список текстов
        model: Обученная FastText модель
        aggregation: Метод агрегации (mean, max, sum)
        normalize_vectors: Нормализовать ли векторы (L2)
    
    Returns:
        Матрица векторов документов
    """
    kv = model.wv
    vector_size = kv.vector_size
    vectors = []
    
    for text in texts:
        tokens = simple_preprocess(text, deacc=False, min_len=1)
        word_vectors = []
        
        for token in tokens:
            # FastText может обрабатывать OOV слова
            if token in kv:
                word_vectors.append(kv[token])
            else:
                # Получаем вектор для OOV слова
                word_vectors.append(kv.get_vector(token))
        
        if not word_vectors:
            vectors.append(np.zeros(vector_size))
            continue
        
        word_vectors = np.array(word_vectors)
        
        if aggregation == "mean":
            doc_vector = np.mean(word_vectors, axis=0)
        elif aggregation == "max":
            doc_vector = np.max(word_vectors, axis=0)
        elif aggregation == "sum":
            doc_vector = np.sum(word_vectors, axis=0)
        else:
            doc_vector = np.mean(word_vectors, axis=0)
        
        vectors.append(doc_vector)
    
    X = np.array(vectors)
    
    if normalize_vectors:
        X = normalize(X, norm='l2')
    
    return X


def vectorize_with_doc2vec(texts: List[str],
                          model: Doc2Vec,
                          normalize_vectors: bool = True) -> np.ndarray:
    """
    Векторизация текстов с помощью Doc2Vec модели из ЛР2.
    
    Args:
        texts: Список текстов
        model: Обученная Doc2Vec модель
        normalize_vectors: Нормализовать ли векторы (L2)
    
    Returns:
        Матрица векторов документов
    """
    vectors = []
    
    for text in texts:
        tokens = simple_preprocess(text, deacc=False, min_len=1)
        if tokens:
            vec = model.infer_vector(tokens)
        else:
            vec = np.zeros(model.vector_size)
        vectors.append(vec)
    
    X = np.array(vectors)
    
    if normalize_vectors:
        X = normalize(X, norm='l2')
    
    return X


def vectorize_with_glove(texts: List[str],
                         model_path: str,
                         aggregation: str = "mean",
                         normalize_vectors: bool = True) -> np.ndarray:
    """
    Векторизация текстов с помощью GloVe модели из ЛР2.
    
    Примечание: GloVe обычно хранится в формате текстового файла или через gensim.
    Здесь предполагается, что модель загружена через gensim или аналогичный интерфейс.
    
    Args:
        texts: Список текстов
        model_path: Путь к модели GloVe
        aggregation: Метод агрегации (mean, max, sum)
        normalize_vectors: Нормализовать ли векторы (L2)
    
    Returns:
        Матрица векторов документов
    """
    # Пробуем загрузить как KeyedVectors (если сохранено через gensim)
    try:
        from gensim.models import KeyedVectors
        kv = KeyedVectors.load(model_path)
    except:
        # Если не получилось, пробуем загрузить как Word2Vec (совместимость)
        try:
            model = Word2Vec.load(model_path)
            kv = model.wv
        except:
            raise ValueError(f"Не удалось загрузить GloVe модель из {model_path}")
    
    vector_size = kv.vector_size
    vectors = []
    
    for text in texts:
        tokens = simple_preprocess(text, deacc=False, min_len=1)
        word_vectors = []
        
        for token in tokens:
            if token in kv:
                word_vectors.append(kv[token])
        
        if not word_vectors:
            vectors.append(np.zeros(vector_size))
            continue
        
        word_vectors = np.array(word_vectors)
        
        if aggregation == "mean":
            doc_vector = np.mean(word_vectors, axis=0)
        elif aggregation == "max":
            doc_vector = np.max(word_vectors, axis=0)
        elif aggregation == "sum":
            doc_vector = np.sum(word_vectors, axis=0)
        else:
            doc_vector = np.mean(word_vectors, axis=0)
        
        vectors.append(doc_vector)
    
    X = np.array(vectors)
    
    if normalize_vectors:
        X = normalize(X, norm='l2')
    
    return X


def vectorize_texts(texts: List[str],
                   method: str = "tfidf",
                   model_path: Optional[str] = None,
                   **kwargs) -> Tuple[np.ndarray, Any]:
    """
    Универсальная функция векторизации текстов.
    
    Args:
        texts: Список текстов
        method: Метод векторизации (tfidf, bm25, w2v, fasttext, doc2vec, glove)
        model_path: Путь к модели (для w2v, fasttext, doc2vec, glove)
        **kwargs: Дополнительные параметры
    
    Returns:
        Матрица векторов и объект векторизатора/модели
    """
    method = method.lower()
    
    if method == "tfidf":
        return vectorize_tfidf(texts, **kwargs)
    
    elif method == "bm25":
        return vectorize_bm25(texts, **kwargs)
    
    elif method == "w2v" or method == "word2vec":
        if model_path is None:
            raise ValueError("Для Word2Vec требуется model_path")
        model = load_embedding_model(model_path)
        X = vectorize_with_word2vec(texts, model, **kwargs)
        return X, model
    
    elif method == "fasttext":
        if model_path is None:
            raise ValueError("Для FastText требуется model_path")
        model = load_embedding_model(model_path)
        X = vectorize_with_fasttext(texts, model, **kwargs)
        return X, model
    
    elif method == "doc2vec" or method == "d2v":
        if model_path is None:
            raise ValueError("Для Doc2Vec требуется model_path")
        model = load_embedding_model(model_path)
        X = vectorize_with_doc2vec(texts, model, **kwargs)
        return X, model
    
    elif method == "glove":
        if model_path is None:
            raise ValueError("Для GloVe требуется model_path")
        X = vectorize_with_glove(texts, model_path, **kwargs)
        return X, None
    
    else:
        raise ValueError(f"Неизвестный метод векторизации: {method}")


if __name__ == "__main__":
    # Тестирование
    sample_texts = [
        "Это первый тестовый текст для проверки векторизации",
        "Второй текст содержит другую информацию",
        "Третий текст также используется для тестирования"
    ]
    
    # TF-IDF
    X_tfidf, vectorizer = vectorize_tfidf(sample_texts)
    print(f"TF-IDF векторы: форма {X_tfidf.shape}")
    
    # BM25 (если доступен)
    if BM25_AVAILABLE:
        X_bm25, bm25 = vectorize_bm25(sample_texts)
        print(f"BM25 векторы: форма {X_bm25.shape}")