# Builds assets/SingeText.obj (+ .mtl):  the Singe logotype from SingeText.jpeg as solid letters,
# extruded with a chamfer on both faces, standing upright in the XY plane with depth along Z,
# resting on Y = 0 and centred on X = 0.  Convert with util/objToGlb.py for Singe.
#
# --print makes the printable variant:  the letters sit on a backing plate that joins them into one
# piece, with the back faces flat on the plate and the chamfer on the front only.  It is split into
# three colour bodies (letter faces, letter walls, plate); --stl writes one STL per body in
# millimetres, --width mm wide, and --check voxelises the result and counts the pieces.
#
# Usage:  python3 util/textModel.py assets/SingeText.jpeg assets/SingeText.obj [--depth 15] [--chamfer 5]
#         python3 util/textModel.py assets/SingeText.jpeg assets/SingeTextPrint.obj --print --stl assets --width 150 --check
import argparse
import os
import numpy as np
from PIL import Image, ImageFilter
from scipy import ndimage
from objWriter import ObjWriterT, writeMtl, pieceCount, reportPieces, DRAGON_MATERIALS
from outline import traceBoundary, simplifyClean, signedArea, offsetRing, earClip

FACE = 6.0                 # thickness of the coloured face layer on printable letters, image pixels
PLATE = 16.0               # thickness of a backing plate, image pixels
PLATE_MARGIN = 12.0        # how far a plate reaches beyond what it carries, image pixels
PLATE_BRIDGE = 40          # closing radius (at UPSCALE) that joins separate letters into one plate outline
UPSCALE = 2                # the logo is traced at twice its size for smoother curves
BODIES = {'face': 'dark', 'walls': 'frame', 'plate': 'plate'}


def identity(p):
    return p


def plateMask(mask, margin, bridge, upscale=1):
    # A plate outline around a mask:  gaps bridged, reaching margin beyond it, one island, spurs removed.
    radius = bridge
    while True:
        plate = ndimage.binary_dilation(mask, iterations=radius)
        plate = ndimage.binary_erosion(plate, iterations=int(radius - margin * upscale))
        plate = ndimage.binary_fill_holes(plate)
        plate = ndimage.binary_opening(plate, iterations=3 * upscale)
        if ndimage.label(plate)[1] == 1:
            return plate
        radius += 10


class LogoT:
    # The traced logotype:  ink mask, one outline ring per letter, and the routines that extrude them.
    def __init__(self, src, scale=0.01, depth=15.0, chamfer=5.0, tolerance=1.5):
        self.scale = scale
        self.depth = depth
        self.chamfer = chamfer
        grey = Image.open(src).convert('L')
        big = grey.resize((grey.width * UPSCALE, grey.height * UPSCALE), Image.LANCZOS)
        big = big.filter(ImageFilter.GaussianBlur(3.5 / (4 / UPSCALE)))
        self.ink = np.asarray(big) < 128
        ys, xs = np.nonzero(self.ink)
        self.setFootprint(xs.min() / UPSCALE, xs.max() / UPSCALE, ys.max() / UPSCALE)
        labels, count = ndimage.label(self.ink)
        self.holes = 0
        self.letters = []
        for index in range(1, count + 1):
            component = labels == index
            if component.sum() < 50 * UPSCALE * UPSCALE:
                continue
            filled = ndimage.binary_fill_holes(component)
            self.holes += int((filled & ~component).sum() > 0)
            self.letters.append(self.ringOf(filled, tolerance))
        self.letters.sort(key=lambda ring: ring[:, 0].min())
        self.writer = None

    def setFootprint(self, x0, x1, y1):
        # Model space:  centred on X between x0 and x1, resting on Y = 0 at image row y1.
        self.x0 = x0
        self.x1 = x1
        self.y1 = y1

    def toModel(self, x, y, z):
        return np.array([(x - (self.x0 + self.x1) / 2) * self.scale, (self.y1 - y) * self.scale, z * self.scale])

    def ringOf(self, mask, tolerance):
        # Traced at UPSCALE, simplified, in image pixels, counter-clockwise as seen from +Z in model space.
        ring = simplifyClean(traceBoundary(mask), tolerance * UPSCALE) / UPSCALE
        xy = np.array([self.toModel(x, y, 0.0)[:2] for x, y in ring])
        return ring[::-1] if signedArea(xy) < 0 else ring

    def prism(self, name, ring, levels, capMaterial, transform=identity):
        # levels:  (ring, z, material of the walls rising from it) from back to front; caps close the ends.
        rings = [[self.writer.vertex(transform(self.toModel(x, y, z))) for x, y in poly] for poly, z, material in levels]
        self.writer.object(name)
        self.writer.material(capMaterial)
        ringCcw = signedArea(ring) > 0
        for a, b, c in earClip(ring):
            tri = [a, b, c] if ringCcw else [c, b, a]
            self.writer.face([rings[0][i] for i in tri[::-1]])
            self.writer.face([rings[-1][i] for i in tri])
        n = len(ring)
        for level in range(len(rings) - 1):
            self.writer.material(levels[level][2])
            a = rings[level]
            b = rings[level + 1]
            for i in range(n):
                j = (i + 1) % n
                self.writer.face([a[i], a[j], b[j], b[i]])

    def buildLetters(self, writer, printable, zBase, transform=identity):
        # Display:  letters from zBase up, chamfered both sides.  Printable:  flat-backed letters from
        # zBase up, walls in one body and the chamfered face in another.
        self.writer = writer
        depth = self.depth
        chamfer = self.chamfer
        for number, ring in enumerate(self.letters):
            inset = offsetRing(ring, -chamfer)
            name = 'letter%02d' % (number + 1)
            if printable:
                self.prism(name + 'Walls', ring, [(ring, zBase, BODIES['walls']), (ring, zBase + 2 * depth - FACE, BODIES['walls'])], BODIES['walls'], transform)
                self.prism(name + 'Face', ring, [(ring, zBase + 2 * depth - FACE, BODIES['face']), (ring, zBase + 2 * depth - chamfer, BODIES['face']), (inset, zBase + 2 * depth, BODIES['face'])], BODIES['face'], transform)
            else:
                levels = [(inset, zBase, 'logo'), (ring, zBase + chamfer, 'logo'), (ring, zBase + 2 * depth - chamfer, 'logo'), (inset, zBase + 2 * depth, 'logo')]
                self.prism(name, ring, levels, 'logo', transform)

    def buildPlate(self, writer, mask, zTop, transform=identity):
        # A backing plate under a mask given at UPSCALE, from zTop - PLATE up to zTop.
        self.buildPlateRing(writer, self.ringOf(mask, 2.0), zTop, transform)

    def buildPlateRing(self, writer, ring, zTop, transform=identity):
        # The plate from an outline ring in this logo's image pixels.
        self.writer = writer
        self.prism('plate', ring, [(ring, zTop - PLATE, BODIES['plate']), (ring, zTop, BODIES['plate'])], BODIES['plate'], transform)


def writeOutputs(writer, out, header, materials, bodies=None, stlDir=None, widthMm=None, modelWidth=None, check=False, voxel=0.01):
    mtlName = os.path.splitext(os.path.basename(out))[0] + '.mtl'
    writer.write(out, mtlName, header)
    writeMtl(os.path.join(os.path.dirname(out), mtlName), materials, 'Materials for %s' % os.path.basename(out))
    if stlDir:
        stem = os.path.splitext(os.path.basename(out))[0]
        unitScale = widthMm / modelWidth
        for body, material in bodies.items():
            path = os.path.join(stlDir, '%s-%s.stl' % (stem, body))
            print('%s: %d triangles' % (path, writer.writeStl(path, {material}, unitScale, '%s %s' % (stem, body))))
    if check:
        reportPieces(*pieceCount(writer, voxel))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('src')
    parser.add_argument('out')
    parser.add_argument('--scale', type=float, default=0.01, help='model units per image pixel')
    parser.add_argument('--depth', type=float, default=15.0, help='half thickness of the letters in image pixels')
    parser.add_argument('--chamfer', type=float, default=5.0, help='edge chamfer in image pixels')
    parser.add_argument('--tolerance', type=float, default=1.5, help='outline simplification in image pixels')
    parser.add_argument('--print', dest='printable', action='store_true', help='one-piece printable variant on a backing plate')
    parser.add_argument('--stl', default=None, help='directory to write one STL per colour body into')
    parser.add_argument('--width', type=float, default=150.0, help='width of the logo in the STL files, millimetres')
    parser.add_argument('--check', action='store_true', help='voxelise the model and count separate pieces')
    args = parser.parse_args()
    logo = LogoT(args.src, args.scale, args.depth, args.chamfer, args.tolerance)
    writer = ObjWriterT()
    if args.printable:
        plate = plateMask(logo.ink, PLATE_MARGIN, PLATE_BRIDGE, UPSCALE)
        # The plate is the footprint now:  it, not the letters, rests on Y = 0.
        ys, xs = np.nonzero(plate)
        logo.setFootprint(xs.min() / UPSCALE, xs.max() / UPSCALE, ys.max() / UPSCALE)
        logo.buildPlate(writer, plate, -args.depth)
        logo.buildLetters(writer, True, -args.depth)
        materials = {BODIES[body]: DRAGON_MATERIALS[BODIES[body]] for body in BODIES}
        bodies = BODIES
    else:
        logo.buildLetters(writer, False, -args.depth)
        materials = {'logo': DRAGON_MATERIALS['dark']}
        bodies = {'logo': 'logo'}
    variant = 'printable ' if args.printable else ''
    header = ['Singe logotype, %sextruded from SingeText.jpeg by util/textModel.py' % variant,
              'Units: %g per image pixel; Y up, front faces toward +Z, resting on Y = 0, centred on X = 0' % args.scale]
    print('%s: %d letters, %d vertices, %d faces%s' % (args.out, len(logo.letters), len(writer.vertices), writer.faceCount, ', %d with holes (filled)' % logo.holes if logo.holes else ''))
    writeOutputs(writer, args.out, header, materials, bodies, args.stl, args.width, (logo.x1 - logo.x0) * args.scale, args.check, args.scale)