#!/usr/bin/env python3
"""exp15 — fertility calibration B1.2 ($0, offline): src chars -> ru output-token coefficients.

Regresses ru output tokens of the 25-chapter rerun finals on (cjk, other) source-char counts, giving
the chunker's est_out = f_cjk*cjk + f_other*other (§B1.1/B1.2). Chunker-spec deliverable feeding the
Q1 size-match. Offline tokenizer per §B1.2: deepseek-v4 tokenizer dir is EMPTY in repo, so we use the
editor-family glm-4.6 tokenizer (present) and label it; Go needs only the COEFFICIENTS, not a tokenizer.

Run: eval/.venv/bin/python eval/exp15/fertility_calib.py
"""
from __future__ import annotations
import json
from pathlib import Path

import numpy as np
from tokenizers import Tokenizer

REC = Path("/home/ubuntu/books/gu-zhenren/rerun/records.json")
TOK = Path("/home/ubuntu/projects/textmachine/eval/tokenizers/glm-4.6/tokenizer.json")
OUT = Path("/home/ubuntu/books/gu-zhenren/exp15/fertility_calib.json")


def is_cjk(ch: str) -> bool:
    o = ord(ch)
    return (0x4E00 <= o <= 0x9FFF or 0x3040 <= o <= 0x30FF or 0x3400 <= o <= 0x4DBF
            or 0xF900 <= o <= 0xFAFF or 0x20000 <= o <= 0x2A6DF)


def char_classes(s: str):
    cjk = sum(1 for c in s if is_cjk(c))
    other = sum(1 for c in s if not c.isspace() and not is_cjk(c))
    return cjk, other


def main():
    tok = Tokenizer.from_file(str(TOK))
    recs = json.load(open(REC))
    rows, meta = [], []
    for r in recs:
        src = r.get("source") or ""
        fin = r.get("final") or ""
        if not src.strip() or not fin.strip():
            continue
        cjk, other = char_classes(src)
        out_tok = len(tok.encode(fin).ids)
        if cjk + other < 200:      # skip tiny/metadata chunks (e.g. title block)
            continue
        rows.append((cjk, other, out_tok))
        meta.append((r.get("chapter"), r.get("chunk_idx")))

    A = np.array([[c, o] for c, o, _ in rows], dtype=float)
    y = np.array([t for _, _, t in rows], dtype=float)
    # non-negative-ish least squares, no intercept (est_out = f_cjk*cjk + f_other*other)
    coef, *_ = np.linalg.lstsq(A, y, rcond=None)
    f_cjk, f_other = float(coef[0]), float(coef[1])
    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 if ss_tot else float("nan")
    mape = float(np.mean(np.abs((y - pred) / y))) * 100

    result = {
        "n_records": len(rows),
        "tokenizer": "glm-4.6 (editor family; deepseek-v4 tokenizer dir empty in repo — §B1.2 caveat)",
        "model": "est_out = f_cjk*cjk_src + f_other*other_src (no intercept)",
        "f_cjk": round(f_cjk, 4),
        "f_other": round(f_other, 4),
        "r2": round(r2, 4),
        "mape_pct": round(mape, 2),
        "placeholder_B1_1": {"cjk_char": 0.75, "other_char": 0.25},
        "note": "per-language-pair config (layer 2); recompute = loud resnapshot. Src-char est is the SECOND "
                "quantity (§B1.2); primary budget is ru-out tokens. Calibrated on 25-ch guzhenren rerun (in-pretrain).",
    }
    OUT.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"fertility calibration (n={len(rows)} chunks, glm-4.6 ru tokens):")
    print(f"  est_out = {f_cjk:.4f}*cjk + {f_other:.4f}*other   R2={r2:.4f}  MAPE={mape:.1f}%")
    print(f"  vs §B1.1 placeholder cjk=0.75/other=0.25")
    print(f"  sample src cjk/other -> out (first 6):")
    for (c, o, t), m in list(zip(rows, meta))[:6]:
        print(f"    ch{m[0]}#{m[1]}: cjk={c} other={o} out_tok={t}  pred={f_cjk*c+f_other*o:.0f}")
    print(f"-> {OUT}")
    # self-review: coefficients must be positive and R2 reasonable
    assert f_cjk > 0, f"non-positive cjk fertility {f_cjk}"
    assert r2 > 0.5, f"poor fit R2={r2} — investigate before trusting coefficients"
    print("OK: positive fertility coefficients, R2>0.5.")


if __name__ == "__main__":
    main()