import hashlib
from io import BytesIO
import logging
import os
from typing import Any, cast, Dict, List, Optional, Sequence, Type, TYPE_CHECKING, Union

from pkg_resources import parse_version
import wandb
from wandb import util

from ._private import MEDIA_TMP
from .base_types.media import BatchableMedia, Media
from .helper_types.bounding_boxes_2d import BoundingBoxes2D
from .helper_types.classes import Classes
from .helper_types.image_mask import ImageMask

if TYPE_CHECKING:  # pragma: no cover
    import matplotlib  # type: ignore
    import numpy as np  # type: ignore
    import PIL  # type: ignore
    import torch  # type: ignore
    from wandb.apis.public import Artifact as PublicArtifact

    from ..wandb_artifacts import Artifact as LocalArtifact
    from ..wandb_run import Run as LocalRun

    ImageDataType = Union[
        "matplotlib.artist.Artist", "PIL.Image", "TorchTensorType", "np.ndarray"
    ]
    ImageDataOrPathType = Union[str, "Image", ImageDataType]
    TorchTensorType = Union["torch.Tensor", "torch.Variable"]


def _server_accepts_image_filenames() -> bool:
    # Newer versions of wandb accept large image filenames arrays
    # but older versions would have issues with this.
    max_cli_version = util._get_max_cli_version()
    if max_cli_version is None:
        return False
    return parse_version("0.12.10") <= parse_version(max_cli_version)


class Image(BatchableMedia):
    """Format images for logging to W&B.

    Arguments:
        data_or_path: (numpy array, string, io) Accepts numpy array of
            image data, or a PIL image. The class attempts to infer
            the data format and converts it.
        mode: (string) The PIL mode for an image. Most common are "L", "RGB",
            "RGBA". Full explanation at https://pillow.readthedocs.io/en/4.2.x/handbook/concepts.html#concept-modes.
        caption: (string) Label for display of image.

    Examples:
        ### Create a wandb.Image from a numpy array
        <!--yeadoc-test:log-image-numpy->
        ```python
        import numpy as np
        import wandb

        wandb.init()
        examples = []
        for i in range(3):
            pixels = np.random.randint(low=0, high=256, size=(100, 100, 3))
            image = wandb.Image(pixels, caption=f"random field {i}")
            examples.append(image)
        wandb.log({"examples": examples})
        ```

        ### Create a wandb.Image from a PILImage
        <!--yeadoc-test:log-image-pil->
        ```python
        import numpy as np
        from PIL import Image as PILImage
        import wandb

        wandb.init()
        examples = []
        for i in range(3):
            pixels = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
            pil_image = PILImage.fromarray(pixels, mode="RGB")
            image = wandb.Image(pil_image, caption=f"random field {i}")
            examples.append(image)
        wandb.log({"examples": examples})
        ```
    """

    MAX_ITEMS = 108

    # PIL limit
    MAX_DIMENSION = 65500

    _log_type = "image-file"

    format: Optional[str]
    _grouping: Optional[int]
    _caption: Optional[str]
    _width: Optional[int]
    _height: Optional[int]
    _image: Optional["PIL.Image"]
    _classes: Optional["Classes"]
    _boxes: Optional[Dict[str, "BoundingBoxes2D"]]
    _masks: Optional[Dict[str, "ImageMask"]]

    def __init__(
        self,
        data_or_path: "ImageDataOrPathType",
        mode: Optional[str] = None,
        caption: Optional[str] = None,
        grouping: Optional[int] = None,
        classes: Optional[Union["Classes", Sequence[dict]]] = None,
        boxes: Optional[Union[Dict[str, "BoundingBoxes2D"], Dict[str, dict]]] = None,
        masks: Optional[Union[Dict[str, "ImageMask"], Dict[str, dict]]] = None,
    ) -> None:
        super(Image, self).__init__()
        # TODO: We should remove grouping, it's a terrible name and I don't
        # think anyone uses it.

        self._grouping = None
        self._caption = None
        self._width = None
        self._height = None
        self._image = None
        self._classes = None
        self._boxes = None
        self._masks = None

        # Allows the user to pass an Image object as the first parameter and have a perfect copy,
        # only overriding additional metdata passed in. If this pattern is compelling, we can generalize.
        if isinstance(data_or_path, Image):
            self._initialize_from_wbimage(data_or_path)
        elif isinstance(data_or_path, str):
            self._initialize_from_path(data_or_path)
        else:
            self._initialize_from_data(data_or_path, mode)

        self._set_initialization_meta(grouping, caption, classes, boxes, masks)

    def _set_initialization_meta(
        self,
        grouping: Optional[int] = None,
        caption: Optional[str] = None,
        classes: Optional[Union["Classes", Sequence[dict]]] = None,
        boxes: Optional[Union[Dict[str, "BoundingBoxes2D"], Dict[str, dict]]] = None,
        masks: Optional[Union[Dict[str, "ImageMask"], Dict[str, dict]]] = None,
    ) -> None:
        if grouping is not None:
            self._grouping = grouping

        if caption is not None:
            self._caption = caption

        total_classes = {}

        if boxes:
            if not isinstance(boxes, dict):
                raise ValueError('Images "boxes" argument must be a dictionary')
            boxes_final: Dict[str, BoundingBoxes2D] = {}
            for key in boxes:
                box_item = boxes[key]
                if isinstance(box_item, BoundingBoxes2D):
                    boxes_final[key] = box_item
                elif isinstance(box_item, dict):
                    # TODO: Consider injecting top-level classes if user-provided is empty
                    boxes_final[key] = BoundingBoxes2D(box_item, key)
                total_classes.update(boxes_final[key]._class_labels)
            self._boxes = boxes_final

        if masks:
            if not isinstance(masks, dict):
                raise ValueError('Images "masks" argument must be a dictionary')
            masks_final: Dict[str, ImageMask] = {}
            for key in masks:
                mask_item = masks[key]
                if isinstance(mask_item, ImageMask):
                    masks_final[key] = mask_item
                elif isinstance(mask_item, dict):
                    # TODO: Consider injecting top-level classes if user-provided is empty
                    masks_final[key] = ImageMask(mask_item, key)
                if hasattr(masks_final[key], "_val"):
                    total_classes.update(masks_final[key]._val["class_labels"])
            self._masks = masks_final

        if classes is not None:
            if isinstance(classes, Classes):
                total_classes.update(
                    {val["id"]: val["name"] for val in classes._class_set}
                )
            else:
                total_classes.update({val["id"]: val["name"] for val in classes})

        if len(total_classes.keys()) > 0:
            self._classes = Classes(
                [
                    {"id": key, "name": total_classes[key]}
                    for key in total_classes.keys()
                ]
            )
        self._width, self._height = self.image.size  # type: ignore
        self._free_ram()

    def _initialize_from_wbimage(self, wbimage: "Image") -> None:
        self._grouping = wbimage._grouping
        self._caption = wbimage._caption
        self._width = wbimage._width
        self._height = wbimage._height
        self._image = wbimage._image
        self._classes = wbimage._classes
        self._path = wbimage._path
        self._is_tmp = wbimage._is_tmp
        self._extension = wbimage._extension
        self._sha256 = wbimage._sha256
        self._size = wbimage._size
        self.format = wbimage.format
        self._artifact_source = wbimage._artifact_source
        self._artifact_target = wbimage._artifact_target

        # We do not want to implicitly copy boxes or masks, just the image-related data.
        # self._boxes = wbimage._boxes
        # self._masks = wbimage._masks

    def _initialize_from_path(self, path: str) -> None:
        pil_image = util.get_module(
            "PIL.Image",
            required='wandb.Image needs the PIL package. To get it, run "pip install pillow".',
        )
        self._set_file(path, is_tmp=False)
        self._image = pil_image.open(path)
        self._image.load()
        ext = os.path.splitext(path)[1][1:]
        self.format = ext

    def _initialize_from_data(self, data: "ImageDataType", mode: str = None,) -> None:
        pil_image = util.get_module(
            "PIL.Image",
            required='wandb.Image needs the PIL package. To get it, run "pip install pillow".',
        )
        if util.is_matplotlib_typename(util.get_full_typename(data)):
            buf = BytesIO()
            util.ensure_matplotlib_figure(data).savefig(buf)
            self._image = pil_image.open(buf)
        elif isinstance(data, pil_image.Image):
            self._image = data
        elif util.is_pytorch_tensor_typename(util.get_full_typename(data)):
            vis_util = util.get_module(
                "torchvision.utils", "torchvision is required to render images"
            )
            if hasattr(data, "requires_grad") and data.requires_grad:
                data = data.detach()
            data = vis_util.make_grid(data, normalize=True)
            self._image = pil_image.fromarray(
                data.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
            )
        else:
            if hasattr(data, "numpy"):  # TF data eager tensors
                data = data.numpy()
            if data.ndim > 2:
                data = data.squeeze()  # get rid of trivial dimensions as a convenience
            self._image = pil_image.fromarray(
                self.to_uint8(data), mode=mode or self.guess_mode(data)
            )

        tmp_path = os.path.join(MEDIA_TMP.name, str(util.generate_id()) + ".png")
        self.format = "png"
        self._image.save(tmp_path, transparency=None)
        self._set_file(tmp_path, is_tmp=True)

    @classmethod
    def from_json(
        cls: Type["Image"], json_obj: dict, source_artifact: "PublicArtifact"
    ) -> "Image":
        classes = None
        if json_obj.get("classes") is not None:
            classes = source_artifact.get(json_obj["classes"]["path"])

        masks = json_obj.get("masks")
        _masks: Optional[Dict[str, ImageMask]] = None
        if masks:
            _masks = {}
            for key in masks:
                _masks[key] = ImageMask.from_json(masks[key], source_artifact)
                _masks[key]._set_artifact_source(source_artifact)
                _masks[key]._key = key

        boxes = json_obj.get("boxes")
        _boxes: Optional[Dict[str, BoundingBoxes2D]] = None
        if boxes:
            _boxes = {}
            for key in boxes:
                _boxes[key] = BoundingBoxes2D.from_json(boxes[key], source_artifact)
                _boxes[key]._key = key

        return cls(
            source_artifact.get_path(json_obj["path"]).download(),
            caption=json_obj.get("caption"),
            grouping=json_obj.get("grouping"),
            classes=classes,
            boxes=_boxes,
            masks=_masks,
        )

    @classmethod
    def get_media_subdir(cls: Type["Image"]) -> str:
        return os.path.join("media", "images")

    def bind_to_run(
        self,
        run: "LocalRun",
        key: Union[int, str],
        step: Union[int, str],
        id_: Optional[Union[int, str]] = None,
        ignore_copy_err: Optional[bool] = None,
    ) -> None:
        super().bind_to_run(run, key, step, id_, ignore_copy_err=ignore_copy_err)
        if self._boxes is not None:
            for i, k in enumerate(self._boxes):
                id_ = "{}{}".format(id_, i) if id_ is not None else None
                self._boxes[k].bind_to_run(
                    run, key, step, id_, ignore_copy_err=ignore_copy_err
                )

        if self._masks is not None:
            for i, k in enumerate(self._masks):
                id_ = "{}{}".format(id_, i) if id_ is not None else None
                self._masks[k].bind_to_run(
                    run, key, step, id_, ignore_copy_err=ignore_copy_err
                )

    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
        json_dict = super(Image, self).to_json(run_or_artifact)
        json_dict["_type"] = Image._log_type
        json_dict["format"] = self.format

        if self._width is not None:
            json_dict["width"] = self._width
        if self._height is not None:
            json_dict["height"] = self._height
        if self._grouping:
            json_dict["grouping"] = self._grouping
        if self._caption:
            json_dict["caption"] = self._caption

        if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_artifacts.Artifact):
            artifact = run_or_artifact
            if (
                self._masks is not None or self._boxes is not None
            ) and self._classes is None:
                raise ValueError(
                    "classes must be passed to wandb.Image which have masks or bounding boxes when adding to artifacts"
                )

            if self._classes is not None:
                class_id = hashlib.md5(
                    str(self._classes._class_set).encode("utf-8")
                ).hexdigest()
                class_name = os.path.join("media", "classes", class_id + "_cls",)
                classes_entry = artifact.add(self._classes, class_name)
                json_dict["classes"] = {
                    "type": "classes-file",
                    "path": classes_entry.path,
                    "digest": classes_entry.digest,
                }

        elif not isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
            raise ValueError("to_json accepts wandb_run.Run or wandb_artifact.Artifact")

        if self._boxes:
            json_dict["boxes"] = {
                k: box.to_json(run_or_artifact) for (k, box) in self._boxes.items()
            }
        if self._masks:
            json_dict["masks"] = {
                k: mask.to_json(run_or_artifact) for (k, mask) in self._masks.items()
            }
        return json_dict

    def guess_mode(self, data: "np.ndarray") -> str:
        """
        Guess what type of image the np.array is representing
        """
        # TODO: do we want to support dimensions being at the beginning of the array?
        if data.ndim == 2:
            return "L"
        elif data.shape[-1] == 3:
            return "RGB"
        elif data.shape[-1] == 4:
            return "RGBA"
        else:
            raise ValueError(
                "Un-supported shape for image conversion %s" % list(data.shape)
            )

    @classmethod
    def to_uint8(cls, data: "np.ndarray") -> "np.ndarray":
        """
        Converts floating point image on the range [0,1] and integer images
        on the range [0,255] to uint8, clipping if necessary.
        """
        np = util.get_module(
            "numpy",
            required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy",
        )

        # I think it's better to check the image range vs the data type, since many
        # image libraries will return floats between 0 and 255

        # some images have range -1...1 or 0-1
        dmin = np.min(data)
        if dmin < 0:
            data = (data - np.min(data)) / np.ptp(data)
        if np.max(data) <= 1.0:
            data = (data * 255).astype(np.int32)

        # assert issubclass(data.dtype.type, np.integer), 'Illegal image format.'
        return data.clip(0, 255).astype(np.uint8)

    @classmethod
    def seq_to_json(
        cls: Type["Image"],
        seq: Sequence["BatchableMedia"],
        run: "LocalRun",
        key: str,
        step: Union[int, str],
    ) -> dict:
        """
        Combines a list of images into a meta dictionary object describing the child images.
        """
        if TYPE_CHECKING:
            seq = cast(Sequence["Image"], seq)

        jsons = [obj.to_json(run) for obj in seq]

        media_dir = cls.get_media_subdir()

        for obj in jsons:
            expected = util.to_forward_slash_path(media_dir)
            if not obj["path"].startswith(expected):
                raise ValueError(
                    "Files in an array of Image's must be in the {} directory, not {}".format(
                        cls.get_media_subdir(), obj["path"]
                    )
                )

        num_images_to_log = len(seq)
        width, height = seq[0].image.size  # type: ignore
        format = jsons[0]["format"]

        def size_equals_image(image: "Image") -> bool:
            img_width, img_height = image.image.size  # type: ignore
            return img_width == width and img_height == height  # type: ignore

        sizes_match = all(size_equals_image(img) for img in seq)
        if not sizes_match:
            logging.warning(
                "Images sizes do not match. This will causes images to be display incorrectly in the UI."
            )

        meta = {
            "_type": "images/separated",
            "width": width,
            "height": height,
            "format": format,
            "count": num_images_to_log,
        }
        if _server_accepts_image_filenames():
            meta["filenames"] = [obj["path"] for obj in jsons]
        else:
            wandb.termwarn(
                "Unable to log image array filenames. In some cases, this can prevent images from being"
                "viewed in the UI. Please upgrade your wandb server",
                repeat=False,
            )

        captions = Image.all_captions(seq)

        if captions:
            meta["captions"] = captions

        all_masks = Image.all_masks(seq, run, key, step)

        if all_masks:
            meta["all_masks"] = all_masks

        all_boxes = Image.all_boxes(seq, run, key, step)

        if all_boxes:
            meta["all_boxes"] = all_boxes

        return meta

    @classmethod
    def all_masks(
        cls: Type["Image"],
        images: Sequence["Image"],
        run: "LocalRun",
        run_key: str,
        step: Union[int, str],
    ) -> Union[List[Optional[dict]], bool]:
        all_mask_groups: List[Optional[dict]] = []
        for image in images:
            if image._masks:
                mask_group = {}
                for k in image._masks:
                    mask = image._masks[k]
                    mask_group[k] = mask.to_json(run)
                all_mask_groups.append(mask_group)
            else:
                all_mask_groups.append(None)
        if all_mask_groups and not all(x is None for x in all_mask_groups):
            return all_mask_groups
        else:
            return False

    @classmethod
    def all_boxes(
        cls: Type["Image"],
        images: Sequence["Image"],
        run: "LocalRun",
        run_key: str,
        step: Union[int, str],
    ) -> Union[List[Optional[dict]], bool]:
        all_box_groups: List[Optional[dict]] = []
        for image in images:
            if image._boxes:
                box_group = {}
                for k in image._boxes:
                    box = image._boxes[k]
                    box_group[k] = box.to_json(run)
                all_box_groups.append(box_group)
            else:
                all_box_groups.append(None)
        if all_box_groups and not all(x is None for x in all_box_groups):
            return all_box_groups
        else:
            return False

    @classmethod
    def all_captions(
        cls: Type["Image"], images: Sequence["Media"]
    ) -> Union[bool, Sequence[Optional[str]]]:
        return cls.captions(images)

    def __ne__(self, other: object) -> bool:
        return not self.__eq__(other)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Image):
            return False
        else:
            self_image = self.image
            other_image = other.image
            if self_image is not None:
                self_image = list(self_image.getdata())
            if other_image is not None:
                other_image = list(other_image.getdata())

            return (
                self._grouping == other._grouping
                and self._caption == other._caption
                and self._width == other._width
                and self._height == other._height
                and self_image == other_image
                and self._classes == other._classes
            )

    def to_data_array(self) -> List[Any]:
        res = []
        if self.image is not None:
            data = list(self.image.getdata())
            for i in range(self.image.height):
                res.append(data[i * self.image.width : (i + 1) * self.image.width])
        self._free_ram()
        return res

    def _free_ram(self) -> None:
        if self._path is not None:
            self._image = None

    @property
    def image(self) -> Optional["PIL.Image"]:
        if self._image is None:
            if self._path is not None:
                pil_image = util.get_module(
                    "PIL.Image",
                    required='wandb.Image needs the PIL package. To get it, run "pip install pillow".',
                )
                self._image = pil_image.open(self._path)
                self._image.load()
        return self._image

[((37, 22, 37, 49), 'wandb.util._get_max_cli_version', 'util._get_max_cli_version', ({}, {}), '()', False, 'from wandb import util\n'), ((40, 11, 40, 35), 'pkg_resources.parse_version', 'parse_version', ({(40, 25, 40, 34): '"""0.12.10"""'}, {}), "('0.12.10')", False, 'from pkg_resources import parse_version\n'), ((40, 39, 40, 69), 'pkg_resources.parse_version', 'parse_version', ({(40, 53, 40, 68): 'max_cli_version'}, {}), '(max_cli_version)', False, 'from pkg_resources import parse_version\n'), ((223, 20, 226, 9), 'wandb.util.get_module', 'util.get_module', (), '', False, 'from wandb import util\n'), ((234, 20, 237, 9), 'wandb.util.get_module', 'util.get_module', (), '', False, 'from wandb import util\n'), ((304, 15, 304, 46), 'os.path.join', 'os.path.join', ({(304, 28, 304, 35): '"""media"""', (304, 37, 304, 45): '"""images"""'}, {}), "('media', 'images')", False, 'import os\n'), ((399, 13, 402, 9), 'wandb.util.get_module', 'util.get_module', (), '', False, 'from wandb import util\n'), ((238, 39, 238, 67), 'wandb.util.get_full_typename', 'util.get_full_typename', ({(238, 62, 238, 66): 'data'}, {}), '(data)', False, 'from wandb import util\n'), ((239, 18, 239, 27), 'io.BytesIO', 'BytesIO', ({}, {}), '()', False, 'from io import BytesIO\n'), ((429, 18, 429, 46), 'typing.cast', 'cast', ({(429, 23, 429, 40): "Sequence['Image']", (429, 42, 429, 45): 'seq'}, {}), "(Sequence['Image'], seq)", False, 'from typing import Any, cast, Dict, List, Optional, Sequence, Type, TYPE_CHECKING, Union\n'), ((436, 23, 436, 60), 'wandb.util.to_forward_slash_path', 'util.to_forward_slash_path', ({(436, 50, 436, 59): 'media_dir'}, {}), '(media_dir)', False, 'from wandb import util\n'), ((454, 12, 456, 13), 'logging.warning', 'logging.warning', ({(455, 16, 455, 105): '"""Images sizes do not match. This will causes images to be display incorrectly in the UI."""'}, {}), "(\n    'Images sizes do not match. This will causes images to be display incorrectly in the UI.'\n    )", False, 'import logging\n'), ((468, 12, 472, 13), 'wandb.termwarn', 'wandb.termwarn', (), '', False, 'import wandb\n'), ((230, 14, 230, 36), 'os.path.splitext', 'os.path.splitext', ({(230, 31, 230, 35): 'path'}, {}), '(path)', False, 'import os\n'), ((356, 29, 356, 81), 'os.path.join', 'os.path.join', ({(356, 42, 356, 49): '"""media"""', (356, 51, 356, 60): '"""classes"""', (356, 62, 356, 79): "class_id + '_cls'"}, {}), "('media', 'classes', class_id + '_cls')", False, 'import os\n'), ((583, 28, 586, 17), 'wandb.util.get_module', 'util.get_module', (), '', False, 'from wandb import util\n'), ((240, 12, 240, 47), 'wandb.util.ensure_matplotlib_figure', 'util.ensure_matplotlib_figure', ({(240, 42, 240, 46): 'data'}, {}), '(data)', False, 'from wandb import util\n'), ((244, 45, 244, 73), 'wandb.util.get_full_typename', 'util.get_full_typename', ({(244, 68, 244, 72): 'data'}, {}), '(data)', False, 'from wandb import util\n'), ((245, 23, 247, 13), 'wandb.util.get_module', 'util.get_module', ({(246, 16, 246, 35): '"""torchvision.utils"""', (246, 37, 246, 79): '"""torchvision is required to render images"""'}, {}), "('torchvision.utils', 'torchvision is required to render images'\n    )", False, 'from wandb import util\n'), ((263, 52, 263, 70), 'wandb.util.generate_id', 'util.generate_id', ({}, {}), '()', False, 'from wandb import util\n')]