---

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'

# Функция для получения агрегированных данных по 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=60, color='blue'):
    x = grouped[x_col].values.reshape(-1, 1)
    y = grouped[y_col].values
    w = grouped[weight_col].values
    
    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='gray', linestyle='--', linewidth=2,
            label=f'Тренд: y = {model.coef_[0]:.3f}x + {model.intercept_:.3f}')
    
    ax.set_xlabel('pr_p_n')
    ax.set_ylabel('Взвешенная доля')
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend()

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

# --- Построение с использованием GridSpec (2 строки, 2 колонки) ---
fig = plt.figure(figsize=(12, 10))
gs = fig.add_gridspec(2, 2, height_ratios=[2, 1], width_ratios=[1, 1])

# Верхний график (занимает обе колонки)
ax1 = fig.add_subplot(gs[0, :])
plot_with_trend(ax1, grouped_all, title='Зависимость взвешенной доли пролонгаций от pr_p_n (весь портфель)',
                point_size=70, color='blue')

# Нижний левый (сегмент PREM)
ax2 = fig.add_subplot(gs[1, 0])
plot_with_trend(ax2, grouped_prem, title='Сегмент PREM', point_size=60, color='blue')

# Нижний правый (сегмент MID)
ax3 = fig.add_subplot(gs[1, 1])
plot_with_trend(ax3, grouped_mid, title='Сегмент MID', point_size=60, color='blue')

plt.tight_layout()
plt.show()