"""
Модуль для вычисления различных метрик оценки качества ответов моделей.
"""
import re
import math
from typing import List, Tuple, Dict, Any
from collections import Counter

def calculate_f1_score(model_response: str, expected: str) -> float:
    """
    Вычисляет F1-score на основе совпадения токенов.

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ

    Returns:
        F1-score (0.0-1.0)
    """
    model_tokens = set(model_response.lower().split())
    expected_tokens = set(expected.lower().split())

    if len(expected_tokens) == 0:
        return 0.0

    intersection = model_tokens.intersection(expected_tokens)
    precision = len(intersection) / len(model_tokens) if model_tokens else 0.0
    recall = len(intersection) / len(expected_tokens) if expected_tokens else 0.0

    if (precision + recall) == 0:
        return 0.0
    f1 = 2 * (precision * recall) / (precision + recall)

    return round(f1, 3)

def calculate_exact_match(model_response: str, expected: str) -> float:
    """
    Вычисляет Exact Match Ratio (EM) - процент тестов с точным совпадением.

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ

    Returns:
        1.0 если ответ точно совпадает, иначе 0.0
    """
    # Удаление лишних пробелов и табуляций
    model_clean = ' '.join(model_response.strip().split())
    expected_clean = ' '.join(expected.strip().split())

    return 1.0 if model_clean == expected_clean else 0.0

def calculate_levenshtein_distance(model_response: str, expected: str) -> int:
    """
    Вычисляет Levenshtein Distance (расстояние редактирования) между двумя строками.

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ

    Returns:
        Количество редактирований (вставок, удалений, замен)
    """
    if len(expected) == 0:
        return len(model_response)
    if len(model_response) == 0:
        return len(expected)

    # Матрица для хранения расстояний
    d = [[0] * (len(expected) + 1) for _ in range(len(model_response) + 1)]

    # Инициализация
    for i in range(len(model_response) + 1):
        d[i][0] = i
    for j in range(len(expected) + 1):
        d[0][j] = j

    # Заполнение матрицы
    for j in range(1, len(expected) + 1):
        for i in range(1, len(model_response) + 1):
            if model_response[i-1] == expected[j-1]:
                substitution_cost = 0
            else:
                substitution_cost = 1
            d[i][j] = min(
                d[i-1][j] + 1,      # удаление
                d[i][j-1] + 1,      # вставка
                d[i-1][j-1] + substitution_cost  # замена
            )

    return d[len(model_response)][len(expected)]

def calculate_normalized_levenshtein(model_response: str, expected: str) -> float:
    """
    Вычисляет нормализованное Levenshtein Distance (0.0-1.0).

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ

    Returns:
        Нормализованное расстояние (0.0 = идентично, 1.0 = полностью разные)
    """
    max_len = max(len(model_response), len(expected))
    if max_len == 0:
        return 0.0
    distance = calculate_levenshtein_distance(model_response, expected)
    return round(1.0 - (distance / max_len), 3)

def calculate_bleu_score(model_response: str, expected: str, ngram_weights: List[float] = None) -> float:
    """
    Вычисляет BLEU Score на основе n-грамм.

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ
        ngram_weights: Веса для разных n-грамм. По умолчанию [0.25, 0.25, 0.25, 0.25] для 1-4 грамм

    Returns:
        BLEU Score (0.0-1.0)
    """
    if ngram_weights is None:
        ngram_weights = [0.25, 0.25, 0.25, 0.25]

    # Токенизация
    model_tokens = model_response.lower().split()
    expected_tokens = expected.lower().split()

    if not model_tokens or not expected_tokens:
        return 0.0

    # Вычисление precision для разных n-грамм
    precisions = []
    for n in range(1, 5):
        if n > len(model_tokens):
            precisions.append(0)
            continue

        # Счетчики n-грамм в ответе модели
        model_ngrams = Counter(
            tuple(model_tokens[i:i+n])
            for i in range(len(model_tokens) - n + 1)
        )

        # Счетчики n-грамм в ожидаемом ответе
        expected_ngrams = Counter(
            tuple(expected_tokens[i:i+n])
            for i in range(len(expected_tokens) - n + 1)
        )

        if not model_ngrams:
            precisions.append(0)
            continue

        # Вычисление precision с smoothing
        clipped_count = 0
        for ngram, count in model_ngrams.items():
            clipped_count += min(count, expected_ngrams.get(ngram, 0))

        precision = clipped_count / len(model_ngrams)
        precisions.append(precision)

    # Взвешенное среднее
    weighted_sum = sum(w * p for w, p in zip(ngram_weights[:len(precisions)], precisions))

    # Бонус за brevity (длина ответа близка к ожидаемой)
    if len(model_tokens) > len(expected_tokens):
        bp = 1.0
    else:
        bp = math.exp(1 - len(expected_tokens) / len(model_tokens))

    bleu_score = bp * math.exp(weighted_sum)
    return round(bleu_score, 3)

def calculate_rouge_scores(model_response: str, expected: str) -> Dict[str, float]:
    """
    Вычисляет ROUGE Scores (ROUGE-1, ROUGE-2, ROUGE-L).

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ

    Returns:
        Словарь с ROUGE метриками: {'rouge1': ..., 'rouge2': ..., 'rougeL': ...}
    """
    model_tokens = model_response.lower().split()
    expected_tokens = expected.lower().split()

    if not model_tokens or not expected_tokens:
        return {'rouge1': 0.0, 'rouge2': 0.0, 'rougeL': 0.0}

    # ROUGE-1: перекрытие unigrams
    model_grams = Counter(model_tokens)
    expected_grams = Counter(expected_tokens)

    intersection = sum((model_grams & expected_grams).values())
    rouge1 = intersection / len(expected_grams) if expected_grams else 0.0

    # ROUGE-2: перекрытие bigrams
    model_bigrams = Counter(
        tuple(model_tokens[i:i+2])
        for i in range(len(model_tokens) - 1)
    )
    expected_bigrams = Counter(
        tuple(expected_tokens[i:i+2])
        for i in range(len(expected_tokens) - 1)
    )

    intersection = sum((model_bigrams & expected_bigrams).values())
    rouge2 = intersection / len(expected_bigrams) if expected_bigrams else 0.0

    # ROUGE-L: Longest Common Subsequence (LCS)
    rougeL = calculate_rouge_lcs(model_tokens, expected_tokens)

    return {
        'rouge1': round(rouge1, 3),
        'rouge2': round(rouge2, 3),
        'rougeL': round(rougeL, 3)
    }

def calculate_rouge_lcs(model_tokens: List[str], expected_tokens: List[str]) -> float:
    """
    Вычисляет ROUGE-L на основе Longest Common Subsequence (LCS).

    Args:
        model_tokens: Токены ответа модели
        expected_tokens: Токены ожидаемого ответа

    Returns:
        ROUGE-L score (0.0-1.0)
    """
    # Матрица для хранения длины LCS
    m = len(model_tokens)
    n = len(expected_tokens)
    dp = [[0] * (n + 1) for _ in range(m + 1)]

    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if model_tokens[i-1] == expected_tokens[j-1]:
                dp[i][j] = dp[i-1][j-1] + 1
            else:
                dp[i][j] = max(dp[i-1][j], dp[i][j-1])

    lcs_length = dp[m][n]

    # Вычисление ROUGE-L
    precision = lcs_length / m if m > 0 else 0.0
    recall = lcs_length / n if n > 0 else 0.0

    if (precision + recall) == 0:
        return 0.0
    f_score = 2 * (precision * recall) / (precision + recall)

    return f_score

def calculate_code_similarity(model_response: str, expected: str) -> float:
    """
    Вычисляет похожесть кода на основе структурных элементов.

    Args:
        model_response: Сгенерированный код
        expected: Ожидаемый код

    Returns:
        Коэффициент похожести (0.0-1.0)
    """
    # Удаление комментариев
    model_clean = remove_comments(model_response)
    expected_clean = remove_comments(expected)

    # Нормализация (удаление лишних пробелов, табуляций)
    model_normalized = normalize_code(model_clean)
    expected_normalized = normalize_code(expected_clean)

    # Сравнение токенов кода
    model_tokens = tokenize_code(model_normalized)
    expected_tokens = tokenize_code(expected_normalized)

    if not expected_tokens:
        return 0.0

    # Простая метрика на основе совпадения ключевых токенов
    intersection = set(model_tokens) & set(expected_tokens)
    precision = len(intersection) / len(model_tokens) if model_tokens else 0.0
    recall = len(intersection) / len(expected_tokens) if expected_tokens else 0.0

    if (precision + recall) == 0:
        return 0.0
    f1 = 2 * (precision * recall) / (precision + recall)

    return round(f1, 3)

def remove_comments(code: str) -> str:
    """
    Удаляет комментарии из кода.

    Args:
        code: Исходный код

    Returns:
        Код без комментариев
    """
    # Удаление однострочных комментариев
    code = re.sub(r'//.*', '', code)
    code = re.sub(r'#.*', '', code)

    # Удаление многострочных комментариев
    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)

    return code

def normalize_code(code: str) -> str:
    """
    Нормализует код (удаляет лишние пробелы, табуляции).

    Args:
        code: Исходный код

    Returns:
        Нормализованный код
    """
    # Замена нескольких пробелов/табуляций на один
    code = re.sub(r'\s+', ' ', code)
    # Удаление пробелов в начале и конце строк
    code = '\n'.join(line.strip() for line in code.split('\n'))
    return code

def tokenize_code(code: str) -> List[str]:
    """
    Токенизирует код, выделяя ключевые элементы.

    Args:
        code: Исходный код

    Returns:
        Список токенов
    """
    # Регулярное выражение для токенизации кода
    token_pattern = r"""
        \w+|          # Идентификаторы и ключевые слова
        [+\-*/%=<>!&|^~.,;(){}[\]]|  # Операторы и знаки препинания
        [0-9]+|       # Числа
        [A-Za-z_][A-Za-z0-9_]*|  # Идентификаторы
        \S            # Любые другие непробельные символы
    """
    tokens = re.findall(token_pattern, code, re.VERBOSE)
    return [token for token in tokens if token.strip()]

def get_all_scores(model_response: str, expected: str) -> Dict[str, Any]:
    """
    Вычисляет все доступные метрики для ответа модели.

    Args:
        model_response: Ответ от модели
        expected: Ожидаемый ответ

    Returns:
        Словарь со всеми метриками
    """
    return {
        'f1_score': calculate_f1_score(model_response, expected),
        'exact_match': calculate_exact_match(model_response, expected),
        'levenshtein_distance': calculate_levenshtein_distance(model_response, expected),
        'normalized_levenshtein': calculate_normalized_levenshtein(model_response, expected),
        'bleu_score': calculate_bleu_score(model_response, expected),
        'rouge_scores': calculate_rouge_scores(model_response, expected),
        'code_similarity': calculate_code_similarity(model_response, expected)
    }
