---

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

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

# ------------------------------------------------------------
# 1. Функции группировки (без apply, без deprecation warning)
# ------------------------------------------------------------

def get_grouped_data(df, weight_col, target, group_col='pr_p_n'):
    """Группировка по pr_p_n с вычислением взвешенного среднего."""
    df_temp = df[[group_col, target, weight_col]].copy()
    df_temp['_weighted_sum'] = df_temp[target] * df_temp[weight_col]
    
    grouped = df_temp.groupby(group_col).agg({
        '_weighted_sum': 'sum',
        weight_col: 'sum'
    }).reset_index()
    
    grouped['weighted_target'] = grouped['_weighted_sum'] / grouped[weight_col]
    grouped = grouped.drop(columns=['_weighted_sum'])
    grouped = grouped[grouped[weight_col] > 0]
    
    return grouped.sort_values(group_col).rename(columns={weight_col: 'total_weight'})


def get_grouped_data_by_delta(df, weight_col, target, step=0.25, min_weight_ratio=0.01):
    """Группировка по new_delta с округлением до заданного шага."""
    df_temp = df[[target, weight_col, 'new_delta']].copy()
    df_temp['new_delta_binned'] = np.round(df_temp['new_delta'] / step) * step
    df_temp['_weighted_sum'] = df_temp[target] * df_temp[weight_col]
    
    grouped = df_temp.groupby('new_delta_binned').agg({
        '_weighted_sum': 'sum',
        weight_col: 'sum'
    }).reset_index()
    
    grouped['weighted_target'] = grouped['_weighted_sum'] / grouped[weight_col]
    grouped = grouped.drop(columns=['_weighted_sum'])
    grouped = grouped[grouped[weight_col] > 0]
    
    total_weight = grouped[weight_col].sum()
    if total_weight > 0:
        grouped = grouped[grouped[weight_col] >= min_weight_ratio * total_weight]
    
    return grouped.sort_values('new_delta_binned').rename(columns={weight_col: 'total_weight'})


def plot_with_trend(ax, grouped, x_col, 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
    
    if len(x) < 2:
        ax.text(0.5, 0.5, 'Недостаточно данных для регрессии', 
                ha='center', va='center', transform=ax.transAxes)
        ax.set_title(title)
        return
    
    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}')
    
    if x_col == 'new_delta_binned':
        ax.set_xlabel('new_delta (шаг 0.25)')
    else:
        ax.set_xlabel(x_col)
    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 add_volume_bars(ax, grouped, x_col, 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_2x2(df, weight_col, target, factor, segments, 
                             factor_type='pr_p_n', step=0.25, figsize=(14, 12)):
    """
    Построение графиков в сетке 2x2:
    - Верхний левый: все сегменты
    - Верхний правый: mass
    - Нижний левый: mid
    - Нижний правый: prem
    """
    # Выбор функции группировки
    if factor_type == 'pr_p_n':
        get_data_func = get_grouped_data
        x_col = 'pr_p_n'
    elif factor_type == 'new_delta':
        get_data_func = lambda d, wc, t: get_grouped_data_by_delta(d, wc, t, step=step)
        x_col = 'new_delta_binned'
    else:
        raise ValueError(f"Неизвестный factor_type: {factor_type}")
    
    # Подготовка данных
    all_groups = {
        'все сегменты': df
    }
    for seg in segments:
        seg_df = df[df['segment'] == seg]
        if len(seg_df) > 0:
            all_groups[seg] = seg_df
    
    # Сетка 2x2
    fig, axes = plt.subplots(2, 2, figsize=figsize)
    axes = axes.flatten()
    
    # Порядок графиков: все сегменты, mass, mid, prem
    order = ['все сегменты'] + segments
    
    for ax, label in zip(axes, order):
        group_data = all_groups.get(label)
        if group_data is None or len(group_data) == 0:
            ax.text(0.5, 0.5, 'Нет данных', ha='center', va='center', transform=ax.transAxes)
            ax.set_title(label)
            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)
            ax.set_title(label)
            continue
        
        plot_with_trend(ax, grouped, x_col=x_col, title=label, point_size=60, color='blue')
        add_volume_bars(ax, grouped, x_col=x_col, weight_col='total_weight')
    
    # Скрываем лишние оси (если сегментов меньше 3, но у нас их ровно 3)
    # для порядка, если axes больше чем order
    for i in range(len(order), len(axes)):
        axes[i].set_visible(False)
    
    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_2x2(
    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_2x2(
    df=df,
    weight_col=weight_col,
    target=target,
    factor='new_delta',
    segments=segments,
    factor_type='new_delta',
    step=0.25,
    figsize=(14, 12)
)