Source code for nlpmed_portal.annotations.api.serializers

# 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 django.db import transaction
from django.db.models import Exists
from django.db.models import OuterRef
from django.db.models import Q
from drf_writable_nested import WritableNestedModelSerializer
from rest_framework import serializers

from nlpmed_portal.annotations.models import Adjudication
from nlpmed_portal.annotations.models import ImportJob
from nlpmed_portal.annotations.models import Lab
from nlpmed_portal.annotations.models import LabelCategory
from nlpmed_portal.annotations.models import LabelValue
from nlpmed_portal.annotations.models import NERAnnotation
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.models import Project
from nlpmed_portal.annotations.models import ProjectMembership
from nlpmed_portal.annotations.models import Section
from nlpmed_portal.annotations.models import Sentence
from nlpmed_portal.annotations.models import Task
from nlpmed_portal.annotations.utils import build_annotation_window_start_map
from nlpmed_portal.nlp.api.serializers import NLPSettingSerializer
from nlpmed_portal.users.models import User




class LabelValueSerializer(serializers.ModelSerializer):


    class Meta:
        model = LabelValue
        fields = ["id", "name"]






class LabelCategorySerializer(WritableNestedModelSerializer):
    label_values = LabelValueSerializer(many=True, required=False)



    class Meta:
        model = LabelCategory
        fields = ["id", "name", "label_values"]






class ProjectListSerializer(serializers.ModelSerializer):
    label_categories = LabelCategorySerializer(many=True, read_only=True)



    class Meta:
        model = Project
        fields = [
            "id",
            "name",
            "annotation_type",
            "required_patient_annotations",
            "required_patient_adjudications",
            "label_categories",
        ]
        read_only_fields = fields






class ProjectDetailSerializer(WritableNestedModelSerializer):
    label_categories = LabelCategorySerializer(many=True, required=False)
    nlp_setting = NLPSettingSerializer(required=False)



    class Meta:
        model = Project
        fields = [
            "id",
            "name",
            "description",
            "disp_range_months_before_idx",
            "disp_range_months_after_idx",
            "run_nlp",
            "annotation_type",
            "required_patient_annotations",
            "required_patient_adjudications",
            "label_categories",
            "nlp_setting",
        ]






class PatientSerializer(serializers.ModelSerializer):
    project_name = serializers.CharField(source="project.name", read_only=True)
    num_distinct_annotators = serializers.IntegerField(read_only=True)
    num_distinct_adjudicators = serializers.IntegerField(read_only=True)



    class Meta:
        model = Patient
        fields = "__all__"






class SentenceSerializer(serializers.ModelSerializer):


    class Meta:
        model = Sentence
        fields = [
            "id",
            "text",
            "start_index",
            "end_index",
            "is_duplicate",
            "is_important",
            "is_expanded",
        ]






class SectionSerializer(WritableNestedModelSerializer):
    section_sentences = SentenceSerializer(many=True, required=False)



    class Meta:
        model = Section
        fields = ["id", "text", "start_index", "end_index", "note", "section_sentences"]






class NoteSerializer(WritableNestedModelSerializer):
    note_sections = SectionSerializer(many=True, required=False)



    class Meta:
        model = Note
        fields = [
            "id",
            "note_id",
            "note_text",
            "note_date",
            "note_type",
            "patient",
            "preprocessed_text",
            "predicted_label",
            "predicted_score",
            "note_sections",
            "important_sentences_count",
            "duplicate_sentences_count",
        ]






class LabSerializer(serializers.ModelSerializer):


    class Meta:
        model = Lab
        fields = "__all__"






class PatientAnnotationListSerializer(serializers.ModelSerializer):
    annotator_name = serializers.CharField(source="annotator.username", read_only=True)
    event_label_cat_values = serializers.SerializerMethodField()



    class Meta:
        model = PatientAnnotation
        fields = [
            "id",
            "patient",
            "annotator",
            "annotator_name",
            "event_date",
            "notes",
            "comment",
            "event_label_values",
            "event_label_cat_values",
        ]
        read_only_fields = fields




    def get_event_label_cat_values(self, obj):
        cat_dict = defaultdict(list)

        for lv in obj.event_label_values.all():
            cat_dict[lv.category].append(lv)

        return [
            {
                "id": cat.pk,
                "name": cat.name,
                "label_values": [{"id": lv.pk, "name": lv.name} for lv in lvs],
            }
            for cat, lvs in cat_dict.items()
        ]






class PatientAnnotationDetailSerializer(serializers.ModelSerializer):
    annotator = serializers.HiddenField(default=serializers.CurrentUserDefault())
    annotator_name = serializers.CharField(source="annotator.username", read_only=True)
    event_label_cat_values = serializers.SerializerMethodField()
    all_label_cat_values = serializers.SerializerMethodField()



    class Meta:
        model = PatientAnnotation
        fields = [
            "id",
            "patient",
            "annotator",
            "annotator_name",
            "event_date",
            "notes",
            "comment",
            "event_label_values",
            "event_label_cat_values",
            "all_label_cat_values",
        ]
        read_only_fields = ["id", "annotator_name"]




    def get_event_label_cat_values(self, obj):
        cat_dict = defaultdict(list)

        for lv in obj.event_label_values.all():
            cat_dict[lv.category].append(lv)

        return [
            {
                "id": cat.pk,
                "name": cat.name,
                "label_values": [{"id": lv.pk, "name": lv.name} for lv in lvs],
            }
            for cat, lvs in cat_dict.items()
        ]




    def get_all_label_cat_values(self, obj):
        categories = obj.patient.project.label_categories.all()
        return LabelCategorySerializer(categories, many=True).data




    def validate(self, attrs):  # ruff: ignore[complex-structure, too-many-branches]
        instance = getattr(self, "instance", None)

        if instance is not None:
            is_referenced = Adjudication.objects.filter(
                source_annotations__id=instance.id,
            ).exists()

            if is_referenced:
                raise serializers.ValidationError(
                    {
                        "detail": "This annotation is referenced by an \
                            adjudication and cannot be edited.",
                    },
                )

        new_event_date = attrs.get(
            "event_date",
            instance.event_date if instance else None,
        )

        if "event_label_values" in attrs:
            new_event_label_values = attrs["event_label_values"]

        else:
            new_event_label_values = instance.event_label_values.all() if instance else []

        if not isinstance(new_event_label_values, list):
            new_event_label_values = list(new_event_label_values)

        both_missing = (not new_event_date) and (len(new_event_label_values) == 0)

        # This is a negative patient annotation
        if both_missing:
            patient = attrs.get("patient", instance.patient if instance else None)
            if not patient:
                raise serializers.ValidationError(
                    {"Error": "No patient specified or found."},
                )

            project = patient.project
            none_categories = project.label_categories.filter(name__iexact="none")
            if not none_categories.exists():
                raise serializers.ValidationError(
                    {"Error": 'No label category named "none" found in this project.'},
                )

            none_category = none_categories.first()
            none_label_values = none_category.label_values.all()
            if none_label_values.count() != 1:
                raise serializers.ValidationError(
                    {
                        "Error": '"none" category must have exactly \
                            one label value in this project.',
                    },
                )

            attrs["event_date"] = None
            attrs["event_label_values"] = [none_label_values.first()]

        # This is a positive patient annotation
        else:
            if "event_date" in attrs and new_event_date in {None, ""}:
                raise serializers.ValidationError(
                    {"event_date": "event_date is required."},
                )

            if "event_label_values" in attrs and len(new_event_label_values) == 0:
                raise serializers.ValidationError(
                    {"event_label_values": "event_label_values is required."},
                )

        attrs = super().validate(attrs)
        patient = attrs.get("patient") or (instance.patient if instance else None)
        if patient is None:
            raise serializers.ValidationError(
                {"patient": "patient is required."},
            )

        # Should have correct annotation type
        if patient.project.annotation_type != Project.CLASSIFICATION:
            raise serializers.ValidationError(
                {"Error": "This is not a classification project."},
            )

        # Notes should belong to the same patient
        for note in attrs.get("notes", []):
            if note.patient_id != patient.id:
                raise serializers.ValidationError(
                    {
                        "Error": f"Note {note.pk} does not belong to this patient.",
                    },
                )

        # Labels should belong to the same project
        for lv in attrs.get("event_label_values", []):
            if lv.category.project_id != patient.project_id:
                raise serializers.ValidationError(
                    {
                        "Error": f"LabelValue {lv} not in this project.",
                    },
                )

        return attrs






class ProjectMembershipListSerializer(serializers.ModelSerializer):
    group_name = serializers.CharField(source="group.name", read_only=True)
    assignee_name = serializers.CharField(source="assignee.username", read_only=True)
    assigner_name = serializers.CharField(source="assigner.username", read_only=True)
    project_name = serializers.CharField(source="project.name", read_only=True)



    class Meta:
        model = ProjectMembership
        fields = [
            "id",
            "assignee",
            "assignee_name",
            "project",
            "project_name",
            "group",
            "group_name",
            "assigner",
            "assigner_name",
            "created_at",
        ]
        read_only_fields = fields






class ProjectMembershipDetailSerializer(serializers.ModelSerializer):
    assigner = serializers.HiddenField(default=serializers.CurrentUserDefault())
    group_name = serializers.CharField(source="group.name", read_only=True)
    assignee_name = serializers.CharField(source="assignee.username", read_only=True)
    assigner_name = serializers.CharField(source="assigner.username", read_only=True)
    project_name = serializers.CharField(source="project.name", read_only=True)
    all_groups = serializers.SerializerMethodField()
    all_users = serializers.SerializerMethodField()



    class Meta:
        model = ProjectMembership
        fields = [
            "id",
            "assignee",
            "assignee_name",
            "project",
            "project_name",
            "group",
            "group_name",
            "assigner",
            "assigner_name",
            "created_at",
            "all_groups",
            "all_users",
        ]
        read_only_fields = [
            "id",
            "assignee_name",
            "project_name",
            "group_name",
            "assigner_name",
            "created_at",
        ]




    def get_all_groups(self, obj):
        return self.context["all_groups"]




    def get_all_users(self, obj):
        return self.context["all_users"]






class TaskListSerializer(serializers.ModelSerializer):
    assignee_name = serializers.CharField(source="assignee.username", read_only=True)
    assigner_name = serializers.CharField(source="assigner.username", read_only=True)
    project_name = serializers.CharField(source="patient.project.name", read_only=True)
    status_name = serializers.CharField(source="get_status_display", read_only=True)



    class Meta:
        model = Task
        fields = [
            "id",
            "assignee",
            "assignee_name",
            "patient",
            "status",
            "status_name",
            "assigner",
            "assigner_name",
            "created_at",
            "task_type",
            "project_name",
        ]
        read_only_fields = fields






class TaskDetailSerializer(serializers.ModelSerializer):
    codenames = [
        "add_patientannotation",
        "add_nerannotation",
        "add_adjudication",
    ]
    assignee = serializers.PrimaryKeyRelatedField(
        queryset=User.objects.filter(
            Q(user_permissions__codename__in=codenames)
            | Q(groups__permissions__codename__in=codenames)
            | Exists(
                ProjectMembership.objects.filter(
                    assignee=OuterRef("pk"),
                    group__permissions__codename__in=codenames,
                ),
            ),
        ).distinct(),
    )

    assignee_name = serializers.CharField(source="assignee.username", read_only=True)
    assigner_name = serializers.CharField(source="assigner.username", read_only=True)
    all_assignees = serializers.SerializerMethodField()



    class Meta:
        model = Task
        fields = [
            "id",
            "assignee",
            "assignee_name",
            "patient",
            "status",
            "assigner",
            "assigner_name",
            "created_at",
            "task_type",
            "all_assignees",
        ]
        read_only_fields = ["id", "assignee_name", "assigner_name", "created_at"]




    def get_all_assignees(self, obj):
        qs = User.objects.filter(
            Q(user_permissions__codename__in=self.codenames)
            | Q(groups__permissions__codename__in=self.codenames)
            | Exists(
                ProjectMembership.objects.filter(
                    assignee=OuterRef("pk"),
                    group__permissions__codename__in=self.codenames,
                ),
            ),
        ).distinct()

        return list(qs.values("id", "username"))






class BulkCreateTaskSerializer(serializers.Serializer):
    codenames = [
        "add_patientannotation",
        "add_nerannotation",
        "add_adjudication",
    ]
    assignee = serializers.PrimaryKeyRelatedField(
        queryset=User.objects.filter(
            Q(user_permissions__codename__in=codenames)
            | Q(groups__permissions__codename__in=codenames)
            | Exists(
                ProjectMembership.objects.filter(
                    assignee=OuterRef("pk"),
                    group__permissions__codename__in=codenames,
                ),
            ),
        ).distinct(),
    )
    patients = serializers.PrimaryKeyRelatedField(
        queryset=Patient.objects.all(),
        many=True,
    )
    task_type = serializers.ChoiceField(choices=Task.TYPE_CHOICES)



    def validate(self, attrs):
        assignee = attrs["assignee"]
        patients = attrs["patients"]
        task_type = attrs["task_type"]

        existing_tasks = Task.objects.filter(
            assignee=assignee,
            patient__in=patients,
            task_type=task_type,
        )

        if existing_tasks.exists():
            existing_patient_ids = set(
                existing_tasks.values_list("patient_id", flat=True),
            )
            duplicate_patients = [p.id for p in patients if p.id in existing_patient_ids]

            raise serializers.ValidationError(
                {
                    "Error": f"{task_type} task already exist for assignee \
                        {assignee.username} and patients {duplicate_patients}.",
                },
            )

        return super().validate(attrs)






class ImportJobSerializer(serializers.ModelSerializer):
    username = serializers.CharField(source="user.username", read_only=True)



    class Meta:
        model = ImportJob
        fields = "__all__"
        read_only_fields = [
            "job_status",
            "error_message",
            "created_at",
            "updated_at",
            "job_progress",
            "new_rows_count",
            "user",
        ]






class AdjudicationSerializer(serializers.ModelSerializer):
    adjudicator = serializers.HiddenField(
        default=serializers.CurrentUserDefault(),
    )
    adjudicator_name = serializers.CharField(
        source="adjudicator.username",
        read_only=True,
    )
    event_label_cat_values = serializers.SerializerMethodField()



    class Meta:
        model = Adjudication
        fields = [
            "id",
            "patient",
            "event_date",
            "event_label_values",
            "event_label_cat_values",
            "adjudicator",
            "adjudicator_name",
            "notes",
            "source_annotations",
        ]
        read_only_fields = ["id", "adjudicator_name", "event_label_cat_values", "notes"]




    def get_event_label_cat_values(self, obj):
        cat_dict = defaultdict(list)

        for lv in obj.event_label_values.all():
            cat_dict[lv.category].append(lv)

        return [
            {
                "id": cat.pk,
                "name": cat.name,
                "label_values": [{"id": lv.pk, "name": lv.name} for lv in lvs],
            }
            for cat, lvs in cat_dict.items()
        ]




    def validate(self, attrs):  # ruff: ignore[complex-structure]
        instance = getattr(self, "instance", None)
        patient = attrs.get("patient") or (instance.patient if instance else None)
        if not patient:
            raise serializers.ValidationError({"patient": "patient is required."})

        event_date = attrs.get("event_date", instance.event_date if instance else None)
        if not event_date:
            raise serializers.ValidationError({"event_date": "event_date is required."})

        source_annotations = attrs.get(
            "source_annotations",
            instance.source_annotations.all() if instance else [],
        )
        source_annotations = list(source_annotations)

        if not source_annotations:
            raise serializers.ValidationError(
                {"source_annotations": "At least one source annotation is required."},
            )

        for note in attrs.get("notes", []):
            if note.patient_id != patient.id:
                raise serializers.ValidationError(
                    {"notes": f"Note {note.pk} not in this patient."},
                )

        for ann in attrs.get("source_annotations", []):
            if ann.patient_id != patient.id:
                raise serializers.ValidationError(
                    {"source_annotations": f"Annotation {ann.pk} not in this patient."},
                )
            if not ann.event_date:
                raise serializers.ValidationError(
                    {"source_annotations": f"Annotation {ann.pk} has no event_date."},
                )

        for lv in attrs.get("event_label_values", []):
            if lv.category.project_id != patient.project_id:
                raise serializers.ValidationError(
                    {"event_label_values": f"LabelValue {lv.pk} not in this project."},
                )

        event_date_to_window = build_annotation_window_start_map(patient.id)

        window_starts = {
            event_date_to_window.get(ann.event_date, ann.event_date) for ann in source_annotations
        }

        if len(window_starts) != 1:
            raise serializers.ValidationError(
                {
                    "source_annotations": (
                        "All source annotations must belong to the same annotation window."
                    ),
                },
            )

        return attrs


    @staticmethod
    def _sync_notes_from_sources(instance: Adjudication):
        note_ids = (
            PatientAnnotation.objects
            .filter(
                id__in=instance.source_annotations.values_list("id", flat=True),
                notes__id__isnull=False,
            )
            .values_list("notes__id", flat=True)
            .distinct()
        )
        instance.notes.set(note_ids)



    @transaction.atomic
    def create(self, validated_data):
        """Create an adjudication and synchronize its related objects."""
        src_anns = validated_data.pop("source_annotations", None)
        ev_vals = validated_data.pop("event_label_values", None)

        instance = super().create(validated_data)

        if ev_vals is not None:
            instance.event_label_values.set(ev_vals)
        if src_anns is not None:
            instance.source_annotations.set(src_anns)
            self._sync_notes_from_sources(instance)

        return instance




    @transaction.atomic
    def update(self, instance, validated_data):
        """Update an adjudication and synchronize its related objects."""
        src_anns = validated_data.pop("source_annotations", None)
        ev_vals = validated_data.pop("event_label_values", None)

        instance = super().update(instance, validated_data)

        if ev_vals is not None:
            instance.event_label_values.set(ev_vals)
        if src_anns is not None:
            instance.source_annotations.set(src_anns)
            self._sync_notes_from_sources(instance)

        return instance






class NERAnnotationSerializer(serializers.ModelSerializer):
    annotator = serializers.HiddenField(default=serializers.CurrentUserDefault())
    annotator_name = serializers.CharField(source="annotator.username", read_only=True)
    spans = serializers.JSONField()



    class Meta:
        model = NERAnnotation
        fields = ["pk", "patient", "note", "annotator", "annotator_name", "spans"]
        read_only_fields = ["pk", "annotator_name"]




    def validate(self, attrs):
        attrs = super().validate(attrs)
        instance = getattr(self, "instance", None)

        note = attrs.get("note") or (instance.note if instance else None)
        patient = attrs.get("patient") or (instance.patient if instance else None)

        if patient is None:
            raise serializers.ValidationError(
                {"patient": "patient is required."},
            )
        if note is None:
            raise serializers.ValidationError(
                {"note": "note is required."},
            )

        # Should have correct annotation type
        if patient.project.annotation_type != Project.NER:
            raise serializers.ValidationError(
                {"Error": "This is not an NER project."},
            )

        # Note should belong to the same patient
        if note.patient_id != patient.id:
            raise serializers.ValidationError(
                {
                    "Error": f"Note {note.pk} does not belong to this patient.",
                },
            )

        return attrs






class ExportAnnotationSerializer(serializers.Serializer):
    project = serializers.PrimaryKeyRelatedField(
        queryset=Project.objects.all(),
    )
    patients = serializers.PrimaryKeyRelatedField(
        queryset=Patient.objects.all(),
        many=True,
        allow_empty=False,
    )
    annotators = serializers.PrimaryKeyRelatedField(
        queryset=User.objects.all(),
        many=True,
        allow_empty=False,
    )
    export_type = serializers.ChoiceField(
        choices=(
            ("patient_level", "Patient-level"),
            ("note_level", "Note-level"),
        ),
    )
    export_file_format = serializers.ChoiceField(
        choices=(
            ("csv", "CSV"),
            ("xlsx", "Excel"),
            ("dta", "Stata"),
            ("parquet", "Parquet"),
            ("feather", "Feather"),
        ),
    )



    def validate(self, attrs):
        project = attrs["project"]
        patients = attrs["patients"]

        if any(p.project_id != project.id for p in patients):
            raise serializers.ValidationError(
                {
                    "patients": "All patients must belong to the selected project.",
                },
            )

        return attrs






class PatientHistoryChangeSerializer(serializers.Serializer):
    field = serializers.CharField()
    operation = serializers.CharField()
    old_value = serializers.CharField(allow_null=True, required=False)
    new_value = serializers.CharField(allow_null=True, required=False)





class PatientHistorySerializer(serializers.Serializer):
    log_id = serializers.IntegerField()
    model_type = serializers.CharField()
    object_id = serializers.IntegerField()
    timestamp = serializers.DateTimeField()
    actor = serializers.CharField(allow_null=True)
    changes = PatientHistoryChangeSerializer(many=True)