#!/usr/bin/env python3
"""WS2 (г) verification — INDEPENDENT re-derivation of the output-token fertility
coefficients (L2-budget-wrong-unit fix), NOT trusting fertility_calib.json.

Reproduces est_out = f_cjk*cjk_src + f_other*other_src (OLS, no intercept) over the
25-chapter guzhenren rerun, and cross-checks the §1.6 finding: fertility_calib's
ord-range CJK classifier differs from the backend EstimateTokens class set
(unicode.Han|Hiragana|Katakana|Hangul). We fit under BOTH classifiers and report the
coefficient delta so the plan can decide (align vs document-invariance).

$0: deterministic, no LLM, no network. Reuses eval/tokenizers/glm-4.6 + rerun/records.json.
Writes result to /home/ubuntu/books/gu-zhenren/design11/ws2_fertility_verify.json (OUT of git).
"""
import json, os, sys
import numpy as np
from tokenizers import Tokenizer

ROOT = "/home/ubuntu/projects/textmachine"
RECORDS = "/home/ubuntu/books/gu-zhenren/rerun/records.json"
TOK = f"{ROOT}/eval/tokenizers/glm-4.6/tokenizer.json"
STORED = "/home/ubuntu/books/gu-zhenren/exp15/fertility_calib.json"
OUT = "/home/ubuntu/books/gu-zhenren/design11/ws2_fertility_verify.json"

# ---- classifier A: fertility_calib.py is_cjk (ord ranges, NO Hangul) ----
def is_cjk_ord(c):
    o = ord(c)
    return (0x4E00 <= o <= 0x9FFF or 0x3040 <= o <= 0x30FF or 0x3400 <= o <= 0x4DBF
            or 0xF900 <= o <= 0xFAFF or 0x20000 <= o <= 0x2A6DF)

# ---- classifier B: backend EstimateTokens (unicode property Han|Hira|Kata|Hangul) ----
import unicodedata
def is_cjk_prop(c):
    # Approximate Go unicode.Han|Hiragana|Katakana|Hangul via Unicode script.
    # Python has no stdlib script table; use unicodedata name prefixes as a proxy that
    # matches the Go ranges for the scripts in question.
    try:
        name = unicodedata.name(c)
    except ValueError:
        return False
    o = ord(c)
    # Han (CJK ideographs incl ext), Hiragana, Katakana, Hangul syllables/jamo
    if (0x3400 <= o <= 0x9FFF or 0xF900 <= o <= 0xFAFF or 0x20000 <= o <= 0x2FA1F):
        return "CJK" in name or "IDEOGRAPH" in name
    if 0x3040 <= o <= 0x309F: return True   # Hiragana
    if 0x30A0 <= o <= 0x30FF: return True   # Katakana (incl. 0x30FB middle dot? Go: Katakana table)
    if 0xAC00 <= o <= 0xD7A3 or 0x1100 <= o <= 0x11FF or 0x3130 <= o <= 0x318F: return True  # Hangul
    return False

def classes(text, is_cjk):
    cjk = other = 0
    for c in text:
        if c.isspace():
            continue
        if is_cjk(c):
            cjk += 1
        else:
            other += 1
    return cjk, other

def fit(rows):
    A = np.array([[c, o] for (c, o, _) in rows], dtype=float)
    y = np.array([t for (_, _, t) in rows], dtype=float)
    coef, *_ = np.linalg.lstsq(A, y, rcond=None)
    pred = A @ coef
    ss_res = float(np.sum((y - pred) ** 2))
    ss_tot = float(np.sum((y - y.mean()) ** 2))
    r2 = 1 - ss_res / ss_tot
    mape = float(np.mean(np.abs(y - pred) / y) * 100)
    return float(coef[0]), float(coef[1]), r2, mape

def main():
    tok = Tokenizer.from_file(TOK)
    recs = json.load(open(RECORDS))
    rows_A, rows_B = [], []
    used = skipped_empty = skipped_tiny = 0
    for r in recs:
        src = r.get("source", "") or ""
        fin = r.get("final", "") or ""
        if not src.strip() or not fin.strip():
            skipped_empty += 1
            continue
        cA, oA = classes(src, is_cjk_ord)
        if cA + oA < 200:
            skipped_tiny += 1
            continue
        out_tok = len(tok.encode(fin).ids)
        cB, oB = classes(src, is_cjk_prop)
        rows_A.append((cA, oA, out_tok))
        rows_B.append((cB, oB, out_tok))
        used += 1

    fa = fit(rows_A)   # classifier A (ord ranges, as fertility_calib)
    fb = fit(rows_B)   # classifier B (backend property set)
    stored = json.load(open(STORED))

    result = {
        "n_records_used": used, "skipped_empty": skipped_empty, "skipped_tiny": skipped_tiny,
        "total_records": len(recs),
        "classifier_A_ord_ranges (fertility_calib)": {
            "f_cjk": round(fa[0], 4), "f_other": round(fa[1], 4),
            "r2": round(fa[2], 4), "mape_pct": round(fa[3], 2)},
        "classifier_B_backend_property (Han|Hira|Kata|Hangul)": {
            "f_cjk": round(fb[0], 4), "f_other": round(fb[1], 4),
            "r2": round(fb[2], 4), "mape_pct": round(fb[3], 2)},
        "stored_artifact": {"f_cjk": stored["f_cjk"], "f_other": stored["f_other"],
                            "r2": stored["r2"], "mape_pct": stored["mape_pct"],
                            "n_records": stored["n_records"]},
        "coef_delta_A_vs_stored": {"f_cjk": round(abs(fa[0] - stored["f_cjk"]), 6),
                                   "f_other": round(abs(fa[1] - stored["f_other"]), 6)},
        "coef_delta_B_vs_A (classifier mismatch impact)": {
            "f_cjk": round(fb[0] - fa[0], 4), "f_other": round(fb[1] - fa[1], 4),
            "pct_est_out_shift_on_median_zh_chunk": None},
    }
    # quantify classifier mismatch: predicted est_out on the median src chunk under A vs B
    med = rows_A[len(rows_A)//2]
    estA = fa[0]*med[0] + fa[1]*med[1]
    # use classifier B counts for the same chunk under fit B
    medB = rows_B[len(rows_B)//2]
    estB = fb[0]*medB[0] + fb[1]*medB[1]
    result["coef_delta_B_vs_A (classifier mismatch impact)"]["pct_est_out_shift_on_median_zh_chunk"] = round((estB-estA)/estA*100, 3)

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    json.dump(result, open(OUT, "w"), ensure_ascii=False, indent=2)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    # PASS assertion: independent re-derivation matches stored to <1e-3
    dA = result["coef_delta_A_vs_stored"]
    ok = dA["f_cjk"] < 1e-3 and dA["f_other"] < 1e-3
    print(f"\nVERDICT: independent re-derivation {'REPRODUCES' if ok else 'DIVERGES FROM'} stored coefficients "
          f"(Δf_cjk={dA['f_cjk']}, Δf_other={dA['f_other']}).")

if __name__ == "__main__":
    main()