from __future__ import annotations

import base64
import io
import logging
import pathlib
from typing import Optional

# TODO: Remove these, as they should be contained in the numpy backend.
import numpy as np
import PIL.Image as PImage

import arlunio.ast as ast
import arlunio.color as color
import arlunio.mask as mask
import arlunio.math as math

logger = logging.getLogger(__name__)


class Image:
    """Our representation of an image, implemented as a wrapper around a standard
    Pillow image."""

    def __init__(self, img: PImage.Image):
        self.img = img
        """The wrapped pillow image object."""

    def __eq__(self, other):

        if not isinstance(other, Image):
            return False

        a = np.asarray(self.img)
        b = np.asarray(other.img)

        return (a == b).all()

    def __add__(self, other):

        if isinstance(other, Image):
            other = other.img

        if not isinstance(other, PImage.Image):
            raise TypeError("Addition is only supported between images.")

        img = self.copy()
        img.alpha_composite(other)

        return img

    @property
    def __array_interface__(self):
        # Ensure that our version of an image also plays nice with numpy.
        return self.img.__array_interface__

    def _repr_png_(self):
        # Give nice previews in jupyter notebooks
        return self.img._repr_png_()

    @property
    def size(self):
        return self.img.size

    def alpha_composite(self, im, *args, **kwargs):
        """Composites an image onto this image.

        See :meth:`pillow:PIL.Image.Image.alpha_composite`
        """

        if isinstance(im, Image):
            im = im.img

        self.img.alpha_composite(im, *args, **kwargs)

    def copy(self):
        """Return a copy of the image.
        See :meth:`pillow:PIL.Image.Image.copy`
        """
        return Image(self.img.copy())

    def paste(self, *args, **kwargs):
        """Paste another image into this image.

        See :meth:`pillow:PIL.Image.Image.paste`
        """
        self.img.paste(*args, **kwargs)

    def save(self, *args, **kwargs):
        """Save the image with the given filename.

        See :meth:`pillow:PIL.Image.Image.save`
        """
        self.img.save(*args, **kwargs)

    def thumbnail(self, *args, **kwargs):
        """Convert this image into a thumbail.

        See :meth:`pillow:PIL.Image.Image.thumbnail`
        """

        self.img.thumbnail(*args, **kwargs)


def new(color) -> Image:
    """Creates a new image with the given background color."""
    return ast.Node.builtin(name="image", color=color)


def fromarray(*args, **kwargs):
    """Create an image from an array

    See :func:`pillow:PIL.Image.fromarray`
    """
    return Image(PImage.fromarray(*args, **kwargs))


def load(*args, **kwargs) -> Image:
    """Load an image from the given file.

    See :func:`pillow:PIL.Image.open`
    """
    return Image(PImage.open(*args, **kwargs))


def save(image: Image, filename: str, mkdirs: bool = False) -> None:
    """Save an image in PNG format.

    :param filename: The filepath to save the image to.
    :param mkdirs: If true, make any parent directories
    """
    path = pathlib.Path(filename)

    if not path.parent.exists() and mkdirs:
        path.parent.mkdir(parents=True)

    with open(filename, "wb") as f:
        image.save(f)


def encode(image: Image) -> bytes:
    """Return the image encoded as a base64 string.

    Parameters
    ----------
    image:
       The image to encode.

    Example
    -------
    ::

       >>> import arlunio.image as image
       >>> img = image.new((8, 8), color='red')
       >>> image.encode(img)
       b'iVBORw0KGgoAAAANSUhEUgAAAAgAAAAICAYAAADED76LAAAAFklEQVR4nGP8z8DwnwEPYMInOXwUAAASWwIOH0pJXQAAAABJRU5ErkJggg=='

    """

    with io.BytesIO() as byte_stream:
        image.save(byte_stream, "PNG")
        image_bytes = byte_stream.getvalue()

        return base64.b64encode(image_bytes)


def decode(bytestring: bytes) -> Image:
    """Decode the image represented by the given bytestring into an image object.

    Parameters
    ----------
    bytestring:
       The bytestring to decode.

    Example
    -------

    .. arlunio-image:: Decode Example
       :include-code:

       ::

          import arlunio.image as image

          bytestring = b'iVBORw0KGgoAAAANSUhEUgAAAAgAAAAICAYAAADED76LAAAAFklEQVR4nGP8z8DwnwEPYMInOXwUAAASWwIOH0pJXQAAAABJRU5ErkJggg=='  # noqa: E501
          img = image.decode(bytestring)
    """

    data = base64.b64decode(bytestring)
    bytes_ = io.BytesIO(data)

    return Image(load(bytes_))


def colorramp(values, start: Optional[str] = None, stop: Optional[str] = None) -> Image:
    """Given a 2d array of values, produce an image gradient based on them.

    .. arlunio-image:: Colorramp Demo
       :align: right

       ::

          import arlunio.image as image
          import arlunio.math as math
          import numpy as np

          cartesian = math.Cartesian()
          p = cartesian(width=256, height=256)
          x, y = p[:, :, 0], p[:, :, 1]

          values = np.sin(2*x*np.pi) * np.sin(2*y* np.pi)
          img = image.colorramp(values)

    First this function will scale the input array so that all values fall in the range
    :math:`[0, 1]`. It will then produce an image with the same dimensions as the
    original array. The color of each pixel will be chosen based on the corresponding
    value of the scaled array.

    - If the value is :math:`0` the color will be given by the :code:`start` parameter

    - If the value is :math:`1` the color will be given by the :code:`stop` parameter

    - Otherwise the color will be some mix between the two.

    Parameters
    ----------
    values:
       The array of values used to decide on the color.
    start:
       The color to use for values near :math:`0` (default, :code:`black`)
    stop:
       The color to use for values near :math:`1` (default, :code:`white`)

    Examples
    --------

    .. arlunio-image:: Colorramp Demo 2
       :include-code:

       ::

          import arlunio.image as image
          import arlunio.math as math
          import numpy as np

          cartesian = math.Cartesian()
          p = cartesian(width=256, height=256)

          x = image.colorramp(p[:, :, 0], start="#0000", stop="#f007")
          y = image.colorramp(p[:, :, 1], start="#0000", stop="#00f7")

          img = x + y
    """

    # Scale all the values so that they fall into the range [0, 1]
    minx = np.min(values)
    vs = np.array(values) - minx
    vs = vs / np.max(vs)

    if start is None:
        start = "black"

    if stop is None:
        stop = "white"

    start = color.getcolor(start, "RGBA")
    stop = color.getcolor(stop, "RGBA")

    funcs = [math.lerp(a, b) for a, b in zip(start, stop)]
    channels = [np.floor(func(vs)) for func in funcs]

    pixels = np.array(np.dstack(channels), dtype=np.uint8)
    return fromarray(pixels)


def fill(
    region,
    foreground: Optional[str] = None,
    background: Optional[str] = None,
    image: Optional[Image] = None,
) -> Image:
    """Apply color to an image, as specified by a mask.

    Parameters
    ----------
    mask:
        The mask that selects the region to be coloured
    foreground:
        A string representation of the color to use, this can be in any format that is
        supported by the :mod:`pillow:PIL.ImageColor` module. If omitted this will
        default to black.
    background:
        In the case where an existing image is not provided this parameter can be used
        to set the background color of the generated image. This can be any string that
        is accepted by the :mod:`pillow:PIL.ImageColor` module. If omitted this will
        default to transparent
    image:
        The image to color in, if omitted a blank image will be used.

    Example
    --------

    .. arlunio-image:: Fill Demo
       :include-code:

       ::

          import arlunio.image as image
          import arlunio.shape as shape

          circle = shape.Circle(x0=-0.5, y0=0.25, r=0.6)
          img = image.fill(circle(width=512, height=256), foreground='red')

          circle.x0, circle.y0 = 0, 0
          img = image.fill(circle(width=512, height=256), foreground='#0f0', image=img)

          circle.x0, circle.y0 = 0.5, -0.25
          img = image.fill(circle(width=512, height=256), foreground='blue', image=img)

    """

    foreground = "#000" if foreground is None else foreground
    fill_color = color.getcolor(foreground, "RGBA")

    if image is None:
        background = "#0000" if background is None else background
        image = new(color=background)

    if not isinstance(region, ast.Node):
        region = region()

    return ast.Node.fill(image, region, fill_color)

[(18, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (107, 'arlunio.ast.Node.builtin', 'ast.Node.builtin', True, 'import arlunio.ast as ast\n'), (132, 'pathlib.Path', 'pathlib.Path', False, 'import pathlib\n'), (189, 'base64.b64decode', 'base64.b64decode', False, 'import base64\n'), (190, 'io.BytesIO', 'io.BytesIO', False, 'import io\n'), (256, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (266, 'arlunio.color.getcolor', 'color.getcolor', True, 'import arlunio.color as color\n'), (267, 'arlunio.color.getcolor', 'color.getcolor', True, 'import arlunio.color as color\n'), (323, 'arlunio.color.getcolor', 'color.getcolor', True, 'import arlunio.color as color\n'), (332, 'arlunio.ast.Node.fill', 'ast.Node.fill', True, 'import arlunio.ast as ast\n'), (34, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (35, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (115, 'PIL.Image.fromarray', 'PImage.fromarray', True, 'import PIL.Image as PImage\n'), (123, 'PIL.Image.open', 'PImage.open', True, 'import PIL.Image as PImage\n'), (160, 'io.BytesIO', 'io.BytesIO', False, 'import io\n'), (164, 'base64.b64encode', 'base64.b64encode', False, 'import base64\n'), (257, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (258, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (269, 'arlunio.math.lerp', 'math.lerp', True, 'import arlunio.math as math\n'), (272, 'numpy.dstack', 'np.dstack', True, 'import numpy as np\n')]