"""Основные функции вычисления метрик качества текста."""

import logging
from typing import Any, Dict, Set, Tuple

import evaluate
import nltk
import pymorphy2

from src.utils import (
    extract_entities,
    extract_json_values_from_markdown,
    lemmatize_text_russian,
    prepare_text_for_analysis,
)

# Настройка логирования
logger = logging.getLogger(__name__)


def initialize_resources() -> Tuple[Any, Any, Any, Any]:
    """Инициализирует и возвращает необходимые ресурсы для анализа текста."""
    nltk.download('punkt')
    rouge_metric = evaluate.load('rouge', quiet=True)
    bert_metric = evaluate.load('bertscore', quiet=True)
    meteor_metric = evaluate.load('meteor', quiet=True)
    morphology = pymorphy2.MorphAnalyzer()
    return rouge_metric, bert_metric, meteor_metric, morphology


try:
    rouge, bertscore, meteor, morph = initialize_resources()
    logger.info('Все модели загружены')
except Exception as error:
    logger.error(f'Ошибка инициализации: {str(error)}')


def calculate_ner_metrics(ref_entities: Set[str], hyp_entities: Set[str]) -> Dict[str, float]:
    """Вычисляет метрики NER на основе множеств сущностей."""
    if not hyp_entities:
        return {'f1': 0.0, 'precision': 0.0, 'recall': 0.0, 'iou': 0.0}

    # True Positives: слова, которые есть и в эталоне, и в гипотезе
    true_positives = len(ref_entities.intersection(hyp_entities))

    # False Positives: слова в гипотезе, но отсутствующие в эталоне
    false_positives = len(hyp_entities - ref_entities)

    # False Negatives: слова в эталоне, но отсутствующие в гипотезе
    false_negatives = len(ref_entities - hyp_entities)

    # Precision и Recall
    precision = true_positives / (true_positives + false_positives) if (true_positives + false_positives) > 0 else 0.0
    recall = true_positives / (true_positives + false_negatives) if (true_positives + false_negatives) > 0 else 0.0

    # F1-мера
    if precision + recall == 0:
        return {'f1': 0.0, 'precision': precision, 'recall': recall, 'iou': 0.0}

    f1 = 2 * (precision * recall) / (precision + recall)

    # IoU (Intersection over Union)
    iou = (
        len(ref_entities & hyp_entities) / len(ref_entities | hyp_entities)
        if len(ref_entities | hyp_entities) > 0
        else 0.0
    )

    return {'f1': f1, 'precision': precision, 'recall': recall, 'iou': iou}


def calculate_word_f1(reference: str, hypothesis: str) -> Dict[str, float]:
    """Вычисляет метрики F1 для сущностей в текстах."""
    # Предобработка текста
    reference = prepare_text_for_analysis(reference)
    hypothesis = prepare_text_for_analysis(hypothesis)

    # Извлечение значений из JSON в markdown
    reference = extract_json_values_from_markdown(reference)
    hypothesis = extract_json_values_from_markdown(hypothesis)

    # Стоп-слова
    stopwords: Set[str] = set([])

    # Извлечение сущностей
    ref_entities = extract_entities(reference, stopwords)
    hyp_entities = extract_entities(hypothesis, stopwords)

    # Расчет метрик
    return calculate_ner_metrics(ref_entities, hyp_entities)


def compute_rouge_metrics(ref_text: str, hyp_text: str) -> Dict[str, float]:
    """Вычисляет метрики ROUGE."""
    rouge_scores = rouge.compute(predictions=[hyp_text], references=[ref_text])
    return {'rouge_l': rouge_scores['rougeL'], 'rouge_2': rouge_scores['rouge2']}


def compute_bert_score(ref_text: str, hyp_text: str) -> float:
    """Вычисляет метрику BERTScore."""
    bert_scores = bertscore.compute(predictions=[hyp_text], references=[ref_text], lang='ru')
    return bert_scores['f1'][0]


def compute_meteor_score(ref_text: str, hyp_text: str) -> float:
    """Вычисляет метрику METEOR на лемматизированных текстах."""
    hyp_lemmatized = lemmatize_text_russian(hyp_text, morph)
    ref_lemmatized = lemmatize_text_russian(ref_text, morph)

    meteor_score = meteor.compute(predictions=[hyp_lemmatized], references=[ref_lemmatized])
    return meteor_score['meteor']


def compute_metrics(reference: str, hypothesis: str) -> Dict[str, Any]:
    """Вычисляет набор метрик для оценки качества генерации."""
    metrics_result = {}
    # ROUGE
    metrics_result['rouge'] = compute_rouge_metrics(reference, hypothesis)  # dict

    # BERTScore
    metrics_result['bert_score'] = compute_bert_score(reference, hypothesis)  # float

    # METEOR
    metrics_result['meteor'] = compute_meteor_score(reference, hypothesis)  # float

    # Word-level F1 score
    metrics_result['ner'] = calculate_word_f1(reference, hypothesis)  # dict

    logger.info('Метрики вычислены')
    return metrics_result
