#!/usr/bin/env python3
"""
typoglyph2.py — 把中文字換成「長得很像但不對」的字（多重特徵版）。

v1 只用單一字型的模糊圖片距離，缺點是「整體平均」會掩蓋局部的巨大差異，
偶爾會換出形狀差很多的字。v2 用六個特徵一起打分：

  1. glob   全域模糊像素距離        — 整體輪廓（主力）
  2. worst  3×3 分區距離的 L4 範數  — 懲罰「只有某一區差很多」的字
  3. coarse 8×8 粗略版面距離        — 大結構（左右分、上下分⋯⋯）
  4. ink    墨水量差異              — 筆畫繁簡程度的替代指標
  5. ids    IDS 部件集合的 Jaccard  — 有沒有共用部件（馬字旁、辶⋯⋯）
  6. struct IDS 頂層結構運算子      — ⿰左右／⿱上下／⿺半包圍 是否相同

前四項用兩套字型（黑體＋明體）各算一次再平均，避免被單一字型的
造型癖好帶偏。所有特徵在候選池內正規化到 0~1，再加權相加。

用法：
    python3 typoglyph2.py -t "今天天氣真好" -r 0.4 -k 3
    python3 typoglyph2.py -c 實驗盤          # 查相似字排行
    python3 typoglyph2.py -c 驗 --explain    # 看每個特徵給了幾分
    cat 文章.md | python3 typoglyph2.py -r 0.5

    -w glob=1,worst=1,coarse=.5,ink=.3,ids=.6,struct=.3    自己調權重

需要：pillow, numpy, scipy。首次執行會抓 IDS 資料並建索引（約 30 秒）。
IDS 資料來自 CHISE / cjkvi-ids（MIT-like 授權）。
"""

import argparse
import hashlib
import os
import random
import re
import sys
import urllib.request

import numpy as np
from PIL import Image, ImageDraw, ImageFont
from scipy.fft import dctn
from scipy.ndimage import gaussian_filter

# ── 參數 ────────────────────────────────────────────────────────────
RENDER = 48
CELL = 24          # 比對解析度（模糊後降採樣）
SIGMA = 1.5        # 模糊強度：調大 → 更看重整體輪廓、更不在意筆畫細節
HASH_N = 16
POOL = 800         # phash 粗篩保留的候選數
GRID = 3           # 分區數（3×3）
PNORM = 4          # 分區聚合用的範數。→∞ 等於取最差區，越小越寬容

IDS_URL = "https://raw.githubusercontent.com/cjkvi/cjkvi-ids/master/ids.txt"
CACHE = os.path.expanduser("~/.cache/typoglyph")
OPS = "⿰⿱⿲⿳⿴⿵⿶⿷⿸⿹⿺⿻⿼⿽⿾⿿"

DEFAULT_FONTS = [
    ("/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", 3),
    ("/usr/share/fonts/opentype/noto/NotoSerifCJK-Regular.ttc", 3),
]
DEFAULT_W = dict(glob=1, worst=0.2, coarse=1, ink=3, ids=15, struct=1)

POPCOUNT = np.unpackbits(np.arange(256, dtype=np.uint8)[:, None], axis=1).sum(1).astype(np.uint8)


# ── IDS 部件資料 ────────────────────────────────────────────────────
def load_ids(verbose=True):
    """回傳 {字: IDS 字串}。優先取台灣（T）的寫法。"""
    path = os.path.join(CACHE, "ids.txt")
    if not os.path.exists(path):
        if verbose:
            print("下載 IDS 部件資料⋯", file=sys.stderr)
        os.makedirs(CACHE, exist_ok=True)
        try:
            urllib.request.urlretrieve(IDS_URL, path)
        except Exception as e:
            print(f"IDS 下載失敗（{e}），改用純圖像模式", file=sys.stderr)
            return {}

    out = {}
    for line in open(path, encoding="utf-8"):
        if line.startswith("#"):
            continue
        p = line.rstrip("\n").split("\t")
        if len(p) < 3:
            continue
        # 多個來源時挑台灣的：欄位長得像 ⿰馬僉[TJKV]
        pick = p[2]
        for x in p[2:]:
            tags = re.findall(r"\[(\w+)\]", x)
            if tags and "T" in tags[0]:
                pick = x
                break
        out[p[1]] = re.sub(r"\[.*?\]", "", pick)
    return out


def components(ch, ids, depth=0, seen=None):
    """遞迴收集所有層級的部件，**含中間節點**。

    只收葉子是不行的：驗 會被拆成 一人口灬，連「馬」都不見，
    跟「駿」的共用部首訊號就整個消失了。
    """
    seen = seen or set()
    if depth > 3 or ch in seen:
        return set()
    d = ids.get(ch)
    out = set()
    if d and d != ch:
        for c in d:
            if c in OPS:
                continue
            out.add(c)
            out |= components(c, ids, depth + 1, seen | {ch})
    return out or {ch}


# ── 索引 ────────────────────────────────────────────────────────────
def build(fonts, pool, verbose=True):
    cands = []
    for cp in range(0x4E00, 0xA000):
        ch = chr(cp)
        if pool == "big5":
            try:
                ch.encode("big5")       # Big5 編得出來 → 順便排除簡體字
            except UnicodeEncodeError:
                continue
        cands.append(ch)

    if verbose:
        print(f"算繪 {len(cands)} 字 × {len(fonts)} 套字型⋯", file=sys.stderr)
    stacks = []
    for path, face in fonts:
        f = ImageFont.truetype(path, int(RENDER * 0.84), index=face)
        arr = np.zeros((len(cands), RENDER, RENDER), np.float32)
        for i, ch in enumerate(cands):
            im = Image.new("L", (RENDER, RENDER), 0)
            ImageDraw.Draw(im).text((RENDER // 2, RENDER // 2), ch,
                                    font=f, fill=255, anchor="mm")
            arr[i] = np.asarray(im, np.float32) / 255
        stacks.append(arr)

    ok = np.all([s.sum((1, 2)) > 5 for s in stacks], axis=0)   # 剔除缺字形的
    chars = "".join(c for c, k in zip(cands, ok) if k)
    stacks = [s[ok] for s in stacks]
    n = len(chars)

    # phash 用第一套字型就好（只是粗篩）
    D = dctn(stacks[0], axes=(1, 2), norm="ortho")[:, :HASH_N, :HASH_N].reshape(n, -1)[:, 1:]
    P = np.packbits(D > np.median(D, axis=1, keepdims=True), axis=1)

    step = RENDER // CELL
    B = np.stack([(gaussian_filter(s, sigma=(0, SIGMA, SIGMA))[:, ::step, ::step] * 255)
                  .astype(np.uint8) for s in stacks])          # (字型, 字, 24, 24)
    ink = stacks[0].sum((1, 2)).astype(np.float32)
    return chars, B, P, ink


def load(fonts, pool, verbose=True):
    key = hashlib.md5(
        f"{fonts}|{pool}|{RENDER}|{CELL}|{SIGMA}|{HASH_N}".encode()).hexdigest()[:12]
    path = os.path.join(CACHE, f"idx{key}.npz")
    if os.path.exists(path):
        z = np.load(path)
        return str(z["chars"]), z["B"], z["P"], z["ink"]
    chars, B, P, ink = build(fonts, pool, verbose)
    os.makedirs(CACHE, exist_ok=True)
    np.savez_compressed(path, chars=np.array(chars), B=B, P=P, ink=ink)
    if verbose:
        print(f"索引完成：{len(chars)} 字 → {path}", file=sys.stderr)
    return chars, B, P, ink


class Glypher:
    def __init__(self, fonts=None, pool="big5", weights=None, verbose=True):
        fonts = fonts or DEFAULT_FONTS
        self.w = dict(DEFAULT_W, **(weights or {}))
        self.chars, B, self.P, self.ink = load(fonts, pool, verbose)
        self.idx = {c: i for i, c in enumerate(self.chars)}
        n = len(self.chars)
        B = B.astype(np.int32)
        g = CELL // GRID
        self.flat = [b.reshape(n, -1) for b in B]
        self.reg = [b.reshape(n, GRID, g, GRID, g)
                     .transpose(0, 1, 3, 2, 4).reshape(n, GRID * GRID, g * g) for b in B]
        self.coarse = [b[:, ::3, ::3].reshape(n, -1) for b in B]

        ids = load_ids(verbose)
        if ids:
            self.comp = [components(c, ids) for c in self.chars]
            self.struct = [next((x for x in ids.get(c, "") if x in OPS), "") for c in self.chars]
        else:
            self.comp = self.struct = None

    # ── 評分 ──
    def _score(self, i, cand):
        def nz(x):
            x = np.asarray(x, np.float64)
            r = x.max() - x.min()
            return (x - x.min()) / r if r > 0 else np.zeros_like(x)

        f = {}
        f["glob"] = nz(np.mean([np.abs(F[cand] - F[i]).sum(1) for F in self.flat], 0))
        f["worst"] = nz(np.mean([(np.abs(R[cand] - R[i]).sum(2) ** PNORM).mean(1) ** (1 / PNORM)
                                 for R in self.reg], 0))
        f["coarse"] = nz(np.mean([np.abs(C[cand] - C[i]).sum(1) for C in self.coarse], 0))
        f["ink"] = nz(np.abs(self.ink[cand] - self.ink[i]) / max(self.ink[i], 1))
        if self.comp:
            ci = self.comp[i]
            f["ids"] = np.array([1 - len(ci & self.comp[j]) / max(len(ci | self.comp[j]), 1)
                                 for j in cand])
            f["struct"] = np.array([0.0 if self.struct[j] == self.struct[i] else 1.0
                                    for j in cand])
        else:
            f["ids"] = f["struct"] = np.zeros(len(cand))
        total = sum(self.w[k] * f[k] for k in f)
        return total, f

    def similar(self, ch, n=8, explain=False):
        i = self.idx.get(ch)
        if i is None:
            return [] if not explain else ([], None)
        ham = POPCOUNT[np.bitwise_xor(self.P, self.P[i])].sum(1)
        cand = np.argpartition(ham, POOL)[:POOL]
        cand = cand[cand != i]
        total, f = self._score(i, cand)
        order = np.argsort(total)[:n]
        picks = [self.chars[cand[o]] for o in order]
        if explain:
            rows = [(self.chars[cand[o]], total[o], {k: f[k][o] for k in f}) for o in order]
            return picks, rows
        return picks

    def swap(self, ch, k=4, rng=random):
        c = self.similar(ch, k)
        if not c:
            return ch
        return rng.choices(c, weights=[1 / (r + 1.5) for r in range(len(c))])[0]

    def transform(self, text, rate=1.0, k=4, seed=None):
        rng = random.Random(seed)
        return "".join(self.swap(c, k, rng) if (c in self.idx and rng.random() < rate) else c
                       for c in text)


def parse_w(s):
    if not s:
        return {}
    return {k: float(v) for k, v in (kv.split("=") for kv in s.split(","))}


def main():
    p = argparse.ArgumentParser(description="把中文字換成長得很像但不對的字（多重特徵版）")
    p.add_argument("-t", "--text")
    p.add_argument("-c", "--char", help="查相似字排行")
    p.add_argument("-n", type=int, default=12)
    p.add_argument("-r", "--rate", type=float, default=1.0, help="替換比例 0~1")
    p.add_argument("-k", type=int, default=4, help="從前 k 個裡挑，越小越好讀")
    p.add_argument("-s", "--seed", type=int)
    p.add_argument("-w", "--weights", help="glob=1,worst=1,ids=.6,⋯")
    p.add_argument("--explain", action="store_true", help="顯示各特徵分數")
    p.add_argument("--pool", choices=["big5", "all"], default="big5")
    p.add_argument("--font", action="append", metavar="路徑[:字面編號]",
                   help="可指定多次；不給則用 Noto 黑體＋明體")
    a = p.parse_args()

    fonts = None
    if a.font:
        fonts = []
        for f in a.font:
            path, _, face = f.rpartition(":")
            fonts.append((path, int(face)) if path and face.isdigit() else (f, 0))

    g = Glypher(fonts, a.pool, parse_w(a.weights))

    if a.char:
        for ch in a.char:
            if a.explain:
                picks, rows = g.similar(ch, a.n, explain=True)
                print(f"\n{ch}  →  {' '.join(picks)}")
                print(f"  {'字':<3} {'總分':>6} " +
                      " ".join(f"{k:>6}" for k in DEFAULT_W))
                for c, tot, f in rows:
                    print(f"  {c:<3} {tot:6.3f} " +
                          " ".join(f"{f[k]:6.3f}" for k in DEFAULT_W))
            else:
                print(f"{ch} → {' '.join(g.similar(ch, a.n)) or '（索引裡沒有）'}")
        return

    text = a.text or sys.stdin.read().rstrip("\n")
    if not text:
        p.error("請用 -t 給文字，或從 stdin 餵進來")
    print(g.transform(text, a.rate, a.k, a.seed))


if __name__ == "__main__":
    main()