# -*- coding: utf-8 -*-
"""Take 2. The previous mask was swamped by the global darkening pass.
Fix: histogram-match live luminance to vanilla per tab FIRST, which cancels the
tone shift, then difference, blur to kill speckle, and cluster what survives.
A removed arrow is a large coherent blob; recolour residue is not.
"""
import os
import struct
import numpy as np
from PIL import Image, ImageFilter
from scipy import ndimage

LIVE_DIR = r"C:\Diablo II\data\global\ui\SPELLS"
VAN_DIR = os.path.join(LIVE_DIR, "originals")
# Output goes beside this script, NOT into a session scratchpad. It used to point at
# %TEMP%\claude\...\arrows2, which Windows can clear at any moment and which is
# meaningless on any machine but the one that wrote it -- so the script silently
# produced nothing findable. D:\_SERVER\CLAUDE.md: scratch is for throwaway only.
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out", "arrows2")
os.makedirs(OUT, exist_ok=True)
pb = open(r"C:\Diablo II\Modding Tools\dc6color\act1.dat", "rb").read()
# D2 palettes are BGR, not RGB. Reading them as RGB reverses every colour.
pal = [(pb[i * 3 + 2], pb[i * 3 + 1], pb[i * 3]) for i in range(256)]

CLASSES = [("skltree_a_back.DC6", "Amazon"), ("skltree_b_back.DC6", "Barbarian"),
           ("skltree_d_back.dc6", "Druid"), ("skltree_i_back.dc6", "Assassin"),
           ("skltree_n_back.DC6", "Necromancer"), ("skltree_p_back.DC6", "Paladin"),
           ("skltree_s_back.DC6", "Sorceress")]
MIN_BLOB = 120     # an arrow shaft+head is well over this


def frames(path):
    d = open(path, "rb").read()
    dirs, fpd = struct.unpack_from("<ii", d, 16)
    ptrs = struct.unpack_from("<%di" % (dirs * fpd), d, 24)
    out = []
    for fp in ptrs:
        flip, w, h, ox, oy, unk, nb, ln = struct.unpack_from("<8i", d, fp)
        out.append((w, h, d[fp + 32:fp + 32 + ln]))
    return out


def decode(w, h, data):
    img = Image.new("RGBA", (max(w, 1), max(h, 1)), (0, 0, 0, 0))
    if w <= 0 or h <= 0:
        return img
    px = img.load()
    x, y, pos, L = 0, h - 1, 0, len(data)
    while pos < L and y >= 0:
        b = data[pos]; pos += 1
        if b == 0x80:
            x = 0; y -= 1
        elif b & 0x80:
            x += b & 0x7F
        else:
            for _ in range(b):
                if pos < L and 0 <= x < w and 0 <= y < h:
                    px[x, y] = pal[data[pos]] + (255,)
                pos += 1; x += 1
    return img


def group(fr, g):
    im = Image.new("RGB", (320, 432), (0, 0, 0))
    for k, xy in enumerate([(0, 0), (256, 0), (0, 256), (256, 256)]):
        t = decode(*fr[4 * g + k])
        im.paste(t, xy, t)
    return im


def hist_match(src, ref):
    """Match src's grey histogram to ref's. Cancels a global tone/contrast change."""
    s = src.ravel()
    sv, sidx, scnt = np.unique(s, return_inverse=True, return_counts=True)
    rv, rcnt = np.unique(ref.ravel(), return_counts=True)
    sq = np.cumsum(scnt).astype(np.float64) / s.size
    rq = np.cumsum(rcnt).astype(np.float64) / ref.size
    interp = np.interp(sq, rq, rv)
    return interp[sidx].reshape(src.shape)


results = []
for fn, cls in CLASSES:
    lp, vp = os.path.join(LIVE_DIR, fn), os.path.join(VAN_DIR, fn)
    lf, vf = frames(lp), frames(vp)
    for g in (1, 2, 3):
        v = group(vf, g); m = group(lf, g)
        va = np.asarray(v.convert("L")).astype(np.float64)
        ma = np.asarray(m.convert("L")).astype(np.float64)
        mm = hist_match(ma, va)
        d = np.abs(va - mm)
        d = np.asarray(Image.fromarray(d.astype(np.uint8)).filter(
            ImageFilter.GaussianBlur(1.4))).astype(np.float64)
        thr = d.mean() + 4.0 * d.std()
        mask = d > max(thr, 28)
        lab, n = ndimage.label(mask)
        if n:
            sizes = ndimage.sum(mask, lab, range(1, n + 1))
            objs = ndimage.find_objects(lab)
            cand = []
            for i, s in enumerate(sizes):
                if s < MIN_BLOB:
                    continue
                sl = objs[i]
                y0, x0 = sl[0].start, sl[1].start
                h, w = sl[0].stop - y0, sl[1].stop - x0
                cand.append((int(s), x0, y0, w, h))
            cand.sort(reverse=True)
        else:
            cand = []
        results.append((cls, g, cand, d.mean(), thr))
        tag = "" if not cand else "   <== %d blob(s)" % len(cand)
        print("%-12s tab%d  diffmean %5.2f thr %5.2f  blobs>=%d: %-3d%s"
              % (cls, g, d.mean(), thr, MIN_BLOB, len(cand), tag))
        for s, x, y, w, h in cand[:6]:
            print("        %6d px at (%3d,%3d) %3dx%-3d" % (s, x, y, w, h))
        if cand:
            sheet = Image.new("RGB", (320 * 3 + 24, 432), (18, 18, 24))
            sheet.paste(v, (0, 0)); sheet.paste(m, (332, 0))
            over = np.array(m).copy()
            over[mask] = (255, 40, 40)
            sheet.paste(Image.fromarray(over), (664, 0))
            sheet.save(os.path.join(OUT, "%s_tab%d.png" % (cls, g)))

tot = sum(len(c) for _, _, c, _, _ in results)
print("\n" + "=" * 72)
print("TOTAL blobs >= %dpx after tone-matching: %d" % (MIN_BLOB, tot))
print("tabs with any: %s" % [(c, g, len(b)) for c, g, b, _, _ in results if b])
print("sheets -> %s" % OUT)