from pydantic import BaseModel, Field


class Size(BaseModel):
    width: int
    height: int
    depth: int


class Source(BaseModel):
    database: str


class BndBox(BaseModel):
    xmin: int | float
    xmax: int | float
    ymin: int | float
    ymax: int | float


class Object(BaseModel):
    name: str
    pose: str | None = Field(default=None)
    truncated: int | None = Field(default=None)
    difficult: int | None = Field(default=None)
    occluded: int | None = Field(default=None)
    bndbox: BndBox | None = Field(default=None)
    polygon: dict[str, int | float] | None = Field(default=None)

    def is_rle(self) -> bool:
        return False

    def is_polygon(self) -> bool:
        return self.polygon is not None

    def is_rectangle(self) -> bool:
        return self.polygon is None and self.bndbox is not None

    def polygon_to_list_coordinates(self) -> list[list[int]]:
        if not self.is_polygon():
            raise ValueError("Not a polygon")

        coords = []
        for i in range(1, 1 + len(self.polygon) // 2):
            x = "x" + str(i)
            y = "y" + str(i)
            if x not in self.polygon or y not in self.polygon:
                raise ValueError(f"{x} or {y} not found in this polygon.")

            coords.append([int(self.polygon[x]), int(self.polygon[y])])

        return coords


class Annotation(BaseModel):
    filename: str
    object: Object | list[Object]
    path: str | None = Field(default=None)
    folder: str | None = Field(default=None)
    source: Source | None = Field(default=None)
    size: Size | None = Field(default=None)
    segmented: int | None = Field(default=None)


class PascalVOCFile(BaseModel):
    annotation: Annotation
