---

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression

# Предполагаем, что df, target, weight_col определены
target = 'renewed_fact_rate'
weight_col = 'portfolio_share'
date_col = 'end_month_plan'   # не используется, но для ясности

# Функция для получения агрегированных данных по pr_p_n
def get_grouped_data(df, weight_col, target, group_col='pr_p_n'):
    grouped = df.groupby(group_col).apply(
        lambda g: pd.Series({
            'weighted_target': np.average(g[target], weights=g[weight_col]),
            'total_weight': g[weight_col].sum()
        })
    ).reset_index()
    # Удаляем группы с нулевым весом (если есть)
    grouped = grouped[grouped['total_weight'] > 0]
    return grouped.sort_values(group_col)

# Функция для построения графика с линейным трендом
def plot_with_trend(ax, grouped, x_col='pr_p_n', y_col='weighted_target', weight_col='total_weight', 
                    title='', point_size=50, color='blue'):
    # Извлекаем данные
    x = grouped[x_col].values.reshape(-1, 1)
    y = grouped[y_col].values
    w = grouped[weight_col].values
    
    # Взвешенная линейная регрессия через sklearn
    model = LinearRegression()
    model.fit(x, y, sample_weight=w)
    y_pred = model.predict(x)
    
    # Строим точки (размер постоянный)
    ax.scatter(x, y, s=point_size, color=color, alpha=0.7, edgecolors='k', linewidth=0.5, label='Данные')
    # Прямая тренда
    ax.plot(x, y_pred, color='red', linestyle='--', linewidth=2, label=f'Тренд: y = {model.coef_[0]:.4f}x + {model.intercept_:.4f}')
    
    ax.set_xlabel('pr_p_n')
    ax.set_ylabel('Взвешенная доля пролонгаций')
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend()

# --- 1. Данные для всего портфеля ---
grouped_all = get_grouped_data(df, weight_col, target)

# --- 2. Данные для сегментов ---
grouped_prem = get_grouped_data(df[df['segment'] == 'prem'], weight_col, target)
grouped_mid = get_grouped_data(df[df['segment'] == 'mid'], weight_col, target)

# --- 3. Построение трёх графиков ---
fig = plt.figure(figsize=(10, 10))
# Создаём сетку: верхний график – 2/3 высоты, нижние – по 1/3 каждый
gs = fig.add_gridspec(3, 1, height_ratios=[2, 1, 1])

ax1 = fig.add_subplot(gs[0])
ax2 = fig.add_subplot(gs[1])
ax3 = fig.add_subplot(gs[2])

# Все сегменты
plot_with_trend(ax1, grouped_all, title='Зависимость взвешенной доли пролонгаций от pr_p_n (весь портфель)', point_size=60, color='blue')

# Сегмент PREM
plot_with_trend(ax2, grouped_prem, title='Сегмент PREM', point_size=50, color='green')

# Сегмент MID
plot_with_trend(ax3, grouped_mid, title='Сегмент MID', point_size=50, color='orange')

plt.tight_layout()
plt.show()