#!/usr/bin/env python3
"""가지치기 후보 3종을 만든다 — 공격성이 다른 마스크의 트레이드오프를 재기 위해.

## 왜 3종인가
전면 마스킹(M3)은 오염을 2.70% -> 0.06% 로 없애지만 한국어 한자 병기까지 죽인다
(개항(開港) -> 개항(개항), 그리고 더 나쁘게는 債權 -> 倖 처럼 **틀린 한자**를 만든다).
반대로 아무것도 안 자르면 오염이 남는다. 그 사이 어디가 맞는지는 **재봐야 안다** —
내가 목록을 고르는 것이 아니라 곡선을 그려 사용자가 고르게 한다.

## 실측이 설계를 정했다 (2026-08-31)
- 오염 문자 1,126자 중 간체전용 22.5% · 가나 5.2% · 번체/공용 72.3%.
  ⇒ "간체만 자르기"로는 1/4 밖에 못 잡는다.
- 그런데 **토큰**으로 보면 다르다: 오염 토큰 498회 중 355회가 **2자 이상 복합**이고
  그 정체가 `您的` `贵公司` `具体时间` `本次会议` 같은 **중국어 단어**다.
  한국어는 저런 걸 통짜로 뱉지 않는다.
- 한국어 한자 병기는 낱글자로 쪼개진다: 개항(開港) -> '開','港' / 채권(債權) -> '債','權'.
  ⇒ **단일 한자 토큰을 살리면 병기가 산다.**
- ⛔ 단 예외가 있다: 목적(目的) -> '目的' 통짜, 사고(思考) -> '思考' 통짜.
  복합 토큰을 다 자르면 이런 한국 한자어도 죽는다. 그 대가가 M2 의 비용이다.

## 세 후보
- **M1 보수적**: 가나 + 간체전용 글자를 포함한 토큰. 한국어 위험 0.
- **M2 중간**  : M1 + 2자 이상 순수 한자 토큰(중국어 단어). 목적/사고류가 대가.
- **M3 전면**  : 한자·가나 포함 토큰 전부. 이미 측정됨(0.06%).
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
HAN = [(0x4E00, 0x9FFF), (0x3400, 0x4DBF), (0xF900, 0xFAFF)]
KANA = [(0x3040, 0x309F), (0x30A0, 0x30FF), (0x31F0, 0x31FF)]
HANGUL = [(0xAC00, 0xD7A3), (0x1100, 0x11FF), (0x3130, 0x318F)]


def inr(cp, rs):
    return any(a <= cp <= b for a, b in rs)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen3.8-27B")
    ap.add_argument("--outdir", default=str(HERE))
    a = ap.parse_args()

    from transformers import AutoTokenizer
    from opencc import OpenCC
    s2t = OpenCC("s2t")
    tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
    n = max(tok.get_vocab().values()) + 1

    m1, m2, m3 = [], [], []
    stats = {"scanned": 0, "hangul_spared": 0}
    for tid in range(n):
        try:
            s = tok.decode([tid], skip_special_tokens=False)
        except Exception:
            continue
        if not s:
            continue
        cps = [ord(c) for c in s]
        has_han = any(inr(c, HAN) for c in cps)
        has_kana = any(inr(c, KANA) for c in cps)
        if not (has_han or has_kana):
            continue
        if any(inr(c, HANGUL) for c in cps):      # ⛔ 한글 섞이면 절대 보존
            stats["hangul_spared"] += 1
            continue
        stats["scanned"] += 1
        m3.append(tid)

        han_chars = [c for c in s if inr(ord(c), HAN)]
        # 간체 전용 = 번체로 바꾸면 글자가 달라지는 것
        simplified = any(s2t.convert(c) != c for c in han_chars)
        if has_kana or simplified:
            m1.append(tid); m2.append(tid); continue
        # 2자 이상 순수 한자 토큰 = 중국어 단어일 가능성 (한국 한자어도 일부 포함)
        if len(han_chars) >= 2 and len(s.strip()) == len(han_chars):
            m2.append(tid)

    out = {}
    for name, ids, desc in (
        ("m1-conservative", m1, "가나 + 간체전용 글자 포함 토큰 (한국어 위험 0)"),
        ("m2-medium", m2, "M1 + 2자이상 순수한자 토큰 (중국어 단어; 목적/사고류가 대가)"),
        ("m3-full", m3, "한자·가나 포함 토큰 전부 (한자 병기 전멸)"),
    ):
        p = Path(a.outdir) / f"mask_{name}.json"
        p.write_text(json.dumps(
            {"name": name, "desc": desc, "model": a.model,
             "target_token_ids": sorted(ids), "n_target": len(ids)},
            ensure_ascii=False), encoding="utf-8")
        out[name] = len(ids)
        print(f"  {name:<18} {len(ids):>7,}개  {desc}")

    print(f"\n  한글 섞여 보존: {stats['hangul_spared']}개")
    # 감각 확인 — 대표 토큰이 어느 마스크에 들어가나
    print("\n  표본 판정:")
    s1, s2_, s3 = set(m1), set(m2), set(m3)
    for probe in ["開", "港", "目的", "思考", "您的", "贵公司", "具体时间", "の", "案"]:
        ids = tok.encode(probe, add_special_tokens=False)
        if len(ids) != 1:
            print(f"    {probe:<8} (복수토큰 {len(ids)}개 — 판정 생략)")
            continue
        t = ids[0]
        tags = [nm for nm, st in (("M1", s1), ("M2", s2_), ("M3", s3)) if t in st]
        print(f"    {probe:<8} 잘림: {tags or ['(없음)']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())