"""Render the game's survivor (the male player body, dressed) walking and standing, into sprites for the studio's
explore mode. Read from the player's own install, like the cars (render_vehicles.py):

  models_X/Skinned/MaleBody.x            the skinned body: bone hierarchy, mesh, skin weights
  models_X/Skinned/Clothes/Bob_Trousers.x  jeans (a separate skinned mesh); the body under them is masked away
  models_X/Skinned/Hair/Bob_Hair_*.x     hair
  anims_X/Bob/Bob_Walk.x, Bob_Idle.x      the walk cycle and the idle pose
  textures/Body, textures/Clothes, textures/F_Hair_White.png
A T-shirt and trainers are texture layers painted over the body texture, as the game does it.

DirectX .x text files use row vectors: a vertex v is posed as v * Offset(bone) * World(bone), World = Local * World(parent).
"""
from pathlib import Path
import re, sys
import numpy as np
from PIL import Image

import render_vehicles as RV

TOKEN = re.compile(r'"[^"]*"|<[^>]*>|[{};,]|[-+]?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?|[A-Za-z_][\w.\-]*')
DIRS = {'E': (1, 0), 'SE': (1, 1), 'S': (0, 1), 'SW': (-1, 1), 'W': (-1, 0), 'NW': (-1, -1), 'N': (0, -1), 'NE': (1, -1)}
HEIGHT = 1.78                 # tiles: a grown survivor next to a 2.45-tile storey, as in the game
WALK_FRAMES = 8


class Node:
    __slots__ = ('kind', 'name', 'data', 'kids', 'refs')

    def __init__(self, kind, name):
        self.kind, self.name, self.data, self.kids, self.refs = kind, name, [], [], []


def parse_x(path):
    text = Path(path).read_text(errors='ignore')
    toks = TOKEN.findall(text.split('\n', 1)[1] if text.startswith('xof') else text)
    i, n = 0, len(toks)
    root = Node('root', '')

    def block(node):
        nonlocal i
        while i < n:
            t = toks[i]
            if t == '}':
                i += 1; return
            if t == '{':                                # a reference: { Name }
                node.refs.append(toks[i + 1]); i += 3; continue
            if t in (';', ',') or t.startswith('<'):
                i += 1; continue
            if t.startswith('"'):
                node.data.append(t[1:-1]); i += 1; continue
            if t[0].isdigit() or t[0] in '+-.':
                node.data.append(float(t)); i += 1; continue
            # an identifier starts a child object: Kind [Name] {
            kind = t; i += 1; name = ''
            if toks[i] != '{':
                name = toks[i]; i += 1
            i += 1                                      # '{'
            if kind == 'template':
                depth = 1
                while depth:
                    depth += {'{': 1, '}': -1}.get(toks[i], 0); i += 1
                continue
            child = Node(kind, name); node.kids.append(child); block(child)
    block(root)
    return root


def walk(node):
    yield node
    for k in node.kids:
        yield from walk(k)


def mat(vals):
    return np.array(vals[:16], np.float64).reshape(4, 4)


def skeleton(root):
    """{bone: (parent, local 4x4)} in hierarchy order."""
    out = {}
    def rec(node, parent):
        for k in node.kids:
            if k.kind == 'Frame':
                ftm = next((c for c in k.kids if c.kind == 'FrameTransformMatrix'), None)
                out[k.name] = (parent, mat(ftm.data) if ftm else np.eye(4))
                rec(k, k.name)
    rec(root, None)
    return out


def meshes(root):
    """Every Mesh: vertices, triangles (fans), uv per vertex, texture file, skin [(bone, idx, w, offset)]."""
    out = []
    for m in walk(root):
        if m.kind != 'Mesh':
            continue
        d = m.data; nv = int(d[0]); V = np.array(d[1:1 + 3 * nv]).reshape(nv, 3)
        k = 1 + 3 * nv; nf = int(d[k]); k += 1; tris = []
        for _ in range(nf):
            c = int(d[k]); idx = [int(x) for x in d[k + 1:k + 1 + c]]; k += 1 + c
            tris += [[idx[0], idx[j], idx[j + 1]] for j in range(1, c - 1)]
        uv = None; tex = None; skin = []
        for c in m.kids:
            if c.kind == 'MeshTextureCoords':
                nt = int(c.data[0]); uv = np.array(c.data[1:1 + 2 * nt]).reshape(nt, 2)
            elif c.kind == 'SkinWeights':
                bone = c.data[0]; cnt = int(c.data[1]); idx = np.array(c.data[2:2 + cnt], int)
                w = np.array(c.data[2 + cnt:2 + 2 * cnt]); off = mat(c.data[2 + 2 * cnt:])
                skin.append((bone, idx, w, off))
        for c in walk(m):
            if c.kind == 'TextureFilename':
                tex = c.data[0]
            if c.kind == 'MeshMaterialList' and c.refs and tex is None:
                mat_ = next((q for q in walk(root) if q.kind == 'Material' and q.name == c.refs[0]), None)
                tf = mat_ and next((q for q in walk(mat_) if q.kind == 'TextureFilename'), None)
                tex = tf.data[0] if tf else None
        out.append({'V': V, 'T': np.array(tris, int), 'uv': uv, 'tex': tex, 'skin': skin})
    return out


def quat(w, x, y, z):
    # rotation key (w, x, y, z) as the matrix the file's own FrameTransformMatrix rows use (checked against the root bone)
    return np.array([[1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y), 0],
                     [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x), 0],
                     [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y), 0],
                     [0, 0, 0, 1]])


def animation(root):
    """{bone: {'R': [(t, q)], 'S': [(t, v)], 'T': [(t, v)]}} and the clip length in ticks."""
    out, end = {}, 0.0
    for a in walk(root):
        if a.kind != 'Animation' or not a.refs:
            continue
        keys = {}
        for k in a.kids:
            if k.kind != 'AnimationKey':
                continue
            d = k.data; kt = int(d[0]); nk = int(d[1]); p = 2; seq = []
            for _ in range(nk):
                t = d[p]; cnt = int(d[p + 1]); seq.append((t, np.array(d[p + 2:p + 2 + cnt]))); p += 2 + cnt; end = max(end, t)
            keys[{0: 'R', 1: 'S', 2: 'T', 4: 'M'}[kt]] = seq
        out[a.refs[0]] = keys
    return out, end


def sample(seq, t):
    if len(seq) == 1 or t <= seq[0][0]:
        return seq[0][1]
    for (t0, a), (t1, b) in zip(seq, seq[1:]):
        if t0 <= t <= t1:
            u = (t - t0) / (t1 - t0) if t1 > t0 else 0
            if len(a) == 4:                              # quaternion: nlerp along the short way
                if np.dot(a, b) < 0:
                    b = -b
                q = a * (1 - u) + b * u; return q / np.linalg.norm(q)
            return a * (1 - u) + b * u
    return seq[-1][1]


def pose(skel, anim, t):
    world = {}
    for bone, (parent, local) in skel.items():
        k = anim.get(bone)
        if k and 'R' in k:
            S = np.diag(list(sample(k['S'], t)) + [1]) if 'S' in k else np.eye(4)
            R = quat(*sample(k['R'], t))
            T = np.eye(4); T[3, :3] = sample(k['T'], t) if 'T' in k else local[3, :3]
            local = S @ R @ T
        elif k and 'M' in k:
            local = sample(k['M'], t).reshape(4, 4)
        world[bone] = local @ world[parent] if parent in world else local
    return world


def skin(mesh, world):
    V = np.column_stack([mesh['V'], np.ones(len(mesh['V']))])
    out = np.zeros((len(V), 3)); tot = np.zeros(len(V))
    for bone, idx, w, off in mesh['skin']:
        if bone not in world:
            continue
        M = off @ world[bone]
        out[idx] += (V[idx] @ M)[:, :3] * w[:, None]; tot[idx] += w
    un = tot < 1e-6
    out[un] = mesh['V'][un]
    return out


# ---------------------------------------------------------------- dressing and rendering
SHIRT = (74, 96, 128)          # a plain blue-grey T-shirt (the game's white T-shirt, tinted)
HAIR = (92, 64, 40)            # brown
TROUSER_MASKS = ['Belt', 'Crotch', 'LeftLeg', 'RightLeg']      # the jeans' masks 14, 15, 7, 9: body under them is not drawn


def rgba(path):
    return np.asarray(Image.open(path).convert('RGBA')).astype(np.float32)


def over(base, top):
    a = top[..., 3:4] / 255.0
    base[..., :3] = top[..., :3] * a + base[..., :3] * (1 - a)
    base[..., 3:4] = np.maximum(base[..., 3:4], top[..., 3:4])
    return base


def outfit(media):
    T = media / 'textures'
    body = rgba(T / 'Body' / 'MaleBody02a.png')
    shirt = rgba(T / 'Clothes' / 'Shirt_Tshirt_Textures' / 'TShirt_White.png'); shirt[..., :3] *= np.array(SHIRT) / 255.0
    over(body, shirt)
    over(body, rgba(T / 'Clothes' / 'Shoes_Socks_Textres' / 'Trainers_White.png'))
    for m in TROUSER_MASKS:
        mk = rgba(T / 'Body' / 'Masks' / f'{m}.png')
        body[mk[..., 3] > 0, 3] = 0
    hair = rgba(T / 'F_Hair_White.png'); hair[..., :3] *= np.array(HAIR) / 190.0; hair[..., 3] = 255
    jeans = rgba(T / 'Clothes' / 'Trousers_Mesh' / 'TrousersMesh_Denim.png')
    return body, jeans, np.clip(hair, 0, 255)


def load(media):
    S = media / 'models_X' / 'Skinned'
    body = parse_x(S / 'MaleBody.x')
    skel = skeleton(body)
    tex_body, tex_jeans, tex_hair = outfit(media)
    parts = [(meshes(body)[0], tex_body)]
    for f, tex in ((S / 'Clothes' / 'Bob_Trousers.x', tex_jeans), (S / 'Hair' / 'Bob_Hair_CrewCut.x', tex_hair)):
        try:
            parts += [(m, tex) for m in meshes(parse_x(f)) if m['skin']]
        except FileNotFoundError:
            pass
    walk_clip, walk_end = animation(parse_x(media / 'anims_X' / 'Bob' / 'Bob_Walk.x'))
    idle_clip, _ = animation(parse_x(media / 'anims_X' / 'Bob' / 'Bob_Idle.x'))
    return skel, parts, (walk_clip, walk_end), idle_clip


def render_pose(skel, parts, clip, t, d):
    """One sprite: the survivor posed at time t of a clip, facing tile direction d. Returns (image, anchor x, anchor y)."""
    world = pose(skel, clip, t)
    posed = [(skin(m, world), m, tex) for m, tex in parts]
    allv = np.vstack([p[0] for p in posed])
    k = HEIGHT / max(1e-6, allv[:, 1].max() - allv[:, 1].min())
    F = np.array(d, float); F /= np.linalg.norm(F); Lft = np.array([F[1], -F[0]])
    tris, pts = [], []
    for V, m, tex in posed:
        # model: +y up, the figure faces -z (turned 180 degrees about the vertical here, which keeps its handedness)
        xy = np.outer(-V[:, 2] * k, F) + np.outer(V[:, 0] * k, Lft)
        P3 = np.column_stack([xy, (V[:, 1] - allv[:, 1].min()) * k])
        sx = (P3[:, 0] - P3[:, 1]) * RV.PX; sy = (P3[:, 0] + P3[:, 1]) * RV.PY - P3[:, 2] * RV.PZ
        dep = .612 * P3[:, 0] + .612 * P3[:, 1] + .5 * P3[:, 2]
        S2 = np.column_stack([sx, sy]); pts.append(S2)
        uv = m['uv'].copy() if m['uv'] is not None else np.zeros((len(V), 2))
        uv[:, 1] = 1 - uv[:, 1]                       # DirectX texture coordinates run top-down; the rasteriser expects bottom-up
        for a in m['T']:
            e1, e2 = P3[a[1]] - P3[a[0]], P3[a[2]] - P3[a[0]]
            n = np.cross(e1, e2); ln = np.linalg.norm(n); n = n / ln if ln > 1e-12 else np.array([0, 0, 1.0])
            if np.dot(n, [.612, .612, .5]) < 0:
                n = -n
            tris.append((S2[a], dep[a], uv[a], tex, 'cut', n))
    allp = np.vstack(pts)
    mn, mx = np.floor(allp.min(0)) - 3, np.ceil(allp.max(0)) + 3
    img = RV.rasterise(tris, (int(mx[0] - mn[0]), int(mx[1] - mn[1])), -mn)
    return img, -mn[0], -mn[1]


# the game's own movement: B42 moves characters by each clip's root motion (the Translation_Data bone), played at the
# speed scales of the player's animation sets (media/AnimSets/player): the set's own scale times its uninjured blend's
GAITS = {'walk': ('Bob_Walk', 'movement/defaultWalk'), 'run': ('Bob_Run', 'run/defaultRun'), 'sprint': ('Bob_Sprint', 'sprint/defaultSprint')}


def speed_scale(media, anim_set, clip):
    t = (media / 'AnimSets' / 'player' / f'{anim_set}.xml').read_text(encoding='utf-8-sig')
    node = re.search(r'<m_SpeedScale>([\d.]+)</m_SpeedScale>', t)
    blend = next((float(m.group(2)) for m in re.finditer(r'<m_2DBlends[^>]*>\s*<m_AnimName>(.*?)</m_AnimName>.*?<m_SpeedScale>([\d.]+)</m_SpeedScale>', t, re.S)
                  if m.group(1) == clip), 1.0)
    return (float(node.group(1)) if node else 1.0) * blend


def gait(media, name, height_units):
    """(clip, clip length in ticks, speed in tiles a second, tiles covered per cycle) for walk, run or sprint."""
    clip, anim_set = GAITS[name]
    root = parse_x(media / 'anims_X' / 'Bob' / f'{clip}.x')
    anim, end = animation(root)
    tps = next((n.data[0] for n in walk(root) if n.kind == 'AnimTicksPerSecond'), 4800.0)
    T = anim.get('Translation_Data', {}).get('T', [])
    dist = float(np.hypot(*(T[-1][1] - T[0][1])[[0, 2]])) if len(T) > 1 else 0.0       # horizontal travel per cycle
    k = HEIGHT / height_units                                                            # model units -> tiles, as drawn
    scale = speed_scale(media, anim_set, clip)
    secs = end / tps / scale
    return anim, end, round(dist * k / secs, 3), round(dist * k, 3)


def render_all(media):
    """{sprite name: (image, ox, oy)} in the atlas's 128 x 256 frame convention, and the character's metadata:
    per gait its speed in tiles a second and the tiles one animation cycle covers (so the feet never slide)."""
    skel, parts, _, idle_clip = load(media)
    rest = np.vstack([skin(m, pose(skel, {}, 0)) for m, _ in parts])
    height_units = rest[:, 1].max() - rest[:, 1].min()
    imgs, meta = {}, {'dirs': list(DIRS), 'frames': WALK_FRAMES, 'walk': WALK_FRAMES, 'height': HEIGHT, 'gaits': {}}
    for g in GAITS:
        anim, end, speed, stride = gait(media, g, height_units)
        meta['gaits'][g] = {'speed': speed, 'stride': stride}
        for name, d in DIRS.items():
            for i in range(WALK_FRAMES):
                img, ax, ay = render_pose(skel, parts, anim, end * i / WALK_FRAMES, d)
                imgs[f'char_{g}_{name}_{i}'] = (img, int(round(64 - ax)), int(round(192 - ay)))
    for name, d in DIRS.items():
        img, ax, ay = render_pose(skel, parts, idle_clip, 0, d)
        imgs[f'char_idle_{name}'] = (img, int(round(64 - ax)), int(round(192 - ay)))
    return imgs, meta


if __name__ == '__main__':                      # contact sheet: python3 render_character.py [media] [out.png]
    import time
    media = Path(sys.argv[1]) if len(sys.argv) > 1 and sys.argv[1] else RV.Path.home() / 'Library/Application Support/Steam/steamapps/common/ProjectZomboid/Project Zomboid.app/Contents/Java/media'
    t = time.time(); imgs, meta = render_all(media); print(f'{len(imgs)} frames in {time.time() - t:.1f}s')
    print(meta['gaits'])
    rows = [[imgs[f'char_idle_{d}'][0]] + [imgs[f'char_{g}_{d}_{i}'][0] for i in range(WALK_FRAMES)] for g in GAITS for d in ('SE', 'SW')]
    W = max(sum(i.width for i in r) + 10 * len(r) for r in rows); H = sum(max(i.height for i in r) + 10 for r in rows)
    sheet = Image.new('RGBA', (W, H), (70, 70, 70, 255)); y = 0
    for r in rows:
        x = 0
        for im in r: sheet.alpha_composite(im, (x, y)); x += im.width + 10
        y += max(i.height for i in r) + 10
    sheet.save(sys.argv[2] if len(sys.argv) > 2 else 'character-check.png')