---

# 5. Наличие нелинейных зависимостей

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

# ------------------------------------------------------------
# 1. Вспомогательные функции
# ------------------------------------------------------------

def get_grouped_data(df, weight_col, target, group_col='pr_p_n'):
    """
    Группировка по pr_p_n с вычислением взвешенного среднего.
    """
    grouped = df.groupby(group_col, group_keys=False).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 get_grouped_data_by_delta(df, weight_col, target, step=0.25, min_weight_ratio=0.01):
    """
    Группировка по new_delta с округлением до заданного шага.
    """
    df_copy = df.copy()
    df_copy['new_delta_binned'] = np.round(df_copy['new_delta'] / step) * step
    
    grouped = df_copy.groupby('new_delta_binned', group_keys=False).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(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(x_col if x_col != 'new_delta_binned' else 'new_delta (шаг 0.25)')
    ax.set_ylabel('Взвешенная доля пролонгаций')
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='lower left')
    
    # Подпись справа с суммой весов
    ax.text(0.95, 0.95, f'Объём: {w.sum():.2e}', transform=ax.transAxes,
            fontsize=9, verticalalignment='top', horizontalalignment='right',
            bbox=dict(boxstyle='round', facecolor='white', alpha=0.7))


def plot_volume_bars(ax, grouped, x_col='pr_p_n', weight_col='total_weight'):
    """
    Добавление столбцов объёма на вторую ось.
    """
    ax2 = ax.twinx()
    ax2.bar(grouped[x_col], grouped[weight_col], alpha=0.3, color='gray', 
            width=0.6, label='Объём')
    ax2.set_ylabel('Суммарный объём (sum_total_out)', color='gray')
    ax2.tick_params(axis='y', labelcolor='gray')


def plot_nonlinear_test(df, weight_col, target, factor, segments, 
                         factor_type='pr_p_n', step=0.25, figsize=(14, 12)):
    """
    Универсальная функция для построения графиков по фактору.
    factor_type: 'pr_p_n' или 'new_delta'
    """
    if factor_type == 'pr_p_n':
        get_data_func = get_grouped_data
        x_col = 'pr_p_n'
        title_prefix = 'pr_p_n'
    else:
        get_data_func = lambda d, wc, t: get_grouped_data_by_delta(d, wc, t, step=step)
        x_col = 'new_delta_binned'
        title_prefix = 'new_delta (шаг 0.25)'
    
    # Подготовка данных для всех сегментов + общего портфеля
    all_groups = {'Все сегменты': df}
    for seg in segments:
        all_groups[f'Сегмент {seg.upper()}'] = df[df['segment'] == seg]
    
    # Создаём сетку: верхний график для всех сегментов, ниже по одному на сегмент
    n_seg = len(segments)
    fig = plt.figure(figsize=figsize)
    gs = fig.add_gridspec(n_seg + 1, 1, height_ratios=[2] + [1] * n_seg)
    
    axes = []
    # Верхний график (все сегменты)
    ax_main = fig.add_subplot(gs[0])
    axes.append(('Все сегменты', ax_main))
    
    # Нижние графики (по сегментам)
    for i, seg in enumerate(segments):
        ax = fig.add_subplot(gs[i + 1])
        axes.append((f'Сегмент {seg.upper()}', ax))
    
    # Строим графики
    for (label, ax) in axes:
        group_data = all_groups[label]
        if len(group_data) == 0:
            ax.text(0.5, 0.5, 'Нет данных', ha='center', va='center', transform=ax.transAxes)
            continue
        
        grouped = get_data_func(group_data, weight_col, target)
        if len(grouped) == 0:
            ax.text(0.5, 0.5, 'Нет данных после агрегации', ha='center', va='center', transform=ax.transAxes)
            continue
        
        # Основной график с трендом
        plot_with_trend(ax, grouped, x_col=x_col, title=label, point_size=60, color='blue')
        # Столбцы объёма
        plot_volume_bars(ax, grouped, x_col=x_col, weight_col='total_weight')
    
    plt.tight_layout()
    plt.show()


# ------------------------------------------------------------
# 2. Основной код теста
# ------------------------------------------------------------

# Параметры
weight_col = 'sum_total_out'
target = 'renewed_fact_rate'
segments = ['mass', 'mid', 'prem']

# 2.1. Графики для pr_p_n
print("=== ЗАВИСИМОСТЬ ОТ pr_p_n ===")
plot_nonlinear_test(
    df=df,
    weight_col=weight_col,
    target=target,
    factor='pr_p_n',
    segments=segments,
    factor_type='pr_p_n',
    figsize=(14, 12)
)

# 2.2. Графики для new_delta (шаг 0.25)
print("\n=== ЗАВИСИМОСТЬ ОТ new_delta (шаг 0.25) ===")
plot_nonlinear_test(
    df=df,
    weight_col=weight_col,
    target=target,
    factor='new_delta',
    segments=segments,
    factor_type='new_delta',
    step=0.25,
    figsize=(14, 12)
)

# ------------------------------------------------------------
# 3. Дополнительно: duration_months (если нужно)
# ------------------------------------------------------------
print("\n=== ЗАВИСИМОСТЬ ОТ duration_months ===")
plot_nonlinear_test(
    df=df,
    weight_col=weight_col,
    target=target,
    factor='duration_months',
    segments=segments,
    factor_type='pr_p_n',  # используем ту же логику, что для pr_p_n
    figsize=(14, 12)
)