---

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

# ------------------------------------------------------------
# 3. Функция расчёта метрик для одного сравнения
# ------------------------------------------------------------
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
    }

# ------------------------------------------------------------
# 4. Сбор данных для всех сегментов + общего портфеля
# ------------------------------------------------------------
segments = ['mass', 'mid', 'prem']
results = []

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

# По сегментам (train и test)
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:
            metrics = calc_metrics(comp, benchmark_value)
            if metrics:
                results.append({
                    'Выборка': f'{label} ({seg})',
                    'Модель wMAE': metrics['model_wmae'],
                    'Модель R2': metrics['model_r2'],
                    'Бенчмарк wMAE': metrics['bench_wmae'],
                    'Бенчмарк R2': metrics['bench_r2']
                })

# ------------------------------------------------------------
# 5. Создание и вывод итоговой таблицы
# ------------------------------------------------------------
df_metrics = pd.DataFrame(results).round(4)

# Переставляем столбцы для удобства
columns_order = ['Выборка', 'Модель wMAE', 'Модель R2', 'Бенчмарк wMAE', 'Бенчмарк R2']
df_metrics = df_metrics[columns_order]

print("=== МЕТРИКИ КАЧЕСТВА ПРОГНОЗА ===")
print(f"Бенчмарк (средневзвешенное на train): {benchmark_value:.5f}")
print()
display(df_metrics)