#!/usr/bin/env python3
"""Minecraft and MIAPI geometry, read out of jars and turned into textured quads.

This is the part of the armour tooling that has to be *right* rather than
merely plausible: everything downstream is just a camera pointed at whatever
this module says the shape is. So the two UV conventions Minecraft uses are
both implemented here, separately and by name, rather than being averaged into
one that is wrong for half the models:

*Entity cubes* (`ModelPart.Cube`) are what vanilla armour is. Boxes are given a
single texture offset and the six faces are laid out around it in the familiar
cross; +y is down, because entity model space is flipped once more at draw
time. `vanilla_armour` builds these.

*JSON model faces* (`FaceBakery`) are what every MIAPI module is - a Blockbench
item model with a `uv` rectangle written out per face. The two conventions
disagree about which end of the rectangle is which on four of the six faces,
which is exactly the sort of difference that survives a careless eyeball on a
symmetrical breastplate and then ruins a pauldron.

Everything is in model pixels: +x is the wearer's left, +y is down, -z is
forward, and the origin of each part is its `HumanoidModel` pivot. Nothing
converts between "item space" and "entity space", because MIAPI does not
either - Armory's slot transforms carry an explicit `"rotation": {"z": 180}`
to flip a model onto a body part, and that flip is the whole conversion.
"""
from __future__ import annotations

import colorsys
import io
import json
import math
import os
import zipfile
from dataclasses import dataclass

import numpy as np
from PIL import Image

# --------------------------------------------------------------------- resources


class Resources:
    """Jars and loose directories, searched in the order they were given.

    A resource pack and a mod jar are the same thing to this: somewhere an
    `assets/<namespace>/...` path can be read from. Directories are listed
    first when the same file is in both, so the repo's own assets win over the
    copy inside a built jar.
    """

    def __init__(self, sources):
        self.dirs, self.zips, self.names = [], [], []
        for src in sources:
            src = os.path.expanduser(str(src))
            if os.path.isdir(src):
                self.dirs.append(src)
            elif os.path.isfile(src):
                self.zips.append(zipfile.ZipFile(src))
            else:
                raise FileNotFoundError(src)
            self.names.append(os.path.basename(src.rstrip('/')))

    def read(self, path):
        """Bytes at an exact `assets/...` or `data/...` path."""
        for d in self.dirs:
            full = os.path.join(d, path)
            if os.path.isfile(full):
                return open(full, 'rb').read()
        for z in self.zips:
            try:
                return z.read(path)
            except KeyError:
                pass
        raise KeyError(path)

    def json(self, path):
        return json.loads(self.read(path))

    # MIAPI writes model references two ways. A module's `model` entry gives a
    # whole path - "miapi:models/item/foo/[material.texture].json" - while a
    # model's own `parent` gives the short form the game uses, "item/generated".
    def model_path(self, ref, variant='default'):
        ns, _, rest = ref.partition(':')
        if not _:
            ns, rest = 'minecraft', ref
        rest = rest.replace('[material.texture]', variant)
        if not rest.startswith('models/'):
            rest = 'models/' + rest
        if not rest.endswith('.json'):
            rest += '.json'
        return f'assets/{ns}/{rest}'

    def model(self, ref, variant='default'):
        """A model with its parent chain already folded in."""
        model = self.json(self.model_path(ref, variant))
        chain = [model]
        seen = set()
        while 'parent' in chain[-1]:
            parent = chain[-1]['parent']
            if parent in seen:
                break
            seen.add(parent)
            try:
                chain.append(self.json(self.model_path(parent, variant)))
            except KeyError:
                # item/generated and item/handheld are builtins with no file.
                break
        out = {'textures': {}, 'parents': [c.get('parent') for c in chain]}
        for part in reversed(chain):
            out['textures'].update(part.get('textures', {}))
            for key in ('elements', 'texture_size', 'display'):
                if key in part:
                    out[key] = part[key]
        return out

    def paths(self, prefix='', suffix='', loose=False):
        """Every path in the sources between a prefix and a suffix, sorted.

        For finding what is there rather than reading what you already knew
        about. Duplicates collapse the way `read` resolves them - one path,
        whichever source answers for it first. `loose` limits the answer to
        files on disk, for a caller that means the working tree rather than
        whatever a build happened to zip up.
        """
        found = set()
        for d in self.dirs:
            root = os.path.join(d, prefix)
            start = root if os.path.isdir(root) else os.path.dirname(root)
            for base, _dirs, files in os.walk(start):
                rel = os.path.relpath(base, d)
                found.update(f'{rel}/{f}' for f in files
                             if f'{rel}/{f}'.startswith(prefix) and f.endswith(suffix))
        for z in ([] if loose else self.zips):
            found.update(n for n in z.namelist()
                         if n.startswith(prefix) and n.endswith(suffix))
        return sorted(found)

    def image(self, ref):
        """A texture as RGBA. `ref` is a namespaced texture id, no extension."""
        ns, _, rest = ref.partition(':')
        if not _:
            ns, rest = 'minecraft', ref
        if rest.startswith('textures/'):
            rest = rest[len('textures/'):]
        data = self.read(f'assets/{ns}/textures/{rest}.png')
        return Image.open(io.BytesIO(data)).convert('RGBA')


# ------------------------------------------------------------------- transforms
#
# MIAPI's Transform, reproduced including the part of it that loses information.

def ident():
    return np.eye(4)


def translate(t):
    m = np.eye(4)
    m[:3, 3] = t
    return m


def scale(s):
    m = np.eye(4)
    m[0, 0], m[1, 1], m[2, 2] = s
    return m


def rot(axis, deg):
    a = math.radians(deg)
    c, s = math.cos(a), math.sin(a)
    m = np.eye(4)
    if axis == 'x':
        m[1, 1], m[1, 2], m[2, 1], m[2, 2] = c, -s, s, c
    elif axis == 'y':
        m[0, 0], m[0, 2], m[2, 0], m[2, 2] = c, s, -s, c
    else:
        m[0, 0], m[0, 1], m[1, 0], m[1, 1] = c, -s, s, c
    return m


def transform_matrix(tr):
    """MIAPI's `Transform.toMatrix`: T * Rx * Ry * Rz * S, translation in pixels."""
    tr = tr or {}

    def g(key, axis, dflt):
        sub = tr.get(key)
        return float(sub.get(axis, dflt)) if isinstance(sub, dict) else dflt

    m = translate([g('translation', a, 0.0) for a in 'xyz'])
    for axis in 'xyz':
        m = m @ rot(axis, g('rotation', axis, 0.0))
    return m @ scale([g('scale', a, 1.0) for a in 'xyz'])


def decompose(m):
    """MIAPI's `Transform.fromMatrix`: translation, XYZ Euler, per-column scale.

    Shear in the product is dropped on the floor. That is not a bug here - it
    is what the mod does, and reproducing it is the only way the preview agrees
    with the game once a rotation meets a non-uniform scale.
    """
    t = m[:3, 3]
    cols = [m[:3, c] for c in range(3)]
    s = [float(np.linalg.norm(c)) or 1.0 for c in cols]
    r = np.column_stack([cols[c] / s[c] for c in range(3)])
    y = math.asin(max(-1.0, min(1.0, float(r[0, 2]))))
    x = math.atan2(-float(r[1, 2]), float(r[2, 2]))
    z = math.atan2(-float(r[0, 1]), float(r[0, 0]))
    return {'translation': {'x': float(t[0]), 'y': float(t[1]), 'z': float(t[2])},
            'rotation': {'x': math.degrees(x), 'y': math.degrees(y), 'z': math.degrees(z)},
            'scale': {'x': s[0], 'y': s[1], 'z': s[2]}}


def merge(parent, child, lossy=False):
    """MIAPI's `Transform.merge`.

    In 1.21 this is a plain matrix multiply - `Transform` holds a `Matrix4f`
    and `merge` returns `new Transform(parent.matrix.mul(child.matrix))`, with
    no round trip through Euler angles. The parent is the transform already
    accumulated and the child the one being added, so the child applies first.

    `lossy` reproduces the older behaviour, where the product was decomposed
    back into translation, Euler angles and scale before being stored - which
    silently drops the shear that appears the moment a rotation meets a
    non-uniform scale. Armory's limb slots are exactly that, so the two answers
    differ there and it is worth being able to see both.
    """
    product = child @ parent
    return transform_matrix(decompose(product)) if lossy else product


# `HumanoidModel.createMesh` pivots - the frame each `origin` resolves against.
PIVOTS = {
    'head': (0.0, 0.0, 0.0),
    'hat': (0.0, 0.0, 0.0),
    'body': (0.0, 0.0, 0.0),
    'item': (0.0, 0.0, 0.0),
    'left_arm': (5.0, 2.0, 0.0),
    'right_arm': (-5.0, 2.0, 0.0),
    'left_leg': (1.9, 12.0, 0.0),
    'right_leg': (-1.9, 12.0, 0.0),
}


# ----------------------------------------------------------------------- quads


@dataclass
class Quad:
    """Four corners, four texture coordinates, and the texture they index.

    `shape` is which box of its model the face belongs to and `face` is which
    of the six it is. Neither matters for drawing a textured model, but both
    are what lets the untextured view and the unwrap template agree on a
    colour, which is the only thing making one a legend for the other.
    """
    pts: np.ndarray          # (4, 3) model pixels
    uv: np.ndarray           # (4, 2) normalised, v measured down from the top
    texture: str | None      # namespaced texture id
    shape: int = 0
    face: str = ''

    def transformed(self, m, offset=(0.0, 0.0, 0.0)):
        pts = np.column_stack([self.pts, np.ones(4)]) @ m.T
        return Quad(pts[:, :3] + np.asarray(offset, float), self.uv, self.texture,
                    self.shape, self.face)


# ------------------------------------------------------------------- colouring
#
# One hue per shape, one shade per face. The hue says which box you are looking
# at and the shade says which side of it, so a patch of texture can be found on
# the model - and a face drawn on the wrong patch shows up as the wrong shade of
# the right colour rather than as something that looks fine.

# Saturation and value per face. Front and back are the vivid pair because they
# are what you look at most; up and down are pushed to the ends of the value
# range so a box read from above or below is never ambiguous. Entity cubes name
# their vertical faces differently, and both names are here rather than
# translated, so neither convention has to know about the other.
FACE_SHADES = {
    'north': (0.90, 0.98), 'south': (0.90, 0.60),
    'east':  (0.55, 0.90), 'west':  (0.55, 0.68),
    'up':    (0.26, 1.00), 'down':  (1.00, 0.42),
    'top':   (0.26, 1.00), 'bottom': (1.00, 0.42),
}


def shape_hues(count):
    """The colour wheel split evenly, one slice per shape."""
    count = max(1, int(count))
    return [i / count for i in range(count)]


def face_colour(shape, count, face):
    """The colour of one face of one shape, as floats in 0..1."""
    hues = shape_hues(count)
    hue = hues[int(shape) % len(hues)]
    sat, val = FACE_SHADES.get(face, (0.70, 0.80))
    return colorsys.hsv_to_rgb(hue, sat, val)


# uv index 0..3 is (u1,v1), (u1,v2), (u2,v2), (u2,v1) - `BlockFaceUV.getU/getV`.
# Each entry picks x, y and z from (from, to) per corner: 0 is `from`, 1 is `to`.
_JSON_FACE = {
    'north': ((1, 1, 0), (1, 0, 0), (0, 0, 0), (0, 1, 0)),
    'south': ((0, 1, 1), (0, 0, 1), (1, 0, 1), (1, 1, 1)),
    'west':  ((0, 1, 0), (0, 0, 0), (0, 0, 1), (0, 1, 1)),
    'east':  ((1, 1, 1), (1, 0, 1), (1, 0, 0), (1, 1, 0)),
    'up':    ((0, 1, 0), (0, 1, 1), (1, 1, 1), (1, 1, 0)),
    'down':  ((0, 0, 1), (0, 0, 0), (1, 0, 0), (1, 0, 1)),
}

# The two axes a face's uv rectangle runs along, for the auto-uv Minecraft
# generates when a face omits `uv`: (u axis, v axis) as (index, flipped).
_AUTO_UV = {
    'north': ((0, True), (1, True)), 'south': ((0, False), (1, True)),
    'west':  ((2, False), (1, True)), 'east':  ((2, True), (1, True)),
    'up':    ((0, False), (2, False)), 'down':  ((0, False), (2, True)),
}


def _resolve_texture(textures, key):
    """Follow `#2` -> `#layer0` -> `miapi:item/...` to a real texture id."""
    seen = 0
    while isinstance(key, str) and key.startswith('#') and seen < 8:
        key = textures.get(key[1:])
        seen += 1
    return key


def _element_matrix(el):
    """A JSON element's own `rotation`, about its own origin."""
    r = el.get('rotation')
    if not r:
        return np.eye(4)
    axis, angle = r.get('axis', 'y'), float(r.get('angle', 0.0))
    origin = np.asarray(r.get('origin', [0, 0, 0]), float)
    m = translate(origin) @ rot(axis, angle) @ translate(-origin)
    if r.get('rescale') and angle:
        f = 1.0 / math.cos(math.radians(abs(angle)))
        s = [f, f, f]
        s['xyz'.index(axis)] = 1.0
        m = m @ (translate(origin) @ scale(s) @ translate(-origin))
    return m


def model_quads(model, res=None, y_up=False):
    """Every drawable face of a parsed JSON model, in the model's own pixels.

    A model with no `elements` whose ancestry runs through `item/generated` is
    a flat sprite, and is turned into a two-sided slab centred on the origin -
    which is what a gemstone is, and why a gem's position is a point rather
    than a plate.
    """
    textures = model.get('textures', {})
    elements = model.get('elements')
    if not elements:
        return _sprite_quads(textures, res, y_up)

    tw, th = (model.get('texture_size') or [16, 16])[:2]
    quads = []
    for index, el in enumerate(elements):
        lo = np.asarray(el['from'], float)
        hi = np.asarray(el['to'], float)
        bounds = np.column_stack([lo, hi])          # (3, 2): axis -> (from, to)
        m = _element_matrix(el)
        for name, face in (el.get('faces') or {}).items():
            picks = _JSON_FACE.get(name)
            if picks is None:
                continue
            uv = face.get('uv')
            if uv is None:
                uv = _auto_uv(name, lo, hi)
            u1, v1, u2, v2 = (float(x) for x in uv)
            corners = [(u1, v1), (u1, v2), (u2, v2), (u2, v1)]
            turns = int(face.get('rotation', 0) // 90) % 4
            if turns:
                corners = corners[turns:] + corners[:turns]
            pts = np.array([[bounds[a][p[a]] for a in range(3)] for p in picks], float)
            pts = (np.column_stack([pts, np.ones(4)]) @ m.T)[:, :3]
            quads.append(Quad(pts,
                              np.array([[u / tw, v / th] for u, v in corners]),
                              _resolve_texture(textures, face.get('texture', '#0')),
                              index, name))
    return quads


def _auto_uv(name, lo, hi):
    (ui, uflip), (vi, vflip) = _AUTO_UV[name]
    a, b = (16 - hi[ui], 16 - lo[ui]) if uflip else (lo[ui], hi[ui])
    c, d = (16 - hi[vi], 16 - lo[vi]) if vflip else (lo[vi], hi[vi])
    return [a, c, b, d]


def _sprite_quads(textures, res=None, y_up=False):
    """An `item/generated` sprite: a two-sided slab, front and back.

    Trimmed to the sprite's opaque pixels rather than left at the full 16x16,
    because that is what the module actually is - a medium gemstone is two
    pixels of gem in the middle of fourteen of nothing, and a placement judged
    against the empty square around it is judged against the wrong thing.

    Centred on the origin, because MIAPI places one by its middle: a gemstone
    sits where its slot's translation points, not 8 pixels down and left of it.

    `y_up` builds it in item space instead. Worn armour is drawn with +y down -
    that is what the `"rotation": {"z": 180}` on every armour slot is for - but
    an inventory icon is a plain item model, where +y is up and the sprite's top
    row is at the top. Armory's belt gem settles which is which: its slot names
    no origin, so it reaches the icon, and `y: 4` puts it on the buckle, which
    the art draws four pixels above the middle.
    """
    tex = _resolve_texture(textures, '#layer0') or _resolve_texture(textures, '#0')
    x0, y0, x1, y1 = -8.0, -8.0, 8.0, 8.0
    u0, v0, u1, v1 = 0.0, 0.0, 1.0, 1.0
    if res is not None and tex:
        try:
            image = res.image(tex)
        except KeyError:
            image = None
        box = image.split()[3].getbbox() if image is not None else None
        if box:
            w, h = image.size
            u0, u1 = box[0] / w, box[2] / w
            v0, v1 = box[1] / h, box[3] / h
            # The sprite spans 16 model pixels whatever its resolution.
            x0, x1 = u0 * 16.0 - 8.0, u1 * 16.0 - 8.0
            y0, y1 = v0 * 16.0 - 8.0, v1 * 16.0 - 8.0

    if y_up:
        y0, y1 = -y1, -y0
        uv_top, uv_bottom = v1, v0          # texture top is now the larger y
    else:
        uv_top, uv_bottom = v0, v1
    v0, v1 = uv_top, uv_bottom

    front = np.array([[x0, y0, -0.5], [x0, y1, -0.5], [x1, y1, -0.5], [x1, y0, -0.5]])
    back = np.array([[x1, y0, 0.5], [x1, y1, 0.5], [x0, y1, 0.5], [x0, y0, 0.5]])
    uv_f = np.array([[u0, v0], [u0, v1], [u1, v1], [u1, v0]])
    uv_b = np.array([[u1, v0], [u1, v1], [u0, v1], [u0, v0]])
    return [Quad(front, uv_f, tex, 0, 'north'), Quad(back, uv_b, tex, 0, 'south')]


# ------------------------------------------------------------- vanilla armour
#
# `ModelPart.Cube`'s layout, which is the other convention and the reference
# this whole tool is checked against.

# Per face: the corner picks, and which slice of the texture cross it takes.
# Regions are named by the offsets vanilla computes - f9..f14 across, f15..f17
# down - so the table can be read against the source it came from.
_CUBE_FACE = {
    'top':    (((1, 0, 1), (0, 0, 1), (0, 0, 0), (1, 0, 0)), ('f10', 'f15', 'f11', 'f16')),
    'bottom': (((1, 1, 0), (0, 1, 0), (0, 1, 1), (1, 1, 1)), ('f11', 'f16', 'f12', 'f15')),
    'west':   (((0, 0, 0), (0, 0, 1), (0, 1, 1), (0, 1, 0)), ('f9', 'f16', 'f10', 'f17')),
    'north':  (((1, 0, 0), (0, 0, 0), (0, 1, 0), (1, 1, 0)), ('f10', 'f16', 'f11', 'f17')),
    'east':   (((1, 0, 1), (1, 0, 0), (1, 1, 0), (1, 1, 1)), ('f11', 'f16', 'f13', 'f17')),
    'south':  (((0, 0, 1), (1, 0, 1), (1, 1, 1), (0, 1, 1)), ('f13', 'f16', 'f14', 'f17')),
}


def cube_quads(origin, size, tex_offs, grow=0.0, mirror=False,
               tex_size=(64, 32), texture=None, shape=0):
    """One entity cube, the way `ModelPart.Cube` builds it.

    `origin` and `size` are the arguments of `CubeListBuilder.addBox`, `grow`
    the `CubeDeformation`. The uv rectangle is measured from the undeformed
    size, which is why an inflated armour layer still lines up with the skin
    it is drawn over.
    """
    ox, oy, oz = (float(v) for v in origin)
    dx, dy, dz = (float(v) for v in size)
    x1, y1, z1 = ox - grow, oy - grow, oz - grow
    x2, y2, z2 = ox + dx + grow, oy + dy + grow, oz + dz + grow
    if mirror:
        x1, x2 = x2, x1
    bounds = ((x1, x2), (y1, y2), (z1, z2))

    u, v = float(tex_offs[0]), float(tex_offs[1])
    reg = {'f9': u, 'f10': u + dz, 'f11': u + dz + dx, 'f12': u + dz + dx + dx,
           'f13': u + dz + dx + dz, 'f14': u + dz + dx + dz + dx,
           'f15': v, 'f16': v + dz, 'f17': v + dz + dy}
    tw, th = tex_size

    quads = []
    for name, (picks, (ua, va, ub, vb)) in _CUBE_FACE.items():
        u0, v0, u1_, v1_ = reg[ua], reg[va], reg[ub], reg[vb]
        # vertex order is [0]->(u1,v0) [1]->(u0,v0) [2]->(u0,v1) [3]->(u1,v1)
        corners = [(u1_, v0), (u0, v0), (u0, v1_), (u1_, v1_)]
        pts = np.array([[bounds[a][p[a]] for a in range(3)] for p in picks], float)
        quads.append(Quad(pts, np.array([[cu / tw, cv / th] for cu, cv in corners]),
                          texture, shape, name))
    return quads


# `HumanoidModel.createMesh`: box origin, box size, texture offset, mirrored.
HUMANOID = {
    'head':      ((-4, -8, -4), (8, 8, 8), (0, 0), False),
    'hat':       ((-4, -8, -4), (8, 8, 8), (32, 0), False),
    'body':      ((-4, 0, -2), (8, 12, 4), (16, 16), False),
    'right_arm': ((-3, -2, -2), (4, 12, 4), (40, 16), False),
    'left_arm':  ((-1, -2, -2), (4, 12, 4), (40, 16), True),
    'right_leg': ((-2, 0, -2), (4, 12, 4), (0, 16), False),
    'left_leg':  ((-2, 0, -2), (4, 12, 4), (0, 16), True),
}

# Which parts each vanilla armour slot draws, and how far each layer is
# inflated - `HumanoidArmorModel`'s inner and outer `CubeDeformation`.
ARMOUR_PIECES = {
    'helmet':     (('head', 'hat'), 1.0, 1),
    'chestplate': (('body', 'left_arm', 'right_arm'), 1.0, 1),
    'leggings':   (('body', 'left_leg', 'right_leg'), 0.5, 2),
    'boots':      (('left_leg', 'right_leg'), 1.0, 1),
}


def vanilla_armour(piece, texture, tex_size=(64, 32)):
    """A vanilla armour piece as world-space quads, ready to draw.

    The point of this is that its answer is already known. Iron leggings look
    like iron leggings or the renderer is wrong, and there is no MIAPI, no
    transform stack and no module JSON in the way of finding out which.
    """
    parts, grow, _ = ARMOUR_PIECES[piece]
    quads = []
    for shape, part in enumerate(parts):
        origin, size, offs, mirror = HUMANOID[part]
        # The hat is a second, slightly larger skin on the head; a helmet is
        # already the inflated head, so it takes vanilla's extra 0.5 as well.
        g = grow + 0.5 if part == 'hat' else grow
        pivot = PIVOTS[part]
        quads += [q.transformed(np.eye(4), pivot)
                  for q in cube_quads(origin, size, offs, g, mirror, tex_size,
                                      texture, shape)]
    return quads


def humanoid_body(texture=None, tex_size=(64, 64)):
    """The wearer, as a reference figure to judge a placement against."""
    quads = []
    for shape, (part, (origin, size, offs, mirror)) in enumerate(HUMANOID.items()):
        if part == 'hat':
            continue
        quads += [q.transformed(np.eye(4), PIVOTS[part])
                  for q in cube_quads(origin, size, offs, 0.0, mirror, tex_size,
                                      texture, shape)]
    return quads


# ------------------------------------------------------------------- unwrapping
#
# Box UV, the layout Blockbench calls "Box UV" and Minecraft's entity cubes use:
#
#         [up ][dn ]
#    [west][north][east][south]
#
# A box needs (2*dz + 2*dx) across and (dz + dy) down. Faces are written back as
# explicit per-face `uv` rectangles, because that is the only thing the JSON
# model format can say - `texture_size` scales them, nothing else.

def texel_size(lo, hi):
    """A box's dimensions in whole texels, which is the only size a patch has.

    Rounded up and never zero: a bezel half a pixel thick still needs a row of
    texture to be painted on, and a face allotted 0.5 of a texel is a face that
    cannot be drawn and, on the way there, a rectangle PIL refuses to fill.
    """
    return tuple(max(1, int(math.ceil(round(abs(float(b) - float(a)), 4))))
                 for a, b in zip(lo, hi))


def net_size(lo, hi):
    """The width and height one box's net needs, in texture pixels."""
    dx, dy, dz = texel_size(lo, hi)
    return (2 * dz + 2 * dx, dz + dy)


def _net_faces(u, v, dx, dy, dz):
    """Where each face lands in a net whose top-left corner is (u, v)."""
    return {
        'up':    (u + dz, v, u + dz + dx, v + dz),
        'down':  (u + dz + dx, v, u + dz + dx + dx, v + dz),
        'west':  (u, v + dz, u + dz, v + dz + dy),
        'north': (u + dz, v + dz, u + dz + dx, v + dz + dy),
        'east':  (u + dz + dx, v + dz, u + dz + dx + dz, v + dz + dy),
        'south': (u + dz + dx + dz, v + dz, u + dz + dx + dz + dx, v + dz + dy),
    }


def unwrap(elements, texture=None, padding=0, atlas_width=None):
    """Give every face of every element its own patch of texture.

    Shelf-packs one net per element, tallest first, and writes the resulting
    rectangles back into each face's `uv`. Returns the atlas size the caller
    should put in `texture_size`, and the placement of each element so a
    template can be drawn to match.

    Nothing here is clever about sharing texture between identical faces. That
    is deliberate: a shared patch is a patch you cannot edit on one face
    without editing the other, and this exists so that faces can be painted.
    """
    nets = []
    for i, el in enumerate(elements):
        lo, hi = el['from'], el['to']
        dx, dy, dz = texel_size(lo, hi)
        w, h = net_size(lo, hi)
        nets.append({'index': i, 'dx': dx, 'dy': dy, 'dz': dz,
                     'w': w + 2 * padding, 'h': h + 2 * padding})

    widest = max((n['w'] for n in nets), default=1.0)
    total = sum(n['w'] * n['h'] for n in nets) or 1.0
    if atlas_width is None:
        # Wide enough for the widest net, and roughly square overall.
        atlas_width = max(widest, math.sqrt(total) * 1.3)
        atlas_width = 1 << max(0, math.ceil(math.log2(max(1.0, atlas_width))))

    # Shelves: tallest first so a short net never strands a tall one.
    shelf_x, shelf_y, shelf_h = 0.0, 0.0, 0.0
    for net in sorted(nets, key=lambda n: -n['h']):
        if shelf_x + net['w'] > atlas_width and shelf_x > 0:
            shelf_y += shelf_h
            shelf_x, shelf_h = 0.0, 0.0
        net['u'], net['v'] = shelf_x + padding, shelf_y + padding
        shelf_x += net['w']
        shelf_h = max(shelf_h, net['h'])
    height = shelf_y + shelf_h
    atlas_height = 1 << max(0, math.ceil(math.log2(max(1.0, height))))

    for net in nets:
        el = elements[net['index']]
        rects = _net_faces(net['u'], net['v'], net['dx'], net['dy'], net['dz'])
        faces = el.setdefault('faces', {})
        for name, rect in rects.items():
            face = faces.get(name)
            if face is None:
                face = {'texture': texture or '#0'}
                faces[name] = face
            face['uv'] = [round(c, 4) for c in rect]
            face.pop('rotation', None)
            if texture:
                face['texture'] = texture
    return (atlas_width, atlas_height), nets


def unwrap_template(size, nets, elements, labelled=True):
    """A painting guide for an unwrap: one coloured, labelled patch per face.

    The colours are the ones the untextured viewport uses - hue per shape,
    shade per face - so the model on screen is the legend for this sheet.
    Find the colour on the model, find the same colour here, and that is the
    patch to paint.

    Doubles as the check on the unwrap itself. Render a model with this as its
    texture and every face should show its own colour, its own letter, and the
    letter the right way up; anything else is a uv routed to the wrong place.
    """
    from PIL import Image, ImageDraw

    img = Image.new('RGBA', (int(size[0]), int(size[1])), (0, 0, 0, 0))
    d = ImageDraw.Draw(img)
    count = len(elements)
    for net in nets:
        rects = _net_faces(net['u'], net['v'], net['dx'], net['dy'], net['dz'])
        for name, (u0, v0, u1, v1) in rects.items():
            if u1 <= u0 or v1 <= v0:
                continue
            tint = tuple(int(round(c * 255))
                         for c in face_colour(net['index'], count, name))
            d.rectangle([u0, v0, u1 - 1, v1 - 1], fill=tint + (255,))
            # A darker top edge and left edge, so the patch has an up and a left.
            shade = tuple(int(c * 0.62) for c in tint)
            d.line([(u0, v0), (u1 - 1, v0)], fill=shade + (255,))
            d.line([(u0, v0), (u0, v1 - 1)], fill=shade + (255,))
            if labelled and (u1 - u0) >= 3 and (v1 - v0) >= 5:
                d.text((u0 + 1, v0 + 1), name[0].upper(), fill=(20, 20, 24, 255))
    return img