#!/usr/bin/env python3
"""exp15 — Q1 arm setup + size-match + cut-alignment (checkpoint item e; design §D2/§1.14).

A0 = greedy-1500 (the validated greedy port reproduces A0's 24 chunks byte-identically). The LOGICAL arm
is the §B1.5 DP chunker sized to A0's ACTUAL mean est_out (±10%), so the ONLY Q1 variable is cut
ALIGNMENT, not size (size = Q2's domain). Reports:
  - est_out distributions (greedy vs logical) + size-match verdict (±10% mean)
  - cut-alignment (secondary Q1 metric): fraction of each arm's cuts landing ON scene-shift seams
    (higher=better) vs continuation/negative seams (lower=better), and fraction of seams hit.
This is the honest "lever engaged / Q1 decidable a priori" evidence.
"""
from __future__ import annotations
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from chunker import greedy_chunks, logical_chunks, est_out, normalize_source

S2 = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime.txt")
SEAMS = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime_seams.json")
LABELS = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime_seam_labels.json")
OUT = Path("/home/ubuntu/books/gu-zhenren/exp15/q1_setup.json")
TOL = 0   # a cut "hits" a seam iff exact rune offset match (both are line/sentence boundaries)


def greedy_cut_offsets(raw):
    """Rune offsets (in NORMALIZED text) where each greedy chunk starts (excl 0)."""
    chs = greedy_chunks(raw)
    offs, cur = [], 0
    for c in chs:
        if cur > 0:
            offs.append(cur)
        cur += len(c)
    return chs, offs


def main():
    raw = S2.read_text(encoding="utf-8")
    norm = normalize_source(raw)
    # NOTE: greedy_chunks re-joins paragraphs with "\n\n" and trims; its cut offsets are in the
    # concatenated-chunk space. For seam alignment we need offsets in the SAME space as the seams
    # (the normalized S2'). We recompute cut offsets by locating each chunk's head in `norm`.
    g_chunks = greedy_chunks(raw)
    g_offs = _locate_starts(g_chunks, norm)
    g_out = [est_out(c) for c in g_chunks]
    mean_g = sum(g_out) / len(g_out)

    # logical arm sized to A0 mean est_out
    plan = logical_chunks(norm, target_out=mean_g, hard_max_out=max(g_out))
    l_chunks = [c[0] for c in plan.chunks]
    l_offs = _locate_starts(l_chunks, norm)
    l_out = [c[1] for c in plan.chunks]
    mean_l = sum(l_out) / len(l_out)

    # seams (already aligned to normalized/current S2'? seams were computed on the LF text = S2'.txt;
    # norm = S2'.txt stripped of leading/trailing ws. leading strip shifts offsets by the stripped prefix.)
    strip_shift = raw.find(norm[:10])   # how many runes were stripped from the front
    seams = json.load(open(SEAMS))["seams"]
    labels = json.load(open(LABELS))["labels"]
    seam_pts = [{"id": s["seam_id"], "off": s["rune_offset"] - strip_shift,
                 "class": labels[str(s["seam_id"])]["class"]} for s in seams]

    def hits(offs):
        offset_set = set(offs)
        res = {"scene-shift": 0, "continuation": 0, "seams_hit": []}
        for sp in seam_pts:
            hit = any(abs(o - sp["off"]) <= TOL for o in offs)
            if hit:
                res[sp["class"]] += 1
                res["seams_hit"].append(sp["id"])
        return res

    g_hits, l_hits = hits(g_offs), hits(l_offs)
    n_scene = sum(1 for s in seam_pts if s["class"] == "scene-shift")
    n_cont = sum(1 for s in seam_pts if s["class"] == "continuation")

    lo, hi = mean_g * 0.9, mean_g * 1.1
    result = {
        "greedy_A0": {"n_chunks": len(g_chunks), "n_cuts": len(g_offs), "mean_est_out": round(mean_g, 1),
                      "min": round(min(g_out), 1), "max": round(max(g_out), 1), "cut_offsets": g_offs},
        "logical": {"n_chunks": len(l_chunks), "n_cuts": len(l_offs), "mean_est_out": round(mean_l, 1),
                    "min": round(min(l_out), 1), "max": round(max(l_out), 1), "cut_offsets": l_offs,
                    "right_boundary_classes": [c[2] for c in plan.chunks[:-1]]},
        "size_match": {"target": round(mean_g, 1), "logical_mean": round(mean_l, 1),
                       "within_10pct": lo <= mean_l <= hi, "band": [round(lo, 1), round(hi, 1)]},
        "seams": {"n_scene_shift": n_scene, "n_continuation": n_cont},
        "cut_alignment": {
            "greedy": {"scene_shift_hits": g_hits["scene-shift"], "continuation_hits": g_hits["continuation"],
                       "seams_hit": g_hits["seams_hit"]},
            "logical": {"scene_shift_hits": l_hits["scene-shift"], "continuation_hits": l_hits["continuation"],
                        "seams_hit": l_hits["seams_hit"]},
            "interpretation": "logical should hit MORE scene-shift seams and FEWER continuation seams than greedy",
        },
    }
    OUT.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"greedy A0: {len(g_chunks)} chunks, mean est_out {mean_g:.0f}")
    print(f"logical:   {len(l_chunks)} chunks, mean est_out {mean_l:.0f}  size-match(±10%)={result['size_match']['within_10pct']}")
    print(f"\nseams: {n_scene} scene-shift, {n_cont} continuation")
    print(f"cut-alignment (exact rune hits):")
    print(f"  greedy : scene-shift {g_hits['scene-shift']}/{n_scene}  continuation {g_hits['continuation']}/{n_cont}  seams_hit={g_hits['seams_hit']}")
    print(f"  logical: scene-shift {l_hits['scene-shift']}/{n_scene}  continuation {l_hits['continuation']}/{n_cont}  seams_hit={l_hits['seams_hit']}")
    print(f"-> {OUT}")


def _locate_starts(chunks, norm):
    """Rune offset in `norm` where each chunk (except the first) starts, via its head."""
    offs, pos = [], 0
    for i, c in enumerate(chunks):
        head = c[:25]
        idx = norm.find(head, max(0, pos - 5))
        if idx < 0:
            idx = norm.find(head)
        if i > 0 and idx >= 0:
            offs.append(idx)
        pos = idx + len(c) if idx >= 0 else pos + len(c)
    return offs


if __name__ == "__main__":
    main()
