from __future__ import annotations import argparse import concurrent.futures from collections import defaultdict from collections.abc import Sequence from pathlib import Path from typing import Any from PIL import Image, ImageDraw from copienator import configuration from copienator import ( EvaluationWorkspace, ExitCode, atomic_write_json, evaluation_parser, execute, read_json, utils, workspace_from_args, ) from copienator.annotation_actions import apply_checkbox_actions, apply_score_overrides from copienator.answer_info import build_answer_info from copienator.annotation_data import AnnotationData, RefaireList, load_annotation_data from copienator.commands import annotating from copienator.commands.reading_annotations import ( concatenate, detect_checks_and_notes, has_significant_notes, ) from copienator.filesystem import staged_files from copienator.return_answers import save_return_answer_options LabelNotes = dict[str, dict[str, Any]] ScanResult = tuple[dict[str, list[dict[str, Any]]], dict[str, LabelNotes]] SCAN_WORKERS = 2 def get_extra_pdfs_as_images( root_dir: str | Path, label: str, annotating_module: Any, all_labels: list[str], ) -> list[Image.Image]: """Convert the context, question and solution PDFs associated with a label.""" paths = [ *utils.pdf_images_of_contexts(root_dir, label, all_labels), utils.pdf_image_of_enonce(root_dir, label), utils.pdf_image_of_solution(root_dir, label), ] images = [] for path in paths: if path: image, _, _ = annotating_module.make_base_image(path) if image is not None: images.append(image) return images def save_paginated_pdf( image_groups: list[list[Image.Image]], output_path: Path ) -> None: """Paginate vertically concatenated image groups and save them as a PDF.""" non_empty = [group for group in image_groups if group] if not non_empty: return max_width = max(image.width for group in non_empty for image in group) max_page_height = int(max_width * 1.414 * 1.25) border = int((0.2 / 2.54) * 100) left_margin = int((0.3 / 2.54) * 100) vertical_margin = int((0.2 / 2.54) * 100) max_content_height = max_page_height - 2 * vertical_margin pages: list[Image.Image] = [] page_images: list[Image.Image] = [] page_height = 0 def finish_page() -> None: nonlocal page_images, page_height if not page_images: return page = Image.new( "RGB", (max_width + left_margin, page_height + 2 * vertical_margin), "white", ) current_y = vertical_margin for image in page_images: page.paste(image, (left_margin, current_y)) current_y += image.height pages.append(page) page_images = [] page_height = 0 for group in non_empty: processed: list[Image.Image] = [] for index, image in enumerate(group): if index in (0, 1): image = image.copy() color = "black" if index == 0 else "blue" ImageDraw.Draw(image).rectangle( [0, 0, image.width - 1, image.height - 1], outline=color, width=border, ) processed.append(image) group_height = sum(image.height for image in processed) if page_images and page_height + group_height > max_content_height: finish_page() page_images.extend(processed) page_height += group_height finish_page() pages[0].save( output_path, "PDF", resolution=100.0, save_all=True, append_images=pages[1:], ) def _scan_annotation_directory( directory: Path, only_ids: set[str] | None = None, default_student_id: str | None = None, *, required: bool = False, ) -> ScanResult: bnote_path = directory / "bnote.json" if not bnote_path.is_file(): raise FileNotFoundError(f"Missing {bnote_path}") bnote = read_json(bnote_path) if not isinstance(bnote, dict): raise TypeError(f"Expected a JSON object in {bnote_path}") images = [item for item in bnote.get("images", []) if isinstance(item, dict)] if only_ids and not any( str(item.get("id", default_student_id)) in only_ids for item in images ): return {}, {} actions, notes_image = detect_checks_and_notes(directory) if notes_image is None: if required: raise ValueError(f"Could not read annotations in {directory}") return {}, {} actions_by_student: dict[str, list[dict[str, Any]]] = defaultdict(list) notes_by_student: dict[str, LabelNotes] = defaultdict(dict) for action in actions: raw_student_id = action.get("student_id", default_student_id) if raw_student_id is not None: actions_by_student[str(raw_student_id)].append(action) for image_info in images: student_id = str(image_info.get("id", default_student_id or "")) label = str(image_info.get("label", "")) hmin = int(image_info.get("hmin", 0)) hmax = int(image_info.get("hmax", 0)) if student_id and label and hmax > hmin: crop = notes_image.crop((0, hmin, notes_image.width, hmax)) if has_significant_notes(crop): notes_by_student[student_id][label] = { "img": crop, "old_header_h": int(image_info.get("header_height", 0)), } return dict(actions_by_student), dict(notes_by_student) def _merge_scan_result( target_actions: dict[str, list[dict[str, Any]]], target_notes: dict[str, LabelNotes], result: ScanResult, ) -> None: actions, notes = result for student_id, student_actions in actions.items(): target_actions[student_id].extend(student_actions) for student_id, student_notes in notes.items(): target_notes[student_id].update(student_notes) def apply_actions_and_regenerate_grouped( workspace: EvaluationWorkspace, data: AnnotationData, student_id: str, actions: list[dict[str, Any]], label_notes: LabelNotes, all_labels: list[str], *, update_score: bool = False, annotation_dir: str = "BGnot", selected_labels: set[str] | None = None, ) -> tuple[ExitCode, str]: """Regenerate a copy, preserving reviewed images outside the redo selection.""" logs = [f"\nProcessing compilation for: Copie{student_id}"] output_dir = workspace.root / annotation_dir / f"Copie{student_id}" labels_data = data.get(student_id, {}) apply_checkbox_actions(labels_data, actions, logs.append) score_path = output_dir / "score.json" preserve_score_file = update_score and score_path.is_file() if update_score: apply_score_overrides( labels_data, score_path, logs.append ) selected_labels = selected_labels if selected_labels is not None else set() simple_layout = None simple_annotated = None if selected_labels and annotation_dir == "Anot": imported = next( ( output_dir / name for name in ("Concat_annotated.jpg", "Concat_annotated.jpeg") if (output_dir / name).is_file() ), None, ) if imported is not None: layout_path = output_dir / "refaire_simple_layout.json" if layout_path.is_file(): simple_layout = read_json(layout_path) else: simple_layout = {"images": {}, "replaced": []} y = 0 for label in sorted(labels_data, key=utils.natural_key): path = output_dir / f"{label}.jpg" if ( path.is_file() and labels_data[label]["result"].get("error") != "empty-answer" ): with Image.open(path) as saved: simple_layout["images"][label] = [y, y + saved.height] y += saved.height with Image.open(imported) as saved: simple_annotated = saved.convert("RGB").copy() expected_height = max( (bounds[1] for bounds in simple_layout["images"].values()), default=0 ) if simple_annotated.height != expected_height: raise ValueError( "Imported simple image height does not match the original copy layout" ) old_scores = ( read_json(output_dir / "score.json") if selected_labels and (output_dir / "score.json").is_file() else {} ) scores = dict.fromkeys(all_labels, "") touched = dict.fromkeys(all_labels, False) answer_labels: list[str] = [] dirty_images: dict[str, Image.Image] = {} concat_images: list[Image.Image] = [] filtered_groups: list[list[Image.Image]] = [] incomplete = False for label, content in sorted( labels_data.items(), key=lambda item: utils.natural_key(item[0]) ): result = content["result"] if ( selected_labels and label not in selected_labels and old_scores.get(label, "") != "" ): result["score"] = old_scores[label] scores[label] = str(result.get("score", 0)) touched[label] = False if result.get("error") == "empty-answer": continue saved_image = output_dir / f"{label}.jpg" if selected_labels and label not in selected_labels and saved_image.is_file(): with Image.open(saved_image) as saved: final_image = saved.convert("RGB").copy() if ( simple_annotated is not None and label in simple_layout["images"] and label not in simple_layout["replaced"] ): hmin, hmax = simple_layout["images"][label] final_image = simple_annotated.crop( (0, hmin, simple_annotated.width, hmax) ) dirty_images[label] = final_image scores[label] = str(old_scores.get(label, scores[label])) concat_images.append(final_image) answer_labels.append(label) # Keep previously reviewed content, including handwriting. if annotation_dir == "BGnot": extras = get_extra_pdfs_as_images( workspace.root, label, annotating, all_labels ) filtered_groups.append([*extras, final_image]) touched[label] = True else: filtered_groups.append([final_image]) continue pdf_path = Path(content["pdf_path"]) if not pdf_path.is_file(): logs.append(f" Missing answer PDF: {pdf_path}") incomplete = True continue base_image, _, _ = annotating.make_base_image(pdf_path) final_image, new_header_height = annotating.compose_label_image( base_image, label, result, content["coordinates"][0], with_error=False, ) if final_image is None: incomplete = True continue has_notes = False if label in label_notes: sub_note = label_notes[label]["img"] old_header_height = int(label_notes[label]["old_header_h"]) has_notes = has_significant_notes(sub_note) if has_notes: width, height = sub_note.size if old_header_height > 0: header = sub_note.crop( (0, 0, width, min(height, old_header_height)) ) final_image.paste(header, (0, 0), mask=header) if height > old_header_height: body = sub_note.crop((0, old_header_height, width, height)) final_image.paste(body, (0, new_header_height), mask=body) # Persist every final block, including unchanged answers, for returns. dirty_images[label] = final_image answer_labels.append(label) concat_images.append(final_image) feedbacks = result.get("feedback", []) perfect = float(scores[label]) >= 4.0 and all( feedback.get("to_delete", False) for feedback in feedbacks ) if not perfect or has_notes: extras = ( get_extra_pdfs_as_images(workspace.root, label, annotating, all_labels) if annotation_dir == "BGnot" else [] ) filtered_groups.append([*extras, final_image]) touched[label] = annotation_dir == "BGnot" concat_image = concatenate(concat_images) if incomplete: return ExitCode.PARTIAL, "\n".join(logs) with staged_files(output_dir, remove=("Concat.jpg", "Concat_F.pdf", "Concat_F.jpg", "touched.json", "answer_labels.json")) as staging: if simple_layout is not None: simple_layout["replaced"] = sorted( set(simple_layout["replaced"]) | selected_labels ) atomic_write_json(staging / "refaire_simple_layout.json", simple_layout) for label, image in dirty_images.items(): image.save(staging / f"{label}.jpg") if not preserve_score_file: atomic_write_json(staging / "score.json", scores) atomic_write_json(staging / "info.json", build_answer_info( scores, labels_data, answer_labels, touched )) if concat_image is not None: concat_image.save(staging / "Concat.jpg") if filtered_groups: if annotation_dir == "BGnot": save_paginated_pdf(filtered_groups, staging / "Concat_F.pdf") else: filtered_image = concatenate( [image for group in filtered_groups for image in group] ) filtered_image.save(staging / "Concat_F.jpg") if preserve_score_file: logs.append(f" Preserved existing score.json in {output_dir}") logs.append(f" Saved regenerated files in {output_dir}") status = ExitCode.PARTIAL if incomplete else ExitCode.SUCCESS return status, "\n".join(logs) def _read_refaire( workspace: EvaluationWorkspace, ) -> tuple[RefaireList, dict[str, list[str]]]: loaded = read_json(workspace.refaire_file) if not isinstance(loaded, list): raise TypeError("refaire.json must contain a JSON array") entries: RefaireList = [] by_student: dict[str, list[str]] = {} for entry in loaded: if ( not isinstance(entry, list) or len(entry) != 2 or not isinstance(entry[1], list) ): raise TypeError(f"Malformed refaire entry: {entry!r}") copy_name, labels = entry student_id = str(copy_name).removeprefix("Copie") normalized_labels = [str(label) for label in labels] entries.append([str(copy_name), normalized_labels]) by_student[student_id] = normalized_labels return entries, by_student def _scan_redo_annotations( directory: Path, expected: dict[str, set[str]], ) -> tuple[dict[str, list[dict[str, Any]]], dict[str, LabelNotes], set[str]]: """Read either grouped or per-copy redo PDFs using their student/label metadata.""" actions: dict[str, list[dict[str, Any]]] = defaultdict(list) notes: dict[str, LabelNotes] = defaultdict(dict) seen: dict[str, set[str]] = defaultdict(set) incomplete: set[str] = set() plans = [] required = ("checkboxes.json", "Reference.jpg", "Concat_annotated.pdf") for path in sorted(directory.iterdir()): if not path.is_dir(): continue default_id = ( path.name.removeprefix("Copie") if path.name.startswith("Copie") else None ) try: metadata = read_json(path / "bnote.json") pairs = [ (str(item.get("id", default_id)), str(item["label"])) for item in metadata["images"] ] except (OSError, ValueError, TypeError, KeyError) as exc: print(f"Warning: unreadable redo metadata in {path}: {exc}") incomplete.update(expected) continue students = { student_id for student_id, _label in pairs if student_id in expected } if not students: continue for student_id, label in pairs: if student_id not in expected: continue if label in seen[student_id]: incomplete.add(student_id) seen[student_id].add(label) if any(not (path / name).is_file() for name in required): print(f"Warning: missing returned redo inputs in {path}") incomplete.update(students) else: plans.append((path, default_id, students)) for student_id, labels in expected.items(): if seen[student_id] != labels: print( f"Warning: redo labels do not match refaire.json for Copie{student_id}; regenerate BRnot" ) incomplete.add(student_id) for path, default_id, students in plans: if students <= incomplete: continue try: result = _scan_annotation_directory( path, default_student_id=default_id, required=True ) _merge_scan_result(actions, notes, result) except (OSError, ValueError, TypeError) as exc: print(f"Warning: could not read redo annotations in {path}: {exc}") incomplete.update(students) return dict(actions), dict(notes), incomplete def run( workspace: EvaluationWorkspace, *, refaire: bool = False, update_score: bool = False, annotation_dir: str = "BGnot", return_answers_context: bool | None = None, return_answers_question: bool | None = None, return_answers_solution: bool | None = None, ) -> ExitCode: workspace.require_files("labels", "correction.json") workspace.require_directories("Copies", "Par label", annotation_dir) refaire_list: RefaireList | None = None refaire_by_student: dict[str, list[str]] = {} if refaire: workspace.require_files("refaire.json") workspace.require_directories("BRnot") refaire_list, refaire_by_student = _read_refaire(workspace) if configuration.RETURN_ANSWERS_ENABLED: save_return_answer_options( workspace.root, context=( configuration.RETURN_ANSWERS_CONTEXT if return_answers_context is None else return_answers_context ), question=( configuration.RETURN_ANSWERS_QUESTION if return_answers_question is None else return_answers_question ), solution=( configuration.RETURN_ANSWERS_SOLUTION if return_answers_solution is None else return_answers_solution ), ) all_labels = utils.read_all_labels(workspace.root) loaded = load_annotation_data(workspace) if refaire_list: # Add explicitly requested answers without filtering out the rest of a copy. selected_data = load_annotation_data(workspace, refaire_list=refaire_list) for student_id, labels in selected_data.data.items(): loaded.data.setdefault(student_id, {}).update(labels) for warning in loaded.warnings: print(f"Warning: {warning}") if not loaded.data: print("No annotation data found.") return ExitCode.PARTIAL actions_by_student: dict[str, list[dict[str, Any]]] = defaultdict(list) notes_by_student: dict[str, LabelNotes] = defaultdict(dict) only_ids = set(refaire_by_student) or None group_dirs = [ path for path in (workspace.root / annotation_dir).iterdir() if annotation_dir == "BGnot" and path.is_dir() and not path.name.startswith("Copie") ] # Each worker decodes a full-height returned group and its reference image. # Keep this stage deliberately narrow; answer regeneration below has its own # parallel executor and a much smaller per-task memory footprint. with concurrent.futures.ThreadPoolExecutor(max_workers=SCAN_WORKERS) as executor: futures = [ executor.submit(_scan_annotation_directory, path, only_ids) for path in group_dirs ] for future in concurrent.futures.as_completed(futures): _merge_scan_result(actions_by_student, notes_by_student, future.result()) if annotation_dir == "Bnot": for student_id in refaire_by_student: directory = workspace.root / annotation_dir / f"Copie{student_id}" if directory.is_dir(): _merge_scan_result( actions_by_student, notes_by_student, _scan_annotation_directory( directory, default_student_id=student_id ), ) skipped_students: set[str] = set() if refaire: expected = { student_id: set(labels or loaded.data.get(student_id, {})) for student_id, labels in refaire_by_student.items() } redo_actions, redo_notes, skipped_students = _scan_redo_annotations( workspace.annotation_dir("refaire"), expected ) for student_id, selected in expected.items(): if student_id in skipped_students: continue actions_by_student[student_id] = [ action for action in actions_by_student[student_id] if str(action.get("label")) not in selected ] for label in selected: notes_by_student[student_id].pop(label, None) actions_by_student[student_id].extend( action for action in redo_actions.get(student_id, []) if str(action.get("label")) in selected ) notes_by_student[student_id].update( { label: note for label, note in redo_notes.get(student_id, {}).items() if label in selected } ) status = ( ExitCode.PARTIAL if loaded.warnings or skipped_students else ExitCode.SUCCESS ) student_ids = ( list(refaire_by_student) if refaire else sorted(loaded.data, key=utils.natural_key) ) with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: futures = { executor.submit( apply_actions_and_regenerate_grouped, workspace, loaded.data, student_id, actions_by_student[student_id], notes_by_student[student_id], all_labels, update_score=update_score, annotation_dir=annotation_dir, selected_labels=( set(refaire_by_student[student_id] or loaded.data[student_id]) if refaire else None ), ): student_id for student_id in student_ids if student_id in loaded.data and student_id not in skipped_students } for future in concurrent.futures.as_completed(futures): result, output = future.result() print(output) if result != ExitCode.SUCCESS: status = ExitCode.PARTIAL return status def build_parser() -> argparse.ArgumentParser: parser = evaluation_parser("Read grouped annotations and regenerate copies") parser.add_argument( "--annotation-dir", choices=("BGnot", "Bnot", "Anot"), default="BGnot", help="Original annotation directory for --refaire (default: BGnot)", ) parser.add_argument( "--refaire", action="store_true", help="Use refaire.json and merge annotations from BRnot", ) parser.add_argument( "--update-score", action="store_true", help=( "Regenerate images with current statement/solution PDFs while " "preserving and applying existing score.json values" ), ) parser.add_argument( "--return-answers-context", action=argparse.BooleanOptionalAction, default=configuration.RETURN_ANSWERS_CONTEXT, help="Include applicable context pages in individual answer exports", ) parser.add_argument( "--return-answers-question", action=argparse.BooleanOptionalAction, default=configuration.RETURN_ANSWERS_QUESTION, help="Include the current question PDF in individual answer exports", ) parser.add_argument( "--return-answers-solution", action=argparse.BooleanOptionalAction, default=configuration.RETURN_ANSWERS_SOLUTION, help="Include the current solution PDF in individual answer exports", ) return parser def main(argv: Sequence[str] | None = None) -> int: parser = build_parser() def handle(args: argparse.Namespace) -> ExitCode: if args.annotation_dir != "BGnot" and not args.refaire: parser.error("--annotation-dir requires --refaire") return run( workspace_from_args(args), refaire=args.refaire, update_score=args.update_score, annotation_dir=args.annotation_dir, return_answers_context=args.return_answers_context, return_answers_question=args.return_answers_question, return_answers_solution=args.return_answers_solution, ) return execute(parser, argv, handle) if __name__ == "__main__": raise SystemExit(main())