chatkit/backend/app/answer_grid.py

"""Canonical answer-grid data model.

Single source of truth for per-student report generation:
- total_questions (N)
- official_answers: 1xN vector of letters
- students: m rows of N letters (or None for blank)
- questions: N question metadata entries for rendering wrong blocks

Single-choice only. Letters A-Z. All comparisons are case-insensitive trim.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Optional

_BLANK_MARKERS = {"", "n/a", "none", "null"}
_CORRECT_MARKER = "="
# Legacy dash markers are accepted as input and converted to CORRECT_MARKER
_DASH_MARKERS = {"-", "–", "—"}
_LETTER_RE = re.compile(r"^[A-Z]$")


@dataclass(frozen=True)
class QuestionMeta:
    number: int
    text: str
    options: tuple[str, ...]


@dataclass(frozen=True)
class StudentRow:
    name: str
    answers: tuple[Optional[str], ...]


@dataclass(frozen=True)
class AnswerGrid:
    total_questions: int
    official_answers: tuple[Optional[str], ...]
    students: tuple[StudentRow, ...]
    questions: tuple[QuestionMeta, ...]


@dataclass(frozen=True)
class WrongAnswer:
    question: QuestionMeta
    student_answer: Optional[str]
    correct_answer: str


def normalize_letter(raw: Optional[str]) -> Optional[str]:
    """Normalize a raw cell to a letter (A-Z), the correct-marker "=", or None.

    - "=" or any legacy dash marker ("-", "–", "—") → "=" (answered correctly)
    - Single letters with optional trailing "A) foo" → uppercase letter
    - Empty/blank/unparseable → None
    """
    if raw is None:
        return None
    s = str(raw).strip()
    if not s or s.casefold() in _BLANK_MARKERS:
        return None
    if s == _CORRECT_MARKER or s in _DASH_MARKERS:
        return _CORRECT_MARKER
    s = s.upper()
    # Tolerate "A)" or "A) foo" -> "A"
    if len(s) > 1 and s[0].isalpha() and s[1] in (")", ".", "、", " "):
        s = s[0]
    return s if _LETTER_RE.match(s) else None


def _normalize_row(
    answers: list, n: int
) -> tuple[Optional[str], ...]:
    """Normalize an answer row: pad/truncate to length n."""
    out: list[Optional[str]] = []
    for i in range(n):
        raw = answers[i] if i < len(answers) else None
        out.append(normalize_letter(raw))
    return tuple(out)


def from_dict(payload: dict) -> AnswerGrid:
    """Parse + validate a client-submitted answer grid."""
    if not isinstance(payload, dict):
        raise ValueError("payload must be a dict")

    n_raw = payload.get("total_questions")
    if not isinstance(n_raw, int) or n_raw < 1:
        raise ValueError("total_questions must be a positive integer")
    n = n_raw

    official = payload.get("official_answers") or []
    if not isinstance(official, list):
        raise ValueError("official_answers must be a list")
    official_tuple = _normalize_row(official, n)

    students_raw = payload.get("students") or []
    if not isinstance(students_raw, list):
        raise ValueError("students must be a list")
    students: list[StudentRow] = []
    for i, s in enumerate(students_raw):
        if not isinstance(s, dict):
            raise ValueError(f"students[{i}] must be a dict")
        name = str(s.get("name") or f"Student {i + 1}")
        answers = s.get("answers") or []
        if not isinstance(answers, list):
            raise ValueError(f"students[{i}].answers must be a list")
        students.append(StudentRow(name=name, answers=_normalize_row(answers, n)))

    questions_raw = payload.get("questions") or []
    if not isinstance(questions_raw, list):
        raise ValueError("questions must be a list")
    # Keyed by "number", not array position — see seed_from_parsed for why.
    q_by_num: dict[int, dict] = {}
    for idx, q in enumerate(questions_raw):
        if not isinstance(q, dict):
            continue
        num = q.get("number") if isinstance(q.get("number"), int) else idx + 1
        q_by_num[num] = q

    questions: list[QuestionMeta] = []
    for i in range(n):
        number = i + 1
        q = q_by_num.get(number, {})
        text = str(q.get("text") or "")
        opts_raw = q.get("options") or []
        opts = tuple(str(o) for o in opts_raw) if isinstance(opts_raw, list) else ()
        questions.append(QuestionMeta(number=number, text=text, options=opts))

    return AnswerGrid(
        total_questions=n,
        official_answers=official_tuple,
        students=tuple(students),
        questions=tuple(questions),
    )


def to_dict(grid: AnswerGrid) -> dict:
    """Serialize an AnswerGrid to JSON-friendly dict."""
    return {
        "total_questions": grid.total_questions,
        "official_answers": list(grid.official_answers),
        "students": [
            {"name": s.name, "answers": list(s.answers)} for s in grid.students
        ],
        "questions": [
            {"number": q.number, "text": q.text, "options": list(q.options)}
            for q in grid.questions
        ],
    }


def seed_from_parsed(
    questions_data: dict,
    student_answers_data: dict,
    teacher_answers_data: dict,
) -> AnswerGrid:
    """Synthesize an initial grid from existing parsed JSON.

    - N is derived from questions_data if present, else from teacher_answers length,
      else from the longest student row.
    - Student rows are padded/truncated to N.
    - Cells are normalized to letters, the correct-marker "=", or None.
    """
    q_list = (questions_data or {}).get("questions") or []
    t_list = (teacher_answers_data or {}).get("answers") or []
    s_list = (student_answers_data or {}).get("students") or []

    if q_list:
        # Use the highest explicit "number" field, not raw array length — if the
        # extraction skipped a question anywhere (gap), len(q_list) undercounts
        # and every question past the gap gets silently truncated from the grid.
        explicit_nums = [q.get("number") for q in q_list if isinstance(q, dict) and isinstance(q.get("number"), int)]
        n = max(explicit_nums) if explicit_nums else len(q_list)
        n = max(n, len(q_list))
    elif t_list:
        n = max((a.get("question_number", 0) for a in t_list), default=0)
    else:
        longest = 0
        for s in s_list:
            ans = s.get("answers") or []
            nums = [a.get("question_number", 0) for a in ans]
            if nums:
                longest = max(longest, max(nums))
        n = longest

    n = max(n, 1)

    # Build official answers by question_number index
    official: list[Optional[str]] = [None] * n
    for a in t_list:
        qn = a.get("question_number")
        if isinstance(qn, int) and 1 <= qn <= n:
            official[qn - 1] = normalize_letter(a.get("correct_answer"))

    # Build student rows by question_number index. Legacy "-" markers are
    # mapped to the canonical "=" correct-marker by normalize_letter.
    students: list[StudentRow] = []
    for i, s in enumerate(s_list):
        row: list[Optional[str]] = [None] * n
        for a in s.get("answers") or []:
            qn = a.get("question_number")
            if isinstance(qn, int) and 1 <= qn <= n:
                row[qn - 1] = normalize_letter(a.get("answer"))
        name = str(s.get("name") or f"Student {i + 1}")
        students.append(StudentRow(name=name, answers=tuple(row)))

    # Build question metadata, keyed by each question's own "number" field
    # (not array position) — the raw extraction can skip/reorder a question
    # in a long document, and indexing by position would silently shift every
    # subsequent question's text/options onto the wrong number.
    q_by_num: dict[int, dict] = {}
    for idx, q in enumerate(q_list):
        if not isinstance(q, dict):
            continue
        num = q.get("number") if isinstance(q.get("number"), int) else idx + 1
        q_by_num[num] = q

    questions: list[QuestionMeta] = []
    for i in range(n):
        number = i + 1
        q = q_by_num.get(number, {})
        text = str(q.get("text") or "")
        opts_raw = q.get("options") or []
        opts = tuple(str(o) for o in opts_raw) if isinstance(opts_raw, list) else ()
        questions.append(QuestionMeta(number=number, text=text, options=opts))

    return AnswerGrid(
        total_questions=n,
        official_answers=tuple(official),
        students=tuple(students),
        questions=tuple(questions),
    )


def score_student(grid: AnswerGrid, student_index: int) -> tuple[int, int, float]:
    """Return (graded_questions, correct_count, score_percent) for one student.

    Only questions with a concrete official letter (not blank, not the "="
    correct-marker) are graded. A student cell of "=" counts as correct.
    """
    graded = sum(
        1 for c in grid.official_answers if c is not None and c != _CORRECT_MARKER
    )
    wrong = len(diff_student(grid, student_index))
    correct = graded - wrong
    score = round(correct / graded * 100, 1) if graded else 0.0
    return graded, correct, score


def diff_student(grid: AnswerGrid, student_index: int) -> list[WrongAnswer]:
    """Return the student's wrong answers ordered by question number.

    - Skips questions where the official key is blank (None) or "=".
    - Skips questions where the student's cell is "=" (answered correctly).
    - Otherwise flags any cell that differs from the official letter.
    """
    if student_index < 0 or student_index >= len(grid.students):
        raise IndexError(f"student_index {student_index} out of range")
    row = grid.students[student_index].answers
    wrongs: list[WrongAnswer] = []
    for i, correct in enumerate(grid.official_answers):
        if correct is None or correct == _CORRECT_MARKER:
            continue
        student_ans = row[i] if i < len(row) else None
        if student_ans == _CORRECT_MARKER:
            continue
        if student_ans != correct:
            wrongs.append(
                WrongAnswer(
                    question=grid.questions[i],
                    student_answer=student_ans,
                    correct_answer=correct,
                )
            )
    return wrongs