#!/usr/bin/env python3
# coding=utf-8
"""
Extension Inkscape : decoupe une image en couches empilees par nuances de gris.

L'image selectionnee est convertie en gris, floutee puis quantifiee en quelques
niveaux ; chaque niveau devient une forme fermee dont le contour suit la
frontiere de l'aplat, a la maniere des courbes de niveau d'une carte.

La logique vit dans `gray_iso_layers_core`, sans dependance a inkex.
"""

import base64
import hashlib
import io
import math
import os
import sys
from urllib.parse import unquote, unquote_to_bytes, urlparse
from urllib.request import url2pathname

import inkex
from inkex.localization import localize

# L'extension vit dans son propre sous-dossier des extensions Inkscape, qui
# n'est pas forcement dans sys.path : on l'y ajoute pour trouver gray_iso_layers_core.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from gray_iso_layers_core import (  # noqa: E402
    add_detail, blend_fields, board_labels, gray_to_hex, guided_filter, iso_boards,
    lines_to_d, normalize, number_lines, parse_color, rings_to_d, scale_lines,
    scale_rings)
from gray_iso_layers_depth import DepthUnavailable, depth_field  # noqa: E402

# Modes Pillow a plus de 8 bits par canal (carte de profondeur en 16 bits).
_DEEP_MODES = ("I", "F")

# Grand cote de l'image envoyee au calcul de profondeur : les modeles
# travaillent a 518 points au plus.
_DEPTH_SIZE = 768

# Installe _() dans les builtins : catalogue « grayisolayers » de la langue
# de l'interface (Inkscape transmet domaine et dossier par l'environnement),
# textes anglais d'origine sinon.
localize()


class GrayIsoLayers(inkex.EffectExtension):
    """Decoupe chaque image selectionnee en planches a empiler, une par
    niveau de gris, chacune dans son propre calque avec sa decoupe et le
    repere de la planche suivante."""

    def add_arguments(self, pars):
        # Un argument par <param> du .inx, memes noms et memes valeurs par defaut.
        pars.add_argument("--tab", default="levels")
        pars.add_argument("--levels", type=int, default=16)
        pars.add_argument("--equalize", type=float, default=0.0)
        pars.add_argument("--shapes", default="light")
        pars.add_argument("--layout", default="stack")
        pars.add_argument("--mark", type=inkex.Boolean, default=True)
        pars.add_argument("--numbers", type=inkex.Boolean, default=True)
        pars.add_argument("--blur", type=float, default=1.0)
        pars.add_argument("--edges", type=inkex.Boolean, default=True)
        pars.add_argument("--resolution", type=int, default=800)
        pars.add_argument("--min_area", type=float, default=0.02)
        pars.add_argument("--simplify", type=float, default=0.4)
        pars.add_argument("--smooth", type=inkex.Boolean, default=True)
        pars.add_argument("--depth_mix", type=float, default=0.0)
        pars.add_argument("--depth_model", default="midas_small")
        pars.add_argument("--depth_detail", type=float, default=100.0)
        pars.add_argument("--fill_mode", default="gray")
        pars.add_argument("--fill_color", default="4294967295")
        pars.add_argument("--stroke_width", type=float, default=0.2)
        pars.add_argument("--unit", default="mm")
        pars.add_argument("--stroke_color", default="4278190335")
        pars.add_argument("--mark_color", default="255")
        pars.add_argument("--keep_original", type=inkex.Boolean, default=True)

    def effect(self):
        images = list(self.svg.selection.get(inkex.Image))
        if not images:
            inkex.errormsg(_("Select at least one bitmap image (File > Import), "
                             "then run the extension again."))
            return
        if self.options.levels < 2:
            inkex.errormsg(_("At least 2 gray levels are needed."))
            return
        # Un calque par planche, commun a toutes les images traitees.
        self.board_layers = {}
        for image in images:
            self.process(image)

    def board_layer(self, number, above):
        """Calque de la planche `number` (1 = la plus au fond), cree a la
        demande au premier plan du document, juste au-dessus de `above`
        (calque precedent ou image)."""
        if number not in self.board_layers:
            layer = inkex.Layer()
            layer.label = _("Board no. {}").format(number)
            root = self.svg
            # Un calque Inkscape est un enfant direct de la racine : on remonte
            # jusqu'a l'ancetre de premier plan pour s'inserer apres lui.
            while above.getparent() is not None and above.getparent() is not root:
                above = above.getparent()
            root.insert(root.index(above) + 1, layer)
            self.board_layers[number] = layer
        return self.board_layers[number]

    def process(self, image):
        """Traite une image : calcul dans gray_iso_layers_core, ecriture dans le SVG."""
        opt = self.options
        try:
            picture, digest = self.load_picture(image)
            field = self.gray_field(picture, opt.resolution)
        except Exception as error:  # fichier absent, format inconnu, lien casse...
            inkex.errormsg(_("Cannot read the image '{}' ({}). Embed it in the "
                             "document (right-click > Embed Image) or check its "
                             "link.").format(image.get_id(), error))
            return
        if opt.depth_mix > 0:
            field = self.add_depth(image, picture, digest, field)

        rows, cols = field.shape
        boards = iso_boards(
            field, levels=opt.levels,
            blur=opt.blur / 100.0 * max(rows, cols),
            preserve_edges=opt.edges,
            min_area=opt.min_area / 100.0,
            tolerance=opt.simplify,
            shapes=opt.shapes,
            equalize=opt.equalize / 100.0)
        if sum(1 for board in boards if board[2]) < 2:
            inkex.errormsg(_("The image '{}' is uniform: only the base level has "
                             "a shape. Reduce the smoothing or pick a more "
                             "contrasted image.").format(image.get_id()))

        # Les couches sont calculees dans le repere propre de l'image, puis
        # ecrites en coordonnees du document : les calques n'ont ainsi aucune
        # transformation, meme pour une image tournee ou etiree.
        left, top, width, height = self.image_box(image, cols, rows)
        scale_x, scale_y = width / (cols - 1.0), height / (rows - 1.0)
        box = (left, top, left + width, top + height)
        base = image.composed_transform().matrix
        shifts = self.board_shifts(base, box, len(boards))

        above = image
        for number, (_level, gray, rings, mark) in enumerate(boards, 1):
            (a, c, e), (b, d, f) = base
            shift_x, shift_y = shifts[number - 1]
            matrix = ((a, c, e + shift_x), (b, d, f + shift_y))
            # Les planches sont numerotees du fond vers l'avant. Calque cree
            # meme si le niveau est vide : autant de calques que de niveaux.
            layer = self.board_layer(number, above)
            above = layer
            if not rings:
                continue
            node = inkex.PathElement()
            node.label = _("Cut")
            node.set("d", rings_to_d(scale_rings(rings, scale_x, scale_y, left, top),
                                     smooth=opt.smooth, box=box, matrix=matrix))
            node.style = self.layer_style(gray)
            layer.add(node)

            # Repere d'assemblage : le bord de la planche suivante, reporte
            # sur celle-ci pour savoir ou la poser.
            mark_style = self.mark_style() if opt.mark and mark else None
            if mark_style is not None and (mark[0] or mark[1]):
                node = inkex.PathElement()
                node.label = _("Marking of board no. {}").format(number + 1)
                node.set("d", lines_to_d(scale_lines(mark, scale_x, scale_y, left, top),
                                         smooth=opt.smooth, box=box, matrix=matrix))
                node.style = mark_style
                layer.add(node)

            # Numeros a graver : celui de la planche, encadre, sur chacune de
            # ses pieces, et celui de la planche suivante la ou elle se pose.
            if opt.numbers and self.mark_style() is not None:
                following = boards[number][2] if mark is not None else None
                own, after = board_labels(
                    rings, following, number, cols, rows,
                    largest=0.03 * max(rows, cols), smallest=0.008 * max(rows, cols))
                parts = [number_lines(number, x, y, size, framed=True)
                         for x, y, size in own]
                parts += [number_lines(number + 1, x, y, size) for x, y, size in after]
                d = " ".join(
                    lines_to_d(scale_lines(part, scale_x, scale_y, left, top),
                               smooth=False, matrix=matrix) for part in parts)
                if d:
                    node = inkex.PathElement()
                    node.label = _("Numbers")
                    node.set("d", d)
                    node.style = self.mark_style()
                    layer.add(node)

        if not opt.keep_original:
            image.delete()

    def board_shifts(self, matrix, box, count):
        """Decalage (x, y) de chaque planche, en coordonnees du document.

        Planches superposees : aucun, chacune reste a la place de l'image.
        Planches juxtaposees : elles sont rangees en grille a droite de
        l'image, de gauche a droite puis de haut en bas, chacune dans une
        case de la taille de l'image, pour etre vues et decoupees une a une.
        """
        if self.options.layout != "side":
            return [(0.0, 0.0)] * count
        (a, c, e), (b, d, f) = matrix
        x0, y0, x1, y1 = box
        # Encombrement de l'image dans le document, meme tournee ou etiree.
        corners = [(a * x + c * y + e, b * x + d * y + f)
                   for x, y in ((x0, y0), (x1, y0), (x1, y1), (x0, y1))]
        width = max(x for x, _y in corners) - min(x for x, _y in corners)
        height = max(y for _x, y in corners) - min(y for _x, y in corners)
        gap = 0.05 * max(width, height)
        columns = int(math.ceil(math.sqrt(count)))
        return [((width + gap) * (1 + k % columns), (height + gap) * (k // columns))
                for k in range(count)]

    def layer_style(self, gray):
        """Style d'une couche : remplissage selon le mode, trait facultatif."""
        opt = self.options
        style = inkex.Style({"fill": "none", "fill-rule": "evenodd", "stroke": "none"})
        if opt.fill_mode == "gray":
            style["fill"] = gray_to_hex(gray)
        elif opt.fill_mode == "color":
            color, opacity = parse_color(opt.fill_color, ("#ffffff", 1.0))
            style["fill"] = color
            style["fill-opacity"] = str(opacity)

        width = self.svg.unittouu("{}{}".format(opt.stroke_width, opt.unit))
        if width > 0:
            color, opacity = parse_color(opt.stroke_color, ("#ff0000", 1.0))
            style["stroke"] = color
            style["stroke-opacity"] = str(opacity)
            style["stroke-width"] = str(round(width, 6))
            style["stroke-linejoin"] = "round"
        return style

    def mark_style(self):
        """Style du repere d'assemblage : un trait sans remplissage, de la
        meme epaisseur que la decoupe. `None` s'il n'y a pas de trait."""
        opt = self.options
        width = self.svg.unittouu("{}{}".format(opt.stroke_width, opt.unit))
        if width <= 0:
            return None
        color, opacity = parse_color(opt.mark_color, ("#000000", 1.0))
        return inkex.Style({
            "fill": "none", "stroke": color, "stroke-opacity": str(opacity),
            "stroke-width": str(round(width, 6)), "stroke-linejoin": "round",
            "stroke-linecap": "round"})

    @staticmethod
    def image_box(image, cols, rows):
        """Rectangle reellement occupe par le bitmap, dans le repere de l'image.

        Sauf preserveAspectRatio="none", le bitmap garde ses proportions et il
        est centre dans le cadre x, y, width, height (cas « meet » par defaut).
        """
        left, top, width, height = image.left, image.top, image.width, image.height
        ratio = image.get("preserveAspectRatio", "xMidYMid meet")
        if "none" not in ratio and "slice" not in ratio:
            zoom = min(width / cols, height / rows)
            fitted_width, fitted_height = cols * zoom, rows * zoom
            left += (width - fitted_width) / 2.0
            top += (height - fitted_height) / 2.0
            width, height = fitted_width, fitted_height
        return left, top, width, height

    def add_depth(self, image, picture, digest, field):
        """Melange aux gris `field` la profondeur estimee de l'image, dans la
        proportion `depth_mix`. Si elle ne peut pas etre calculee, le dit et
        rend les gris seuls."""
        import numpy
        from PIL import Image

        opt = self.options
        try:
            depth = depth_field(self.color_picture(picture), digest, opt.depth_model)
        except DepthUnavailable as error:
            inkex.errormsg(_("No depth map for the image '{}' ({}). Install the "
                             "depth models with 'deploy.ps1 -Depth'. The grays "
                             "of the image are used alone.").format(
                                 image.get_id(), error))
            return field
        rows, cols = field.shape
        depth = Image.fromarray(numpy.asarray(depth, dtype=numpy.float32))
        depth = normalize(numpy.asarray(depth.resize((cols, rows), Image.BICUBIC),
                                        dtype=float))
        # La carte sort du modele en basse resolution : ses bords flous sont
        # recales sur ceux de la photo.
        depth = guided_filter(field, depth, max(2, int(round(max(rows, cols) / 100.0))))
        # Les yeux et la bouche, absents de la profondeur, viennent des gris.
        depth = add_detail(depth, field, opt.depth_detail / 100.0,
                           max(rows, cols) / 40.0)
        return blend_fields(field, depth, opt.depth_mix / 100.0)

    def load_picture(self, image):
        """Lit le bitmap : `(image PIL, empreinte de ses octets)`."""
        from PIL import Image

        href = image.get("xlink:href") or image.get("href") or ""
        if href.startswith("data:"):
            header, _sep, payload = href.partition(",")
            data = base64.b64decode(payload) if ";base64" in header \
                else unquote_to_bytes(payload)
        else:
            if href.startswith("file:"):
                path = url2pathname(urlparse(href).path)
            else:
                path = unquote(href)
            # Lien relatif : il part du dossier du document.
            with open(self.absolute_href(path), "rb") as handle:
                data = handle.read()
        picture = Image.open(io.BytesIO(data))
        picture.load()
        return picture, hashlib.sha1(data).hexdigest()

    @staticmethod
    def gray_picture(picture):
        """Gris de l'image, en image PIL : mode « L » (0 a 255), ou « F »
        (0 a 1) pour une image a plus de 8 bits par canal."""
        import numpy
        from PIL import Image

        if picture.mode.startswith(_DEEP_MODES):
            # Carte de profondeur en 16 bits : Pillow l'ecreterait a 255 en la
            # convertissant, on l'etire nous-memes sur 0..1.
            values = normalize(numpy.asarray(picture, dtype=float))
            return Image.fromarray(values.astype(numpy.float32))
        # La transparence se lit comme du blanc (fond de page).
        rgba = picture.convert("RGBA")
        white = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
        return Image.alpha_composite(white, rgba).convert("L")

    @classmethod
    def gray_field(cls, picture, resolution):
        """Gris de l'image (0 noir, 1 blanc) en tableau 2D, ramene a
        `resolution` points sur son grand cote."""
        import numpy
        from PIL import Image

        gray = cls.gray_picture(picture)
        zoom = resolution / float(max(gray.size))
        size = (max(2, int(round(gray.width * zoom))),
                max(2, int(round(gray.height * zoom))))
        gray = gray.resize(size, Image.LANCZOS)
        scale = 1.0 if gray.mode == "F" else 255.0
        return numpy.asarray(gray, dtype=float) / scale

    @classmethod
    def color_picture(cls, picture):
        """Image RGB envoyee au calcul de profondeur, transparence en blanc,
        reduite a `_DEPTH_SIZE` points sur son grand cote."""
        import numpy
        from PIL import Image

        if picture.mode.startswith(_DEEP_MODES):
            gray = numpy.asarray(cls.gray_picture(picture), dtype=float)
            color = Image.fromarray((gray * 255.0).round().astype(numpy.uint8))
            color = color.convert("RGB")
        else:
            rgba = picture.convert("RGBA")
            white = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
            color = Image.alpha_composite(white, rgba).convert("RGB")
        color.thumbnail((_DEPTH_SIZE, _DEPTH_SIZE), Image.LANCZOS)
        return color


if __name__ == "__main__":
    GrayIsoLayers().run()