import numpy as np
from skimage.filters import threshold_otsu
from skimage.morphology import remove_small_objects
from skimage.feature import hog
from skimage.measure import moments_hu, label
from skimage.morphology import skeletonize, binary_dilation, disk
from scipy.ndimage import convolve
from skimage.morphology import binary_opening
def _prune_spurs(skel: np.ndarray, iters: int = 2) -> np.ndarray:
    """
    迭代剪掉骨架上长度很短的端点分支（spur）。
    iters 表示最多向内剪掉几步（像素）。推荐 1~3。
    """
    s = skel.copy().astype(bool)
    # 用 3x3 邻域统计端点：中心权重10，其它1；“10+1=11”即1个邻居的端点
    K = np.array([[1,1,1],
                  [1,10,1],
                  [1,1,1]], dtype=np.uint8)
    for _ in range(iters):
        nb = convolve(s.astype(np.uint8), K, mode="constant", cval=0)
        endpoints = (nb == 11)          # 只有 1 个邻居
        # 只剪 endpoints，不动分叉/主干
        s = s & ~endpoints
    return s
def _ensure_ink_true(bw_bool: np.ndarray) -> np.ndarray:
    bw = bw_bool.astype(bool)
    if bw.mean() > 0.5:
        bw = ~bw
    return bw
def stroke_normalize(bw: np.ndarray, target_px: int = 2) -> np.ndarray:
    if bw.dtype != bool:
        bw = (bw > 0)
    if bw.mean() > 0.5:
        bw = ~bw

    skel = skeletonize(bw)
    skel = _prune_spurs(skel, iters=2)      # ← 新增：剪短刺，去笔锋小尖

    if target_px <= 1:
        return skel.astype(np.float32)
    rad   = max(1, int(round(target_px/2)))
    thick = binary_dilation(skel, disk(rad))
    #thick = binary_opening(thick, disk(1))#optional

    return (thick & bw).astype(np.float32)
def to_64_gray(imgPIL):
    return np.array(imgPIL, dtype=np.uint8)

def binarize(gray64, min_size: int = 4, keep_largest: bool = False):
    if gray64.ndim == 3:
        gray64 = gray64[..., 0]
    g = gray64.astype(float)
    if g.max() > 1:
        g /= 255.0

    t  = threshold_otsu(g)
    bw = (g <= t)
    bw = remove_small_objects(bw.astype(bool), min_size=min_size).astype(bool)

    if keep_largest:
        lab = label(bw)
        if lab.max() > 0:
            areas = np.bincount(lab.ravel()); areas[0] = 0
            bw = (lab == areas.argmax())

    return bw

def proj_features(bw):
    # normalization
    hp = bw.sum(axis=0); vp = bw.sum(axis=1)
    if hp.max()>0: hp = hp/hp.max()
    if vp.max()>0: vp = vp/vp.max()
    # fix dimension
    def pool(v, m=32):
        idx = np.linspace(0, len(v), m+1, endpoint=True).astype(int)
        return np.array([v[idx[i]:idx[i+1]].mean() for i in range(m)], dtype=np.float32)
    return np.concatenate([pool(hp), pool(vp)])

def feat_vec(bw):

    hu = moments_hu(bw).astype(np.float32)
    hu = np.sign(hu)*np.log1p(np.abs(hu))

    hogv = hog(bw, orientations=9, pixels_per_cell=(8,8), cells_per_block=(2,2),
           block_norm='L2-Hys', feature_vector=True).astype(np.float32)
    proj = proj_features(bw).astype(np.float32)
    v = np.concatenate([hu, hogv, proj]).astype(np.float32)

    n = np.linalg.norm(v) + 1e-8
    return v / n

def cosine_sim(a, B):
    return (B @ a) / (np.linalg.norm(a)+1e-8) / (np.linalg.norm(B,axis=1)+1e-8)