---

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

# Предполагается, что df уже загружен, 
# определены date_col, target, weight_col, volume_col

date_col = 'end_month_plan'
target = 'renewed_fact_rate'
weight_col = 'portfolio_share'      # столбец весов для взвешенного среднего
pred_col = 'prediction'             # столбец с прогнозом

# Функция агрегации (взвешенное среднее)
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))

# Получаем агрегированные данные для всего портфеля
fact_all = aggregate_weighted(df, date_col, target, weight_col)
pred_all = aggregate_weighted(df, date_col, pred_col, weight_col)
comp_all = fact_all.merge(pred_all, on=date_col, suffixes=('_fact', '_pred'))

# Расчёт метрик
y_true = comp_all[target].values
y_pred = comp_all[pred_col].values
mae = mean_absolute_error(y_true, y_pred)
r2 = r2_score(y_true, y_pred)

# Построение графика
fig, ax = plt.subplots(figsize=(12, 5))

# Основная ось – линии факта и прогноза
ax.plot(comp_all[date_col], comp_all[target], marker='o', label='Факт', color='blue')
ax.plot(comp_all[date_col], comp_all[pred_col], marker='x', label='Прогноз', color='red')
ax.set_title('Все сегменты (взвешенная по portfolio_share доля)')
ax.set_ylabel('Доля пролонгаций')
ax.grid(True, alpha=0.3)
ax.tick_params(axis='x', rotation=45)

# Вторая ось – столбцы суммарного веса (объём)
ax2 = ax.twinx()
volume = comp_all[f'{weight_col}_fact'].values
ax2.bar(comp_all[date_col], volume, alpha=0.3, color='gray', width=20, label='Объём (сумма весов)')
ax2.set_ylabel('Суммарный вес', color='gray')
ax2.tick_params(axis='y', labelcolor='gray')

# Добавляем метрики в виде текста
textstr = f'MAE = {mae:.4f}\nR² = {r2:.4f}'
ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=10,
        verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

ax.legend(loc='lower left')
plt.tight_layout()
plt.show()

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

# Предполагается, что df, date_col, target, weight_col, pred_col уже определены

# Функция агрегации (та же, что выше)
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))

# Сегменты для построения
segments = ['prem', 'mid']

# Собираем агрегированные данные по каждому сегменту
comp_seg = {}
for seg in segments:
    df_seg = df[df['segment'] == seg].copy()
    fact_seg = aggregate_weighted(df_seg, date_col, target, weight_col)
    pred_seg = aggregate_weighted(df_seg, date_col, pred_col, weight_col)
    comp_seg[seg] = fact_seg.merge(pred_seg, on=date_col, suffixes=('_fact', '_pred'))

# Расчёт метрик для каждого сегмента
metrics_seg = {}
for seg in segments:
    comp = comp_seg[seg]
    y_true = comp[target].values
    y_pred = comp[pred_col].values
    metrics_seg[seg] = (mean_absolute_error(y_true, y_pred), r2_score(y_true, y_pred))

# Построение двух графиков
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
fig.suptitle('Сравнение фактической и прогнозной доли пролонгаций по сегментам (взвешенной по portfolio_share)', fontsize=14)

for ax, seg in zip(axes, segments):
    comp = comp_seg[seg]
    mae, r2 = metrics_seg[seg]

    # Основная ось
    ax.plot(comp[date_col], comp[target], marker='o', label='Факт', color='blue')
    ax.plot(comp[date_col], comp[pred_col], marker='x', label='Прогноз', color='red')
    ax.set_title(f'Сегмент {seg.upper()}')
    ax.set_ylabel('Доля пролонгаций')
    ax.grid(True, alpha=0.3)
    ax.tick_params(axis='x', rotation=45)

    # Вторая ось – столбцы объёма
    ax2 = ax.twinx()
    volume = comp[f'{weight_col}_fact'].values
    ax2.bar(comp[date_col], volume, alpha=0.3, color='gray', width=20, label='Объём (сумма весов)')
    ax2.set_ylabel('Суммарный вес', color='gray')
    ax2.tick_params(axis='y', labelcolor='gray')

    # Метрики
    textstr = f'MAE = {mae:.4f}\nR² = {r2:.4f}'
    ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

    ax.legend(loc='lower left')

plt.tight_layout()
plt.show()