---

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

# ---------------------------------------------------------------------
# 1. Задаём основные переменные (при необходимости измените)
# ---------------------------------------------------------------------
date_col = 'end_month_plan'
target = 'renewed_fact_rate'
pred_col = 'prediction'
weight_col = 'sum_total_out'   # или 'portfolio_share' – укажите нужный

# ---------------------------------------------------------------------
# 2. Функция агрегации (взвешенное среднее)
# ---------------------------------------------------------------------
def aggregate_weighted(df, group_col, value_col, weight_col):
    return (df.groupby(group_col)
              .apply(lambda g: pd.Series({
                  value_col: np.average(g[value_col], weights=g[weight_col]),
                  weight_col: g[weight_col].sum()
              }))
              .reset_index()
              .sort_values(group_col))

# ---------------------------------------------------------------------
# 3. Создаём comparison (факт + прогноз) для всего портфеля
# ---------------------------------------------------------------------
# Факт (взвешенный по weight_col)
fact_all = aggregate_weighted(df, date_col, target, weight_col)
# Прогноз (взвешенный по weight_col)
pred_all = aggregate_weighted(df, date_col, pred_col, weight_col)
# Объединяем
comparison = fact_all.merge(pred_all, on=date_col, suffixes=('_fact', '_pred'))

# Проверяем, что сравнение создалось
print("comparison создан, строк:", len(comparison))
print(comparison.head())

# ---------------------------------------------------------------------
# 4. Расчёт метрик
# ---------------------------------------------------------------------
y_true = comparison[target].values
y_pred = comparison[pred_col].values
weights = comparison[f'{weight_col}_fact'].values   # сумма весов по месяцам

# Обычные метрики
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)

# Взвешенные метрики (без R2)
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

# ---------------------------------------------------------------------
# 5. Формируем DataFrame с результатами
# ---------------------------------------------------------------------
metrics_df = pd.DataFrame({
    'MAE': [mae, wmae],
    'RMSE': [rmse, wrmse],
    'MAPE (%)': [mape, wmape],
    'R²': [r2, np.nan]
}, index=['Обычные', 'Взвешенные'])

metrics_df = metrics_df.round(5)

# ---------------------------------------------------------------------
# 6. Вывод
# ---------------------------------------------------------------------
print("\nМетрики качества прогноза (вес = {})".format(weight_col))
display(metrics_df)   # в Jupyter – красивая таблица
# Если в обычном скрипте, замените на: print(metrics_df.to_string())