Copies/copienator/annotation_data.py

from __future__ import annotations

import copy
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from PIL import Image

from .json_io import read_json
from .feedback_boxes import valid_feedback_box
from .workspace import EvaluationWorkspace

AnnotationData = dict[str, dict[str, dict[str, Any]]]
RefaireList = list[list[Any]]


@dataclass(frozen=True, slots=True)
class GroupCoordinates:
    minimum: int
    maximum: int
    width: int
    height: int


@dataclass(slots=True)
class AnnotationLoadResult:
    data: AnnotationData
    warnings: list[str]


def _coordinate_index(
    workspace: EvaluationWorkspace,
) -> tuple[dict[tuple[str, str], GroupCoordinates], list[str]]:
    index: dict[tuple[str, str], GroupCoordinates] = {}
    warnings: list[str] = []
    if not workspace.groups_dir.is_dir():
        return index, [f"Group directory not found: {workspace.groups_dir}"]

    # A redo appends a new numbered group; its coordinates supersede the old group.
    for metadata_path in sorted(
        workspace.groups_dir.glob("*/Group_*.json"),
        key=lambda path: int(path.stem.removeprefix("Group_")),
        reverse=True,
    ):
        image_path = metadata_path.with_suffix(".jpg")
        try:
            entries = read_json(metadata_path)
            if not isinstance(entries, list):
                raise TypeError("expected a JSON array")
            with Image.open(image_path) as image:
                width, height = image.size
            for entry in entries:
                copy_id = str(entry[0])
                minimum = int(entry[1])
                maximum = int(entry[2])
                label = str(entry[4])
                index.setdefault(
                    (label, copy_id),
                    GroupCoordinates(minimum, maximum, width, height),
                )
        except (IndexError, OSError, TypeError, ValueError) as exc:
            warnings.append(f"Could not read group metadata {metadata_path}: {exc}")
    return index, warnings


def _scaled_result(result: dict[str, Any], coordinates: GroupCoordinates | None):
    scaled = copy.deepcopy(result)
    if coordinates is None:
        return scaled
    for feedback in scaled.get("feedback", []):
        box = feedback.get("box_2d")
        if not box or not valid_feedback_box(box):
            continue
        box[0] = int(box[0] * coordinates.height) // 1000
        box[2] = int(box[2] * coordinates.height) // 1000
        box[1] = int(box[1] * coordinates.width) // 1000
        box[3] = int(box[3] * coordinates.width) // 1000
    return scaled


def _answer_pdf(
    workspace: EvaluationWorkspace,
    copy_id: str,
    label: str,
    suffix: str = "",
) -> Path:
    copy_dir = workspace.copies_dir / f"Copie{copy_id}"
    preferred = copy_dir / f"{label}{suffix}.pdf"
    if preferred.exists():
        return preferred
    for candidate in (
        copy_dir / f"{label}.pdf",
        copy_dir / f"{label}_new.pdf",
    ):
        if candidate.exists():
            return candidate
    return preferred


def _dummy_entry(
    workspace: EvaluationWorkspace,
    copy_id: str,
    label: str,
) -> dict[str, Any]:
    return {
        "pdf_path": _answer_pdf(workspace, copy_id, label),
        "result": {
            "score": 0.0,
            "feedback": [],
            "error": "non traité",
        },
        "coordinates": (0, 0),
        "issues": ["No correction result was available for this answer."],
    }


def _apply_refaire_filter(
    workspace: EvaluationWorkspace,
    data: AnnotationData,
    refaire_list: RefaireList,
    warnings: list[str],
) -> AnnotationData:
    filtered: AnnotationData = {}
    for raw_entry in refaire_list:
        if not isinstance(raw_entry, list) or len(raw_entry) != 2:
            warnings.append(f"Ignoring malformed refaire entry: {raw_entry!r}")
            continue
        copy_name, requested_labels = raw_entry
        copy_id = str(copy_name).removeprefix("Copie")
        available = data.get(copy_id, {})
        if not requested_labels:
            filtered[copy_id] = dict(available)
            continue
        selected: dict[str, dict[str, Any]] = {}
        for raw_label in requested_labels:
            label = str(raw_label)
            selected[label] = (
                available[label]
                if label in available
                else _dummy_entry(workspace, copy_id, label)
            )
        filtered[copy_id] = selected
    return filtered


def load_annotation_data(
    workspace: EvaluationWorkspace,
    *,
    refaire_list: RefaireList | None = None,
    copy_id: str | None = None,
) -> AnnotationLoadResult:
    workspace.require_files("correction.json")
    corrections = read_json(workspace.correction_file)
    if not isinstance(corrections, dict):
        raise TypeError("correction.json must contain a JSON object")

    coordinate_index, warnings = _coordinate_index(workspace)
    data: AnnotationData = {}
    for label, raw_batches in corrections.items():
        if not isinstance(raw_batches, list):
            warnings.append(f"Ignoring malformed correction batches for {label!r}")
            continue
        for raw_batch in raw_batches:
            if not isinstance(raw_batch, list):
                warnings.append(f"Ignoring malformed correction batch for {label!r}")
                continue
            for item in raw_batch:
                if not isinstance(item, dict) or not isinstance(
                    item.get("result"), dict
                ):
                    warnings.append(f"Ignoring malformed correction item for {label!r}")
                    continue
                student_id = str(item.get("id", ""))
                if not student_id:
                    warnings.append(
                        f"Ignoring correction item without an id for {label!r}"
                    )
                    continue
                result = item["result"]
                suffix = str(result.get("suffix", ""))
                if suffix == "_old":
                    continue
                coordinates = coordinate_index.get((str(label), student_id))
                issues: list[str] = []
                if coordinates is None:
                    issues.append("Group coordinates were not found.")
                pdf_path = _answer_pdf(workspace, student_id, str(label), suffix)
                if not pdf_path.exists():
                    issues.append(f"Answer PDF not found: {pdf_path}")
                data.setdefault(student_id, {})[str(label)] = {
                    "pdf_path": pdf_path,
                    "result": _scaled_result(result, coordinates),
                    "coordinates": (
                        (coordinates.minimum, coordinates.maximum)
                        if coordinates is not None
                        else (0, 0)
                    ),
                    "issues": issues,
                }
                warnings.extend(
                    f"Copie{student_id} {label}: {issue}" for issue in issues
                )

    if refaire_list is not None:
        data = _apply_refaire_filter(workspace, data, refaire_list, warnings)
    if copy_id is not None:
        data = {copy_id: data[copy_id]} if copy_id in data else {}
        if not data:
            warnings.append(f"Copy id {copy_id} was not found in correction.json")
    return AnnotationLoadResult(data, warnings)