"""1人ずつ描いた見本を、顔の大きさをそろえて横一列に並べる台本（ブログ「複数人をセミリアルで描く」のサンプル）。

使い方：
  python lineup.py compose --size 1536x864 --out lineup.png  a.png a_mask.png  b.png b_mask.png  ...
      見本の絵と、その切り抜きの面（白＝人物）を、左から順に組で渡す。並べた絵と、顔の位置の faces.json を書く
      --background bg.png を付けると、無地の灰色の代わりに背景の上に並べる（人物の色味を背景へ --tone だけ寄せる）
      --tone 0.5（0＝寄せない・1＝背景と同じ色味）・--face 0.22（顔の高さの上限・絵の高さに対する割合）・--top 0.12（顔の上端の位置）
      例：お花見の絵は --face 0.11 --top 0.36、色味は朝 0.5・昼 0.4・夜 0.75
  python lineup.py lock --lineup lineup.png --blended blended.png --faces faces.json --out final.png
      なじませた絵に、並べた絵の顔の周りを貼り戻す

要るもの：Python 3・Pillow・numpy・opencv-python（4.5.4 以上）。
顔の検出は OpenCV の YuNet（MIT ライセンス）。モデルは opencv_zoo の
models/face_detection_yunet/face_detection_yunet_2023mar.onnx を、この台本と同じ場所に置く。
"""
import argparse
import json
from pathlib import Path

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFilter

YUNET = Path(__file__).with_name("face_detection_yunet_2023mar.onnx")


def detect_face(path):
    """一番大きい顔の枠（x0, y0, x1, y1）。"""
    img = cv2.imread(str(path))
    h, w = img.shape[:2]
    _, found = cv2.FaceDetectorYN.create(str(YUNET), "", (w, h), 0.6, 0.3, 5000).detect(img)
    if found is None:
        raise SystemExit(f"顔が見つからない：{path}")
    f = max(found, key=lambda r: r[2] * r[3])
    return [int(f[0]), int(f[1]), int(f[0] + f[2]), int(f[1] + f[3])]


def layout(people, W, H, top=0.12, face=0.22):
    """顔の高さをそろえ、全員の体が下端まで届く大きさを選ぶ。返り値は人ごとの (x, y, 倍率)。"""
    n = len(people)
    slot = W / n
    base = min(H * face, slot * 0.55 * 1.25)
    need = max((p["face"][3] - p["face"][1]) * (H * (1 - top)) / max(1, p["bottom"] - p["face"][1]) for p in people)
    target = min(max(base, need), slot * 0.8 * 1.25)
    place = []
    for i, p in enumerate(people):
        x0, y0, x1, y1 = p["face"]
        s = target / (y1 - y0)
        place.append((int(slot * (i + 0.5) - (x0 + x1) / 2 * s), int(H * (top + (0.02 if i % 2 else 0)) - y0 * s), s))
    return place


def tone_match(px, wt, mean, std, k=0.5):
    """人物の色の平均と散らばりを、背景の値へ k だけ寄せる。"""
    w = wt[:, None] / max(1e-6, float(wt.sum()))
    mu = (px * w).sum(0)
    sd = np.sqrt(((px - mu) ** 2 * w).sum(0)) + 1e-6
    return np.clip((px - mu) * (1 + k * (std / sd - 1)) + mu + k * (mean - mu), 0, 255)


def compose(args):
    W, H = (int(v) for v in args.size.split("x"))
    pairs = list(zip(args.files[0::2], args.files[1::2]))
    people = []
    for img, mask in pairs:
        bbox = Image.open(mask).convert("L").point(lambda v: 255 if v > 128 else 0).getbbox()
        people.append({"img": img, "mask": mask, "face": detect_face(img), "bottom": bbox[3]})
    place = layout(people, W, H, args.top, args.face)
    canvas = Image.open(args.background).convert("RGB").resize((W, H)) if args.background else Image.new("RGB", (W, H), (205, 205, 207))
    faces = []
    n = len(people)
    for i in sorted(range(n), key=lambda i: -abs(i - (n - 1) / 2)):     # 外側から貼る＝真ん中が手前
        p, (x, y, s) = people[i], place[i]
        im = Image.open(p["img"]).convert("RGB")
        im = im.resize((int(im.width * s), int(im.height * s)), Image.LANCZOS)
        mk = Image.open(p["mask"]).convert("L").resize(im.size, Image.LANCZOS)
        if args.background:
            behind = np.asarray(canvas.crop((max(0, x), max(0, y), min(W, x + im.width), min(H, y + im.height))), np.float32).reshape(-1, 3)
            px = tone_match(np.asarray(im, np.float32).reshape(-1, 3), np.asarray(mk, np.float32).reshape(-1) / 255, behind.mean(0), behind.std(0), args.tone)
            im = Image.fromarray(px.reshape(im.height, im.width, 3).astype("uint8"))
        canvas.paste(im, (x, y), mk)
    for p, (x, y, s) in zip(people, place):
        x0, y0, x1, y1 = p["face"]
        faces.append([int(x + x0 * s), int(y + y0 * s), int(x + x1 * s), int(y + y1 * s)])
    canvas.save(args.out)
    Path(args.out).with_name("faces.json").write_text(json.dumps(faces))
    print(f"{args.out} と faces.json を書いた（{n}人）")


def lock(args):
    fin = Image.open(args.blended).convert("RGB")
    lineup_img = Image.open(args.lineup).convert("RGB").resize(fin.size)
    boxes = json.loads(Path(args.faces).read_text())
    mask = Image.new("L", fin.size, 0)
    d = ImageDraw.Draw(mask)
    for x0, y0, x1, y1 in boxes:
        cx, cy, w, h = (x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0) * 1.45, (y1 - y0) * 1.45
        d.ellipse([cx - w / 2, cy - h / 2 - (y1 - y0) * 0.12, cx + w / 2, cy + h / 2], fill=255)
    mask = mask.filter(ImageFilter.GaussianBlur(max(4, int(min(b[2] - b[0] for b in boxes) * 0.12))))
    fin.paste(lineup_img, (0, 0), mask)
    fin.save(args.out)
    print(f"{args.out} を書いた")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    c = sub.add_parser("compose")
    c.add_argument("--size", default="1536x864")
    c.add_argument("--out", default="lineup.png")
    c.add_argument("--background")
    c.add_argument("--tone", type=float, default=0.5)
    c.add_argument("--face", type=float, default=0.22)
    c.add_argument("--top", type=float, default=0.12)
    c.add_argument("files", nargs="+")
    k = sub.add_parser("lock")
    k.add_argument("--lineup", required=True)
    k.add_argument("--blended", required=True)
    k.add_argument("--faces", required=True)
    k.add_argument("--out", default="final.png")
    a = ap.parse_args()
    compose(a) if a.cmd == "compose" else lock(a)