---

# 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 с вычислением взвешенного среднего.
    Без apply – использует обычную агрегацию.
    """
    # Создаём временные колонки для взвешенных сумм
    df_temp = df.copy()
    df_temp['_weighted_sum'] = df_temp[target] * df_temp[weight_col]
    
    # Группировка и агрегация
    grouped = df_temp.groupby(group_col, as_index=False).agg({
        '_weighted_sum': 'sum',
        weight_col: 'sum'
    })
    
    # Вычисляем взвешенное среднее
    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 с округлением до заданного шага.
    Без apply.
    """
    df_temp = df.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', as_index=False).agg({
        '_weighted_sum': 'sum',
        weight_col: 'sum'
    })
    
    # Вычисляем взвешенное среднее
    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()
    grouped = grouped[grouped[weight_col] >= min_weight_ratio * total_weight]
    
    return grouped.sort_values('new_delta_binned').rename(columns={weight_col: 'total_weight'})


def get_grouped_data_duration(df, weight_col, target, group_col='duration_months'):
    """
    Группировка по duration_months.
    """
    df_temp = df.copy()
    df_temp['_weighted_sum'] = df_temp[target] * df_temp[weight_col]
    
    grouped = df_temp.groupby(group_col, as_index=False).agg({
        '_weighted_sum': 'sum',
        weight_col: 'sum'
    })
    
    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'})

# ------------------------------------------------------------
# 2. Функции построения графиков
# ------------------------------------------------------------

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)
        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')

# ------------------------------------------------------------
# 3. Универсальная функция построения для всех сегментов
# ------------------------------------------------------------

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', 'duration_months'
    """
    # Выбор функции группировки
    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'
    elif factor_type == 'duration_months':
        get_data_func = get_grouped_data_duration
        x_col = 'duration_months'
    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[f'Сегмент {seg.upper()}'] = seg_df
    
    # Создаём сетку: верхний график для всех сегментов, ниже по одному на сегмент
    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.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)
            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')
        # Столбцы объёма
        add_volume_bars(ax, grouped, x_col=x_col, weight_col='total_weight')
    
    plt.tight_layout()
    plt.show()

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

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

# 4.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)
)

# 4.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)
)

# 4.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='duration_months',
    figsize=(14, 12)
)

# ------------------------------------------------------------
# 5. (Опционально) Если нужно разделить на train/test
# ------------------------------------------------------------
# Для этого просто передайте train_df или test_df вместо df:
# plot_nonlinear_test(df=train_df, ...)
# plot_nonlinear_test(df=test_df, ...)