---

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

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

# ------------------------------------------------------------
# 1. Подготовка данных (используем существующую aggregate_weighted)
# ------------------------------------------------------------

weight_col = 'sum_total_out'
date_col = 'end_month_plan'
target = 'renewed_fact_rate'
pred_col = 'prediction'

# Разделение на train и test
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()


def get_comparison(df_part, date_col, target, pred_col, weight_col):
    """Возвращает объединённый 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):
    """Построение одного графика прогноза."""
    y_true = comp[target].values
    y_pred = comp[pred_col].values
    weights = comp[f'{weight_col}_sum'].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.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')


# ------------------------------------------------------------
# 2. Построение сетки 2x2
# ------------------------------------------------------------

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()
    if len(train_seg) > 0:
        comparisons[(seg, 'train')] = get_comparison(train_seg, date_col, target, pred_col, weight_col)
    if len(test_seg) > 0:
        comparisons[(seg, 'test')] = get_comparison(test_seg, date_col, target, pred_col, weight_col)

# Блок 1: Train
fig1, axes1 = plt.subplots(2, 2, figsize=(16, 12))
axes1 = axes1.flatten()
order = ['все сегменты'] + segments

for ax, label in zip(axes1, order):
    comp = comparisons.get((label, 'train'))
    if comp is not None and len(comp) > 0:
        plot_forecast(ax, comp, date_col, target, pred_col, weight_col, title=f'{label} (train)')
    else:
        ax.text(0.5, 0.5, 'Нет данных', ha='center', va='center', transform=ax.transAxes)
        ax.set_title(f'{label} (train)')

plt.suptitle('Прогнозные результаты модели (train, до 02.2025)', fontsize=14)
plt.tight_layout()
plt.show()

# Блок 2: Test
fig2, axes2 = plt.subplots(2, 2, figsize=(16, 12))
axes2 = axes2.flatten()

for ax, label in zip(axes2, order):
    comp = comparisons.get((label, 'test'))
    if comp is not None and len(comp) > 0:
        plot_forecast(ax, comp, date_col, target, pred_col, weight_col, title=f'{label} (test)')
    else:
        ax.text(0.5, 0.5, 'Нет данных', ha='center', va='center', transform=ax.transAxes)
        ax.set_title(f'{label} (test)')

plt.suptitle('Прогнозные результаты модели (test, с 02.2025)', fontsize=14)
plt.tight_layout()
plt.show()