"""Parser for CVAT Images 1.1 keypoint annotations

see CVAT Documentation for format information: https://opencv.github.io/cvat/docs/manual/advanced/xml_format/

This file was created by labeling the example images using the flow that is described in the repo.
Then this xml was converted to JSON using xmltodict
Then an initial version of the parser was created using https://pydantic-docs.helpmanual.io/datamodel_code_generator/:
`datamodel-codegen --input annotations.json --input-file-type json --output cvat_keypoints_parser.py --class-name CVATKeypointsParser`

And the parser was then further finetuned.

Note that the attributes cannot have _ as prefix for Pydantic.
"""

# generated by datamodel-codegen:
#   filename:  annotations.json
#   timestamp: 2022-08-24T12:15:52+00:00
# then further changed by @tlips

from __future__ import annotations

from typing import Any, List, Optional, Union

from pydantic import BaseModel, field_validator


class Segment(BaseModel):
    id: str
    start: str
    stop: str
    url: str


class Segments(BaseModel):
    segment: Segment


class Owner(BaseModel):
    username: str
    email: str


class LabelItem(BaseModel):
    name: str
    color: str
    type: str
    attributes: Any


class Labels(BaseModel):
    label: Union[List[LabelItem], LabelItem]


class Job(BaseModel):
    id: str
    name: Optional[str] = None
    size: str
    mode: str
    overlap: str
    bugtracker: Any
    created: str
    updated: str
    subset: str
    start_frame: str
    stop_frame: str
    frame_filter: Any
    segments: Segments
    owner: Owner
    assignee: Any
    labels: Labels


class Meta(BaseModel):
    job: Optional[Job] = None
    task: Optional[Job] = None
    dumped: str

    def get_job_or_task(self) -> Job:
        if self.job:
            return self.job
        elif self.task:
            return self.task
        else:
            raise ValueError("No task or job found in meta. This should not happen")


class AnnotationItem(BaseModel):
    label: str
    source: str
    z_order: str
    occluded: str
    group_id: Optional[str] = "1"  # set default group id to 1.


class Point(AnnotationItem):
    points: str

    @field_validator("points")
    @classmethod
    def has_two_coordinates(cls, v: str) -> str:
        assert len(v.split(",")) == 2, "each point must be a single 2D coordinate for the AIRO flow."
        return v


class Polygon(AnnotationItem):
    points: str


class Box(AnnotationItem):
    xtl: str
    ytl: str
    xbr: str
    ybr: str


class ImageItem(BaseModel):
    id: str
    name: str
    width: str
    height: str
    points: Optional[Union[List[Point], Point]] = None
    polygon: Optional[Union[List[Polygon], Polygon]] = None
    box: Optional[Union[List[Box], Box]] = None


class Annotations(BaseModel):
    version: str
    meta: Meta
    image: List[ImageItem]


class CVATImagesParser(BaseModel):
    annotations: Annotations


if __name__ == "__main__":
    import json
    import pathlib

    path = pathlib.Path(__file__).parent.absolute()
    cvat_dict = json.load(open(str(path.parent / "cvat_labeling" / "example" / "annotations.json")))
    cvat_keypoints_parser = CVATImagesParser(**cvat_dict)
