src/ml_pipeline/visualization/model_visualizer.py

import os
import traceback
from typing import Dict, List
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
from scipy.stats import norm
import logging
from pathlib import Path

import shap

from src.ml_pipeline.utils.validators import SHAP_AVAILABLE


from ..config.ml_config import MLConfig

logger = logging.getLogger(__name__)



# ============================================
# КЛАСС 8: ВИЗУАЛИЗАЦИЯ И ДИАГНОСТИКА
# ============================================

class ModelVisualizer:
    """Расширенный класс для визуализации результатов"""

    def __init__(self, config: MLConfig):
        self.config = config
        self.plots_created = {}

        self.texts = {
            'ru': {
                'actual': 'Фактические значения',
                'predicted': 'Предсказанные значения',
                'residuals': 'Остатки',
                'frequency': 'Частота',
                'density': 'Плотность',
                'error': 'Ошибка',
                'importance': 'Важность признака',
                'features': 'Признаки',
                'metrics': 'Метрики',
                'comparison': 'Сравнение моделей',
                'training': 'Тренировочная',
                'validation': 'Валидационная',
                'test': 'Тестовая'
            },
            'en': {
                'actual': 'Actual values',
                'predicted': 'Predicted values',
                'residuals': 'Residuals',
                'frequency': 'Frequency',
                'density': 'Density',
                'error': 'Error',
                'importance': 'Feature importance',
                'features': 'Features',
                'metrics': 'Metrics',
                'comparison': 'Model comparison',
                'training': 'Training',
                'validation': 'Validation',
                'test': 'Test'
            }
        }

        self.text = self.texts.get(self.config.plot_language, self.texts['ru'])

    def get_translation(self, key: str) -> str:
        return self.text.get(key, key)

    def plot_predictions_vs_actual(self, y_true, y_pred, model_name: str,
                                  dataset_name: str = 'test', save_path: str = None):
        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/predictions/{model_name}_{dataset_name}_predictions.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        fig, axes = plt.subplots(1, 2, figsize=(16, 7))

        scatter = axes[0].scatter(y_true, y_pred, alpha=0.6, s=30,
                                 c=np.abs(y_true - y_pred), cmap='viridis',
                                 edgecolors='black', linewidth=0.5)

        min_val = min(y_true.min(), y_pred.min())
        max_val = max(y_true.max(), y_pred.max())
        axes[0].plot([min_val, max_val], [min_val, max_val], 'r--',
                    linewidth=2, alpha=0.7, label='Идеальная линия')

        try:
            from scipy import stats
            slope, intercept, r_value, p_value, std_err = stats.linregress(y_true, y_pred)
            x_line = np.linspace(min_val, max_val, 100)
            y_line = slope * x_line + intercept
            axes[0].plot(x_line, y_line, 'g-', linewidth=2, alpha=0.7,
                        label=f'Линия регрессии (R²={r_value**2:.3f})')
        except:
            pass

        axes[0].set_xlabel(self.get_translation('actual'))
        axes[0].set_ylabel(self.get_translation('predicted'))
        axes[0].set_title(f'{model_name}: {self.get_translation("predicted")} vs {self.get_translation("actual")} ({dataset_name})')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)

        cbar = plt.colorbar(scatter, ax=axes[0])
        cbar.set_label('Абсолютная ошибка', rotation=270, labelpad=15)

        errors = y_pred - y_true
        axes[1].hist(errors, bins=50, edgecolor='black', alpha=0.7, density=True)

        x = np.linspace(errors.min(), errors.max(), 100)
        normal_pdf = norm.pdf(x, errors.mean(), errors.std())
        axes[1].plot(x, normal_pdf, 'r-', linewidth=2, alpha=0.8,
                    label='Нормальное распределение')

        axes[1].axvline(x=0, color='green', linestyle='--', linewidth=2,
                       alpha=0.7, label='Нулевая ошибка')
        axes[1].axvline(x=errors.mean(), color='blue', linestyle=':',
                       linewidth=2, alpha=0.7, label=f'Средняя ошибка ({errors.mean():.2f})')

        axes[1].set_xlabel(self.get_translation('error'))
        axes[1].set_ylabel(self.get_translation('density'))
        axes[1].set_title(f'{model_name}: Распределение ошибок ({dataset_name})')
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)

        stats_text = f'''
        Средняя ошибка: {errors.mean():.2f}
        СКО ошибки: {errors.std():.2f}
        Медиана ошибки: {np.median(errors):.2f}
        Min ошибка: {errors.min():.2f}
        Max ошибка: {errors.max():.2f}
        '''

        axes[1].text(0.02, 0.98, stats_text, transform=axes[1].transAxes,
                    fontsize=9, verticalalignment='top',
                    bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

        plt.suptitle(f'Результаты модели: {model_name} ({dataset_name})', fontsize=16)
        plt.tight_layout()

        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close(fig)

        self.plots_created[f'predictions_{model_name}_{dataset_name}'] = save_path
        logger.debug(f"✓ График предсказаний сохранен: {save_path}")

    def plot_residuals_analysis(self, residuals, model_name: str, dataset_name: str = 'test', save_path: str = None):
        """Анализ остатков"""
        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/residuals/{model_name}_{dataset_name}_residuals.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        fig, axes = plt.subplots(2, 2, figsize=(15, 12))

        # 1. Гистограмма остатков
        axes[0, 0].hist(residuals, bins=50, edgecolor='black', alpha=0.7, density=True)
        axes[0, 0].axvline(x=0, color='red', linestyle='--', linewidth=2, alpha=0.7)
        axes[0, 0].set_xlabel('Остатки')
        axes[0, 0].set_ylabel('Плотность')
        axes[0, 0].set_title(f'Распределение остатков ({model_name})')
        axes[0, 0].grid(True, alpha=0.3)

        # 2. QQ-plot
        stats.probplot(residuals, dist="norm", plot=axes[0, 1])
        axes[0, 1].set_title(f'QQ-plot остатков ({model_name})')
        axes[0, 1].grid(True, alpha=0.3)

        # 3. Остатки по времени (если есть индекс времени)
        axes[1, 0].scatter(range(len(residuals)), residuals, alpha=0.6, s=20)
        axes[1, 0].axhline(y=0, color='red', linestyle='--', linewidth=2, alpha=0.7)
        axes[1, 0].set_xlabel('Наблюдение')
        axes[1, 0].set_ylabel('Остатки')
        axes[1, 0].set_title(f'Остатки по порядку ({model_name})')
        axes[1, 0].grid(True, alpha=0.3)

        # 4. Автокорреляция остатков
        from pandas.plotting import autocorrelation_plot
        autocorrelation_plot(residuals, ax=axes[1, 1])
        axes[1, 1].set_title(f'Автокорреляция остатков ({model_name})')
        axes[1, 1].grid(True, alpha=0.3)

        plt.suptitle(f'Анализ остатков: {model_name} ({dataset_name})', fontsize=16)
        plt.tight_layout()

        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close(fig)

        logger.debug(f"✓ График анализа остатков сохранен: {save_path}")

    def plot_feature_importance(self, feature_importances: Dict, model_name: str, save_path: str = None):
        """Визуализация важности признаков"""
        if not feature_importances:
            logger.warning(f"Нет данных о важности признаков для {model_name}")
            return

        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/feature_importance/{model_name}_importance.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        # Сортируем признаки по важности
        sorted_features = sorted(feature_importances.items(), key=lambda x: x[1], reverse=True)
        features, importance = zip(*sorted_features[:20])  # Топ-20 признаков

        fig, axes = plt.subplots(1, 2, figsize=(16, 8))

        # 1. Горизонтальная бар-диаграмма
        y_pos = np.arange(len(features))
        axes[0].barh(y_pos, importance, align='center', alpha=0.7, color='skyblue', edgecolor='black')
        axes[0].set_yticks(y_pos)
        axes[0].set_yticklabels(features)
        axes[0].invert_yaxis()  # Самые важные сверху
        axes[0].set_xlabel('Важность')
        axes[0].set_title(f'Топ-20 важных признаков ({model_name})')
        axes[0].grid(True, alpha=0.3, axis='x')

        # 2. Круговая диаграмма для топ-10
        if len(features) > 10:
            top_features = features[:10]
            top_importance = importance[:10]
            # Рассчитываем проценты
            total = sum(top_importance)
            percentages = [imp/total*100 for imp in top_importance]

            wedges, texts, autotexts = axes[1].pie(top_importance, labels=top_features, autopct='%1.1f%%',
                                                  startangle=90, colors=plt.cm.Set3.colors)
            axes[1].axis('equal')
            axes[1].set_title(f'Распределение важности (Топ-10, {model_name})')

            # Улучшаем читаемость текста
            for autotext in autotexts:
                autotext.set_color('black')
                autotext.set_fontsize(9)

        plt.suptitle(f'Важность признаков: {model_name}', fontsize=16)
        plt.tight_layout()

        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close(fig)

        logger.debug(f"✓ График важности признаков сохранен: {save_path}")

    def create_model_dashboard(self, model, X_test, y_test, model_name: str, feature_names: List[str], save_path: str = None):
        """Создание дашборда для модели"""
        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/diagnostics/{model_name}_dashboard.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        fig = plt.figure(figsize=(20, 15))

        # Сетка для графиков
        gs = fig.add_gridspec(3, 3)

        # 1. Предсказания vs Фактические значения
        ax1 = fig.add_subplot(gs[0, :2])
        y_pred = model.predict(X_test)
        ax1.scatter(y_test, y_pred, alpha=0.6, s=30)
        ax1.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
        ax1.set_xlabel('Фактические значения')
        ax1.set_ylabel('Предсказанные значения')
        ax1.set_title(f'Предсказания vs Фактические ({model_name})')
        ax1.grid(True, alpha=0.3)

        # 2. Распределение ошибок
        ax2 = fig.add_subplot(gs[0, 2])
        errors = y_pred - y_test
        ax2.hist(errors, bins=50, edgecolor='black', alpha=0.7)
        ax2.axvline(x=0, color='red', linestyle='--', linewidth=2)
        ax2.set_xlabel('Ошибка')
        ax2.set_ylabel('Частота')
        ax2.set_title('Распределение ошибок')
        ax2.grid(True, alpha=0.3)

        # 3. Важность признаков (если доступна)
        ax3 = fig.add_subplot(gs[1, :])
        try:
            feature_importance = model.get_feature_importance(feature_names)
            if feature_importance:
                sorted_features = sorted(feature_importance.items(), key=lambda x: x[1], reverse=True)[:15]
                features, importance = zip(*sorted_features)
                y_pos = np.arange(len(features))
                ax3.barh(y_pos, importance, align='center', alpha=0.7)
                ax3.set_yticks(y_pos)
                ax3.set_yticklabels(features)
                ax3.invert_yaxis()
                ax3.set_xlabel('Важность')
                ax3.set_title('Топ-15 важных признаков')
                ax3.grid(True, alpha=0.3, axis='x')
        except:
            ax3.text(0.5, 0.5, 'Важность признаков не доступна',
                    horizontalalignment='center', verticalalignment='center',
                    transform=ax3.transAxes)

        # 4. Кривая обучения (если доступна)
        ax4 = fig.add_subplot(gs[2, 0])
        try:
            learning_curve_data = model.get_learning_curve(X_test, y_test)
            if learning_curve_data:
                train_sizes = learning_curve_data['train_sizes']
                train_scores = learning_curve_data['train_scores']
                test_scores = learning_curve_data['test_scores']

                train_scores_mean = np.mean(train_scores, axis=1)
                train_scores_std = np.std(train_scores, axis=1)
                test_scores_mean = np.mean(test_scores, axis=1)
                test_scores_std = np.std(test_scores, axis=1)

                ax4.plot(train_sizes, train_scores_mean, 'o-', color='blue', label='Train')
                ax4.fill_between(train_sizes, train_scores_mean - train_scores_std,
                               train_scores_mean + train_scores_std, alpha=0.1, color='blue')
                ax4.plot(train_sizes, test_scores_mean, 'o-', color='green', label='Test')
                ax4.fill_between(train_sizes, test_scores_mean - test_scores_std,
                               test_scores_mean + test_scores_std, alpha=0.1, color='green')
                ax4.set_xlabel('Размер тренировочной выборки')
                ax4.set_ylabel('Score')
                ax4.set_title('Кривая обучения')
                ax4.legend()
                ax4.grid(True, alpha=0.3)
        except:
            ax4.text(0.5, 0.5, 'Кривая обучения не доступна',
                    horizontalalignment='center', verticalalignment='center',
                    transform=ax4.transAxes)

        # 5. Метрики модели
        ax5 = fig.add_subplot(gs[2, 1])
        metrics_text = f'''
        Модель: {model_name}
        R²: {model.test_metrics.get('r2', 'N/A'):.4f}
        RMSE: {model.test_metrics.get('rmse', 'N/A'):.4f}
        MAE: {model.test_metrics.get('mae', 'N/A'):.4f}
        MAPE: {model.test_metrics.get('mape', 'N/A'):.2f}%
        Время обучения: {model.training_time:.2f} сек
        '''
        ax5.text(0.1, 0.5, metrics_text, fontsize=12, verticalalignment='center',
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        ax5.axis('off')
        ax5.set_title('Метрики модели')

        # 6. Остатки по времени
        ax6 = fig.add_subplot(gs[2, 2])
        if hasattr(model, 'residuals') and 'test' in model.residuals:
            residuals = model.residuals['test']
            ax6.plot(range(len(residuals)), residuals, 'o', alpha=0.6, markersize=4)
            ax6.axhline(y=0, color='red', linestyle='--', linewidth=2)
            ax6.set_xlabel('Наблюдение')
            ax6.set_ylabel('Остатки')
            ax6.set_title('Остатки по порядку')
            ax6.grid(True, alpha=0.3)
        else:
            ax6.text(0.5, 0.5, 'Данные об остатках не доступны',
                    horizontalalignment='center', verticalalignment='center',
                    transform=ax6.transAxes)

        plt.suptitle(f'Дашборд модели: {model_name}', fontsize=18, y=1.02)
        plt.tight_layout()

        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close(fig)

        logger.debug(f"✓ Дашборд модели сохранен: {save_path}")

    def plot_learning_curves(self, model, X_train, y_train, model_name: str, save_path: str = None):
        """Улучшенная визуализация кривых обучения"""
        from sklearn.model_selection import learning_curve

        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/learning_curves/{model_name}_learning_curve.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        # Список моделей, для которых используем упрощенный метод
        problem_models = ['lightgbm', 'xgboost', 'catboost', 'gaussian_process']

        if model_name in problem_models:
            logger.debug(f"Используем упрощенный метод для {model_name}")
            try:
                self.plot_simple_learning_curve(model, X_train, y_train, model_name, save_path)
            except Exception as e:
                logger.debug(f"Упрощенная кривая обучения для {model_name} не удалась: {e}")
            return

        try:
            # Ваш существующий код для нормальных моделей
            train_sizes, train_scores, test_scores = learning_curve(
                model.model, X_train, y_train,
                train_sizes=np.linspace(0.1, 1.0, 10),
                cv=5,
                scoring='neg_mean_squared_error',
                n_jobs=min(2, self.config.n_jobs),
                shuffle=True,
                random_state=self.config.random_state
            )

            train_scores_mean = -np.mean(train_scores, axis=1)
            train_scores_std = np.std(train_scores, axis=1)
            test_scores_mean = -np.mean(test_scores, axis=1)
            test_scores_std = np.std(test_scores, axis=1)

            fig, ax = plt.subplots(figsize=(10, 6))

            ax.plot(train_sizes, train_scores_mean, 'o-', color='blue',
                  label='Train', linewidth=2)
            ax.fill_between(train_sizes, train_scores_mean - train_scores_std,
                          train_scores_mean + train_scores_std, alpha=0.1, color='blue')

            ax.plot(train_sizes, test_scores_mean, 'o-', color='green',
                  label='Test', linewidth=2)
            ax.fill_between(train_sizes, test_scores_mean - test_scores_std,
                          test_scores_mean + test_scores_std, alpha=0.1, color='green')

            ax.set_xlabel('Размер тренировочной выборки')
            ax.set_ylabel('MSE')
            ax.set_title(f'Кривая обучения: {model_name}')
            ax.legend(loc='best')
            ax.grid(True, alpha=0.3)

            plt.tight_layout()
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
            plt.close(fig)

            logger.debug(f"✓ Кривая обучения сохранена: {save_path}")

        except Exception as e:
            logger.warning(f"Не удалось построить кривую обучения для {model_name}: {e}")
            # Пробуем упрощенный метод как fallback
            try:
                logger.debug(f"Пробуем упрощенный метод как fallback")
                self.plot_simple_learning_curve(model, X_train, y_train, model_name,
                                              save_path.replace('.png', '_simple.png'))
            except Exception as fallback_error:
                logger.debug(f"Fallback также не удался: {fallback_error}")

    def plot_shap_summary(self, model, model_name: str, save_path: str = None):
        """Визуализация SHAP значений с исправлениями"""
        if not SHAP_AVAILABLE:
            logger.warning("SHAP не установлен")
            return

        if not hasattr(model, 'shap_values') or model.shap_values is None:
            logger.warning(f"SHAP значения не рассчитаны для модели {model_name}")
            return

        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/shap/{model_name}_shap_summary.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        try:
            # Получаем данные, для которых были рассчитаны SHAP значения
            X_sample = getattr(model, 'X_sample_shap', None)

            if X_sample is None:
                logger.warning(f"Нет данных для построения SHAP графика {model_name}")
                return

            shap_values = model.shap_values

            # Проверяем размерности
            if isinstance(shap_values, list):
                shap_values = shap_values[0]

            if X_sample.shape[0] != shap_values.shape[0]:
                logger.warning(f"Несоответствие размеров: X_sample {X_sample.shape}, shap_values {shap_values.shape}")
                return

            feature_names = X_sample.columns.tolist() if hasattr(X_sample, 'columns') else [f'feature_{i}' for i in range(X_sample.shape[1])]

            # Summary plot
            fig = plt.figure(figsize=(12, 8))
            shap.summary_plot(shap_values, X_sample, feature_names=feature_names, show=False)
            plt.title(f'SHAP Summary Plot: {model_name}')
            plt.tight_layout()
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
            plt.close(fig)

            # Bar plot
            bar_path = save_path.replace('.png', '_bar.png')
            fig = plt.figure(figsize=(10, 6))
            shap.summary_plot(shap_values, X_sample, plot_type="bar", feature_names=feature_names, show=False)
            plt.title(f'SHAP Feature Importance: {model_name}')
            plt.tight_layout()
            plt.savefig(bar_path, dpi=self.config.dpi, bbox_inches='tight')
            plt.close(fig)

            logger.info(f"✓ SHAP графики сохранены для {model_name}")

        except Exception as e:
            logger.warning(f"Не удалось построить SHAP графики для {model_name}: {e}")
            if self.config.debug:
                logger.debug(traceback.format_exc())


    def plot_model_comparison(self, models_results: Dict, metric: str = 'rmse', save_path: str = None):
        """Сравнение моделей по метрикам"""
        if save_path is None:
            save_path = f"{self.config.ml_results_dir}/plots/comparison/model_comparison_{metric}.png"

        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        # Собираем метрики
        model_names = []
        train_metrics = []
        val_metrics = []
        test_metrics = []

        for model_name, results in models_results.items():
            if 'train_metrics' in results and 'val_metrics' in results and 'test_metrics' in results:
                model_names.append(model_name)
                train_metrics.append(results['train_metrics'].get(metric, np.nan))
                val_metrics.append(results['val_metrics'].get(metric, np.nan))
                test_metrics.append(results['test_metrics'].get(metric, np.nan))

        if not model_names:
            logger.warning("Нет данных для сравнения моделей")
            return

        fig, axes = plt.subplots(1, 2, figsize=(16, 8))

        # 1. Бар-диаграмма сравнения
        x = np.arange(len(model_names))
        width = 0.25

        axes[0].bar(x - width, train_metrics, width, label='Train', alpha=0.7)
        axes[0].bar(x, val_metrics, width, label='Validation', alpha=0.7)
        axes[0].bar(x + width, test_metrics, width, label='Test', alpha=0.7)

        axes[0].set_xlabel('Модели')
        axes[0].set_ylabel(metric.upper())
        axes[0].set_title(f'Сравнение моделей по {metric.upper()}')
        axes[0].set_xticks(x)
        axes[0].set_xticklabels(model_names, rotation=45, ha='right')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3, axis='y')

        # 2. Тепловая карта метрик
        metrics_df = pd.DataFrame({
            'Train': train_metrics,
            'Validation': val_metrics,
            'Test': test_metrics
        }, index=model_names)

        im = axes[1].imshow(metrics_df.values, cmap='viridis', aspect='auto')
        axes[1].set_xticks(np.arange(len(metrics_df.columns)))
        axes[1].set_xticklabels(metrics_df.columns)
        axes[1].set_yticks(np.arange(len(model_names)))
        axes[1].set_yticklabels(model_names)

        # Добавляем значения в ячейки
        for i in range(len(model_names)):
            for j in range(len(metrics_df.columns)):
                text = axes[1].text(j, i, f'{metrics_df.iloc[i, j]:.4f}',
                                  ha="center", va="center", color="w", fontsize=9)

        axes[1].set_title(f'Тепловая карта метрик ({metric.upper()})')
        plt.colorbar(im, ax=axes[1])

        plt.suptitle(f'Сравнение моделей', fontsize=16)
        plt.tight_layout()

        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close(fig)

        logger.debug(f"✓ График сравнения моделей сохранен: {save_path}")