#!/usr/bin/env python3
"""Render where a gem lands on Armory's worn armour, without starting the game.

MIAPI places a module by composing its slot's transform with the ones above it
and then with the model part it names, so the only way to know whether a gem sits
on a pauldron or in mid air is to do that composition and look. This does it: it
reads Armory's own models out of its jar, applies the same chain MIAPI does, and
draws the result from as many angles as asked for.

    python3 tools/preview_armour.py <armory.jar> [out.png]

Two things it does not simplify away, because both change the answer:

*Perspective.* An orthographic view flatters a placement - a gem floating a pixel
off a plate looks welded to it head on. Views are taken from several yaws and
pitches with a real camera so the gap shows.

*MIAPI's arithmetic.* `Transform.merge` does not compose matrices; it multiplies
them and then decomposes the product back into translation, Euler angles and
scale. That is lossy the moment a rotation meets a non-uniform scale, which is
exactly what Armory's limb slots are, so the shear is dropped and the gem comes
out turned and stretched differently from what the matrix says. `--exact` renders
the composition MIAPI does not do, for comparison.

Everything is in model pixels, the units the module JSON is written in: +x is the
wearer's left, +y is down, -z is forward, and the origin is the base of the neck.
"""
import argparse, json, math, zipfile

# ---------------------------------------------------------------- linear algebra

def ident():
    return [[1.0 if r == c else 0.0 for c in range(4)] for r in range(4)]

def matmul(a, b):
    return [[sum(a[r][k] * b[k][c] for k in range(4)) for c in range(4)] for r in range(4)]

def apply(m, p):
    v = list(p) + [1.0]
    return tuple(sum(m[r][c] * v[c] for c in range(4)) for r in range(3))

def translate(t):
    m = ident()
    for i in range(3):
        m[i][3] = t[i]
    return m

def scale(s):
    m = ident()
    for i in range(3):
        m[i][i] = s[i]
    return m

def rot(axis, deg):
    a = math.radians(deg)
    c, s = math.cos(a), math.sin(a)
    m = ident()
    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."""
    g = lambda d, k, dflt: float(d.get(k, dflt)) if isinstance(d, dict) else dflt
    t, r, s = tr.get('translation', {}), tr.get('rotation', {}), tr.get('scale', {})
    m = translate((g(t, 'x', 0), g(t, 'y', 0), g(t, 'z', 0)))
    for axis in 'xyz':
        m = matmul(m, rot(axis, g(r, axis, 0)))
    return matmul(m, scale((g(s, 'x', 1), g(s, 'y', 1), g(s, 'z', 1))))

# ------------------------------------------------- MIAPI's lossy merge, reproduced

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

    Any shear in the product is silently discarded, which is the whole reason
    this function is here rather than a plain matrix multiply.
    """
    t = (m[0][3], m[1][3], m[2][3])
    cols = [[m[r][c] for r in range(3)] for c in range(3)]
    s = [math.sqrt(sum(v * v for v in col)) or 1.0 for col in cols]
    r = [[cols[c][row] / s[c] for c in range(3)] for row in range(3)]
    # R = Rx*Ry*Rz, so sin(y) is r[0][2] and the other two follow from it.
    y = math.asin(max(-1.0, min(1.0, r[0][2])))
    x = math.atan2(-r[1][2], r[2][2])
    z = math.atan2(-r[0][1], r[0][0])
    return {'translation': {'x': t[0], 'y': t[1], 'z': 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, exact=False):
    """MIAPI's Transform.merge: the parent applies first, then the product is
    decomposed and rebuilt. `exact` keeps the product instead."""
    product = matmul(child, parent)
    return product if exact else transform_matrix(decompose(product))

# ------------------------------------------------------------------- the figure

# HumanoidModel part pivots, the frame each origin resolves against. Checked
# against HumanoidModel.createMesh rather than remembered.
PIVOTS = {
    'head': (0, 0, 0), 'hat': (0, 0, 0), 'body': (0, 0, 0),
    'left_arm': (5, 2, 0), 'right_arm': (-5, 2, 0),
    'left_leg': (1.9, 12, 0), 'right_leg': (-1.9, 12, 0),
}

ARMORY = 'assets/miapi/models/item/armor/model'

def boxes_from_model(jar, path):
    with zipfile.ZipFile(jar) as z:
        data = json.loads(z.read(path))
    return [(tuple(e['from']), tuple(e['to'])) for e in data.get('elements', [])]

def place(boxes, chain, pivot):
    """Local boxes through the chain and into the part's frame, as 8 corners."""
    out = []
    for lo, hi in boxes:
        pts = []
        for i in range(8):
            corner = (hi[0] if i & 1 else lo[0], hi[1] if i & 2 else lo[1], hi[2] if i & 4 else lo[2])
            x, y, z = apply(chain, corner)
            pts.append((x + pivot[0], y + pivot[1], z + pivot[2]))
        out.append(pts)
    return out

def gem_box(half_xy, half_z=0.5):
    """A gem sprite as a slab centred on its own origin, the way MIAPI bakes it."""
    return [((-half_xy, -half_xy, -half_z), (half_xy, half_xy, half_z))]

# ---------------------------------------------------------------------- drawing

# The 6 faces of a box as indices into the 8 corners produced by place().
FACES = [(0, 1, 3, 2), (4, 5, 7, 6), (0, 1, 5, 4), (2, 3, 7, 6), (0, 2, 6, 4), (1, 3, 7, 5)]

def camera(yaw, pitch, dist, target=(0, 8, 0)):
    """World -> eye, looking at the figure's middle from yaw/pitch degrees out.

    Yaw 0 is the wearer's front. Minecraft faces a model down -z, so the camera
    starts on that side rather than on +z, which would label the back "front" and
    is exactly the sort of quiet inversion this tool exists to avoid.
    """
    m = translate((0, 0, -dist))
    m = matmul(m, rot('x', pitch))
    m = matmul(m, rot('y', yaw + 180))
    return matmul(m, translate((-target[0], -target[1], -target[2])))

def render(groups, out_path, views=None, size=300, fov=38.0):
    """One panel per view. `groups` is (fill, outline, label, [boxes])."""
    from PIL import Image, ImageDraw
    views = views or [('front', 0, 0), ('front-left', 35, 10), ('left', 90, 0),
                      ('above-left', 45, 35), ('back', 180, 0), ('below-left', 40, -30)]
    cols = min(3, len(views))
    rows = (len(views) + cols - 1) // cols
    img = Image.new('RGB', (size * cols, size * rows), (26, 26, 30))
    d = ImageDraw.Draw(img, 'RGBA')
    f = (size / 2) / math.tan(math.radians(fov) / 2)

    for i, (label, yaw, pitch) in enumerate(views):
        ox, oy = (i % cols) * size, (i // cols) * size
        view = camera(yaw, pitch, 46.0)

        def project(p):
            x, y, z = apply(view, p)
            # +y is down in model space, and the camera looks down -z.
            depth = max(0.1, -z)
            return (ox + size / 2 + x * f / depth, oy + size / 2 + y * f / depth), depth

        polys = []
        for fill, outline, _, boxes in groups:
            for pts in boxes:
                eye = [apply(view, p) for p in pts]
                if all(e[2] > -0.1 for e in eye):
                    continue  # entirely behind the camera
                for face in FACES:
                    quad = [project(pts[k])[0] for k in face]
                    z = sum(project(pts[k])[1] for k in face) / 4
                    polys.append((z, quad, fill, outline))
        for _, quad, fill, outline in sorted(polys, key=lambda t: -t[0]):
            d.polygon(quad, fill=fill, outline=outline)
        d.text((ox + 8, oy + 6), f"{label}  yaw {yaw} pitch {pitch}", fill=(190, 190, 200))
    img.save(out_path)
    return out_path

# ------------------------------------------------------------------- the scene

def scene(jar, repo, exact=False):
    """Armory's heavy pieces plus a gem in each of this mod's four sockets."""
    def socket(name):
        with open(f'{repo}/src/main/resources/packs/armory/data/cmmodular/miapi'
                  f'/modules/armor/socket/{name}.json') as fh:
            return json.load(fh)['data']['replace']['slots']['gem']['transform']

    with zipfile.ZipFile(jar) as z:
        chest = json.loads(z.read('data/tm_armory/miapi/modules/armor/chestplate.json'))['slots']
        pants = json.loads(z.read('data/tm_armory/miapi/modules/armor/pants.json'))['slots']

    plate, gems = [], []
    for part, chain_json, model, sock in (
        ('left_arm',  chest['arm_left'],  'arm_left/heavy',  'arm_left'),
        ('right_arm', chest['arm_right'], 'arm_right/heavy', 'arm_right'),
        ('left_leg',  pants['leg_left'],  'leg_left/heavy',  'leg_left'),
        ('right_leg', pants['leg_right'], 'leg_right/heavy', 'leg_right'),
    ):
        limb = transform_matrix(chain_json['transform'])
        plate += place(boxes_from_model(jar, f'{ARMORY}/{model}/default.json'), limb, PIVOTS[part])
        gems += place(gem_box(1.0), merge(transform_matrix(socket(sock)), limb, exact), PIVOTS[part])

    # Armory's own socketed chest, as a control: its gem is known to sit on the
    # sternum, so if this lands there the chain above is being modelled right.
    body = transform_matrix(chest['chest_front']['transform'])
    plate += place(boxes_from_model(jar, f'{ARMORY}/chest_front/socket/default.json'), body, PIVOTS['body'])
    gems += place(gem_box(1.0), merge(transform_matrix(
        {'translation': {'y': -1}, 'scale': {'x': 1.1, 'y': 1.1, 'z': 1.1}}), body, exact), PIVOTS['body'])
    return plate, gems

def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('jar')
    ap.add_argument('out', nargs='?', default='armour-preview.png')
    ap.add_argument('--exact', action='store_true',
                    help="compose matrices properly instead of reproducing MIAPI's lossy merge")
    ap.add_argument('--repo', default='.')
    args = ap.parse_args()
    plate, gems = scene(args.jar, args.repo, args.exact)
    print(render([((78, 62, 104, 255), (120, 104, 150, 255), 'plate', plate),
                  ((86, 216, 122, 255), (235, 255, 240, 255), 'gem', gems)], args.out))

if __name__ == '__main__':
    main()