"""
Модуль для интерпретации моделей классификации: SHAP, LIME, важность признаков,
визуализация внимания для нейросетей и трансформеров.
"""

from __future__ import annotations

from typing import List, Dict, Any, Optional, Tuple
import numpy as np
import pandas as pd

try:
    import shap
    SHAP_AVAILABLE = True
except ImportError:
    SHAP_AVAILABLE = False
    print("⚠️ SHAP не установлен. Установите: pip install shap")

try:
    from lime import lime_text
    from lime.lime_text import LimeTextExplainer
    LIME_AVAILABLE = True
except ImportError:
    LIME_AVAILABLE = False
    print("⚠️ LIME не установлен. Установите: pip install lime")

try:
    import matplotlib.pyplot as plt
    import seaborn as sns
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    print("⚠️ Matplotlib не установлен. Визуализация недоступна.")


def get_feature_importance_linear(model, feature_names: Optional[List[str]] = None) -> pd.DataFrame:
    """
    Извлекает важность признаков для линейных моделей (LR, SVM).
    
    Args:
        model: Обученная модель
        feature_names: Названия признаков
    
    Returns:
        DataFrame с важностью признаков
    """
    if hasattr(model, 'coef_'):
        coef = model.coef_
        if len(coef.shape) > 1:
            # Многоклассовая классификация - берем среднее по классам
            importance = np.abs(coef).mean(axis=0)
        else:
            importance = np.abs(coef)
        
        if feature_names is None:
            feature_names = [f"Признак {i}" for i in range(len(importance))]
        
        df = pd.DataFrame({
            "Признак": feature_names,
            "Важность": importance
        }).sort_values("Важность", ascending=False)
        
        return df
    
    return pd.DataFrame()


def get_feature_importance_tree(model, feature_names: Optional[List[str]] = None) -> pd.DataFrame:
    """
    Извлекает важность признаков для tree-based моделей (RF, XGBoost, etc.).
    
    Args:
        model: Обученная модель
        feature_names: Названия признаков
    
    Returns:
        DataFrame с важностью признаков
    """
    if hasattr(model, 'feature_importances_'):
        importance = model.feature_importances_
        
        if feature_names is None:
            feature_names = [f"Признак {i}" for i in range(len(importance))]
        
        df = pd.DataFrame({
            "Признак": feature_names,
            "Важность": importance
        }).sort_values("Важность", ascending=False)
        
        return df
    
    return pd.DataFrame()


def get_tfidf_important_words(vectorizer, model, class_idx: int = 0, top_k: int = 20) -> pd.DataFrame:
    """
    Извлекает наиболее важные слова для TF-IDF векторизации.
    
    Args:
        vectorizer: Обученный векторизатор
        model: Обученная модель
        class_idx: Индекс класса
        top_k: Количество топ-слов
    
    Returns:
        DataFrame с важными словами
    """
    if not hasattr(model, 'coef_'):
        return pd.DataFrame()
    
    coef = model.coef_[class_idx] if len(model.coef_.shape) > 1 else model.coef_
    
    if hasattr(vectorizer, 'get_feature_names_out'):
        feature_names = vectorizer.get_feature_names_out()
    elif hasattr(vectorizer, 'get_feature_names'):
        feature_names = vectorizer.get_feature_names()
    else:
        return pd.DataFrame()
    
    # Сортируем по важности
    indices = np.argsort(np.abs(coef))[-top_k:][::-1]
    
    df = pd.DataFrame({
        "Слово": [feature_names[i] for i in indices],
        "Коэффициент": [coef[i] for i in indices],
        "Абсолютное значение": [np.abs(coef[i]) for i in indices]
    })
    
    return df


def explain_with_shap(model, X: np.ndarray, 
                     feature_names: Optional[List[str]] = None,
                     max_samples: int = 100) -> Optional[shap.Explanation]:
    """
    Объяснение предсказаний модели с помощью SHAP.
    
    Args:
        model: Обученная модель с методом predict_proba
        X: Признаки для объяснения
        feature_names: Названия признаков
        max_samples: Максимальное количество образцов для объяснения
    
    Returns:
        SHAP Explanation объект или None
    """
    if not SHAP_AVAILABLE:
        print("SHAP не установлен. Установите: pip install shap")
        return None
    
    # Ограничиваем количество образцов для производительности
    if len(X) > max_samples:
        indices = np.random.choice(len(X), max_samples, replace=False)
        X_sample = X[indices]
    else:
        X_sample = X
    
    try:
        # Создаем explainer в зависимости от типа модели
        if hasattr(model, 'predict_proba'):
            explainer = shap.Explainer(model, X_sample)
        else:
            # Для моделей без predict_proba используем KernelExplainer
            explainer = shap.KernelExplainer(model.predict, X_sample)
        
        shap_values = explainer(X_sample)
        
        if feature_names is not None:
            shap_values.feature_names = feature_names
        
        return shap_values
    
    except Exception as e:
        print(f"Ошибка при создании SHAP объяснения: {e}")
        return None


def explain_with_lime_text(model, texts: List[str],
                          vectorizer: Any,
                          class_names: Optional[List[str]] = None,
                          num_features: int = 10) -> List[Dict[str, Any]]:
    """
    Объяснение предсказаний модели с помощью LIME для текста.
    
    Args:
        model: Обученная модель
        texts: Тексты для объяснения
        vectorizer: Векторизатор текстов
        class_names: Названия классов
        num_features: Количество важных признаков для показа
    
    Returns:
        Список объяснений для каждого текста
    """
    if not LIME_AVAILABLE:
        print("LIME не установлен. Установите: pip install lime")
        return []
    
    explainer = LimeTextExplainer(class_names=class_names)
    
    def predict_proba_wrapper(texts_list):
        """Обертка для predict_proba с векторизацией."""
        X = vectorizer.transform(texts_list)
        if hasattr(model, 'predict_proba'):
            return model.predict_proba(X)
        else:
            # Для моделей без predict_proba
            predictions = model.predict(X)
            # Создаем псевдо-вероятности
            proba = np.zeros((len(predictions), len(np.unique(predictions))))
            for i, pred in enumerate(predictions):
                proba[i, pred] = 1.0
            return proba
    
    explanations = []
    for text in texts:
        try:
            explanation = explainer.explain_instance(
                text,
                predict_proba_wrapper,
                num_features=num_features
            )
            
            # Извлекаем важные слова
            exp_list = explanation.as_list()
            explanations.append({
                "text": text,
                "important_words": exp_list,
                "prediction": explanation.predict_proba.argmax() if hasattr(explanation, 'predict_proba') else None
            })
        except Exception as e:
            print(f"Ошибка при объяснении текста: {e}")
            explanations.append({
                "text": text,
                "important_words": [],
                "prediction": None
            })
    
    return explanations


def visualize_attention_weights(attention_weights: np.ndarray,
                               tokens: List[str],
                               save_path: Optional[str] = None) -> None:
    """
    Визуализация весов внимания для трансформерных моделей.
    
    Args:
        attention_weights: Матрица весов внимания (n_heads, seq_len, seq_len) или (seq_len, seq_len)
        tokens: Список токенов
        save_path: Путь для сохранения изображения
    """
    if not MATPLOTLIB_AVAILABLE:
        print("Matplotlib не установлен. Визуализация недоступна.")
        return
    
    # Если несколько голов внимания, усредняем
    if len(attention_weights.shape) == 3:
        attention_weights = attention_weights.mean(axis=0)
    
    # Ограничиваем длину для визуализации
    max_len = min(50, len(tokens))
    attention_weights = attention_weights[:max_len, :max_len]
    tokens = tokens[:max_len]
    
    plt.figure(figsize=(12, 10))
    sns.heatmap(
        attention_weights,
        xticklabels=tokens,
        yticklabels=tokens,
        cmap='Blues',
        cbar=True
    )
    plt.title("Визуализация весов внимания")
    plt.xlabel("Токены")
    plt.ylabel("Токены")
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def analyze_error_cases(y_true: np.ndarray, y_pred: np.ndarray,
                       texts: Optional[List[str]] = None,
                       top_k: int = 10) -> pd.DataFrame:
    """
    Анализ случаев, где модель ошибается.
    
    Args:
        y_true: Истинные метки
        y_pred: Предсказанные метки
        texts: Тексты (опционально)
        top_k: Количество примеров для показа
    
    Returns:
        DataFrame с примерами ошибок
    """
    errors = y_true != y_pred
    error_indices = np.where(errors)[0]
    
    if len(error_indices) == 0:
        return pd.DataFrame({"Сообщение": ["Ошибок не найдено"]})
    
    # Ограничиваем количество
    if len(error_indices) > top_k:
        error_indices = np.random.choice(error_indices, top_k, replace=False)
    
    results = []
    for idx in error_indices:
        result = {
            "Индекс": int(idx),
            "Истинный класс": int(y_true[idx]),
            "Предсказанный класс": int(y_pred[idx])
        }
        if texts is not None:
            result["Текст"] = texts[idx][:200] + "..." if len(texts[idx]) > 200 else texts[idx]
        results.append(result)
    
    return pd.DataFrame(results)


if __name__ == "__main__":
    # Тестирование
    from sklearn.datasets import make_classification
    from sklearn.linear_model import LogisticRegression
    from sklearn.feature_extraction.text import TfidfVectorizer
    
    # Создаем тестовые данные
    texts = [
        "Это положительный отзыв о продукте",
        "Отрицательный отзыв не понравилось",
        "Нейтральный отзыв нормально",
    ] * 10
    
    vectorizer = TfidfVectorizer()
    X = vectorizer.fit_transform(texts).toarray()
    y = np.array([0, 1, 2] * 10)
    
    # Обучение модели
    model = LogisticRegression(max_iter=1000, random_state=42)
    model.fit(X, y)
    
    # Важность признаков
    feature_importance = get_feature_importance_linear(model)
    print("Важность признаков (топ-10):")
    print(feature_importance.head(10))
    
    # Важные слова для TF-IDF
    important_words = get_tfidf_important_words(vectorizer, model, class_idx=0, top_k=10)
    print("\nВажные слова для класса 0:")
    print(important_words)
    
    # SHAP (если доступен)
    if SHAP_AVAILABLE:
        shap_values = explain_with_shap(model, X[:5], max_samples=5)
        if shap_values is not None:
            print("\nSHAP объяснение создано успешно")
    
    # LIME (если доступен)
    if LIME_AVAILABLE:
        lime_explanations = explain_with_lime_text(model, texts[:3], vectorizer)
        print(f"\nLIME объяснения: {len(lime_explanations)} создано")