#!/usr/bin/env python3
"""exp16 — bank-mining substrate (research/20 §B1/§D). Deterministic, $0. miner-v1.

Shared layer for the detector arms A1–A3: candidate generation (Han char n-grams 1–6, NO word
segmentation — §B1/Qiu&Zhang), general-zh contrast corpus (jieba dict.txt, versioned artifact),
GT loading + occurrence counting (Aho-Corasick-equivalent WITHOUT suppressContained — D24.2), frequency
strata, the 60/40 tuning/test chapter split, and the catastrophe screen.

Reuses the VALIDATED backend port eval/exp15/mem_select.py for memnorm normalization (byte-faithful to
memnorm.go). Candidate normalization is memnorm-symmetric so candidate keys and GT keys live in the same
space (trad2simp-folded, NFKC, lowercased) — and comparable to jieba's simplified-zh dict.

Self-review by execution: `python exp16_common.py --selftest`.
"""
from __future__ import annotations

import json
import math
import sys
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path

import regex

HERE = Path(__file__).resolve().parent
# APPEND exp15 (not insert(0)) so exp15's colliding module names (arms.py/chunker.py) do NOT shadow
# exp16's own modules; only unique names like mem_select are resolved from exp15.
sys.path.append(str(HERE.parent / "exp15"))
import mem_select as MS   # validated memnorm port

BOOK = Path("/home/ubuntu/books/gu-zhenren")
SEED = BOOK / "guzhenren-seed-v2.yaml"
RECORDS = BOOK / "rerun" / "records.json"
CONTRAST = HERE / "data" / "jieba_dict_general_zh.txt"   # versioned general-zh word-freq artifact

_RE_HAN = regex.compile(r"\p{Han}")

# ── frozen substrate params (miner-v1); thresholds tuned on the tuning half are set in detectors.py ──
NGRAM_MIN, NGRAM_MAX = 1, 6
FREQ_FLOOR = 3                       # §A1/§B1: below this V-A does not score (conscious blind spot)
CATASTROPHE = ["方源", "蛊", "蛊师", "古月"]   # §D4-в: must be in top-50 of the winner
# 60/40 chapter split (§D1): thresholds chosen on tuning, frozen, only test reported.
TUNING_CH = set(range(1, 16))        # chapters 1–15
TEST_CH = set(range(16, 26))         # chapters 16–25
ANNOTATION_CHUNK = (1, 0)            # ch1/chunk0 = cover blurb + metadata (spoiler-region for blurb terms)


def norm(s: str) -> str:
    return MS.normalize_source_key(s)


def han_only(s: str) -> bool:
    return bool(s) and all(_RE_HAN.match(c) for c in s)


# ── materials ───────────────────────────────────────────────────────────────────────────────────
@dataclass
class Chunk:
    chapter: int
    chunk_idx: int
    source: str          # raw zh
    draft: str           # raw ru draft (deepseek-v4-flash, WITH seed injection — confound §D1)
    nsource: str = ""    # memnorm-normalized source

    @property
    def is_annotation(self) -> bool:
        return (self.chapter, self.chunk_idx) == ANNOTATION_CHUNK


def load_chunks() -> list[Chunk]:
    recs = json.load(open(RECORDS, encoding="utf-8"))
    out = []
    for r in recs:
        c = Chunk(chapter=r["chapter"], chunk_idx=r["chunk_idx"],
                  source=r["source"], draft=r.get("draft", ""))
        c.nsource = norm(c.source)
        out.append(c)
    out.sort(key=lambda c: (c.chapter, c.chunk_idx))
    return out


@dataclass
class GTEntity:
    src: str
    dst: str
    typ: str
    status: str
    since_ch: int
    until_ch: int
    aliases: list = field(default_factory=list)     # alias surface strings
    gender: str = ""

    @property
    def surfaces(self) -> list[str]:
        return [self.src] + list(self.aliases)

    @property
    def norm_surfaces(self) -> list[str]:
        return [norm(s) for s in self.surfaces if norm(s)]


def load_gt() -> list[GTEntity]:
    rows = MS.load_seed_rows(str(SEED))
    out = []
    for r in rows:
        out.append(GTEntity(
            src=r["src"], dst=r.get("dst", "") or "", typ=r.get("type", "") or "",
            status=r.get("status") or "approved",
            since_ch=int(r.get("since_ch", 0) or 0), until_ch=int(r.get("until_ch", 0) or 0),
            aliases=[a["alias"] for a in (r.get("aliases") or []) if a.get("alias")],
            gender=r.get("gender", "") or ""))
    return out


# ── occurrence counting (AC-equivalent, overlapping, NO suppressContained — D24.2) ────────────────
def count_occurrences(needle_norm: str, chunks: list[Chunk], include_annotation: bool = True) -> int:
    """Total overlapping occurrences of a normalized needle across chunk sources."""
    if not needle_norm:
        return 0
    total = 0
    for c in chunks:
        if c.is_annotation and not include_annotation:
            continue
        hay = c.nsource
        pos = 0
        while True:
            i = hay.find(needle_norm, pos)
            if i < 0:
                break
            total += 1
            pos = i + 1
    return total


def gt_occurrences(ent: GTEntity, chunks: list[Chunk], include_annotation: bool = True) -> int:
    """Entity occurrence count = sum over its normalized surfaces (src + aliases), overlapping."""
    return sum(count_occurrences(ns, chunks, include_annotation) for ns in ent.norm_surfaces)


def gt_chapters(ent: GTEntity, chunks: list[Chunk], include_annotation: bool = True) -> set[int]:
    chs = set()
    for c in chunks:
        if c.is_annotation and not include_annotation:
            continue
        if any(ns and ns in c.nsource for ns in ent.norm_surfaces):
            chs.add(c.chapter)
    return chs


def freq_stratum(f: int) -> str:
    if f >= 10:
        return "f>=10"
    if f >= 3:
        return "f3-9"
    return "f<3"


# ── general-zh contrast corpus (jieba dict.txt: 'word freq POS' per line) ──────────────────────────
class Contrast:
    """General-domain zh reference: word frequencies + derived char frequencies. Versioned artifact."""

    def __init__(self, path: Path = CONTRAST):
        self.word_freq: dict[str, int] = {}
        self.char_freq: Counter = Counter()
        total_w = 0
        for line in path.read_text(encoding="utf-8").splitlines():
            parts = line.split(" ")
            if len(parts) < 2:
                continue
            w, fr = parts[0], parts[1]
            try:
                fr = int(fr)
            except ValueError:
                continue
            wn = norm(w)                 # normalize the dict key into the same space as candidates
            if not wn:
                continue
            # keep the max freq if normalization collides (trad/simp variants folding together)
            if wn not in self.word_freq or fr > self.word_freq[wn]:
                self.word_freq[wn] = fr
            total_w += fr
            for ch in wn:
                self.char_freq[ch] += fr
        self.total_word = total_w
        self.total_char = sum(self.char_freq.values())

    def word_rel(self, ngram_norm: str) -> float:
        """P_general(ngram) as a word (0 if OOV)."""
        return self.word_freq.get(ngram_norm, 0) / self.total_word if self.total_word else 0.0

    def char_indep_rel(self, ngram_norm: str) -> float:
        """Expected P_general(ngram) under char-independence (smoothing for OOV multiword; §A1)."""
        p = 1.0
        for ch in ngram_norm:
            pc = (self.char_freq.get(ch, 0) + 1) / (self.total_char + len(self.char_freq))
            p *= pc
        return p

    def char_rarity(self, ngram_norm: str) -> float:
        """Mean surprisal of constituent chars in general zh (rare chars like 蛊/窍 → high)."""
        if not ngram_norm:
            return 0.0
        s = 0.0
        for ch in ngram_norm:
            pc = (self.char_freq.get(ch, 0) + 1) / (self.total_char + len(self.char_freq))
            s += -math.log(pc)
        return s / len(ngram_norm)


# ── candidate generation: Han char n-grams 1–6 over normalized source (§B1) ────────────────────────
def han_runs(ntext: str):
    """Maximal runs of Han characters in normalized text."""
    runs = []
    cur = []
    for ch in ntext:
        if _RE_HAN.match(ch):
            cur.append(ch)
        else:
            if cur:
                runs.append("".join(cur))
                cur = []
    if cur:
        runs.append("".join(cur))
    return runs


def enumerate_candidates(chunks: list[Chunk], nmin=NGRAM_MIN, nmax=NGRAM_MAX) -> Counter:
    """All Han n-grams (length nmin..nmax) with overlapping counts across chunk sources.
    Annotation chunk INCLUDED (blurb terms live there; GT filter handles the blurb rule separately)."""
    cnt = Counter()
    for c in chunks:
        for run in han_runs(c.nsource):
            L = len(run)
            for n in range(nmin, nmax + 1):
                for i in range(0, L - n + 1):
                    cnt[run[i:i + n]] += 1
    return cnt


def candidate_chapters(cand_norm: str, chunks: list[Chunk]) -> set[int]:
    return {c.chapter for c in chunks if cand_norm in c.nsource}


if __name__ == "__main__":
    if "--selftest" not in sys.argv:
        print("usage: python exp16_common.py --selftest"); sys.exit(0)

    chunks = load_chunks()
    gt = load_gt()
    print(f"chunks: {len(chunks)}  (chapters {min(c.chapter for c in chunks)}–{max(c.chapter for c in chunks)})")
    print(f"GT entities: {len(gt)}  aliases: {sum(len(e.aliases) for e in gt)}")

    # GT occurrence filter (§D1): every GT term must have >=1 occurrence. Blurb rule both ways.
    print("\n== GT occurrence filter (incl. vs excl. annotation chunk) ==")
    n_pass_incl = n_pass_excl = 0
    blurb_terms = []
    for e in gt:
        oi = gt_occurrences(e, chunks, include_annotation=True)
        oe = gt_occurrences(e, chunks, include_annotation=False)
        n_pass_incl += oi >= 1
        n_pass_excl += oe >= 1
        if oi >= 1 and oe == 0:
            blurb_terms.append((e.src, e.dst, e.since_ch, oi))
    print(f"  pass (incl annotation): {n_pass_incl}/{len(gt)}")
    print(f"  pass (excl annotation): {n_pass_excl}/{len(gt)}")
    print(f"  annotation-only terms ({len(blurb_terms)}): "
          + ", ".join(f"{s}({dst},since{sc},occ{o})" for s, dst, sc, o in blurb_terms))

    # frequency strata of GT (incl annotation)
    strata = Counter(freq_stratum(gt_occurrences(e, chunks)) for e in gt)
    print(f"\n== GT freq strata (incl annotation): {dict(strata)} ==")

    # candidate space
    cands = enumerate_candidates(chunks)
    above = {k: v for k, v in cands.items() if v >= FREQ_FLOOR}
    print(f"\ncandidates: {len(cands)} distinct Han n-grams; {len(above)} at freq>={FREQ_FLOOR}")

    # catastrophe entities present + their counts
    print("\n== catastrophe screen entities ==")
    for s in CATASTROPHE:
        print(f"  {s}: source-occ={count_occurrences(norm(s), chunks)}  in-candidate-set={norm(s) in cands} "
              f"(freq {cands.get(norm(s),0)})")

    # contrast smoke
    C = Contrast()
    print(f"\ncontrast: {len(C.word_freq)} words, {len(C.char_freq)} chars, total_word={C.total_word}")
    for w in ["方源", "蛊", "蛊师", "古月", "自己", "一个", "时候"]:
        wn = norm(w)
        print(f"  {w:<4} word_rel={C.word_rel(wn):.2e} char_rarity={C.char_rarity(wn):.2f} "
              f"char_indep={C.char_indep_rel(wn):.2e}")
    print("\nselftest OK")