# WL-814x2 팔레트 「최종 화면 색」 역보정 (닫힌 루프)
#   측정 지점을 알베도가 아니라 **Play 화면 픽셀**로 옮긴다.
#   칸마다  원본색 O(부위 보정 역산) · 화면색 S(측정)  를 비교해
#   목표 T = O × (밝기 S / 밝기 O)  가 되도록 팔레트 칸 색에 채널별 이득을 곱한다.
#   python AgentScripts/WL814x2_calib.py <측정에쓸 렌더이름> [apply]
import sys, os, json
import numpy as np
from PIL import Image

RAW = 'Screenshots_WL/WL814x/2nd/raw'
PAL = 'Assets/WL/Look/Character/Textures/WLPalette_FINAL.png'
TGT = 'AgentScripts/WL814x2_target.json'
GAIN = {'M05': (0.80, 1.00), 'Skin': (0.86, 0.93), 'Hair': (0.85, 0.93)}
PARTMAP = {'Body_m05': 'M05', 'Arm_m05': 'M05', 'Leg_M05': 'M05', 'Head': 'Skin', 'Hair05': 'Hair'}
REN = sys.argv[1] if len(sys.argv) > 1 else 'a_now'
APPLY = len(sys.argv) > 2 and sys.argv[2] == 'apply'


def L(n):
    return np.asarray(Image.open(RAW + '/' + n + '.png').convert('RGB'), dtype=np.uint8)


def lum(c):
    return 0.299 * c[0] + 0.587 * c[1] + 0.114 * c[2]


def hue(c):
    r, g, b = [float(x) for x in c]
    mx, mn = max(r, g, b), min(r, g, b); d = mx - mn
    if d < 1e-6: return 0.0
    if mx == r: h = ((g - b) / d) % 6
    elif mx == g: h = (b - r) / d + 2
    else: h = (r - g) / d + 4
    return h * 60.0


def dh(a, b):
    return (b - a + 540) % 360 - 180


idm = L('idmap'); ma = L('mask_all')[:, :, 0] > 127
palimg = np.array(Image.open(PAL).convert('RGB'), dtype=np.uint8)
S = palimg.shape[0]; G = S // 8
idref = np.array([[9 + j * 15, 240 - j * 11, 40 + j * 7] for j in range(G * G)], np.float32)
flat = idm.reshape(-1, 3).astype(np.float32)
d2 = ((flat[:, None, :] - idref[None, :, :]) ** 2).sum(2)
near = d2.argmin(1).reshape(idm.shape[:2])
kid = np.where((d2.min(1).reshape(idm.shape[:2]) < 400) & (idm.sum(2) > 12), near, -1)

masks = {}
for p in PARTMAP:
    f = RAW + '/mask_' + p + '.png'
    if os.path.exists(f):
        masks[p] = L('mask_' + p)[:, :, 0] > 127

# 목표색(원본) — 최초 1회만 저장하고 이후 회차에서는 고정해 쓴다
if os.path.exists(TGT):
    target = {int(k): v for k, v in json.load(open(TGT)).items()}
else:
    target = {}
    for j in range(G * G):
        gx, gy = j % G, j // G
        p = palimg[gy * 8 + 4, gx * 8 + 4].astype(np.float32)
        sel = ma & (kid == j)
        if sel.sum() < 200:
            target[j] = [float(x) for x in p]; continue
        part = max(((k, int((sel & v).sum())) for k, v in masks.items()), key=lambda t: t[1])[0]
        gs, gv = GAIN[PARTMAP[part]]
        mx = p.max() / 255.0; mn = p.min() / 255.0
        s = (mx - mn) / mx if mx > 1e-6 else 0
        s2 = min(s / gs, 1.0); v2 = min(mx / gv, 1.0)
        # H 유지 · S,V 만 역보정
        import colorsys
        h = hue(p) / 360.0
        r, g, b = colorsys.hsv_to_rgb(h, s2, v2)
        target[j] = [float(r * 255), float(g * 255), float(b * 255)]
    json.dump({str(k): v for k, v in target.items()}, open(TGT, 'w'))

img = L(REN)
newpal = palimg.copy()
print('%3s %-10s %7s | %-16s %-16s %-16s | %-16s %6s -> %6s' %
      ('칸', '부위', 'px', '원본 O', '팔레트 P', '화면 S', '새 팔레트 P\'', 'dHue', 'dHue예상'))
rows = []
for j in range(G * G):
    gx, gy = j % G, j // G
    p = palimg[gy * 8 + 4, gx * 8 + 4].astype(np.float32)
    sel = ma & (kid == j)
    n = int(sel.sum())
    if n < 200:
        continue
    part = max(((k, int((sel & v).sum())) for k, v in masks.items()), key=lambda t: t[1])[0]
    o = np.array(target[j], np.float32)
    s = img[sel].mean(0).astype(np.float32)
    t = o * (max(lum(s), 1.0) / max(lum(o), 1.0))          # 같은 밝기에서의 「원본 색조」
    gain = np.clip(t / np.maximum(s, 4.0), 0.55, 1.8)
    np_ = np.clip(p * gain, 0, 255)
    rows.append((j, part, n, o, p, s, t, np_))
    print('%3d %-10s %7d | %-16s %-16s %-16s | %-16s %+6.1f -> %+6.1f' % (
        j, part, n, tuple(int(x) for x in o), tuple(int(x) for x in p), tuple(int(x) for x in s),
        tuple(int(x) for x in np_), dh(hue(o), hue(s)), dh(hue(o), hue(t))))
    if APPLY:
        newpal[gy * 8:gy * 8 + 8, gx * 8:gx * 8 + 8] = np.round(np_).astype(np.uint8)

err = [abs(dh(hue(r[3]), hue(r[5]))) for r in rows]
w = np.array([r[2] for r in rows], np.float64)
print('화면 Hue 오차: 면적가중 평균 %.1f° · 단순평균 %.1f° · 최대 %.1f°' %
      ((np.array(err) * w).sum() / w.sum(), np.mean(err), np.max(err)))
if APPLY:
    Image.fromarray(newpal).save(PAL)
    print('팔레트 갱신 →', PAL)