import pandas as pd
import torch
import numpy as np
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
    DataCollatorWithPadding
)
from sklearn.model_selection import train_test_split
from sklearn.metrics import (
    accuracy_score, 
    precision_recall_fscore_support,
    confusion_matrix,
    classification_report
)
from datasets import Dataset
import matplotlib.pyplot as plt
import seaborn as sns

# === 1. Загрузка данных ===
data = pd.read_csv(r"путь к файлу с данными")

# Разделяем данные
train_texts, val_texts, train_labels, val_labels = train_test_split(
    data["text"], data["real_damage"], test_size=0.2, random_state=42, stratify=data["real_damage"]
)

train_data = Dataset.from_dict({"text": train_texts.tolist(), "label": train_labels.tolist()})
val_data = Dataset.from_dict({"text": val_texts.tolist(), "label": val_labels.tolist()})

# === 2. Токенизация ===
model_name = "DeepPavlov/rubert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_name)

def tokenize_data(batch):
    return tokenizer(batch["text"], padding=False, truncation=True, max_length=512)  # Dynamic padding

train_data = train_data.map(tokenize_data, batched=True, remove_columns=["text"])
val_data = val_data.map(tokenize_data, batched=True, remove_columns=["text"])

train_data.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
val_data.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

# === 3. Модель ===
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# === 4. Метрики ===
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    
    accuracy = accuracy_score(labels, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, predictions, average='binary'
    )
    
    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1
    }

# === 5. Настройки обучения ===
training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    greater_is_better=True,
    save_total_limit=2  # Сохранять только 2 лучших чекпоинта
)

# Data collator для динамического паддинга
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=val_data,
    compute_metrics=compute_metrics,
    data_collator=data_collator,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]
)

# === 6. Обучение ===
print("\n Начало обучения...")
trainer.train()

# === 7. Оценка модели ===
eval_results = trainer.evaluate()
for key, value in eval_results.items():
    print(f"{key}: {value:.4f}")

# Детальная оценка
predictions = trainer.predict(val_data)
preds = np.argmax(predictions.predictions, axis=-1)

print("\n Classification Report:")
print(classification_report(
    val_labels, 
    preds,
    target_names=["Не реальный ущерб", "Реальный ущерб"],
    digits=4
))

# Confusion Matrix
cm = confusion_matrix(val_labels, preds)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
            xticklabels=["Не реальный ущерб", "Реальный ущерб"],
            yticklabels=["Не реальный ущерб", "Реальный ущерб"])
plt.title('Confusion Matrix')
plt.ylabel('True Label')
plt.xlabel('Predicted Label')
plt.tight_layout()
plt.savefig('confusion_matrix.png')

# === 8. Сохранение модели ===
model.save_pretrained("./binary_classification_model")
tokenizer.save_pretrained("./binary_classification_model")

class BERTClassifier:
    """Классификатор на основе собственной обученной BERT модели"""

    def __init__(self, model_path: str):
        print("Загрузка BERT модели...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.model.eval()  # Режим inference
        print("BERT модель загружена")

    def classify(self, text: str) -> Tuple[int, float, float]:
        """
        Классифицирует текст

        Возврат:
            prediction: 0 или 1
            confidence: уверенность в предсказании (0-1)
            time_taken: время выполнения в секундах
        """
        start_time = time.time()

        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)

        with torch.no_grad():
            outputs = self.model(**inputs)

        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=-1)[0]
        prediction = torch.argmax(logits, dim=-1).item()
        confidence = probabilities[prediction].item()

        time_taken = time.time() - start_time

        return prediction, confidence, time_taken

    def batch_classify(self, texts: List[str]) -> List[Dict]:
        """Классифицирует список текстов"""
        results = []
        for text in texts:
            pred, conf, time_taken = self.classify(text)
            results.append({
                "prediction": pred,
                "confidence": conf,
                "time": time_taken,
                "label": "Реальный ущерб" if pred == 1 else "Не реальный ущерб"
            })
        return results

import openai
import traceback

class OpenAIClassifier:
    """Классификатор на основе OpenAI GPT API"""

    def __init__(self, api_key: str, use_explanation: bool = True):
        self.api_key = api_key
        self.use_explanation = use_explanation
        self.prompt_template = CLASSIFICATION_PROMPT_WITH_EXPLANATION if use_explanation else CLASSIFICATION_PROMPT_WITHOUT_EXPLANATION
        self.client = openai.OpenAI(api_key=api_key)
        
    def classify(self, text: str) -> Tuple[int, float, float]:
        """Классифицирует текст через OpenAI API"""
        # Запуск таймера выполнения теста
        start_time = time.time()

        try:
            system_message = "Ты эксперт по российскому гражданскому праву." if self.use_explanation else "Ты помощник для классификации текстов."
            # Получение ответа от модели
            response = self.client.chat.completions.create(
                model="gpt-5-mini",
                messages=[
                    {"role": "system", "content": system_message},
                    {"role": "user", "content": self.prompt_template.format(text=text)}
                ],
                max_completion_tokens=500  # Установлено для стабильной работы gpt-5-mini
            )
            # Извлечение текста ответа
            message_content = response.choices[0].message.content
            # Очистка текста от лишних пробелов
            answer = message_content.strip()
            # Приведение текста к верхнему регистру
            answer_upper = answer.upper()
            
            # Парсинг ответа
            # Ищем "ДА" или "YES" в начале ответа или как отдельное слово
            if answer_upper.startswith("ДА") or answer_upper.startswith("YES") or " ДА" in answer_upper or " YES" in answer_upper:
                prediction = 1
                confidence = 0.9
            # Ищем "НЕТ" или "NO" в начале ответа или как отдельное слово
            elif answer_upper.startswith("НЕТ") or answer_upper.startswith("NO") or " НЕТ" in answer_upper or " NO" in answer_upper:
                prediction = 0
                confidence = 0.9
            # Дополнительная проверка: если в ответе есть оба слова, выбираем первое встречающееся
            elif "ДА" in answer_upper and "НЕТ" in answer_upper:
                da_pos = answer_upper.find("ДА")
                net_pos = answer_upper.find("НЕТ")
                if da_pos < net_pos:
                    prediction = 1
                    confidence = 0.7
                else:
                    prediction = 0
                    confidence = 0.7
            else:
                # Если ответ непонятен, выводим для отладки
                print(f"Неоднозначный ответ от OpenAI: '{answer}' для текста: '{text[:50]}...'")
                print(f"Finish reason: {response.choices[0].finish_reason}")
                prediction = -1
                confidence = 0.0
            # Вычисление итогового времени на получение классификации
            time_taken = time.time() - start_time
            return prediction, confidence, time_taken

        except Exception as e:
            print(f"Ошибка OpenAI API: {e}")
            traceback.print_exc()
            return -1, 0.0, 0.0

    def batch_classify(self, texts: List[str]) -> List[Dict]:
        """Классифицирует список текстов"""
        results = []
        for text in texts:
            pred, conf, time_taken = self.classify(text)
            results.append({
                "prediction": pred,
                "confidence": conf,
                "time": time_taken,
                "label": "Реальный ущерб" if pred == 1 else ("Не реальный ущерб" if pred == 0 else "Ошибка")
            })
        return results

import requests

class YandexGPTClassifier:
    """Классификатор на основе YandexGPT API"""

    def __init__(self, api_key: str, folder_id: str, use_explanation: bool = True):
        self.api_key = api_key
        self.folder_id = folder_id
        self.use_explanation = use_explanation
        self.prompt_template = CLASSIFICATION_PROMPT_WITH_EXPLANATION if use_explanation else CLASSIFICATION_PROMPT_WITHOUT_EXPLANATION      
        self.requests = requests

    def classify(self, text: str) -> Tuple[int, float, float]:
        """Классифицирует текст через YandexGPT API"""

        # Запуск таймера выполнения теста
        start_time = time.time()

        # Запрос на выполнение классификации
        try:
            url = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"

            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Api-Key {self.api_key}",
                "x-folder-id": self.folder_id
            }

            system_message = "Ты эксперт по российскому гражданскому праву." if self.use_explanation else "Ты помощник для классификации текстов."

            payload = {
                "modelUri": f"gpt://{self.folder_id}/yandexgpt/latest",
                "completionOptions": {
                    "stream": False,
                    "temperature": 0.0,
                    "maxTokens": 10
                },
                "messages": [
                    {
                        "role": "system",
                        "text": system_message
                    },
                    {
                        "role": "user",
                        "text": self.prompt_template.format(text=text)
                    }
                ]
            }
            # Объект ответа
            response = self.requests.post(url, headers=headers, json=payload, timeout=30)
            
            # Проверка HTTP на предмет наличия положительного статуса
            # Если положительного статуса нет (200-299), выводит исключение
            response.raise_for_status()

            # Преобразование json в словарь Python
            result = response.json()
            # Получение текста ответа
            answer = result["result"]["alternatives"][0]["message"]["text"].strip().upper()

            # Парсинг ответа
            if "ДА" in answer or "YES" in answer:
                prediction = 1
                confidence = 0.9
            elif "НЕТ" in answer or "NO" in answer:
                prediction = 0
                confidence = 0.9
            else:
                prediction = -1
                confidence = 0.0

            # Вычисление итогового времени на получение классификации    
            time_taken = time.time() - start_time
            return prediction, confidence, time_taken

        except Exception as e:
            print(f"Ошибка YandexGPT API: {e}")
            return -1, 0.0, 0.0

    def batch_classify(self, texts: List[str]) -> List[Dict]:
        """Классифицирует список текстов"""
        results = []
        for text in texts:
            pred, conf, time_taken = self.classify(text)
            results.append({
                "prediction": pred,
                "confidence": conf,
                "time": time_taken,
                "label": "Реальный ущерб" if pred == 1 else ("Не реальный ущерб" if pred == 0 else "Ошибка")
            })
        return results

import requests

class YandexZeroShotClassifier:
    """
    Классификатор на основе Yandex Classifier API (Zero-shot)
    Использует специализированный API для классификации без примеров
    """

    def __init__(self, api_key: str, folder_id: str):
        self.api_key = api_key
        self.folder_id = folder_id
        self.first_call = True  # Флаг для отслеживания первого вызова
        self.requests = requests

    def classify(self, text: str) -> Tuple[int, float, float]:
        """Классифицирует текст через Yandex Classifier API (zero-shot)"""

        # Добавляем задержку перед первым запросом
        if self.first_call:
            time.sleep(2.0)
            self.first_call = False
        
        # Запуск таймера выполнения теста
        start_time = time.time()
        
        # Запрос на выполнение классификации
        try:
            url = "https://llm.api.cloud.yandex.net/foundationModels/v1/fewShotTextClassification"

            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Api-Key {self.api_key}",
                "x-folder-id": self.folder_id
            }

            payload = {
                "modelUri": f"cls://{self.folder_id}/yandexgpt/latest",
                "taskDescription": """Определи, относится ли описанный случай к понятию "реальный ущерб" согласно статье 15 ГК РФ.

Реальный ущерб включает:
1. Утрату или повреждение имущества
2. Расходы на восстановление нарушенного права

НЕ является реальным ущербом:
- Упущенная выгода (неполученные доходы)
- Моральный вред""",
                "labels": [
                    "Не реальный ущерб",
                    "Реальный ущерб"
                ],
                "text": text
            }
            
            # Объект ответа
            response = self.requests.post(url, headers=headers, json=payload, timeout=30)
            
            # Проверка HTTP на предмет наличия положительного статуса
            # Если положительного статуса нет (200-299), выводит исключение
            response.raise_for_status()
            
            # Преобразование json в словарь Python
            result = response.json()

            # Получаем предсказания
            predictions = result.get("predictions", [])
            
            # Возврат в случае отсутствия предсказания
            if not predictions:
                return -1, 0.0, 0.0

            # Сортируем по confidence и берем лучший результат
            best_prediction = max(predictions, key=lambda x: x.get("confidence", 0))
            label = best_prediction.get("label", "")
            confidence = best_prediction.get("confidence", 0.0)

            # Преобразуем label в 0/1
            if "Реальный ущерб" in label:
                prediction = 1
            else:
                prediction = 0

            # Вычисление итогового времени на получение классификации      
            time_taken = time.time() - start_time
            return prediction, confidence, time_taken

        except Exception as e:
            print(f"Ошибка Yandex Classifier API: {e}")
            return -1, 0.0, 0.0

    def batch_classify(self, texts: List[str]) -> List[Dict]:
        """Классифицирует список текстов"""
        results = []
        for i, text in enumerate(texts):
            pred, conf, time_taken = self.classify(text)
            results.append({
                "prediction": pred,
                "confidence": conf,
                "time": time_taken,
                "label": "Реальный ущерб" if pred == 1 else ("Не реальный ущерб" if pred == 0 else "Ошибка")
            })
            # Добавляем задержку между запросами для избежания rate limit
            if i < len(texts) - 1:  # Не ждём после последнего запроса
                time.sleep(0.5)  # 500ms задержка
        return results

import requests


class YandexFewShotClassifier:
    """
    Классификатор на основе Yandex Classifier API (Few-shot)
    Использует примеры для обучения классификатора
    """

    def __init__(self, api_key: str, folder_id: str):
        self.api_key = api_key
        self.folder_id = folder_id
        self.first_call = True  # Флаг для отслеживания первого вызова
        self.requests = requests
        
    def classify(self, text: str) -> Tuple[int, float, float]:
        """Классифицирует текст через Yandex Classifier API (few-shot)"""

        # Добавляем задержку перед первым запросом
        if self.first_call:
            time.sleep(2.0)
            self.first_call = False
            
        # Запуск таймера выполнения теста
        start_time = time.time()
        
        # Запрос на выполнение классификации
        try:
            url = "https://llm.api.cloud.yandex.net/foundationModels/v1/fewShotTextClassification"

            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Api-Key {self.api_key}",
                "x-folder-id": self.folder_id
            }

            # Few-shot примеры для обучения
            samples = [
                # Примеры реального ущерба
                {"text": "Автомобиль был поврежден в результате ДТП", "label": "Реальный ущерб"},
                {"text": "Затопление квартиры привело к повреждению мебели и ремонта", "label": "Реальный ущерб"},
                {"text": "Расходы на оплату услуг юриста для восстановления права собственности", "label": "Реальный ущерб"},

                # Примеры НЕ реального ущерба
                {"text": "Потеря прибыли из-за срыва контракта", "label": "Не реальный ущерб"},
                {"text": "Моральные страдания от оскорблений", "label": "Не реальный ущерб"},
                {"text": "Недополученный доход от бизнеса", "label": "Не реальный ущерб"},
            ]

            payload = {
                "modelUri": f"cls://{self.folder_id}/yandexgpt/latest",
                "taskDescription": "Определи, относится ли случай к понятию 'реальный ущерб' согласно российскому законодательству",
                "labels": [
                    "Не реальный ущерб",
                    "Реальный ущерб"
                ],
                "text": text,
                "samples": samples
            }
            # Объект ответа
            response = self.requests.post(url, headers=headers, json=payload, timeout=30)
            
            # Проверка HTTP на предмет наличия положительного статуса
            # Если положительного статуса нет (200-299), выводит исключение 
            response.raise_for_status()

            # Преобразование json в словарь Python
            result = response.json()

            # Получаем предсказания
            predictions = result.get("predictions", [])

            # Возврат в случае отсутствия предсказания
            if not predictions:
                return -1, 0.0, 0.0

            # Сортируем по confidence и берем лучший результат
            best_prediction = max(predictions, key=lambda x: x.get("confidence", 0))
            label = best_prediction.get("label", "")
            confidence = best_prediction.get("confidence", 0.0)

            # Преобразуем label в 0/1
            if "Реальный ущерб" in label:
                prediction = 1
            else:
                prediction = 0

            # Вычисление итогового времени на получение классификации    
            time_taken = time.time() - start_time
            return prediction, confidence, time_taken

        except Exception as e:
            print(f"Ошибка Yandex Few-shot Classifier API: {e}")
            return -1, 0.0, 0.0

    def batch_classify(self, texts: List[str]) -> List[Dict]:
        """Классифицирует список текстов"""
        results = []
        for i, text in enumerate(texts):
            pred, conf, time_taken = self.classify(text)
            results.append({
                "prediction": pred,
                "confidence": conf,
                "time": time_taken,
                "label": "Реальный ущерб" if pred == 1 else ("Не реальный ущерб" if pred == 0 else "Ошибка")
            })
            # Добавляем задержку между запросами для избежания rate limit
            if i < len(texts) - 1:  # Не ждём после последнего запроса
                time.sleep(0.5)  # 500ms задержка
        return results

# Создаем папку для результатов
RESULTS_DIR = Path("relation_classification_comparison_results")
RESULTS_DIR.mkdir(exist_ok=True)

CONFIG = {
    "bert_model_path": r"путь к модели BERT",
    "openai_api_key": os.getenv("OPENAI_API_KEY", ""),
    "yandex_api_key": os.getenv("YANDEX_API_KEY", ""),
    "yandex_folder_id": os.getenv("YANDEX_FOLDER_ID", ""),
}

def calculate_metrics(predictions: List[int], labels: List[int]) -> Dict[str, float]:
    """Вычисляет метрики качества классификации"""

    valid_pairs = [(p, l) for p, l in zip(predictions, labels) if p != -1]

    if not valid_pairs:
        return {
            "accuracy": 0.0,
            "precision": 0.0,
            "recall": 0.0,
            "f1": 0.0,
            "valid_predictions": 0
        }

    predictions_valid = [p for p, l in valid_pairs]
    labels_valid = [l for p, l in valid_pairs]

    tp = sum(1 for p, l in zip(predictions_valid, labels_valid) if p == 1 and l == 1)
    fp = sum(1 for p, l in zip(predictions_valid, labels_valid) if p == 1 and l == 0)
    tn = sum(1 for p, l in zip(predictions_valid, labels_valid) if p == 0 and l == 0)
    fn = sum(1 for p, l in zip(predictions_valid, labels_valid) if p == 0 and l == 1)

    accuracy = (tp + tn) / len(predictions_valid) if predictions_valid else 0
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0

    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "valid_predictions": len(predictions_valid),
        "tp": tp,
        "fp": fp,
        "tn": tn,
        "fn": fn
    }


def compare_methods(test_cases: List[Dict]) -> Tuple[pd.DataFrame, Dict]:
    """Сравнивает все методы классификации"""

    classifiers = {
        "BERT (собственная модель)": BERTClassifier(CONFIG["bert_model_path"]),
        "OpenAI GPT (с объяснением)": OpenAIClassifier(CONFIG["openai_api_key"], use_explanation=True),
        "OpenAI GPT (без объяснения)": OpenAIClassifier(CONFIG["openai_api_key"], use_explanation=False),
        "YandexGPT (с объяснением)": YandexGPTClassifier(CONFIG["yandex_api_key"], CONFIG["yandex_folder_id"], use_explanation=True),
        "YandexGPT (без объяснения)": YandexGPTClassifier(CONFIG["yandex_api_key"], CONFIG["yandex_folder_id"], use_explanation=False),
        "Yandex Classifier (zero-shot)": YandexZeroShotClassifier(CONFIG["yandex_api_key"], CONFIG["yandex_folder_id"]),
        "Yandex Classifier (few-shot)": YandexFewShotClassifier(CONFIG["yandex_api_key"], CONFIG["yandex_folder_id"]),
    }

    texts = [case["text"] for case in test_cases]
    true_labels = [case["label"] for case in test_cases]
    categories = [case["category"] for case in test_cases]

    all_results = {}

    for method_name, classifier in classifiers.items():

        results = classifier.batch_classify(texts)
        all_results[method_name] = results

        predictions = [r["prediction"] for r in results]
        metrics = calculate_metrics(predictions, true_labels)

        print(f"✓ Обработано примеров: {len(texts)}")
        print(f"✓ Валидных предсказаний: {metrics['valid_predictions']}")
        print(f"✓ Accuracy: {metrics['accuracy']:.2%}")
        print(f"✓ Precision: {metrics['precision']:.2%}")
        print(f"✓ Recall: {metrics['recall']:.2%}")
        print(f"✓ F1-score: {metrics['f1']:.2%}")

        avg_time = sum(r["time"] for r in results) / len(results) if results else 0
        print(f"✓ Среднее время на пример: {avg_time:.3f} сек")

    comparison_data = []

    for i, case in enumerate(test_cases):
        row = {
            "Текст": case["text"][:60] + "..." if len(case["text"]) > 60 else case["text"],
            "Категория": case["category"],
            "Истина": "Да" if case["label"] == 1 else "Нет"
        }

        for method_name, results in all_results.items():
            pred = results[i]["prediction"]
            conf = results[i]["confidence"]

            if pred == -1:
                row[f"{method_name}"] = "Ошибка"
                row[f"{method_name} (уверенность)"] = "N/A"
            else:
                row[f"{method_name}"] = "Да" if pred == 1 else "Нет"
                row[f"{method_name} (уверенность)"] = f"{conf:.2%}"

        comparison_data.append(row)

    df = pd.DataFrame(comparison_data)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = RESULTS_DIR / f"classification_comparison_{timestamp}.csv"
    df.to_csv(output_file, index=False, encoding="utf-8-sig")
    print(f"\n Результаты сохранены в файл: {output_file}")

    return df, all_results


def visualize_results(all_results: Dict, test_cases: List[Dict]):
    """Создает визуализации результатов сравнения"""

    true_labels = [case["label"] for case in test_cases]

    methods = []
    accuracies = []
    f1_scores = []
    avg_times = []
    colors_list = []

    for method_name, results in all_results.items():
        predictions = [r["prediction"] for r in results]
        metrics = calculate_metrics(predictions, true_labels)

        if metrics["valid_predictions"] > 0:
            short_name = method_name.replace(" (собственная модель)", "")

            if "BERT" in method_name:
                color = '#2ecc71'
            elif "без объяснения" in method_name:
                color = '#e74c3c'
            elif "с объяснением" in method_name:
                color = '#3498db'
            else:
                color = '#95a5a6'

            methods.append(short_name)
            accuracies.append(metrics["accuracy"] * 100)
            f1_scores.append(metrics["f1"] * 100)
            avg_times.append(sum(r["time"] for r in results) / len(results))
            colors_list.append(color)

    # График 1: Accuracy и F1-score
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

    x = range(len(methods))

    bars1 = ax1.bar(x, accuracies, color=colors_list, alpha=0.8, edgecolor='black', linewidth=1.2)
    ax1.set_ylabel('Accuracy (%)', fontsize=12)
    ax1.set_title('Сравнение точности методов классификации отношений', fontsize=14, fontweight='bold')
    ax1.set_xticks(x)
    ax1.set_xticklabels(methods, rotation=25, ha='right', fontsize=9)
    ax1.grid(axis='y', alpha=0.3)
    ax1.set_ylim(0, 105)

    for i, (bar, acc) in enumerate(zip(bars1, accuracies)):
        height = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width()/2., height + 1,
                f'{acc:.1f}%', ha='center', va='bottom', fontsize=9, fontweight='bold')

    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor='#2ecc71', label='BERT (обученная модель)'),
        Patch(facecolor='#3498db', label='LLM с объяснением'),
        Patch(facecolor='#e74c3c', label='LLM без объяснения')
    ]
    ax1.legend(handles=legend_elements, loc='lower right', fontsize=9)

    bars2 = ax2.bar(x, f1_scores, color=colors_list, alpha=0.8, edgecolor='black', linewidth=1.2)
    ax2.set_ylabel('F1-score (%)', fontsize=12)
    ax2.set_title('F1-score для каждого метода', fontsize=14, fontweight='bold')
    ax2.set_xticks(x)
    ax2.set_xticklabels(methods, rotation=25, ha='right', fontsize=9)
    ax2.grid(axis='y', alpha=0.3)
    ax2.set_ylim(0, 105)

    for i, (bar, f1) in enumerate(zip(bars2, f1_scores)):
        height = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width()/2., height + 1,
                f'{f1:.1f}%', ha='center', va='bottom', fontsize=9, fontweight='bold')

    plt.tight_layout()
    output_file = RESULTS_DIR / 'classification_methods_comparison.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"График сохранен: {output_file}")

    # График 2: Время выполнения
    fig, ax3 = plt.subplots(figsize=(12, 6))
    time_colors = ['#2ecc71' if t < 0.1 else '#f39c12' if t < 1.0 else '#e74c3c' for t in avg_times]
    bars3 = ax3.bar(x, avg_times, color=time_colors, alpha=0.8, edgecolor='black', linewidth=1.2)
    ax3.set_ylabel('Время (секунды)', fontsize=12)
    ax3.set_title('Среднее время классификации одного примера', fontsize=14, fontweight='bold')
    ax3.set_xticks(x)
    ax3.set_xticklabels(methods, rotation=25, ha='right', fontsize=9)
    ax3.grid(axis='y', alpha=0.3)

    for i, (bar, t) in enumerate(zip(bars3, avg_times)):
        height = bar.get_height()
        ax3.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                f'{t:.3f}s', ha='center', va='bottom', fontsize=9, fontweight='bold')

    time_legend = [
        Patch(facecolor='#2ecc71', label='Быстро (<0.1с)'),
        Patch(facecolor='#f39c12', label='Средне (0.1-1с)'),
        Patch(facecolor='#e74c3c', label='Медленно (>1с)')
    ]
    ax3.legend(handles=time_legend, loc='upper right', fontsize=9)

    plt.tight_layout()
    output_file = RESULTS_DIR / 'classification_time_comparison.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"График времени сохранен: {output_file}")

    # График 3: Confusion Matrix
    # Создаем mapping для безопасного доступа к axes
    method_to_idx = {method_name: idx for idx, method_name in enumerate(methods)}

    fig, axes = plt.subplots(1, len(methods), figsize=(5*len(methods), 4))

    if len(methods) == 1:
        axes = [axes]

    for method_name, results in all_results.items():
        predictions = [r["prediction"] for r in results if r["prediction"] != -1]
        labels_valid = [true_labels[i] for i, r in enumerate(results) if r["prediction"] != -1]

        method_short = method_name.replace(" (собственная модель)", "")

        # Проверяем, что метод есть в списке methods (т.е. имеет валидные предсказания)
        if method_short in method_to_idx and predictions:
            idx = method_to_idx[method_short]
            from sklearn.metrics import confusion_matrix
            cm = confusion_matrix(labels_valid, predictions)

            sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=axes[idx],
                       xticklabels=["Нет", "Да"],
                       yticklabels=["Нет", "Да"])
            axes[idx].set_title(method_short, fontweight='bold', fontsize=10)
            axes[idx].set_ylabel('Истинная метка')
            axes[idx].set_xlabel('Предсказание')

    plt.tight_layout()
    output_file = RESULTS_DIR / 'confusion_matrices_comparison.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"Confusion matrices сохранены: {output_file}")

    plt.close('all')