"""Prediction and evaluation views for any CEFR checkpoint (ours or a baseline).

Two views are reported, matching ADR 0003:
- "document": doc/paragraph/dialogue-format items, chunk probabilities
  aggregated per doc_id by expected-rank mean (the headline product metric);
- "sentence": sentence-format items scored individually.

`predict_probs` reorders model outputs to the canonical level order using the
checkpoint's id2label, so a published baseline with a different label order is
evaluated correctly — or fails loudly if its labels are not CEFR levels.
"""

from collections import defaultdict
from collections.abc import Sequence

from tutor.ml.cefr.aggregation import aggregate_chunk_probs
from tutor.ml.cefr.metrics import classification_report
from tutor.ml.cefr.preprocessing import CANONICAL_LEVELS, Passage

LEVEL_TO_RANK = {level: rank for rank, level in enumerate(CANONICAL_LEVELS)}
DOC_FORMATS = {"document-level", "paragraph-level", "dialogue-level"}


def canonical_column_order(id2label: dict[int, str]) -> list[int]:
    """For each canonical level, the model output column that carries it."""
    label_to_column = {str(label).strip().upper(): int(col) for col, label in id2label.items()}
    missing = [level for level in CANONICAL_LEVELS if level not in label_to_column]
    if missing:
        msg = (
            f"model labels {sorted(label_to_column)} do not cover the canonical "
            f"CEFR levels (missing {missing}); cannot evaluate this checkpoint"
        )
        raise ValueError(msg)
    return [label_to_column[level] for level in CANONICAL_LEVELS]


def predict_probs(
    model,
    tokenizer,
    passages: Sequence[Passage],
    *,
    max_length: int,
    batch_size: int = 64,
) -> list[list[float]]:
    """One probability row per passage, columns in canonical level order."""
    import torch  # lazy: keeps this module importable without the train group

    order = canonical_column_order(model.config.id2label)
    device = next(model.parameters()).device
    model.eval()
    probs: list[list[float]] = []
    with torch.no_grad():
        for start in range(0, len(passages), batch_size):
            texts = [p.text for p in passages[start : start + batch_size]]
            encoded = tokenizer(
                texts,
                truncation=True,
                max_length=max_length,
                padding=True,
                return_tensors="pt",
            ).to(device)
            raw = torch.softmax(model(**encoded).logits, dim=-1).cpu().tolist()
            probs.extend([[row[column] for column in order] for row in raw])
    return probs


def evaluate_views(
    passages: Sequence[Passage],
    probs: Sequence[Sequence[float]],
) -> dict[str, dict[str, float]]:
    """Compute the metric bundle for the document and sentence views."""
    doc_rows: dict[str, list[Sequence[float]]] = defaultdict(list)
    doc_levels: dict[str, str] = {}
    sentence_true: list[int] = []
    sentence_pred: list[int] = []

    for passage, row in zip(passages, probs, strict=True):
        if passage.source_format in DOC_FORMATS:
            doc_rows[passage.doc_id].append(row)
            doc_levels[passage.doc_id] = passage.level
        else:
            predicted_level, _ = aggregate_chunk_probs([row])
            sentence_true.append(LEVEL_TO_RANK[passage.level])
            sentence_pred.append(LEVEL_TO_RANK[predicted_level])

    document_true: list[int] = []
    document_pred: list[int] = []
    for doc_id, rows in doc_rows.items():
        predicted_level, _ = aggregate_chunk_probs(rows)
        document_true.append(LEVEL_TO_RANK[doc_levels[doc_id]])
        document_pred.append(LEVEL_TO_RANK[predicted_level])

    views: dict[str, dict[str, float]] = {}
    if document_true:
        views["document"] = classification_report(document_true, document_pred)
    if sentence_true:
        views["sentence"] = classification_report(sentence_true, sentence_pred)
    return views


def lang_filtered(
    passages: Sequence[Passage],
    probs: Sequence[Sequence[float]],
    lang: str,
) -> tuple[list[Passage], list[list[float]]]:
    """Subset (passages, probs) to one language, keeping rows aligned."""
    pairs = [(p, list(row)) for p, row in zip(passages, probs, strict=True) if p.lang == lang]
    return [p for p, _ in pairs], [row for _, row in pairs]