---

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

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

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

def get_grouped_data(df, weight_col, target, pred_col, group_col='pr_p_n'):
    """
    Группировка по pr_p_n с вычислением взвешенного среднего для факта и прогноза.
    """
    df_temp = df[[group_col, target, pred_col, weight_col]].copy()
    df_temp['_weighted_sum_target'] = df_temp[target] * df_temp[weight_col]
    df_temp['_weighted_sum_pred'] = df_temp[pred_col] * df_temp[weight_col]
    
    grouped = df_temp.groupby(group_col).agg({
        '_weighted_sum_target': 'sum',
        '_weighted_sum_pred': 'sum',
        weight_col: 'sum'
    }).reset_index()
    
    grouped['weighted_target'] = grouped['_weighted_sum_target'] / grouped[weight_col]
    grouped['weighted_pred'] = grouped['_weighted_sum_pred'] / grouped[weight_col]
    grouped = grouped.drop(columns=['_weighted_sum_target', '_weighted_sum_pred'])
    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, pred_col, step=0.1):
    """
    Группировка по new_delta с округлением до заданного шага.
    """
    df_temp = df[[target, pred_col, weight_col, 'new_delta']].copy()
    df_temp['new_delta_binned'] = np.round(df_temp['new_delta'] / step) * step
    df_temp['_weighted_sum_target'] = df_temp[target] * df_temp[weight_col]
    df_temp['_weighted_sum_pred'] = df_temp[pred_col] * df_temp[weight_col]
    
    grouped = df_temp.groupby('new_delta_binned').agg({
        '_weighted_sum_target': 'sum',
        '_weighted_sum_pred': 'sum',
        weight_col: 'sum'
    }).reset_index()
    
    grouped['weighted_target'] = grouped['_weighted_sum_target'] / grouped[weight_col]
    grouped['weighted_pred'] = grouped['_weighted_sum_pred'] / grouped[weight_col]
    grouped = grouped.drop(columns=['_weighted_sum_target', '_weighted_sum_pred'])
    grouped = grouped[grouped[weight_col] > 0]
    
    return grouped.sort_values('new_delta_binned').rename(columns={weight_col: 'total_weight'})


def plot_nonlin(ax, grouped, x_col, title, point_size=50, bar_width=None):
    """
    Построение графика с точками для факта и прогноза.
    """
    if len(grouped) == 0:
        ax.text(0.5, 0.5, 'Нет данных', ha='center', va='center', transform=ax.transAxes)
        ax.set_title(title)
        return
    
    # Точки факта (синие)
    ax.scatter(grouped[x_col], grouped['weighted_target'], 
               s=point_size, color='blue', alpha=0.7, 
               edgecolors='k', linewidth=0.5, label='Факт')
    
    # Точки прогноза (красные)
    ax.scatter(grouped[x_col], grouped['weighted_pred'], 
               s=point_size, color='red', alpha=0.7, 
               edgecolors='k', linewidth=0.5, label='Прогноз')
    
    # Подписи
    if x_col == 'new_delta_binned':
        ax.set_xlabel('new_delta (шаг 0.1)')
    else:
        ax.set_xlabel(x_col)
    ax.set_ylabel('Взвешенная доля пролонгаций')
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='best')
    
    # Подпись справа с объёмом
    ax.text(0.95, 0.95, f'Объём: {grouped["total_weight"].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, bar_width=None):
    """
    Добавление столбцов объёма на вторую ось.
    bar_width – ширина столбцов. Если None, вычисляется автоматически.
    """
    ax2 = ax.twinx()
    
    # Если ширина не задана – вычисляем как шаг между значениями
    if bar_width is None:
        if len(grouped) > 1:
            # Средний шаг между соседними значениями
            diffs = np.diff(grouped[x_col].values)
            bar_width = np.mean(diffs) * 0.8
        else:
            bar_width = 0.6
    
    ax2.bar(grouped[x_col], grouped['total_weight'], 
            alpha=0.3, color='gray', width=bar_width, 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, pred_col, factor, segments, 
                             factor_type='pr_p_n', step=0.1, figsize=(14, 12)):
    """
    Построение графиков в сетке 2x2.
    """
    # Выбор функции группировки
    if factor_type == 'pr_p_n':
        get_data_func = lambda d: get_grouped_data(d, weight_col, target, pred_col, group_col='pr_p_n')
        x_col = 'pr_p_n'
        bar_width = 0.4  # для pr_p_n ширина столбцов
    elif factor_type == 'new_delta':
        get_data_func = lambda d: get_grouped_data_by_delta(d, weight_col, target, pred_col, step=step)
        x_col = 'new_delta_binned'
        bar_width = step * 0.8  # ширина столбцов = 80% от шага (0.08 для step=0.1)
    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()
    
    # Порядок графиков
    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)
        if len(grouped) == 0:
            ax.text(0.5, 0.5, 'Нет данных после агрегации', ha='center', va='center', transform=ax.transAxes)
            ax.set_title(label)
            continue
        
        plot_nonlin(ax, grouped, x_col=x_col, title=label, point_size=50)
        add_volume_bars(ax, grouped, x_col=x_col, bar_width=bar_width)
    
    # Скрываем лишние оси
    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'
pred_col = 'prediction'
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,
    pred_col=pred_col,
    factor='pr_p_n',
    segments=segments,
    factor_type='pr_p_n',
    figsize=(14, 12)
)

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