# SPDX-FileCopyrightText: Copyright (C) 2026 Omid Jafari <omidjafari.com>
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
from collections import defaultdict
from operator import attrgetter
from typing import Any
from django.db.models import Prefetch
from django.db.models import Q
from django.http import HttpResponse
from rest_framework import status
from rest_framework.decorators import action
from rest_framework.exceptions import ValidationError
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework.viewsets import ViewSet
from nlpmed_portal.annotations.api.serializers import ExportAnnotationSerializer
from nlpmed_portal.annotations.models import Adjudication
from nlpmed_portal.annotations.models import LabelValue
from nlpmed_portal.annotations.models import Note
from nlpmed_portal.annotations.models import Patient
from nlpmed_portal.annotations.models import PatientAnnotation
from nlpmed_portal.annotations.permissions import MethodBasePermission
from nlpmed_portal.annotations.permissions import has_permission
from nlpmed_portal.annotations.utils import annotations_to_buffer
from nlpmed_portal.users.models import User
[docs]
class ExportAnnotationViewSet(ViewSet):
serializer_class = ExportAnnotationSerializer
permission_classes = [IsAuthenticated, MethodBasePermission]
required_perms = {
"GET": "view_exportannotation",
"HEAD": "view_exportannotation",
"OPTIONS": "view_exportannotation",
"POST": "add_exportannotation",
}
def _build_output( # ruff: ignore[complex-structure, too-many-branches, too-many-locals]
self,
patients: list[Patient],
annotators: list[User],
export_type: str,
):
# Load PatientAnnotation objects
pa_qs = (
PatientAnnotation.objects
.filter(
patient__in=patients,
annotator__in=annotators,
)
.select_related("patient", "annotator")
.order_by("event_date", "id")
.prefetch_related(
Prefetch(
"notes",
queryset=Note.objects.select_related("patient"),
),
Prefetch(
"event_label_values",
queryset=LabelValue.objects.select_related("category"),
),
Prefetch(
"adjudications",
queryset=Adjudication.objects.select_related(
"patient",
"adjudicator",
).prefetch_related(
Prefetch(
"event_label_values",
queryset=LabelValue.objects.select_related("category"),
),
Prefetch(
"source_annotations",
queryset=PatientAnnotation.objects.only("id"),
),
),
),
)
)
pa_list: list[PatientAnnotation] = list(pa_qs)
pas_by_patient: defaultdict[int, list[PatientAnnotation]] = defaultdict(list)
pas_by_note: defaultdict[int, list[PatientAnnotation]] = defaultdict(list)
notes_map: dict[int, Note] = {}
adj_by_note: defaultdict[int, set[Adjudication]] = defaultdict(set)
adj_note_ids: defaultdict[int, set[str]] = defaultdict(set)
for pa in pa_list:
pas_by_patient[pa.patient_id].append(pa)
pa_notes = list(pa.notes.all())
for note in pa_notes:
notes_map[note.id] = note
pas_by_note[note.id].append(pa)
for adj in pa.adjudications.all():
for note in pa_notes:
adj_by_note[note.id].add(adj)
adj_note_ids[adj.id].add(note.note_id)
def _labels_text(label_values) -> str:
lvs = sorted(label_values, key=lambda lv: (lv.category.name, lv.name))
return "\n".join(f"{lv.category.name}: {lv.name}" for lv in lvs)
def _pivot_annotation(pa: PatientAnnotation) -> dict:
return {
"pk": pa.id,
"annotator": pa.annotator.username,
"event_date": pa.event_date,
"labels": _labels_text(pa.event_label_values.all()),
"comment": pa.comment,
}
def _pivot_adjudication(adj: Adjudication) -> dict:
return {
"pk": adj.id,
"annotator": adj.adjudicator.username,
"event_date": adj.event_date,
"labels": _labels_text(adj.event_label_values.all()),
}
# Build output
output: list[dict] = []
if export_type == "note_level":
note_ids = list(pas_by_note.keys())
note_ids.sort(key=lambda nid: (notes_map[nid].patient_id, nid))
for note_id in note_ids:
note = notes_map[note_id]
patient = note.patient
row = {
# Internal Django IDs
"patient_pk": patient.id,
"note_pk": note.id,
# Data IDs
"pat_id": patient.pat_id,
"index_date": patient.index_date,
"note_id": note.note_id,
"note_date": note.note_date,
"note_type": note.note_type,
"note_text": note.note_text,
"preprocessed_text": note.preprocessed_text,
"predicted_label": note.predicted_label,
"predicted_score": note.predicted_score,
}
# Pivot annotations of this note
annotation_fields: dict[str, Any] = {}
note_pas = sorted(
pas_by_note[note_id],
key=attrgetter("id"),
)
for idx, pa in enumerate(note_pas, start=1):
pivot = _pivot_annotation(pa)
for k, v in pivot.items():
annotation_fields[f"annotation{idx}_{k}"] = v
# Pivot adjudications of this note
adjudication_fields: dict[str, Any] = {}
adjs_for_note = adj_by_note.get(note_id, set())
for jdx, adj in enumerate(
sorted(adjs_for_note, key=attrgetter("event_date", "id")),
start=1,
):
pivot = _pivot_adjudication(adj)
for k, v in pivot.items():
adjudication_fields[f"adjudication{jdx}_{k}"] = v
row.update(annotation_fields)
row.update(adjudication_fields)
output.append(row)
# patient_level
else:
for patient_id in sorted(pas_by_patient):
pas = sorted(
pas_by_patient[patient_id],
key=attrgetter("id"),
)
patient = pas[0].patient
row = {
"patient_pk": patient.id,
"pat_id": patient.pat_id,
"index_date": patient.index_date,
}
# Pivot annotations of this patient
annotation_fields = {}
for idx, pa in enumerate(pas, start=1):
pivot = _pivot_annotation(pa)
for k, v in pivot.items():
annotation_fields[f"annotation{idx}_{k}"] = v
# Pivot adjudications of this patient
adjudication_fields = {}
all_adjs = {adj for pa in pas for adj in pa.adjudications.all()}
for jdx, adj in enumerate(
sorted(all_adjs, key=attrgetter("event_date", "id")),
start=1,
):
pivot = _pivot_adjudication(adj)
for k, v in pivot.items():
adjudication_fields[f"adjudication{jdx}_{k}"] = v
row.update(annotation_fields)
row.update(adjudication_fields)
output.append(row)
return output
[docs]
def create(self, request: Any) -> Response:
serializer = self.serializer_class(data=request.data)
serializer.is_valid(raise_exception=True)
project = serializer.validated_data["project"]
patients = serializer.validated_data["patients"]
annotators = serializer.validated_data["annotators"]
export_type = serializer.validated_data["export_type"]
if not has_permission(
user=request.user,
codename="add_exportannotation",
project=project,
):
return Response(
{"error": "No permission to export annotations from this project."},
status=status.HTTP_403_FORBIDDEN,
)
return Response(
self._build_output(patients, annotators, export_type),
status=status.HTTP_200_OK,
)
[docs]
@action(detail=False, methods=["GET", "POST"])
def download(self, request):
serializer = self.serializer_class(data=request.data)
serializer.is_valid(raise_exception=True)
project = serializer.validated_data["project"]
patients = serializer.validated_data["patients"]
annotators = serializer.validated_data["annotators"]
export_type = serializer.validated_data["export_type"]
export_file_format = serializer.validated_data["export_file_format"]
if not has_permission(
user=request.user,
codename="add_exportannotation",
project=project,
):
return Response(
{"error": "No permission to export annotations from this project."},
status=status.HTTP_403_FORBIDDEN,
)
output = self._build_output(patients, annotators, export_type)
buf, content_type, filename = annotations_to_buffer(
output,
export_file_format,
f"project{project.id}_annotations_{export_type}.{export_file_format}",
)
resp = HttpResponse(buf.getvalue(), content_type=content_type)
resp["Content-Disposition"] = f'attachment; filename="{filename}"'
return resp