---

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'

# Функция для получения агрегированных данных по new_delta с заданным шагом
def get_grouped_data_by_delta(df, weight_col, target, step=0.25, min_weight_ratio=0.01):
    """
    Группирует данные по округлённому new_delta (шаг step),
    вычисляет взвешенное среднее и сумму весов,
    отфильтровывает группы с весом < min_weight_ratio от общего веса.
    """
    df_copy = df.copy()
    df_copy['new_delta_binned'] = np.round(df_copy['new_delta'] / step) * step
    
    grouped = df_copy.groupby('new_delta_binned').apply(
        lambda g: pd.Series({
            'weighted_target': np.average(g[target], weights=g[weight_col]),
            'total_weight': g[weight_col].sum()
        })
    ).reset_index()
    
    # Фильтр по весу
    total_weight = grouped['total_weight'].sum()
    grouped = grouped[grouped['total_weight'] >= min_weight_ratio * total_weight]
    
    return grouped.sort_values('new_delta_binned')

# Функция построения графика с линейным трендом (общий цвет точек, серая линия)
def plot_with_trend_delta(ax, grouped, x_col='new_delta_binned', 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('new_delta (шаг 0.25)')
    ax.set_ylabel('Взвешенная доля')
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend()

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

# --- Построение с использованием 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_delta(ax1, grouped_all, title='Зависимость взвешенной доли пролонгаций от new_delta (весь портфель)',
                      point_size=70, color='blue')

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

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

plt.tight_layout()
plt.show()