---

# Визуализация прогнозных результатов модели

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, r2_score

# ------------------------------------------------------------
# 1. Подготовка данных (используем существующие переменные)
# ------------------------------------------------------------
# weight_col, date_col, target, pred_col уже определены в ноутбуке
# df уже загружен

train_end = pd.Timestamp('2025-02-01')
train_df = df[df[date_col] < train_end].copy()
test_df = df[df[date_col] >= train_end].copy()

# ------------------------------------------------------------
# 2. Функция получения веса (если есть суффикс)
# ------------------------------------------------------------
def get_weight_col(comp, weight_col):
    candidates = [weight_col, f'{weight_col}_sum']
    for col in candidates:
        if col in comp.columns:
            return col
    return weight_col

# ------------------------------------------------------------
# 3. Расчёт бенчмарка на TRAIN (средневзвешенное)
# ------------------------------------------------------------
benchmark_value = np.average(
    train_df[target].values,
    weights=train_df[weight_col].values
)
print(f"Бенчмарк (средневзвешенное на train): {benchmark_value:.5f}")

# ------------------------------------------------------------
# 4. Функция построения одного графика (на весь экран)
# ------------------------------------------------------------
def plot_forecast_full(ax, comp, date_col, target, pred_col, weight_col, 
                        benchmark_value, title, split_date):
    """Построение одного графика прогноза на всю ширину."""
    if len(comp) == 0:
        ax.text(0.5, 0.5, 'Нет данных', ha='center', va='center', transform=ax.transAxes)
        ax.set_title(title)
        return
    
    y_true = comp[target].values
    y_pred = comp[pred_col].values
    weight_sum_col = get_weight_col(comp, weight_col)
    weights = comp[weight_sum_col].values
    
    # Линии
    ax.plot(comp[date_col], y_true, marker='o', label='Факт', color='blue', linewidth=2)
    ax.plot(comp[date_col], y_pred, marker='s', label='Прогноз модели', color='red', linewidth=2)
    ax.axhline(y=benchmark_value, color='green', linestyle='--', linewidth=2, label=f'Бенчмарк (ср.={benchmark_value:.3f})')
    
    # Вертикальная линия разделения train/test
    ax.axvline(x=split_date, color='black', linestyle=':', linewidth=2, alpha=0.7, label='Разделение train/test')
    
    ax.set_title(title)
    ax.set_ylabel('Доля пролонгаций')
    ax.grid(True, alpha=0.3)
    ax.tick_params(axis='x', rotation=45)
    
    # Столбцы объёма (вторая ось)
    ax2 = ax.twinx()
    ax2.bar(comp[date_col], weights, alpha=0.3, color='gray', width=15, label='Объём (sum_total_out)')
    ax2.set_ylabel('Суммарный объём', color='gray')
    ax2.tick_params(axis='y', labelcolor='gray')
    
    # Метрики
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    ax.text(0.05, 0.95, f'MAE = {mae:.3f}\nR² = {r2:.3f}', 
            transform=ax.transAxes, fontsize=11,
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    # Легенда
    ax.legend(loc='lower left')

# ------------------------------------------------------------
# 5. Сбор данных для всех сегментов (train + test вместе)
# ------------------------------------------------------------
segments = ['mass', 'mid', 'prem']
comparisons = {}

# Общий портфель (все данные)
comparisons['все сегменты'] = get_comparison(df, date_col, target, pred_col, weight_col)

# По сегментам (все данные)
for seg in segments:
    df_seg = df[df['segment'] == seg].copy()
    comparisons[seg] = get_comparison(df_seg, date_col, target, pred_col, weight_col)

# ------------------------------------------------------------
# 6. Построение вертикальных графиков (каждый на весь экран)
# ------------------------------------------------------------
split_date = train_end

# Список заголовков и ключей
order = ['все сегменты'] + segments
titles = {
    'все сегменты': 'Все сегменты (прогноз vs факт)',
    'mass': 'Масс-сегмент (прогноз vs факт)',
    'mid': 'Mid-сегмент (прогноз vs факт)',
    'prem': 'Prem-сегмент (прогноз vs факт)'
}

# Создаём фигуру с вертикальными subplot
n_plots = len(order)
fig, axes = plt.subplots(n_plots, 1, figsize=(16, 5 * n_plots))

# Если один график – axes не массив
if n_plots == 1:
    axes = [axes]

for ax, key in zip(axes, order):
    comp = comparisons.get(key, pd.DataFrame())
    plot_forecast_full(
        ax=ax,
        comp=comp,
        date_col=date_col,
        target=target,
        pred_col=pred_col,
        weight_col=weight_col,
        benchmark_value=benchmark_value,
        title=titles.get(key, key),
        split_date=split_date
    )

plt.suptitle('Сравнение фактической и прогнозной доли пролонгаций с бенчмарком', fontsize=16, y=1.02)
plt.tight_layout()
plt.show()