methods/utils.py - master - maratgaliulin/landcode_classifier

methods/utils.py
1 147 строк47 KB
maratgaliulin
23 июн 2026, 15:04
23 июн 2026, 15:04488c495
import pickle
import torch
import pandas as pd
import os
import requests
import numpy as np
import time

from urllib.parse import urlencode

from torch.utils.data import DataLoader
from torch.optim import AdamW

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score, davies_bouldin_score
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    f1_score,
    precision_score,
    recall_score,
)

from transformers import AutoTokenizer, AutoModel
from tqdm.auto import tqdm

from typing import List, Tuple, Dict, Optional, Any
import matplotlib.pyplot as plt
import json
import re
from datetime import datetime
from collections import Counter

from methods.classes.BertLandUseDataset import BertLandUseDataset
from methods.classes.BertWithNumeric import BertWithNumeric
from methods.classes.ABTestRunner import ABTestResult
from methods.classes.AreaClusterizer import AreaClusterizer
from methods.classes.BertOKSDatasetWithCluster import BertOKSDatasetWithCluster
from methods.classes.ModelEvaluator import ModelEvaluator
from methods.oks_training_config import OKS_MAX_LENGTH
from methods.oks_training import (
    train_new_model,
    test_existing_model,
    load_kody_oks_mapping,
    add_group_labels,
    build_subgroup_mask_index,
)
def display_metrics(metrics, save_dir:str = 'files/oks_model_metrics.txt') -> None:
    print(f"\n📊 БАЗОВЫЕ МЕТРИКИ:")
    print(f"   ├── Loss: {metrics['loss']:.4f}")
    print(f"   ├── Accuracy: {metrics['accuracy']:.4f} ({metrics['accuracy']*100:.2f}%)")
    print(f"   ├── Balanced Accuracy: {metrics['balanced_accuracy']:.4f} ({metrics['balanced_accuracy']*100:.2f}%)")
    print(f"   ├── MCC: {metrics['mcc']:.4f}")
    print(f"   ├── Cohen's Kappa: {metrics['kappa']:.4f}")
    print(f"   └── Top-3 Accuracy: {metrics['top_3_accuracy']:.4f} ({metrics['top_3_accuracy']*100:.2f}%)")

    # Precision/Recall/F1
    print(f"\n🎯 PRECISION/RECALL/F1:")
    print(f"   ├── Macro:  P={metrics['precision_macro']:.4f}, R={metrics['recall_macro']:.4f}, F1={metrics['f1_macro']:.4f}")
    print(f"   └── Weighted: P={metrics['precision_weighted']:.4f}, R={metrics['recall_weighted']:.4f}, F1={metrics['f1_weighted']:.4f}")
    
    # Анализ уверенности
    print(f"\n🎲 АНАЛИЗ УВЕРЕННОСТИ:")
    print(f"   ├── Средняя уверенность: {metrics['confidence_analysis']['mean_confidence']:.2%}")
    print(f"   ├── Уверенность при правильных: {metrics['confidence_analysis']['confidence_by_correct']['correct']:.2%}")
    print(f"   ├── Уверенность при ошибках: {metrics['confidence_analysis']['confidence_by_correct']['incorrect']:.2%}")
    
    # Анализ по кластерам
    print(f"\n📌 АНАЛИЗ ПО КЛАСТЕРАМ:")
    for cluster, info in metrics['cluster_analysis'].items():
        print(f"   ├── Кластер {cluster}: {info['count']} объектов, accuracy={info['accuracy']:.2%}")
    
    # Анализ ошибок
    print(f"\n⚠️ АНАЛИЗ ОШИБОК:")
    print(f"   ├── Всего ошибок: {metrics['error_analysis']['error_count']} ({metrics['error_analysis']['error_rate']:.2%})")
    print(f"   ├── Средняя уверенность при ошибках: {metrics['error_analysis']['mean_error_confidence']:.2%}")
    print(f"   ├── Ошибки с высокой уверенностью (>80%): {metrics['error_analysis']['high_confidence_errors']}")
    
    # Classification report
    print(f"\n📋 CLASSIFICATION REPORT:")
    print(metrics['classification_report'])
    
    # Сохранение метрик в файл
    with open(save_dir, 'w', encoding='utf-8') as f:
        f.write("ДЕТАЛЬНЫЙ АНАЛИЗ МОДЕЛИ ОКС\n")
        f.write("="*60 + "\n\n")
        f.write(f"Accuracy: {metrics['accuracy']:.4f}\n")
        f.write(f"Balanced Accuracy: {metrics['balanced_accuracy']:.4f}\n")
        f.write(f"MCC: {metrics['mcc']:.4f}\n")
        f.write(f"F1 Macro: {metrics['f1_macro']:.4f}\n")
        f.write(f"F1 Weighted: {metrics['f1_weighted']:.4f}\n")
        f.write(f"Top-3 Accuracy: {metrics['top_3_accuracy']:.4f}\n\n")
        f.write(metrics['classification_report'])
    
    print(f"\n✅ Метрики сохранены в {save_dir}")

def analyze_optimal_clusters(df, area_column='Площадь, кв.м', 
                             subgroup_column='код подгруппы',
                             max_clusters=15):
    """
    Анализ оптимального числа кластеров для площади
    """
    # Подготовка данных
    # Для каждого уникального кода подгруппы собираем статистику по площади
    subgroup_stats = df.groupby(subgroup_column)[area_column].agg([
        'mean', 'std', 'min', 'max', 'count'
    ]).reset_index()
    
    # Логарифмируем среднюю площадь (лучше для кластеризации)
    X = np.log1p(subgroup_stats['mean'].values).reshape(-1, 1)
    
    # Нормализуем
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # Анализ разных методов
    inertia = []
    silhouette_scores = []
    davies_bouldin_scores = []
    
    for k in range(2, max_clusters + 1):
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        labels = kmeans.fit_predict(X_scaled)
        
        inertia.append(kmeans.inertia_)
        silhouette_scores.append(silhouette_score(X_scaled, labels))
        davies_bouldin_scores.append(davies_bouldin_score(X_scaled, labels))
        
        print(f"k={k}: Inertia={kmeans.inertia_:.0f}, "
              f"Silhouette={silhouette_scores[-1]:.3f}, "
              f"DB={davies_bouldin_scores[-1]:.3f}")
    
    # Визуализация
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    axes[0].plot(range(2, max_clusters + 1), inertia, 'bo-')
    axes[0].set_xlabel('Number of clusters')
    axes[0].set_ylabel('Inertia')
    axes[0].set_title('Elbow Method')
    axes[0].grid(True)
    
    axes[1].plot(range(2, max_clusters + 1), silhouette_scores, 'go-')
    axes[1].set_xlabel('Number of clusters')
    axes[1].set_ylabel('Silhouette Score')
    axes[1].set_title('Silhouette Analysis')
    axes[1].grid(True)
    
    axes[2].plot(range(2, max_clusters + 1), davies_bouldin_scores, 'ro-')
    axes[2].set_xlabel('Number of clusters')
    axes[2].set_ylabel('Davies-Bouldin Score')
    axes[2].set_title('Davies-Bouldin Index')
    axes[2].grid(True)
    
    plt.tight_layout()
    plt.savefig('area_cluster_analysis.png', dpi=150)
    plt.show()
    
    # Рекомендация
    # Ищем "локоть" (изменение инерции)
    diffs = np.diff(inertia)
    diffs2 = np.diff(diffs)
    optimal_k = np.argmax(diffs2) + 3  # +3 потому что diff сдвигает
    
    print(f"\n📊 Рекомендуемое число кластеров: {optimal_k}")
    
    return optimal_k, scaler

def return_is_in_text_columns(dataframe_columns:list, text_columns:list) -> bool:
    is_in_text_columns = False
    for col in text_columns:
        if col in dataframe_columns:
            is_in_text_columns = True
        else:
            is_in_text_columns = False
            break

    return is_in_text_columns

def save_full_model(model, tokenizer, label_encoder, text_columns, model_path, components_path):
        
    torch.save({
        'model_state_dict': model.state_dict(),
        'text_columns': text_columns,
        'model_config': {
            'bert_model_name': 'DeepPavlov/rubert-base-cased',
            'num_labels': len(label_encoder.classes_)
        }
    }, model_path)
    
    with open(components_path, 'wb') as f:
        pickle.dump({
            'tokenizer': tokenizer,
            'label_encoder': label_encoder
        }, f)
    
    print(f"Model saved to {model_path}")
    print(f"Components saved to {components_path}")


def return_single_dataframe_from_pickle_or_from_excel(pickle_df_dir:str, excel_df_dir:str, cols_to_drop:list, text_columns:list, target_column:str) -> pd.DataFrame:
    """Метод считывает файл Excel и сохраняет датафрейм в Pickle формате. Если уже есть Piclke файл, то считывает с него датафрейм"""
    
    if(os.path.isfile(pickle_df_dir)):
        df = pd.read_pickle(pickle_df_dir)
    else:
        

        df = pd.read_excel(excel_df_dir)     
        
        df.drop(columns=cols_to_drop, inplace=True)        
        
        df['combined_text'] = df[text_columns].fillna('').astype(str).agg(' | '.join, axis=1).str.lower()   
        
        df[text_columns] = df[text_columns].fillna('')

        df['Площадь, кв.м'] = df['Площадь, кв.м'].fillna(0.0)
        
        df.dropna(axis=0, subset=['combined_text', target_column], inplace=True)        
        
        df.to_pickle(pickle_df_dir)

    return df

def return_single_dataframe_from_pickle_or_from_excel_with_clusters(
        pickle_df_dir: str, 
        excel_df_dir: str, 
        clusterizer_dir: str,
        text_columns: list, 
        target_column: str
) -> pd.DataFrame:
    """
    Считывает файл Excel, добавляет кластеры площади и сохраняет в Pickle
    """
    from methods.classes.AreaClusterizer import AreaClusterizer
    
    if os.path.isfile(pickle_df_dir):
        df = pd.read_pickle(pickle_df_dir)
    else:
        # Загружаем кластеризатор
        area_clusterizer = AreaClusterizer()
        area_clusterizer.load(clusterizer_dir)
        
        # Читаем Excel
        df = pd.read_excel(excel_df_dir)
        
        # Удаляем ненужные колонки
        columns_all = df.columns.to_list()

        for col in columns_all:
            if (
                (col != target_column) and 
                (col != 'Площадь, кв.м') and 
                (col not in text_columns)
                ):
                
                df.drop(columns=[col], inplace=True, errors='ignore')
        
        # Заполняем пропуски
        df[text_columns] = df[text_columns].astype('str').fillna('')
        df['Площадь, кв.м'] = df['Площадь, кв.м'].fillna(0.0)
        
        # Добавляем кластеры
        df['area_cluster'] = area_clusterizer.predict(df['Площадь, кв.м'].values)
        
        # Добавляем текст кластера в комбинированный текст
        df['cluster_text'] = 'кластер_' + df['area_cluster'].astype(str)
        
        # Создаём комбинированный текст с учётом кластера
        text_columns_with_cluster = text_columns + ['cluster_text']
        df['combined_text'] = df[text_columns_with_cluster].fillna('').astype(str).agg(' | '.join, axis=1).str.lower()
        
        # Удаляем строки с пропущенными целевыми значениями
        df.dropna(axis=0, subset=['combined_text', target_column], inplace=True)
        df[target_column] = df[target_column].astype(str).str.strip()

        # Сохраняем
        df.to_pickle(pickle_df_dir)
    
    return df


def return_dataframes_from_pickle_or_from_excel(pickle_df_ids_dir:str, pickle_df_intermed_dir:str, excel_df_ids_dir:str, excel_df_intermed_dir:str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Метод считывает файл Excel и сохраняет датафрейм в Pickle формате. Если уже есть Piclke файл, то считывает с него датафрейм"""
    
    cols_to_drop = [
    'Код вида разрешенного использования (в соответствии с классификатором видов разрешенного использования земельных участков)',
    'Условный сегмент',
    'Сегмент',
    'метод расчета'
    ]

    if(os.path.isfile(pickle_df_ids_dir) and os.path.isfile(pickle_df_intermed_dir)):
        df_ids = pd.read_pickle(pickle_df_ids_dir)
        df_intermed = pd.read_pickle(pickle_df_intermed_dir)
    else:
        
        text_columns = [
        'Вид земельного участка',
        'Категория земель',
        'Вид использования участка по документу (САМЫЙ ГЛАВНЫЙ АТРИБУТ - Приоритет 1)',
        'Разрешенное использование (текстовое описание) - Приоритет 2'
    ]
        df_ids = pd.read_excel(excel_df_ids_dir, sheet_name='кодыВРИ')
        df_intermed = pd.read_excel(excel_df_intermed_dir, sheet_name='данные')        
        
        df_ids.drop(columns=cols_to_drop, axis=1, inplace=True)
        df_intermed.drop(columns=['Ранее учтенный'], axis=1, inplace=True)
        
        
        
        df_intermed['combined_text'] = df_intermed[text_columns].fillna('').astype(str).agg(' | '.join, axis=1).str.lower()
        df_ids['combined_text'] = df_ids['Наименование вида использования'].fillna('').astype(str).str.lower()
        
        df_intermed.dropna(axis=0, subset=['combined_text', 'Код расчёта вида использования (ГБУ)'], inplace=True)
        df_ids.dropna(axis=0, subset=['combined_text', 'Код расчета вида использования'], inplace=True)
        
        df_ids.to_pickle(pickle_df_ids_dir)
        df_intermed.to_pickle(pickle_df_intermed_dir)

    return df_ids, df_intermed

def download_model_and_components_from_cloud(public_key:str, file_directory:str)->None:
    base_url = 'https://cloud-api.yandex.net/v1/disk/public/resources/download?'

    # Получаем загрузочную ссылку
    final_url = base_url + urlencode(dict(public_key=public_key))
    response = requests.get(final_url)
    download_url = response.json()['href']

    # Загружаем файл и сохраняем его
    download_response = requests.get(download_url)
    with open(file_directory, 'wb') as f:   # Здесь укажите нужный путь к файлу
        f.write(download_response.content)

def downsample_frequent_codes(df, frequent_codes=['13:021', '02:010'], sample_ratio=0.1, random_state=42):
    """
    Даунсемплинг частых кодов для создания сбалансированного датасета
    
    Parameters:
    -----------
    df : pd.DataFrame
        Исходный датасет
    frequent_codes : list
        Список кодов для даунсемплинга
    sample_ratio : float
        Доля записей для сохранения (от 0 до 1)
    random_state : int
        Seed для воспроизводимости
    
    Returns:
    --------
    pd.DataFrame : Сбалансированный датасет
    """
    from sklearn.utils import shuffle
    
    # 1. Выделяем датасет без частых кодов
    df_without_frequent = df[~df['Код расчёта вида использования (ГБУ)'].isin(frequent_codes)].copy()
    print(f"Записей без частых кодов: {len(df_without_frequent)}")
    
    # 2. Работа с каждым частым кодом отдельно
    sampled_dfs = []
    
    for code in frequent_codes:
        code_df = df[df['Код расчёта вида использования (ГБУ)'] == code].copy()
        code_df = shuffle(code_df, random_state=random_state)
        sample_size = max(1, int(len(code_df) * sample_ratio))
        sampled_code_df = code_df.head(sample_size)
        sampled_dfs.append(sampled_code_df)
        
        print(f"Код {code}: было {len(code_df)} записей, взято {sample_size} ({sample_size/len(code_df)*100:.1f}%)")
    
    # 3. Объединяем и перемешиваем
    df_balanced = pd.concat([df_without_frequent] + sampled_dfs, ignore_index=True)
    df_balanced = shuffle(df_balanced, random_state=random_state)
    
    print(f"\nИтоговый сбалансированный датасет: {len(df_balanced)} записей")
    
    return df_balanced


def prepare_balanced_dataset(df, text_columns, target_column='Код расчёта вида использования (ГБУ)', 
                             test_size=0.2, random_state=42):
    """
    Подготовка сбалансированного датасета для обучения
    
    Parameters:
    -----------
    df : pd.DataFrame
        Сбалансированный датасет
    text_columns : list
        Список текстовых колонок для объединения
    target_column : str
        Название целевой колонки
    test_size : float
        Доля тестовой выборки
    random_state : int
        Seed для воспроизводимости
    
    Returns:
    --------
    tuple: (train_texts, val_texts, train_labels, val_labels, label_encoder)
    """
    # Создаём комбинированный текст
    df['combined_text'] = df[text_columns].fillna('').astype(str).agg(' | '.join, axis=1)
    
    # Кодируем метки
    label_encoder = LabelEncoder()
    df['label'] = label_encoder.fit_transform(df[target_column])
    
    # Разделяем на train/val
    train_texts, val_texts, train_labels, val_labels = train_test_split(
        df['combined_text'].values,
        df['label'].values,
        test_size=test_size,
        random_state=random_state,
        # stratify=df['label'].values
    )
    
    print(f"Train size: {len(train_texts)}, Val size: {len(val_texts)}")
    print(f"Number of classes: {len(label_encoder.classes_)}")
    
    return train_texts, val_texts, train_labels, val_labels, label_encoder


def train_second_model(train_texts, val_texts, train_labels, val_labels, 
                       label_encoder, model_save_path, components_save_path,
                       model_name='DeepPavlov/rubert-base-cased', 
                       batch_size=16, epochs=5, learning_rate=2e-5,
                       device=None):
    """
    Обучение второй модели на сбалансированном датасете
    
    Parameters:
    -----------
    train_texts, val_texts : array-like
        Текстовые данные для обучения и валидации
    train_labels, val_labels : array-like
        Метки для обучения и валидации
    label_encoder : LabelEncoder
        Энкодер для декодирования предсказаний
    model_save_path : str
        Путь для сохранения модели (.pth)
    components_save_path : str
        Путь для сохранения компонентов (.pkl)
    model_name : str
        Название предобученной модели
    batch_size : int
        Размер батча
    epochs : int
        Количество эпох
    learning_rate : float
        Скорость обучения
    device : torch.device
        Устройство для обучения (cuda/cpu)
    
    Returns:
    --------
    tuple: (model, tokenizer, history)
    """
    
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    print(f"Using device: {device}")
    
    # Загрузка токенизатора
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # Создание датасетов
    train_dataset = BertLandUseDataset(train_texts, train_labels, tokenizer)
    val_dataset = BertLandUseDataset(val_texts, val_labels, tokenizer)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    
    # Загрузка BERT модели
    bert_model = AutoModel.from_pretrained(model_name)
    
    # Создание классификатора (без числовых признаков)
    model = BertWithNumeric(bert_model, num_labels=len(label_encoder.classes_))
    model.to(device)
    
    # Оптимизатор и функция потерь
    optimizer = AdamW(model.parameters(), lr=learning_rate)
    criterion = torch.nn.CrossEntropyLoss()
    
    # История обучения
    history = {
        'train_loss': [],
        'val_loss': [],
        'val_accuracy': []
    }
    
    # Обучение
    best_val_accuracy = 0
    
    for epoch in range(epochs):
        print(f"\n{'='*50}")
        print(f"Epoch {epoch+1}/{epochs}")
        print(f"{'='*50}")
        
        # Training
        model.train()
        total_train_loss = 0
        
        train_pbar = tqdm(train_loader, desc=f"Training Epoch {epoch+1}")
        for batch in train_pbar:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            
            optimizer.zero_grad()
            outputs = model(input_ids, attention_mask)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            total_train_loss += loss.item()
            train_pbar.set_postfix({'loss': loss.item()})
        
        avg_train_loss = total_train_loss / len(train_loader)
        history['train_loss'].append(avg_train_loss)
        
        # Validation
        model.eval()
        total_val_loss = 0
        correct = 0
        total = 0
        
        val_pbar = tqdm(val_loader, desc=f"Validation Epoch {epoch+1}")
        with torch.no_grad():
            for batch in val_pbar:
                input_ids = batch['input_ids'].to(device)
                attention_mask = batch['attention_mask'].to(device)
                labels = batch['labels'].to(device)
                
                outputs = model(input_ids, attention_mask)
                loss = criterion(outputs, labels)
                
                total_val_loss += loss.item()
                
                predictions = torch.argmax(outputs, dim=1)
                correct += (predictions == labels).sum().item()
                total += len(labels)
                
                val_pbar.set_postfix({'loss': loss.item()})
        
        avg_val_loss = total_val_loss / len(val_loader)
        val_accuracy = correct / total
        
        history['val_loss'].append(avg_val_loss)
        history['val_accuracy'].append(val_accuracy)
        
        print(f"\nEpoch {epoch+1} Summary:")
        print(f"  Train Loss: {avg_train_loss:.4f}")
        print(f"  Val Loss: {avg_val_loss:.4f}")
        print(f"  Val Accuracy: {val_accuracy:.4f} ({val_accuracy*100:.2f}%)")
        
        # Сохраняем лучшую модель
        if val_accuracy > best_val_accuracy:
            best_val_accuracy = val_accuracy
            print(f"  ✓ New best model! Saving...")
            
            # Сохраняем модель
            torch.save({
                'model_state_dict': model.state_dict(),
                'text_columns': None,  # Будет заполнено при сохранении компонентов
                'model_config': {
                    'bert_model_name': model_name,
                    'num_labels': len(label_encoder.classes_)
                }
            }, model_save_path)
    
    # Сохраняем компоненты
    components = {
        'tokenizer': tokenizer,
        'label_encoder': label_encoder,
        'text_columns': None  # Будет заполнено позже
    }
    
    with open(components_save_path, 'wb') as f:
        pickle.dump(components, f)
    
    print(f"\n✅ Model saved to: {model_save_path}")
    print(f"✅ Components saved to: {components_save_path}")
    print(f"\nBest validation accuracy: {best_val_accuracy:.4f} ({best_val_accuracy*100:.2f}%)")
    
    return model, tokenizer, history



def prepare_test_data(df: pd.DataFrame, 
                      text_columns: List[str],
                      label_encoder,
                      target_column: str = 'Код расчёта вида использования (ГБУ)') -> pd.DataFrame:
    """
    Подготовка тестовых данных для A/B тестирования
    
    Parameters:
    -----------
    df : pd.DataFrame
        Исходный датасет
    text_columns : List[str]
        Список текстовых колонок
    label_encoder : LabelEncoder
        Энкодер меток (должен быть уже обучен)
    target_column : str
        Название целевой колонки
    
    Returns:
    --------
    pd.DataFrame: Подготовленный датасет с колонками 'combined_text' и 'label'
    """
    
    df_test = df.copy()
    
    # 1. Удаляем строки с пропущенными значениями в целевой колонке
    initial_count = len(df_test)
    df_test = df_test.dropna(subset=[target_column])
    dropped_count = initial_count - len(df_test)
    
    if dropped_count > 0:
        print(f"⚠️ Удалено {dropped_count} строк с пропущенными значениями в колонке '{target_column}'")
    
    if df_test.empty:
        raise ValueError("После удаления NaN в целевой колонке не осталось данных для тестирования")
    
    # 2. Проверяем, что все значения в целевой колонке известны энкодеру
    unique_values = df_test[target_column].unique()
    known_values = set(label_encoder.classes_)
    unknown_values = set(unique_values) - known_values
    
    if unknown_values:
        print(f"⚠️ Найдены неизвестные значения: {unknown_values}")
        # Удаляем строки с неизвестными значениями
        df_test = df_test[df_test[target_column].isin(known_values)]
        print(f"   После удаления неизвестных: {len(df_test)} строк")
    
    # 3. Создаём комбинированный текст
    df_test['combined_text'] = df_test[text_columns].fillna('').astype(str).agg(' | '.join, axis=1)
    
    # 4. Кодируем метки
    df_test['label'] = label_encoder.transform(df_test[target_column])
    
    print(f"✅ Подготовлено {len(df_test)} записей для тестирования")
    
    return df_test


def analyze_switched_cases(result_df: pd.DataFrame, 
                           test_df: pd.DataFrame,
                           label_encoder) -> pd.DataFrame:
    """
    Анализ случаев, где произошла замена на модель B
    
    Parameters:
    -----------
    result_df : pd.DataFrame
        DataFrame с результатами A/B тестирования (из ABTestResult.predictions_df)
    test_df : pd.DataFrame
        Исходный тестовый датасет
    label_encoder : LabelEncoder
        Энкодер для декодирования меток
    
    Returns:
    --------
    pd.DataFrame: Детальный анализ замен
    """
    
    # Фильтруем только заменённые случаи
    switched_df = result_df[result_df['source'] == 'model_b'].copy()
    
    if switched_df.empty:
        print("Нет замен на модель B")
        return switched_df
    
    # Декодируем метки
    switched_df['true_code'] = label_encoder.inverse_transform(switched_df['true_label'])
    switched_df['pred_a_code'] = label_encoder.inverse_transform(switched_df['pred_a'])
    switched_df['pred_b_code'] = label_encoder.inverse_transform(switched_df['pred_b'])
    switched_df['final_code'] = label_encoder.inverse_transform(switched_df['final_pred'])
    
    # Добавляем информацию из исходного датасета
    switched_df = switched_df.merge(
        test_df[['Код расчёта вида использования (ГБУ)', 'combined_text']],
        left_on='code',
        right_on='Код расчёта вида использования (ГБУ)',
        how='left'
    )
    
    # Анализ качества замен
    switched_df['switch_was_correct'] = switched_df['is_correct_final']
    switched_df['improved'] = (~switched_df['is_correct_a']) & switched_df['is_correct_final']
    switched_df['worsened'] = switched_df['is_correct_a'] & (~switched_df['is_correct_final'])
    switched_df['unchanged'] = (switched_df['is_correct_a'] == switched_df['is_correct_final'])
    
    return switched_df


def plot_ab_test_results(result: ABTestResult, save_path: str = None):
    """
    Визуализация результатов A/B тестирования
    
    Parameters:
    -----------
    result : ABTestResult
        Результат A/B тестирования
    save_path : str
        Путь для сохранения графика (опционально)
    """
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # 1. Сравнение точности
    models = ['Model A', 'Model B', 'Hybrid']
    accuracies = [result.accuracy_model_a, result.accuracy_model_b, result.accuracy_hybrid]
    
    axes[0, 0].bar(models, accuracies, color=['#3498db', '#e74c3c', '#2ecc71'])
    axes[0, 0].set_ylabel('Accuracy')
    axes[0, 0].set_title('Сравнение точности моделей')
    axes[0, 0].set_ylim([0, 1])
    for i, v in enumerate(accuracies):
        axes[0, 0].text(i, v + 0.01, f'{v:.3f}', ha='center', fontweight='bold')
    
    # 2. Precision/Recall/F1
    metrics = ['Precision', 'Recall', 'F1']
    model_a_scores = [result.precision_model_a, result.recall_model_a, result.f1_model_a]
    model_b_scores = [result.precision_model_b, result.recall_model_b, result.f1_model_b]
    hybrid_scores = [result.precision_hybrid, result.recall_hybrid, result.f1_hybrid]
    
    x = np.arange(len(metrics))
    width = 0.25
    
    axes[0, 1].bar(x - width, model_a_scores, width, label='Model A', color='#3498db')
    axes[0, 1].bar(x, model_b_scores, width, label='Model B', color='#e74c3c')
    axes[0, 1].bar(x + width, hybrid_scores, width, label='Hybrid', color='#2ecc71')
    axes[0, 1].set_ylabel('Score')
    axes[0, 1].set_title('Precision, Recall, F1 Score')
    axes[0, 1].set_xticks(x)
    axes[0, 1].set_xticklabels(metrics)
    axes[0, 1].legend()
    axes[0, 1].set_ylim([0, 1])
    
    # 3. Статистика замен
    labels = ['Handled by A', 'Switched to B']
    counts = [result.samples_handled_by_model_a, result.samples_switched_to_b]
    
    axes[1, 0].pie(counts, labels=labels, autopct='%1.1f%%', 
                   colors=['#3498db', '#e74c3c'], startangle=90)
    axes[1, 0].set_title(f'Распределение предсказаний (всего: {result.total_samples})')
    
    # 4. Качество замен
    if result.samples_switched_to_b > 0:
        switch_labels = ['Correct switches', 'Incorrect switches']
        switch_counts = [result.samples_switched_correctly, result.samples_switched_incorrectly]
        
        axes[1, 1].bar(switch_labels, switch_counts, color=['#2ecc71', '#e74c3c'])
        axes[1, 1].set_ylabel('Count')
        axes[1, 1].set_title('Качество замен на модель B')
        for i, v in enumerate(switch_counts):
            axes[1, 1].text(i, v + 1, str(v), ha='center', fontweight='bold')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"График сохранён: {save_path}")
    
    plt.show()


def plot_threshold_analysis(threshold_results: Dict[float, ABTestResult], save_path: str = None):
    """
    Визуализация анализа порогов
    
    Parameters:
    -----------
    threshold_results : Dict[float, ABTestResult]
        Результаты для разных порогов
    save_path : str
        Путь для сохранения графика
    """
    
    thresholds = sorted(threshold_results.keys())
    accuracies_a = [threshold_results[t].accuracy_model_a for t in thresholds]
    accuracies_b = [threshold_results[t].accuracy_model_b for t in thresholds]
    accuracies_hybrid = [threshold_results[t].accuracy_hybrid for t in thresholds]
    switch_rates = [threshold_results[t].samples_switched_to_b / threshold_results[t].total_samples * 100 for t in thresholds]
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # График точности
    axes[0].plot(thresholds, accuracies_a, 'o-', label='Model A', color='#3498db', linewidth=2, markersize=8)
    axes[0].plot(thresholds, accuracies_b, 's-', label='Model B', color='#e74c3c', linewidth=2, markersize=8)
    axes[0].plot(thresholds, accuracies_hybrid, '^-', label='Hybrid', color='#2ecc71', linewidth=2, markersize=8)
    axes[0].set_xlabel('Threshold')
    axes[0].set_ylabel('Accuracy')
    axes[0].set_title('Accuracy vs Threshold')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # График частоты замен
    axes[1].plot(thresholds, switch_rates, 'd-', label='Switch Rate', color='#e67e22', linewidth=2, markersize=8)
    axes[1].set_xlabel('Threshold')
    axes[1].set_ylabel('Switch Rate (%)')
    axes[1].set_title('Switch Rate vs Threshold')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"График сохранён: {save_path}")
    
    plt.show()


# --- OKS A/B test helpers (evaluation only, no training) ---

AB_TEST_LOGS_DIR = 'files/ab_test_logs'


def get_oks_ab_test_config() -> Dict[str, Any]:
    return {
        'excel_df_dir_long': 'files/excel/Перечень ОКС (актуальные).xlsx',
        'pickle_df_dir_long': 'files/pickle/df_oks_long.pkl',
        'pickle_df_dir_no_cluster': 'files/pickle/df_oks_long_no_cluster_ab.pkl',
        'text_columns_long': [
            'вид ОН', 'Наименование объекта', 'Назначение здания',
            'Назначение сооружения', 'Назначение помещения',
            'Материал стен, в оценку', 'Материал стен (квартир), в оценку',
            'Количество этажей', 'Количество подземных этажей',
        ],
        'target_column': 'код подгруппы',
        'clusterizer_dir_long': 'files/pickle/clusterizer_long.pkl',
        'num_clusters': 6,
        'model_directory_long': 'files/pickle/ml_models/oks/rubert_landuse_model_oks_long.pth',
        'model_components_directory_long': 'files/pickle/ml_models/oks/model_components_oks_long.pkl',
        'without_area_dir': 'files/pickle/ml_models/oks/old/without_area',
        'old_cluster_dir': 'files/pickle/ml_models/oks/old',
        'use_best_checkpoint': False,
        'eval_max_rows': None,
    }


def resolve_checkpoint_paths(model_path: str, components_path: str, use_best: bool = False) -> Tuple[str, str]:
    if use_best:
        best_model = model_path.replace('.pth', '_best.pth')
        best_components = components_path.replace('.pkl', '_best.pkl')
        if os.path.isfile(best_model) and os.path.isfile(best_components):
            return best_model, best_components
    return model_path, components_path


def inspect_oks_checkpoint(model_path: str, components_path: str) -> Dict[str, Any]:
    info = {
        'model_path': model_path,
        'components_path': components_path,
        'model_exists': os.path.isfile(model_path),
        'components_exists': os.path.isfile(components_path),
        'model_size_mb': None,
        'components_size_mb': None,
        'text_columns': None,
        'num_labels': None,
        'bert_model_name': None,
        'label_encoder_classes': None,
        'components_keys': None,
        'best_variant_exists': False,
        'error': None,
    }
    if info['model_exists']:
        info['model_size_mb'] = round(os.path.getsize(model_path) / (1024 * 1024), 2)
        best_path = model_path.replace('.pth', '_best.pth')
        info['best_variant_exists'] = os.path.isfile(best_path)
        try:
            checkpoint = torch.load(model_path, map_location='cpu')
            info['text_columns'] = checkpoint.get('text_columns')
            config = checkpoint.get('model_config', {})
            info['num_labels'] = config.get('num_labels')
            info['bert_model_name'] = config.get('bert_model_name')
        except Exception as e:
            info['error'] = str(e)
    if info['components_exists']:
        info['components_size_mb'] = round(os.path.getsize(components_path) / (1024 * 1024), 2)
        try:
            with open(components_path, 'rb') as f:
                components = pickle.load(f)
            info['components_keys'] = list(components.keys())
            if 'label_encoder' in components:
                info['label_encoder_classes'] = len(components['label_encoder'].classes_)
        except Exception as e:
            info['error'] = info['error'] or str(e)
    return info


def detect_oks_checkpoint_architecture(model_path: str) -> str:
    """Return 'cluster', 'numeric_with_area', or 'numeric' from checkpoint weights."""
    checkpoint = torch.load(model_path, map_location='cpu')
    state = checkpoint['model_state_dict']
    if 'cluster_embedding.weight' in state:
        return 'cluster'
    in_features = int(state['classifier.weight'].shape[1])
    if in_features == 769:
        return 'numeric_with_area'
    return 'numeric'


def normalize_oks_labels_for_eval(series, mode: str = 'full_string') -> pd.Series:
    s = pd.Series(series).astype(str).str.strip()
    if mode == 'code_only':
        return s.str.split().str[0]
    return s


def evaluate_oks_string_predictions(y_true, y_pred, confidences=None) -> Dict[str, Any]:
    y_true = pd.Series(y_true).astype(str).str.strip().reset_index(drop=True)
    y_pred = pd.Series(y_pred).astype(str).str.strip().reset_index(drop=True)
    valid = y_true.notna() & (y_true != '') & (y_true.str.lower() != 'nan')
    y_true = y_true[valid].reset_index(drop=True)
    y_pred = y_pred[valid].reset_index(drop=True)
    if confidences is not None:
        confidences = pd.Series(confidences).reset_index(drop=True).iloc[valid.values].reset_index(drop=True)

    n = len(y_true)
    if n == 0:
        return {'n_samples': 0, 'accuracy': 0.0}

    correct = (y_true == y_pred)
    metrics = {
        'n_samples': n,
        'n_correct': int(correct.sum()),
        'accuracy': float(accuracy_score(y_true, y_pred)),
        'balanced_accuracy': float(balanced_accuracy_score(y_true, y_pred)),
        'precision_macro': float(precision_score(y_true, y_pred, average='macro', zero_division=0)),
        'recall_macro': float(recall_score(y_true, y_pred, average='macro', zero_division=0)),
        'f1_macro': float(f1_score(y_true, y_pred, average='macro', zero_division=0)),
        'precision_weighted': float(precision_score(y_true, y_pred, average='weighted', zero_division=0)),
        'recall_weighted': float(recall_score(y_true, y_pred, average='weighted', zero_division=0)),
        'f1_weighted': float(f1_score(y_true, y_pred, average='weighted', zero_division=0)),
    }

    if confidences is not None and len(confidences) == n:
        metrics['mean_confidence'] = float(confidences.mean())
        metrics['mean_confidence_correct'] = float(confidences[correct].mean()) if correct.any() else None
        metrics['mean_confidence_incorrect'] = float(confidences[~correct].mean()) if (~correct).any() else None

    pairs = Counter(zip(y_true[~correct], y_pred[~correct]))
    metrics['top_errors'] = [{'true': t, 'pred': p, 'count': c} for (t, p), c in pairs.most_common(10)]
    return metrics


def _format_metrics_block(metrics: Dict[str, Any]) -> str:
    if not metrics or metrics.get('n_samples', 0) == 0:
        return '  (нет данных)\n'
    lines = [
        f"  samples: {metrics['n_samples']}",
        f"  accuracy: {metrics['accuracy']:.4f} ({metrics['accuracy']*100:.2f}%)",
        f"  balanced_accuracy: {metrics['balanced_accuracy']:.4f}",
        f"  f1_macro: {metrics['f1_macro']:.4f}",
        f"  f1_weighted: {metrics['f1_weighted']:.4f}",
    ]
    if 'mean_confidence' in metrics:
        lines.append(f"  mean_confidence: {metrics['mean_confidence']:.4f}")
    if metrics.get('mean_confidence_correct') is not None:
        lines.append(f"  mean_confidence_correct: {metrics['mean_confidence_correct']:.4f}")
    if metrics.get('mean_confidence_incorrect') is not None:
        lines.append(f"  mean_confidence_incorrect: {metrics['mean_confidence_incorrect']:.4f}")
    if metrics.get('top_errors'):
        lines.append('  top_errors:')
        for err in metrics['top_errors']:
            lines.append(f"    {err['true']} -> {err['pred']}: {err['count']}")
    return '\n'.join(lines) + '\n'


def write_ab_test_log(path: str, title: str, sections: Dict[str, Any]) -> None:
    os.makedirs(os.path.dirname(path) or '.', exist_ok=True)
    lines = [
        '=' * 70,
        title,
        f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        '=' * 70,
        '',
    ]
    for section_name, content in sections.items():
        lines.append('-' * 70)
        lines.append(section_name)
        lines.append('-' * 70)
        if isinstance(content, dict):
            if 'accuracy' in content and 'n_samples' in content:
                lines.append(_format_metrics_block(content))
            else:
                for k, v in content.items():
                    if isinstance(v, dict) and 'n_samples' in v:
                        lines.append(f"{k}:")
                        lines.append(_format_metrics_block(v))
                    elif isinstance(v, list):
                        lines.append(f"{k}:")
                        for item in v:
                            lines.append(f"  {item}")
                    else:
                        lines.append(f"  {k}: {v}")
                lines.append('')
        elif isinstance(content, list):
            for item in content:
                lines.append(f"  {item}")
            lines.append('')
        else:
            lines.append(str(content))
            lines.append('')
    with open(path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(lines))
    print(f"Log saved: {path}")


def _normalize_text_for_oks_mapping(value) -> str:
    if pd.isna(value):
        return ''
    text = str(value).lower().strip()
    text = re.sub(r'\s+', ' ', text)
    text = re.sub(r'[.!?,;:]$', '', text)
    return text


def get_oks_exact_mapping_mask(
    df: pd.DataFrame,
    search_column: str = 'Наименование объекта',
    mapping_path: str = 'files/mapping_oks_1.json',
) -> pd.Series:
    if not os.path.isfile(mapping_path):
        return pd.Series(False, index=df.index)
    with open(mapping_path, 'r', encoding='utf-8') as f:
        mapping = json.load(f)
    if search_column not in df.columns:
        return pd.Series(False, index=df.index)
    normalized = df[search_column].map(_normalize_text_for_oks_mapping)
    return normalized.isin(mapping)


def prepare_oks_eval_df_without_cluster(
    pickle_df_dir: str,
    excel_df_dir: str,
    text_columns: list,
    target_column: str,
) -> pd.DataFrame:
    if os.path.isfile(pickle_df_dir):
        df = pd.read_pickle(pickle_df_dir)
    else:
        df = pd.read_excel(excel_df_dir)
        for col in df.columns.to_list():
            if col != target_column and col != 'Площадь, кв.м' and col not in text_columns:
                df.drop(columns=[col], inplace=True, errors='ignore')
        df[text_columns] = df[text_columns].astype('str').fillna('')
        df['Площадь, кв.м'] = df['Площадь, кв.м'].fillna(0.0)
        df.dropna(axis=0, subset=[target_column], inplace=True)
        df.to_pickle(pickle_df_dir)

    for col in text_columns:
        if col not in df.columns:
            df[col] = ''
    df['combined_text'] = df[text_columns].fillna('').astype(str).agg(' | '.join, axis=1).str.lower()
    return df


def run_oks_predictor_eval(
    predictor,
    df: pd.DataFrame,
    target_column: str,
    mapping_mask: Optional[pd.Series] = None,
    eval_max_rows: Optional[int] = None,
    target_eval_mode: str = 'full_string',
) -> Tuple[Dict[str, Any], Dict[str, Any], pd.DataFrame]:
    eval_df = df.copy()
    if eval_max_rows is not None and eval_max_rows > 0:
        eval_df = eval_df.iloc[:eval_max_rows].copy()
    if mapping_mask is not None:
        mapping_mask = mapping_mask.reindex(eval_df.index).fillna(False)

    result = predictor.predict(eval_df, return_proba=True, show_progress=True)
    y_true = normalize_oks_labels_for_eval(eval_df[target_column], target_eval_mode)
    y_pred = normalize_oks_labels_for_eval(result['predicted_code'], target_eval_mode)
    conf = result['prediction_confidence'] if 'prediction_confidence' in result.columns else None

    metrics_all = evaluate_oks_string_predictions(y_true, y_pred, conf)
    if mapping_mask is not None:
        model_only = ~mapping_mask
        metrics_model_only = evaluate_oks_string_predictions(
            y_true[model_only], y_pred[model_only], conf[model_only] if conf is not None else None,
        )
    else:
        metrics_model_only = {'n_samples': 0}

    return metrics_all, metrics_model_only, result


def pick_checkpoint_pair(pairs: List[Dict[str, str]], use_best: bool = False) -> Dict[str, str]:
    if not pairs:
        raise FileNotFoundError('No checkpoint pairs found')
    if use_best:
        for p in pairs:
            if '_best.pth' in p['model_path']:
                return p
    for p in pairs:
        if '_best.pth' not in p['model_path']:
            return p
    return pairs[0]


def discover_pth_pkl_pairs(directory: str) -> List[Dict[str, str]]:
    if not os.path.isdir(directory):
        return []
    pairs = []
    pkl_files = [f for f in os.listdir(directory) if f.endswith('.pkl')]
    for name in sorted(os.listdir(directory)):
        if not name.endswith('.pth'):
            continue
        full = os.path.join(directory, name)
        stem = name[:-4]
        comp = None
        for pfn in pkl_files:
            if stem.split('_')[-1] in pfn or 'components' in pfn:
                comp = os.path.join(directory, pfn)
                break
        if comp is None and pkl_files:
            comp = os.path.join(directory, pkl_files[0])
        pairs.append({'model_path': full, 'components_path': comp or '(not found)'})
    return pairs


def format_checkpoint_audit_lines(pairs: List[Dict[str, str]]) -> List[str]:
    lines = []
    for pair in pairs:
        if pair['components_path'] == '(not found)':
            lines.append(f"model: {pair['model_path']} | components: NOT FOUND")
            continue
        info = inspect_oks_checkpoint(pair['model_path'], pair['components_path'])
        lines.append(
            f"model: {info['model_path']} ({info['model_size_mb']} MB) | "
            f"labels: {info['num_labels']} | best_exists: {info['best_variant_exists']}"
        )
        lines.append(f"  components: {info['components_path']} | encoder_classes: {info['label_encoder_classes']}")
        lines.append(f"  text_columns: {info['text_columns']}")
        if info['error']:
            lines.append(f"  ERROR: {info['error']}")
    return lines