"""102 · AI 文字水印实验（本机 Mac mini，Qwen3-4B-Instruct-2507，MPS）

两种水印：
  green   —— Kirchenbauer et al. 2023 的「绿名单」软水印：用密钥 + 前一个词决定这一步哪一半词表是绿的，
             绿词的分数加 delta；检测 = 数绿词，算 z 分数。
  synthid —— Google DeepMind SynthID-Text（Nature 2024）的开源实现（transformers 自带），
             Anthropic 说 Claude 用的就是它的一个版本；检测 = 平均 g 值（没水印时约 0.5）。

用法：
  python wm_lab.py gen  <name> <mode none|green|synthid> <prompt_file> [--seed N] [--max 300]
  python wm_lab.py score <text_file>     （两种检测都打分，输出 json）
"""
import json, math, sys, os, hashlib, argparse, time
import torch
from transformers import (AutoModelForCausalLM, AutoTokenizer, LogitsProcessor, LogitsProcessorList,
                          SynthIDTextWatermarkingConfig, SynthIDTextWatermarkLogitsProcessor)

HERE = os.path.dirname(os.path.abspath(__file__))
MODEL = os.path.join(HERE, 'models', 'qwen3-4b')
OUT = os.path.join(HERE, 'out')
os.makedirs(OUT, exist_ok=True)

# 固定的「密钥」：实验前定好，全程不换
GREEN_KEY = 20261102
GAMMA = 0.5
DELTA = 2.0
SYNTH_KEYS = [int(k, 16) % 1000 for k in (hashlib.sha256(b'mia-102-%d' % i).hexdigest()[:6] for i in range(30))]
NGRAM = 5
# transformers 的 SynthID 随机表按设备生成（MPS 和 CPU 不一样），生成在 MPS 上做，检测也必须在 MPS 上算
SYNTH_DEVICE = 'mps'

_tok = None
_model = None
_green_cache = {}


def tok():
    global _tok
    if _tok is None:
        _tok = AutoTokenizer.from_pretrained(MODEL)
    return _tok


def model():
    global _model
    if _model is None:
        _model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to('mps')
        _model.eval()
    return _model


VOCAB = None


def green_set(prev_id, vocab):
    """前一个词 + 密钥 → 这一步的绿名单（bool 向量）。"""
    k = (prev_id, vocab)
    if k in _green_cache:
        return _green_cache[k]
    g = torch.Generator(device='cpu')
    g.manual_seed((GREEN_KEY * 1_000_003 + int(prev_id)) % (2 ** 62))
    perm = torch.randperm(vocab, generator=g)
    mask = torch.zeros(vocab, dtype=torch.bool)
    mask[perm[: int(GAMMA * vocab)]] = True
    _green_cache[k] = mask
    return mask


class GreenList(LogitsProcessor):
    def __call__(self, input_ids, scores):
        vocab = scores.shape[-1]
        for b in range(input_ids.shape[0]):
            m = green_set(int(input_ids[b, -1]), vocab).to(scores.device)
            scores[b] = scores[b] + DELTA * m.to(scores.dtype)
        return scores


def synth_config():
    return SynthIDTextWatermarkingConfig(keys=SYNTH_KEYS, ngram_len=NGRAM)


def synth_proc():
    return SynthIDTextWatermarkLogitsProcessor(ngram_len=NGRAM, keys=SYNTH_KEYS, sampling_table_size=2 ** 16,
                                               sampling_table_seed=0, context_history_size=1024, device=SYNTH_DEVICE)


def chat_ids(prompt):
    msgs = [{"role": "user", "content": prompt}]
    s = tok().apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    return tok()(s, return_tensors='pt').input_ids


def generate(prompt, mode, seed=0, max_new=300, temp=0.7, top_p=0.8, top_k=20):
    torch.manual_seed(seed)
    torch.mps.manual_seed(seed)
    ids = chat_ids(prompt).to('mps')
    kw = dict(max_new_tokens=max_new, do_sample=True, temperature=temp, top_p=top_p, top_k=top_k)
    if mode == 'green':
        kw['logits_processor'] = LogitsProcessorList([GreenList()])
    elif mode == 'synthid':
        kw['watermarking_config'] = synth_config()
    t0 = time.time()
    with torch.no_grad():
        out = model().generate(ids, **kw)
    gen = out[0, ids.shape[1]:].tolist()
    eos = tok().eos_token_id
    gen = [t for t in gen if t not in (eos, tok().pad_token_id) and t < len(tok())]
    text = tok().decode(gen, skip_special_tokens=True)
    return text, gen, time.time() - t0


def text_ids(text):
    return tok()(text, add_special_tokens=False).input_ids


def score_green(text):
    ids = text_ids(text)
    vocab = model().config.vocab_size if _model is not None else 151936
    toks = []
    green = 0
    n = 0
    seen = set()
    for i in range(len(ids)):
        piece = tok().decode([ids[i]])
        if i == 0:
            toks.append({'t': piece, 'g': None})
            continue
        pair = (ids[i - 1], ids[i])
        isg = bool(green_set(ids[i - 1], vocab)[ids[i]])
        rep = pair in seen
        seen.add(pair)
        toks.append({'t': piece, 'g': isg, 'rep': rep})
        if not rep:
            n += 1
            green += int(isg)
    z = (green - GAMMA * n) / math.sqrt(n * GAMMA * (1 - GAMMA)) if n else 0.0
    p = 0.5 * math.erfc(z / math.sqrt(2))
    return {'n': n, 'green': green, 'frac': green / n if n else 0, 'z': z, 'p': p, 'tokens': toks}


def score_synth(text):
    ids = torch.tensor([text_ids(text)]).to(SYNTH_DEVICE)
    proc = synth_proc()
    if ids.shape[1] < NGRAM + 1:
        return {'n': 0, 'mean': 0.5, 'z': 0.0, 'tokens': []}
    g = proc.compute_g_values(ids)[0].float().cpu()            # (L-4, 30)
    mask = proc.compute_context_repetition_mask(ids)[0].cpu()  # (L-4,)
    sel = g[mask.bool()]
    n = sel.shape[0]
    mean = float(sel.mean()) if n else 0.5
    z = (mean - 0.5) / (0.5 / math.sqrt(n * g.shape[1])) if n else 0.0
    per = g.mean(-1).tolist()
    pieces = [tok().decode([t]) for t in ids[0].cpu().tolist()]
    toks = [{'t': pieces[i], 'g': None} for i in range(NGRAM - 1)]
    for j, v in enumerate(per):
        toks.append({'t': pieces[j + NGRAM - 1], 'g': round(v, 3), 'rep': not bool(mask[j])})
    return {'n': n, 'mean': mean, 'z': z, 'p': 0.5 * math.erfc(z / math.sqrt(2)), 'tokens': toks}


def score_both(text):
    g = score_green(text)
    s = score_synth(text)
    return {'green': g, 'synthid': s}


def prefix_curve(text, steps=(10, 20, 40, 60, 80, 100, 150, 200, 250, 300)):
    ids = text_ids(text)
    rows = []
    for k in steps:
        if k > len(ids):
            break
        sub = tok().decode(ids[:k])
        g = score_green(sub)
        s = score_synth(sub)
        rows.append({'tokens': k, 'green_z': g['z'], 'green_frac': g['frac'], 'synth_z': s['z'], 'synth_mean': s['mean']})
    return rows


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('cmd')
    ap.add_argument('args', nargs='*')
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--max', type=int, default=300)
    a = ap.parse_args()
    if a.cmd == 'gen':
        name, mode, pf = a.args
        prompt = open(pf).read().strip()
        text, gen, dt = generate(prompt, mode, a.seed, a.max)
        res = {'name': name, 'mode': mode, 'seed': a.seed, 'prompt': prompt, 'text': text, 'secs': dt,
               'n_gen': len(gen), 'score': score_both(text)}
        json.dump(res, open(os.path.join(OUT, f'{name}.json'), 'w'), ensure_ascii=False, indent=1)
        sc = res['score']
        print(f"{name}\t{mode}\tseed={a.seed}\t{len(gen)} tok\t{dt:.0f}s\tgreen {sc['green']['frac']:.2f} z={sc['green']['z']:.1f}\tsynth {sc['synthid']['mean']:.3f} z={sc['synthid']['z']:.1f}")
        print(text[:400])
    elif a.cmd == 'score':
        for pf in a.args:
            text = open(pf).read().strip()
            sc = score_both(text)
            name = os.path.splitext(os.path.basename(pf))[0]
            json.dump({'name': name, 'text': text, 'score': sc}, open(os.path.join(OUT, f'score-{name}.json'), 'w'), ensure_ascii=False, indent=1)
            print(f"{name}\tgreen {sc['green']['frac']:.2f} z={sc['green']['z']:.1f} (n={sc['green']['n']})\tsynth {sc['synthid']['mean']:.3f} z={sc['synthid']['z']:.1f} (n={sc['synthid']['n']})")
    elif a.cmd == 'curve':
        for pf in a.args:
            d = json.load(open(pf))
            rows = prefix_curve(d['text'])
            d['curve'] = rows
            json.dump(d, open(pf, 'w'), ensure_ascii=False, indent=1)
            for r in rows:
                print(r)