"""
Created on Dec 18, 2015
@author: Stephen O'Hara

This module defines the pyvision image class, which wraps
the image data (array) in a class with convenient methods.

One very useful aspect of the image class is that annotations
(such as bounding boxes that might be drawn around detection regions)
are kept on a separate layer, as opposed to drawing directly
onto the image array itself.
"""
# The following prevents a bunch of pylint no-member errors
# with the cv2 module.
# pylint: disable=E1101

import cv2
import numpy as np
import io

try:
    import matplotlib.pyplot as plot
except ImportError:
    print("Error importing matplotlib.")
    print("Matplotlib integration will not work")

try:
    import shapely.geometry as sg
except ImportError:
    print("Error importing shapely.")
    print("Shapely is required for annotating shapes (polygons) on images.")
    print("Shapely is also used to determine if a crop is in bounds, etc.")

from .pv_exceptions import OutOfBoundsError, ImageAnnotationError
from .geometry import in_bounds, integer_bounds, Rect


class Image(object):
    """
    A pyvision3 Image object contains the image data, an
    annotations layer, and many convenient methods.
    Supports 1 channel and 3 channel images.
    """

    def __init__(self, source, *args, desc="Pyvision Image", **kwargs):
        """
        The constructor wraps a cv2.imread(...) function,
        passing in the args and kwargs appropriately, and
        then allocating data for the annotations layer.

        Parameters
        ----------
        desc: string
            Provide a short description of this image, that will be used
            by default in window titles and other functions
        source: string, file object, or cv2 image array
            If string, this is the full path to the image file to load.
            If file object, this is an open file handle from which to load
            the image.
            If ndarray, then we assume this is a cv2 image array which we will
            just wrap.
        args: variable
            Other args will be passed through to cv2.imread, the first arg
            should be the image source, like a file name. See the cv2 docs
            on imread for more details.
        kwargs: variable
            Keyword args will be passed through to cv2.imread

        Examples
        --------
        #loading from files
        img = pv3.Image('mypic.jpg')
        img2 = pv3.Image('mypic.jpg', cv2.IMREAD_GRAYSCALE)
        
        #loading from a file handle...this example is silly, but
        # this capability is good for example when loading from a file streamed
        # over a network connection, or S3, etc.
        with open('somepath/somefile.png', 'rb') as infile:
            img3 = pv3.Image(infile)
        
        #Wrapping of a numpy/cv2 ndarray
        img4 = pv3.Image( np.zeros( (480,640), dtype='uint8' ) )
        """
        self.desc = desc
        # metadata dictionary can be used to pass arbitrary info with the image
        self.metadata = {}

        if isinstance(source, np.ndarray):
            self.data = source
            self.metadata["source"] = "np.ndarray"
        elif type(source) == str:
            self.data = cv2.imread(source, *args, **kwargs)
            self.metadata["source"] = "file"
            self.metadata["filename"] = source
        else:
            # assume a file object
            buf = source.read()
            x = np.fromstring(buf, dtype="uint8")
            self.data = cv2.imdecode(x, cv2.IMREAD_UNCHANGED)
            self.metadata["source"] = "file object or buffer"

        self.height, self.width = self.data.shape[0:2]
        self.size = (self.width, self.height)
        self.nchannels = self.data.shape[2] if len(self.data.shape) == 3 else 1

        # Annotation data is a separate BGR image array.
        self.annotation_data = (
            self.data.copy()
            if self.nchannels == 3
            else cv2.cvtColor(self.data, cv2.COLOR_GRAY2BGR)
        )

    def __str__(self):
        txt = "Pyvision3 Image: {}".format(self.desc)
        txt += "\nWidth: {}, Height: {}, Channels: {}, Depth: {}".format(
            self.width, self.height, self.nchannels, self.data.dtype
        )
        return txt

    def __repr__(self):
        return str(self)

    def __getitem__(self, slc):
        return self.data[slc]

    def as_grayscale(self, as_type="PV"):
        """
        Parameters
        ----------
        as_type: str in ("CV", "PV"), default is "PV"

        Returns
        -------
        A copy of the image (data only, not annotations) as a single channel opencv numpy array,
        or, if as_type is "PV", then a pyvision3 Image wrapped around the same.
        """
        if self.nchannels == 3:
            img_gray = cv2.cvtColor(self.data, cv2.COLOR_BGR2GRAY)
        else:
            img_gray = self.data.copy()

        if as_type == "CV":
            return img_gray
        else:
            return Image(img_gray)

    def as_annotated(self, alpha=0.5, as_type="PV"):
        """
        Provides an array which represents the merging of the image data
        with the annotations layer.

        Parameters
        ----------
        alpha: float
            Specify a number between 0.0 and 1.0
            This is the alpha blend when merging the annotations layer onto the image data.
            1.0 means that the annotation will be completely opaque and 0.0 completely transparent.
        as_type: string
            Specify either "CV" or "PV" (default) to indicate the return type of the annotated
            image. If "CV", then an ndarray in the normal opencv format is returned.
            If "PV", then a new pyvision3 image is returned with the annotations baked-in.

        Returns
        -------
        A copy of the image array (possibly expanded to 3 channels if the
        original image is single channel) with the color annotations baked in by
        replacing those pixels in the image with the corresponding non-zero pixels
        in the annotation data.

        Return type is either an opencv ndarray or a pyvision3 image depending on as_type.
        """
        # TODO: What if self.data is a floating point image and the annotations
        # are uint8 BGR? We should probably call a normalizing routine of some
        # sort that copies/converts self.data into a 3-channel BGR 8-bit image
        if self.nchannels == 1:
            tmp_img = cv2.cvtColor(self.data, cv2.COLOR_GRAY2BGR)
        else:
            tmp_img = self.data.copy()

        # this works because the annotation_data was initialized as a copy of the
        # source data. Annotations draw on this copy, and when we alpha-blend, those pixels
        # that were not changed by an annotation will blend back to full intensity.
        # i.e., if there were no annotations on pixel x, (1-alpha)*I(x) + (alpha)*A(x) = I(x)
        # because A(x) == I(x) where not otherwise annotated.
        tmp_img = cv2.addWeighted(
            tmp_img, 1.0 - alpha, self.annotation_data, alpha, 0.0
        )

        if as_type == "PV":
            return Image(tmp_img)

        return tmp_img

    def annotate_shape(
        self, shape, color=(255, 0, 0), fill_color=None, *args, **kwargs
    ):
        """
        Draws the specified shape on the annotation data layer.
        Currently supports LineString, MultiLineString, LinearRing,
        and Polygons.

        Parameters
        ----------
        shape: shapely.geometry shape object
        color:
            An RGB tuple indicating the color, red is (255,0,0)
        fill_color:
            The color used to fill a closed shape (polygon). None means
            that there will be no fill (default).

        *args and **kwargs are for optional line parameters that will
        be passed onto cv2.line(...) which is used at the core for
        drawing the line segments of the shape.

        Note
        ----
        OpenCV stores images in BGR not RGB order, thus the annotation
        layer will be constructed this way. However, to make it easier for the
        user, colors to be used are specified in the normal (r,g,b) tuple order,
        and internally, we handle it.
        """
        # TODO: annotate_shape should support all shapely geometries
        # TODO: annotations should support alpha-channel fills for partial transparency
        if isinstance(shape, sg.LinearRing) or isinstance(shape, sg.LineString):
            self._draw_segments(shape, color, *args, **kwargs)
        elif isinstance(shape, sg.MultiLineString):
            for line_string in shape:
                self._draw_segments(line_string, color, *args, **kwargs)
        elif isinstance(shape, sg.Polygon):
            if fill_color is not None:
                # use cv2.fillPoly
                c = self._fix_color_tuple(fill_color)
                exterior = np.array(shape.exterior.coords, dtype="int")
                interiors = [np.array(x.coords, dtype="int") for x in shape.interiors]
                cv2.fillPoly(self.annotation_data, [exterior] + interiors, color=c)
            # draw external ring of polygon
            self._draw_segments(shape.exterior, color, *args, **kwargs)
            # draw interior rings (holes) if any
            for interior_ring in shape.interiors:
                self._draw_segments(interior_ring, color, *args, **kwargs)

    def annotate_point(self, point, color=(255, 0, 0)):
        """
        Annotates a point by drawing a filled circle of radius-3 pix in the annotation
        layer.

        Parameters
        ----------
        point:  tuple (int: x, int: y) or shapely Point object
        color:  tuple (r,g,b)
        """
        pt = (
            (int(point.x), int(point.y))
            if isinstance(point, sg.point.Point)
            else tuple(point)
        )
        self.annotate_circle(pt, 3, color, thickness=-1)

    def annotate_circle(self, ctr, radius, color=(255, 0, 0), *args, **kwargs):
        """
        Draws a circle on the annotation layer

        Parameters
        ----------
        ctr: (int: x, int: y)
            the center points of the circle

        radius: int
            the radius in pixels of the circle

        color: (r,g,b)

        *args and **kwargs will be passed onto cv2.circle function, and
        can be used to control the thickness and line type. Note that a negative
        thickness indicates that the circle will be filled.
        """
        c = self._fix_color_tuple(color)
        ctr = (int(ctr[0]), int(ctr[1]))
        cv2.circle(self.annotation_data, ctr, radius, c, *args, **kwargs)

    def annotate_line(self, pt1, pt2, color, *args, **kwargs):
        """
        Draws a line segment between two points on the annotation layer.

        Parameters
        ----------
        pt1, pt2: tuples, (int: x, int: y)
            The start and end points of the line segment

        color: tuple (r,g,b)

        *args and **kwargs will be passed onto cv2.line, which can be
        used to control line thickness and style
        """
        c = self._fix_color_tuple(color)
        cv2.line(self.annotation_data, tuple(pt1), tuple(pt2), c, *args, **kwargs)

    def annotate_rect(self, pt1, pt2, color=(255, 0, 0), *args, **kwargs):
        """
        Draws a rectangle using upper left (pt1) and lower right (pt2) coordinates.
        To fill the rectangle, specify thickness=-1 in the kwargs. If you have a
        rectangle as a shapely object (polygon), use the annotate_shape method instead.

        Parameters
        ----------
        pt1, pt2: tuple (int:x, int:y)
            The upper left and lower right corners of the rectangle
        color:  tuple (r,g,b)
            The rgb color of the rectangle
        *args and **kwargs will be passed onto cv2.rectangle, which can be
        used to control line thickness and style.

        Note
        ----
        If your rectangle is a shapely Polygon (which it will be if you use
        pv3.Rect(...) to create it), then use the annotate_shape method instead.
        """
        c = self._fix_color_tuple(color)
        cv2.rectangle(
            self.annotation_data, tuple(pt1), tuple(pt2), color=c, *args, **kwargs
        )

    def annotate_text(
        self,
        txt,
        point,
        color=(0, 0, 0),
        bg_color=None,
        font_face=cv2.FONT_HERSHEY_PLAIN,
        font_scale=1,
        *args,
        **kwargs
    ):
        """
        Draws the specified text string to the annotations layer, uses cv2.putText() function
        as the basis, so *args and **kwargs will be passed onto that function.

        Parameters
        ----------
        txt:    str
            The text string to be drawn to the image
        point:  tuple (int: x, int: y)
            The location of the start of the text string
        color:  tuple (r,g,b)
            The color of the text
        bg_color: tuple (r,g,b) or None
            Optional. If specified then a filled box of this color will
            serve as a background for the text. This is handy to ensure
            text legibility.
        font_face: cv2 font selection, defaults to cv2.FONT_HERSHEY_PLAIN
        font_scale: cv2 font Scale, defaults to 1

        *args and **kwargs will be passed onto the underlying cv2.putText
        function.
        """
        c = self._fix_color_tuple(color)
        if bg_color is not None:
            # we will draw a filled rectangle of this color to be behind
            # the annotated text
            ((w, h), _) = cv2.getTextSize(txt, font_face, font_scale, thickness=1)
            point1 = (point[0] - 1, point[1] + 1)
            point2 = (point1[0] + w + 2, point1[1] - h - 2)
            self.annotate_rect(point1, point2, color=bg_color, thickness=-1)

        cv2.putText(
            self.annotation_data,
            txt,
            tuple(point),
            fontFace=font_face,
            fontScale=font_scale,
            color=c,
            *args,
            **kwargs
        )

    def annotate_mask(self, mask_img, transparency=(0, 0, 0)):
        """
        Replaces any and all current annotations on the image with the provided
        3-channel 8-bit mask image. To change what mask pixels are considered
        transparent, use the set_annotation_transparency method.

        Parameters
        ----------
        mask_img:       ndarray (3 channel, uint8 image, same size as self.data, BGR color order)
        transparency:   pixels in the input mask_img that match this color will be set as transparent,
                        and thus the background image's pixel value will be used instead.
                        None means there are no fully transparent pixels. Default is black (0, 0, 0),
                        so anything black in the mask_img will be replaced with the corresponding color
                        value from the underlying image.
        """
        if mask_img.shape[0:2] != self.data.shape[0:2]:
            raise ImageAnnotationError("Invalid mask. Must be same (w,h) as image.")
        self.annotation_data = mask_img.copy()
        if transparency is not None:
            pix = np.nonzero((self.annotation_data == transparency).all(axis=2))
            self.annotation_data[pix] = self.data[pix]

    def annotate_inset_image(self, inset_image, pos, size=None):
        """
        On the annotation layer, composites another image, typically a thumbnail or other graphic,
        at the specified inset position and size. This is a way that an icon image could be drawn
        to a specific location as an annotation, or for an inset graphic like a performance meter,
        to be drawn as an inset onto a corner of the image, etc.

        Parameters
        ----------
        inset_image:    A pyvision3 image to be drawn as an annotation on the current image. This may be
                        a single channel or 3 channel image.
        pos:            The (x, y) coordinates, relative to the current image, where the top-left corner
                        of the inset image will be placed.
        size:           The (width, height) of the inset image in pixels. The inset image will be resized
                        to this specified size before drawing it to the annotation layer. Specify None
                        if no resizing is required of the inset image.

        Note
        -----
        The pixels from the inset image will replace whatever is already at that location in the annotation
        layer. The inset image will be subject to the same alpha-blending as the other annotations.
        """
        if size is not None:
            if inset_image.size != size:
                inset_image = inset_image.resize(size, as_type="PV")
        else:
            size = inset_image.size

        cv_tile = inset_image.data
        roi = Rect(pos[0], pos[1], size[0], size[1])
        depth = cv_tile.shape[-1] if len(cv_tile.shape) == 3 else 1

        if depth == 1:
            cv_tile_bgr = cv2.cvtColor(cv_tile, cv2.COLOR_GRAY2BGR)
        else:
            cv_tile_bgr = cv_tile

        # copy pixels of tile onto appropriate location in annotation image
        dest = self.annotation_data
        (minx, miny, maxx, maxy) = integer_bounds(roi)
        dest[miny : (maxy + 1), minx : (maxx + 1), :] = cv_tile_bgr

    def _draw_segments(self, simple_shape, color, *args, **kwargs):
        """
        Internal method for drawing a "simple" shapely geometric object,
        i.e., one that has a single set of coordinates which will be
        connected in sequence. Examples are LineStrings and LinearRings,
        the latter being a closed sequence of points and the former is not.

        Parameters
        ----------
        simple_shape:   shape
            A shape object that provides a coords member
        color:  tuple (r,g,b)
        *args, **kwargs:
            will be passed on to self.annotate_line for each
            line segment in the simple shape.
        """
        points = [(int(x), int(y)) for (x, y) in simple_shape.coords]
        for i in range(1, len(points)):
            pt1 = points[i - 1]
            pt2 = points[i]
            self.annotate_line(pt1, pt2, color, *args, **kwargs)

    def _fix_color_tuple(self, color):
        """
        Puts the color tuple (r,g,b) into (b,g,r) order.

        Parameters
        ----------
        color:  tuple (r,g,b)

        Returns
        -------
        tuple (b,g,r)
        """
        (r, g, b) = color
        return (b, g, r)

    def copy(self):
        """
        Returns
        -------
        Returns a pyvision3 image which is a 'deep copy' of this image which can be
        freely modified without affecting the source of the copy.
        """