"""N-gram ROUGE 점수 (계획서 성능지표 No.7, p.24 수식).

계획서 평가수식:
    ROUGE-N = Σ_S∈ref Σ_gram_n Count_match(gram_n) / Σ_S∈ref Σ_gram_n Count(gram_n)

즉 레퍼런스 n-gram 기준 재현율(recall) 형태. 본 모듈은 No.7 평가를 위해
ROUGE-1 / ROUGE-2 / ROUGE-L 을 자체 구현한다 (외부 라이브러리 의존 0).

토큰화는 두 가지 지원:
  - "lemma": kiwi 형태소 기본형 (한국어 어미 변화에 강건, 권장)
  - "char" : 공백/문자 기준 단순 토큰 (의존성 없이 동작)
"""

from __future__ import annotations

import re
from collections import Counter
from collections.abc import Sequence
from dataclasses import dataclass


def _char_tokens(text: str) -> list[str]:
    return re.findall(r"[가-힣A-Za-z0-9]+", text.lower())


def tokenize(text: str, mode: str = "lemma") -> list[str]:
    if mode == "lemma":
        try:
            from app.engine.structural import extract_lemmas
            toks = extract_lemmas(text)
            if toks:
                return toks
        except Exception:
            pass
    return _char_tokens(text)


def _ngrams(tokens: list[str], n: int) -> Counter:
    if len(tokens) < n:
        return Counter()
    return Counter(tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1))


@dataclass
class RougeScore:
    precision: float
    recall: float
    f1: float

    def as_dict(self) -> dict[str, float]:
        return {"precision": round(self.precision, 4), "recall": round(self.recall, 4), "f1": round(self.f1, 4)}


def _prf(match: int, sys_total: int, ref_total: int) -> RougeScore:
    precision = match / sys_total if sys_total else 0.0
    recall = match / ref_total if ref_total else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
    return RougeScore(precision, recall, f1)


def _as_references(reference: str | Sequence[str]) -> list[str]:
    """참조를 항상 리스트로 다룬다. 문자열 하나도 참조 1개짜리 리스트다."""
    if isinstance(reference, str):
        return [reference]
    return [r for r in reference if r and r.strip()]


def rouge_n(
    system: str, reference: str | Sequence[str], n: int = 1, mode: str = "lemma"
) -> RougeScore:
    """계획서 수식 그대로의 ROUGE-N. 참조는 1개도, 여러 개도 받는다.

        분자 = Σ_S∈refs Σ_gram Count_match(gram)
        분모 = Σ_S∈refs Σ_gram Count(gram)

    참조가 1개면 기존 단일 참조 계산과 완전히 같은 값이 나온다.
    """
    refs = _as_references(reference)
    if not refs:
        return RougeScore(0.0, 0.0, 0.0)
    sys_g = _ngrams(tokenize(system, mode), n)
    sys_total = sum(sys_g.values())
    match = ref_total = 0
    for ref in refs:
        ref_g = _ngrams(tokenize(ref, mode), n)
        match += sum((sys_g & ref_g).values())
        ref_total += sum(ref_g.values())
    # 분모의 recall 은 계획서 수식 그대로다. precision 은 수식에 정의가 없어
    # 참조 수만큼 시스템 n-gram 을 세는 관행(pooled)을 따르며, f1 은 그 파생값이다.
    # **보고 대상은 recall** 이다.
    return _prf(match, sys_total * len(refs), ref_total)


def _lcs_length(a: list[str], b: list[str]) -> int:
    if not a or not b:
        return 0
    prev = [0] * (len(b) + 1)
    for x in a:
        cur = [0] * (len(b) + 1)
        for j, y in enumerate(b, 1):
            cur[j] = prev[j - 1] + 1 if x == y else max(prev[j], cur[j - 1])
        prev = cur
    return prev[-1]


def rouge_l(
    system: str, reference: str | Sequence[str], mode: str = "lemma"
) -> RougeScore:
    """ROUGE-L. ROUGE-N 과 같은 방식으로 참조를 합산한다."""
    refs = _as_references(reference)
    if not refs:
        return RougeScore(0.0, 0.0, 0.0)
    s_tok = tokenize(system, mode)
    lcs = ref_total = 0
    for ref in refs:
        r_tok = tokenize(ref, mode)
        lcs += _lcs_length(s_tok, r_tok)
        ref_total += len(r_tok)
    return _prf(lcs, len(s_tok) * len(refs), ref_total)


def evaluate_pairs(
    pairs: list[tuple[str, str | Sequence[str]]],
    mode: str = "lemma",
) -> dict[str, dict[str, float]]:
    """(system, reference) 페어 리스트 → 코퍼스 평균 ROUGE-1/2/L.

    reference 자리에 문자열 하나 또는 참조 여러 개의 리스트를 넣을 수 있다.
    계획서 No.2-1년차 목표: N-gram ROUGE 65점 (gpt-4o 줄글 요약 64 대비).
    **목표와 대조하는 값은 recall** 이다(계획서 수식의 분모가 참조 n-gram 수).
    """
    if not pairs:
        return {}
    acc = {"rouge1": [], "rouge2": [], "rougeL": []}
    for system, reference in pairs:
        acc["rouge1"].append(rouge_n(system, reference, 1, mode))
        acc["rouge2"].append(rouge_n(system, reference, 2, mode))
        acc["rougeL"].append(rouge_l(system, reference, mode))

    def avg(scores: list[RougeScore]) -> dict[str, float]:
        k = len(scores)
        return {
            "precision": round(sum(s.precision for s in scores) / k, 4),
            "recall": round(sum(s.recall for s in scores) / k, 4),
            "f1": round(sum(s.f1 for s in scores) / k, 4),
        }

    return {metric: avg(scores) for metric, scores in acc.items()}