#!/usr/bin/env python3
"""exp15 — chunkers for Q1: a faithful zh source-sentence splitter + the §B1.5 LOGICAL DP chunker.

Q1 contrast (design §D2): A0 (greedy-1500, from the real backend) vs a LOGICAL-boundary arm sized to
A0's ACTUAL est_out distribution (±10% mean). The ONLY variable is cut ALIGNMENT (greedy is blind to
scene structure; logical aligns cuts to scene-shift boundaries and avoids mid-scene/dialogue cuts).
Size is held equal so Q1 does not confound boundary class with chunk size (that is Q2's domain).

Ports (backend/internal/pipeline/chunker.go):
  splitSourceSentences / isOpenQuote / isCloseQuote / isSourceTerminator / isCJKTerminator /
  isClosingBracket / isASCIISpace   (chunker.go:196-357) — quote-depth-aware, lossless tiling.
Token math:
  est_tokens = cjk + other//3           (EstimateTokens units — the greedy 1500 budget space)
  est_out    = 1.1978*cjk + 0.3852*other (fertility_calib.json — glm-4.6 ru output tokens; §B1.2)

The logical chunker implements §B1.4 boundary classes (B0 chapter/separator, B1 scene-shift lexicon,
B2 TextTiling cohesion valley, B3 paragraph, B4 sentence) and §B1.5 DP packing (Knuth-Plass shortest
path over candidate cuts, INTEGER-scaled costs, asymmetric size penalty). It does NOT peek at the
ground-truth seam offsets — scene-shifts are DISCOVERED (that is the Q1 test).

Self-review by execution: `python chunker.py --selftest` checks tiling losslessness + a size sweep.
"""
from __future__ import annotations

import math
import sys
import unicodedata
from dataclasses import dataclass
from pathlib import Path

import regex

_RE_CJK = regex.compile(r"[\p{Han}\p{Hiragana}\p{Katakana}\p{Hangul}]")
TARGET_CHUNK_TOKENS = 1500  # chunker.go:47 const

F_CJK, F_OTHER = 1.1978, 0.3852  # fertility_calib.json (glm-4.6 ru out tokens from src char classes)

# --- char-class token estimators -------------------------------------------------------------------
def _classes(text: str):
    cjk = other = 0
    for ch in text:
        if _RE_CJK.match(ch):
            cjk += 1
        elif ch.isspace():
            pass
        else:
            other += 1
    return cjk, other


def est_tokens(text: str) -> int:
    """EstimateTokens units (source side): cjk + other//3, floored at 16 (chunker.go estTokensFrom)."""
    cjk, other = _classes(text)
    v = cjk + other // 3
    return max(v, 16)


def est_out(text: str) -> float:
    """Predicted glm-4.6 ru OUTPUT tokens for a source slice (fertility model; §B1.2)."""
    cjk, other = _classes(text)
    return F_CJK * cjk + F_OTHER * other


# --- GREEDY chunker port (chunker.go:82-194 + render.go:285 NormalizeSource) -----------------------
def normalize_source(s: str) -> str:
    """render.go:285 — BOM strip, CRLF/CR->LF, valid-UTF8 (Python strs already are), NFC, TrimSpace."""
    s = s.lstrip("﻿")
    s = s.replace("\r\n", "\n").replace("\r", "\n")
    s = unicodedata.normalize("NFC", s)
    return s.strip()   # Go strings.TrimSpace strips Unicode whitespace incl. U+3000 (== Python str.strip())


def _class_counts(s: str):
    cjk = other = 0
    for ch in s:
        if _RE_CJK.match(ch):
            cjk += 1
        elif ch.isspace():
            pass
        else:
            other += 1
    return cjk, other


def _est_tokens_from(cjk, other):
    return max(cjk + other // 3, 16)   # chunker.go estTokensFrom (integer // 3, floor 16)


def split_paragraphs(chapter: str):
    return [p.strip() for p in chapter.split("\n\n") if p.strip()]


def _pack_sentences(paragraph: str):
    segs = split_source_sentences(paragraph)
    chunks, buf, bc, bo = [], [], 0, 0
    def flush():
        nonlocal buf, bc, bo
        if buf:
            chunks.append("".join(buf)); buf, bc, bo = [], 0, 0
    for s in segs:
        sc, so = _class_counts(s)
        if buf and _est_tokens_from(bc + sc, bo + so) > TARGET_CHUNK_TOKENS:
            flush()
        buf.append(s); bc, bo = bc + sc, bo + so
    flush()
    return chunks


def greedy_chunks(raw_text: str):
    """Port of chunker.go SplitChunks for ONE chapter (S2' has no chapter markers). Returns chunk texts."""
    chapter = normalize_source(raw_text)
    paras = split_paragraphs(chapter)
    chunks, buf, bc, bo = [], [], 0, 0
    def flush():
        nonlocal buf, bc, bo
        if buf:
            text = "\n\n".join(buf).strip()
            if text:
                chunks.append(text)
            buf, bc, bo = [], 0, 0
    for p in paras:
        pc, po = _class_counts(p)
        if _est_tokens_from(pc, po) > TARGET_CHUNK_TOKENS:
            flush()
            for sub in _pack_sentences(p):
                t = sub.strip()
                if t:
                    chunks.append(t)
            continue
        if buf and _est_tokens_from(bc + pc, bo + po) > TARGET_CHUNK_TOKENS:
            flush()
        buf.append(p); bc, bo = bc + pc, bo + po
    flush()
    return chunks


# --- source sentence splitter (port of chunker.go:196-357) -----------------------------------------
_OPEN_Q = set("「『（【《〈［｛“‘")
_CLOSE_Q = set("」』）】》〉］｝”’")
_TERM = set("。！？.!?…")
_CJK_TERM = set("。！？")
_CLOSE_BRK = set("」』】）)》〉］]｝}\"'”’")
_ASCII_SPACE = set(" \t\n\r\f\v　")


def split_source_sentences(paragraph: str):
    runes = list(paragraph)
    segs = []
    start = i = depth = 0
    n = len(runes)
    while i < n:
        r = runes[i]
        if r in _OPEN_Q:
            depth += 1
            i += 1
            continue
        if r in _CLOSE_Q:
            if depth > 0:
                depth -= 1
            i += 1
            continue
        if r not in _TERM:
            i += 1
            continue
        j = i
        cjk = False
        while j < n and runes[j] in _TERM:
            if runes[j] in _CJK_TERM:
                cjk = True
            j += 1
        depth_here = depth
        e = j
        while e < n and runes[e] in _CLOSE_BRK:
            if runes[e] in _CLOSE_Q and depth > 0:
                depth -= 1
            e += 1
        if cjk:
            boundary = depth_here == 0
        else:
            boundary = depth_here == 0 and (e >= n or runes[e] in _ASCII_SPACE)
            # abbrev guard omitted for zh source (CJK terminators dominate; en not in S2')
        if not boundary:
            i = e
            continue
        w = e
        while w < n and runes[w] in _ASCII_SPACE:
            w += 1
        segs.append("".join(runes[start:w]))
        start = i = w
    if start < n:
        segs.append("".join(runes[start:]))
    return segs


# --- units: split the text into atomic sentence-ish spans that TILE it -----------------------------
def units_from_text(text: str):
    """Atomic units for the DP = source sentences. Preserves whether each unit ENDS a paragraph
    (\\n\\n) or a line (\\n) for boundary classification. Tiling: ''.join(u.text) == text."""
    out = []
    # split into paragraph blocks on blank lines, remember the separator that followed each block
    # (so paragraph-end vs line-end vs sentence-end is recoverable). We tile exactly by re-attaching
    # the separators to the preceding unit.
    for seg, sep in _iter_blocks(text):
        sents = split_source_sentences(seg)
        for k, s in enumerate(sents):
            ends_para = (k == len(sents) - 1) and sep == "\n\n"
            ends_line = (k == len(sents) - 1) and sep == "\n"
            out.append((s, ends_para, ends_line))
        # re-attach the block separator to the last unit so tiling is exact
        if out and sep:
            s, ep, el = out[-1]
            out[-1] = (s + sep, ep, el)
    return out


def _iter_blocks(text: str):
    """Yield (block_text, separator) where separator is '\\n\\n' | '\\n' | ''. Tiles the input."""
    i = 0
    n = len(text)
    buf_start = 0
    while i < n:
        if text[i] == "\n":
            if i + 1 < n and text[i + 1] == "\n":
                yield text[buf_start:i], "\n\n"
                i += 2
                buf_start = i
                continue
            else:
                yield text[buf_start:i], "\n"
                i += 1
                buf_start = i
                continue
        i += 1
    if buf_start < n:
        yield text[buf_start:], ""


# --- boundary classification (§B1.4) ---------------------------------------------------------------
# B1 scene-shift lexicon: STRONG time/place/POV shift markers in the first chars of the NEXT unit.
# Curated generic zh markers (NOT the S2' seam text — scene-shifts must be DISCOVERED). Kept to strong
# temporal/spatial jumps + narrative shift particles; weak mid-scene connectors (顿时/忽然) excluded.
_B1_MARKERS = [
    # temporal jumps
    "次日", "翌日", "第二天", "第三天", "隔日", "数日后", "数天后", "几天后", "三天后", "三天三夜",
    "一个星期", "一周", "几个月", "半年", "一年后", "三年后", "五年后", "多年后",
    "当天", "这一天", "那一天", "这一日", "翌朝", "清晨", "早晨", "拂晓", "五更", "天亮", "天明",
    "傍晚", "黄昏", "夜里", "半夜", "深夜", "入夜", "朝阳", "红日", "夕阳", "日落", "日出", "正午",
    "转眼", "转眼间", "很快", "不久", "不多时", "没过多久", "光阴荏苒", "时光飞逝",
    # narrative-shift particles / POV or place switch
    "却说", "再说", "话说", "单说", "且说", "与此同时", "另一边", "另一方", "一方", "此时此刻",
    "回到", "镜头", "视线",
]


def _b1_scene_shift(next_unit_text: str) -> bool:
    head = next_unit_text.lstrip("　 \n\t")[:12]
    return any(m in head for m in _B1_MARKERS)


def _dialogue_boundary_penalty(prev_unit: str, next_unit: str) -> int:
    """Penalize cutting between a reply and its attribution or inside a dialogue exchange (§A.6):
    if prev ends inside/with a quote and next continues speech attribution, cutting is bad."""
    p = prev_unit.rstrip("　 \n\t")
    nx = next_unit.lstrip("　 \n\t")
    pen = 0
    # prev ends with a closing quote AND next starts with attribution-ish (说道/道/问/答 ...) -> keep together
    if p and p[-1] in _CLOSE_Q and nx[:2] in ("说道", "说", "道", "问道", "答道", "喝道", "笑道"):
        pen += 1
    # next starts a quote while prev did not close one cleanly (mid-exchange) — mild
    if nx[:1] in _OPEN_Q and p and p[-1] in _OPEN_Q:
        pen += 1
    return pen


# --- TextTiling cohesion valley (B2) ---------------------------------------------------------------
def _char_bigrams(text: str):
    t = [c for c in text if not c.isspace()]
    return [t[k] + t[k + 1] for k in range(len(t) - 1)]


def _cosine(a: dict, b: dict) -> float:
    if not a or not b:
        return 0.0
    dot = sum(v * b.get(k, 0) for k, v in a.items())
    na = math.sqrt(sum(v * v for v in a.values()))
    nb = math.sqrt(sum(v * v for v in b.values()))
    return dot / (na * nb) if na and nb else 0.0


def _bag(units, lo, hi):
    d = {}
    for u in units[lo:hi]:
        for bg in _char_bigrams(u[0]):
            d[bg] = d.get(bg, 0) + 1
    return d


def texttiling_depths(units, w=4):
    """Depth score of the lexical-cohesion valley at each inter-unit gap (higher = stronger boundary).
    Sliding window of w units either side; cosine of char-bigram bags; depth = local valley prominence."""
    n = len(units)
    gaps = n - 1
    sims = [0.0] * gaps
    for g in range(gaps):
        left = _bag(units, max(0, g + 1 - w), g + 1)
        right = _bag(units, g + 1, min(n, g + 1 + w))
        sims[g] = _cosine(left, right)
    # depth score: for each gap, (peak_left - sim) + (peak_right - sim)
    depths = [0.0] * gaps
    for g in range(gaps):
        lpk = sims[g]
        k = g
        while k > 0 and sims[k - 1] >= sims[k]:
            k -= 1
            lpk = max(lpk, sims[k])
        rpk = sims[g]
        k = g
        while k < gaps - 1 and sims[k + 1] >= sims[k]:
            k += 1
            rpk = max(rpk, sims[k])
        depths[g] = (lpk - sims[g]) + (rpk - sims[g])
    return depths


# --- §B1.5 DP logical chunker ----------------------------------------------------------------------
@dataclass
class ChunkPlan:
    chunks: list          # list of (text, est_out, right_boundary_class)
    cut_offsets: list     # rune offsets in the ORIGINAL text where a cut falls (chunk starts, excl 0)


# boundary-class penalty (§B1.5): B0=0 < B1 < B2 < B3 < B4 (prefer cutting at strong boundaries)
_CLASS_PEN = {"B0": 0, "B1": 1, "B2": 3, "B3": 6, "B4": 10}


def logical_chunks(text: str, target_out: float, hard_max_out: float, min_out: float = None,
                   w_size=3, w_bound=60, w_dialog=200, b2_z=0.8,
                   corridor_lo=0.62, corridor_hi=1.12) -> ChunkPlan:
    """§B1.5 DP: minimize Σ [w_size·size_pen + w_bound·class_pen(right boundary) + w_dialog·dialogue].
    Cuts only at unit (sentence) boundaries — never mid-sentence (B5 forbidden). Costs INTEGER-scaled.
    target_out/hard_max_out = A0's actual est_out mean/max (size-match)."""
    units = units_from_text(text)
    n = len(units)
    if min_out is None:
        min_out = 0.45 * target_out
    # per-unit est_out and cumulative rune offsets
    out_of = [est_out(u[0]) for u in units]
    cum_out = [0.0] * (n + 1)
    cum_rune = [0] * (n + 1)
    for k in range(n):
        cum_out[k + 1] = cum_out[k] + out_of[k]
        cum_rune[k + 1] = cum_rune[k] + len(units[k][0])

    # boundary class at each inter-unit gap g (cut AFTER unit g, i.e. between unit g and g+1)
    depths = texttiling_depths(units)
    if depths:
        mu = sum(depths) / len(depths)
        sd = math.sqrt(sum((d - mu) ** 2 for d in depths) / len(depths)) or 1e-9
    gap_class = []
    gap_dialog = []
    for g in range(n - 1):
        cls = "B4"                                   # default: sentence boundary
        if units[g][2] or units[g][1]:
            cls = "B3"                                # ends a line/paragraph
        if depths and (depths[g] - mu) / sd >= b2_z:
            cls = "B2"                                # TextTiling cohesion valley
        if _b1_scene_shift(units[g + 1][0]):
            cls = "B1"                                # scene-shift lexicon in next unit
        gap_class.append(cls)
        gap_dialog.append(_dialogue_boundary_penalty(units[g][0], units[g + 1][0]))

    lo_band, hi_band = corridor_lo * target_out, corridor_hi * target_out
    def size_pen(a, b):
        """Corridor size penalty (§B1.5 intent: size defines a BUDGET CORRIDOR; boundary class decides
        WITHIN it). Flat (0) inside [lo_band, hi_band]; steep outside so chunks stay near target size."""
        o = cum_out[b] - cum_out[a]
        if o > hard_max_out:
            return 100000 + int(40 * (o - hard_max_out))   # hard ceiling — never exceed
        if o > hi_band:
            return int(30 * (o - hi_band))                 # over corridor
        if o < lo_band:
            return int(30 * (lo_band - o))                 # under corridor
        return 0                                           # inside corridor: size is free, boundary decides

    INF = float("inf")
    # dp[b] = min cost to segment units[0:b]; back[b] = previous cut index a
    dp = [INF] * (n + 1)
    back = [0] * (n + 1)
    dp[0] = 0
    for b in range(1, n + 1):
        # try all chunk starts a < b (bounded to a window where size is plausible)
        a = b - 1
        while a >= 0:
            o = cum_out[b] - cum_out[a]
            if o > hard_max_out * 1.6 and a < b - 1:
                break   # window: chunks far over hardMax are never optimal
            cost = dp[a] + w_size * size_pen(a, b)
            if b < n:                                  # boundary penalty on the RIGHT cut (gap b-1)
                cost += w_bound * _CLASS_PEN[gap_class[b - 1]] + w_dialog * gap_dialog[b - 1]
            if cost < dp[b]:
                dp[b] = cost
                back[b] = a
            a -= 1
    # reconstruct
    cuts = []
    b = n
    while b > 0:
        a = back[b]
        cuts.append((a, b))
        b = a
    cuts.reverse()
    chunks = []
    cut_offsets = []
    for (a, b) in cuts:
        txt = "".join(units[k][0] for k in range(a, b))
        rcls = gap_class[b - 1] if b < n else "B0"
        chunks.append((txt, est_out(txt), rcls))
        if a > 0:
            cut_offsets.append(cum_rune[a])
    return ChunkPlan(chunks=chunks, cut_offsets=cut_offsets)


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        text = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime.txt").read_text(encoding="utf-8")
        units = units_from_text(text)
        tiled = "".join(u[0] for u in units)
        print("units:", len(units), "tiling lossless:", tiled == text)
        assert tiled == text, "UNIT TILING NOT LOSSLESS"
        total_out = est_out(text)
        print(f"S2' total est_out={total_out:.0f}  est_tokens={est_tokens(text)}")
        # size sweep: how many logical chunks at a few targets
        for tgt in (1500, 1797, 2200):
            plan = logical_chunks(text, target_out=tgt, hard_max_out=tgt * 1.25)
            outs = [c[1] for c in plan.chunks]
            classes = [c[2] for c in plan.chunks[:-1]]  # right-boundary classes (excl last)
            from collections import Counter
            print(f"  target_out={tgt}: {len(plan.chunks)} chunks  "
                  f"mean_out={sum(outs)/len(outs):.0f}  min={min(outs):.0f} max={max(outs):.0f}  "
                  f"cut-classes={dict(Counter(classes))}")
        print("chunker selftest OK")
