---

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

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 и 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()

# Функция агрегации (без apply)
def aggregate_weighted(df, group_col, value_col, weight_col):
    df_temp = df[[group_col, value_col, weight_col]].copy()
    df_temp['_weighted_sum'] = df_temp[value_col] * df_temp[weight_col]
    
    grouped = df_temp.groupby(group_col).agg({
        '_weighted_sum': 'sum',
        weight_col: 'sum'
    }).reset_index()
    
    grouped[value_col] = grouped['_weighted_sum'] / grouped[weight_col]
    grouped = grouped.drop(columns=['_weighted_sum'])
    grouped = grouped[grouped[weight_col] > 0]
    
    return grouped.sort_values(group_col).rename(columns={weight_col: f'{weight_col}_sum'})


def get_comparison(df_part, date_col, target, pred_col, weight_col):
    """Возвращает объединённый DataFrame с фактом и прогнозом."""
    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 для train и test
# ------------------------------------------------------------

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)

# ------------------------------------------------------------
# 3. Построение двух блоков по 4 графика (train + test)
# ------------------------------------------------------------

# Блок 1: Train (все сегменты + mass + mid + prem)
fig1, axes1 = plt.subplots(2, 2, figsize=(16, 12))
axes1 = axes1.flatten()

order_train = ['все сегменты'] + segments
for ax, label in zip(axes1, order_train):
    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 (все сегменты + mass + mid + prem)
fig2, axes2 = plt.subplots(2, 2, figsize=(16, 12))
axes2 = axes2.flatten()

order_test = ['все сегменты'] + segments
for ax, label in zip(axes2, order_test):
    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()