"""骨の動きを確かめる：決めた姿勢を付けて描く。
python3 nurse_pose.py <nurse_clay.blend> <出力フォルダ> [姿勢の名前...]
姿勢：rest give receive drawer point bow worried talk
（Unity では Humanoid の動き・腕の IK で動かす。ここの姿勢は、骨と重みで服や肌が破れないかの確認用）
"""
import math
import os
import sys

import bpy  # noqa: I001
from mathutils import Matrix, Vector

src, out = sys.argv[1], sys.argv[2]
want = sys.argv[3:]
os.makedirs(out, exist_ok=True)


def setup_scene():
    sc = bpy.context.scene
    sc.render.engine = "CYCLES"
    sc.cycles.device = "CPU"
    sc.cycles.samples = 20
    sc.cycles.use_denoising = True
    sc.view_settings.view_transform = "Standard"
    w = bpy.data.worlds.new("W")
    sc.world = w
    w.use_nodes = True
    bg = w.node_tree.nodes["Background"]
    bg.inputs["Color"].default_value = (0.80, 0.90, 0.89, 1)
    bg.inputs["Strength"].default_value = 0.65
    for nm, en, sz, loc in (("K", 380, 3, (-2.2, -2.8, 3.0)), ("F", 130, 4, (2.8, -2.0, 1.4))):
        L = bpy.data.lights.new(nm, "AREA")
        L.energy, L.size = en, sz
        o = bpy.data.objects.new(nm, L)
        o.location = loc
        o.rotation_euler = (Vector((0, 0, 1.0)) - Vector(loc)).to_track_quat("-Z", "Y").to_euler()
        sc.collection.objects.link(o)
    return sc


def rot_world(rig, bone, axis, deg, about=None):
    """骨を、ワールドの軸 axis のまわりに deg 度回す（回す中心は骨の根元）"""
    pb = rig.pose.bones[bone]
    h = Vector(about) if about is not None else pb.head.copy()
    R = Matrix.Rotation(math.radians(deg), 4, Vector(axis).normalized())
    pb.matrix = Matrix.Translation(h) @ R @ Matrix.Translation(-h) @ pb.matrix
    bpy.context.view_layer.update()


def twist(rig, bone, deg):
    """骨を自分の向きのまわりにねじる"""
    pb = rig.pose.bones[bone]
    rot_world(rig, bone, pb.tail - pb.head, deg)


def curl(rig, side, deg, fingers=("Index", "Middle", "Ring", "Little"), thumb=None):
    """指を手のひらのほうへ曲げる（骨の z 軸を手のひらに向けてあるので、x 軸まわりに回す）"""
    for f in fingers:
        for seg, k in (("Proximal", 1.0), ("Intermediate", 1.2)):
            pb = rig.pose.bones[side + f + seg]
            pb.rotation_mode = "XYZ"
            pb.rotation_euler.x = math.radians(deg * k)
    if thumb is not None:
        for seg in ("Proximal", "Intermediate"):
            pb = rig.pose.bones[side + "Thumb" + seg]
            pb.rotation_mode = "XYZ"
            pb.rotation_euler.x = math.radians(thumb)
    bpy.context.view_layer.update()


def mouth(kind):
    for ob in bpy.data.objects:
        if ob.name.startswith("Nurse_Mouth_"):
            on = ob.name == "Nurse_Mouth_" + kind
            ob.hide_render = not on
            ob.hide_viewport = not on


def brows(rig, deg, lift=0.0):
    """眉：deg＞0 で眉頭が上がる（困り顔）。lift は上へ（m）"""
    for tag, sgn in (("L", 1), ("R", -1)):
        pb = rig.pose.bones["Brow_" + tag]
        rot_world(rig, "Brow_" + tag, (0, 1, 0), deg * sgn)
        pb.matrix = Matrix.Translation((0, 0, lift)) @ pb.matrix
        bpy.context.view_layer.update()


def reset(rig):
    for pb in rig.pose.bones:
        pb.matrix_basis = Matrix.Identity(4)
    bpy.context.view_layer.update()
    mouth("smile")


X, Y, Z = (1, 0, 0), (0, 1, 0), (0, 0, 1)


def arm_forward(rig, S, up_deg, elbow_deg, palm_up=0.0):
    """腕を前へ上げる（up_deg）。ひじを曲げて前腕をさらに上げる（elbow_deg）。palm_up で手のひらを上へ"""
    rot_world(rig, S + "UpperArm", X, -up_deg)
    rot_world(rig, S + "LowerArm", X, -elbow_deg)
    if palm_up:
        sgn = 1 if S == "Right" else -1
        twist(rig, S + "LowerArm", sgn * palm_up * 0.5)
        twist(rig, S + "Hand", sgn * palm_up * 0.5)


POSES = {}


def pose(fn):
    POSES[fn.__name__] = fn
    return fn


@pose
def rest(rig):
    pass


@pose
def give(rig):      # 片手で物を差し出す
    arm_forward(rig, "Right", 55, 30, palm_up=80)
    curl(rig, "Right", 12, thumb=5)
    rot_world(rig, "Head", X, 8)
    rot_world(rig, "LeftUpperArm", Y, 4)


@pose
def receive(rig):   # 両手で受け取る
    for S in ("Left", "Right"):
        arm_forward(rig, S, 40, 35, palm_up=80)
        curl(rig, S, 15, thumb=5)
    rot_world(rig, "Chest", X, 6)
    rot_world(rig, "Head", X, 10)


@pose
def drawer(rig):    # 前かがみで引き出しの中の物をつかむ
    rot_world(rig, "Spine", X, 12)
    rot_world(rig, "Chest", X, 12)
    for S in ("Left", "Right"):
        arm_forward(rig, S, 35, 10)
        curl(rig, S, 55, thumb=30)
    rot_world(rig, "Head", X, 12)


@pose
def point(rig):     # 指で示す（話しながら）
    arm_forward(rig, "Right", 75, 5)
    rot_world(rig, "RightUpperArm", Z, -18)
    curl(rig, "Right", 85, fingers=("Middle", "Ring", "Little"), thumb=40)
    rot_world(rig, "Head", Z, -12)
    mouth("talk")


@pose
def bow(rig):       # おじぎ
    rot_world(rig, "Spine", X, 18)
    rot_world(rig, "Chest", X, 14)
    rot_world(rig, "Head", X, 6)
    for S, sgn in (("Left", -1), ("Right", 1)):
        rot_world(rig, S + "UpperArm", Y, sgn * -6)
        rot_world(rig, S + "UpperArm", X, -12)


@pose
def worried(rig):   # 困った：首をかしげ、手を口元へ
    rot_world(rig, "Head", Y, 10)
    rot_world(rig, "Head", X, 6)
    arm_forward(rig, "Right", 25, 115)
    rot_world(rig, "RightUpperArm", Z, -8)
    curl(rig, "Right", 30, thumb=10)
    brows(rig, 16, lift=0.002)
    mouth("worried")


@pose
def talk(rig):      # 話しかける：手のひらを見せる
    arm_forward(rig, "Left", 30, 55, palm_up=60)
    curl(rig, "Left", 10, thumb=0)
    rot_world(rig, "Head", Z, 8)
    brows(rig, 0, lift=0.003)
    mouth("talk")


bpy.ops.wm.open_mainfile(filepath=src)
rig = bpy.data.objects["Nurse_Oshiro"]
sc = setup_scene()
H = 1.65
views = {"full": ((1.6, -3.0, 1.15), (0, -0.1, 0.95), 50, (520, 820)),
         "close": ((0.75, -1.25, 1.42), (0, -0.15, 1.22), 50, (560, 560))}
for name in (want or list(POSES)):
    reset(rig)
    POSES[name](rig)
    for vname, (loc, tgt, lens, res) in views.items():
        for o in [o for o in bpy.data.objects if o.type == "CAMERA"]:
            bpy.data.objects.remove(o)
        cam = bpy.data.cameras.new("C")
        cam.lens = lens
        co = bpy.data.objects.new("C", cam)
        co.location = loc
        co.rotation_euler = (Vector(tgt) - Vector(loc)).to_track_quat("-Z", "Y").to_euler()
        sc.collection.objects.link(co)
        sc.camera = co
        sc.render.resolution_x, sc.render.resolution_y = res
        sc.render.filepath = os.path.join(out, f"pose_{name}_{vname}.png")
        bpy.ops.render.render(write_still=True)
    print("wrote", name)