---

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

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 = 'sum_total_out'        # уже определено
date_col = 'end_month_plan'         # уже определено
target = 'renewed_fact_rate'        # уже определено
pred_col = 'prediction'             # уже определено

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. Функции построения (без дублирования aggregate_weighted)
# ------------------------------------------------------------
def get_comparison(df_part, date_col, target, pred_col, weight_col):
    """Возвращает объединённый DataFrame с фактом и прогнозом."""
    if len(df_part) == 0:
        return pd.DataFrame()
    
    # Используем существующую функцию aggregate_weighted
    fact = aggregate_weighted(df_part, date_col, target, weight_col)
    pred = aggregate_weighted(df_part, date_col, pred_col, weight_col)
    comp = fact.merge(pred, on=date_col, suffixes=('_fact', '_pred'))
    return comp


def plot_forecast(ax, comp, date_col, target, pred_col, weight_col, title):
    """Построение одного графика прогноза."""
    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
    weights = comp[weight_col].values   # используем weight_col без суффикса
    
    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.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='Объём')
    ax2.set_ylabel('Суммарный объём (sum_total_out)', 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=10,
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    ax.legend(loc='lower left')

# ------------------------------------------------------------
# 3. Сбор данных и построение
# ------------------------------------------------------------
segments = ['mass', 'mid', 'prem']
comparisons = {}

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

# По сегментам
for seg in segments:
    train_seg = train_df[train_df['segment'] == seg].copy()
    test_seg = test_df[test_df['segment'] == seg].copy()
    comparisons[(seg, 'train')] = get_comparison(train_seg, date_col, target, pred_col, weight_col)
    comparisons[(seg, 'test')] = get_comparison(test_seg, date_col, target, pred_col, weight_col)

# Графики
order = ['все сегменты'] + segments

fig1, axes1 = plt.subplots(2, 2, figsize=(16, 12))
for ax, label in zip(axes1.flatten(), order):
    comp = comparisons.get((label, 'train'), pd.DataFrame())
    plot_forecast(ax, comp, date_col, target, pred_col, weight_col, title=f'{label} (train)')
plt.suptitle('Прогнозные результаты модели (train, до 02.2025)', fontsize=14)
plt.tight_layout()
plt.show()

fig2, axes2 = plt.subplots(2, 2, figsize=(16, 12))
for ax, label in zip(axes2.flatten(), order):
    comp = comparisons.get((label, 'test'), pd.DataFrame())
    plot_forecast(ax, comp, date_col, target, pred_col, weight_col, title=f'{label} (test)')
plt.suptitle('Прогнозные результаты модели (test, с 02.2025)', fontsize=14)
plt.tight_layout()
plt.show()