from .boundingbox import BoundingBox, BoxFormat
from .errors import ParsingError, FileParsingError
from .atomic import open_atomic
from .file_utils import PathLike

from pathlib import Path
from typing import Mapping, Optional, Union
import xml.etree.ElementTree as et
import json
from warnings import warn


class Annotation:
    """Stores bounding box annotations for one image.

    The image should be uniquely identified by the `image_id` str. The
    image size is necessary for some operations."""

    def __init__(
        self,
        image_id: str,
        image_size: Optional["tuple[int, int]"] = None,
        boxes: Optional["list[BoundingBox]"] = None,
    ) -> None:
        if image_size is not None:
            img_w, img_h = image_size
            assert img_w > 0 and img_h > 0
            assert int(img_w) == img_w and int(img_h) == img_h

        self.image_id = image_id
        self.image_size = image_size
        self.boxes = boxes or []

    @property
    def image_width(self) -> Optional[int]:
        return self.image_size[0] if self.image_size is not None else None

    @property
    def image_height(self) -> Optional[int]:
        return self.image_size[1] if self.image_size is not None else None

    def add(self, box: BoundingBox):
        self.boxes.append(box)

    def map_labels(self, mapping: Mapping[str, str]) -> "Annotation":
        """Change all the bounding box annotation labels according to
        the provided mapping (act as a translation)."""
        for box in self.boxes:
            box.label = mapping[box.label]
        return self

    def _labels(self) -> "set[str]":
        return {b.label for b in self.boxes}

    @staticmethod
    def from_txt(
        file_path: PathLike,
        *,
        image_id: Optional[str] = None,
        image_extension: str = ".jpg",
        box_format: BoxFormat = BoxFormat.LTRB,
        relative: bool = False,
        image_size: Optional["tuple[int, int]"] = None,
        separator: str = " ",
        conf_last: bool = False,
    ) -> "Annotation":
        path = Path(file_path).expanduser().resolve()

        if image_id is None:
            assert image_extension.startswith(".")
            image_id = path.with_suffix(image_extension).name

        try:
            lines = path.read_text().splitlines()
        except OSError:
            raise FileParsingError(path, reason="cannot read file")

        try:
            boxes = [
                BoundingBox.from_txt(
                    l,
                    box_format=box_format,
                    relative=relative,
                    image_size=image_size,
                    separator=separator,
                    conf_last=conf_last,
                )
                for l in lines
            ]
        except ParsingError as e:
            raise FileParsingError(path, e.reason)

        return Annotation(image_id, image_size, boxes)

    @staticmethod
    def _from_yolo(
        file_path: PathLike,
        *,
        image_size: "tuple[int, int]",
        image_id: Optional[str] = None,
        image_extension: str = ".jpg",
        conf_last: bool = False,
    ) -> "Annotation":
        return Annotation.from_txt(
            file_path,
            image_id=image_id,
            image_extension=image_extension,
            box_format=BoxFormat.XYWH,
            relative=True,
            image_size=image_size,
            separator=" ",
            conf_last=conf_last,
        )

    @staticmethod
    def from_yolo(
        file_path: PathLike,
        *,
        image_size: "tuple[int, int]",
        image_id: Optional[str] = None,
        image_extension: str = ".jpg",
        conf_last: bool = False,
    ) -> "Annotation":
        warn(
            "'from_yolo' is deprecated. Please use `from_yolo_darknet` or `from_yolo_v5`",
            category=DeprecationWarning,
            stacklevel=2,
        )

        return Annotation._from_yolo(
            file_path,
            image_size=image_size,
            image_id=image_id,
            image_extension=image_extension,
            conf_last=conf_last,
        )

    @staticmethod
    def from_yolo_darknet(
        file_path: PathLike,
        *,
        image_size: "tuple[int, int]",
        image_id: Optional[str] = None,
        image_extension: str = ".jpg",
    ) -> "Annotation":
        return Annotation._from_yolo(
            file_path,
            image_size=image_size,
            image_id=image_id,
            image_extension=image_extension,
            conf_last=False,
        )

    @staticmethod
    def from_yolo_v5(
        file_path: PathLike,
        *,
        image_size: "tuple[int, int]",
        image_id: Optional[str] = None,
        image_extension: str = ".jpg",
    ) -> "Annotation":
        return Annotation._from_yolo(
            file_path,
            image_size=image_size,
            image_id=image_id,
            image_extension=image_extension,
            conf_last=True,
        )

    @staticmethod
    def from_yolo_v7(
        file_path: PathLike,
        *,
        image_size: "tuple[int, int]",
        image_id: Optional[str] = None,
        image_extension: str = ".jpg",
    ) -> "Annotation":
        return Annotation.from_yolo_v5(
            file_path,
            image_size=image_size,
            image_id=image_id,
            image_extension=image_extension,
        )

    @staticmethod
    def from_xml(file_path: PathLike) -> "Annotation":
        path = Path(file_path).expanduser().resolve()

        try:
            with path.open() as f:
                root = et.parse(f).getroot()
        except (OSError, et.ParseError):
            raise ParsingError("Syntax error in imagenet annotation file.")

        image_id = root.findtext("filename")
        size_node = root.find("size")

        if (image_id is None) or (size_node is None):
            raise ParsingError("Syntax error in imagenet annotation file.")

        width = size_node.findtext("width")
        height = size_node.findtext("height")

        if (width is None) or (height is None):
            raise ParsingError("Syntax error in imagenet annotation file.")

        try:
            image_size = int(width), int(height)
        except ValueError:
            raise ParsingError("Syntax error in imagenet annotation file.")

        boxes = [BoundingBox.from_xml(n) for n in root.iter("object")]

        return Annotation(image_id, image_size, boxes)

    @staticmethod
    def from_pascal_voc(file_path: PathLike) -> "Annotation":
        return Annotation.from_xml(file_path)

    @staticmethod
    def from_imagenet(file_path: PathLike) -> "Annotation":
        return Annotation.from_xml(file_path)

    @staticmethod
    def from_labelme(file_path: PathLike) -> "Annotation":
        path = Path(file_path).expanduser().resolve()
        try:
            with path.open() as f:
                content = json.load(f)
                if "imageData" in content:
                    del content["imageData"]
        except (OSError, json.JSONDecodeError):
            raise ParsingError("Syntax error in labelme annotation file.")

        try:
            image_id = str(content["imagePath"])
            width = int(content["imageWidth"])
            height = int(content["imageHeight"])
            boxes = [
                BoundingBox.from_labelme(n)
                for n in content["shapes"]
                if n["shape_type"] == "rectangle"
            ]
        except (KeyError, ValueError):
            raise ParsingError("Syntax error in labelme annotation file.")

        return Annotation(image_id, image_size=(width, height), boxes=boxes)

    @staticmethod
    def _from_coco_partial(node: dict) -> "Annotation":
        try:
            image_id = str(node["file_name"])
            image_size = int(node["width"]), int(node["height"])
        except (ValueError, KeyError):
            raise ParsingError("Syntax error in COCO annotation file.")

        return Annotation(image_id, image_size)

    @staticmethod
    def from_cvat(node: et.Element) -> "Annotation":
        image_id = node.get("name")
        width, height = node.get("width"), node.get("height")

        if (image_id is None) or (width is None) or (height is None):
            raise ParsingError("Syntax error in CVAT annotation file.")

        try:
            img_size = int(width), int(height)
        except ValueError:
            raise ParsingError("Syntax error in CVAT annotation file.")

        boxes = [BoundingBox.from_cvat(n) for n in node.iter("box")]

        return Annotation(image_id, image_size=img_size, boxes=boxes)

    @staticmethod
    def from_via_json(
        annotation: dict,
        *,
        label_key: str = "label_id",
        confidence_key: str = "confidence",
        image_size: "Optional[tuple[int, int]]" = None,
    ) -> "Annotation":
        try:
            filename = annotation["filename"]
            regions = annotation["regions"]

            bboxes = [
                BoundingBox.from_via_json(
                    region, label_key=label_key, confidence_key=confidence_key
                )
                for region in regions
                if region["shape_attributes"]["name"] == "rect"
            ]
        except KeyError:
            raise ParsingError("Syntax error in VIA JSON annotation file.")

        return Annotation(
            image_id=filename,
            image_size=image_size,
            boxes=bboxes,
        )

    def to_txt(
        self,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        box_format: BoxFormat = BoxFormat.LTRB,
        relative=False,
        image_size: Optional["tuple[int, int]"] = None,
        separator: str = " ",
        conf_last: bool = False,
    ) -> str:
        image_size = image_size or self.image_size

        return "\n".join(
            box.to_txt(
                label_to_id=label_to_id,
                box_format=box_format,
                relative=relative,
                image_size=image_size,
                separator=separator,
                conf_last=conf_last,
            )
            for box in self.boxes
        )

    def _to_yolo(
        self,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
        conf_last: bool = False,
    ) -> str:
        image_size = image_size or self.image_size

        if image_size is None:
            raise ValueError(
                "Either `image_size` shoud be provided as argument or stored in the Annotation object for conversion to YOLO format."
            )

        return "\n".join(
            box.to_yolo(
                image_size=image_size, label_to_id=label_to_id, conf_last=conf_last
            )
            for box in self.boxes
        )

    def to_yolo(
        self,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
        conf_last: bool = False,
    ) -> str:
        warn(
            "'to_yolo' is deprecated. Please use `to_yolo_darknet` or `to_yolo_v5`",
            category=DeprecationWarning,
            stacklevel=2,
        )

        return self._to_yolo(
            label_to_id=label_to_id, image_size=image_size, conf_last=conf_last
        )

    def to_yolo_darknet(
        self,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
    ) -> str:
        return self._to_yolo(
            label_to_id=label_to_id, image_size=image_size, conf_last=False
        )

    def to_yolo_v5(
        self,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
    ) -> str:
        return self._to_yolo(
            label_to_id=label_to_id, image_size=image_size, conf_last=True
        )

    def to_yolo_v7(
        self,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
    ) -> str:
        return self.to_yolo_v5(label_to_id=label_to_id, image_size=image_size)

    def save_txt(
        self,
        path: PathLike,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        box_format: BoxFormat = BoxFormat.LTRB,
        relative: bool = False,
        image_size: Optional["tuple[int, int]"] = None,
        separator: str = " ",
        conf_last: bool = False,
    ):
        content = self.to_txt(
            label_to_id=label_to_id,
            box_format=box_format,
            relative=relative,
            image_size=image_size,
            separator=separator,
            conf_last=conf_last,
        )

        with open_atomic(path, "w") as f:
            f.write(content)

    def _save_yolo(
        self,
        path: PathLike,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
        conf_last: bool = False,
    ):
        content = self._to_yolo(
            label_to_id=label_to_id, image_size=image_size, conf_last=conf_last
        )

        with open_atomic(path, "w") as f:
            f.write(content)

    def save_yolo(
        self,
        path: PathLike,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
        conf_last: bool = False,
    ):
        warn(
            "'save_yolo' is deprecated. Please use `save_yolo_darknet` or `save_yolo_v5`",
            category=DeprecationWarning,
            stacklevel=2,
        )

        self._save_yolo(
            path, label_to_id=label_to_id, image_size=image_size, conf_last=conf_last
        )

    def save_yolo_darknet(
        self,
        path: PathLike,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
    ):
        self._save_yolo(
            path, label_to_id=label_to_id, image_size=image_size, conf_last=False
        )

    def save_yolo_v5(
        self,
        path: PathLike,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
    ):
        self._save_yolo(
            path, label_to_id=label_to_id, image_size=image_size, conf_last=True
        )

    def save_yolo_v7(
        self,
        path: PathLike,
        *,
        label_to_id: Optional[Mapping[str, Union[int, str]]] = None,
        image_size: Optional["tuple[int, int]"] = None,
    ):
        self.save_yolo_v5(path, label_to_id=label_to_id, image_size=image_size)

    def to_labelme(self, *, image_size: Optional["tuple[int, int]"] = None) -> dict:
        image_size = image_size or self.image_size
        assert (
            image_size is not None
        ), "An image size should be provided either by argument or by `self.image_size`."

        return {
            "imagePath": self.image_id,
            "imageWidth": image_size[0],
            "imageHeight": image_size[1],
            "imageData": None,
            "shapes": [b.to_labelme() for b in self.boxes],
        }

    def save_labelme(
        self, path: PathLike, *, image_size: Optional["tuple[int, int]"] = None
    ):
        content = self.to_labelme(image_size=image_size)
        with open_atomic(path, "w") as f:
            json.dump(content, fp=f, allow_nan=False)

    def to_xml(self, *, image_size: Optional["tuple[int, int]"] = None) -> et.Element:
        image_size = image_size or self.image_size
        assert (
            image_size is not None
        ), "An image size should be provided either by argument or by `self.image_size`."

        ann_node = et.Element("annotation")
        et.SubElement(ann_node, "filename").text = self.image_id

        size_node = et.SubElement(ann_node, "size")
        et.SubElement(size_node, "width").text = f"{image_size[0]}"
        et.SubElement(size_node, "height").text = f"{image_size[1]}"

        for box in self.boxes:
            ann_node.append(box.to_xml())

        return ann_node

    def to_pascal_voc(
        self, *, image_size: Optional["tuple[int, int]"] = None
    ) -> et.Element:
        return self.to_xml(image_size=image_size)

    def to_imagenet(
        self, *, image_size: Optional["tuple[int, int]"] = None
    ) -> et.Element:
        return self.to_xml(image_size=image_size)

    def save_xml(
        self, path: PathLike, *, image_size: Optional["tuple[int, int]"] = None
    ):
        content = self.to_xml(image_size=image_size)
        content = et.tostring(content, encoding="unicode")

        with open_atomic(path, "w") as f:
            f.write(content)

    def save_pascal_voc(
        self, path: PathLike, *, image_size: Optional["tuple[int, int]"] = None
    ):
        self.save_xml(path, image_size=image_size)

    def save_imagenet(
        self, path: PathLike, *, image_size: Optional["tuple[int, int]"] = None
    ):
        self.save_xml(path, image_size=image_size)

    def to_cvat(self, *, image_size: Optional["tuple[int, int]"] = None) -> et.Element:
        image_size = image_size or self.image_size
        assert (
            image_size is not None
        ), "An image size should be provided either by argument or by `self.image_size`."

        img_node = et.Element(
            "image",
            attrib={
                "name": self.image_id,
                "width": f"{image_size[0]}",
                "height": f"{image_size[1]}",
            },
        )

        img_node.extend(box.to_cvat() for box in self.boxes)

        return img_node

    def to_via_json(
        self,
        *,
        image_folder: PathLike,
        label_key: str = "label_id",
        confidence_key: str = "confidence",
    ) -> dict:
        path = Path(image_folder).expanduser().resolve()

        assert path.is_dir()

        image_id = self.image_id
        image_path = path / image_id
        file_size = image_path.stat().st_size

        regions = [
            box.to_via_json(label_key=label_key, confidence_key=confidence_key)
            for box in self.boxes
        ]

        return {"filename": image_id, "size": file_size, "regions": regions}

    def __repr__(self) -> str:
        return f"Annotation(image_id: {self.image_id}, image_size: {self.image_size}, boxes: {self.boxes})"
