---

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

weight_col = 'sum_total_out'

def calc_metrics_for_comp(comp, weight_col):
    """Рассчитывает метрики для одного comparison DataFrame"""
    y_true = comp[target].values
    y_pred = comp[pred_col].values
    weights = comp[f'{weight_col}_fact'].values
    
    benchmark = np.average(y_true, weights=weights)
    comp['benchmark'] = benchmark
    
    mae = mean_absolute_error(y_true, y_pred)
    rmse = np.sqrt(mean_squared_error(y_true, y_pred))
    mape = np.mean(np.abs((y_true - y_pred) / np.maximum(np.abs(y_true), 1e-6))) * 100
    r2 = r2_score(y_true, y_pred)
    
    b_mae = mean_absolute_error(y_true, comp['benchmark'])
    b_rmse = np.sqrt(mean_squared_error(y_true, comp['benchmark']))
    b_mape = np.mean(np.abs((y_true - comp['benchmark']) / np.maximum(np.abs(y_true), 1e-6))) * 100
    b_r2 = r2_score(y_true, comp['benchmark'])
    
    wmae = np.average(np.abs(y_true - y_pred), weights=weights)
    wrmse = np.sqrt(np.average((y_true - y_pred)**2, weights=weights))
    wmape = np.average(np.abs((y_true - y_pred) / np.maximum(np.abs(y_true), 1e-6)), weights=weights) * 100
    
    b_wmae = np.average(np.abs(y_true - comp['benchmark']), weights=weights)
    b_wrmse = np.sqrt(np.average((y_true - comp['benchmark'])**2, weights=weights))
    b_wmape = np.average(np.abs((y_true - comp['benchmark']) / np.maximum(np.abs(y_true), 1e-6)), weights=weights) * 100
    
    return {
        'MAE': [mae, b_mae],
        'RMSE': [rmse, b_rmse],
        'MAPE (%)': [mape, b_mape],
        'R2': [r2, b_r2],
        'wMAE': [wmae, b_wmae],
        'wRMSE': [wrmse, b_wrmse],
        'wMAPE (%)': [wmape, b_wmape]
    }

# ---- Сбор данных по всем выборкам и сегментам ----
results = []

# 1. Общий портфель (train и test)
for label, df_part in [('Train', train_df), ('Test', test_df)]:
    comp = aggregate_weighted(df_part, date_col, target, weight_col).merge(
        aggregate_weighted(df_part, date_col, pred_col, weight_col),
        on=date_col, suffixes=('_fact', '_pred')
    )
    metrics = calc_metrics_for_comp(comp, weight_col)
    # Добавляем строки с метриками модели и бенчмарка
    results.append({'Выборка': f'{label} (Model)', **{k: v[0] for k, v in metrics.items()}})
    results.append({'Выборка': f'{label} (Benchmark)', **{k: v[1] for k, v in metrics.items()}})

# 2. По сегментам (train и test)
for seg in ['prem', 'mid']:
    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:
            comp = aggregate_weighted(df_seg, date_col, target, weight_col).merge(
                aggregate_weighted(df_seg, date_col, pred_col, weight_col),
                on=date_col, suffixes=('_fact', '_pred')
            )
            metrics = calc_metrics_for_comp(comp, weight_col)
            results.append({'Выборка': f'{seg.upper()} {label} (Model)', **{k: v[0] for k, v in metrics.items()}})
            results.append({'Выборка': f'{seg.upper()} {label} (Benchmark)', **{k: v[1] for k, v in metrics.items()}})

# ---- Создаём единую таблицу ----
final_metrics_df = pd.DataFrame(results)
final_metrics_df = final_metrics_df.round(3)

print("=== МЕТРИКИ КАЧЕСТВА ПРОГНОЗА (по выборкам и сегментам) ===")
display(final_metrics_df)