---

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

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

# ------------------------------------------------------------
# 1. Подготовка данных
# ------------------------------------------------------------
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 с фактом и прогнозом."""
    if len(df_part) == 0:
        return pd.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

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

# ------------------------------------------------------------
# 3. Функция построения графика
# ------------------------------------------------------------
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
    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.axhline(y=benchmark_value, color='green', linestyle='--', linewidth=2, 
               label=f'Бенчмарк (ср.={benchmark_value:.3f})')
    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='Объём')
    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=11,
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    ax.legend(loc='lower left')

# ------------------------------------------------------------
# 4. Сбор данных
# ------------------------------------------------------------
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)

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

n_plots = len(order)
fig, axes = plt.subplots(n_plots, 1, figsize=(16, 5 * n_plots))
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()