---

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

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

# ------------------------------------------------------------
# 1. Функция получения сравнения
# ------------------------------------------------------------
def get_comparison(df_part):
    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. Функция построения графика
# ------------------------------------------------------------
def plot_forecast_full(ax, comp, title, benchmark_value):
    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[f'{weight_col}_fact'].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})')
    ax.axvline(x=train_end, 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')

# ------------------------------------------------------------
# 3. Подготовка данных
# ------------------------------------------------------------
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()

segments = ['mass', 'mid', 'prem']
comparisons = {}
benchmarks = {}

# Общий портфель
comparisons['все сегменты'] = get_comparison(df)
benchmarks['все сегменты'] = np.average(
    train_df[target].values,
    weights=train_df[weight_col].values
)

# По сегментам
for seg in segments:
    train_seg = train_df[train_df['segment'] == seg].copy()
    if len(train_seg) > 0:
        benchmarks[seg] = np.average(
            train_seg[target].values,
            weights=train_seg[weight_col].values
        )
    else:
        benchmarks[seg] = None
    
    df_seg = df[df['segment'] == seg].copy()
    comparisons[seg] = get_comparison(df_seg)

# ------------------------------------------------------------
# 4. Построение графиков
# ------------------------------------------------------------
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())
    benchmark_value = benchmarks.get(key, 0)
    if benchmark_value is None:
        benchmark_value = 0
    plot_forecast_full(ax, comp, titles.get(key, key), benchmark_value)

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

# 1. Качество прогноза модели

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

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

# ------------------------------------------------------------
# 2. Функция расчёта метрик
# ------------------------------------------------------------
def calc_metrics(comp, benchmark_value):
    if len(comp) == 0:
        return None
    
    y_true = comp[target].values
    y_pred = comp[pred_col].values
    weights = comp[f'{weight_col}_fact'].values
    
    model_wmae = np.average(np.abs(y_true - y_pred), weights=weights)
    model_r2 = r2_score(y_true, y_pred)
    
    benchmark_pred = np.full_like(y_true, benchmark_value)
    bench_wmae = np.average(np.abs(y_true - benchmark_pred), weights=weights)
    bench_r2 = r2_score(y_true, benchmark_pred)
    
    return {
        'model_wmae': model_wmae,
        'model_r2': model_r2,
        'bench_wmae': bench_wmae,
        'bench_r2': bench_r2
    }

# ------------------------------------------------------------
# 3. Сбор данных
# ------------------------------------------------------------
segments = ['mass', 'mid', 'prem']
results = []

def get_comparison(df_part):
    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

# Общий портфель
for label, df_part in [('Train', train_df), ('Test', test_df)]:
    comp = get_comparison(df_part)
    if len(comp) > 0:
        # Бенчмарк на train ЭТОЙ же выборки
        bench_val = np.average(
            train_df[target].values,
            weights=train_df[weight_col].values
        )
        metrics = calc_metrics(comp, bench_val)
        if metrics:
            results.append({
                'Выборка': f'{label} (весь портфель)',
                'Модель wMAE': metrics['model_wmae'],
                'Модель R2': metrics['model_r2'],
                'Бенчмарк wMAE': metrics['bench_wmae'],
                'Бенчмарк R2': metrics['bench_r2']
            })

# По сегментам
for seg in segments:
    for label, df_part in [('Train', train_df), ('Test', test_df)]:
        df_seg = df_part[df_part['segment'] == seg].copy()
        if len(df_seg) == 0:
            continue
        
        comp = get_comparison(df_seg)
        if len(comp) == 0:
            continue
        
        # Бенчмарк на train ЭТОГО сегмента
        train_seg = train_df[train_df['segment'] == seg].copy()
        if len(train_seg) > 0:
            bench_val = np.average(
                train_seg[target].values,
                weights=train_seg[weight_col].values
            )
        else:
            bench_val = 0
        
        metrics = calc_metrics(comp, bench_val)
        if metrics:
            results.append({
                'Выборка': f'{label} ({seg})',
                'Модель wMAE': metrics['model_wmae'],
                'Модель R2': metrics['model_r2'],
                'Бенчмарк wMAE': metrics['bench_wmae'],
                'Бенчмарк R2': metrics['bench_r2']
            })

# ------------------------------------------------------------
# 4. Вывод таблицы
# ------------------------------------------------------------
df_metrics = pd.DataFrame(results).round(4)
columns_order = ['Выборка', 'Модель wMAE', 'Модель R2', 'Бенчмарк wMAE', 'Бенчмарк R2']
df_metrics = df_metrics[columns_order]

print("=== МЕТРИКИ КАЧЕСТВА ПРОГНОЗА ===")
print("Бенчмарк = средневзвешенное на train соответствующей выборки")
print()
display(df_metrics)